Skip to content

Commit

Permalink
feat: Initial support for R language (#257)
Browse files Browse the repository at this point in the history
* feat: Initial support for R language

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

* fix groupID

Signed-off-by: terrytangyuan <terrytangyuan@gmail.com>
  • Loading branch information
terrytangyuan committed Jun 10, 2022
1 parent ba64556 commit 66d83d6
Show file tree
Hide file tree
Showing 9 changed files with 135 additions and 34 deletions.
6 changes: 6 additions & 0 deletions examples/r-basic/build.envd
Original file line number Diff line number Diff line change
@@ -0,0 +1,6 @@
def build():
base(os="ubuntu20.04", language="r")
install.r_packages([
"remotes",
"rlang",
])
1 change: 1 addition & 0 deletions pkg/lang/frontend/starlark/install/const.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
)
24 changes: 24 additions & 0 deletions pkg/lang/frontend/starlark/install/install.go
Original file line number Diff line number Diff line change
Expand Up @@ -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),
Expand Down Expand Up @@ -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
Expand Down
76 changes: 45 additions & 31 deletions pkg/lang/ir/compile.go
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,7 @@ func NewGraph() *Graph {
CUDNN: nil,

PyPIPackages: []string{},
RPackages: []string{},
SystemPackages: []string{},
Exec: []string{},
Shell: shellBASH,
Expand Down Expand Up @@ -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
Expand All @@ -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.
Expand Down
4 changes: 4 additions & 0 deletions pkg/lang/ir/interface.go
Original file line number Diff line number Diff line change
Expand Up @@ -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...)
}
Expand Down
42 changes: 42 additions & 0 deletions pkg/lang/ir/r.go
Original file line number Diff line number Diff line change
@@ -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()
}
14 changes: 11 additions & 3 deletions pkg/lang/ir/system.go
Original file line number Diff line number Diff line change
Expand Up @@ -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())
}
Expand Down
1 change: 1 addition & 0 deletions pkg/lang/ir/types.go
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,7 @@ type Graph struct {

BuiltinSystemPackages []string
PyPIPackages []string
RPackages []string
SystemPackages []string
VSCodePlugins []vscode.Plugin

Expand Down
1 change: 1 addition & 0 deletions pkg/types/envd.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down

0 comments on commit 66d83d6

Please sign in to comment.