feat(CLI): Support runner in context (#961)

* WIP

Signed-off-by: Ce Gao <cegao@tensorchord.ai>

* fix: Update

Signed-off-by: Ce Gao <cegao@tensorchord.ai>

* fix: Add check

Signed-off-by: Ce Gao <cegao@tensorchord.ai>

* fix: Update

Signed-off-by: Ce Gao <cegao@tensorchord.ai>

* fix: Update

Signed-off-by: Ce Gao <cegao@tensorchord.ai>

* fix: Update

Signed-off-by: Ce Gao <cegao@tensorchord.ai>

Signed-off-by: Ce Gao <cegao@tensorchord.ai>

e2e/cli/context_test.go

@@ -40,19 +40,28 @@ var _ = Describe("home context", func() {
 	})
 
 	Describe("add a new context", Ordered, func() {
-		contextName := "tcp-test"
-		builder := types.BuilderTypeTCP
-		socketAddr := "0.0.0.0:12345"
+		testContext := "envd_home_test"
+		testBuilderAddress := "0.0.0.0:12345"
+		testBuilder := types.BuilderTypeTCP
+		testRunner := types.RunnerTypeEnvdServer
+		testRunnerAddress := "http://localhost"
+		c := types.Context{
+			Name:           testContext,
+			Builder:        testBuilder,
+			BuilderAddress: testBuilderAddress,
+			Runner:         testRunner,
+			RunnerAddress:  &testRunnerAddress,
+		}
 
 		BeforeAll(func() {
-			err := home.GetManager().ContextCreate(contextName, builder, socketAddr, true)
+			err := home.GetManager().ContextCreate(c, true)
 			Expect(err).NotTo(HaveOccurred())
 		})
 
-		It("should found a new context", func() {
+		It("should find a new context", func() {
 			contexts, err := home.GetManager().ContextList()
 			Expect(err).NotTo(HaveOccurred())
-			Expect(contexts.Current).To(Equal(contextName))
+			Expect(contexts.Current).To(Equal(testContext))
 		})
 
 		Describe("connect buildkit through TCP", Ordered, func() {
@@ -85,14 +94,14 @@ var _ = Describe("home context", func() {
 		})
 
 		It("fail to delete the current context", func() {
-			err := home.GetManager().ContextRemove(contextName)
+			err := home.GetManager().ContextRemove(testContext)
 			Expect(err).To(HaveOccurred())
 		})
 
 		AfterAll(func() {
 			err := home.GetManager().ContextUse(defaultContext)
 			Expect(err).NotTo(HaveOccurred())
-			err = home.GetManager().ContextRemove(contextName)
+			err = home.GetManager().ContextRemove(testContext)
 			Expect(err).NotTo(HaveOccurred())
 			contexts, err := home.GetManager().ContextList()
 			Expect(err).NotTo(HaveOccurred())

pkg/app/bootstrap.go

@@ -192,14 +192,14 @@ func buildkit(clicontext *cli.Context) error {
 		return nil
 	}
 
-	currentDriver, currentSocket, err := home.GetManager().ContextGetCurrent()
+	c, 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,
-		currentDriver, currentSocket, clicontext.String("dockerhub-mirror"))
+		c.Builder, c.BuilderAddress, clicontext.String("dockerhub-mirror"))
 	if err != nil {
 		return errors.Wrap(err, "failed to create buildkit client")
 	}

pkg/app/context_create.go

@@ -39,10 +39,19 @@ var CommandContextCreate = &cli.Command{
 			Value: string(types.BuilderTypeDocker),
 		},
 		&cli.StringFlag{
-			Name:  "builder-socket",
-			Usage: "Builder socket",
+			Name:  "builder-address",
+			Usage: "Builder address",
 			Value: "envd_buildkitd",
 		},
+		&cli.StringFlag{
+			Name:  "runner",
+			Usage: "Runner to use(docker, envd-server)",
+			Value: string(types.RunnerTypeDocker),
+		},
+		&cli.StringFlag{
+			Name:  "runner-address",
+			Usage: "Runner address",
+		},
 		&cli.BoolFlag{
 			Name:  "use",
 			Usage: "Use the context",
@@ -54,11 +63,22 @@ var CommandContextCreate = &cli.Command{
 func contextCreate(clicontext *cli.Context) error {
 	name := clicontext.String("name")
 	builder := clicontext.String("builder")
-	builderSocket := clicontext.String("builder-socket")
+	builderAddress := clicontext.String("builder-address")
+	runner := clicontext.String("runner")
+	runnerAddress := clicontext.String("runner-address")
 	use := clicontext.Bool("use")
 
-	err := home.GetManager().ContextCreate(name,
-		types.BuilderType(builder), builderSocket, use)
+	c := types.Context{
+		Name:           name,
+		Builder:        types.BuilderType(builder),
+		BuilderAddress: builderAddress,
+		Runner:         types.RunnerType(runner),
+	}
+	if runnerAddress != "" {
+		c.RunnerAddress = &runnerAddress
+	}
+
+	err := home.GetManager().ContextCreate(c, use)
 	if err != nil {
 		return errors.Wrap(err, "failed to create context")
 	}

pkg/app/context_ls.go

@@ -44,7 +44,7 @@ func contextList(clicontext *cli.Context) error {
 
 func renderContext(contexts types.EnvdContext, w io.Writer) {
 	table := tablewriter.NewWriter(w)
-	table.SetHeader([]string{"context", "builder", "socket"})
+	table.SetHeader([]string{"context", "builder", "builder addr", "runner", "runner addr"})
 
 	table.SetAutoWrapText(false)
 	table.SetAutoFormatHeaders(true)
@@ -59,14 +59,18 @@ func renderContext(contexts types.EnvdContext, w io.Writer) {
 	table.SetNoWhiteSpace(true)
 
 	for _, p := range contexts.Contexts {
-		envRow := make([]string, 3)
+		envRow := make([]string, 5)
 		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)
+		envRow[2] = fmt.Sprintf("%s://%s", p.Builder, p.BuilderAddress)
+		envRow[3] = string(p.Runner)
+		if p.RunnerAddress != nil {
+			envRow[4] = stringOrNone(*p.RunnerAddress)
+		}
 		table.Append(envRow)
 	}
 	table.Render()

pkg/app/prune.go

@@ -66,12 +66,12 @@ func prune(clicontext *cli.Context) error {
 	filter := clicontext.StringSlice("filter")
 	verbose := clicontext.Bool("verbose")
 
-	currentDriver, currentSocket, err := home.GetManager().ContextGetCurrent()
+	c, 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, "")
+		c.Builder, c.BuilderAddress, "")
 	if err != nil {
 		return errors.Wrap(err, "failed to create buildkit client")
 	}

pkg/builder/builder.go

@@ -126,11 +126,11 @@ func New(ctx context.Context, opt Options) (Builder, error) {
 		}),
 	}
 
-	currentDriver, currentSocket, err := home.GetManager().ContextGetCurrent()
+	c, 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, "")
+	cli, err := buildkitd.NewClient(ctx, c.Builder, c.BuilderAddress, "")
 	if err != nil {
 		return nil, errors.Wrap(err, "failed to create buildkit client")
 	}

pkg/home/context.go

@@ -29,9 +29,8 @@ 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
+	ContextGetCurrent() (*types.Context, error)
+	ContextCreate(c types.Context, use bool) error
 	ContextRemove(name string) error
 }
 
@@ -80,38 +79,36 @@ func (m generalManager) ContextFile() string {
 	return m.contextFile
 }
 
-func (m generalManager) ContextGetCurrent() (types.BuilderType, string, error) {
-	var driver types.BuilderType
-	var socket string
+func (m generalManager) ContextGetCurrent() (*types.Context, error) {
 	for _, c := range m.context.Contexts {
 		if m.context.Current == c.Name {
-			driver = c.Builder
-			socket = c.BuilderSocket
-			return driver, socket, nil
+			return &c, nil
 		}
 	}
-	return "", "", errors.New("no current context")
+	return nil, errors.New("no current context")
 }
 
-func (m *generalManager) ContextCreate(
-	name string, builder types.BuilderType, socket string, use bool) error {
+func (m *generalManager) ContextCreate(ctx types.Context, use bool) error {
 	for _, c := range m.context.Contexts {
-		if c.Name == name {
-			return errors.Newf("context \"%s\" already exists", name)
+		if c.Name == ctx.Name {
+			return errors.Newf("context \"%s\" already exists", ctx.Name)
 		}
 	}
-	switch builder {
+	switch ctx.Builder {
 	case types.BuilderTypeDocker, types.BuilderTypeKubernetes, types.BuilderTypeTCP:
-		m.context.Contexts = append(m.context.Contexts, types.Context{
-			Name:          name,
-			Builder:       builder,
-			BuilderSocket: socket,
-		})
+		break
 	default:
 		return errors.New("unknown builder type")
 	}
+	switch ctx.Runner {
+	case types.RunnerTypeDocker, types.RunnerTypeEnvdServer:
+		break
+	default:
+		return errors.New("unknown runner type")
+	}
+	m.context.Contexts = append(m.context.Contexts, ctx)
 	if use {
-		return m.ContextUse(name)
+		return m.ContextUse(ctx.Name)
 	}
 	return m.dumpContext()
 }

pkg/home/context_test.go

@@ -24,27 +24,38 @@ import (
 var _ = Describe("context test", func() {
 	defaultContext := "default"
 	testContext := "envd_home_test"
-	testSocket := "0.0.0.0:12345"
+	testBuilderAddress := "0.0.0.0:12345"
 	testBuilder := types.BuilderTypeTCP
+	testRunner := types.RunnerTypeEnvdServer
+	testRunnerAddress := "http://localhost"
+	c := types.Context{
+		Name:           testContext,
+		Builder:        testBuilder,
+		BuilderAddress: testBuilderAddress,
+		Runner:         testRunner,
+		RunnerAddress:  &testRunnerAddress,
+	}
 
 	BeforeEach(func() {
 		Expect(Initialize()).To(Succeed())
 	})
 
 	Describe("create with use", Ordered, func() {
 		BeforeAll(func() {
-			err := GetManager().ContextCreate(testContext, testBuilder, testSocket, true)
+			err := GetManager().ContextCreate(c, true)
 			Expect(err).NotTo(HaveOccurred())
 		})
 
 		It("current context should be the new one", func() {
 			contexts, err := GetManager().ContextList()
 			Expect(err).NotTo(HaveOccurred())
 			Expect(contexts.Current).To(Equal(testContext))
-			builder, socket, err := GetManager().ContextGetCurrent()
+			c, err := GetManager().ContextGetCurrent()
 			Expect(err).NotTo(HaveOccurred())
-			Expect(builder).To(Equal(testBuilder))
-			Expect(socket).To(Equal(testSocket))
+			Expect(c.Builder).To(Equal(testBuilder))
+			Expect(c.BuilderAddress).To(Equal(testBuilderAddress))
+			Expect(c.Runner).To(Equal(testRunner))
+			Expect(*c.RunnerAddress).To(Equal(testRunnerAddress))
 		})
 
 		It("cannot delete the current context", func() {
@@ -63,12 +74,12 @@ var _ = Describe("context test", func() {
 
 	Describe("create without use", Ordered, func() {
 		BeforeAll(func() {
-			err := GetManager().ContextCreate(testContext, testBuilder, testSocket, false)
+			err := GetManager().ContextCreate(c, false)
 			Expect(err).NotTo(HaveOccurred())
 		})
 
 		It("should not be able to create the same context", func() {
-			err := GetManager().ContextCreate(testContext, testBuilder, testSocket, false)
+			err := GetManager().ContextCreate(c, false)
 			Expect(err).To(HaveOccurred())
 		})
 

pkg/home/manager.go

@@ -60,9 +60,11 @@ func Initialize() error {
 				Current: "default",
 				Contexts: []types.Context{
 					{
-						Name:          "default",
-						Builder:       types.BuilderTypeDocker,
-						BuilderSocket: "envd_buildkitd",
+						Name:           "default",
+						Builder:        types.BuilderTypeDocker,
+						BuilderAddress: "envd_buildkitd",
+						Runner:         types.RunnerTypeDocker,
+						RunnerAddress:  nil,
 					},
 				},
 			},

pkg/home/manager_test.go

@@ -34,9 +34,9 @@ var _ = Describe("home manager", func() {
 					Current: "default",
 					Contexts: []types.Context{
 						{
-							Name:          "default",
-							Builder:       types.BuilderTypeDocker,
-							BuilderSocket: "envd_buildkitd",
+							Name:           "default",
+							Builder:        types.BuilderTypeDocker,
+							BuilderAddress: "envd_buildkitd",
 						},
 					},
 				},
@@ -46,10 +46,11 @@ var _ = Describe("home manager", func() {
 			Expect(m.CacheDir()).To(Equal(filepath.Join(fileutil.DefaultCacheDir)))
 			Expect(m.ConfigFile()).To(Equal(filepath.Join(fileutil.DefaultConfigDir, "config.envd")))
 			Expect(m.ContextFile()).To(Equal(filepath.Join(fileutil.DefaultConfigDir, "contexts")))
-			driver, socket, err := m.ContextGetCurrent()
+			c, err := m.ContextGetCurrent()
 			Expect(err).NotTo(HaveOccurred())
-			Expect(driver).To(Equal(types.BuilderTypeDocker))
-			Expect(socket).To(Equal("envd_buildkitd"))
+			Expect(c.Builder).To(Equal(types.BuilderTypeDocker))
+			Expect(c.BuilderAddress).To(Equal("envd_buildkitd"))
+			Expect(c.Runner).To(Equal(types.RunnerTypeDocker))
 		})
 		It("should return the cache status", func() {
 			Expect(os.RemoveAll(filepath.Join(fileutil.DefaultCacheDir, "cache.status"))).NotTo(HaveOccurred())

pkg/types/envd.go

@@ -107,9 +107,11 @@ type EnvdContext struct {
 }
 
 type Context struct {
-	Name          string      `json:"name,omitempty"`
-	Builder       BuilderType `json:"builder,omitempty"`
-	BuilderSocket string      `json:"builder_socket,omitempty"`
+	Name           string      `json:"name,omitempty"`
+	Builder        BuilderType `json:"builder,omitempty"`
+	BuilderAddress string      `json:"builder_address,omitempty"`
+	Runner         RunnerType  `json:"runner,omitempty"`
+	RunnerAddress  *string     `json:"runner_address,omitempty"`
 }
 
 type BuilderType string
@@ -120,6 +122,13 @@ const (
 	BuilderTypeTCP        BuilderType = "tcp"
 )
 
+type RunnerType string
+
+const (
+	RunnerTypeDocker     RunnerType = "docker"
+	RunnerTypeEnvdServer RunnerType = "envd-server"
+)
+
 type Dependency struct {
 	APTPackages  []string `json:"apt_packages,omitempty"`
 	PyPIPackages []string `json:"pypi_packages,omitempty"`