Support new sandbox-bundled features (#381)

* Create a docker volume to persist database and object store Signed-off-by: Jeev B <[email protected]> * Cleanup sandbox configuration directory and mount Signed-off-by: Jeev B <[email protected]> * Add logic for backward-compatible demo reload Signed-off-by: Jeev B <[email protected]> * Fix working directory for exec Signed-off-by: Jeev B <[email protected]> * Generate mocks Signed-off-by: Jeev B <[email protected]> * Get existing tests passing Signed-off-by: Jeev B <[email protected]> * Cleanup volume creation and add a dryRun message Signed-off-by: Jeev B <[email protected]> * Use fmt.Errorf in place of errors.New Signed-off-by: Jeev B <[email protected]> * Fix build Signed-off-by: Jeev B <[email protected]> * Add tests for volume creation Signed-off-by: Jeev B <[email protected]> * Add test for volume teardown Signed-off-by: Jeev B <[email protected]> * Include code path for volume creation in sandbox start test Signed-off-by: Jeev B <[email protected]> * Add tests for demo reload with backward compatibility Signed-off-by: Jeev B <[email protected]> * Suppress output of `which` when testing for sandbox version during demo reload --------- Signed-off-by: Jeev B <[email protected]>
flyteorg · Feb 9, 2023 · 1810234 · 1810234
1 parent 049be2b
commit 1810234
Show file tree

Hide file tree

Showing 16 changed files with 427 additions and 44 deletions.
diff --git a/cmd/config/subcommand/sandbox/teardown.go b/cmd/config/subcommand/sandbox/teardown.go
@@ -0,0 +1,21 @@
+package sandbox
+
+import (
+	"fmt"
+
+	"github.com/spf13/pflag"
+)
+
+type TeardownFlags struct {
+	Volume bool
+}
+
+var (
+	DefaultTeardownFlags = &TeardownFlags{}
+)
+
+func (f *TeardownFlags) GetPFlagSet(prefix string) *pflag.FlagSet {
+	cmdFlags := pflag.NewFlagSet("TeardownFlags", pflag.ExitOnError)
+	cmdFlags.BoolVarP(&f.Volume, fmt.Sprintf("%v%v", prefix, "volume"), "v", f.Volume, "Optional. Clean up Docker volume. This will result in a permanent loss of all data within the database and object store. Use with caution!")
+	return cmdFlags
+}
diff --git a/cmd/demo/demo.go b/cmd/demo/demo.go
@@ -56,8 +56,10 @@ func CreateDemoCommand() *cobra.Command {
 			Short: reloadShort,
 			Long:  reloadLong, PFlagProvider: sandboxCmdConfig.DefaultConfig, DisableFlyteClient: true},
 		"teardown": {CmdFunc: teardownDemoCluster, Aliases: []string{}, ProjectDomainNotRequired: true,
-			Short: teardownShort,
-			Long:  teardownLong, DisableFlyteClient: true},
+			Short:              teardownShort,
+			Long:               teardownLong,
+			PFlagProvider:      sandboxCmdConfig.DefaultTeardownFlags,
+			DisableFlyteClient: true},
 		"status": {CmdFunc: demoClusterStatus, Aliases: []string{}, ProjectDomainNotRequired: true,
 			Short: statusShort,
 			Long:  statusLong},

