Added flags for using local images in sandbox (flyteorg#216)

* Added flags for local images in sandbox Signed-off-by: Yuvraj <[email protected]> * debug Signed-off-by: Yuvraj <[email protected]> * Added pull policy in sandbox Signed-off-by: Yuvraj <[email protected]> * fix pflag manually Signed-off-by: Yuvraj <[email protected]> * more changes Signed-off-by: Yuvraj <[email protected]> * Allow setting ImagePullPolicy in cmd line (flyteorg#218) Signed-off-by: Haytham Abuelfutuh <[email protected]> * fix test Signed-off-by: Yuvraj <[email protected]> * Fix docs for ImagePullPolicy values Signed-off-by: Haytham Abuelfutuh <[email protected]> Co-authored-by: Haytham Abuelfutuh <[email protected]>
eapolinario · Nov 24, 2021 · be931d7 · be931d7
1 parent 68db2d5
commit be931d7
Show file tree

Hide file tree

Showing 9 changed files with 208 additions and 11 deletions.
diff --git a/cmd/config/subcommand/sandbox/config_flags.go b/cmd/config/subcommand/sandbox/config_flags.go
diff --git a/cmd/config/subcommand/sandbox/config_flags_test.go b/cmd/config/subcommand/sandbox/config_flags_test.go
diff --git a/cmd/config/subcommand/sandbox/imagepullpolicy_enumer.go b/cmd/config/subcommand/sandbox/imagepullpolicy_enumer.go
diff --git a/cmd/config/subcommand/sandbox/sandbox_config.go b/cmd/config/subcommand/sandbox/sandbox_config.go
@@ -1,11 +1,36 @@
 package sandbox
 
+//go:generate enumer -type=ImagePullPolicy -trimprefix=ImagePullPolicy --json
+type ImagePullPolicy int
+
+const (
+	ImagePullPolicyAlways ImagePullPolicy = iota
+	ImagePullPolicyIfNotPresent
+	ImagePullPolicyNever
+)
+
+// Set implements PFlag's Value interface to attempt to set the value of the flag from string.
+func (i *ImagePullPolicy) Set(val string) error {
+	policy, err := ImagePullPolicyString(val)
+	if err != nil {
+		return err
+	}
+
+	*i = policy
+	return nil
+}
+
+// Type implements PFlag's Value interface to return type name.
+func (i ImagePullPolicy) Type() string {
+	return "ImagePullPolicy"
+}
+
 //go:generate pflags Config --default-var DefaultConfig --bind-default-var
 var (
 	DefaultConfig = &Config{}
 )
 
-//Config
+//Config holds configuration flags for sandbox command.
 type Config struct {
 	Source string `json:"source" pflag:",Path of your source code"`
 
@@ -18,4 +43,8 @@ type Config struct {
 	// 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."`
+
+	// Optionally it is possible to use local sandbox image
+	// If local flag pass then flytectl will not pull image from registry. Usually useful, if you want to test your local images without pushing them to a registry
+	ImagePullPolicy ImagePullPolicy `json:"imagePullPolicy" pflag:",Optional. Defines the image pull behavior [Always/IfNotPresent/Never]"`
 }
diff --git a/cmd/sandbox/start.go b/cmd/sandbox/start.go
@@ -56,6 +56,11 @@ Specify a Flyte Sandbox compliant image with the registry. This is useful, in ca
 
   flytectl sandbox start --image docker.io/my-override:latest
 
+	
+Specify a Flyte Sandbox image pull policy. Possible pull policy values are Always, IfNotPresent, or Never
+::
+
+ flytectl sandbox start  --image docker.io/my-override:latest --imagePullPolicy Always
 Usage
 `
 	k8sEndpoint             = "https://127.0.0.1:30086"
@@ -143,7 +148,8 @@ func startSandbox(ctx context.Context, cli docker.Docker, reader io.Reader) (*bu
 		return nil, err
 	}
 	fmt.Printf("%v pulling docker image for release %s\n", emoji.Whale, image)
-	if err := docker.PullDockerImage(ctx, cli, image); err != nil {
+
+	if err := docker.PullDockerImage(ctx, cli, image, sandboxConfig.DefaultConfig.ImagePullPolicy); err != nil {
 		return nil, err
 	}
 

diff --git a/pkg/docker/docker.go b/pkg/docker/docker.go
@@ -24,6 +24,7 @@ type Docker interface {
 	ContainerExecCreate(ctx context.Context, container string, config types.ExecConfig) (types.IDResponse, error)
 	ContainerExecAttach(ctx context.Context, execID string, config types.ExecStartCheck) (types.HijackedResponse, error)
 	ContainerExecInspect(ctx context.Context, execID string) (types.ContainerExecInspect, error)
+	ImageList(ctx context.Context, listOption types.ImageListOptions) ([]types.ImageSummary, error)
 }
 
 type FlyteDocker struct {

diff --git a/pkg/docker/docker_util.go b/pkg/docker/docker_util.go
@@ -9,6 +9,8 @@ import (
 	"os"
 	"strings"
 
+	sandboxConfig "github.com/flyteorg/flytectl/cmd/config/subcommand/sandbox"
+
 	"github.com/flyteorg/flytectl/clierrors"
 
 	"github.com/docker/docker/api/types"
@@ -88,14 +90,30 @@ func GetSandboxPorts() (map[nat.Port]struct{}, map[nat.Port][]nat.PortBinding, e
 }
 
 // PullDockerImage will Pull docker image
-func PullDockerImage(ctx context.Context, cli Docker, image string) error {
-	r, err := cli.ImagePull(ctx, image, types.ImagePullOptions{})
-	if err != nil {
+func PullDockerImage(ctx context.Context, cli Docker, image string, pullPolicy sandboxConfig.ImagePullPolicy) error {
+	if pullPolicy == sandboxConfig.ImagePullPolicyAlways || pullPolicy == sandboxConfig.ImagePullPolicyIfNotPresent {
+		if pullPolicy == sandboxConfig.ImagePullPolicyIfNotPresent {
+			imageSummary, err := cli.ImageList(ctx, types.ImageListOptions{})
+			if err != nil {
+				return err
+			}
+			for _, img := range imageSummary {
+				for _, tags := range img.RepoTags {
+					if image == tags {
+						return nil
+					}
+				}
+			}
+		}
+		r, err := cli.ImagePull(ctx, image, types.ImagePullOptions{})
+		if err != nil {
+			return err
+		}
+
+		_, err = io.Copy(os.Stdout, r)
 		return err
 	}
-
-	_, err = io.Copy(os.Stdout, r)
-	return err
+	return nil
 }
 
 //StartContainer will create and start docker container

diff --git a/pkg/docker/docker_util_test.go b/pkg/docker/docker_util_test.go
@@ -8,6 +8,8 @@ import (
 	"strings"
 	"testing"
 
+	sandboxConfig "github.com/flyteorg/flytectl/cmd/config/subcommand/sandbox"
+
 	f "github.com/flyteorg/flytectl/pkg/filesystemutils"
 
 	"github.com/docker/docker/api/types/container"
@@ -101,13 +103,13 @@ func TestRemoveSandboxWithNoReply(t *testing.T) {
 }
 
 func TestPullDockerImage(t *testing.T) {
-	t.Run("Successfully pull image", func(t *testing.T) {
+	t.Run("Successfully pull image Always", func(t *testing.T) {
 		setupSandbox()
 		mockDocker := &mocks.Docker{}
 		context := context.Background()
 		// Verify the attributes
 		mockDocker.OnImagePullMatch(context, mock.Anything, types.ImagePullOptions{}).Return(os.Stdin, nil)
-		err := PullDockerImage(context, mockDocker, "nginx")
+		err := PullDockerImage(context, mockDocker, "nginx:latest", sandboxConfig.ImagePullPolicyAlways)
 		assert.Nil(t, err)
 	})
 
@@ -117,10 +119,28 @@ func TestPullDockerImage(t *testing.T) {
 		context := context.Background()
 		// Verify the attributes
 		mockDocker.OnImagePullMatch(context, mock.Anything, types.ImagePullOptions{}).Return(os.Stdin, fmt.Errorf("error"))
-		err := PullDockerImage(context, mockDocker, "nginx")
+		err := PullDockerImage(context, mockDocker, "nginx:latest", sandboxConfig.ImagePullPolicyAlways)
 		assert.NotNil(t, err)
 	})
 
+	t.Run("Successfully pull image IfNotPresent", func(t *testing.T) {
+		setupSandbox()
+		mockDocker := &mocks.Docker{}
+		context := context.Background()
+		// Verify the attributes
+		mockDocker.OnImagePullMatch(context, mock.Anything, types.ImagePullOptions{}).Return(os.Stdin, nil)
+		mockDocker.OnImageListMatch(context, types.ImageListOptions{}).Return([]types.ImageSummary{}, nil)
+		err := PullDockerImage(context, mockDocker, "nginx:latest", sandboxConfig.ImagePullPolicyIfNotPresent)
+		assert.Nil(t, err)
+	})
+
+	t.Run("Successfully pull image Never", func(t *testing.T) {
+		setupSandbox()
+		mockDocker := &mocks.Docker{}
+		context := context.Background()
+		err := PullDockerImage(context, mockDocker, "nginx:latest", sandboxConfig.ImagePullPolicyNever)
+		assert.Nil(t, err)
+	})
 }
 
 func TestStartContainer(t *testing.T) {

diff --git a/pkg/docker/mocks/docker.go b/pkg/docker/mocks/docker.go