Repository: tensorchord/envd (language: python)

Issue:
feat(lang): Support runtime execution targets
## Description

Users want to define the entrypoints e.g. `MIDI serve` to run some user-defined logics.


Discussion / hints:
One draft:

```python
def build():
    runtime.scripts({
        "train": "python main.py",
        "test": "python main.py",
        "predict": "python predict.py",
        "show": "python streamlit.py",
    })
```

```
$ envd run --target train
...
$ envd run --target test
...
$ envd run --cmd "python ./ad-hoc.py"
...
```
One thing to add: The image might be directly used on k8s. Therefore the entrypoint functionality should also work without envd.

Fix patch:
diff --git a/pkg/app/app.go b/pkg/app/app.go
index 16eabed84..2d02da8ef 100644
--- a/pkg/app/app.go
+++ b/pkg/app/app.go
@@ -60,6 +60,7 @@ func New() EnvdApp {
 		CommandInit,
 		CommandPause,
 		CommandPrune,
+		CommandRun,
 		CommandResume,
 		CommandUp,
 		CommandVersion,
diff --git a/pkg/app/bootstrap.go b/pkg/app/bootstrap.go
index 35f4c18b9..5252e26a0 100644
--- a/pkg/app/bootstrap.go
+++ b/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")
diff --git a/pkg/app/build.go b/pkg/app/build.go
index e9cac1d39..9164603c1 100644
--- a/pkg/app/build.go
+++ b/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{
diff --git a/pkg/app/env_describe.go b/pkg/app/env_describe.go
index b485d714e..f37216ab3 100644
--- a/pkg/app/env_describe.go
+++ b/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,
 }
diff --git a/pkg/app/run.go b/pkg/app/run.go
new file mode 100644
index 000000000..caab8e78f
--- /dev/null
+++ b/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
+}
diff --git a/pkg/app/up.go b/pkg/app/up.go
index a1df97817..d48ae1026 100644
--- a/pkg/app/up.go
+++ b/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")
 		}
diff --git a/pkg/ssh/config/key.go b/pkg/ssh/config/key.go
index bc6a3627d..8151096c8 100644
--- a/pkg/ssh/config/key.go
+++ b/pkg/ssh/config/key.go
@@ -155,7 +155,25 @@ 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")
@@ -163,8 +181,8 @@ func GetPublicKey() string {
 	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")
diff --git a/pkg/ssh/config/ssh_config.go b/pkg/ssh/config/ssh_config.go
index 7ca5ee56b..7d193f82c 100644
--- a/pkg/ssh/config/ssh_config.go
+++ b/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
 	}
diff --git a/pkg/ssh/ssh.go b/pkg/ssh/ssh.go
index 15789118e..d8b4ca5d2 100644
--- a/pkg/ssh/ssh.go
+++ b/pkg/ssh/ssh.go
@@ -37,6 +37,7 @@ import (
 	"golang.org/x/term"
 
 	"github.com/tensorchord/envd/pkg/lang/ir"
+	"github.com/tensorchord/envd/pkg/ssh/config"
 )
 
 type Client interface {
@@ -45,14 +46,47 @@ type Client interface {
 	Close() error
 }
 
+type Options struct {
+	Server         string
+	User           string
+	Port           int
+	Auth           bool
+	PrivateKeyPath string
+	PrivateKeyPwd  string
+}
+
+func DefaultOptions() Options {
+	return Options{
+		Server:        "localhost",
+		User:          "envd",
+		Auth:          true,
+		PrivateKeyPwd: "",
+	}
+}
+
+func GetOptions(entry string) (*Options, error) {
+	path, err := config.GetPrivateKey()
+	if err != nil {
+		return nil, errors.Wrap(err, "getting private key failed")
+	}
+	port, err := config.GetPort(entry)
+	if err != nil {
+		return nil, errors.Wrap(err, "getting port failed")
+	}
+	// TODO(gaocegege): Make it configurable.
+	opt := DefaultOptions()
+	opt.Port = port
+	opt.PrivateKeyPath = path
+	return &opt, nil
+}
+
 type generalClient struct {
 	cli *ssh.Client
 }
 
-func NewClient(server, user string,
-	port int, auth bool, privateKeyPath, privateKeyPwd string) (Client, error) {
+func NewClient(opt Options) (Client, error) {
 	config := &ssh.ClientConfig{
-		User: user,
+		User: opt.User,
 		HostKeyCallback: func(hostname string, remote net.Addr, key ssh.PublicKey) error {
 			// use OpenSSH's known_hosts file if you care about host validation
 			return nil
@@ -61,14 +95,15 @@ func NewClient(server, user string,
 
 	var cli *ssh.Client
 
-	if auth {
+	if opt.Auth {
 		// read private key file
-		pemBytes, err := ioutil.ReadFile(privateKeyPath)
+		pemBytes, err := ioutil.ReadFile(opt.PrivateKeyPath)
 		if err != nil {
-			return nil, errors.Wrapf(err, "reading private key %s failed", privateKeyPath)
+			return nil, errors.Wrapf(
+				err, "reading private key %s failed", opt.PrivateKeyPath)
 		}
 		// create signer
-		signer, err := signerFromPem(pemBytes, []byte(privateKeyPwd))
+		signer, err := signerFromPem(pemBytes, []byte(opt.PrivateKeyPwd))
 		if err != nil {
 			return nil, errors.Wrap(err, "creating signer from private key failed")
 		}
@@ -77,7 +112,7 @@ func NewClient(server, user string,
 		}
 	}
 
-	host := fmt.Sprintf("%s:%d", server, port)
+	host := fmt.Sprintf("%s:%d", opt.Server, opt.Port)
 	// open connection
 	conn, err := ssh.Dial("tcp", host, config)
 	if err != nil {


Test patch:
diff --git a/e2e/e2e_helper.go b/e2e/e2e_helper.go
index 363a1ad5c..07bdbc9ad 100644
--- a/e2e/e2e_helper.go
+++ b/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
-}
diff --git a/e2e/quick_start_test.go b/e2e/quick_start_test.go
index 6b71460c4..b4567034e 100644
--- a/e2e/quick_start_test.go
+++ b/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())
diff --git a/pkg/builder/builder_test.go b/pkg/builder/builder_test.go
index 8f5d3e3a5..e109c7938 100644
--- a/pkg/builder/builder_test.go
+++ b/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,