feat(lang): Support image in base (#595)

feat(lang): Support base image customization Signed-off-by: Ce Gao <cegao@tensorchord.ai>
tensorchord · Jul 13, 2022 · 9b3fbe3 · 9b3fbe3
1 parent b467279
commit 9b3fbe3
Show file tree

Hide file tree

Showing 6 changed files with 30 additions and 11 deletions.
diff --git a/examples/custom-image/build.envd b/examples/custom-image/build.envd
@@ -0,0 +1,5 @@
+def build():
+    base(language="python", image="python:3.8")
+    install.python_packages(name=[
+      "ormb",
+    ])
diff --git a/pkg/lang/frontend/starlark/universe/universe.go b/pkg/lang/frontend/starlark/universe/universe.go
@@ -40,20 +40,21 @@ func RegisterBuildContext(buildContextDir string) {
 
 func ruleFuncBase(thread *starlark.Thread, _ *starlark.Builtin,
 	args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
-	var os, language starlark.String
+	var os, language, image starlark.String
 
 	if err := starlark.UnpackArgs(ruleBase, args, kwargs,
-		"os?", &os, "language?", &language); err != nil {
+		"os?", &os, "language?", &language, "image?", &image); err != nil {
 		return nil, err
 	}
 
 	osStr := os.GoString()
 	langStr := language.GoString()
+	imageStr := image.GoString()
 
-	logger.Debugf("rule `%s` is invoked, os=%s, language=%s",
-		ruleBase, osStr, langStr)
+	logger.Debugf("rule `%s` is invoked, os=%s, language=%s, image=%s",
+		ruleBase, osStr, langStr, imageStr)
 
-	err := ir.Base(osStr, langStr)
+	err := ir.Base(osStr, langStr, imageStr)
 	return starlark.None, err
 }
 

diff --git a/pkg/lang/ir/compile.go b/pkg/lang/ir/compile.go
@@ -140,6 +140,11 @@ func (g Graph) ExposedPorts() (map[string]struct{}, error) {
 }
 
 func (g Graph) Entrypoint(buildContextDir string) ([]string, error) {
+	// Do not set entrypoint if the image is customized.
+	if g.Image != nil {
+		return []string{}, nil
+	}
+
 	ep := []string{
 		"tini",
 		"--",

diff --git a/pkg/lang/ir/interface.go b/pkg/lang/ir/interface.go
@@ -21,7 +21,7 @@ import (
 	"github.com/tensorchord/envd/pkg/lang/ir/parser"
 )
 
-func Base(os, language string) error {
+func Base(os, language, image string) error {
 	l, version, err := parseLanguage(language)
 	if err != nil {
 		return err
@@ -31,6 +31,9 @@ func Base(os, language string) error {
 		Version: version,
 	}
 	DefaultGraph.OS = os
+	if image != "" {
+		DefaultGraph.Image = &image
+	}
 	return nil
 }
 

diff --git a/pkg/lang/ir/system.go b/pkg/lang/ir/system.go
@@ -98,7 +98,9 @@ func (g *Graph) compileBase() llb.State {
 	logger.Debug("compile base image")
 
 	var base llb.State
-	if g.CUDA == nil && g.CUDNN == nil {
+	if g.Image != nil {
+		base = llb.Image(*g.Image)
+	} else if g.CUDA == nil && g.CUDNN == nil {
 		switch g.Language.Name {
 		case "r":
 			base = llb.Image("docker.io/tensorchord/r-base:4.2")
@@ -134,7 +136,7 @@ func (g *Graph) compileBase() llb.State {
 				llb.WithCustomName("[internal] set root home dir to /home/envd")).
 			Run(llb.Shlex("sed -i \"s/envd:x:1001/envd:x:0/g\" /etc/group"),
 				llb.WithCustomName("[internal] set envd group to 0 as root group"))
-		if g.Language.Name == "python" {
+		if g.Image == nil && g.Language.Name == "python" {
 			res = res.Run(llb.Shlex("chown -R root:root /opt/conda"),
 				llb.WithCustomName("[internal] configure user permissions"))
 		}
@@ -148,7 +150,7 @@ func (g *Graph) compileBase() llb.State {
 				llb.WithCustomName("[internal] add user envd to sudoers")).
 			Run(llb.Shlex("chown -R envd:envd /usr/local/lib"),
 				llb.WithCustomName("[internal] configure user permissions"))
-		if g.Language.Name == "python" {
+		if g.Image == nil && g.Language.Name == "python" {
 			res = res.Run(llb.Shlex("chown -R envd:envd /opt/conda"),
 				llb.WithCustomName("[internal] configure user permissions"))
 		}
@@ -157,15 +159,17 @@ func (g *Graph) compileBase() llb.State {
 }
 
 func (g Graph) copySSHKey(root llb.State) (llb.State, error) {
-	// TODO(gaocegege): Remove global var ssh image.
 	public := DefaultGraph.PublicKeyPath
 	bdat, err := os.ReadFile(public)
 	dat := strings.TrimSuffix(string(bdat), "\n")
 	if err != nil {
 		return llb.State{}, errors.Wrap(err, "Cannot read public SSH key")
 	}
 	run := root.
+		File(llb.Mkdir("/var/envd", 0755, llb.WithParents(true),
+			llb.WithUIDGID(g.uid, g.gid))).
 		File(llb.Mkfile(config.ContainerAuthorizedKeysPath,
-			0644, []byte(dat+" envd"), llb.WithUIDGID(g.uid, g.gid)), llb.WithCustomName("install ssh keys"))
+			0644, []byte(dat+" envd"), llb.WithUIDGID(g.uid, g.gid)),
+			llb.WithCustomName("install ssh keys"))
 	return run, nil
 }
diff --git a/pkg/lang/ir/types.go b/pkg/lang/ir/types.go
@@ -27,6 +27,7 @@ type Graph struct {
 
 	OS string
 	Language
+	Image *string
 
 	Shell   string
 	CUDA    *string