From 66d83d699651cf5b20b4acafe0c2369f1945c561 Mon Sep 17 00:00:00 2001 From: Yuan Tang Date: Thu, 9 Jun 2022 20:52:48 -0400 Subject: [PATCH] feat: Initial support for R language (#257) * feat: Initial support for R language Signed-off-by: terrytangyuan * fix groupID Signed-off-by: terrytangyuan --- examples/r-basic/build.envd | 6 ++ pkg/lang/frontend/starlark/install/const.go | 1 + pkg/lang/frontend/starlark/install/install.go | 24 ++++++ pkg/lang/ir/compile.go | 76 +++++++++++-------- pkg/lang/ir/interface.go | 4 + pkg/lang/ir/r.go | 42 ++++++++++ pkg/lang/ir/system.go | 14 +++- pkg/lang/ir/types.go | 1 + pkg/types/envd.go | 1 + 9 files changed, 135 insertions(+), 34 deletions(-) create mode 100644 examples/r-basic/build.envd create mode 100644 pkg/lang/ir/r.go diff --git a/examples/r-basic/build.envd b/examples/r-basic/build.envd new file mode 100644 index 000000000..ef835b1ec --- /dev/null +++ b/examples/r-basic/build.envd @@ -0,0 +1,6 @@ +def build(): + base(os="ubuntu20.04", language="r") + install.r_packages([ + "remotes", + "rlang", + ]) diff --git a/pkg/lang/frontend/starlark/install/const.go b/pkg/lang/frontend/starlark/install/const.go index 52e9b8364..9ed3d9799 100644 --- a/pkg/lang/frontend/starlark/install/const.go +++ b/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" ) diff --git a/pkg/lang/frontend/starlark/install/install.go b/pkg/lang/frontend/starlark/install/install.go index b6bd06486..4dc0d3d11 100644 --- a/pkg/lang/frontend/starlark/install/install.go +++ b/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 diff --git a/pkg/lang/ir/compile.go b/pkg/lang/ir/compile.go index cd79ecc4a..69c359612 100644 --- a/pkg/lang/ir/compile.go +++ b/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. diff --git a/pkg/lang/ir/interface.go b/pkg/lang/ir/interface.go index 5f8cd8ce6..15dac716e 100644 --- a/pkg/lang/ir/interface.go +++ b/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...) } diff --git a/pkg/lang/ir/r.go b/pkg/lang/ir/r.go new file mode 100644 index 000000000..721ef64e3 --- /dev/null +++ b/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() +} diff --git a/pkg/lang/ir/system.go b/pkg/lang/ir/system.go index 9665e8662..1069e6796 100644 --- a/pkg/lang/ir/system.go +++ b/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()) } diff --git a/pkg/lang/ir/types.go b/pkg/lang/ir/types.go index 9b1ffac4d..f21e733f9 100644 --- a/pkg/lang/ir/types.go +++ b/pkg/lang/ir/types.go @@ -36,6 +36,7 @@ type Graph struct { BuiltinSystemPackages []string PyPIPackages []string + RPackages []string SystemPackages []string VSCodePlugins []vscode.Plugin diff --git a/pkg/types/envd.go b/pkg/types/envd.go index f9b100a71..b96207e5d 100644 --- a/pkg/types/envd.go +++ b/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"