flytectl sandbox start --image xyz (#195)

* flytectl sandbox start --image xyz Sandbox can now use a custom image Signed-off-by: Ketan Umare <[email protected]> * fixed unit test Signed-off-by: Ketan Umare <[email protected]> * updated comments Signed-off-by: Ketan Umare <[email protected]>
flyteorg · Oct 6, 2021 · 6f9a2aa · 6f9a2aa
1 parent bdf47b8
commit 6f9a2aa
Show file tree

Hide file tree

Showing 5 changed files with 56 additions and 19 deletions.
diff --git a/flytectl/cmd/config/subcommand/sandbox/config_flags.go b/flytectl/cmd/config/subcommand/sandbox/config_flags.go
diff --git a/flytectl/cmd/config/subcommand/sandbox/config_flags_test.go b/flytectl/cmd/config/subcommand/sandbox/config_flags_test.go
diff --git a/flytectl/cmd/config/subcommand/sandbox/sandbox_config.go b/flytectl/cmd/config/subcommand/sandbox/sandbox_config.go
@@ -9,8 +9,13 @@ var (
 type Config struct {
 	Source string `json:"source" pflag:",Path of your source code"`
 
-	// Flytectl sandbox only support flyte version available in Github release https://github.com/flyteorg/flyte/tags
+	// Flytectl sandbox only supports flyte version available in Github release https://github.com/flyteorg/flyte/tags
 	// Flytectl sandbox will only work for v0.10.0+
-	// Default value dind represent the latest release
-	Version string `json:"version" pflag:",Version of flyte. Only support v0.10.0+ flyte release"`
+	// Default value dind represents the latest release
+	Version string `json:"version" pflag:",Version of flyte. Only supports flyte releases greater than v0.10.0"`
+
+	// Optionally it is possible to specify a specific fqn for the docker image with the tag. This should be
+	// Flyte compliant sandbox image. Usually useful, if you want to push the image to your own registry and relaunch
+	// from there.
+	Image string `json:"image" pflag:",Optional. Provide a fully qualified path to a Flyte compliant docker image."`
 }
diff --git a/flytectl/cmd/sandbox/start.go b/flytectl/cmd/sandbox/start.go
@@ -37,22 +37,27 @@ The Flyte Sandbox is a fully standalone minimal environment for running Flyte. p
 Start sandbox cluster without any source code
 ::
 
- bin/flytectl sandbox start
+ flytectl sandbox start
 	
 Mount your source code repository inside sandbox 
 ::
 
- bin/flytectl sandbox start --source=$HOME/flyteorg/flytesnacks 
+ flytectl sandbox start --source=$HOME/flyteorg/flytesnacks 
 	
 Run specific version of flyte. flytectl sandbox only support flyte version available in Github release https://github.com/flyteorg/flyte/tags
 ::
 
- bin/flytectl sandbox start  --version=v0.14.0
+ flytectl sandbox start  --version=v0.14.0
+
+Note: Flytectl sandbox is only supported for Flyte versions > v0.10.0
+
+Specify a Flyte Sandbox compliant image with the registry. This is useful, in case you want to use an image from your registry.
+::
+
+  flytectl sandbox start --image docker.io/my-override:latest
 
-Note: Flytectl sandbox will only work for v0.10.0+
-	
 Usage
-	`
+`
 	k8sEndpoint             = "https://127.0.0.1:30086"
 	flyteNamespace          = "flyte"
 	flyteRepository         = "flyte"
@@ -133,7 +138,7 @@ func startSandbox(ctx context.Context, cli docker.Docker, reader io.Reader) (*bu
 		volumes = append(volumes, *vol)
 	}
 
-	image, err := getSandboxImage(sandboxConfig.DefaultConfig.Version)
+	image, err := getSandboxImage(sandboxConfig.DefaultConfig.Version, sandboxConfig.DefaultConfig.Image)
 	if err != nil {
 		return nil, err
 	}
@@ -158,9 +163,14 @@ func startSandbox(ctx context.Context, cli docker.Docker, reader io.Reader) (*bu
 	return logReader, nil
 }
 
-func getSandboxImage(version string) (string, error) {
-	// Latest release will use image cr.flyte.org/flyteorg/flyte-sandbox:dind
-	// In case of version flytectl will use cr.flyte.org/flyteorg/flyte-sandbox:dind-{SHA}
+// Returns the alternate image if specified, else
+// if no version is specified then the Latest release of cr.flyte.org/flyteorg/flyte-sandbox:dind is used
+// else cr.flyte.org/flyteorg/flyte-sandbox:dind-{SHA}, where sha is derived from the version.
+func getSandboxImage(version string, alternateImage string) (string, error) {
+
+	if len(alternateImage) > 0 {
+		return alternateImage, nil
+	}
 
 	var tag = dind
 	if len(version) > 0 {

diff --git a/flytectl/cmd/sandbox/start_test.go b/flytectl/cmd/sandbox/start_test.go
@@ -678,26 +678,33 @@ func TestGetNodeTaintStatus(t *testing.T) {
 
 func TestGetSandboxImage(t *testing.T) {
 	t.Run("Get Latest sandbox", func(t *testing.T) {
-		image, err := getSandboxImage("")
+		image, err := getSandboxImage("", "")
 		assert.Nil(t, err)
 		assert.Equal(t, docker.GetSandboxImage(dind), image)
 	})
 
 	t.Run("Get sandbox image with version ", func(t *testing.T) {
-		image, err := getSandboxImage("v0.14.0")
+		image, err := getSandboxImage("v0.14.0", "")
 		assert.Nil(t, err)
 		assert.Equal(t, true, strings.HasPrefix(image, docker.ImageName))
 	})
 	t.Run("Get sandbox image with wrong version ", func(t *testing.T) {
-		_, err := getSandboxImage("v100.1.0")
+		_, err := getSandboxImage("v100.1.0", "")
 		assert.NotNil(t, err)
 	})
 	t.Run("Get sandbox image with wrong version ", func(t *testing.T) {
-		_, err := getSandboxImage("aaaaaa")
+		_, err := getSandboxImage("aaaaaa", "")
 		assert.NotNil(t, err)
 	})
 	t.Run("Get sandbox image with version that is not supported", func(t *testing.T) {
-		_, err := getSandboxImage("v0.10.0")
+		_, err := getSandboxImage("v0.10.0", "")
 		assert.NotNil(t, err)
 	})
+
+	t.Run("Get sandbox image with version that is not supported", func(t *testing.T) {
+		img := "docker.io/my-override:latest"
+		i, err := getSandboxImage("v0.11.0", img)
+		assert.Nil(t, err)
+		assert.Equal(t, i, img)
+	})
 }