feat(CLI): Support run command (#701)

* feat(CLI): Support run

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

* fix: Fix test cases

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>

* fix: Update

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

e2e/e2e_helper.go

@@ -15,15 +15,15 @@
 package e2e
 
 import (
+	"bytes"
 	"context"
 	"strings"
 
+	"github.com/pkg/errors"
 	"github.com/sirupsen/logrus"
 
 	"github.com/tensorchord/envd/pkg/app"
 	"github.com/tensorchord/envd/pkg/docker"
-	"github.com/tensorchord/envd/pkg/ssh"
-	sshconfig "github.com/tensorchord/envd/pkg/ssh/config"
 )
 
 func (e *Example) BuildImage(force bool) func() {
@@ -68,25 +68,32 @@ func GetDockerClient(ctx context.Context) docker.Client {
 type Example struct {
 	Name string
 	Tag  string
-	app  app.EnvdApp
+	app  *app.EnvdApp
 }
 
 func NewExample(name string, testcaseAbbr string) *Example {
 	tag := name + ":" + testcaseAbbr
+	app := app.New()
 	return &Example{
 		Name: name,
 		Tag:  tag,
-		app:  app.New(),
+		app:  &app,
 	}
 }
 
-func (e *Example) Exec(cmd string) string {
-	sshClient := e.getSSHClient()
-	ret, err := sshClient.ExecWithOutput(cmd)
+func (e *Example) Exec(cmd string) (string, error) {
+	args := []string{
+		"envd.test", "run", "--name", e.Name, "--command", cmd,
+	}
+
+	buffer := new(bytes.Buffer)
+	e.app.Writer = buffer
+
+	err := e.app.Run(args)
 	if err != nil {
-		panic(err)
+		return "", errors.Wrap(err, "failed to start `run` command")
 	}
-	return strings.Trim(string(ret), "\n")
+	return strings.Trim(buffer.String(), "\n"), nil
 }
 
 func (e *Example) RunContainer() func() {
@@ -114,18 +121,3 @@ func (e *Example) DestroyContainer() func() {
 		}
 	}
 }
-
-func (e *Example) getSSHClient() ssh.Client {
-	localhost := "127.0.0.1"
-	port, err := sshconfig.GetPort(e.Name)
-	if err != nil {
-		panic(err)
-	}
-	priv_path := sshconfig.GetPrivateKey()
-	sshClient, err := ssh.NewClient(
-		localhost, "envd", port, true, priv_path, "")
-	if err != nil {
-		panic(err)
-	}
-	return sshClient
-}

e2e/quick_start_test.go

@@ -26,7 +26,9 @@ var _ = Describe("e2e quickstart", Ordered, func() {
 	BeforeAll(e.BuildImage(true))
 	BeforeEach(e.RunContainer())
 	It("execute python demo.py", func() {
-		Expect(e.Exec("python demo.py")).To(Equal("[2 3 4]"))
+		res, err := e.Exec("python demo.py")
+		Expect(err).To(BeNil())
+		Expect(res).To(Equal("[2 3 4]"))
 	})
 	AfterEach(e.DestroyContainer())
 	AfterAll(e.RemoveImage())

pkg/app/app.go

@@ -60,6 +60,7 @@ func New() EnvdApp {
 		CommandInit,
 		CommandPause,
 		CommandPrune,
+		CommandRun,
 		CommandResume,
 		CommandUp,
 		CommandVersion,

pkg/app/bootstrap.go

@@ -52,8 +52,9 @@ var CommandBootstrap = &cli.Command{
 			Aliases: []string{"m"},
 		},
 		&cli.StringSliceFlag{
-			Name:    "ssh-keypair",
-			Usage:   fmt.Sprintf("Manually specify ssh key pair as `publicKey,privateKey`. Envd will generate a keypair at %s and %s if not specified", sshconfig.GetPublicKey(), sshconfig.GetPrivateKey()),
+			Name: "ssh-keypair",
+			Usage: fmt.Sprintf("Manually specify ssh key pair as `publicKey,privateKey`. Envd will generate a keypair at %s and %s if not specified",
+				sshconfig.GetPublicKeyOrPanic(), sshconfig.GetPrivateKeyOrPanic()),
 			Aliases: []string{"k"},
 		},
 	},
@@ -85,19 +86,27 @@ func bootstrap(clicontext *cli.Context) error {
 		if err != nil {
 			return errors.Wrap(err, "Cannot get default key status")
 		}
+
+		path, err := sshconfig.GetPrivateKey()
+		if err != nil {
+			return errors.Wrap(err, "Cannot get private key path")
+		}
+
 		if keyExists {
 			var exists bool
 			var newPrivateKeyName string
+
 			for ok := true; ok; ok = exists {
-				newPrivateKeyName = filepath.Join(filepath.Dir(sshconfig.GetPrivateKey()), fmt.Sprintf("envd_%s.pk", namesgenerator.GetRandomName(0)))
+				newPrivateKeyName = filepath.Join(filepath.Dir(path),
+					fmt.Sprintf("envd_%s.pk", namesgenerator.GetRandomName(0)))
 				exists, err = fileutil.FileExists(newPrivateKeyName)
 				if err != nil {
 					return err
 				}
 			}
 			logrus.Debugf("New key name: %s", newPrivateKeyName)
-			err := sshconfig.ReplaceKeyManagedByEnvd(sshconfig.GetPrivateKey(), newPrivateKeyName)
-			if err != nil {
+			if err := sshconfig.ReplaceKeyManagedByEnvd(
+				path, newPrivateKeyName); err != nil {
 				return err
 			}
 		}
@@ -106,7 +115,7 @@ func bootstrap(clicontext *cli.Context) error {
 		if err != nil {
 			return errors.Wrap(err, "Cannot open public key")
 		}
-		err = ioutil.WriteFile(sshconfig.GetPublicKey(), pubKey, 0644)
+		err = ioutil.WriteFile(path, pubKey, 0644)
 		if err != nil {
 			return errors.Wrap(err, "Cannot write public key")
 		}
@@ -115,7 +124,7 @@ func bootstrap(clicontext *cli.Context) error {
 		if err != nil {
 			return errors.Wrap(err, "Cannot open private key")
 		}
-		err = ioutil.WriteFile(sshconfig.GetPrivateKey(), priKey, 0600)
+		err = ioutil.WriteFile(path, priKey, 0600)
 
 		if err != nil {
 			return errors.Wrap(err, "Cannot write private key")

pkg/app/build.go

@@ -64,7 +64,7 @@ To build and push the image to a registry:
 			Name:    "public-key",
 			Usage:   "Path to the public key",
 			Aliases: []string{"pubk"},
-			Value:   sshconfig.GetPublicKey(),
+			Value:   sshconfig.GetPublicKeyOrPanic(),
 			Hidden:  true,
 		},
 		&cli.StringFlag{

pkg/app/env_describe.go

@@ -23,7 +23,6 @@ import (
 	"github.com/urfave/cli/v2"
 
 	"github.com/tensorchord/envd/pkg/envd"
-	sshconfig "github.com/tensorchord/envd/pkg/ssh/config"
 	"github.com/tensorchord/envd/pkg/types"
 )
 
@@ -37,12 +36,6 @@ var CommandDescribeEnvironment = &cli.Command{
 			Usage:   "Specify the envd environment to use",
 			Aliases: []string{"e"},
 		},
-		&cli.PathFlag{
-			Name:    "private-key",
-			Usage:   "Path to the private key",
-			Aliases: []string{"k"},
-			Value:   sshconfig.GetPrivateKey(),
-		},
 	},
 	Action: getEnvironmentDependency,
 }

pkg/app/run.go

@@ -0,0 +1,77 @@
+// 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"
+
+	"github.com/cockroachdb/errors"
+	"github.com/urfave/cli/v2"
+
+	"github.com/tensorchord/envd/pkg/docker"
+	"github.com/tensorchord/envd/pkg/ssh"
+)
+
+var CommandRun = &cli.Command{
+	Name:  "run",
+	Usage: "Spawns a command installed into the environment.",
+	Flags: []cli.Flag{
+		&cli.PathFlag{
+			Name:    "name",
+			Usage:   "Name of the environment",
+			Aliases: []string{"n"},
+		},
+		&cli.StringFlag{
+			Name:    "command",
+			Usage:   "Command to execute",
+			Aliases: []string{"c"},
+		},
+	},
+
+	Action: run,
+}
+
+func run(clicontext *cli.Context) error {
+	name := clicontext.String("name")
+
+	// Check if the container is running.
+	dockerClient, err := docker.NewClient(clicontext.Context)
+	if err != nil {
+		return errors.Wrap(err, "failed to create the docker client")
+	}
+	if isRunning, err :=
+		dockerClient.IsRunning(clicontext.Context, name); err != nil {
+		return errors.Wrapf(
+			err, "failed to check if the environment %s is running", name)
+	} else if !isRunning {
+		return errors.Newf("the environment %s is not running", name)
+	}
+
+	opt, err := ssh.GetOptions(name)
+	if err != nil {
+		return errors.Wrap(err, "failed to get the ssh options")
+	}
+	// SSH into the container and execute the command.
+	sshClient, err := ssh.NewClient(*opt)
+	if err != nil {
+		return errors.Wrap(err, "failed to get the ssh client")
+	}
+	if bytes, err := sshClient.ExecWithOutput(clicontext.String("command")); err != nil {
+		return errors.Wrap(err, "failed to execute the command")
+	} else {
+		fmt.Fprint(clicontext.App.Writer, string(bytes))
+	}
+	return nil
+}

pkg/app/up.go

@@ -69,14 +69,14 @@ var CommandUp = &cli.Command{
 			Name:    "private-key",
 			Usage:   "Path to the private key",
 			Aliases: []string{"k"},
-			Value:   sshconfig.GetPrivateKey(),
+			Value:   sshconfig.GetPrivateKeyOrPanic(),
 			Hidden:  true,
 		},
 		&cli.PathFlag{
 			Name:    "public-key",
 			Usage:   "Path to the public key",
 			Aliases: []string{"pubk"},
-			Value:   sshconfig.GetPublicKey(),
+			Value:   sshconfig.GetPublicKeyOrPanic(),
 			Hidden:  true,
 		},
 		&cli.DurationFlag{
@@ -234,8 +234,10 @@ func up(clicontext *cli.Context) error {
 	}
 
 	if !detach {
-		sshClient, err := ssh.NewClient(
-			localhost, "envd", sshPortInHost, true, clicontext.Path("private-key"), "")
+		opt := ssh.DefaultOptions()
+		opt.PrivateKeyPath = clicontext.Path("private-key")
+		opt.Port = sshPortInHost
+		sshClient, err := ssh.NewClient(opt)
 		if err != nil {
 			return errors.Wrap(err, "failed to create the ssh client")
 		}

pkg/builder/builder_test.go

@@ -50,7 +50,8 @@ var _ = Describe("Builder", func() {
 			BeforeEach(func() {
 				ctrl := gomock.NewController(GinkgoT())
 				ctrlStarlark := gomock.NewController(GinkgoT())
-				pub := sshconfig.GetPublicKey()
+				pub, err := sshconfig.GetPublicKey()
+				Expect(err).NotTo(HaveOccurred())
 				b = &generalBuilder{
 					Options: Options{
 						ManifestFilePath: manifestFilePath,

pkg/ssh/config/key.go

@@ -155,16 +155,34 @@ func DefaultKeyExists() (bool, error) {
 }
 
 // GetPublicKey returns the path to the public key
-func GetPublicKey() string {
+func GetPublicKey() (string, error) {
+	pub, _, err := getDefaultKeyPaths()
+	if err != nil {
+		return "", err
+	}
+	return pub, nil
+}
+
+// GetPrivateKey returns the path to the private key
+func GetPrivateKey() (string, error) {
+	_, pri, err := getDefaultKeyPaths()
+	if err != nil {
+		return "", err
+	}
+	return pri, nil
+}
+
+// GetPublicKeyOrPanic returns the path to the public key or panic.
+func GetPublicKeyOrPanic() string {
 	pub, _, err := getDefaultKeyPaths()
 	if err != nil {
 		logrus.Fatal("Cannot get public key path")
 	}
 	return pub
 }
 
-// GetPrivateKey returns the path to the private key
-func GetPrivateKey() string {
+// GetPrivateKeyOrPanic returns the path to the private key or panic.
+func GetPrivateKeyOrPanic() string {
 	_, pri, err := getDefaultKeyPaths()
 	if err != nil {
 		logrus.Fatal("Cannot get private key path")

pkg/ssh/config/ssh_config.go

@@ -348,7 +348,12 @@ func ReplaceKeyManagedByEnvd(oldKey string, newKey string) error {
 		}
 	}
 
-	err = os.Rename(GetPrivateKey(), newKey)
+	path, err := GetPrivateKey()
+	if err != nil {
+		return err
+	}
+
+	err = os.Rename(path, newKey)
 	if err != nil {
 		return err
 	}