Flyte Add Enable/Disable Options for Agent Services (flyteorg#420)

--------- Signed-off-by: Future Outlier <[email protected]> Signed-off-by: Future-Outlier <[email protected]> Co-authored-by: Future Outlier <[email protected]> Co-authored-by: Kevin Su <[email protected]>
eapolinario · Oct 17, 2023 · 64e0d7b · 64e0d7b
1 parent 1152c5f
commit 64e0d7b
Show file tree

Hide file tree

Showing 5 changed files with 25 additions and 0 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/sandbox_config.go b/cmd/config/subcommand/sandbox/sandbox_config.go
@@ -19,6 +19,9 @@ type Config struct {
 	// Default value false represents that Flytectl will not use the latest pre-release if it exists.
 	Prerelease bool `json:"pre" pflag:",Optional. Pre release Version of flyte will be used for sandbox."`
 
+	// Agent Service
+	DisableAgent bool `json:"disable-agent" pflag:",Optional. Disable the agent service."`
+
 	// Optionally it is possible to pass in environment variables to sandbox container.
 	Env []string `json:"env" pflag:",Optional. Provide Env variable in key=value format which can be passed to sandbox container."`
 

diff --git a/pkg/sandbox/start.go b/pkg/sandbox/start.go
@@ -229,6 +229,10 @@ func startSandbox(ctx context.Context, cli docker.Docker, g github.GHRepoService
 		sandboxEnv = append(sandboxEnv, "FLYTE_DEV=True")
 	}
 
+	if sandboxConfig.DisableAgent {
+		sandboxEnv = append(sandboxEnv, "DISABLE_AGENT=True")
+	}
+
 	ID, err := docker.StartContainer(ctx, cli, volumes, exposedPorts, portBindings, docker.FlyteSandboxClusterName,
 		sandboxImage, sandboxEnv, sandboxConfig.DryRun)
 
@@ -393,6 +397,7 @@ func StartDemoCluster(ctx context.Context, args []string, sandboxConfig *sandbox
 		return err
 	}
 	// K3s will automatically write the file specified by this var, which is mounted from user's local state dir.
+
 	sandboxConfig.Env = append(sandboxConfig.Env, k3sKubeConfigEnvVar)
 	err = StartCluster(ctx, args, sandboxConfig, demoImageName, sandboxImagePrefix, exposedPorts, portBindings, util.DemoConsolePort)
 	if err != nil {

diff --git a/pkg/sandbox/start_test.go b/pkg/sandbox/start_test.go
@@ -106,6 +106,8 @@ func TestStartFunc(t *testing.T) {
 		RegistryAuth: "",
 		Platform:     "",
 	}
+	config.Dev = true
+	config.DisableAgent = true
 	assert.Nil(t, util.SetupFlyteDir())
 	assert.Nil(t, os.MkdirAll(f.FilePathJoin(f.UserHomeDir(), ".flyte", "state"), os.ModePerm))
 	assert.Nil(t, ioutil.WriteFile(docker.Kubeconfig, []byte(content), os.ModePerm))