feat: Initial support for R language (#257)

* feat: Initial support for R language

Signed-off-by: terrytangyuan <terrytangyuan@gmail.com>

* fix groupID

Signed-off-by: terrytangyuan <terrytangyuan@gmail.com>

examples/r-basic/build.envd

@@ -0,0 +1,6 @@
+def build():
+    base(os="ubuntu20.04", language="r")
+    install.r_packages([
+            "remotes",
+            "rlang",
+        ])

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

@@ -17,6 +17,7 @@ package install
 const (
 	ruleSystemPackage = "install.system_packages"
 	rulePyPIPackage   = "install.python_packages"
+	ruleRPackage      = "install.r_packages"
 	ruleCUDA          = "install.cuda"
 	ruleVSCode        = "install.vscode_extensions"
 )

pkg/lang/frontend/starlark/install/install.go

@@ -30,6 +30,8 @@ var Module = &starlarkstruct.Module{
 	Members: starlark.StringDict{
 		"python_packages": starlark.NewBuiltin(
 			rulePyPIPackage, ruleFuncPyPIPackage),
+		"r_packages": starlark.NewBuiltin(
+			ruleRPackage, ruleFuncRPackage),
 		"system_packages": starlark.NewBuiltin(
 			ruleSystemPackage, ruleFuncSystemPackage),
 		"cuda":              starlark.NewBuiltin(ruleCUDA, ruleFuncCUDA),
@@ -60,6 +62,28 @@ func ruleFuncPyPIPackage(thread *starlark.Thread, _ *starlark.Builtin,
 	return starlark.None, nil
 }
 
+func ruleFuncRPackage(thread *starlark.Thread, _ *starlark.Builtin,
+	args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
+	var name *starlark.List
+
+	if err := starlark.UnpackArgs(ruleRPackage,
+		args, kwargs, "name", &name); err != nil {
+		return nil, err
+	}
+
+	nameList := []string{}
+	if name != nil {
+		for i := 0; i < name.Len(); i++ {
+			nameList = append(nameList, name.Index(i).(starlark.String).GoString())
+		}
+	}
+
+	logger.Debugf("rule `%s` is invoked, name=%v", ruleRPackage, nameList)
+	ir.RPackage(nameList)
+
+	return starlark.None, nil
+}
+
 func ruleFuncSystemPackage(thread *starlark.Thread, _ *starlark.Builtin,
 	args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
 	var name *starlark.List

pkg/lang/ir/compile.go

@@ -34,6 +34,7 @@ func NewGraph() *Graph {
 		CUDNN:    nil,
 
 		PyPIPackages:   []string{},
+		RPackages:      []string{},
 		SystemPackages: []string{},
 		Exec:           []string{},
 		Shell:          shellBASH,
@@ -86,6 +87,11 @@ func (g Graph) Labels() (map[string]string, error) {
 		return nil, err
 	}
 	labels[types.ImageLabelPyPI] = string(str)
+	str, err = json.Marshal(g.RPackages)
+	if err != nil {
+		return nil, err
+	}
+	labels[types.ImageLabelR] = string(str)
 	if g.GPUEnabled() {
 		labels[types.ImageLabelGPU] = "true"
 		labels[types.ImageLabelCUDA] = *g.CUDA
@@ -102,41 +108,49 @@ func (g Graph) Compile() (llb.State, error) {
 	// TODO(gaocegege): Support more OS and langs.
 	base := g.compileBase()
 	aptStage := g.compileUbuntuAPT(base)
-	pypiMirrorStage := g.compilePyPIIndex(aptStage)
-
-	g.compileJupyter()
-	builtinSystemStage := pypiMirrorStage
-	sshStage, err := g.copySSHKey(builtinSystemStage)
-	if err != nil {
-		return llb.State{}, errors.Wrap(err, "failed to copy ssh keys")
-	}
-	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"))
-	diffSSHStage := llb.Diff(builtinSystemStage, sshStage, llb.WithCustomName("install ssh keys"))
-	pypiStage := llb.Diff(builtinSystemStage, g.compilePyPIPackages(builtinSystemStage), llb.WithCustomName("install PyPI packages"))
-	systemStage := llb.Diff(builtinSystemStage, g.compileSystemPackages(builtinSystemStage), llb.WithCustomName("install system packages"))
-
-	if err != nil {
-		return llb.State{}, errors.Wrap(err, "failed to copy SSH key")
-	}
-
-	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 {
+	if g.Language == "r" {
+		// TODO(terrytangyuan): Support RStudio local server
+		rPackageInstallStage := llb.Diff(aptStage, g.installRPackages(aptStage), llb.WithCustomName("install R packages"))
 		merged = llb.Merge([]llb.State{
-			builtinSystemStage, systemStage, diffSSHStage, pypiStage, *vscodeStage, diffShellStage,
+			aptStage, rPackageInstallStage,
 		}, llb.WithCustomName("merging all components into one"))
 	} else {
-		merged = llb.Merge([]llb.State{
-			builtinSystemStage, systemStage, diffSSHStage, pypiStage, diffShellStage,
-		}, llb.WithCustomName("merging all components into one"))
+		pypiMirrorStage := g.compilePyPIIndex(aptStage)
+
+		g.compileJupyter()
+		builtinSystemStage := pypiMirrorStage
+		sshStage, err := g.copySSHKey(builtinSystemStage)
+		if err != nil {
+			return llb.State{}, errors.Wrap(err, "failed to copy ssh keys")
+		}
+		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"))
+		diffSSHStage := llb.Diff(builtinSystemStage, sshStage, llb.WithCustomName("install ssh keys"))
+		pypiStage := llb.Diff(builtinSystemStage, g.compilePyPIPackages(builtinSystemStage), llb.WithCustomName("install PyPI packages"))
+		systemStage := llb.Diff(builtinSystemStage, g.compileSystemPackages(builtinSystemStage), llb.WithCustomName("install system packages"))
+
+		if err != nil {
+			return llb.State{}, errors.Wrap(err, "failed to copy SSH key")
+		}
+
+		vscodeStage, err := g.compileVSCode()
+		if err != nil {
+			return llb.State{}, errors.Wrap(err, "failed to get vscode plugins")
+		}
+
+		if vscodeStage != nil {
+			merged = llb.Merge([]llb.State{
+				builtinSystemStage, systemStage, diffSSHStage, pypiStage, *vscodeStage, diffShellStage,
+			}, llb.WithCustomName("merging all components into one"))
+		} else {
+			merged = llb.Merge([]llb.State{
+				builtinSystemStage, systemStage, diffSSHStage, pypiStage, diffShellStage,
+			}, llb.WithCustomName("merging all components into one"))
+		}
 	}
 
 	// TODO(gaocegege): Support order-based exec.

pkg/lang/ir/interface.go

@@ -29,6 +29,10 @@ func PyPIPackage(deps []string) {
 	DefaultGraph.PyPIPackages = append(DefaultGraph.PyPIPackages, deps...)
 }
 
+func RPackage(deps []string) {
+	DefaultGraph.RPackages = append(DefaultGraph.RPackages, deps...)
+}
+
 func SystemPackage(deps []string) {
 	DefaultGraph.SystemPackages = append(DefaultGraph.SystemPackages, deps...)
 }

pkg/lang/ir/r.go

@@ -0,0 +1,42 @@
+// 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"
+	"strings"
+
+	"github.com/moby/buildkit/client/llb"
+)
+
+func (g Graph) installRPackages(root llb.State) llb.State {
+	// TODO(terrytangyuan): Support different CRAN mirrors
+	var sb strings.Builder
+	sb.WriteString(`R -e 'install.packages(c(`)
+	for i, pkg := range g.RPackages {
+		sb.WriteString(fmt.Sprintf(`"%s"`, pkg))
+		if i != len(g.RPackages)-1 {
+			sb.WriteString(", ")
+		}
+	}
+	sb.WriteString(`))'`)
+
+	// TODO(terrytangyuan): Support cache.
+	cmd := sb.String()
+	root = llb.User("envd")(root)
+	run := root.
+		Run(llb.Shlex(cmd), llb.WithCustomNamef("R package install"))
+	return run.Root()
+}

pkg/lang/ir/system.go

@@ -84,15 +84,23 @@ func (g Graph) compileSystemPackages(root llb.State) llb.State {
 
 func (g *Graph) compileBase() llb.State {
 	var base llb.State
+	var groupID string
 	if g.CUDA == nil && g.CUDNN == nil {
-		base = llb.Image("docker.io/gaocegege/python:3.8-ubuntu20.04")
+		if g.Language == "r" {
+			base = llb.Image("docker.io/r-base:4.2.0")
+			// r-base image already has GID 1000.
+			groupID = "1001"
+		} else {
+			base = llb.Image("docker.io/gaocegege/python:3.8-ubuntu20.04")
+			groupID = "1000"
+		}
 	} else {
 		base = g.compileCUDAPackages()
 	}
 	// TODO(gaocegege): Refactor user to a seperate stage.
 	res := base.
-		Run(llb.Shlex("groupadd -g 1000 envd"), llb.WithCustomName("create user group envd")).
-		Run(llb.Shlex("useradd -p \"\" -u 1000 -g envd -s /bin/sh -m envd"), llb.WithCustomName("create user envd")).
+		Run(llb.Shlex(fmt.Sprintf("groupadd -g %s envd", groupID)), llb.WithCustomName("create user group envd")).
+		Run(llb.Shlex(fmt.Sprintf("useradd -p \"\" -u %s -g envd -s /bin/sh -m envd", groupID)), llb.WithCustomName("create user envd")).
 		Run(llb.Shlex("adduser envd sudo"), llb.WithCustomName("add user envd to sudoers")).Run(llb.Shlex("chown -R envd:envd /usr/local/lib"))
 	return llb.User("envd")(res.Root())
 }

pkg/lang/ir/types.go

@@ -36,6 +36,7 @@ type Graph struct {
 
 	BuiltinSystemPackages []string
 	PyPIPackages          []string
+	RPackages             []string
 	SystemPackages        []string
 	VSCodePlugins         []vscode.Plugin
 

pkg/types/envd.go

@@ -55,6 +55,7 @@ const (
 	ImageLabelGPU     = "ai.tensorchord.envd.gpu"
 	ImageLabelAPT     = "ai.tensorchord.envd.apt.packages"
 	ImageLabelPyPI    = "ai.tensorchord.envd.pypi.packages"
+	ImageLabelR       = "ai.tensorchord.envd.r.packages"
 	ImageLabelCUDA    = "ai.tensorchord.envd.gpu.cuda"
 	ImageLabelCUDNN   = "ai.tensorchord.envd.gpu.cudnn"
 	ImageLabelContext = "ai.tensorchord.envd.build.context"