feat(lang): Support R language (#491)

Signed-off-by: Ce Gao <cegao@tensorchord.ai>
tensorchord · Jun 22, 2022 · 7127365 · 7127365
1 parent 3f86086
commit 7127365
Show file tree

Hide file tree

Showing 5 changed files with 108 additions and 51 deletions.
diff --git a/base-images/build.sh b/base-images/build.sh
@@ -37,7 +37,7 @@ docker buildx build \
     --build-arg ENVD_SSH_IMAGE=ghcr.io/tensorchord/envd-ssh-from-scratch \
     --build-arg HTTP_PROXY=${HTTP_PROXY} \
     --build-arg HTTPS_PROXY=${HTTPS_PROXY} \
-    -t terrytangyuan/r4.2-envd:0.0.1  \
+    -t ${DOCKER_HUB_ORG}/r-base:4.2  \
     --pull --push --platform linux/x86_64,linux/arm64 \
     -f r4.2.Dockerfile .
 

diff --git a/examples/r-basic/build.envd b/examples/r-basic/build.envd
@@ -4,3 +4,4 @@ def build():
             "remotes",
             "rlang",
         ])
+    shell("zsh")
diff --git a/pkg/lang/ir/compile.go b/pkg/lang/ir/compile.go
@@ -131,58 +131,16 @@ func (g Graph) Compile(uid, gid int) (llb.State, error) {
 	base := g.compileBase()
 	aptStage := g.compileUbuntuAPT(base)
 	var merged llb.State
+	var err error
 	if g.Language.Name == "r" {
-		// TODO(terrytangyuan): Support RStudio local server
-		rPackageInstallStage := llb.Diff(aptStage, g.installRPackages(aptStage), llb.WithCustomName("install R packages"))
-		merged = llb.Merge([]llb.State{
-			aptStage, rPackageInstallStage,
-		}, llb.WithCustomName("merging all components into one"))
-	} else {
-		condaChanelStage := g.compileCondaChannel(aptStage)
-		pypiMirrorStage := g.compilePyPIIndex(condaChanelStage)
-
-		g.compileJupyter()
-		builtinSystemStage := pypiMirrorStage
-
-		sshStage, err := g.copySSHKey(builtinSystemStage)
-		if err != nil {
-			return llb.State{}, errors.Wrap(err, "failed to copy ssh keys")
-		}
-		diffSSHStage := llb.Diff(builtinSystemStage, sshStage, llb.WithCustomName("install ssh keys"))
-
-		// Conda affects shell and python, thus we cannot do it in parallel.
-		shellStage, err := g.compileShell(builtinSystemStage)
+		merged, err = g.compileRLang(aptStage)
 		if err != nil {
-			return llb.State{}, errors.Wrap(err, "failed to compile shell")
+			return llb.State{}, errors.Wrap(err, "failed to compile r language")
 		}
-
-		condaEnvStage := g.setCondaENV(shellStage)
-
-		condaStage := llb.Diff(builtinSystemStage,
-			g.compileCondaPackages(condaEnvStage),
-			llb.WithCustomName("install conda packages"))
-
-		pypiStage := llb.Diff(condaEnvStage,
-			g.compilePyPIPackages(condaEnvStage),
-			llb.WithCustomName("install PyPI packages"))
-		systemStage := llb.Diff(builtinSystemStage, g.compileSystemPackages(builtinSystemStage),
-			llb.WithCustomName("install system packages"))
-
-		vscodeStage, err := g.compileVSCode()
+	} else {
+		merged, err = g.compilePython(aptStage)
 		if err != nil {
-			return llb.State{}, errors.Wrap(err, "failed to get vscode plugins")
-		}
-
-		if vscodeStage != nil {
-			merged = llb.Merge([]llb.State{
-				builtinSystemStage, systemStage, condaStage,
-				diffSSHStage, pypiStage, *vscodeStage,
-			}, llb.WithCustomName("merging all components into one"))
-		} else {
-			merged = llb.Merge([]llb.State{
-				builtinSystemStage, systemStage, condaStage,
-				diffSSHStage, pypiStage,
-			}, llb.WithCustomName("merging all components into one"))
+			return llb.State{}, errors.Wrap(err, "failed to compile python")
 		}
 	}
 
@@ -195,3 +153,98 @@ func (g Graph) Compile(uid, gid int) (llb.State, error) {
 	g.Writer.Finish()
 	return finalStage, nil
 }
+
+func (g Graph) compileRLang(aptStage llb.State) (llb.State, error) {
+	g.compileJupyter()
+	builtinSystemStage := aptStage
+
+	sshStage, err := g.copySSHKey(builtinSystemStage)
+	if err != nil {
+		return llb.State{}, errors.Wrap(err, "failed to copy ssh keys")
+	}
+	diffSSHStage := llb.Diff(builtinSystemStage, sshStage, llb.WithCustomName("install ssh keys"))
+
+	// Conda affects shell and python, thus we cannot do it in parallel.
+	shellStage, err := g.compileShell(builtinSystemStage)
+	if err != nil {
+		return llb.State{}, errors.Wrap(err, "failed to compile shell")
+	}
+	diffShellStage := llb.Diff(builtinSystemStage, shellStage, llb.WithCustomName("install shell"))
+
+	systemStage := llb.Diff(builtinSystemStage, g.compileSystemPackages(builtinSystemStage),
+		llb.WithCustomName("install system packages"))
+
+	// TODO(terrytangyuan): Support RStudio local server
+	rPackageInstallStage := llb.Diff(builtinSystemStage,
+		g.installRPackages(builtinSystemStage), llb.WithCustomName("install R packages"))
+
+	vscodeStage, err := g.compileVSCode()
+	if err != nil {
+		return llb.State{}, errors.Wrap(err, "failed to get vscode plugins")
+	}
+
+	var merged llb.State
+	if vscodeStage != nil {
+		merged = llb.Merge([]llb.State{
+			builtinSystemStage, systemStage, diffShellStage,
+			diffSSHStage, rPackageInstallStage, *vscodeStage,
+		}, llb.WithCustomName("merging all components into one"))
+	} else {
+		merged = llb.Merge([]llb.State{
+			builtinSystemStage, systemStage, diffShellStage,
+			diffSSHStage, rPackageInstallStage,
+		}, llb.WithCustomName("merging all components into one"))
+	}
+	return merged, nil
+}
+
+func (g Graph) compilePython(aptStage llb.State) (llb.State, error) {
+	condaChanelStage := g.compileCondaChannel(aptStage)
+	pypiMirrorStage := g.compilePyPIIndex(condaChanelStage)
+
+	g.compileJupyter()
+	builtinSystemStage := pypiMirrorStage
+
+	sshStage, err := g.copySSHKey(builtinSystemStage)
+	if err != nil {
+		return llb.State{}, errors.Wrap(err, "failed to copy ssh keys")
+	}
+	diffSSHStage := llb.Diff(builtinSystemStage, sshStage, llb.WithCustomName("install ssh keys"))
+
+	// Conda affects shell and python, thus we cannot do it in parallel.
+	shellStage, err := g.compileShell(builtinSystemStage)
+	if err != nil {
+		return llb.State{}, errors.Wrap(err, "failed to compile shell")
+	}
+
+	condaEnvStage := g.setCondaENV(shellStage)
+
+	condaStage := llb.Diff(builtinSystemStage,
+		g.compileCondaPackages(condaEnvStage),
+		llb.WithCustomName("install conda packages"))
+
+	pypiStage := llb.Diff(condaEnvStage,
+		g.compilePyPIPackages(condaEnvStage),
+		llb.WithCustomName("install PyPI packages"))
+	systemStage := llb.Diff(builtinSystemStage, g.compileSystemPackages(builtinSystemStage),
+		llb.WithCustomName("install system packages"))
+
+	vscodeStage, err := g.compileVSCode()
+	if err != nil {
+		return llb.State{}, errors.Wrap(err, "failed to get vscode plugins")
+	}
+
+	var merged llb.State
+	if vscodeStage != nil {
+		merged = llb.Merge([]llb.State{
+			builtinSystemStage, systemStage, condaStage,
+			diffSSHStage, pypiStage, *vscodeStage,
+		}, llb.WithCustomName("merging all components into one"))
+	} else {
+		merged = llb.Merge([]llb.State{
+			builtinSystemStage, systemStage, condaStage,
+			diffSSHStage, pypiStage,
+		}, llb.WithCustomName("merging all components into one"))
+	}
+	return merged, nil
+}
diff --git a/pkg/lang/ir/r.go b/pkg/lang/ir/r.go
@@ -41,6 +41,6 @@ func (g Graph) installRPackages(root llb.State) llb.State {
 	cmd := sb.String()
 	root = llb.User("envd")(root)
 	run := root.
-		Run(llb.Shlex(cmd), llb.WithCustomNamef("R package install"))
+		Run(llb.Shlex(cmd), llb.WithCustomNamef("install R packages"))
 	return run.Root()
 }
diff --git a/pkg/lang/ir/system.go b/pkg/lang/ir/system.go
@@ -98,12 +98,15 @@ func (g *Graph) compileBase() llb.State {
 	var base llb.State
 	if g.CUDA == nil && g.CUDNN == nil {
 		if g.Language.Name == "r" {
-			base = llb.Image("docker.io/terrytangyuan/r4.2-envd:0.0.1")
+			base = llb.Image("docker.io/tensorchord/r-base:4.2")
 			// r-base image already has GID 1000.
 			// It is a trick, we actually use GID 1000
 			if g.gid == 1000 {
 				g.gid = 1001
 			}
+			if g.uid == 1000 {
+				g.uid = 1001
+			}
 		} else {
 			base = llb.Image("docker.io/tensorchord/python:3.8-ubuntu20.04")
 		}