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.
I will work on this.

Fix patch:
diff --git a/Makefile b/Makefile
index 261f0c903..2f92836ee 100644
--- a/Makefile
+++ b/Makefile
@@ -178,7 +178,7 @@ e2e-test:
 		-X $(ROOT)/pkg/version.gitTreeState=$(GIT_TREE_STATE)                     \
 		-X $(ROOT)/pkg/version.gitTag="$(shell git describe --tags --abbrev=0)" \
 		-X $(ROOT)/pkg/version.developmentFlag=true" \
-		-race -v -coverpkg=./pkg/app -coverprofile=e2e-coverage.out ./e2e
+		-race -v -timeout 20m -coverpkg=./pkg/app -coverprofile=e2e-coverage.out ./e2e
 
 clean:  ## Clean the outputs and artifacts
 	@-rm -vrf ${OUTPUT_DIR}
diff --git a/examples/python-basic/build.envd b/examples/python-basic/build.envd
index fb2f67c51..e8d345710 100644
--- a/examples/python-basic/build.envd
+++ b/examples/python-basic/build.envd
@@ -1,7 +1,10 @@
 def build():
-    base(os="ubuntu20.04", language="python3.8")
+    base(os="ubuntu20.04", language="python")
     #config.pip_index(url = "https://pypi.tuna.tsinghua.edu.cn/simple")
     install.python_packages([
         "via",
     ])
     io.copy(src="./build.envd", dest="/")
