feature: support runtime init script (#1085)

* stash

Signed-off-by: Jinjing.Zhou <allenzhou@tensorchord.ai>

* add

Signed-off-by: Jinjing.Zhou <allenzhou@tensorchord.ai>

* fix

Signed-off-by: Jinjing.Zhou <allenzhou@tensorchord.ai>

* add example

Signed-off-by: Jinjing.Zhou <allenzhou@tensorchord.ai>

* fix

Signed-off-by: Jinjing.Zhou <allenzhou@tensorchord.ai>

* fix

Signed-off-by: Jinjing.Zhou <allenzhou@tensorchord.ai>

Signed-off-by: Jinjing.Zhou <allenzhou@tensorchord.ai>

envd/api/runtime/__init__.py

@@ -93,3 +93,11 @@ def mount(host_path: str, envd_path: str):
         host_path (str): source path in the host machine
         envd_path (str): destination path in the envd container
     """
+
+
+def init(commands: List[str]):
+    """Commands to be executed when start the container
+
+    Args:
+        commands (List[str]): list of commands
+    """

examples/dpgen2/build.envd

@@ -0,0 +1,40 @@
+def install_kubectl_kind(env_name):
+    run(
+        [
+            'curl -LO "https://dl.k8s.io/release/$(curl -L -s https://dl.k8s.io/release/stable.txt)/bin/linux/amd64/kubectl"',
+            "sudo install -o root -g root -m 0755 kubectl /usr/local/bin/kubectl",
+        ]
+    )
+    run(
+        [
+            "curl -Lo ./kind https://kind.sigs.k8s.io/dl/v0.16.0/kind-linux-amd64",
+            "chmod +x ./kind",
+            "sudo mv ./kind /usr/local/bin/kind",
+        ]
+    )
+    runtime.init(
+        [
+            "sudo kind create cluster",
+            "sudo docker network connect kind {}".format(env_name),
+            'sudo kind get kubeconfig --name "kind" --internal | sed "s/kind-control-plane/$(sudo docker inspect "kind-control-plane" --format "{{ .NetworkSettings.Networks.kind.IPAddress }}")/g" | tee /home/envd/.kindkubeconfig',
+        ]
+    )
+    runtime.environ(env={"KUBECONFIG": "/home/envd/.kindkubeconfig"})
+
+
+def build():
+    base(os="ubuntu20.04", language="python3.9")
+    install.python_packages(
+        name=[
+            "pydflow",
+        ]
+    )
+    shell("zsh")
+    run(
+        [
+            "pip install git+https://github.com/deepmodeling/dpgen2 -i https://pypi.tuna.tsinghua.edu.cn/simple"
+        ]
+    )
+    run(["curl -sSL https://get.daocloud.io/docker | sh"])
+    runtime.mount("/var/run/docker.sock", "/var/run/docker.sock")
+    install_kubectl_kind("dpgen2")

examples/include_pkg/build.envd

@@ -1,4 +1,4 @@
-envdlib = include("https://github.com/kemingy/envdlib")
+envdlib = include("https://github.com/tensorchord/envdlib")
 
 
 def build():

pkg/lang/frontend/starlark/runtime/const.go

@@ -15,9 +15,10 @@
 package runtime
 
 const (
-	ruleCommand = "runtime.command"
-	ruleExpose  = "runtime.expose"
-	ruleDaemon  = "runtime.daemon"
-	ruleEnviron = "runtime.environ"
-	ruleMount   = "runtime.mount"
+	ruleCommand    = "runtime.command"
+	ruleExpose     = "runtime.expose"
+	ruleDaemon     = "runtime.daemon"
+	ruleEnviron    = "runtime.environ"
+	ruleMount      = "runtime.mount"
+	ruleInitScript = "runtime.init"
 )

pkg/lang/frontend/starlark/runtime/runtime.go

@@ -27,6 +27,7 @@ import (
 	"github.com/tensorchord/envd/pkg/lang/frontend/starlark/data"
 	"github.com/tensorchord/envd/pkg/lang/ir"
 	"github.com/tensorchord/envd/pkg/util/fileutil"
+	"github.com/tensorchord/envd/pkg/util/starlarkutil"
 )
 
 var (
@@ -41,6 +42,7 @@ var Module = &starlarkstruct.Module{
 		"expose":  starlark.NewBuiltin(ruleExpose, ruleFuncExpose),
 		"environ": starlark.NewBuiltin(ruleEnviron, ruleFuncEnviron),
 		"mount":   starlark.NewBuiltin(ruleMount, ruleFuncMount),
+		"init":    starlark.NewBuiltin(ruleInitScript, ruleFuncInitScript),
 	},
 }
 
@@ -199,3 +201,24 @@ func ruleFuncMount(thread *starlark.Thread, _ *starlark.Builtin,
 
 	return starlark.None, nil
 }
+
+func ruleFuncInitScript(thread *starlark.Thread, _ *starlark.Builtin,
+	args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
+	var commands *starlark.List
+
+	if err := starlark.UnpackArgs(ruleCommand, args, kwargs,
+		"commands?", &commands); err != nil {
+		return nil, err
+	}
+
+	commandsSlice, err := starlarkutil.ToStringSlice(commands)
+	if err != nil {
+		return nil, err
+	}
+
+	logger.Debugf("rule `%s` is invoked, commands: %v",
+		ruleInitScript, commandsSlice)
+
+	ir.RuntimeInitScript(commandsSlice)
+	return starlark.None, nil
+}

pkg/lang/ir/interface.go

@@ -225,3 +225,7 @@ func RuntimeEnviron(env map[string]string) {
 		DefaultGraph.RuntimeEnviron[k] = v
 	}
 }
+
+func RuntimeInitScript(commands []string) {
+	DefaultGraph.RuntimeInitScript = append(DefaultGraph.RuntimeInitScript, commands)
+}

pkg/lang/ir/supervisor.go

@@ -28,11 +28,14 @@ import (
 const (
 	horustTemplate = `
 name = "%[1]s"
-command = "%[2]s"
-stdout = "/var/logs/%[1]s_stdout.log"
-stderr = "/var/logs/%[1]s_stderr.log"
+command = """
+%[2]s
+"""
+stdout = "/var/log/horust/%[1]s_stdout.log"
+stderr = "/var/log/horust/%[1]s_stderr.log"
 user = "${USER}"
 working-directory = "${%[3]s}"
+%[4]s
 
 [environment]
 keep-env = true
@@ -41,15 +44,26 @@ re-export = [ "PATH", "SHELL", "USER", "%[3]s", "ENVD_AUTHORIZED_KEYS_PATH", "EN
 [restart]
 strategy = "on-failure"
 backoff = "1s"
-attempts = 5
+attempts = 2
 
 [termination]
 wait = "5s"
 `
 )
 
-func (g Graph) addNewProcess(root llb.State, name, command string) llb.State {
-	template := fmt.Sprintf(horustTemplate, name, command, types.EnvdWorkDir)
+func (g Graph) addNewProcess(root llb.State, name, command string, depends []string) llb.State {
+	var sb strings.Builder
+	if len(depends) != 0 {
+		sb.WriteString("start-after = [")
+		for _, d := range depends {
+			sb.WriteString("\"")
+			sb.WriteString(d)
+			sb.WriteString("\",")
+		}
+		sb.WriteString("]\n")
+	}
+	template := fmt.Sprintf(horustTemplate, name, command, types.EnvdWorkDir, sb.String())
+
 	filename := filepath.Join(types.HorustServiceDir, fmt.Sprintf("%s.toml", name))
 	supervisor := root.File(llb.Mkfile(filename, 0644, []byte(template), llb.WithUIDGID(g.uid, g.gid)), llb.WithCustomNamef("[internal] create file %s", filename))
 	return supervisor
@@ -60,22 +74,29 @@ func (g Graph) compileEntrypoint(root llb.State) llb.State {
 		return root
 	}
 	cmd := fmt.Sprintf("/var/envd/bin/envd-sshd --port %d --shell %s", config.SSHPortInContainer, g.Shell)
-	entrypoint := g.addNewProcess(root, "sshd", cmd)
+	entrypoint := g.addNewProcess(root, "sshd", cmd, nil)
+	var deps []string
+	if g.RuntimeInitScript != nil {
+		for i, command := range g.RuntimeInitScript {
+			entrypoint = g.addNewProcess(entrypoint, fmt.Sprintf("init_%d", i), fmt.Sprintf("/bin/bash -c 'set -euo pipefail\n%s'", strings.Join(command, "\n")), nil)
+			deps = append(deps, fmt.Sprintf("init_%d", i))
+		}
+	}
 
 	if g.RuntimeDaemon != nil {
 		for i, command := range g.RuntimeDaemon {
-			entrypoint = g.addNewProcess(entrypoint, fmt.Sprintf("daemon_%d", i), fmt.Sprintf("%s &\n", strings.Join(command, " ")))
+			entrypoint = g.addNewProcess(entrypoint, fmt.Sprintf("daemon_%d", i), fmt.Sprintf("%s &\n", strings.Join(command, " ")), deps)
 		}
 	}
 
 	if g.JupyterConfig != nil {
 		jupyterCmd := g.generateJupyterCommand("")
-		entrypoint = g.addNewProcess(entrypoint, "jupyter", strings.Join(jupyterCmd, " "))
+		entrypoint = g.addNewProcess(entrypoint, "jupyter", strings.Join(jupyterCmd, " "), deps)
 	}
 
 	if g.RStudioServerConfig != nil {
 		rstudioCmd := g.generateRStudioCommand("")
-		entrypoint = g.addNewProcess(entrypoint, "rstudio", strings.Join(rstudioCmd, " "))
+		entrypoint = g.addNewProcess(entrypoint, "rstudio", strings.Join(rstudioCmd, " "), deps)
 	}
 
 	return entrypoint

pkg/lang/ir/types.go

@@ -77,10 +77,11 @@ type Graph struct {
 
 // The results during runtime should be maintained here
 type RuntimeGraph struct {
-	RuntimeCommands map[string]string
-	RuntimeDaemon   [][]string
-	RuntimeEnviron  map[string]string
-	RuntimeExpose   []ExposeItem
+	RuntimeCommands   map[string]string
+	RuntimeDaemon     [][]string
+	RuntimeInitScript [][]string
+	RuntimeEnviron    map[string]string
+	RuntimeExpose     []ExposeItem
 }
 
 type CopyInfo struct {

pkg/types/envd.go

@@ -41,7 +41,7 @@ const (
 	// supervisor
 	HorustImage      = "tensorchord/horust:v0.1.0"
 	HorustServiceDir = "/etc/horust/services"
-	HorustLogDir     = "/var/logs"
+	HorustLogDir     = "/var/log/horust"
 	// env
 	EnvdWorkDir = "ENVD_WORKDIR"
 )