feat: use horust as the dev container supervisor (#1051)

* feat: use horust as the dev container supervisor Signed-off-by: Keming <kemingyang@tensorchord.ai> * fix daemon Signed-off-by: Keming <kemingyang@tensorchord.ai> * change to ENVD_WORKDIR Signed-off-by: Keming <kemingyang@tensorchord.ai> * keep env Signed-off-by: Keming <kemingyang@tensorchord.ai> * bypass custom image Signed-off-by: Keming <kemingyang@tensorchord.ai> * fix lint Signed-off-by: Keming <kemingyang@tensorchord.ai> * make env as a const Signed-off-by: Keming <kemingyang@tensorchord.ai> * change to tensorchord/horust:v0.1.0 Signed-off-by: Keming <kemingyang@tensorchord.ai> Signed-off-by: Keming <kemingyang@tensorchord.ai>
tensorchord · Oct 20, 2022 · cd5cf52 · cd5cf52
1 parent 044d89e
commit cd5cf52
Show file tree

Hide file tree

Showing 7 changed files with 133 additions and 55 deletions.
diff --git a/pkg/lang/ir/compile.go b/pkg/lang/ir/compile.go
@@ -20,7 +20,6 @@ import (
 	"fmt"
 	"os"
 	"path/filepath"
-	"strings"
 
 	"github.com/cockroachdb/errors"
 	"github.com/moby/buildkit/client/llb"
@@ -201,48 +200,12 @@ func (g Graph) DefaultCacheImporter() (*string, error) {
 	return &res, nil
 }
 
-func (g Graph) GetEntrypoint(buildContextDir string) ([]string, error) {
+func (g *Graph) GetEntrypoint(buildContextDir string) ([]string, error) {
 	if g.Image != nil {
 		return g.Entrypoint, nil
 	}
-
-	ep := []string{
-		"tini",
-		"--",
-		"bash",
-		"-c",
-	}
-
-	template := `set -euo pipefail
-/var/envd/bin/envd-sshd --port %d --shell %s &
-%s
-wait -n`
-
-	// Generate jupyter and rstudio server commands.
-	var customCmd strings.Builder
-	workingDir := fileutil.EnvdHomeDir(filepath.Base(buildContextDir))
-	if g.RuntimeDaemon != nil {
-		for _, command := range g.RuntimeDaemon {
-			customCmd.WriteString(fmt.Sprintf("%s &\n", strings.Join(command, " ")))
-		}
-	}
-	if g.JupyterConfig != nil {
-		jupyterCmd := g.generateJupyterCommand(workingDir)
-		customCmd.WriteString(strings.Join(jupyterCmd, " "))
-		customCmd.WriteString("\n")
-	}
-	if g.RStudioServerConfig != nil {
-		rstudioCmd := g.generateRStudioCommand(workingDir)
-		customCmd.WriteString(strings.Join(rstudioCmd, " "))
-		customCmd.WriteString("\n")
-	}
-
-	cmd := fmt.Sprintf(template,
-		config.SSHPortInContainer, g.Shell, customCmd.String())
-	ep = append(ep, cmd)
-
-	logrus.WithField("entrypoint", ep).Debug("generate entrypoint")
-	return ep, nil
+	g.RuntimeEnviron[types.EnvdWorkDir] = fileutil.EnvdHomeDir(filepath.Base(buildContextDir))
+	return []string{"horust"}, nil
 }
 
 func (g Graph) Compile(uid, gid int) (llb.State, error) {
@@ -297,7 +260,8 @@ func (g Graph) Compile(uid, gid int) (llb.State, error) {
 	// TODO(gaocegege): Support order-based exec.
 	run := g.compileRun(copy)
 	git := g.compileGit(run)
-	finalStage := g.compileUserOwn(git)
+	user := g.compileUserOwn(git)
+	entrypoint := g.compileEntrypoint(user)
 	g.Writer.Finish()
-	return finalStage, nil
+	return entrypoint, nil
 }
diff --git a/pkg/lang/ir/editor.go b/pkg/lang/ir/editor.go
@@ -15,6 +15,7 @@
 package ir
 
 import (
+	"fmt"
 	"strconv"
 
 	"github.com/cockroachdb/errors"
@@ -24,6 +25,7 @@ import (
 	"github.com/tensorchord/envd/pkg/editor/vscode"
 	"github.com/tensorchord/envd/pkg/flag"
 	"github.com/tensorchord/envd/pkg/progress/compileui"
+	"github.com/tensorchord/envd/pkg/types"
 	"github.com/tensorchord/envd/pkg/util/fileutil"
 )
 
@@ -79,6 +81,11 @@ func (g Graph) generateJupyterCommand(workingDir string) []string {
 		g.JupyterConfig.Token = "''"
 	}
 
+	// get from env if not set
+	if len(workingDir) == 0 {
+		workingDir = fmt.Sprintf("${%s}", types.EnvdWorkDir)
+	}
+
 	cmd := []string{
 		"python3", "-m", "notebook",
 		"--ip", "0.0.0.0", "--notebook-dir", workingDir,
@@ -99,6 +106,11 @@ func (g Graph) generateRStudioCommand(workingDir string) []string {
 		return nil
 	}
 
+	// get from env if not set
+	// if len(workingDir) == 0 {
+	// 	workingDir = fmt.Sprintf("${%s}", types.EnvdWorkDir)
+	// }
+
 	return []string{
 		// TODO(gaocegege): Remove root permission here.
 		"sudo",

diff --git a/pkg/lang/ir/shell.go b/pkg/lang/ir/shell.go
@@ -46,8 +46,10 @@ disabled = false
 
 func (g *Graph) compileShell(root llb.State) (llb.State, error) {
 	if g.Shell == shellZSH {
+		g.RuntimeEnviron["SHELL"] = "/usr/bin/zsh"
 		return g.compileZSH(root)
 	}
+	g.RuntimeEnviron["SHELL"] = "/usr/bin/bash"
 	return root, nil
 }
 

diff --git a/pkg/lang/ir/supervisor.go b/pkg/lang/ir/supervisor.go
@@ -0,0 +1,82 @@
+// 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 ir
+
+import (
+	"fmt"
+	"path/filepath"
+	"strings"
+
+	"github.com/moby/buildkit/client/llb"
+
+	"github.com/tensorchord/envd/pkg/config"
+	"github.com/tensorchord/envd/pkg/types"
+)
+
+const (
+	horustTemplate = `
+name = "%[1]s"
+command = "%[2]s"
+stdout = "/var/logs/%[1]s_stdout.log"
+stderr = "/var/logs/%[1]s_stderr.log"
+user = "${USER}"
+working-directory = "${%[3]s}"
+
+[environment]
+keep-env = true
+re-export = [ "PATH", "SHELL", "USER", "%[3]s" ]
+
+[restart]
+strategy = "on-failure"
+backoff = "1s"
+attempts = 5
+
+[termination]
+wait = "5s"
+`
+)
+
+func (g Graph) addNewProcess(root llb.State, name, command string) llb.State {
+	template := fmt.Sprintf(horustTemplate, name, command, types.EnvdWorkDir)
+	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)))
+	return supervisor
+}
+
+func (g Graph) compileEntrypoint(root llb.State) llb.State {
+	if g.Image != nil {
+		return root
+	}
+	cmd := fmt.Sprintf("/var/envd/bin/envd-sshd --port %d --shell %s", config.SSHPortInContainer, g.Shell)
+	entrypoint := g.addNewProcess(root, "sshd", cmd)
+
+	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, " ")))
+		}
+	}
+
+	if g.JupyterConfig != nil {
+		jupyterCmd := g.generateJupyterCommand("")
+		entrypoint = g.addNewProcess(entrypoint, "jupyter", strings.Join(jupyterCmd, " "))
+	}
+
+	if g.RStudioServerConfig != nil {
+		rstudioCmd := g.generateRStudioCommand("")
+		entrypoint = g.addNewProcess(entrypoint, "rstudio", strings.Join(rstudioCmd, " "))
+	}
+
+	return entrypoint
+}
diff --git a/pkg/lang/ir/system.go b/pkg/lang/ir/system.go
@@ -218,7 +218,17 @@ func (g *Graph) compileBase() (llb.State, error) {
 	if err != nil {
 		return llb.State{}, errors.Wrap(err, "failed to install conda")
 	}
-	return g.compileSshd(condaStage), nil
+	supervisor := g.installHorust(condaStage)
+	return g.compileSshd(supervisor), nil
+}
+
+func (g Graph) installHorust(root llb.State) llb.State {
+	horust := root.
+		File(llb.Copy(llb.Image(types.HorustImage), "/", "/usr/local/bin", llb.WithUIDGID(g.uid, g.gid)),
+			llb.WithCustomName("[internal] install horust")).
+		File(llb.Mkdir(types.HorustServiceDir, 0755, llb.WithParents(true), llb.WithUIDGID(g.uid, g.gid))).
+		File(llb.Mkdir(types.HorustLogDir, 0755, llb.WithParents(true), llb.WithUIDGID(g.uid, g.gid)))
+	return horust
 }
 
 func (g Graph) copySSHKey(root llb.State) (llb.State, error) {

diff --git a/pkg/lang/ir/user.go b/pkg/lang/ir/user.go
@@ -23,8 +23,10 @@ import (
 // compileUserOwn chown related directories
 func (g *Graph) compileUserOwn(root llb.State) llb.State {
 	if g.Image != nil || g.uid == 0 {
+		g.RuntimeEnviron["USER"] = "root"
 		return root
 	}
+	g.RuntimeEnviron["USER"] = "envd"
 	if len(g.UserDirectories) == 0 {
 		return root.User("envd")
 	}

diff --git a/pkg/types/envd.go b/pkg/types/envd.go
@@ -27,17 +27,24 @@ import (
 	"github.com/tensorchord/envd/pkg/version"
 )
 
-// DefaultPathEnvUnix is unix style list of directories to search for
-// executables. Each directory is separated from the next by a colon
-// ':' character .
-const DefaultPathEnvUnix = "/opt/conda/envs/envd/bin:/opt/conda/bin:/home/envd/.local/bin:/usr/local/julia/bin:" + system.DefaultPathEnvUnix
-
-// DefaultPathEnvWindows is windows style list of directories to search for
-// executables. Each directory is separated from the next by a colon
-// ';' character .
-const DefaultPathEnvWindows = system.DefaultPathEnvWindows
-
-const PythonBaseImage = "ubuntu:20.04"
+const (
+	// DefaultPathEnvUnix is unix style list of directories to search for
+	// executables. Each directory is separated from the next by a colon
+	// ':' character .
+	DefaultPathEnvUnix = "/opt/conda/envs/envd/bin:/opt/conda/bin:/home/envd/.local/bin:/usr/local/julia/bin:" + system.DefaultPathEnvUnix
+	// DefaultPathEnvWindows is windows style list of directories to search for
+	// executables. Each directory is separated from the next by a colon
+	// ';' character .
+	DefaultPathEnvWindows = system.DefaultPathEnvWindows
+	// image
+	PythonBaseImage = "ubuntu:20.04"
+	// supervisor
+	HorustImage      = "tensorchord/horust:v0.1.0"
+	HorustServiceDir = "/etc/horust/services"
+	HorustLogDir     = "/var/logs"
+	// env
+	EnvdWorkDir = "ENVD_WORKDIR"
+)
 
 var EnvdSshdImage = fmt.Sprintf(
 	"tensorchord/envd-sshd-from-scratch:%s",
@@ -71,7 +78,6 @@ var BaseAptPackage = []string{
 	"curl",
 	"openssh-client",
 	"git",
-	"tini",
 	"sudo",
 	"vim",
 	"zsh",