diff --git a/cmd/demo/reload.go b/cmd/demo/reload.go
@@ -12,7 +12,8 @@ import (
 )
 
 const (
-	labelSelector = "app=flyte"
+	internalBootstrapAgent = "flyte-sandbox-bootstrap"
+	labelSelector          = "app.kubernetes.io/name=flyte-binary"
 )
 const (
 	reloadShort = "Power cycle the Flyte executable pod, effectively picking up an updated config."
@@ -28,8 +29,68 @@ Usage
 `
 )
 
-// reloadDemoCluster will kill the flyte binary pod so the new one can pick up a new config file
+func isLegacySandbox(ctx context.Context, cli docker.Docker, containerID string) (bool, error) {
+	var result bool
+
+	// Check if sandbox is compatible with new bootstrap mechanism
+	exec, err := docker.ExecCommend(
+		ctx,
+		cli,
+		containerID,
+		[]string{"sh", "-c", fmt.Sprintf("which %s > /dev/null", internalBootstrapAgent)},
+	)
+	if err != nil {
+		return result, err
+	}
+	if err = docker.InspectExecResp(ctx, cli, exec.ID); err != nil {
+		return result, err
+	}
+	res, err := cli.ContainerExecInspect(ctx, exec.ID)
+	if err != nil {
+		return result, err
+	}
+
+	result = res.ExitCode != 0
+	return result, nil
+}
+
 func reloadDemoCluster(ctx context.Context, args []string, cmdCtx cmdCore.CommandContext) error {
+	cli, err := docker.GetDockerClient()
+	if err != nil {
+		return err
+	}
+	c, err := docker.GetSandbox(ctx, cli)
+	if err != nil {
+		return err
+	}
+	if c == nil {
+		return fmt.Errorf("reload failed - could not find an active sandbox")
+	}
+
+	// Working with a legacy sandbox - fallback to legacy reload mechanism
+	useLegacyMethod, err := isLegacySandbox(ctx, cli, c.ID)
+	if err != nil {
+		return err
+	}
+	if useLegacyMethod {
+		return legacyReloadDemoCluster(ctx)
+	}
+
+	// At this point we know that we are on a modern sandbox, and we can use the
+	// internal bootstrap agent to reload the cluster
+	exec, err := docker.ExecCommend(ctx, cli, c.ID, []string{internalBootstrapAgent})
+	if err != nil {
+		return err
+	}
+	if err = docker.InspectExecResp(ctx, cli, exec.ID); err != nil {
+		return err
+	}
+
+	return nil
+}
+
+// legacyReloadDemoCluster will kill the flyte binary pod so the new one can pick up a new config file
+func legacyReloadDemoCluster(ctx context.Context) error {
 	k8sClient, err := k8s.GetK8sClient(docker.Kubeconfig, K8sEndpoint)
 	if err != nil {
 		fmt.Println("Could not get K8s client")

diff --git a/cmd/demo/reload_test.go b/cmd/demo/reload_test.go
@@ -1,10 +1,16 @@
 package demo
 
 import (
+	"bufio"
+	"bytes"
 	"context"
+	"fmt"
 	"testing"
 
+	"github.com/docker/docker/api/types"
 	cmdCore "github.com/flyteorg/flytectl/cmd/core"
+	"github.com/flyteorg/flytectl/pkg/docker"
+	"github.com/flyteorg/flytectl/pkg/docker/mocks"
 	"github.com/flyteorg/flytectl/pkg/k8s"
 	"github.com/stretchr/testify/assert"
 	corev1 "k8s.io/api/core/v1"
@@ -20,14 +26,76 @@ var fakePod = corev1.Pod{
 	},
 	ObjectMeta: metav1.ObjectMeta{
 		Name:   "dummyflytepod",
-		Labels: map[string]string{"app": "flyte"},
+		Labels: map[string]string{"app.kubernetes.io/name": "flyte-binary"},
 	},
 }
 
-func TestDemoReload(t *testing.T) {
+func sandboxSetup(ctx context.Context, legacy bool) {
+	mockDocker := &mocks.Docker{}
+	docker.Client = mockDocker
+	mockDocker.OnContainerList(ctx, types.ContainerListOptions{All: true}).Return([]types.Container{
+		{
+			ID: docker.FlyteSandboxClusterName,
+			Names: []string{
+				docker.FlyteSandboxClusterName,
+			},
+		},
+	}, nil)
+
+	// This first set of mocks is for the check for the bootstrap agent. This is
+	// Expected to fail in legacy sandboxes
+	var checkLegacySandboxExecExitCode int
+	if legacy {
+		checkLegacySandboxExecExitCode = 1
+	}
+	mockDocker.OnContainerExecCreateMatch(
+		ctx,
+		docker.FlyteSandboxClusterName,
+		types.ExecConfig{
+			AttachStderr: true,
+			Tty:          true,
+			WorkingDir:   "/",
+			AttachStdout: true,
+			Cmd:          []string{"sh", "-c", fmt.Sprintf("which %s > /dev/null", internalBootstrapAgent)},
+		},
+	).Return(types.IDResponse{ID: "0"}, nil)
+	mockDocker.OnContainerExecAttachMatch(ctx, "0", types.ExecStartCheck{}).Return(types.HijackedResponse{
+		Reader: bufio.NewReader(bytes.NewReader([]byte{})),
+	}, nil)
+	mockDocker.OnContainerExecInspectMatch(ctx, "0").Return(types.ContainerExecInspect{ExitCode: checkLegacySandboxExecExitCode}, nil)
+
+	// Register additional mocks for the actual execution of the bootstrap agent
+	// in non-legacy sandboxes
+	if !legacy {
+		mockDocker.OnContainerExecCreateMatch(
+			ctx,
+			docker.FlyteSandboxClusterName,
+			types.ExecConfig{
+				AttachStderr: true,
+				Tty:          true,
+				WorkingDir:   "/",
+				AttachStdout: true,
+				Cmd:          []string{internalBootstrapAgent},
+			},
+		).Return(types.IDResponse{ID: "1"}, nil)
+		mockDocker.OnContainerExecAttachMatch(ctx, "1", types.ExecStartCheck{}).Return(types.HijackedResponse{
+			Reader: bufio.NewReader(bytes.NewReader([]byte{})),
+		}, nil)
+	}
+}
+
+func TestReloadLegacy(t *testing.T) {
 	ctx := context.Background()
 	commandCtx := cmdCore.CommandContext{}
+	sandboxSetup(ctx, false)
+	err := reloadDemoCluster(ctx, []string{}, commandCtx)
+	assert.Nil(t, err)
+}
 
+func TestDemoReloadLegacy(t *testing.T) {
+	ctx := context.Background()
+	commandCtx := cmdCore.CommandContext{}
+	sandboxSetup(ctx, true)
 	t.Run("No errors", func(t *testing.T) {
 		client := testclient.NewSimpleClientset()
 		_, err := client.CoreV1().Pods("flyte").Create(ctx, &fakePod, v1.CreateOptions{})

diff --git a/cmd/demo/teardown.go b/cmd/demo/teardown.go
@@ -3,10 +3,10 @@ package demo
 import (
 	"context"
 
-	"github.com/flyteorg/flytectl/pkg/sandbox"
-
 	"github.com/flyteorg/flytectl/pkg/docker"
+	"github.com/flyteorg/flytectl/pkg/sandbox"
 
+	sandboxCmdConfig "github.com/flyteorg/flytectl/cmd/config/subcommand/sandbox"
 	cmdCore "github.com/flyteorg/flytectl/cmd/core"
 )
 
@@ -28,5 +28,5 @@ func teardownDemoCluster(ctx context.Context, args []string, cmdCtx cmdCore.Comm
 	if err != nil {
 		return err
 	}
-	return sandbox.Teardown(ctx, cli)
+	return sandbox.Teardown(ctx, cli, sandboxCmdConfig.DefaultTeardownFlags)
 }
diff --git a/cmd/demo/teardown_test.go b/cmd/demo/teardown_test.go
@@ -6,6 +6,7 @@ import (
 	"testing"
 
 	"github.com/docker/docker/api/types"
+	sandboxCmdConfig "github.com/flyteorg/flytectl/cmd/config/subcommand/sandbox"
 	"github.com/flyteorg/flytectl/cmd/testutils"
 	"github.com/flyteorg/flytectl/pkg/configutil"
 	"github.com/flyteorg/flytectl/pkg/docker"
@@ -29,31 +30,47 @@ func TestTearDownFunc(t *testing.T) {
 	}
 	containers = append(containers, container1)
 
-	t.Run("Success", func(t *testing.T) {
+	t.Run("SuccessKeepVolume", func(t *testing.T) {
 		ctx := context.Background()
 		mockDocker := &mocks.Docker{}
 		mockDocker.OnContainerList(ctx, types.ContainerListOptions{All: true}).Return(containers, nil)
 		mockDocker.OnContainerRemove(ctx, mock.Anything, types.ContainerRemoveOptions{Force: true}).Return(nil)
 		mockK8sContextMgr := &k8sMocks.ContextOps{}
 		k8s.ContextMgr = mockK8sContextMgr
 		mockK8sContextMgr.OnRemoveContextMatch(mock.Anything).Return(nil)
-		err := sandbox.Teardown(ctx, mockDocker)
+		err := sandbox.Teardown(ctx, mockDocker, sandboxCmdConfig.DefaultTeardownFlags)
 		assert.Nil(t, err)
 	})
-	t.Run("Error", func(t *testing.T) {
+	t.Run("SuccessRemoveVolume", func(t *testing.T) {
+		ctx := context.Background()
+		mockDocker := &mocks.Docker{}
+		mockDocker.OnContainerList(ctx, types.ContainerListOptions{All: true}).Return(containers, nil)
+		mockDocker.OnContainerRemove(ctx, mock.Anything, types.ContainerRemoveOptions{Force: true}).Return(nil)
+		mockDocker.OnVolumeRemove(ctx, docker.FlyteSandboxVolumeName, true).Return(nil)
+		mockK8sContextMgr := &k8sMocks.ContextOps{}
+		k8s.ContextMgr = mockK8sContextMgr
+		mockK8sContextMgr.OnRemoveContextMatch(mock.Anything).Return(nil)
+		err := sandbox.Teardown(
+			ctx,
+			mockDocker,
+			&sandboxCmdConfig.TeardownFlags{Volume: true},
+		)
+		assert.Nil(t, err)
+	})
+	t.Run("ErrorOnContainerRemove", func(t *testing.T) {
 		ctx := context.Background()
 		mockDocker := &mocks.Docker{}
 		mockDocker.OnContainerList(ctx, types.ContainerListOptions{All: true}).Return(containers, nil)
 		mockDocker.OnContainerRemove(ctx, mock.Anything, types.ContainerRemoveOptions{Force: true}).Return(fmt.Errorf("err"))
-		err := sandbox.Teardown(ctx, mockDocker)
+		err := sandbox.Teardown(ctx, mockDocker, sandboxCmdConfig.DefaultTeardownFlags)
 		assert.NotNil(t, err)
 	})
 
-	t.Run("Error", func(t *testing.T) {
+	t.Run("ErrorOnContainerList", func(t *testing.T) {
 		ctx := context.Background()
 		mockDocker := &mocks.Docker{}
 		mockDocker.OnContainerList(ctx, types.ContainerListOptions{All: true}).Return(nil, fmt.Errorf("err"))
-		err := sandbox.Teardown(ctx, mockDocker)
+		err := sandbox.Teardown(ctx, mockDocker, sandboxCmdConfig.DefaultTeardownFlags)
 		assert.NotNil(t, err)
 	})
 

diff --git a/cmd/sandbox/teardown.go b/cmd/sandbox/teardown.go
@@ -6,6 +6,7 @@ import (
 	"github.com/flyteorg/flytectl/pkg/docker"
 	"github.com/flyteorg/flytectl/pkg/sandbox"
 
+	sandboxCmdConfig "github.com/flyteorg/flytectl/cmd/config/subcommand/sandbox"
 	cmdCore "github.com/flyteorg/flytectl/cmd/core"
 )
 
@@ -27,5 +28,5 @@ func teardownSandboxCluster(ctx context.Context, args []string, cmdCtx cmdCore.C
 	if err != nil {
 		return err
 	}
-	return sandbox.Teardown(ctx, cli)
+	return sandbox.Teardown(ctx, cli, sandboxCmdConfig.DefaultTeardownFlags)
 }
diff --git a/pkg/docker/docker.go b/pkg/docker/docker.go
@@ -6,7 +6,9 @@ import (
 
 	"github.com/docker/docker/api/types"
 	"github.com/docker/docker/api/types/container"
+	"github.com/docker/docker/api/types/filters"
 	"github.com/docker/docker/api/types/network"
+	"github.com/docker/docker/api/types/volume"
 	"github.com/docker/docker/client"
 	specs "github.com/opencontainers/image-spec/specs-go/v1"
 )
@@ -27,6 +29,9 @@ type Docker interface {
 	ImageList(ctx context.Context, listOption types.ImageListOptions) ([]types.ImageSummary, error)
 	ContainerStatPath(ctx context.Context, containerID, path string) (types.ContainerPathStat, error)
 	CopyFromContainer(ctx context.Context, containerID, srcPath string) (io.ReadCloser, types.ContainerPathStat, error)
+	VolumeCreate(ctx context.Context, options volume.VolumeCreateBody) (types.Volume, error)
+	VolumeList(ctx context.Context, filter filters.Args) (volume.VolumeListOKBody, error)
+	VolumeRemove(ctx context.Context, volumeID string, force bool) error
 }
 
 type FlyteDocker struct {