Repository: tensorchord/envd (language: python)

Issue:
feat(CLI): Support context
## Description

context is used to switch between the backends (e.g. docker, kubernetes)


Discussion / hints:
cc @FogDong 
It seems we could save the contexts into the config file, something like:
```
$ cat config.yaml
current-context: foo
contexts:
- name: foo
  docker:
    config-path: /path/to/config.json (Defaults to ~/.docker/config.json)
- name: bar
  kubernetes:
    config-path: /path/to/kubeconfig (Defaults to ~/.kube/config)
```

We could add a new subcommand, say `context`, to manage contexts in the config.
BTW this issue reminds me that docker cli supports context now: https://docs.docker.com/engine/context/working-with-contexts/
> It seems we could save the contexts into the config file, something like:
> 
> ```
> $ cat config.yaml
> current-context: foo
> contexts:
> - name: foo
>   docker:
>     config-path: /path/to/config.json (Defaults to ~/.docker/config.json)
> - name: bar
>   kubernetes:
>     config-path: /path/to/kubeconfig (Defaults to ~/.kube/config)
> ```
> 
> We could add a new subcommand, say `context`, to manage contexts in the config.

Yeah, the design LGTM although envd does not support k8s now. We can support different docker contexts currently.
We may need to support two different context types: build and run.

The context for the build is to determine what buildkit will be used in envd. For example, buildkit container in the docker daemon `docker-container://envd_buildkitd` can be used. Or `kube-pod://<pod-name-for-buildkit>` also can be used if there is a buildkitd running on the Kubernetes.

Users may expect to use the buildkit on Kubernetes for CI.  Ref https://github.com/tensorchord/envd/issues/510

The context for the run will be used when running `envd up`. The container or pod will be created after the image is built successfully with buildkitd. We can use Docker or Kubernetes to run it. Thus we also need to support context here.
The user experience for the build may look like this:

```
envd context create --name docker-buildkit --builder docker-container://envd_buildkitd
envd context create --name kubernetes-buildkit --builder kube-pod://envd-buildkitd
envd context use --name kubernetes-buildkit
envd build 
```

The user experience for the run may look like this:

```
envd context create --name docker --docker /var/run/docker.sock
envd context create --name kubernetes --kubernetes <kubeconfig-path>
envd context use kubernetes 
envd up
```

The default builder should be `docker-container://envd_buildkitd`, the default runner should be `--docker /var/run/docker.sock` (Maybe different in macOS)
@rudeigerc Hi are you still interested in the issue? Users expect it in the next release to support #510 
@gaocegege I am a little bit running out of time recently. Please feel free to unassign me if other ones are interested in this issue.

Fix patch:
diff --git a/pkg/app/app.go b/pkg/app/app.go
index 1fade75ad..d5d8c6cf4 100644
--- a/pkg/app/app.go
+++ b/pkg/app/app.go
@@ -48,16 +48,12 @@ func New() EnvdApp {
 			Usage: "docker image to use for buildkitd",
 			Value: "docker.io/moby/buildkit:v0.10.3",
 		},