+    runtime.command(commands={
+        "test": "ls /",
+    })
diff --git a/go.mod b/go.mod
index bd4031312..2210e78a4 100644
--- a/go.mod
+++ b/go.mod
@@ -24,7 +24,6 @@ require (
 	github.com/onsi/gomega v1.20.0
 	github.com/opencontainers/go-digest v1.0.0
 	github.com/opencontainers/image-spec v1.0.3-0.20211202183452-c5a74bcca799
-	github.com/pkg/errors v0.9.1
 	github.com/pkg/sftp v1.13.5
 	github.com/sirupsen/logrus v1.9.0
 	github.com/spf13/viper v1.12.0
@@ -39,7 +38,10 @@ require (
 	golang.org/x/time v0.0.0-20210723032227-1f47c861a9ac
 )
 
-require github.com/op/go-logging v0.0.0-20160211212156-b2cb9fa56473 // indirect
+require (
+	github.com/op/go-logging v0.0.0-20160211212156-b2cb9fa56473 // indirect
+	github.com/pkg/errors v0.9.1 // indirect
+)
 
 require (
 	github.com/Azure/go-ansiterm v0.0.0-20210617225240-d185dfc1b5a1 // indirect
diff --git a/pkg/app/build.go b/pkg/app/build.go
index 895224bce..f876cd8a7 100644
--- a/pkg/app/build.go
+++ b/pkg/app/build.go
@@ -106,7 +106,7 @@ func build(clicontext *cli.Context) error {
 		return errors.Wrap(err, "failed to get absolute path of the build file")
 	}
 	if manifest == "" {
-		return errors.New("file does not exist")
+		return errors.Newf("build file %s does not exist", fileName)
 	}
 
 	cfg := home.GetManager().ConfigFile()
diff --git a/pkg/app/run.go b/pkg/app/run.go
index 08b705f34..9a5cfaffc 100644
--- a/pkg/app/run.go
+++ b/pkg/app/run.go
@@ -16,12 +16,16 @@ package app
 
 import (
 	"fmt"
+	"path/filepath"
 
 	"github.com/cockroachdb/errors"
 	"github.com/urfave/cli/v2"
 
+	"github.com/tensorchord/envd/pkg/builder"
 	"github.com/tensorchord/envd/pkg/docker"
+	"github.com/tensorchord/envd/pkg/lang/ir"
 	"github.com/tensorchord/envd/pkg/ssh"
+	"github.com/tensorchord/envd/pkg/util/fileutil"
 )
 
 var CommandRun = &cli.Command{
@@ -36,9 +40,26 @@ var CommandRun = &cli.Command{
 		},
 		&cli.StringFlag{
 			Name:    "command",
-			Usage:   "Command to execute",
+			Usage:   "Command defined in build.envd to execute",
 			Aliases: []string{"c"},
 		},
+		&cli.PathFlag{
+			Name:    "from",
+			Usage:   "Function to execute, format `file:func`",
+			Aliases: []string{"f"},
+			Value:   "build.envd:build",
+		},
+		&cli.PathFlag{
+			Name:    "path",
+			Usage:   "Path to the directory containing the build.envd",
+			Aliases: []string{"p"},
+			Value:   ".",
+		},
+		&cli.StringFlag{
+			Name:    "raw",
+			Usage:   "Raw command to execute",
+			Aliases: []string{"r"},
+		},
 	},
 
 	Action: run,
@@ -46,6 +67,53 @@ var CommandRun = &cli.Command{
 
 func run(clicontext *cli.Context) error {
 	name := clicontext.String("name")
+	command := clicontext.String("command")
+	rawCommand := clicontext.String("raw")
+	path := clicontext.String("path")
+
+	if command != "" && rawCommand != "" {
+		return errors.New("--raw and --command are mutually exclusive and may only be used once")
+	}
+
+	resultCommand := rawCommand
+	if command != "" {
+		buildContext, err := filepath.Abs(path)
+		if err != nil {
+			return errors.Wrap(err, "failed to get absolute path of the build context")
+		}
+		fileName, funcName, err := builder.ParseFromStr(clicontext.String("from"))
+		if err != nil {
+			return err
+		}
+		manifest, err := filepath.Abs(filepath.Join(buildContext, fileName))
+		if err != nil {
+			return errors.Wrap(err, "failed to get absolute path of the build file")
+		}
+		if manifest == "" {
+			return errors.Newf("build file %s does not exist", fileName)
+		}
+		opt := builder.Options{
+			ManifestFilePath: manifest,
+			BuildContextDir:  buildContext,
+			BuildFuncName:    funcName,
+		}
+		builder, err := builder.New(clicontext.Context, opt)
+		if err != nil {
+			return errors.Wrap(err, "failed to create the builder")
+		}
+		if err := builder.Interpret(); err != nil {
+			return errors.Wrap(err, "failed to interpret the build file")
+		}
+		if cmd, ok := ir.DefaultGraph.RuntimeCommands[command]; !ok {
+			return errors.Newf("command %s does not exist", command)
+		} else {
+			resultCommand = cmd
+		}
+		// Get the environment name if `name` is not specified.
+		if name == "" {
+			name = fileutil.Base(path)
+		}
+	}
 
 	// Check if the container is running.
 	dockerClient, err := docker.NewClient(clicontext.Context)
@@ -69,8 +137,10 @@ func run(clicontext *cli.Context) error {
 	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")
+	if bytes, err := sshClient.ExecWithOutput(resultCommand); err != nil {
+		fmt.Fprintln(clicontext.App.Writer, string(bytes))
+		return errors.Wrapf(err,
+			"failed to execute the command `%s`", resultCommand)
 	} else {
 		fmt.Fprint(clicontext.App.Writer, string(bytes))
 	}
diff --git a/pkg/builder/builder.go b/pkg/builder/builder.go
index b6eeb0fbd..c761e3da0 100644
--- a/pkg/builder/builder.go
+++ b/pkg/builder/builder.go
@@ -40,6 +40,7 @@ import (
 
 type Builder interface {
 	Build(ctx context.Context, force bool) error
+	Interpret() error
 	GPUEnabled() bool
 	NumGPUs() int
 }
@@ -159,20 +160,23 @@ func (b generalBuilder) Build(ctx context.Context, force bool) error {
 	return nil
 }
 
-func (b generalBuilder) interpret() error {
+func (b generalBuilder) Interpret() error {
 	// Evaluate config first.
-	if _, err := b.ExecFile(b.ConfigFilePath, ""); err != nil {
-		return errors.Wrap(err, "failed to exec starlark file")
+	if b.ConfigFilePath != "" {
+		b.logger.Debug("evaluating config file")
+		if _, err := b.ExecFile(b.ConfigFilePath, ""); err != nil {
+			return errors.Wrapf(err, "failed to exec starlark file %s", b.ConfigFilePath)
+		}
 	}
 
 	if _, err := b.ExecFile(b.ManifestFilePath, b.BuildFuncName); err != nil {
-		return errors.Wrap(err, "failed to exec starlark file")
+		return errors.Wrapf(err, "failed to exec starlark file %s", b.ManifestFilePath)
 	}
 	return nil
 }
 
 func (b generalBuilder) compile(ctx context.Context) (*llb.Definition, error) {
-	if err := b.interpret(); err != nil {
+	if err := b.Interpret(); err != nil {
 		return nil, errors.Wrap(err, "failed to interpret")
 	}
 	def, err := ir.Compile(ctx, fileutil.Base(b.BuildContextDir), b.PubKeyPath)
diff --git a/pkg/lang/frontend/starlark/interpreter.go b/pkg/lang/frontend/starlark/interpreter.go
index 16b619bea..dbdd7292d 100644
--- a/pkg/lang/frontend/starlark/interpreter.go
+++ b/pkg/lang/frontend/starlark/interpreter.go
@@ -28,6 +28,7 @@ import (
 	"github.com/tensorchord/envd/pkg/lang/frontend/starlark/config"
 	"github.com/tensorchord/envd/pkg/lang/frontend/starlark/install"
 	"github.com/tensorchord/envd/pkg/lang/frontend/starlark/io"
+	"github.com/tensorchord/envd/pkg/lang/frontend/starlark/runtime"
 	"github.com/tensorchord/envd/pkg/lang/frontend/starlark/universe"
 )
 
@@ -55,6 +56,7 @@ func NewInterpreter(buildContextDir string) Interpreter {
 			"install": install.Module,
 			"config":  config.Module,
 			"io":      io.Module,
+			"runtime": runtime.Module,
 		},
 		buildContextDir: buildContextDir,
 	}
diff --git a/pkg/lang/frontend/starlark/runtime/const.go b/pkg/lang/frontend/starlark/runtime/const.go
new file mode 100644
index 000000000..8feeba927
--- /dev/null
+++ b/pkg/lang/frontend/starlark/runtime/const.go
@@ -0,0 +1,19 @@
+// 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 runtime
+
+const (
+	ruleCommand = "runtime.command"
+)
diff --git a/pkg/lang/frontend/starlark/runtime/runtime.go b/pkg/lang/frontend/starlark/runtime/runtime.go
new file mode 100644
index 000000000..31abb554f
--- /dev/null
+++ b/pkg/lang/frontend/starlark/runtime/runtime.go
@@ -0,0 +1,62 @@
+// 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 runtime
+
+import (
+	"fmt"
+
+	"github.com/sirupsen/logrus"
+	"go.starlark.net/starlark"
+	"go.starlark.net/starlarkstruct"
+
+	"github.com/tensorchord/envd/pkg/lang/ir"
+)
+
+var (
+	logger = logrus.WithField("frontend", "starlark")
+)
+
+var Module = &starlarkstruct.Module{
+	Name: "runtime",
+	Members: starlark.StringDict{
+		"command": starlark.NewBuiltin(ruleCommand, ruleFuncCommand),
+	},
+}
+
+func ruleFuncCommand(thread *starlark.Thread, _ *starlark.Builtin,
+	args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
+	var commands starlark.IterableMapping
+
+	if err := starlark.UnpackArgs(ruleCommand, args, kwargs,
+		"commands?", &commands); err != nil {
+		return nil, err
+	}
+
+	commandsMap := make(map[string]string)
+	for _, tuple := range commands.Items() {
+		if len(tuple) != 2 {
+			return nil, fmt.Errorf("invalid command in %s", ruleCommand)
+		}
+
+		commandsMap[tuple[0].(starlark.String).GoString()] =
+			tuple[1].(starlark.String).GoString()
+	}
+
+	logger.Debugf("rule `%s` is invoked, commands: %v",
+		ruleCommand, commandsMap)
+
+	ir.RuntimeCommands(commandsMap)
+	return starlark.None, nil
+}
diff --git a/pkg/lang/ir/compile.go b/pkg/lang/ir/compile.go
index 4f4d3d974..cb190eae0 100644
--- a/pkg/lang/ir/compile.go
+++ b/pkg/lang/ir/compile.go
@@ -42,12 +42,13 @@ func NewGraph() *Graph {
 		CUDNN:   nil,
 		NumGPUs: -1,
 
-		PyPIPackages:   []string{},
-		RPackages:      []string{},
-		JuliaPackages:  []string{},
-		SystemPackages: []string{},
-		Exec:           []string{},
-		Shell:          shellBASH,
+		PyPIPackages:    []string{},
+		RPackages:       []string{},
+		JuliaPackages:   []string{},
+		SystemPackages:  []string{},
+		Exec:            []string{},
+		RuntimeCommands: make(map[string]string),
+		Shell:           shellBASH,
 	}
 }
 
diff --git a/pkg/lang/ir/interface.go b/pkg/lang/ir/interface.go
index 601538215..177e27181 100644
--- a/pkg/lang/ir/interface.go
+++ b/pkg/lang/ir/interface.go
@@ -203,3 +203,9 @@ func Mount(src, dest string) {
 		Destination: dest,
 	})
 }
+
+func RuntimeCommands(commands map[string]string) {
+	for k, v := range commands {
+		DefaultGraph.RuntimeCommands[k] = v
+	}
+}
diff --git a/pkg/lang/ir/types.go b/pkg/lang/ir/types.go
index 990bf6986..8f7396daa 100644
--- a/pkg/lang/ir/types.go
+++ b/pkg/lang/ir/types.go
@@ -50,9 +50,10 @@ type Graph struct {
 
 	VSCodePlugins []vscode.Plugin
 
-	Exec  []string
-	Copy  []CopyInfo
-	Mount []MountInfo
+	Exec            []string
+	Copy            []CopyInfo
+	Mount           []MountInfo
+	RuntimeCommands map[string]string
 
 	*JupyterConfig
 	*GitConfig
diff --git a/pkg/ssh/ssh.go b/pkg/ssh/ssh.go
index d8b4ca5d2..70f103c55 100644
--- a/pkg/ssh/ssh.go
+++ b/pkg/ssh/ssh.go
@@ -161,7 +161,7 @@ func (c generalClient) ExecWithOutput(cmd string) ([]byte, error) {
 		return nil, errors.Wrap(err, "requesting agent forwarding failed")
 	}
 
-	return session.Output(cmd)
+	return session.CombinedOutput(cmd)
 }
 
 func (c generalClient) Attach() error {


Test patch:
diff --git a/e2e/e2e_helper.go b/e2e/e2e_helper.go
index e5de41596..97d129b2e 100644
--- a/e2e/e2e_helper.go
+++ b/e2e/e2e_helper.go
@@ -83,7 +83,23 @@ func NewExample(name string, testcaseAbbr string) *Example {
 
 func (e *Example) Exec(cmd string) (string, error) {
 	args := []string{
-		"envd.test", "run", "--name", e.Name, "--command", cmd,
+		"envd.test", "run", "--name", e.Name, "--raw", cmd,
+	}
+
+	buffer := new(bytes.Buffer)
+	e.app.Writer = buffer
+
+	err := e.app.Run(args)
+	if err != nil {
+		return "", errors.Wrap(err, "failed to start `run` command")
+	}
+	return strings.Trim(buffer.String(), "\n"), nil
+}
+
+func (e *Example) ExecRuntimeCommand(cmd string) (string, error) {
+	buildContext := "testdata/" + e.Name
+	args := []string{
+		"envd.test", "run", "-p", buildContext, "--command", cmd,
 	}
 
 	buffer := new(bytes.Buffer)
diff --git a/e2e/runtime_test.go b/e2e/runtime_test.go
new file mode 100644
index 000000000..45d2dbe64
--- /dev/null
+++ b/e2e/runtime_test.go
@@ -0,0 +1,35 @@
+// 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 e2e
+
+import (
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+)
+
+var _ = Describe("runtime", Ordered, func() {
+	exampleName := "runtime"
+	testcase := "e2e"
+	e := NewExample(exampleName, testcase)
+	BeforeAll(e.BuildImage(true))
+	BeforeEach(e.RunContainer())
+	It("execute runtime command `numpy`", func() {
+		res, err := e.ExecRuntimeCommand("numpy")
+		Expect(err).To(BeNil())
+		Expect(res).To(Equal("[2 3 4]"))
+	})
+	AfterEach(e.DestroyContainer())
+	AfterAll(e.RemoveImage())
+})
diff --git a/e2e/testdata/runtime/build.envd b/e2e/testdata/runtime/build.envd
new file mode 100644
index 000000000..d2032151e
--- /dev/null
+++ b/e2e/testdata/runtime/build.envd
@@ -0,0 +1,9 @@
+def build():
+    base(os="ubuntu20.04", language="python3")
+    install.python_packages(name = [
+        "numpy",
+    ])
+    shell("zsh")
+    runtime.command(commands={
+        "numpy": "python demo.py"
+    })
diff --git a/e2e/testdata/runtime/demo.py b/e2e/testdata/runtime/demo.py
new file mode 100644
index 000000000..5ea2b3da4
--- /dev/null
+++ b/e2e/testdata/runtime/demo.py
@@ -0,0 +1,5 @@
+import numpy as np
+
+a = np.array([2, 3, 4])
+
+print(a)
\ No newline at end of file