From 8be30356b22e8fa222bbde32fa808d1b790f5924 Mon Sep 17 00:00:00 2001 From: Ce Gao Date: Tue, 14 Jun 2022 19:59:30 +0800 Subject: [PATCH 1/3] WIP Signed-off-by: Ce Gao --- examples/conda/build.envd | 26 ++++++++ pkg/app/destroy.go | 2 +- pkg/editor/jupyter/util.go | 2 +- pkg/lang/frontend/starlark/install/install.go | 15 +++-- .../frontend/starlark/universe/universe.go | 5 +- pkg/lang/ir/compile.go | 23 ++++--- pkg/lang/ir/conda.go | 53 ++++++++++++--- pkg/lang/ir/consts.go | 2 +- pkg/lang/ir/interface.go | 31 +++++++-- pkg/lang/ir/python.go | 4 +- pkg/lang/ir/system.go | 11 +++- pkg/lang/ir/types.go | 25 ++++++-- pkg/lang/ir/util.go | 46 +++++++++++++ pkg/lang/ir/util_test.go | 64 +++++++++++++++++++ 14 files changed, 268 insertions(+), 41 deletions(-) create mode 100644 examples/conda/build.envd create mode 100644 pkg/lang/ir/util.go create mode 100644 pkg/lang/ir/util_test.go diff --git a/examples/conda/build.envd b/examples/conda/build.envd new file mode 100644 index 000000000..fcd1e52ae --- /dev/null +++ b/examples/conda/build.envd @@ -0,0 +1,26 @@ +def build(): + config.conda_channel(channel=""" +channels: + - defaults +show_channel_urls: true +default_channels: + - https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main + - https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/r + - https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/msys2 + - https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/free +custom_channels: + conda-forge: https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud + msys2: https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud + bioconda: https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud + menpo: https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud + pytorch: https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud + pytorch-lts: https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud + simpleitk: https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud +""") + # install.conda_packages(name = [ + # "pytorch=0.4.1", "cuda90" + # ], channel= ["pytorch"]) + base(os="ubuntu20.04", language="python3.5") + install.python_packages(name = [ + "flask" + ]) diff --git a/pkg/app/destroy.go b/pkg/app/destroy.go index 07251838d..6ac151664 100644 --- a/pkg/app/destroy.go +++ b/pkg/app/destroy.go @@ -71,7 +71,7 @@ func destroy(clicontext *cli.Context) error { } if ctrName, err := dockerClient.Destroy(clicontext.Context, ctrName); err != nil { return errors.Wrapf(err, "failed to destroy the environment: %s", ctrName) - } else if name != "" { + } else if ctrName != "" { logrus.Infof("%s is destroyed", ctrName) } diff --git a/pkg/editor/jupyter/util.go b/pkg/editor/jupyter/util.go index 2cc2fd380..3aea62e52 100644 --- a/pkg/editor/jupyter/util.go +++ b/pkg/editor/jupyter/util.go @@ -27,7 +27,7 @@ func GenerateCommand(g ir.Graph, notebookDir string) []string { var cmd []string // Use python in conda env. - if len(g.CondaPackages) != 0 { + if g.CondaEnabled() { cmd = append(cmd, "/opt/conda/bin/python3") } else { cmd = append(cmd, "python3") diff --git a/pkg/lang/frontend/starlark/install/install.go b/pkg/lang/frontend/starlark/install/install.go index 50dc9ae0b..f859aff1e 100644 --- a/pkg/lang/frontend/starlark/install/install.go +++ b/pkg/lang/frontend/starlark/install/install.go @@ -158,10 +158,10 @@ func ruleFuncVSCode(thread *starlark.Thread, _ *starlark.Builtin, func ruleFuncConda(thread *starlark.Thread, _ *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { - var name *starlark.List + var name, channel *starlark.List if err := starlark.UnpackArgs(ruleConda, - args, kwargs, "name", &name); err != nil { + args, kwargs, "name", &name, "channel?", &channel); err != nil { return nil, err } @@ -172,8 +172,15 @@ func ruleFuncConda(thread *starlark.Thread, _ *starlark.Builtin, } } - logger.Debugf("rule `%s` is invoked, name=%v", ruleConda, nameList) - ir.CondaPackage(nameList) + channelList := []string{} + if channel != nil { + for i := 0; i < channel.Len(); i++ { + channelList = append(channelList, channel.Index(i).(starlark.String).GoString()) + } + } + + logger.Debugf("rule `%s` is invoked, name=%v, channel=%v", ruleConda, nameList, channelList) + ir.CondaPackage(nameList, channelList) return starlark.None, nil } diff --git a/pkg/lang/frontend/starlark/universe/universe.go b/pkg/lang/frontend/starlark/universe/universe.go index 66ece88e7..8b4ef8f37 100644 --- a/pkg/lang/frontend/starlark/universe/universe.go +++ b/pkg/lang/frontend/starlark/universe/universe.go @@ -52,7 +52,10 @@ func ruleFuncBase(thread *starlark.Thread, _ *starlark.Builtin, logger.Debugf("rule `%s` is invoked, os=%s, language=%s", ruleBase, osStr, langStr) - ir.Base(osStr, langStr) + err := ir.Base(osStr, langStr) + if err != nil { + return starlark.None, err + } return starlark.None, nil } diff --git a/pkg/lang/ir/compile.go b/pkg/lang/ir/compile.go index 7de5bf2fc..6bccb090d 100644 --- a/pkg/lang/ir/compile.go +++ b/pkg/lang/ir/compile.go @@ -28,10 +28,12 @@ import ( func NewGraph() *Graph { return &Graph{ - OS: osDefault, - Language: languageDefault, - CUDA: nil, - CUDNN: nil, + OS: osDefault, + Language: Language{ + Name: languageDefault, + }, + CUDA: nil, + CUDNN: nil, PyPIPackages: []string{}, RPackages: []string{}, @@ -109,7 +111,7 @@ func (g Graph) Compile() (llb.State, error) { base := g.compileBase() aptStage := g.compileUbuntuAPT(base) var merged llb.State - if g.Language == "r" { + 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{ @@ -133,12 +135,15 @@ func (g Graph) Compile() (llb.State, error) { if err != nil { return llb.State{}, errors.Wrap(err, "failed to compile shell") } + + condaEnvStage := g.setCondaENV(shellStage) + condaStage := llb.Diff(builtinSystemStage, - g.compileCondaPackages(shellStage), - llb.WithCustomName("install PyPI packages")) + g.compileCondaPackages(condaEnvStage), + llb.WithCustomName("install conda packages")) - pypiStage := llb.Diff(builtinSystemStage, - g.compilePyPIPackages(builtinSystemStage), + pypiStage := llb.Diff(condaEnvStage, + g.compilePyPIPackages(condaEnvStage), llb.WithCustomName("install PyPI packages")) systemStage := llb.Diff(builtinSystemStage, g.compileSystemPackages(builtinSystemStage), llb.WithCustomName("install system packages")) diff --git a/pkg/lang/ir/conda.go b/pkg/lang/ir/conda.go index d0151f07f..235c29274 100644 --- a/pkg/lang/ir/conda.go +++ b/pkg/lang/ir/conda.go @@ -26,8 +26,12 @@ const ( condarc = "/home/envd/.condarc" ) +func (g Graph) CondaEnabled() bool { + return g.CondaConfig != nil +} + func (g Graph) compileCondaChannel(root llb.State) llb.State { - if g.CondaChannel != nil { + if g.CondaConfig != nil && g.CondaConfig.CondaChannel != nil { logrus.WithField("conda-channel", *g.CondaChannel).Debug("using custom connda channel") stage := root. File(llb.Mkfile(condarc, @@ -38,7 +42,7 @@ func (g Graph) compileCondaChannel(root llb.State) llb.State { } func (g Graph) compileCondaPackages(root llb.State) llb.State { - if len(g.CondaPackages) == 0 { + if !g.CondaEnabled() || len(g.CondaConfig.CondaPackages) == 0 { return root } @@ -46,8 +50,17 @@ func (g Graph) compileCondaPackages(root llb.State) llb.State { // Compose the package install command. var sb strings.Builder - sb.WriteString("/opt/conda/bin/conda install") - for _, pkg := range g.CondaPackages { + if len(g.CondaConfig.AdditionalChannels) == 0 { + sb.WriteString("/opt/conda/bin/conda install -n envd") + + } else { + sb.WriteString("/opt/conda/bin/conda install -n envd -c") + for _, channel := range g.CondaConfig.AdditionalChannels { + sb.WriteString(fmt.Sprintf(" %s", channel)) + } + } + + for _, pkg := range g.CondaConfig.CondaPackages { sb.WriteString(fmt.Sprintf(" %s", pkg)) } @@ -62,16 +75,38 @@ func (g Graph) compileCondaPackages(root llb.State) llb.State { strings.Join(g.CondaPackages, " "))) run.AddMount(cacheDir, cache, llb.AsPersistentCacheDir(g.CacheID(cacheDir), llb.CacheMountShared), llb.SourcePath("/cache")) - return g.setCondaENV(run.Root()) + return run.Root() } func (g Graph) setCondaENV(root llb.State) llb.State { + if !g.CondaEnabled() { + return root + } + root = llb.User("envd")(root) // Always init bash since we will use it to create jupyter notebook service. - run := root.Run(llb.Shlex("/opt/conda/bin/conda init bash"), llb.WithCustomName("[internal] initialize conda bash environment")) - if g.Shell != shellBASH { - run = run.Run(llb.Shlex(fmt.Sprintf("/opt/conda/bin/conda init %s", g.Shell)), - llb.WithCustomNamef("[internal] initialize conda %s environment", g.Shell)) + run := root.Run(llb.Shlex("bash -c \"/opt/conda/bin/conda init bash\""), llb.WithCustomName("[internal] initialize conda bash environment")) + + pythonVersion := "3.9" + if g.Language.Version != nil { + pythonVersion = *g.Language.Version + } + cmd := fmt.Sprintf( + "bash -c \"/opt/conda/bin/conda create -n envd python=%s\"", pythonVersion) + // Create a conda environment. + run = run.Run(llb.Shlex(cmd), llb.WithCustomName("[internal] create conda environment")) + + switch g.Shell { + case shellBASH: + run = run.Run( + llb.Shlex(`echo "source /opt/conda/bin/activate envd" >> /home/envd/.bashrc`), + llb.WithCustomName("[internal] add conda environment to bashrc")) + case shellZSH: + run = run.Run( + llb.Shlex(fmt.Sprintf("bash -c \"/opt/conda/bin/conda init %s\"", g.Shell)), + llb.WithCustomNamef("[internal] initialize conda %s environment", g.Shell)).Run( + llb.Shlex(`echo "source /opt/conda/bin/activate envd" >> /home/envd/.zshrc`), + llb.WithCustomName("[internal] add conda environment to zshrc")) } return run.Root() } diff --git a/pkg/lang/ir/consts.go b/pkg/lang/ir/consts.go index c6c2db93e..48aaae6d6 100644 --- a/pkg/lang/ir/consts.go +++ b/pkg/lang/ir/consts.go @@ -16,7 +16,7 @@ package ir const ( osDefault = "ubuntu20.04" - languageDefault = "python3.8" + languageDefault = "python" pypiIndexModeAuto = "auto" aptSourceFilePath = "/etc/apt/sources.list" diff --git a/pkg/lang/ir/interface.go b/pkg/lang/ir/interface.go index dff78ba92..f4a3ce88d 100644 --- a/pkg/lang/ir/interface.go +++ b/pkg/lang/ir/interface.go @@ -20,9 +20,17 @@ import ( "github.com/tensorchord/envd/pkg/editor/vscode" ) -func Base(os, language string) { - DefaultGraph.Language = language +func Base(os, language string) error { + l, version, err := parseLanguage(language) + if err != nil { + return err + } + DefaultGraph.Language = Language{ + Name: l, + Version: version, + } DefaultGraph.OS = os + return nil } func PyPIPackage(deps []string) { @@ -116,10 +124,23 @@ func CondaChannel(channel string) error { return errors.New("channel is required") } - DefaultGraph.CondaChannel = &channel + if !DefaultGraph.CondaEnabled() { + DefaultGraph.CondaConfig = &CondaConfig{} + } + + DefaultGraph.CondaConfig.CondaChannel = &channel return nil } -func CondaPackage(deps []string) { - DefaultGraph.CondaPackages = append(DefaultGraph.CondaPackages, deps...) +func CondaPackage(deps []string, channel []string) { + if !DefaultGraph.CondaEnabled() { + DefaultGraph.CondaConfig = &CondaConfig{} + } + DefaultGraph.CondaConfig.CondaPackages = append( + DefaultGraph.CondaConfig.CondaPackages, deps...) + + if len(channel) != 0 { + DefaultGraph.CondaConfig.AdditionalChannels = append( + DefaultGraph.CondaConfig.AdditionalChannels, channel...) + } } diff --git a/pkg/lang/ir/python.go b/pkg/lang/ir/python.go index ff3076064..9c042167c 100644 --- a/pkg/lang/ir/python.go +++ b/pkg/lang/ir/python.go @@ -32,8 +32,8 @@ func (g Graph) compilePyPIPackages(root llb.State) llb.State { // Compose the package install command. var sb strings.Builder - if len(g.CondaPackages) != 0 { - sb.WriteString("/opt/conda/bin/pip install --no-warn-script-location") + if g.CondaEnabled() { + sb.WriteString("/opt/conda/bin/conda run -n envd pip install") } else { sb.WriteString("pip install --no-warn-script-location") } diff --git a/pkg/lang/ir/system.go b/pkg/lang/ir/system.go index 39c08ed93..c7b001009 100644 --- a/pkg/lang/ir/system.go +++ b/pkg/lang/ir/system.go @@ -86,10 +86,19 @@ func (g Graph) compileSystemPackages(root llb.State) llb.State { } func (g *Graph) compileBase() llb.State { + logger := logrus.WithFields(logrus.Fields{ + "os": g.OS, + "language": g.Language.Name, + }) + if g.Language.Version != nil { + logger = logger.WithField("version", *g.Language.Version) + } + logger.Debug("compile base image") + var base llb.State var groupID string = "1000" if g.CUDA == nil && g.CUDNN == nil { - if g.Language == "r" { + if g.Language.Name == "r" { base = llb.Image("docker.io/r-base:4.2.0") // r-base image already has GID 1000. groupID = "1001" diff --git a/pkg/lang/ir/types.go b/pkg/lang/ir/types.go index 3579f4e37..74c3a6cba 100644 --- a/pkg/lang/ir/types.go +++ b/pkg/lang/ir/types.go @@ -22,22 +22,21 @@ import ( // A Graph contains the state, // such as its call stack and thread-local storage. type Graph struct { - OS string - Language string - Shell string - CUDA *string - CUDNN *string + OS string + Language + + Shell string + CUDA *string + CUDNN *string UbuntuAPTSource *string PyPIIndexURL *string PyPIExtraIndexURL *string - CondaChannel *string PublicKeyPath string PyPIPackages []string RPackages []string - CondaPackages []string SystemPackages []string VSCodePlugins []vscode.Plugin @@ -45,11 +44,23 @@ type Graph struct { Exec []string *JupyterConfig *GitConfig + *CondaConfig Writer compileui.Writer CachePrefix string } +type Language struct { + Name string + Version *string +} + +type CondaConfig struct { + CondaPackages []string + AdditionalChannels []string + CondaChannel *string +} + type GitConfig struct { Name string Email string diff --git a/pkg/lang/ir/util.go b/pkg/lang/ir/util.go new file mode 100644 index 000000000..fe2b1a811 --- /dev/null +++ b/pkg/lang/ir/util.go @@ -0,0 +1,46 @@ +// 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 ( + "errors" + "fmt" + "regexp" +) + +func parseLanguage(l string) (string, *string, error) { + var language, version string + if l == "" { + return "", nil, errors.New("language is required") + } + + re := regexp.MustCompile(`\d[\d,]*[\.]?[\d{2}]*[\.]?[\d{2}]*`) + if !re.MatchString(l) { + language = l + } else { + loc := re.FindStringIndex(l) + language = l[:loc[0]] + version = l[loc[0]:] + } + + switch language { + case "python": + return "python", &version, nil + case "r": + return "r", &version, nil + default: + return "", nil, fmt.Errorf("language %s is not supported", language) + } +} diff --git a/pkg/lang/ir/util_test.go b/pkg/lang/ir/util_test.go new file mode 100644 index 000000000..a87153a81 --- /dev/null +++ b/pkg/lang/ir/util_test.go @@ -0,0 +1,64 @@ +package ir + +import "testing" + +func TestParseLanguage(t *testing.T) { + tcs := []struct { + l string + ExpectedLanguage string + ExpectedVersion string + ExpectedError bool + }{ + { + l: "python", + ExpectedLanguage: "python", + ExpectedVersion: "", + ExpectedError: false, + }, + { + l: "python3.7", + ExpectedLanguage: "python", + ExpectedVersion: "3.7", + ExpectedError: false, + }, + { + l: "python3.7.1", + ExpectedLanguage: "python", + ExpectedVersion: "3.7.1", + ExpectedError: false, + }, + { + l: "python-3.7.1", + ExpectedError: true, + }, + { + l: "r", + ExpectedLanguage: "r", + ExpectedVersion: "", + ExpectedError: false, + }, + } + + for _, tc := range tcs { + language, version, err := parseLanguage(tc.l) + if err != nil { + if !tc.ExpectedError { + t.Errorf("parseLanguage(%s) returned error: %v", tc.l, err) + } + } else { + if language != tc.ExpectedLanguage { + t.Errorf("parseLanguage(%s) returned language %s, expected %s", tc.l, language, tc.ExpectedLanguage) + } + if version == nil { + if tc.ExpectedVersion != "" { + t.Errorf("parseLanguage(%s) returned version nil, expected %s", tc.l, tc.ExpectedVersion) + } + } else { + if *version != tc.ExpectedVersion { + t.Errorf("parseLanguage(%s) returned version %s, expected %s", tc.l, *version, tc.ExpectedVersion) + } + } + } + + } +} From 476e2af760a829637b81e50c9df248ae4d004b07 Mon Sep 17 00:00:00 2001 From: Ce Gao Date: Tue, 14 Jun 2022 20:32:30 +0800 Subject: [PATCH 2/3] fix: Use conda envd env Signed-off-by: Ce Gao --- pkg/lang/ir/conda.go | 5 ++--- pkg/lang/ir/util_test.go | 14 ++++++++++++++ 2 files changed, 16 insertions(+), 3 deletions(-) diff --git a/pkg/lang/ir/conda.go b/pkg/lang/ir/conda.go index 235c29274..b610f8ff1 100644 --- a/pkg/lang/ir/conda.go +++ b/pkg/lang/ir/conda.go @@ -99,13 +99,12 @@ func (g Graph) setCondaENV(root llb.State) llb.State { switch g.Shell { case shellBASH: run = run.Run( - llb.Shlex(`echo "source /opt/conda/bin/activate envd" >> /home/envd/.bashrc`), - llb.WithCustomName("[internal] add conda environment to bashrc")) + llb.Shlex(`bash -c 'echo "source /opt/conda/bin/activate envd" >> /home/envd/.bashrc'`)) case shellZSH: run = run.Run( llb.Shlex(fmt.Sprintf("bash -c \"/opt/conda/bin/conda init %s\"", g.Shell)), llb.WithCustomNamef("[internal] initialize conda %s environment", g.Shell)).Run( - llb.Shlex(`echo "source /opt/conda/bin/activate envd" >> /home/envd/.zshrc`), + llb.Shlex(`bash -c 'echo "source /opt/conda/bin/activate envd" >> /home/envd/.zshrc'`), llb.WithCustomName("[internal] add conda environment to zshrc")) } return run.Root() diff --git a/pkg/lang/ir/util_test.go b/pkg/lang/ir/util_test.go index a87153a81..765fa368b 100644 --- a/pkg/lang/ir/util_test.go +++ b/pkg/lang/ir/util_test.go @@ -1,3 +1,17 @@ +// 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 "testing" From 2309aec8e09441f485afe814a17eb7d7d06ded14 Mon Sep 17 00:00:00 2001 From: Ce Gao Date: Tue, 14 Jun 2022 20:49:12 +0800 Subject: [PATCH 3/3] fix: Add name Signed-off-by: Ce Gao --- examples/conda/build.envd | 7 ++++--- pkg/lang/ir/conda.go | 3 ++- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/examples/conda/build.envd b/examples/conda/build.envd index fcd1e52ae..267894d77 100644 --- a/examples/conda/build.envd +++ b/examples/conda/build.envd @@ -17,10 +17,11 @@ custom_channels: pytorch-lts: https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud simpleitk: https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud """) - # install.conda_packages(name = [ - # "pytorch=0.4.1", "cuda90" - # ], channel= ["pytorch"]) + install.conda_packages(name = [ + "pytorch=0.4.1", "cuda90" + ], channel= ["pytorch"]) base(os="ubuntu20.04", language="python3.5") install.python_packages(name = [ "flask" ]) + install.cuda(version="11.6", cudnn="8") diff --git a/pkg/lang/ir/conda.go b/pkg/lang/ir/conda.go index b610f8ff1..08f618bd7 100644 --- a/pkg/lang/ir/conda.go +++ b/pkg/lang/ir/conda.go @@ -99,7 +99,8 @@ func (g Graph) setCondaENV(root llb.State) llb.State { switch g.Shell { case shellBASH: run = run.Run( - llb.Shlex(`bash -c 'echo "source /opt/conda/bin/activate envd" >> /home/envd/.bashrc'`)) + llb.Shlex(`bash -c 'echo "source /opt/conda/bin/activate envd" >> /home/envd/.bashrc'`), + llb.WithCustomName("[internal] add conda environment to bashrc")) case shellZSH: run = run.Run( llb.Shlex(fmt.Sprintf("bash -c \"/opt/conda/bin/conda init %s\"", g.Shell)),