-		&cli.StringFlag{
-			Name:  flag.FlagBuildkitdContainer,
-			Usage: "buildkitd container to use for buildkitd",
-			Value: "envd_buildkitd",
-		},
 	}
 
 	internalApp.Commands = []*cli.Command{
 		CommandBootstrap,
 		CommandBuild,
+		CommandContext,
 		CommandDestroy,
 		CommandGet,
 		CommandPause,
@@ -83,7 +79,6 @@ func New() EnvdApp {
 		}
 
 		// TODO(gaocegege): Add a config struct to keep them.
-		viper.Set(flag.FlagBuildkitdContainer, context.String(flag.FlagBuildkitdContainer))
 		viper.Set(flag.FlagBuildkitdImage, context.String(flag.FlagBuildkitdImage))
 		viper.Set(flag.FlagDebug, debugEnabled)
 		return nil
diff --git a/pkg/app/bootstrap.go b/pkg/app/bootstrap.go
index 6082cc4f0..d9c456844 100644
--- a/pkg/app/bootstrap.go
+++ b/pkg/app/bootstrap.go
@@ -26,6 +26,7 @@ import (
 
 	ac "github.com/tensorchord/envd/pkg/autocomplete"
 	"github.com/tensorchord/envd/pkg/buildkitd"
+	"github.com/tensorchord/envd/pkg/home"
 	sshconfig "github.com/tensorchord/envd/pkg/ssh/config"
 	"github.com/tensorchord/envd/pkg/util/fileutil"
 )
@@ -142,8 +143,13 @@ func bootstrap(clicontext *cli.Context) error {
 	buildkit := clicontext.Bool("buildkit")
 
 	if buildkit {
+		currentDriver, currentSocket, err := home.GetManager().ContextGetCurrent()
+		if err != nil {
+			return errors.Wrap(err, "failed to get the current context")
+		}
 		logrus.Debug("bootstrap the buildkitd container")
-		bkClient, err := buildkitd.NewClient(clicontext.Context, clicontext.String("dockerhub-mirror"))
+		bkClient, err := buildkitd.NewClient(clicontext.Context,
+			currentDriver, currentSocket, clicontext.String("dockerhub-mirror"))
 		if err != nil {
 			return errors.Wrap(err, "failed to create buildkit client")
 		}
diff --git a/pkg/app/build.go b/pkg/app/build.go
index dabaae404..d9b581bda 100644
--- a/pkg/app/build.go
+++ b/pkg/app/build.go
@@ -93,12 +93,11 @@ func build(clicontext *cli.Context) error {
 	}
 
 	logger := logrus.WithFields(logrus.Fields{
-		"build-context":             buildContext,
-		"build-file":                manifest,
-		"config":                    config,
-		"tag":                       tag,
-		flag.FlagBuildkitdImage:     viper.GetString(flag.FlagBuildkitdImage),
-		flag.FlagBuildkitdContainer: viper.GetString(flag.FlagBuildkitdContainer),
+		"build-context":         buildContext,
+		"build-file":            manifest,
+		"config":                config,
+		"tag":                   tag,
+		flag.FlagBuildkitdImage: viper.GetString(flag.FlagBuildkitdImage),
 	})
 	logger.Debug("starting build command")
 	debug := clicontext.Bool("debug")
diff --git a/pkg/app/context.go b/pkg/app/context.go
new file mode 100644
index 000000000..5dc40b6ee
--- /dev/null
+++ b/pkg/app/context.go
@@ -0,0 +1,30 @@
+// Copyright 2022 The envd Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package app
+
+import (
+	"github.com/urfave/cli/v2"
+)
+
+var CommandContext = &cli.Command{
+	Name:  "context",
+	Usage: "Manage envd contexts",
+	Subcommands: []*cli.Command{
+		CommandContextCreate,
+		CommandContextList,
+		CommandContextRemove,
+		CommandContextUse,
+	},
+}
diff --git a/pkg/app/context_create.go b/pkg/app/context_create.go
new file mode 100644
index 000000000..81261675b
--- /dev/null
+++ b/pkg/app/context_create.go
@@ -0,0 +1,68 @@
+// Copyright 2022 The envd Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package app
+
+import (
+	"github.com/cockroachdb/errors"
+	"github.com/sirupsen/logrus"
+	"github.com/tensorchord/envd/pkg/home"
+	"github.com/tensorchord/envd/pkg/types"
+	"github.com/urfave/cli/v2"
+)
+
+var CommandContextCreate = &cli.Command{
+	Name:  "create",
+	Usage: "Create envd context",
+	Flags: []cli.Flag{
+		&cli.StringFlag{
+			Name:  "name",
+			Usage: "Name of the context",
+			Value: "",
+		},
+		&cli.StringFlag{
+			Name:  "builder",
+			Usage: "Builder to use (docker-container, kube-pod)",
+			Value: string(types.BuilderTypeDocker),
+		},
+		&cli.StringFlag{
+			Name:  "builder-name",
+			Usage: "Builder name",
+			Value: "envd_buildkitd",
+		},
+		&cli.BoolFlag{
+			Name:  "use",
+			Usage: "Use the context",
+		},
+	},
+	Action: contextCreate,
+}
+
+func contextCreate(clicontext *cli.Context) error {
+	name := clicontext.String("name")
+	builder := clicontext.String("builder")
+	builderName := clicontext.String("builder-name")
+	use := clicontext.Bool("use")
+
+	err := home.GetManager().ContextCreate(name,
+		types.BuilderType(builder), builderName, use)
+	if err != nil {
+		return errors.Wrap(err, "failed to create context")
+	}
+	logrus.Infof("Context %s is created", name)
+	if use {
+		logrus.Infof("Current context is now \"%s\"", name)
+	}
+	return nil
+}
diff --git a/pkg/app/context_ls.go b/pkg/app/context_ls.go
new file mode 100644
index 000000000..3e2f1272d
--- /dev/null
+++ b/pkg/app/context_ls.go
@@ -0,0 +1,73 @@
+// Copyright 2022 The envd Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package app
+
+import (
+	"fmt"
+	"io"
+	"os"
+
+	"github.com/cockroachdb/errors"
+	"github.com/olekukonko/tablewriter"
+	"github.com/urfave/cli/v2"
+
+	"github.com/tensorchord/envd/pkg/home"
+	"github.com/tensorchord/envd/pkg/types"
+)
+
+var CommandContextList = &cli.Command{
+	Name:   "ls",
+	Usage:  "List envd contexts",
+	Action: contextList,
+}
+
+func contextList(clicontext *cli.Context) error {
+	contexts, err := home.GetManager().ContextList()
+	if err != nil {
+		return errors.Wrap(err, "failed to list context")
+	}
+	renderContext(contexts, os.Stdout)
+	return nil
+}
+
+func renderContext(contexts types.EnvdContext, w io.Writer) {
+	table := tablewriter.NewWriter(w)
+	table.SetHeader([]string{"context", "builder", "socket"})
+
+	table.SetAutoWrapText(false)
+	table.SetAutoFormatHeaders(true)
+	table.SetHeaderAlignment(tablewriter.ALIGN_LEFT)
+	table.SetAlignment(tablewriter.ALIGN_LEFT)
+	table.SetCenterSeparator("")
+	table.SetColumnSeparator("")
+	table.SetRowSeparator("")
+	table.SetHeaderLine(false)
+	table.SetBorder(false)
+	table.SetTablePadding("\t") // pad with tabs
+	table.SetNoWhiteSpace(true)
+
+	for _, p := range contexts.Contexts {
+		envRow := make([]string, 3)
+		if p.Name == contexts.Current {
+			envRow[0] = fmt.Sprintf("%s (current)", p.Name)
+		} else {
+			envRow[0] = p.Name
+		}
+		envRow[1] = string(p.Builder)
+		envRow[2] = fmt.Sprintf("%s://%s", p.Builder, p.BuilderSocket)
+		table.Append(envRow)
+	}
+	table.Render()
+}
diff --git a/pkg/app/context_rm.go b/pkg/app/context_rm.go
new file mode 100644
index 000000000..19d3fbebf
--- /dev/null
+++ b/pkg/app/context_rm.go
@@ -0,0 +1,46 @@
+// Copyright 2022 The envd Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package app
+
+import (
+	"github.com/cockroachdb/errors"
+	"github.com/sirupsen/logrus"
+	"github.com/tensorchord/envd/pkg/home"
+	"github.com/urfave/cli/v2"
+)
+
+var CommandContextRemove = &cli.Command{
+	Name:  "rm",
+	Usage: "Remove envd context",
+	Flags: []cli.Flag{
+		&cli.StringFlag{
+			Name:  "name",
+			Usage: "Name of the context",
+			Value: "",
+		},
+	},
+	Action: contextRemove,
+}
+
+func contextRemove(clicontext *cli.Context) error {
+	name := clicontext.String("name")
+
+	err := home.GetManager().ContextRemove(name)
+	if err != nil {
+		return errors.Wrap(err, "failed to remove context")
+	}
+	logrus.Infof("Context %s is removed", name)
+	return nil
+}
diff --git a/pkg/app/context_use.go b/pkg/app/context_use.go
new file mode 100644
index 000000000..60a63bfc0
--- /dev/null
+++ b/pkg/app/context_use.go
@@ -0,0 +1,46 @@
+// Copyright 2022 The envd Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package app
+
+import (
+	"github.com/cockroachdb/errors"
+	"github.com/sirupsen/logrus"
+	"github.com/tensorchord/envd/pkg/home"
+	"github.com/urfave/cli/v2"
+)
+
+var CommandContextUse = &cli.Command{
+	Name:  "use",
+	Usage: "Use the specified envd context",
+	Flags: []cli.Flag{
+		&cli.StringFlag{
+			Name:  "name",
+			Usage: "Name of the context",
+			Value: "",
+		},
+	},
+	Action: contextUse,
+}
+
+func contextUse(clicontext *cli.Context) error {
+	name := clicontext.String("name")
+
+	err := home.GetManager().ContextUse(name)
+	if err != nil {
+		return errors.Wrapf(err, "failed to use the specified context %s", name)
+	}
+	logrus.Infof("Current context is now \"%s\"", name)
+	return nil
+}
diff --git a/pkg/app/prune.go b/pkg/app/prune.go
index fb7faf9c4..ac835a025 100644
--- a/pkg/app/prune.go
+++ b/pkg/app/prune.go
@@ -64,7 +64,13 @@ func prune(clicontext *cli.Context) error {
 	keepStorage := clicontext.Float64("keep-storage")
 	filter := clicontext.StringSlice("filter")
 	verbose := clicontext.Bool("verbose")
-	bkClient, err := buildkitd.NewClient(clicontext.Context, "")
+
+	currentDriver, currentSocket, err := home.GetManager().ContextGetCurrent()
+	if err != nil {
+		return errors.Wrap(err, "failed to get the current context")
+	}
+	bkClient, err := buildkitd.NewClient(clicontext.Context,
+		currentDriver, currentSocket, "")
 	if err != nil {
 		return errors.Wrap(err, "failed to create buildkit client")
 	}
diff --git a/pkg/app/up.go b/pkg/app/up.go
index 23741f748..35066524e 100644
--- a/pkg/app/up.go
+++ b/pkg/app/up.go
@@ -126,14 +126,13 @@ func up(clicontext *cli.Context) error {
 	detach := clicontext.Bool("detach")
 
 	logger := logrus.WithFields(logrus.Fields{
-		"build-context":             buildContext,
-		"build-file":                manifest,
-		"config":                    config,
-		"tag":                       tag,
-		"container-name":            ctr,
-		"detach":                    detach,
-		flag.FlagBuildkitdImage:     viper.GetString(flag.FlagBuildkitdImage),
-		flag.FlagBuildkitdContainer: viper.GetString(flag.FlagBuildkitdContainer),
+		"build-context":         buildContext,
+		"build-file":            manifest,
+		"config":                config,
+		"tag":                   tag,
+		"container-name":        ctr,
+		"detach":                detach,
+		flag.FlagBuildkitdImage: viper.GetString(flag.FlagBuildkitdImage),
 	})
 	logger.Debug("starting up command")
 	debug := clicontext.Bool("debug")
diff --git a/pkg/builder/builder.go b/pkg/builder/builder.go
index 86e4591b2..79ce11a90 100644
--- a/pkg/builder/builder.go
+++ b/pkg/builder/builder.go
@@ -80,7 +80,11 @@ func New(ctx context.Context, configFilePath, manifestFilePath, buildContextDir,
 		}),
 	}
 
-	cli, err := buildkitd.NewClient(ctx, "")
+	currentDriver, currentSocket, err := home.GetManager().ContextGetCurrent()
+	if err != nil {
+		return nil, errors.Wrap(err, "failed to get the current context")
+	}
+	cli, err := buildkitd.NewClient(ctx, currentDriver, currentSocket, "")
 	if err != nil {
 		return nil, errors.Wrap(err, "failed to create buildkit client")
 	}
diff --git a/pkg/buildkitd/buildkitd.go b/pkg/buildkitd/buildkitd.go
index 14c4aa2e2..78873a551 100644
--- a/pkg/buildkitd/buildkitd.go
+++ b/pkg/buildkitd/buildkitd.go
@@ -30,6 +30,7 @@ import (
 
 	"github.com/tensorchord/envd/pkg/docker"
 	"github.com/tensorchord/envd/pkg/flag"
+	"github.com/tensorchord/envd/pkg/types"
 )
 
 var (
@@ -55,19 +56,27 @@ type generalClient struct {
 	image         string
 	mirror        string
 
+	driver types.BuilderType
+	socket string
+
 	*client.Client
 	logger *logrus.Entry
 }
 
-func NewClient(ctx context.Context, mirror string) (Client, error) {
+func NewClient(ctx context.Context, driver types.BuilderType,
+	socket, mirror string) (Client, error) {
 	c := &generalClient{
-		containerName: viper.GetString(flag.FlagBuildkitdContainer),
+		containerName: socket,
 		image:         viper.GetString(flag.FlagBuildkitdImage),
 		mirror:        mirror,
 	}
+	c.socket = socket
+	c.driver = driver
 	c.logger = logrus.WithFields(logrus.Fields{
 		"container": c.containerName,
 		"image":     c.image,
+		"socket":    c.socket,
+		"driver":    c.driver,
 	})
 
 	cli, err := client.New(ctx, c.BuildkitdAddr(), client.WithFailFast())
@@ -99,32 +108,36 @@ func (c generalClient) Close() error {
 // that can be used to connect to it.
 func (c *generalClient) maybeStart(ctx context.Context,
 	runningTimeout, connectingTimeout time.Duration) (string, error) {
-	dockerClient, err := docker.NewClient(ctx)
-	if err != nil {
-		return "", err
-	}
-
-	created, err := dockerClient.IsCreated(ctx, c.containerName)
-	if err != nil {
-		return "", err
-	}
-
-	if !created {
-		c.logger.Debug("container not created, creating...")
-		if _, err := dockerClient.StartBuildkitd(ctx, c.image, c.containerName, c.mirror); err != nil {
+	if c.driver == types.BuilderTypeDocker {
+		dockerClient, err := docker.NewClient(ctx)
+		if err != nil {
 			return "", err
 		}
-		if err := dockerClient.WaitUntilRunning(ctx, c.containerName, runningTimeout); err != nil {
+
+		created, err := dockerClient.IsCreated(ctx, c.containerName)
+		if err != nil {
 			return "", err
 		}
-	}
-	running, _ := dockerClient.IsRunning(ctx, c.containerName)
-	if created && !running {
-		c.logger.Warnf("please remove or restart the container %s", c.containerName)
-		return "", errors.Errorf("container %s is stopped", c.containerName)
-	}
 
-	c.logger.Debugf("container is running, check if it's ready at %s...", c.BuildkitdAddr())
+		if !created {
+			c.logger.Debug("container not created, creating...")
+			if _, err := dockerClient.StartBuildkitd(
+				ctx, c.image, c.containerName, c.mirror); err != nil {
+				return "", err
+			}
+			if err := dockerClient.WaitUntilRunning(
+				ctx, c.containerName, runningTimeout); err != nil {
+				return "", err
+			}
+		}
+		running, _ := dockerClient.IsRunning(ctx, c.containerName)
+		if created && !running {
+			c.logger.Warnf("please remove or restart the container %s", c.containerName)
+			return "", errors.Errorf("container %s is stopped", c.containerName)
+		}
+
+		c.logger.Debugf("container is running, check if it's ready at %s...", c.BuildkitdAddr())
+	}
 
 	if err := c.waitUntilConnected(ctx, connectingTimeout); err != nil {
 		return "", errors.Wrap(err, "failed to connect to buildkitd")
@@ -217,5 +230,5 @@ func (c generalClient) Prune(ctx context.Context, keepDuration time.Duration,
 }
 
 func (c generalClient) BuildkitdAddr() string {
-	return fmt.Sprintf("docker-container://%s", c.containerName)
+	return fmt.Sprintf("%s://%s", c.driver, c.socket)
 }
diff --git a/pkg/config/config.go b/pkg/config/config.go
index 91cf59ea8..5301dd06b 100644
--- a/pkg/config/config.go
+++ b/pkg/config/config.go
@@ -13,21 +13,10 @@
 
 package config
 
-import (
-	"path/filepath"
-
-	"github.com/adrg/xdg"
-)
-
 type UpState string
 
 const (
 	PrivateKeyFile              = "id_rsa_envd"
 	PublicKeyFile               = "id_rsa_envd.pub"
 	ContainerAuthorizedKeysPath = "/var/envd/authorized_keys"
-	envdFolderName              = ".envd"
 )
-
-func GetEnvdHome() string {
-	return filepath.Join(xdg.ConfigHome, envdFolderName)
-}
diff --git a/pkg/flag/consts.go b/pkg/flag/consts.go
index 7b42b18a0..04c98d0af 100644
--- a/pkg/flag/consts.go
+++ b/pkg/flag/consts.go
@@ -15,9 +15,8 @@
 package flag
 
 const (
-	FlagCacheDir           = "cache-dir"
-	FlagContextDir         = "context-dir"
-	FlagBuildkitdImage     = "buildkitd-image"
-	FlagBuildkitdContainer = "buildkitd-container-name"
-	FlagDebug              = "debug"
+	FlagCacheDir       = "cache-dir"
+	FlagContextDir     = "context-dir"
+	FlagBuildkitdImage = "buildkitd-image"
+	FlagDebug          = "debug"
 )
diff --git a/pkg/home/cache.go b/pkg/home/cache.go
new file mode 100644
index 000000000..cb5ea5eaa
--- /dev/null
+++ b/pkg/home/cache.go
@@ -0,0 +1,108 @@
+// Copyright 2022 The envd Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package home
+
+import (
+	"encoding/gob"
+	"os"
+	"path/filepath"
+
+	"github.com/adrg/xdg"
+	"github.com/cockroachdb/errors"
+	"github.com/sirupsen/logrus"
+)
+
+type cacheManager interface {
+	CacheDir() string
+	MarkCache(string, bool) error
+	Cached(string) bool
+	CleanCache() error
+}
+
+func (m *generalManager) initCache() error {
+	// Create $XDG_CACHE_HOME/envd
+	_, err := xdg.CacheFile("envd/cache")
+	if err != nil {
+		return errors.Wrap(err, "failed to get cache")
+	}
+	m.cacheDir = filepath.Join(xdg.CacheHome, "envd")
+
+	m.cacheStatusFile = filepath.Join(m.cacheDir, "cache.status")
+	_, err = os.Stat(m.cacheStatusFile)
+	if err != nil {
+		if os.IsNotExist(err) {
+			logrus.WithField("filename", m.cacheStatusFile).Debug("Creating file")
+			file, err := os.Create(m.cacheStatusFile)
+			if err != nil {
+				return errors.Wrap(err, "failed to create file")
+			}
+			err = file.Close()
+			if err != nil {
+				return errors.Wrap(err, "failed to close file")
+			}
+			if err := m.dumpCacheStatus(); err != nil {
+				return errors.Wrap(err, "failed to dump cache status")
+			}
+		} else {
+			return errors.Wrap(err, "failed to stat file")
+		}
+	}
+
+	file, err := os.Open(m.cacheStatusFile)
+	if err != nil {
+		return errors.Wrap(err, "failed to open cache status file")
+	}
+	defer file.Close()
+	e := gob.NewDecoder(file)
+	if err := e.Decode(&m.cacheMap); err != nil {
+		return errors.Wrap(err, "failed to decode cache map")
+	}
+	return nil
+}
+
+func (m generalManager) MarkCache(key string, cached bool) error {
+	m.cacheMap[key] = cached
+	return m.dumpCacheStatus()
+}
+
+func (m generalManager) Cached(key string) bool {
+	return m.cacheMap[key]
+}
+
+func (m *generalManager) dumpCacheStatus() error {
+	file, err := os.Create(m.cacheStatusFile)
+	if err != nil {
+		return errors.Wrap(err, "failed to create cache status file")
+	}
+	defer file.Close()
+
+	e := gob.NewEncoder(file)
+	if err := e.Encode(m.cacheMap); err != nil {
+		return errors.Wrap(err, "failed to encode cache map")
+	}
+	return nil
+}
+
+func (m generalManager) CacheDir() string {
+	return m.cacheDir
+}
+
+func (m generalManager) CleanCache() error {
+	if m.cacheDir == "" {
+		return nil
+	}
+	logrus.Debug("cleaning up host cache directory")
+	return os.RemoveAll(m.cacheDir)
+}
diff --git a/pkg/home/config.go b/pkg/home/config.go
new file mode 100644
index 000000000..c277c6bf5
--- /dev/null
+++ b/pkg/home/config.go
@@ -0,0 +1,44 @@
+// Copyright 2022 The envd Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package home
+
+import (
+	"github.com/adrg/xdg"
+	"github.com/cockroachdb/errors"
+
+	"github.com/tensorchord/envd/pkg/util/fileutil"
+)
+
+type configManager interface {
+	ConfigFile() string
+}
+
+func (m *generalManager) initConfig() error {
+	// Create $XDG_CONFIG_HOME/envd/config.envd
+	config, err := xdg.ConfigFile("envd/config.envd")
+	if err != nil {
+		return errors.Wrap(err, "failed to get config file")
+	}
+
+	if err := fileutil.CreateIfNotExist(config); err != nil {
+		return errors.Wrap(err, "failed to create config file")
+	}
+	m.configFile = config
+	return nil
+}
+
+func (m generalManager) ConfigFile() string {
+	return m.configFile
+}
diff --git a/pkg/home/context.go b/pkg/home/context.go
new file mode 100644
index 000000000..cfb210592
--- /dev/null
+++ b/pkg/home/context.go
@@ -0,0 +1,156 @@
+// Copyright 2022 The envd Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package home
+
+import (
+	"encoding/gob"
+	"os"
+
+	"github.com/adrg/xdg"
+	"github.com/cockroachdb/errors"
+	"github.com/sirupsen/logrus"
+
+	"github.com/tensorchord/envd/pkg/types"
+)
+
+type contextManager interface {
+	ContextFile() string
+	ContextList() (types.EnvdContext, error)
+	ContextUse(name string) error
+	ContextGetCurrent() (types.BuilderType, string, error)
+	ContextCreate(name string,
+		builder types.BuilderType, socket string, use bool) error
+	ContextRemove(name string) error
+}
+
+func (m *generalManager) initContext() error {
+	contextfile, err := xdg.ConfigFile("envd/contexts")
+	if err != nil {
+		return errors.Wrap(err, "failed to get context file")
+	}
+	m.contextFile = contextfile
+
+	// Create default context.
+
+	_, err = os.Stat(m.contextFile)
+	if err != nil {
+		if os.IsNotExist(err) {
+			logrus.WithField("filename", m.contextFile).Debug("Creating file")
+			file, err := os.Create(m.contextFile)
+			if err != nil {
+				return errors.Wrap(err, "failed to create file")
+			}
+			err = file.Close()
+			if err != nil {
+				return errors.Wrap(err, "failed to close file")
+			}
+			if err := m.dumpContext(); err != nil {
+				return errors.Wrap(err, "failed to dump context")
+			}
+		} else {
+			return errors.Wrap(err, "failed to stat file")
+		}
+	}
+
+	file, err := os.Open(m.contextFile)
+	if err != nil {
+		return errors.Wrap(err, "failed to open context file")
+	}
+	defer file.Close()
+	e := gob.NewDecoder(file)
+	if err := e.Decode(&m.context); err != nil {
+		return errors.Wrap(err, "failed to decode context file")
+	}
+	return nil
+}
+
+func (m generalManager) ContextFile() string {
+	return m.contextFile
+}
+
+func (m generalManager) ContextGetCurrent() (types.BuilderType, string, error) {
+	var driver types.BuilderType
+	var socket string
+	for _, c := range m.context.Contexts {
+		if m.context.Current == c.Name {
+			driver = c.Builder
+			socket = c.BuilderSocket
+			return driver, socket, nil
+		}
+	}
+	return "", "", errors.New("no current context")
+}
+
+func (m generalManager) ContextCreate(
+	name string, builder types.BuilderType, socket string, use bool) error {
+	for _, c := range m.context.Contexts {
+		if c.Name == name {
+			return errors.New("context already exists")
+		}
+	}
+	switch builder {
+	case types.BuilderTypeDocker, types.BuilderTypeKubernetes:
+		m.context.Contexts = append(m.context.Contexts, types.Context{
+			Name:          name,
+			Builder:       builder,
+			BuilderSocket: socket,
+		})
+	default:
+		return errors.New("unknown builder type")
+	}
+	return m.ContextUse(name)
+}
+
+func (m generalManager) ContextRemove(name string) error {
+	for i, c := range m.context.Contexts {
+		if c.Name == name {
+			if m.context.Current == name {
+				return errors.Newf("cannot remove current context \"%s\"", name)
+			}
+			m.context.Contexts = append(
+				m.context.Contexts[:i], m.context.Contexts[i+1:]...)
+			return m.dumpContext()
+		}
+	}
+	return nil
+}
+
+func (m generalManager) ContextList() (types.EnvdContext, error) {
+	return m.context, nil
+}
+
+func (m generalManager) ContextUse(name string) error {
+	for _, c := range m.context.Contexts {
+		if c.Name == name {
+			m.context.Current = name
+			return m.dumpContext()
+		}
+	}
+	return errors.New("context does not exist")
+}
+
+func (m *generalManager) dumpContext() error {
+	file, err := os.Create(m.contextFile)
+	if err != nil {
+		return errors.Wrap(err, "failed to create cache status file")
+	}
+	defer file.Close()
+
+	e := gob.NewEncoder(file)
+	if err := e.Encode(m.context); err != nil {
+		return errors.Wrap(err, "failed to encode cache map")
+	}
+	return nil
+}
diff --git a/pkg/home/manager.go b/pkg/home/manager.go
index 7dfb2b360..c7e00051d 100644
--- a/pkg/home/manager.go
+++ b/pkg/home/manager.go
@@ -15,34 +15,30 @@
 package home
 
 import (
-	"encoding/gob"
-	"os"
-	"path/filepath"
 	"sync"
 
-	"github.com/adrg/xdg"
 	"github.com/cockroachdb/errors"
 	"github.com/sirupsen/logrus"
 
 	sshconfig "github.com/tensorchord/envd/pkg/ssh/config"
-	"github.com/tensorchord/envd/pkg/util/fileutil"
+	"github.com/tensorchord/envd/pkg/types"
 )
 
 type Manager interface {
-	CacheDir() string
-	MarkCache(string, bool) error
-	Cached(string) bool
-	CleanCache() error
-	ConfigFile() string
+	configManager
+	contextManager
+	cacheManager
 }
 
 type generalManager struct {
 	cacheDir        string
 	cacheStatusFile string
 	configFile      string
+	contextFile     string
 
 	// TODO(gaocegege): Abstract CacheManager.
 	cacheMap map[string]bool
+	context  types.EnvdContext
 
 	logger *logrus.Entry
 }
@@ -56,6 +52,16 @@ func Initialize() error {
 	once.Do(func() {
 		defaultManager = &generalManager{
 			cacheMap: make(map[string]bool),
+			context: types.EnvdContext{
+				Current: "default",
+				Contexts: []types.Context{
+					{
+						Name:          "default",
+						Builder:       types.BuilderTypeDocker,
+						BuilderSocket: "envd_buildkitd",
+					},
+				},
+			},
 		}
 	})
 	return defaultManager.init()
@@ -65,104 +71,30 @@ func GetManager() Manager {
 	return defaultManager
 }
 
-func (m generalManager) CacheDir() string {
-	return m.cacheDir
-}
-
-func (m generalManager) CleanCache() error {
-	if m.cacheDir == "" {
-		return nil
-	}
-	logrus.Debug("cleaning up host cache directory")
-	return os.RemoveAll(m.cacheDir)
-}
-
-func (m generalManager) ConfigFile() string {
-	return m.configFile
-}
-
-func (m generalManager) MarkCache(key string, cached bool) error {
-	m.cacheMap[key] = cached
-	return m.dumpCacheStatus()
-}
-
-func (m generalManager) Cached(key string) bool {
-	return m.cacheMap[key]
-}
-
-func (m *generalManager) dumpCacheStatus() error {
-	file, err := os.Create(m.cacheStatusFile)
-	if err != nil {
-		return errors.Wrap(err, "failed to create cache status file")
-	}
-	defer file.Close()
-
-	e := gob.NewEncoder(file)
-	if err := e.Encode(m.cacheMap); err != nil {
-		return errors.Wrap(err, "failed to encode cache map")
-	}
-	return nil
-}
-
 func (m *generalManager) init() error {
-	// Create $XDG_CONFIG_HOME/envd/config.envd
-	config, err := xdg.ConfigFile("envd/config.envd")
-	if err != nil {
-		return errors.Wrap(err, "failed to get config file")
+	if err := m.initConfig(); err != nil {
+		return errors.Wrap(err, "failed to initialize config")
 	}
 
-	if err := fileutil.CreateIfNotExist(config); err != nil {
-		return errors.Wrap(err, "failed to create config file")
+	if err := m.initContext(); err != nil {
+		return errors.Wrap(err, "failed to initialize context")
 	}
-	m.configFile = config
 
-	// Create $XDG_CACHE_HOME/envd
-	_, err = xdg.CacheFile("envd/cache")
-	if err != nil {
-		return errors.Wrap(err, "failed to get cache")
-	}
-	m.cacheDir = filepath.Join(xdg.CacheHome, "envd")
-
-	m.cacheStatusFile = filepath.Join(m.cacheDir, "cache.status")
-	_, err = os.Stat(m.cacheStatusFile)
-	if err != nil {
-		if os.IsNotExist(err) {
-			logrus.WithField("filename", m.cacheStatusFile).Debug("Creating file")
-			file, err := os.Create(m.cacheStatusFile)
-			if err != nil {
-				return errors.Wrap(err, "failed to create file")
-			}
-			err = file.Close()
-			if err != nil {
-				return errors.Wrap(err, "failed to close file")
-			}
-			if err := m.dumpCacheStatus(); err != nil {
-				return errors.Wrap(err, "failed to dump cache status")
-			}
-		} else {
-			return errors.Wrap(err, "failed to stat file")
-		}
-	}
-
-	file, err := os.Open(m.cacheStatusFile)
-	if err != nil {
-		return errors.Wrap(err, "failed to open cache status file")
-	}
-	defer file.Close()
-	e := gob.NewDecoder(file)
-	if err := e.Decode(&m.cacheMap); err != nil {
-		return errors.Wrap(err, "failed to decode cache map")
+	if err := m.initCache(); err != nil {
+		return errors.Wrap(err, "failed to initialize cache")
 	}
 
-	// Generate SSH keys when init
 	if err := sshconfig.GenerateKeys(); err != nil {
 		return errors.Wrap(err, "failed to generate ssh key")
 	}
 
 	m.logger = logrus.WithFields(logrus.Fields{
 		"cache-dir":    m.cacheDir,
-		"config":       m.configFile,
+		"config-file":  m.configFile,
 		"cache-status": m.cacheStatusFile,
+		"context-file": m.contextFile,
+		"cache-map":    m.cacheMap,
+		"context":      m.context,
 	})
 
 	m.logger.Debug("home manager initialized")
diff --git a/pkg/types/envd.go b/pkg/types/envd.go
index 544f5470a..5adb37daa 100644
--- a/pkg/types/envd.go
+++ b/pkg/types/envd.go
@@ -46,6 +46,24 @@ type EnvdInfo struct {
 	types.Info
 }
 
+type EnvdContext struct {
+	Current  string    `json:"current,omitempty"`
+	Contexts []Context `json:"contexts,omitempty"`
+}
+
+type Context struct {
+	Name          string      `json:"name,omitempty"`
+	Builder       BuilderType `json:"builder,omitempty"`
+	BuilderSocket string      `json:"builder_socket,omitempty"`
+}
+
+type BuilderType string
+
+const (
+	BuilderTypeDocker     BuilderType = "docker-container"
+	BuilderTypeKubernetes BuilderType = "kube-pod"
+)
+
 type Dependency struct {
 	APTPackages  []string `json:"apt_packages,omitempty"`
 	PyPIPackages []string `json:"pypi_packages,omitempty"`


Test patch:
diff --git a/pkg/builder/builder_test.go b/pkg/builder/builder_test.go
index 3663484fe..e812eb8de 100644
--- a/pkg/builder/builder_test.go
+++ b/pkg/builder/builder_test.go
@@ -24,10 +24,8 @@ import (
 	. "github.com/onsi/ginkgo/v2"
 	. "github.com/onsi/gomega"
 	"github.com/sirupsen/logrus"
-	"github.com/spf13/viper"
 
 	mockbuildkitd "github.com/tensorchord/envd/pkg/buildkitd/mock"
-	"github.com/tensorchord/envd/pkg/flag"
 	"github.com/tensorchord/envd/pkg/home"
 	mockstarlark "github.com/tensorchord/envd/pkg/lang/frontend/starlark/mock"
 	"github.com/tensorchord/envd/pkg/lang/ir"
@@ -39,22 +37,15 @@ import (
 
 var _ = Describe("Builder", func() {
 	Describe("building image", Label("buildkitd"), func() {
-		var buildkitdSocket, configFilePath, manifestFilePath, buildContext, tag string
+		var configFilePath, manifestFilePath, buildContext, tag string
 		BeforeEach(func() {
-			buildkitdSocket = "docker-container://envd_buildkitd"
 			configFilePath = "config.envd"
 			manifestFilePath = "build.envd"
 			buildContext = "testdata"
 			tag = "envd-dev:test"
-			viper.Set(flag.FlagBuildkitdContainer, "envd_buildkitd")
-			os.Setenv("DOCKER_API_VERSION", "1.41")
-			DeferCleanup(func() {
-				viper.Set(flag.FlagBuildkitdContainer, "")
-			})
+			Expect(home.Initialize()).NotTo(HaveOccurred())
 		})
 		When("getting the wrong builtkitd address", func() {
-			buildkitdSocket = "wrong"
-			viper.Set(flag.FlagBuildkitdContainer, buildkitdSocket)
 			It("should return an error", func() {
 				_, err := New(context.TODO(), configFilePath, manifestFilePath, buildContext, tag, "", false)
 				Expect(err).To(HaveOccurred())
diff --git a/pkg/home/manager_test.go b/pkg/home/manager_test.go
index fda1013d8..73d643257 100644
--- a/pkg/home/manager_test.go
+++ b/pkg/home/manager_test.go
@@ -21,6 +21,7 @@ import (
 	"github.com/adrg/xdg"
 	. "github.com/onsi/ginkgo/v2"
 	. "github.com/onsi/gomega"
+	"github.com/tensorchord/envd/pkg/types"
 )
 
 var _ = Describe("home manager", func() {
@@ -30,6 +31,11 @@ var _ = Describe("home manager", func() {
 			m := GetManager()
 			Expect(m.CacheDir()).To(Equal(filepath.Join(xdg.CacheHome, "envd")))
 			Expect(m.ConfigFile()).To(Equal(filepath.Join(xdg.ConfigHome, "envd/config.envd")))
+			Expect(m.ContextFile()).To(Equal(filepath.Join(xdg.ConfigHome, "envd/contexts")))
+			driver, socket, err := m.ContextGetCurrent()
+			Expect(err).NotTo(HaveOccurred())
+			Expect(driver).To(Equal(types.BuilderTypeDocker))
+			Expect(socket).To(Equal("envd_buildkitd"))
 		})
 		It("should return the cache status", func() {
 			Expect(os.RemoveAll(filepath.Join(xdg.CacheHome, "envd/cache.status"))).NotTo(HaveOccurred())