From ee8eed47d20d3a0eb485ff727524f2543c530c4d Mon Sep 17 00:00:00 2001 From: Ce Gao Date: Wed, 7 Dec 2022 12:24:27 +0800 Subject: [PATCH] fix Signed-off-by: Ce Gao --- examples/python-basic/build.envd | 20 +- pkg/app/create.go | 2 +- pkg/app/run.go | 2 +- pkg/app/up.go | 2 +- pkg/builder/builder.go | 8 +- pkg/builder/builder_test.go | 4 +- pkg/builder/util.go | 4 +- pkg/lang/frontend/starlark/interpreter.go | 6 + .../frontend/starlark/v0/builtin/builtin.go | 20 + .../frontend/starlark/v0/config/config.go | 228 ++++++++++++ pkg/lang/frontend/starlark/v0/config/const.go | 28 ++ pkg/lang/frontend/starlark/v0/data/const.go | 21 ++ pkg/lang/frontend/starlark/v0/data/rule.go | 54 +++ pkg/lang/frontend/starlark/v0/data/util.go | 53 +++ .../frontend/starlark/v0/install/const.go | 25 ++ .../frontend/starlark/v0/install/install.go | 204 +++++++++++ pkg/lang/frontend/starlark/v0/interpreter.go | 209 +++++++++++ .../frontend/starlark/v0/interpreter_test.go | 29 ++ pkg/lang/frontend/starlark/v0/io/const.go | 20 + pkg/lang/frontend/starlark/v0/io/io.go | 70 ++++ pkg/lang/frontend/starlark/v0/mock/mock.go | 64 ++++ .../frontend/starlark/v0/runtime/const.go | 24 ++ .../frontend/starlark/v0/runtime/runtime.go | 238 ++++++++++++ .../starlark/v0/starlark_suite_test.go | 27 ++ .../frontend/starlark/v0/testdata/test.envd | 8 + .../frontend/starlark/v0/universe/const.go | 25 ++ .../frontend/starlark/v0/universe/universe.go | 141 ++++++++ .../frontend/starlark/v1/config/config.go | 9 +- .../frontend/starlark/v1/install/const.go | 17 +- .../frontend/starlark/v1/install/install.go | 72 +++- pkg/lang/frontend/starlark/v1/interpreter.go | 8 +- .../frontend/starlark/v1/universe/universe.go | 11 +- pkg/lang/ir/graph.go | 1 + pkg/lang/ir/v0/cache.go | 32 ++ pkg/lang/ir/v0/checker.go | 66 ++++ pkg/lang/ir/v0/compile.go | 341 ++++++++++++++++++ pkg/lang/ir/v0/conda.go | 158 ++++++++ pkg/lang/ir/v0/consts.go | 42 +++ pkg/lang/ir/{v1 => v0}/custom.go | 0 pkg/lang/ir/v0/editor.go | 120 ++++++ pkg/lang/ir/v0/editor_test.go | 97 +++++ pkg/lang/ir/v0/fs.go | 25 ++ pkg/lang/ir/v0/git.go | 45 +++ pkg/lang/ir/{v1 => v0}/install-conda.sh | 0 pkg/lang/ir/{v1 => v0}/install-mamba.sh | 0 pkg/lang/ir/v0/interface.go | 298 +++++++++++++++ pkg/lang/ir/v0/julia.go | 101 ++++++ pkg/lang/ir/v0/python.go | 226 ++++++++++++ pkg/lang/ir/v0/r.go | 96 +++++ pkg/lang/ir/v0/shell.go | 119 ++++++ pkg/lang/ir/v0/supervisor.go | 117 ++++++ pkg/lang/ir/v0/system.go | 289 +++++++++++++++ pkg/lang/ir/v0/types.go | 70 ++++ pkg/lang/ir/v0/user.go | 73 ++++ pkg/lang/ir/v0/util.go | 118 ++++++ pkg/lang/ir/v0/util_test.go | 78 ++++ pkg/lang/ir/v1/compile.go | 146 ++++---- pkg/lang/ir/v1/conda.go | 77 +++- pkg/lang/ir/v1/consts.go | 2 +- pkg/lang/ir/v1/editor.go | 11 +- pkg/lang/ir/v1/editor_test.go | 18 +- pkg/lang/ir/v1/get_conda.sh | 18 + pkg/lang/ir/v1/install_conda.sh | 9 + pkg/lang/ir/v1/interface.go | 64 +++- pkg/lang/ir/v1/julia.go | 47 +-- pkg/lang/ir/v1/python.go | 152 +++----- pkg/lang/ir/v1/r.go | 47 +-- pkg/lang/ir/v1/shell.go | 49 ++- pkg/lang/ir/v1/supervisor.go | 3 - pkg/lang/ir/v1/system.go | 114 +++--- pkg/lang/ir/v1/types.go | 19 +- pkg/lang/ir/v1/user.go | 33 +- pkg/lang/ir/v1/util.go | 13 + pkg/util/fileutil/file.go | 6 + 74 files changed, 4534 insertions(+), 459 deletions(-) create mode 100644 pkg/lang/frontend/starlark/interpreter.go create mode 100644 pkg/lang/frontend/starlark/v0/builtin/builtin.go create mode 100644 pkg/lang/frontend/starlark/v0/config/config.go create mode 100644 pkg/lang/frontend/starlark/v0/config/const.go create mode 100644 pkg/lang/frontend/starlark/v0/data/const.go create mode 100644 pkg/lang/frontend/starlark/v0/data/rule.go create mode 100644 pkg/lang/frontend/starlark/v0/data/util.go create mode 100644 pkg/lang/frontend/starlark/v0/install/const.go create mode 100644 pkg/lang/frontend/starlark/v0/install/install.go create mode 100644 pkg/lang/frontend/starlark/v0/interpreter.go create mode 100644 pkg/lang/frontend/starlark/v0/interpreter_test.go create mode 100644 pkg/lang/frontend/starlark/v0/io/const.go create mode 100644 pkg/lang/frontend/starlark/v0/io/io.go create mode 100644 pkg/lang/frontend/starlark/v0/mock/mock.go create mode 100644 pkg/lang/frontend/starlark/v0/runtime/const.go create mode 100644 pkg/lang/frontend/starlark/v0/runtime/runtime.go create mode 100644 pkg/lang/frontend/starlark/v0/starlark_suite_test.go create mode 100644 pkg/lang/frontend/starlark/v0/testdata/test.envd create mode 100644 pkg/lang/frontend/starlark/v0/universe/const.go create mode 100644 pkg/lang/frontend/starlark/v0/universe/universe.go create mode 100644 pkg/lang/ir/v0/cache.go create mode 100644 pkg/lang/ir/v0/checker.go create mode 100644 pkg/lang/ir/v0/compile.go create mode 100644 pkg/lang/ir/v0/conda.go create mode 100644 pkg/lang/ir/v0/consts.go rename pkg/lang/ir/{v1 => v0}/custom.go (100%) create mode 100644 pkg/lang/ir/v0/editor.go create mode 100644 pkg/lang/ir/v0/editor_test.go create mode 100644 pkg/lang/ir/v0/fs.go create mode 100644 pkg/lang/ir/v0/git.go rename pkg/lang/ir/{v1 => v0}/install-conda.sh (100%) rename pkg/lang/ir/{v1 => v0}/install-mamba.sh (100%) create mode 100644 pkg/lang/ir/v0/interface.go create mode 100644 pkg/lang/ir/v0/julia.go create mode 100644 pkg/lang/ir/v0/python.go create mode 100644 pkg/lang/ir/v0/r.go create mode 100644 pkg/lang/ir/v0/shell.go create mode 100644 pkg/lang/ir/v0/supervisor.go create mode 100644 pkg/lang/ir/v0/system.go create mode 100644 pkg/lang/ir/v0/types.go create mode 100644 pkg/lang/ir/v0/user.go create mode 100644 pkg/lang/ir/v0/util.go create mode 100644 pkg/lang/ir/v0/util_test.go create mode 100644 pkg/lang/ir/v1/get_conda.sh create mode 100644 pkg/lang/ir/v1/install_conda.sh diff --git a/examples/python-basic/build.envd b/examples/python-basic/build.envd index a006b104a..0e65ad2ab 100644 --- a/examples/python-basic/build.envd +++ b/examples/python-basic/build.envd @@ -1,16 +1,6 @@ def build(): - base(os="ubuntu20.04", language="python") - # config.pip_index(url = "https://pypi.tuna.tsinghua.edu.cn/simple") - install.python_packages( - [ - "via", - ] - ) - io.copy(host_path="./build.envd", envd_path="/") - runtime.command( - commands={ - "test": "ls /", - } - ) - runtime.environ(env={"ENVD_MODE": "DEV"}) - config.jupyter() + base(dev=True) + install.conda() + install.python() + install.python_packages(name=["ipython"]) + install.apt_packages(name=["ripgrep"]) diff --git a/pkg/app/create.go b/pkg/app/create.go index d7e9c0352..322bb23b3 100644 --- a/pkg/app/create.go +++ b/pkg/app/create.go @@ -27,7 +27,7 @@ import ( "github.com/tensorchord/envd/pkg/envd" "github.com/tensorchord/envd/pkg/home" - v1 "github.com/tensorchord/envd/pkg/lang/ir/v1" + v1 "github.com/tensorchord/envd/pkg/lang/ir/v0" "github.com/tensorchord/envd/pkg/ssh" sshconfig "github.com/tensorchord/envd/pkg/ssh/config" "github.com/tensorchord/envd/pkg/types" diff --git a/pkg/app/run.go b/pkg/app/run.go index 119bd97fb..3a45741bc 100644 --- a/pkg/app/run.go +++ b/pkg/app/run.go @@ -24,7 +24,7 @@ import ( "github.com/tensorchord/envd/pkg/builder" "github.com/tensorchord/envd/pkg/envd" "github.com/tensorchord/envd/pkg/home" - v1 "github.com/tensorchord/envd/pkg/lang/ir/v1" + v1 "github.com/tensorchord/envd/pkg/lang/ir/v0" "github.com/tensorchord/envd/pkg/ssh" "github.com/tensorchord/envd/pkg/util/fileutil" ) diff --git a/pkg/app/up.go b/pkg/app/up.go index 8ca7a5bd0..328658a40 100644 --- a/pkg/app/up.go +++ b/pkg/app/up.go @@ -26,7 +26,7 @@ import ( "github.com/tensorchord/envd/pkg/app/telemetry" "github.com/tensorchord/envd/pkg/envd" "github.com/tensorchord/envd/pkg/home" - v1 "github.com/tensorchord/envd/pkg/lang/ir/v1" + v1 "github.com/tensorchord/envd/pkg/lang/ir/v0" sshconfig "github.com/tensorchord/envd/pkg/ssh/config" "github.com/tensorchord/envd/pkg/types" ) diff --git a/pkg/builder/builder.go b/pkg/builder/builder.go index 022ba39cf..9cb5e15a8 100644 --- a/pkg/builder/builder.go +++ b/pkg/builder/builder.go @@ -33,7 +33,8 @@ import ( "github.com/tensorchord/envd/pkg/docker" "github.com/tensorchord/envd/pkg/flag" "github.com/tensorchord/envd/pkg/home" - starlark "github.com/tensorchord/envd/pkg/lang/frontend/starlark/v1" + "github.com/tensorchord/envd/pkg/lang/frontend/starlark" + starlarkv1 "github.com/tensorchord/envd/pkg/lang/frontend/starlark/v1" "github.com/tensorchord/envd/pkg/lang/ir" v1 "github.com/tensorchord/envd/pkg/lang/ir/v1" "github.com/tensorchord/envd/pkg/progress/progresswriter" @@ -139,7 +140,7 @@ func New(ctx context.Context, opt Options) (Builder, error) { } b.Client = cli - b.Interpreter = starlark.NewInterpreter(opt.BuildContextDir) + b.Interpreter = starlarkv1.NewInterpreter(opt.BuildContextDir) return b, nil } @@ -225,8 +226,9 @@ func (b generalBuilder) imageConfig(ctx context.Context) (string, error) { b.logger.Debugf("final entrypoint: {%s}\n", ep) env := b.graph.GetEnviron() + user := b.graph.GetUser() - data, err := ImageConfigStr(labels, ports, ep, env) + data, err := ImageConfigStr(labels, ports, ep, env, user) if err != nil { return "", errors.Wrap(err, "failed to get image config") } diff --git a/pkg/builder/builder_test.go b/pkg/builder/builder_test.go index ff9699d40..2a5970a5b 100644 --- a/pkg/builder/builder_test.go +++ b/pkg/builder/builder_test.go @@ -27,8 +27,8 @@ import ( mockbuildkitd "github.com/tensorchord/envd/pkg/buildkitd/mock" "github.com/tensorchord/envd/pkg/home" - mockstarlark "github.com/tensorchord/envd/pkg/lang/frontend/starlark/v1/mock" - v1 "github.com/tensorchord/envd/pkg/lang/ir/v1" + mockstarlark "github.com/tensorchord/envd/pkg/lang/frontend/starlark/v0/mock" + v1 "github.com/tensorchord/envd/pkg/lang/ir/v0" "github.com/tensorchord/envd/pkg/progress/compileui" compileuimock "github.com/tensorchord/envd/pkg/progress/compileui/mock" "github.com/tensorchord/envd/pkg/progress/progresswriter" diff --git a/pkg/builder/util.go b/pkg/builder/util.go index 03588df44..49eafb084 100644 --- a/pkg/builder/util.go +++ b/pkg/builder/util.go @@ -36,12 +36,12 @@ const ( ) func ImageConfigStr(labels map[string]string, ports map[string]struct{}, - entrypoint []string, env []string) (string, error) { + entrypoint []string, env []string, user string) (string, error) { pl := platforms.Normalize(platforms.DefaultSpec()) img := v1.Image{ Config: v1.ImageConfig{ Labels: labels, - User: "envd", + User: user, WorkingDir: "/", Env: env, ExposedPorts: ports, diff --git a/pkg/lang/frontend/starlark/interpreter.go b/pkg/lang/frontend/starlark/interpreter.go new file mode 100644 index 000000000..b0e906d82 --- /dev/null +++ b/pkg/lang/frontend/starlark/interpreter.go @@ -0,0 +1,6 @@ +package starlark + +type Interpreter interface { + Eval(script string) (interface{}, error) + ExecFile(filename string, funcname string) (interface{}, error) +} diff --git a/pkg/lang/frontend/starlark/v0/builtin/builtin.go b/pkg/lang/frontend/starlark/v0/builtin/builtin.go new file mode 100644 index 000000000..148ba725c --- /dev/null +++ b/pkg/lang/frontend/starlark/v0/builtin/builtin.go @@ -0,0 +1,20 @@ +// 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 builtin + +const ( + // BuildContextDir is the name of the directory that contains the build context. + BuildContextDir = "_build_context_dir" +) diff --git a/pkg/lang/frontend/starlark/v0/config/config.go b/pkg/lang/frontend/starlark/v0/config/config.go new file mode 100644 index 000000000..360d3bd45 --- /dev/null +++ b/pkg/lang/frontend/starlark/v0/config/config.go @@ -0,0 +1,228 @@ +// 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 config + +import ( + "github.com/cockroachdb/errors" + "github.com/sirupsen/logrus" + "go.starlark.net/starlark" + "go.starlark.net/starlarkstruct" + + ir "github.com/tensorchord/envd/pkg/lang/ir/v0" + "github.com/tensorchord/envd/pkg/util/starlarkutil" +) + +var ( + logger = logrus.WithField("frontend", "starlark") +) + +var Module = &starlarkstruct.Module{ + Name: "config", + Members: starlark.StringDict{ + "apt_source": starlark.NewBuiltin(ruleUbuntuAptSource, ruleFuncUbuntuAptSource), + "gpu": starlark.NewBuiltin(ruleGPU, ruleFuncGPU), + "jupyter": starlark.NewBuiltin(ruleJupyter, ruleFuncJupyter), + "cran_mirror": starlark.NewBuiltin( + ruleCRANMirror, ruleFuncCRANMirror), + "pip_index": starlark.NewBuiltin( + rulePyPIIndex, ruleFuncPyPIIndex), + "conda_channel": starlark.NewBuiltin( + ruleCondaChannel, ruleFuncCondaChannel), + "julia_pkg_server": starlark.NewBuiltin( + ruleJuliaPackageServer, ruleFuncJuliaPackageServer), + "rstudio_server": starlark.NewBuiltin(ruleRStudioServer, ruleFuncRStudioServer), + "entrypoint": starlark.NewBuiltin(ruleEntrypoint, ruleFuncEntrypoint), + "repo": starlark.NewBuiltin(ruleRepo, ruleFuncRepo), + }, +} + +func ruleFuncGPU(thread *starlark.Thread, _ *starlark.Builtin, + args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { + var numGPUs starlark.Int + + if err := starlark.UnpackArgs(ruleGPU, args, kwargs, + "count?", &numGPUs); err != nil { + return nil, err + } + + numGPUsInt, ok := numGPUs.Int64() + if ok { + ir.GPU(int(numGPUsInt)) + logger.Debugf("Using %d GPUs", int(numGPUsInt)) + } else { + logger.Debugf("Failed to convert gpu count to int64") + } + return starlark.None, nil +} + +func ruleFuncJupyter(thread *starlark.Thread, _ *starlark.Builtin, + args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { + var token starlark.String + var port starlark.Int + + if err := starlark.UnpackArgs(ruleJupyter, args, kwargs, + "token?", &token, "port?", &port); err != nil { + return nil, err + } + + pwdStr := token.GoString() + + portInt, ok := port.Int64() + if !ok { + return nil, errors.New("port must be an integer") + } + logger.Debugf("rule `%s` is invoked, password=%s, port=%d", + ruleJupyter, pwdStr, portInt) + if err := ir.Jupyter(pwdStr, portInt); err != nil { + return nil, err + } + + return starlark.None, nil +} + +func ruleFuncPyPIIndex(thread *starlark.Thread, _ *starlark.Builtin, + args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { + var url, extraURL starlark.String + + if err := starlark.UnpackArgs(rulePyPIIndex, args, kwargs, + "url?", &url, "extra_url?", &extraURL); err != nil { + return nil, err + } + + indexStr := url.GoString() + extraIndexStr := extraURL.GoString() + + logger.Debugf("rule `%s` is invoked, index=%s, extraIndex=%s", + rulePyPIIndex, indexStr, extraIndexStr) + if err := ir.PyPIIndex(indexStr, extraIndexStr); err != nil { + return nil, err + } + + return starlark.None, nil +} + +func ruleFuncCRANMirror(thread *starlark.Thread, _ *starlark.Builtin, + args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { + var url starlark.String + + if err := starlark.UnpackArgs(ruleCRANMirror, args, kwargs, + "url?", &url); err != nil { + return nil, err + } + + urlStr := url.GoString() + + logger.Debugf("rule `%s` is invoked, url=%s", ruleCRANMirror, urlStr) + if err := ir.CRANMirror(urlStr); err != nil { + return nil, err + } + return starlark.None, nil +} + +func ruleFuncJuliaPackageServer(thread *starlark.Thread, _ *starlark.Builtin, + args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { + var url starlark.String + + if err := starlark.UnpackArgs(ruleJuliaPackageServer, args, kwargs, + "url?", &url); err != nil { + return nil, err + } + + urlStr := url.GoString() + + logger.Debugf("rule `%s` is invoked, url=%s", ruleJuliaPackageServer, urlStr) + if err := ir.JuliaPackageServer(urlStr); err != nil { + return nil, err + } + return starlark.None, nil +} + +func ruleFuncUbuntuAptSource(thread *starlark.Thread, _ *starlark.Builtin, + args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { + var source starlark.String + + if err := starlark.UnpackArgs(ruleUbuntuAptSource, args, kwargs, + "source?", &source); err != nil { + return nil, err + } + + sourceStr := source.GoString() + + logger.Debugf("rule `%s` is invoked, source=%s", ruleUbuntuAptSource, sourceStr) + if err := ir.UbuntuAPT(sourceStr); err != nil { + return nil, err + } + + return starlark.None, nil +} + +func ruleFuncRStudioServer(thread *starlark.Thread, _ *starlark.Builtin, + args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { + if err := ir.RStudioServer(); err != nil { + return nil, err + } + + return starlark.None, nil +} + +func ruleFuncCondaChannel(thread *starlark.Thread, _ *starlark.Builtin, + args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { + var channel string + var useMamba bool + + if err := starlark.UnpackArgs(ruleCondaChannel, args, kwargs, + "channel?", &channel, "use_mamba?", &useMamba); err != nil { + return nil, err + } + + logger.Debugf("rule `%s` is invoked, channel=%s, use_mamba=%t\n", + ruleCondaChannel, channel, useMamba) + if err := ir.CondaChannel(channel, useMamba); err != nil { + return nil, err + } + + return starlark.None, nil +} + +func ruleFuncEntrypoint(thread *starlark.Thread, _ *starlark.Builtin, + args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { + var argv *starlark.List + + if err := starlark.UnpackArgs(ruleEntrypoint, args, kwargs, "args", &argv); err != nil { + return nil, err + } + + argList, err := starlarkutil.ToStringSlice(argv) + if err != nil { + return nil, err + } + + logger.Debugf("user defined entrypoints: {%s}\n", argList) + ir.Entrypoint(argList) + return starlark.None, nil +} + +func ruleFuncRepo(thread *starlark.Thread, _ *starlark.Builtin, + args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { + var url, description string + + if err := starlark.UnpackArgs(ruleRepo, args, kwargs, "url", &url, "description?", &description); err != nil { + return nil, err + } + + logger.Debugf("repo info: url=%s, description=%s", url, description) + ir.Repo(url, description) + return starlark.None, nil +} diff --git a/pkg/lang/frontend/starlark/v0/config/const.go b/pkg/lang/frontend/starlark/v0/config/const.go new file mode 100644 index 000000000..3e084f794 --- /dev/null +++ b/pkg/lang/frontend/starlark/v0/config/const.go @@ -0,0 +1,28 @@ +// 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 config + +const ( + ruleUbuntuAptSource = "config.apt_source" + rulePyPIIndex = "config.pip_index" + ruleCRANMirror = "config.cran_mirror" + ruleJupyter = "config.jupyter" + ruleCondaChannel = "config.conda_channel" + ruleGPU = "config.gpu" + ruleJuliaPackageServer = "config.julia_pkg_server" + ruleRStudioServer = "config.rstudio_server" + ruleEntrypoint = "config.entrypoint" + ruleRepo = "config.repo" +) diff --git a/pkg/lang/frontend/starlark/v0/data/const.go b/pkg/lang/frontend/starlark/v0/data/const.go new file mode 100644 index 000000000..9af123827 --- /dev/null +++ b/pkg/lang/frontend/starlark/v0/data/const.go @@ -0,0 +1,21 @@ +// 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 data + +const ( + ruleEnvdManagedDataSource = "data.envd" + huggingFaceDatasetPath = "~/.cache/huggingface" + dglFaceDatasetPath = "~/.dgl" +) diff --git a/pkg/lang/frontend/starlark/v0/data/rule.go b/pkg/lang/frontend/starlark/v0/data/rule.go new file mode 100644 index 000000000..44c804b29 --- /dev/null +++ b/pkg/lang/frontend/starlark/v0/data/rule.go @@ -0,0 +1,54 @@ +// 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 data + +import ( + "github.com/sirupsen/logrus" + "go.starlark.net/starlark" + "go.starlark.net/starlarkstruct" + + envddata "github.com/tensorchord/envd/pkg/data" +) + +var ( + logger = logrus.WithField("frontend", "starlark") +) + +var Module = &starlarkstruct.Module{ + Name: "data", + Members: starlark.StringDict{ + "envd": starlark.NewBuiltin(ruleEnvdManagedDataSource, ruleValueEnvdManagedDataSource), + "path": &starlarkstruct.Module{ + Name: "path", + Members: starlark.StringDict{ + "huggingface": starlark.String(huggingFaceDatasetPath), + "dgl": starlark.String(dglFaceDatasetPath), + }}, + }, +} + +func ruleValueEnvdManagedDataSource(thread *starlark.Thread, _ *starlark.Builtin, + args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { + var name starlark.String + + if err := starlark.UnpackArgs(ruleEnvdManagedDataSource, args, kwargs, + "name?", &name); err != nil { + return nil, err + } + logger.Debugf("rule `%s` is invoked, name=%s", + ruleEnvdManagedDataSource, name) + + return NewDataSourceValue(envddata.NewEnvdManagedDataSource(name.GoString())), nil +} diff --git a/pkg/lang/frontend/starlark/v0/data/util.go b/pkg/lang/frontend/starlark/v0/data/util.go new file mode 100644 index 000000000..8f8692c73 --- /dev/null +++ b/pkg/lang/frontend/starlark/v0/data/util.go @@ -0,0 +1,53 @@ +// 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 data + +import ( + "go.starlark.net/starlark" + + envddata "github.com/tensorchord/envd/pkg/data" +) + +type DataSourceValue struct { + source envddata.DataSource +} + +func (d DataSourceValue) Init() error { + return d.source.Init() +} + +func (d DataSourceValue) GetHostDir() (string, error) { + return d.source.GetHostDir() +} + +func (d DataSourceValue) String() string { + return d.source.Type() +} + +func (d DataSourceValue) Type() string { + return d.source.Type() +} + +func (d DataSourceValue) Freeze() {} + +func (d DataSourceValue) Truth() starlark.Bool { return true } + +func (d DataSourceValue) Hash() (uint32, error) { + return d.source.Hash() +} + +func NewDataSourceValue(source envddata.DataSource) *DataSourceValue { + return &DataSourceValue{source: source} +} diff --git a/pkg/lang/frontend/starlark/v0/install/const.go b/pkg/lang/frontend/starlark/v0/install/const.go new file mode 100644 index 000000000..3f39a8ff8 --- /dev/null +++ b/pkg/lang/frontend/starlark/v0/install/const.go @@ -0,0 +1,25 @@ +// 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 install + +const ( + ruleSystemPackage = "install.apt_packages" + rulePyPIPackage = "install.python_packages" + ruleRPackage = "install.r_packages" + ruleCUDA = "install.cuda" + ruleVSCode = "install.vscode_extensions" + ruleConda = "install.conda_packages" + ruleJulia = "install.julia_packages" +) diff --git a/pkg/lang/frontend/starlark/v0/install/install.go b/pkg/lang/frontend/starlark/v0/install/install.go new file mode 100644 index 000000000..feea5f48a --- /dev/null +++ b/pkg/lang/frontend/starlark/v0/install/install.go @@ -0,0 +1,204 @@ +// 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 install + +import ( + "github.com/cockroachdb/errors" + "github.com/sirupsen/logrus" + "go.starlark.net/starlark" + "go.starlark.net/starlarkstruct" + + ir "github.com/tensorchord/envd/pkg/lang/ir/v0" + "github.com/tensorchord/envd/pkg/util/starlarkutil" +) + +var ( + logger = logrus.WithField("frontend", "starlark") +) + +var Module = &starlarkstruct.Module{ + Name: "install", + Members: starlark.StringDict{ + "python_packages": starlark.NewBuiltin(rulePyPIPackage, ruleFuncPyPIPackage), + "r_packages": starlark.NewBuiltin(ruleRPackage, ruleFuncRPackage), + "apt_packages": starlark.NewBuiltin(ruleSystemPackage, ruleFuncSystemPackage), + "cuda": starlark.NewBuiltin(ruleCUDA, ruleFuncCUDA), + "vscode_extensions": starlark.NewBuiltin(ruleVSCode, ruleFuncVSCode), + "conda_packages": starlark.NewBuiltin(ruleConda, ruleFuncConda), + "julia_packages": starlark.NewBuiltin(ruleJulia, ruleFuncJulia), + }, +} + +func ruleFuncPyPIPackage(thread *starlark.Thread, _ *starlark.Builtin, + args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { + var name *starlark.List + var requirementsFile starlark.String + var wheels *starlark.List + + if err := starlark.UnpackArgs(rulePyPIPackage, args, kwargs, + "name?", &name, "requirements?", &requirementsFile, "local_wheels?", &wheels); err != nil { + return nil, err + } + + nameList, err := starlarkutil.ToStringSlice(name) + if err != nil { + return nil, err + } + + requirementsFileStr := requirementsFile.GoString() + + localWheels, err := starlarkutil.ToStringSlice(wheels) + if err != nil { + return nil, err + } + + logger.Debugf("rule `%s` is invoked, name=%v, requirements=%s, local_wheels=%s", + rulePyPIPackage, nameList, requirementsFileStr, localWheels) + + err = ir.PyPIPackage(nameList, requirementsFileStr, localWheels) + return starlark.None, err +} + +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, err := starlarkutil.ToStringSlice(name) + if err != nil { + return nil, err + } + + logger.Debugf("rule `%s` is invoked, name=%v", ruleRPackage, nameList) + ir.RPackage(nameList) + + return starlark.None, nil +} + +func ruleFuncJulia(thread *starlark.Thread, _ *starlark.Builtin, + args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { + var name *starlark.List + + if err := starlark.UnpackArgs(ruleJulia, + args, kwargs, "name", &name); err != nil { + return nil, err + } + + nameList, err := starlarkutil.ToStringSlice(name) + if err != nil { + return nil, err + } + logger.Debugf("rule `%s` is invoked, name=%v", ruleJulia, nameList) + ir.JuliaPackage(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 + + if err := starlark.UnpackArgs(ruleSystemPackage, + args, kwargs, "name?", &name); err != nil { + return nil, err + } + + nameList, err := starlarkutil.ToStringSlice(name) + if err != nil { + return nil, err + } + + logger.Debugf("rule `%s` is invoked, name=%v", ruleSystemPackage, nameList) + ir.SystemPackage(nameList) + + return starlark.None, nil +} + +func ruleFuncCUDA(thread *starlark.Thread, _ *starlark.Builtin, + args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { + var version, cudnn string + + if err := starlark.UnpackArgs(ruleCUDA, args, kwargs, + "version", &version, "cudnn?", &cudnn); err != nil { + return nil, err + } + + logger.Debugf("rule `%s` is invoked, version=%s, cudnn=%s", + ruleCUDA, version, cudnn) + ir.CUDA(version, cudnn) + + return starlark.None, nil +} + +func ruleFuncVSCode(thread *starlark.Thread, _ *starlark.Builtin, + args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { + var plugins *starlark.List + + if err := starlark.UnpackArgs(ruleVSCode, + args, kwargs, "name", &plugins); err != nil { + return nil, err + } + + pluginList, err := starlarkutil.ToStringSlice(plugins) + if err != nil { + return nil, err + } + + logger.Debugf("rule `%s` is invoked, plugins=%v", ruleVSCode, pluginList) + if err := ir.VSCodePlugins(pluginList); err != nil { + return starlark.None, err + } + + return starlark.None, nil +} + +func ruleFuncConda(thread *starlark.Thread, _ *starlark.Builtin, + args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { + var name, channel *starlark.List + var envFile starlark.String + + if err := starlark.UnpackArgs(ruleConda, + args, kwargs, "name?", &name, "channel?", &channel, "env_file?", &envFile); err != nil { + return nil, err + } + + nameList, err := starlarkutil.ToStringSlice(name) + if err != nil { + return nil, err + } + + channelList, err := starlarkutil.ToStringSlice(channel) + if err != nil { + return nil, err + } + + envFileStr := envFile.GoString() + if envFileStr != "" { + if (len(nameList) != 0) || (len(channelList) != 0) { + return nil, errors.New("env_file and name/channel are mutually exclusive") + } + } + + logger.Debugf("rule `%s` is invoked, name=%v, channel=%v, env_file=%s", ruleConda, nameList, channelList, envFileStr) + if err := ir.CondaPackage(nameList, channelList, envFileStr); err != nil { + return starlark.None, err + } + + return starlark.None, nil +} diff --git a/pkg/lang/frontend/starlark/v0/interpreter.go b/pkg/lang/frontend/starlark/v0/interpreter.go new file mode 100644 index 000000000..0ce5f2c15 --- /dev/null +++ b/pkg/lang/frontend/starlark/v0/interpreter.go @@ -0,0 +1,209 @@ +// 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 v1 + +import ( + "bytes" + "hash/fnv" + "io/fs" + "os" + "path/filepath" + "strconv" + "strings" + + "github.com/cockroachdb/errors" + "github.com/sirupsen/logrus" + "go.starlark.net/starlark" + + interp "github.com/tensorchord/envd/pkg/lang/frontend/starlark" + "github.com/tensorchord/envd/pkg/lang/frontend/starlark/v0/config" + "github.com/tensorchord/envd/pkg/lang/frontend/starlark/v0/data" + "github.com/tensorchord/envd/pkg/lang/frontend/starlark/v0/install" + "github.com/tensorchord/envd/pkg/lang/frontend/starlark/v0/io" + "github.com/tensorchord/envd/pkg/lang/frontend/starlark/v0/runtime" + "github.com/tensorchord/envd/pkg/lang/frontend/starlark/v0/universe" + "github.com/tensorchord/envd/pkg/util/fileutil" +) + +type entry struct { + globals starlark.StringDict + err error +} + +// generalInterpreter is the interpreter implementation for Starlark. +// Please refer to https://github.com/google/starlark-go +type generalInterpreter struct { + predeclared starlark.StringDict + buildContextDir string + cache map[string]*entry +} + +func NewInterpreter(buildContextDir string) interp.Interpreter { + // Register envd rules and built-in variables to Starlark. + universe.RegisterEnvdRules() + universe.RegisterBuildContext(buildContextDir) + + return &generalInterpreter{ + predeclared: starlark.StringDict{ + "install": install.Module, + "config": config.Module, + "io": io.Module, + "runtime": runtime.Module, + "data": data.Module, + }, + buildContextDir: buildContextDir, + cache: make(map[string]*entry), + } +} + +func (s *generalInterpreter) NewThread(module string) *starlark.Thread { + thread := &starlark.Thread{ + Name: module, + Load: s.load, + } + return thread +} + +func (s *generalInterpreter) load(thread *starlark.Thread, module string) (starlark.StringDict, error) { + return s.exec(thread, module) +} + +func (s *generalInterpreter) exec(thread *starlark.Thread, module string) (starlark.StringDict, error) { + // There are two cases: + // 1. module exists + // 2. there's an explicit `nil` placeholder for module in s.cache + // 3. module does not exist in s.cache + e, ok := s.cache[module] + + // Case 1. + if e != nil { + return e.globals, e.err + } + + // Case 2. + // There is an explicit `nil` for module, which means we are in the middle of exec module. + if ok { + return nil, errors.Newf("Detected cycle import during parsing %s", module) + } + + // Case 3. + // Add a placeholder to indicate "load in progress". + s.cache[module] = nil + + if !strings.HasPrefix(module, universe.GitPrefix) { + var data interface{} + globals, err := starlark.ExecFile(thread, module, data, s.predeclared) + e = &entry{globals, err} + } else { + // exec remote git repo + url := module[len(universe.GitPrefix):] + path, err := fileutil.DownloadOrUpdateGitRepo(url) + if err != nil { + return nil, err + } + globals, err := s.loadGitModule(thread, path) + e = &entry{globals, err} + } + + // Update the cache. + s.cache[module] = e + + return e.globals, e.err +} + +func (s *generalInterpreter) loadGitModule(thread *starlark.Thread, path string) (globals starlark.StringDict, err error) { + var src interface{} + globals = starlark.StringDict{} + logger := logrus.WithField("file", thread.Name) + logger.Debugf("load git module from: %s", path) + err = filepath.WalkDir(path, func(path string, d fs.DirEntry, err error) error { + if err != nil { + return err + } + if d.IsDir() || !strings.HasSuffix(d.Name(), ".envd") { + return nil + } + dict, err := starlark.ExecFile(thread, path, src, s.predeclared) + if err != nil { + return err + } + for key, val := range dict { + if _, exist := globals[key]; exist { + return errors.Newf("found duplicated object name: %s in %s", key, path) + } + if !strings.HasPrefix(key, "_") { + globals[key] = val + } + } + return nil + }) + return +} + +func (s generalInterpreter) ExecFile(filename string, funcname string) (interface{}, error) { + logrus.WithField("filename", filename).Debug("interprete the file") + thread := s.NewThread(filename) + globals, err := s.exec(thread, filename) + if err != nil { + return nil, err + } + if funcname != "" { + logrus.Debugf("Execute %s func", funcname) + if globals.Has(funcname) { + buildVar := globals[funcname] + if fn, ok := buildVar.(*starlark.Function); ok { + _, err := starlark.Call(thread, fn, nil, nil) + if err != nil { + return nil, errors.Wrapf(err, "Exception when exec %s func", funcname) + } + } else { + return nil, errors.Errorf("%s is not a function", funcname) + } + } else { + return nil, errors.Errorf("envd file doesn't has %s function", funcname) + } + + } + return globals, nil +} + +func (s generalInterpreter) Eval(script string) (interface{}, error) { + thread := s.NewThread(script) + return starlark.ExecFile(thread, "", script, s.predeclared) +} + +func GetEnvdProgramHash(filename string) (string, error) { + envdSrc, err := os.ReadFile(filename) + if err != nil { + return "", err + } + // No Check builtin or predeclared for now + funcAlwaysHas := func(x string) bool { + return true + } + _, prog, err := starlark.SourceProgram(filename, envdSrc, funcAlwaysHas) + if err != nil { + return "", err + } + buf := new(bytes.Buffer) + err = prog.Write(buf) + if err != nil { + return "", err + } + h := fnv.New64a() + h.Write(buf.Bytes()) + hashsum := h.Sum64() + return strconv.FormatUint(hashsum, 16), nil +} diff --git a/pkg/lang/frontend/starlark/v0/interpreter_test.go b/pkg/lang/frontend/starlark/v0/interpreter_test.go new file mode 100644 index 000000000..5872d83c4 --- /dev/null +++ b/pkg/lang/frontend/starlark/v0/interpreter_test.go @@ -0,0 +1,29 @@ +// 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 v1 + +import ( + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("Starlark", func() { + It("should be able to compile envd", func() { + filename := "testdata/test.envd" + hash, err := GetEnvdProgramHash(filename) + Expect(err).NotTo(HaveOccurred()) + Expect(hash).To(Equal("cff1c81818116d42")) + }) +}) diff --git a/pkg/lang/frontend/starlark/v0/io/const.go b/pkg/lang/frontend/starlark/v0/io/const.go new file mode 100644 index 000000000..d46769e9e --- /dev/null +++ b/pkg/lang/frontend/starlark/v0/io/const.go @@ -0,0 +1,20 @@ +// 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 io + +const ( + ruleCopy = "io.copy" + ruleHTTP = "io.http" +) diff --git a/pkg/lang/frontend/starlark/v0/io/io.go b/pkg/lang/frontend/starlark/v0/io/io.go new file mode 100644 index 000000000..0367a1a23 --- /dev/null +++ b/pkg/lang/frontend/starlark/v0/io/io.go @@ -0,0 +1,70 @@ +// 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 io + +import ( + "github.com/sirupsen/logrus" + "go.starlark.net/starlark" + "go.starlark.net/starlarkstruct" + + ir "github.com/tensorchord/envd/pkg/lang/ir/v0" +) + +var ( + logger = logrus.WithField("frontend", "starlark") +) + +var Module = &starlarkstruct.Module{ + Name: "io", + Members: starlark.StringDict{ + "copy": starlark.NewBuiltin(ruleCopy, ruleFuncCopy), + "http": starlark.NewBuiltin(ruleHTTP, ruleFuncHTTP), + }, +} + +func ruleFuncCopy(thread *starlark.Thread, _ *starlark.Builtin, + args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { + var source, destination starlark.String + + if err := starlark.UnpackArgs(ruleCopy, args, kwargs, + "host_path?", &source, "envd_path?", &destination); err != nil { + return nil, err + } + + sourceStr := source.GoString() + destinationStr := destination.GoString() + + logger.Debugf("rule `%s` is invoked, src=%s, dest=%s\n", + ruleCopy, sourceStr, destinationStr) + ir.Copy(sourceStr, destinationStr) + + return starlark.None, nil +} + +func ruleFuncHTTP(thread *starlark.Thread, _ *starlark.Builtin, + args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { + var url, checksum, filename string + if err := starlark.UnpackArgs(ruleHTTP, args, kwargs, + "url", &url, "checksum?", &checksum, "filename?", &filename); err != nil { + return nil, err + } + + logger.Debugf("rule `%s` is invoked, ruleHTTP, url=%s, checksum=%s, filename=%s\n", + ruleHTTP, url, checksum, filename) + if err := ir.HTTP(url, checksum, filename); err != nil { + return nil, err + } + return starlark.None, nil +} diff --git a/pkg/lang/frontend/starlark/v0/mock/mock.go b/pkg/lang/frontend/starlark/v0/mock/mock.go new file mode 100644 index 000000000..665620614 --- /dev/null +++ b/pkg/lang/frontend/starlark/v0/mock/mock.go @@ -0,0 +1,64 @@ +// Code generated by MockGen. DO NOT EDIT. +// Source: pkg/lang/frontend/starlark/interpreter.go + +// Package mock is a generated GoMock package. +package mock + +import ( + reflect "reflect" + + gomock "github.com/golang/mock/gomock" +) + +// MockInterpreter is a mock of Interpreter interface. +type MockInterpreter struct { + ctrl *gomock.Controller + recorder *MockInterpreterMockRecorder +} + +// MockInterpreterMockRecorder is the mock recorder for MockInterpreter. +type MockInterpreterMockRecorder struct { + mock *MockInterpreter +} + +// NewMockInterpreter creates a new mock instance. +func NewMockInterpreter(ctrl *gomock.Controller) *MockInterpreter { + mock := &MockInterpreter{ctrl: ctrl} + mock.recorder = &MockInterpreterMockRecorder{mock} + return mock +} + +// EXPECT returns an object that allows the caller to indicate expected use. +func (m *MockInterpreter) EXPECT() *MockInterpreterMockRecorder { + return m.recorder +} + +// Eval mocks base method. +func (m *MockInterpreter) Eval(script string) (interface{}, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "Eval", script) + ret0, _ := ret[0].(interface{}) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// Eval indicates an expected call of Eval. +func (mr *MockInterpreterMockRecorder) Eval(script interface{}) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Eval", reflect.TypeOf((*MockInterpreter)(nil).Eval), script) +} + +// ExecFile mocks base method. +func (m *MockInterpreter) ExecFile(filename, funcname string) (interface{}, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "ExecFile", filename, funcname) + ret0, _ := ret[0].(interface{}) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// ExecFile indicates an expected call of ExecFile. +func (mr *MockInterpreterMockRecorder) ExecFile(filename, funcname interface{}) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ExecFile", reflect.TypeOf((*MockInterpreter)(nil).ExecFile), filename, funcname) +} diff --git a/pkg/lang/frontend/starlark/v0/runtime/const.go b/pkg/lang/frontend/starlark/v0/runtime/const.go new file mode 100644 index 000000000..584e4c069 --- /dev/null +++ b/pkg/lang/frontend/starlark/v0/runtime/const.go @@ -0,0 +1,24 @@ +// 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 runtime + +const ( + ruleCommand = "runtime.command" + ruleExpose = "runtime.expose" + ruleDaemon = "runtime.daemon" + ruleEnviron = "runtime.environ" + ruleMount = "runtime.mount" + ruleInitScript = "runtime.init" +) diff --git a/pkg/lang/frontend/starlark/v0/runtime/runtime.go b/pkg/lang/frontend/starlark/v0/runtime/runtime.go new file mode 100644 index 000000000..b3b26b4ee --- /dev/null +++ b/pkg/lang/frontend/starlark/v0/runtime/runtime.go @@ -0,0 +1,238 @@ +// 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 runtime + +import ( + "net" + "os/user" + "path/filepath" + "strings" + + "github.com/cockroachdb/errors" + "github.com/sirupsen/logrus" + "go.starlark.net/starlark" + "go.starlark.net/starlarkstruct" + + "github.com/tensorchord/envd/pkg/lang/frontend/starlark/v0/data" + ir "github.com/tensorchord/envd/pkg/lang/ir/v0" + "github.com/tensorchord/envd/pkg/util/fileutil" + "github.com/tensorchord/envd/pkg/util/starlarkutil" +) + +var ( + logger = logrus.WithField("frontend", "starlark") +) + +var Module = &starlarkstruct.Module{ + Name: "runtime", + Members: starlark.StringDict{ + "command": starlark.NewBuiltin(ruleCommand, ruleFuncCommand), + "daemon": starlark.NewBuiltin(ruleDaemon, ruleFuncDaemon), + "expose": starlark.NewBuiltin(ruleExpose, ruleFuncExpose), + "environ": starlark.NewBuiltin(ruleEnviron, ruleFuncEnviron), + "mount": starlark.NewBuiltin(ruleMount, ruleFuncMount), + "init": starlark.NewBuiltin(ruleInitScript, ruleFuncInitScript), + }, +} + +func ruleFuncCommand(thread *starlark.Thread, _ *starlark.Builtin, + args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { + var commands starlark.IterableMapping + + if err := starlark.UnpackArgs(ruleCommand, args, kwargs, + "commands?", &commands); err != nil { + return nil, err + } + + commandsMap := make(map[string]string) + for _, tuple := range commands.Items() { + if len(tuple) != 2 { + return nil, errors.Newf("invalid command in %s", ruleCommand) + } + + commandsMap[tuple[0].(starlark.String).GoString()] = + tuple[1].(starlark.String).GoString() + } + + logger.Debugf("rule `%s` is invoked, commands: %v", + ruleCommand, commandsMap) + + ir.RuntimeCommands(commandsMap) + return starlark.None, nil +} + +func ruleFuncDaemon(thread *starlark.Thread, _ *starlark.Builtin, + args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { + var commands *starlark.List + + if err := starlark.UnpackArgs(ruleDaemon, args, kwargs, "commands", &commands); err != nil { + return nil, err + } + + commandList := [][]string{} + if commands != nil { + for i := 0; i < commands.Len(); i++ { + args, ok := commands.Index(i).(*starlark.List) + if !ok { + return nil, errors.Newf("invalid daemon commands (%s)", commands.Index(i).String()) + } + argList := []string{} + for j := 0; j < args.Len(); j++ { + argList = append(argList, args.Index(j).(starlark.String).GoString()) + } + commandList = append(commandList, argList) + } + + logger.Debugf("rule `%s` is invoked, commands=%v", ruleDaemon, commandList) + ir.RuntimeDaemon(commandList) + } + return starlark.None, nil +} + +func ruleFuncExpose(thread *starlark.Thread, _ *starlark.Builtin, + args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { + var ( + envdPort starlark.Int + hostPort = starlark.MakeInt(0) // 0 means envd can randomly choose a free port + serviceName = starlark.String("") + listeningAddr = starlark.String("127.0.0.1") // default to lisen only on local loopback interface + ) + + if err := starlark.UnpackArgs(ruleExpose, + args, kwargs, "envd_port", &envdPort, "host_port?", &hostPort, "service?", &serviceName, "listen_addr?", &listeningAddr); err != nil { + return nil, err + } + envdPortInt, ok := envdPort.Int64() + if !ok || envdPortInt < 1 || envdPortInt > 65535 { + return nil, errors.New("envd_port must be a positive integer less than 65535") + } + hostPortInt, ok := hostPort.Int64() + if !ok || hostPortInt < 0 || hostPortInt > 65535 { + return nil, errors.New("host_port must be a positive integer less than 65535") + } + serviceNameStr := serviceName.GoString() + listeningAddrStr := listeningAddr.GoString() + if net.ParseIP(listeningAddrStr) == nil { + return nil, errors.New("listening_addr must be a valid IP address") + } + + logger.Debugf("rule `%s` is invoked, envd_port=%d, host_port=%d, service=%s", ruleExpose, envdPortInt, hostPortInt, serviceNameStr) + err := ir.RuntimeExpose(int(envdPortInt), int(hostPortInt), serviceNameStr, listeningAddrStr) + return starlark.None, err +} + +func ruleFuncEnviron(thread *starlark.Thread, _ *starlark.Builtin, + args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { + var env starlark.IterableMapping + var path *starlark.List + + if err := starlark.UnpackArgs(ruleCommand, args, kwargs, + "env?", &env, "extra_path?", &path); err != nil { + return nil, err + } + + envMap := make(map[string]string) + if env != nil { + for _, tuple := range env.Items() { + if len(tuple) != 2 { + return nil, errors.Newf("invalid env (%s)", tuple.String()) + } + envMap[tuple[0].(starlark.String).GoString()] = tuple[1].(starlark.String).GoString() + } + } + + pathList, err := starlarkutil.ToStringSlice(path) + if err != nil { + return nil, err + } + + logger.Debugf("rule `%s` is invoked, env: %v, extra_path: %v", ruleEnviron, envMap, pathList) + ir.RuntimeEnviron(envMap, pathList) + return starlark.None, nil +} + +func ruleFuncMount(thread *starlark.Thread, _ *starlark.Builtin, + args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { + var source starlark.Value + var destination starlark.String + + if err := starlark.UnpackArgs(ruleMount, args, kwargs, + "host_path?", &source, "envd_path?", &destination); err != nil { + return nil, err + } + + var sourceStr string + var err error + + if v, ok := source.(*data.DataSourceValue); ok { + err = v.Init() + if err != nil { + return starlark.None, err + } + sourceStr, err = v.GetHostDir() + if err != nil { + return starlark.None, err + } + } else if vs, ok := source.(starlark.String); ok { + sourceStr = vs.GoString() + } else { + return starlark.None, errors.New("invalid data source") + } + + destinationStr := destination.GoString() + + logger.Debugf("rule `%s` is invoked, src=%s, dest=%s", + ruleMount, sourceStr, destinationStr) + + // Expand source directory based on host user + usr, _ := user.Current() + dir := usr.HomeDir + if sourceStr == "~" { + sourceStr = dir + } else if strings.HasPrefix(sourceStr, "~/") { + sourceStr = filepath.Join(dir, sourceStr[2:]) + } + // Expand dest directory based on container user envd + dir = fileutil.EnvdHomeDir() + if destinationStr == "~" { + destinationStr = dir + } else if strings.HasPrefix(destinationStr, "~/") { + destinationStr = fileutil.EnvdHomeDir(destinationStr[2:]) + } + ir.Mount(sourceStr, destinationStr) + + return starlark.None, nil +} + +func ruleFuncInitScript(thread *starlark.Thread, _ *starlark.Builtin, + args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { + var commands *starlark.List + + if err := starlark.UnpackArgs(ruleCommand, args, kwargs, + "commands?", &commands); err != nil { + return nil, err + } + + commandsSlice, err := starlarkutil.ToStringSlice(commands) + if err != nil { + return nil, err + } + + logger.Debugf("rule `%s` is invoked, commands: %v", + ruleInitScript, commandsSlice) + + ir.RuntimeInitScript(commandsSlice) + return starlark.None, nil +} diff --git a/pkg/lang/frontend/starlark/v0/starlark_suite_test.go b/pkg/lang/frontend/starlark/v0/starlark_suite_test.go new file mode 100644 index 000000000..3fbdd42a5 --- /dev/null +++ b/pkg/lang/frontend/starlark/v0/starlark_suite_test.go @@ -0,0 +1,27 @@ +// 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 v1 + +import ( + "testing" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +func TestManager(t *testing.T) { + RegisterFailHandler(Fail) + RunSpecs(t, "Starlark Suite") +} diff --git a/pkg/lang/frontend/starlark/v0/testdata/test.envd b/pkg/lang/frontend/starlark/v0/testdata/test.envd new file mode 100644 index 000000000..9c0e3cb02 --- /dev/null +++ b/pkg/lang/frontend/starlark/v0/testdata/test.envd @@ -0,0 +1,8 @@ +def build(): + base(os="ubuntu20.04", language="python3") + install.python_packages( + name=[ + "numpy", + ] + ) + shell("zsh") diff --git a/pkg/lang/frontend/starlark/v0/universe/const.go b/pkg/lang/frontend/starlark/v0/universe/const.go new file mode 100644 index 000000000..7e2b74e68 --- /dev/null +++ b/pkg/lang/frontend/starlark/v0/universe/const.go @@ -0,0 +1,25 @@ +// 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 universe + +const ( + ruleBase = "base" + ruleShell = "shell" + ruleRun = "run" + ruleGitConfig = "git_config" + ruleInclude = "include" + + GitPrefix = "git@" +) diff --git a/pkg/lang/frontend/starlark/v0/universe/universe.go b/pkg/lang/frontend/starlark/v0/universe/universe.go new file mode 100644 index 000000000..2b990327b --- /dev/null +++ b/pkg/lang/frontend/starlark/v0/universe/universe.go @@ -0,0 +1,141 @@ +// 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 universe + +import ( + "fmt" + + "github.com/sirupsen/logrus" + "go.starlark.net/starlark" + "go.starlark.net/starlarkstruct" + + "github.com/tensorchord/envd/pkg/lang/frontend/starlark/v0/builtin" + ir "github.com/tensorchord/envd/pkg/lang/ir/v0" + "github.com/tensorchord/envd/pkg/util/starlarkutil" +) + +var ( + logger = logrus.WithField("frontend", "starlark") +) + +// RegisterEnvdRules registers built-in envd rules into the global namespace. +func RegisterEnvdRules() { + starlark.Universe[ruleBase] = starlark.NewBuiltin(ruleBase, ruleFuncBase) + starlark.Universe[ruleShell] = starlark.NewBuiltin(ruleShell, ruleFuncShell) + starlark.Universe[ruleRun] = starlark.NewBuiltin(ruleRun, ruleFuncRun) + starlark.Universe[ruleGitConfig] = starlark.NewBuiltin(ruleGitConfig, ruleFuncGitConfig) + starlark.Universe[ruleInclude] = starlark.NewBuiltin(ruleInclude, ruleFuncInclude) +} + +func RegisterBuildContext(buildContextDir string) { + starlark.Universe[builtin.BuildContextDir] = starlark.String(buildContextDir) +} + +func ruleFuncBase(thread *starlark.Thread, _ *starlark.Builtin, + args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { + var os, language, image string + + if err := starlark.UnpackArgs(ruleBase, args, kwargs, + "os?", &os, "language?", &language, "image?", &image); err != nil { + return nil, err + } + + logger.Debugf("rule `%s` is invoked, os=%s, language=%s, image=%s\n", + ruleBase, os, language, image) + + err := ir.Base(os, language, image) + return starlark.None, err +} + +func ruleFuncRun(thread *starlark.Thread, _ *starlark.Builtin, + args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { + var commands *starlark.List + mountHost := false + + if err := starlark.UnpackArgs(ruleRun, + args, kwargs, "commands", &commands, "mount_host?", &mountHost); err != nil { + return nil, err + } + + goCommands, err := starlarkutil.ToStringSlice(commands) + if err != nil { + return nil, err + } + + logger.Debugf("rule `%s` is invoked, commands=%v, mount_host=%t", ruleRun, goCommands, mountHost) + if err := ir.Run(goCommands, mountHost); err != nil { + return nil, err + } + + return starlark.None, nil +} + +func ruleFuncShell(thread *starlark.Thread, _ *starlark.Builtin, + args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { + var shell starlark.String + + if err := starlark.UnpackPositionalArgs(ruleShell, args, kwargs, 1, &shell); err != nil { + return nil, err + } + + shellStr := shell.GoString() + + logger.Debugf("rule `%s` is invoked, shell=%s", ruleShell, shellStr) + + err := ir.Shell(shellStr) + return starlark.None, err +} + +func ruleFuncGitConfig(thread *starlark.Thread, _ *starlark.Builtin, + args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { + var name, email, editor starlark.String + + if err := starlark.UnpackArgs(ruleGitConfig, + args, kwargs, "name?", &name, "email?", &email, "editor?", &editor); err != nil { + return nil, err + } + + nameStr := name.GoString() + emailStr := email.GoString() + editorStr := editor.GoString() + + logger.Debugf("rule `%s` is invoked, name=%s, email=%s, editor=%s", + ruleGitConfig, nameStr, emailStr, editorStr) + + err := ir.Git(nameStr, emailStr, editorStr) + return starlark.None, err +} + +func ruleFuncInclude(thread *starlark.Thread, _ *starlark.Builtin, + args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { + var gitRepo string + + if err := starlark.UnpackArgs(ruleInclude, + args, kwargs, "git?", &gitRepo); err != nil { + return nil, err + } + + logger.Debugf("rule `%s` is invoked, git=%s", ruleInclude, gitRepo) + + globals, err := thread.Load(thread, fmt.Sprintf("%s%s", GitPrefix, gitRepo)) + if err != nil { + return nil, err + } + module := &starlarkstruct.Module{ + Name: gitRepo, + Members: globals, + } + return module, nil +} diff --git a/pkg/lang/frontend/starlark/v1/config/config.go b/pkg/lang/frontend/starlark/v1/config/config.go index f1d27c2be..67029475b 100644 --- a/pkg/lang/frontend/starlark/v1/config/config.go +++ b/pkg/lang/frontend/starlark/v1/config/config.go @@ -180,16 +180,15 @@ func ruleFuncRStudioServer(thread *starlark.Thread, _ *starlark.Builtin, func ruleFuncCondaChannel(thread *starlark.Thread, _ *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { var channel string - var useMamba bool if err := starlark.UnpackArgs(ruleCondaChannel, args, kwargs, - "channel?", &channel, "use_mamba?", &useMamba); err != nil { + "channel?", &channel); err != nil { return nil, err } - logger.Debugf("rule `%s` is invoked, channel=%s, use_mamba=%t\n", - ruleCondaChannel, channel, useMamba) - if err := ir.CondaChannel(channel, useMamba); err != nil { + logger.Debugf("rule `%s` is invoked, channel=%s\n", + ruleCondaChannel, channel) + if err := ir.CondaChannel(channel); err != nil { return nil, err } diff --git a/pkg/lang/frontend/starlark/v1/install/const.go b/pkg/lang/frontend/starlark/v1/install/const.go index 3f39a8ff8..2f79810f1 100644 --- a/pkg/lang/frontend/starlark/v1/install/const.go +++ b/pkg/lang/frontend/starlark/v1/install/const.go @@ -15,11 +15,20 @@ package install const ( + // language + rulePython = "install.python" + ruleConda = "install.conda" + ruleRLang = "install.r_lang" + ruleJulia = "install.julia" + + // packages ruleSystemPackage = "install.apt_packages" rulePyPIPackage = "install.python_packages" + ruleCondaPackages = "install.conda_packages" ruleRPackage = "install.r_packages" - ruleCUDA = "install.cuda" - ruleVSCode = "install.vscode_extensions" - ruleConda = "install.conda_packages" - ruleJulia = "install.julia_packages" + ruleJuliaPackages = "install.julia_packages" + + // others + ruleCUDA = "install.cuda" + ruleVSCode = "install.vscode_extensions" ) diff --git a/pkg/lang/frontend/starlark/v1/install/install.go b/pkg/lang/frontend/starlark/v1/install/install.go index d82a01390..2d616e425 100644 --- a/pkg/lang/frontend/starlark/v1/install/install.go +++ b/pkg/lang/frontend/starlark/v1/install/install.go @@ -31,16 +31,66 @@ var ( var Module = &starlarkstruct.Module{ Name: "install", Members: starlark.StringDict{ - "python_packages": starlark.NewBuiltin(rulePyPIPackage, ruleFuncPyPIPackage), - "r_packages": starlark.NewBuiltin(ruleRPackage, ruleFuncRPackage), - "apt_packages": starlark.NewBuiltin(ruleSystemPackage, ruleFuncSystemPackage), + // language + "python": starlark.NewBuiltin(rulePython, ruleFuncPython), + "conda": starlark.NewBuiltin(ruleConda, ruleFuncConda), + "r_lang": starlark.NewBuiltin(ruleRLang, ruleFuncRLang), + "julia": starlark.NewBuiltin(ruleJulia, ruleFuncJulia), + // packages + "apt_packages": starlark.NewBuiltin(ruleSystemPackage, ruleFuncSystemPackage), + "python_packages": starlark.NewBuiltin(rulePyPIPackage, ruleFuncPyPIPackage), + "conda_packages": starlark.NewBuiltin(ruleCondaPackages, ruleFuncCondaPackage), + "r_packages": starlark.NewBuiltin(ruleRPackage, ruleFuncRPackage), + "julia_packages": starlark.NewBuiltin(ruleJuliaPackages, ruleFuncJuliaPackage), + // others "cuda": starlark.NewBuiltin(ruleCUDA, ruleFuncCUDA), "vscode_extensions": starlark.NewBuiltin(ruleVSCode, ruleFuncVSCode), - "conda_packages": starlark.NewBuiltin(ruleConda, ruleFuncConda), - "julia_packages": starlark.NewBuiltin(ruleJulia, ruleFuncJulia), }, } +func ruleFuncPython(thread *starlark.Thread, _ *starlark.Builtin, + args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { + version := ir.PythonVersionDefault + + if err := starlark.UnpackArgs(rulePython, args, kwargs, "version?", &version); err != nil { + return nil, err + } + + logger.Debugf("rule `%s` is invoked, version=%s", rulePython, version) + if err := ir.Python(version); err != nil { + return nil, err + } + + return starlark.None, nil +} + +func ruleFuncConda(thread *starlark.Thread, _ *starlark.Builtin, + args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { + useMamba := false + + if err := starlark.UnpackArgs(rulePython, args, kwargs, "use_mamba?", &useMamba); err != nil { + return nil, err + } + + logger.Debugf("rule `%s` is invoked: use_mamba=%t", ruleConda, useMamba) + ir.Conda(useMamba) + return starlark.None, nil +} + +func ruleFuncRLang(thread *starlark.Thread, _ *starlark.Builtin, + args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { + logger.Debugf("rule `%s` is invoked", ruleRLang) + ir.RLang() + return starlark.None, nil +} + +func ruleFuncJulia(thread *starlark.Thread, _ *starlark.Builtin, + args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { + logger.Debugf("rule `%s` is invoked", ruleJulia) + ir.Julia() + return starlark.None, nil +} + func ruleFuncPyPIPackage(thread *starlark.Thread, _ *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { var name *starlark.List @@ -91,11 +141,11 @@ func ruleFuncRPackage(thread *starlark.Thread, _ *starlark.Builtin, return starlark.None, nil } -func ruleFuncJulia(thread *starlark.Thread, _ *starlark.Builtin, +func ruleFuncJuliaPackage(thread *starlark.Thread, _ *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { var name *starlark.List - if err := starlark.UnpackArgs(ruleJulia, + if err := starlark.UnpackArgs(ruleJuliaPackages, args, kwargs, "name", &name); err != nil { return nil, err } @@ -104,7 +154,7 @@ func ruleFuncJulia(thread *starlark.Thread, _ *starlark.Builtin, if err != nil { return nil, err } - logger.Debugf("rule `%s` is invoked, name=%v", ruleJulia, nameList) + logger.Debugf("rule `%s` is invoked, name=%v", ruleJuliaPackages, nameList) ir.JuliaPackage(nameList) return starlark.None, nil @@ -168,12 +218,12 @@ func ruleFuncVSCode(thread *starlark.Thread, _ *starlark.Builtin, return starlark.None, nil } -func ruleFuncConda(thread *starlark.Thread, _ *starlark.Builtin, +func ruleFuncCondaPackage(thread *starlark.Thread, _ *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { var name, channel *starlark.List var envFile starlark.String - if err := starlark.UnpackArgs(ruleConda, + if err := starlark.UnpackArgs(ruleCondaPackages, args, kwargs, "name?", &name, "channel?", &channel, "env_file?", &envFile); err != nil { return nil, err } @@ -195,7 +245,7 @@ func ruleFuncConda(thread *starlark.Thread, _ *starlark.Builtin, } } - logger.Debugf("rule `%s` is invoked, name=%v, channel=%v, env_file=%s", ruleConda, nameList, channelList, envFileStr) + logger.Debugf("rule `%s` is invoked, name=%v, channel=%v, env_file=%s", ruleCondaPackages, nameList, channelList, envFileStr) if err := ir.CondaPackage(nameList, channelList, envFileStr); err != nil { return starlark.None, err } diff --git a/pkg/lang/frontend/starlark/v1/interpreter.go b/pkg/lang/frontend/starlark/v1/interpreter.go index 119c26099..ca403edd8 100644 --- a/pkg/lang/frontend/starlark/v1/interpreter.go +++ b/pkg/lang/frontend/starlark/v1/interpreter.go @@ -27,6 +27,7 @@ import ( "github.com/sirupsen/logrus" "go.starlark.net/starlark" + interp "github.com/tensorchord/envd/pkg/lang/frontend/starlark" "github.com/tensorchord/envd/pkg/lang/frontend/starlark/v1/config" "github.com/tensorchord/envd/pkg/lang/frontend/starlark/v1/data" "github.com/tensorchord/envd/pkg/lang/frontend/starlark/v1/install" @@ -36,11 +37,6 @@ import ( "github.com/tensorchord/envd/pkg/util/fileutil" ) -type Interpreter interface { - Eval(script string) (interface{}, error) - ExecFile(filename string, funcname string) (interface{}, error) -} - type entry struct { globals starlark.StringDict err error @@ -54,7 +50,7 @@ type generalInterpreter struct { cache map[string]*entry } -func NewInterpreter(buildContextDir string) Interpreter { +func NewInterpreter(buildContextDir string) interp.Interpreter { // Register envd rules and built-in variables to Starlark. universe.RegisterEnvdRules() universe.RegisterBuildContext(buildContextDir) diff --git a/pkg/lang/frontend/starlark/v1/universe/universe.go b/pkg/lang/frontend/starlark/v1/universe/universe.go index de8d94e84..a03eb079e 100644 --- a/pkg/lang/frontend/starlark/v1/universe/universe.go +++ b/pkg/lang/frontend/starlark/v1/universe/universe.go @@ -45,17 +45,16 @@ func RegisterBuildContext(buildContextDir string) { func ruleFuncBase(thread *starlark.Thread, _ *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { - var os, language, image string + var image string + var dev bool - if err := starlark.UnpackArgs(ruleBase, args, kwargs, - "os?", &os, "language?", &language, "image?", &image); err != nil { + if err := starlark.UnpackArgs(ruleBase, args, kwargs, "image?", &image, "dev?", &dev); err != nil { return nil, err } - logger.Debugf("rule `%s` is invoked, os=%s, language=%s, image=%s\n", - ruleBase, os, language, image) + logger.Debugf("rule `%s` is invoked, image=%s, dev=%t\n", ruleBase, image, dev) - err := ir.Base(os, language, image) + err := ir.Base(image, dev) return starlark.None, err } diff --git a/pkg/lang/ir/graph.go b/pkg/lang/ir/graph.go index a37293cb6..4863168c3 100644 --- a/pkg/lang/ir/graph.go +++ b/pkg/lang/ir/graph.go @@ -34,4 +34,5 @@ type graphVisitor interface { GetEnviron() []string GetHTTP() []HTTPInfo GetRuntimeCommands() map[string]string + GetUser() string } diff --git a/pkg/lang/ir/v0/cache.go b/pkg/lang/ir/v0/cache.go new file mode 100644 index 000000000..1364a6193 --- /dev/null +++ b/pkg/lang/ir/v0/cache.go @@ -0,0 +1,32 @@ +// 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 v1 + +import ( + "fmt" + + "github.com/sirupsen/logrus" +) + +func (g generalGraph) CacheID(filename string) string { + var cacheID string + if g.CUDA != nil { + cacheID = fmt.Sprintf("%s/%s-gpu", filename, g.EnvironmentName) + } else { + cacheID = fmt.Sprintf("%s/%s-cpu", filename, g.EnvironmentName) + } + logrus.Debugf("apt/pypi calculated cacheID: %s", cacheID) + return cacheID +} diff --git a/pkg/lang/ir/v0/checker.go b/pkg/lang/ir/v0/checker.go new file mode 100644 index 000000000..dbc4e12d3 --- /dev/null +++ b/pkg/lang/ir/v0/checker.go @@ -0,0 +1,66 @@ +package v1 + +import ( + "reflect" + "regexp" + "strings" +) + +func GetDepsFiles(deps []string) []string { + rawGraph := DefaultGraph.(*generalGraph) + tHandle := reflect.TypeOf(*rawGraph) + vHandle := reflect.ValueOf(*rawGraph) + deps = searchFileInGraph(tHandle, vHandle, deps) + return deps +} + +// Match all filed in ir.Graph with the given keyword +func likeFileFiled(str string) bool { + nameKeyword := []string{ + "File", + "Path", + "Wheels", + } + if len(nameKeyword) == 0 { + return true + } + re := regexp.MustCompile(strings.Join(nameKeyword, "|")) + return re.MatchString(str) +} + +// search all files in Graph +func searchFileInGraph(tHandle reflect.Type, vHandle reflect.Value, deps []string) []string { + for i := 0; i < vHandle.NumField(); i++ { + v := vHandle.Field(i) + if v.Type().Kind() == reflect.Struct { + t := v.Type() + deps = searchFileInGraph(t, v, deps) + } else if v.Type().Kind() == reflect.Ptr { + if v.Type().Elem().Kind() == reflect.Struct { + if v.Elem().CanAddr() { + t := v.Type().Elem() + deps = searchFileInGraph(t, v.Elem(), deps) + } + } + } else { + t := tHandle.Field(i) + fieldName := t.Name + if likeFileFiled(fieldName) { + typeName := t.Type.String() + if v.Interface() != nil { + if typeName == "string" { + deps = append(deps, v.Interface().(string)) + } + if typeName == "*string" { + deps = append(deps, *(v.Interface().(*string))) + } + if typeName == "[]string" { + filesList := v.Interface().([]string) + deps = append(deps, filesList...) + } + } + } + } + } + return deps +} diff --git a/pkg/lang/ir/v0/compile.go b/pkg/lang/ir/v0/compile.go new file mode 100644 index 000000000..de52b3290 --- /dev/null +++ b/pkg/lang/ir/v0/compile.go @@ -0,0 +1,341 @@ +// 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 v1 + +import ( + "context" + "encoding/json" + "fmt" + "os" + "path/filepath" + "strings" + + "github.com/cockroachdb/errors" + "github.com/moby/buildkit/client/llb" + "github.com/sirupsen/logrus" + "github.com/spf13/viper" + servertypes "github.com/tensorchord/envd-server/api/types" + + "github.com/tensorchord/envd/pkg/config" + "github.com/tensorchord/envd/pkg/flag" + "github.com/tensorchord/envd/pkg/lang/ir" + "github.com/tensorchord/envd/pkg/progress/compileui" + "github.com/tensorchord/envd/pkg/types" + "github.com/tensorchord/envd/pkg/util/fileutil" + "github.com/tensorchord/envd/pkg/version" +) + +func NewGraph() ir.Graph { + runtimeGraph := ir.RuntimeGraph{ + RuntimeCommands: make(map[string]string), + RuntimeEnviron: make(map[string]string), + } + langVersion := languageVersionDefault + conda := &ir.CondaConfig{} + return &generalGraph{ + OS: osDefault, + Language: ir.Language{ + Name: languageDefault, + Version: &langVersion, + }, + CUDA: nil, + CUDNN: CUDNNVersionDefault, + NumGPUs: 0, + + PyPIPackages: []string{}, + RPackages: []string{}, + JuliaPackages: []string{}, + SystemPackages: []string{}, + Exec: []ir.RunBuildCommand{}, + UserDirectories: []string{}, + RuntimeEnvPaths: []string{types.DefaultPathEnv()}, + Shell: shellBASH, + CondaConfig: conda, + RuntimeGraph: runtimeGraph, + } +} + +var DefaultGraph = NewGraph() + +func (g generalGraph) GetNumGPUs() int { + return g.NumGPUs +} + +func (g generalGraph) GetShell() string { + return g.Shell +} + +func (g generalGraph) GetMount() []ir.MountInfo { + return g.Mount +} + +func (g generalGraph) GetEnvironmentName() string { + return g.EnvironmentName +} + +func (g generalGraph) GetJupyterConfig() *ir.JupyterConfig { + return g.JupyterConfig +} + +func (g generalGraph) GetRStudioServerConfig() *ir.RStudioServerConfig { + return g.RStudioServerConfig +} + +func (g generalGraph) GetExposedPorts() []ir.ExposeItem { + return g.RuntimeExpose +} + +func (g generalGraph) GetRuntimeCommands() map[string]string { + return g.RuntimeCommands +} + +func (g generalGraph) GetUser() string { + return "envd" +} + +func (g *generalGraph) Compile(ctx context.Context, envName string, pub string) (*llb.Definition, error) { + w, err := compileui.New(ctx, os.Stdout, "auto") + if err != nil { + return nil, errors.Wrap(err, "failed to create compileui") + } + g.Writer = w + g.EnvironmentName = envName + g.PublicKeyPath = pub + + uid, gid, err := getUIDGID() + if err != nil { + return nil, errors.Wrap(err, "failed to get uid/gid") + } + state, err := g.CompileLLB(uid, gid) + if err != nil { + return nil, errors.Wrap(err, "failed to compile the graph") + } + // TODO(gaocegege): Support multi platform. + def, err := state.Marshal(ctx, llb.LinuxAmd64) + if err != nil { + return nil, errors.Wrap(err, "failed to marshal the llb definition") + } + return def, nil +} + +func (g generalGraph) GPUEnabled() bool { + return g.CUDA != nil +} + +func (g generalGraph) Labels() (map[string]string, error) { + labels := make(map[string]string) + str, err := json.Marshal(g.SystemPackages) + if err != nil { + return nil, err + } + labels[types.ImageLabelAPT] = string(str) + str, err = json.Marshal(g.PyPIPackages) + if err != nil { + 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 + labels[types.ImageLabelCUDNN] = g.CUDNN + } + labels[types.ImageLabelVendor] = types.ImageVendorEnvd + code, err := g.RuntimeGraph.Dump() + if err != nil { + return labels, err + } + labels[types.RuntimeGraphCode] = code + + ports := []servertypes.EnvironmentPort{} + ports = append(ports, servertypes.EnvironmentPort{ + Name: "ssh", + Port: config.SSHPortInContainer, + }) + if g.JupyterConfig != nil { + ports = append(ports, servertypes.EnvironmentPort{ + Name: "jupyter", + Port: config.JupyterPortInContainer, + }) + } + if g.RStudioServerConfig != nil { + ports = append(ports, servertypes.EnvironmentPort{ + Name: "rstudio-server", + Port: config.RStudioServerPortInContainer, + }) + } + + if g.RuntimeExpose != nil && len(g.RuntimeExpose) > 0 { + for _, item := range g.RuntimeExpose { + ports = append(ports, servertypes.EnvironmentPort{ + Name: item.ServiceName, + Port: int32(item.EnvdPort), + }) + } + } + + portsData, err := json.Marshal(ports) + if err != nil { + return labels, err + } + labels[types.ImageLabelPorts] = string(portsData) + + repoInfo, err := json.Marshal(g.Repo) + if err != nil { + return labels, err + } + labels[types.ImageLabelRepo] = string(repoInfo) + + labels[types.ImageLabelContainerName] = string(g.EnvironmentName) + return labels, nil +} + +func (g generalGraph) ExposedPorts() (map[string]struct{}, error) { + ports := make(map[string]struct{}) + + // do not expose ports for custom images + if g.Image != nil { + return ports, nil + } + + ports[fmt.Sprintf("%d/tcp", config.SSHPortInContainer)] = struct{}{} + if g.JupyterConfig != nil { + ports[fmt.Sprintf("%d/tcp", config.JupyterPortInContainer)] = struct{}{} + } + if g.RStudioServerConfig != nil { + ports[fmt.Sprintf("%d/tcp", config.RStudioServerPortInContainer)] = struct{}{} + } + + if g.RuntimeExpose != nil && len(g.RuntimeExpose) > 0 { + for _, item := range g.RuntimeExpose { + ports[fmt.Sprintf("%d/tcp", item.EnvdPort)] = struct{}{} + } + } + + return ports, nil +} + +func (g generalGraph) EnvString() []string { + var envs []string + for k, v := range g.RuntimeEnviron { + envs = append(envs, fmt.Sprintf("%s=%s", k, v)) + } + return envs +} + +func (g generalGraph) GetEnviron() []string { + if g.Image != nil { + return g.EnvString() + } + // Add PATH and LC_ALL. + return append(g.EnvString(), + "PATH="+strings.Join(g.RuntimeEnvPaths, ":"), + "LC_ALL=en_US.UTF-8", + ) +} + +func (g *generalGraph) SetWriter(w compileui.Writer) { + g.Writer = w +} + +func (g generalGraph) GetHTTP() []ir.HTTPInfo { + return g.HTTP +} + +func (g generalGraph) DefaultCacheImporter() (*string, error) { + // The base remote cache should work for all languages. + var res string + if g.CUDA != nil { + res = fmt.Sprintf( + "type=registry,ref=docker.io/%s/python-cache:envd-%s-cuda-%s-cudnn-%s", + viper.GetString(flag.FlagDockerOrganization), + version.GetVersionForImageTag(), *g.CUDA, g.CUDNN) + } else { + res = fmt.Sprintf( + "type=registry,ref=docker.io/%s/python-cache:envd-%s", + viper.GetString(flag.FlagDockerOrganization), + version.GetVersionForImageTag()) + } + return &res, nil +} + +func (g *generalGraph) GetEntrypoint(buildContextDir string) ([]string, error) { + if g.Image != nil { + return g.Entrypoint, nil + } + g.RuntimeEnviron[types.EnvdWorkDir] = fileutil.EnvdHomeDir(filepath.Base(buildContextDir)) + return []string{"horust"}, nil +} + +func (g generalGraph) CompileLLB(uid, gid int) (llb.State, error) { + g.uid = uid + + // TODO(gaocegege): Remove the hack for https://github.com/tensorchord/envd/issues/370 + g.gid = 1001 + logrus.WithFields(logrus.Fields{ + "uid": g.uid, + "gid": g.gid, + }).Debug("compile LLB") + + // TODO(gaocegege): Support more OS and langs. + aptStage, err := g.compileBase() + if err != nil { + return llb.State{}, errors.Wrap(err, "failed to get the base image") + } + var merged llb.State + // Use custom logic when image is specified. + if g.Image != nil { + merged, err = g.compileCustomPython(aptStage) + if err != nil { + return llb.State{}, errors.Wrap(err, "failed to compile custom python image") + } + } else { + switch g.Language.Name { + case "r": + merged, err = g.compileRLang(aptStage) + if err != nil { + return llb.State{}, errors.Wrap(err, "failed to compile r language") + } + case "python": + merged, err = g.compilePython(aptStage) + if err != nil { + return llb.State{}, errors.Wrap(err, "failed to compile python") + } + case "julia": + merged, err = g.compileJulia(aptStage) + if err != nil { + return llb.State{}, errors.Wrap(err, "failed to compile julia") + } + } + } + + prompt := g.compilePrompt(merged) + copy := g.compileCopy(prompt) + // TODO(gaocegege): Support order-based exec. + run := g.compileRun(copy) + git := g.compileGit(run) + user := g.compileUserOwn(git) + mount := g.compileMountDir(user) + entrypoint, err := g.compileEntrypoint(mount) + if err != nil { + return llb.State{}, errors.Wrap(err, "failed to compile entrypoint") + } + g.Writer.Finish() + return entrypoint, nil +} diff --git a/pkg/lang/ir/v0/conda.go b/pkg/lang/ir/v0/conda.go new file mode 100644 index 000000000..3f2d6c8b6 --- /dev/null +++ b/pkg/lang/ir/v0/conda.go @@ -0,0 +1,158 @@ +// 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 v1 + +import ( + _ "embed" + "fmt" + "path/filepath" + "strings" + + "github.com/cockroachdb/errors" + "github.com/moby/buildkit/client/llb" + "github.com/sirupsen/logrus" + + "github.com/tensorchord/envd/pkg/flag" +) + +const ( + condaVersionDefault = "py39_4.11.0" + // check the issue https://github.com/mamba-org/mamba/issues/1975 + mambaVersionDefault = "0.25.1" + condaRootPrefix = "/opt/conda" + condaBinDir = "/opt/conda/bin" +) + +var ( + // this file can be used by both conda and mamba + // https://mamba.readthedocs.io/en/latest/user_guide/configuration.html#multiple-rc-files + condarc = "/opt/conda/.condarc" + //go:embed install-conda.sh + installCondaBash string + //go:embed install-mamba.sh + installMambaBash string +) + +func (g generalGraph) compileCondaChannel(root llb.State) llb.State { + if g.CondaConfig.CondaChannel != nil { + logrus.WithField("conda-channel", *g.CondaChannel).Debug("using custom conda channel") + stage := root. + File(llb.Mkfile(condarc, + 0644, []byte(*g.CondaChannel), llb.WithUIDGID(g.uid, g.gid)), llb.WithCustomName("[internal] setting conda channel")) + return stage + } + return root +} + +func (g generalGraph) condaCommandPath() string { + if g.CondaConfig.UseMicroMamba { + return filepath.Join(condaBinDir, "micromamba") + } + return filepath.Join(condaBinDir, "conda") +} + +func (g generalGraph) condaInitShell(shell string) string { + path := g.condaCommandPath() + if g.CondaConfig.UseMicroMamba { + return fmt.Sprintf("%s shell init -p %s -s %s", path, condaRootPrefix, shell) + } + return fmt.Sprintf("%s init %s", path, shell) +} + +func (g generalGraph) condaUpdateFromFile() string { + args := fmt.Sprintf("update -n envd --file %s", g.CondaEnvFileName) + if g.CondaConfig.UseMicroMamba { + return fmt.Sprintf("%s %s", g.condaCommandPath(), args) + } + return fmt.Sprintf("%s env %s", g.condaCommandPath(), args) +} + +func (g *generalGraph) compileCondaPackages(root llb.State) llb.State { + if len(g.CondaConfig.CondaPackages) == 0 && len(g.CondaEnvFileName) == 0 { + logrus.Debug("Conda packages not enabled") + return root + } + + cacheDir := filepath.Join(condaRootPrefix, "pkgs") + // Refer to https://github.com/moby/buildkit/blob/31054718bf775bf32d1376fe1f3611985f837584/frontend/dockerfile/dockerfile2llb/convert_runmount.go#L46 + cacheMount := root.File(llb.Mkdir("/cache-conda", 0755, llb.WithParents(true)), + llb.WithCustomName("[internal] setting conda cache mount permissions")) + + // Compose the package install command. + var sb strings.Builder + var run llb.ExecState + + if len(g.CondaEnvFileName) > 0 { + sb.WriteString(g.condaUpdateFromFile()) + } else { + if len(g.CondaConfig.AdditionalChannels) == 0 { + sb.WriteString(fmt.Sprintf("%s install -n envd", g.condaCommandPath())) + } else { + sb.WriteString(fmt.Sprintf("%s install -n envd", g.condaCommandPath())) + for _, channel := range g.CondaConfig.AdditionalChannels { + sb.WriteString(fmt.Sprintf(" -c %s", channel)) + } + } + for _, pkg := range g.CondaConfig.CondaPackages { + sb.WriteString(fmt.Sprintf(" %s", pkg)) + } + } + + cmd := sb.String() + run = root.Dir(g.getWorkingDir()). + Run(llb.Shlex(cmd), llb.WithCustomNamef("[internal] %s %s", + cmd, strings.Join(g.CondaPackages, " "))) + run.AddMount(g.getWorkingDir(), llb.Local(flag.FlagBuildContext)) + run.AddMount(cacheDir, cacheMount, + llb.AsPersistentCacheDir(g.CacheID(cacheDir), llb.CacheMountShared), llb.SourcePath("/cache-conda")) + return run.Root() +} + +func (g generalGraph) compileCondaEnvironment(root llb.State) (llb.State, error) { + // Always init bash since we will use it to create jupyter notebook service. + run := root.Run( + llb.Shlex(fmt.Sprintf("bash -c \"%s\"", g.condaInitShell("bash"))), + llb.WithCustomName("[internal] initialize conda bash environment"), + ) + + pythonVersion, err := g.getAppropriatePythonVersion() + if err != nil { + return llb.State{}, errors.Wrap(err, "failed to get python version") + } + // Create a conda environment. + cmd := fmt.Sprintf("bash -c \"%s create -n envd python=%s\"", g.condaCommandPath(), pythonVersion) + run = run.Run(llb.Shlex(cmd), + llb.WithCustomNamef("[internal] create conda environment: %s", cmd)) + + return run.Root(), nil +} + +// nolint:unparam +func (g generalGraph) installConda(root llb.State) (llb.State, error) { + if g.CondaConfig.UseMicroMamba { + run := root.AddEnv("MAMBA_BIN_DIR", condaBinDir). + AddEnv("MAMBA_ROOT_PREFIX", condaRootPrefix). + AddEnv("MAMBA_VERSION", mambaVersionDefault). + Run(llb.Shlex(fmt.Sprintf("bash -c '%s'", installMambaBash)), + llb.WithCustomName("[internal] install micro mamba")) + return run.Root(), nil + } + run := root.AddEnv("CONDA_VERSION", condaVersionDefault). + File(llb.Mkdir(condaRootPrefix, 0755, llb.WithParents(true)), + llb.WithCustomName("[internal] create conda directory")). + Run(llb.Shlex(fmt.Sprintf("bash -c '%s'", installCondaBash)), + llb.WithCustomName("[internal] install conda")) + return run.Root(), nil +} diff --git a/pkg/lang/ir/v0/consts.go b/pkg/lang/ir/v0/consts.go new file mode 100644 index 000000000..9f154ea66 --- /dev/null +++ b/pkg/lang/ir/v0/consts.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 v1 + +import "github.com/tensorchord/envd/pkg/util/fileutil" + +const ( + osDefault = "ubuntu20.04" + languageDefault = "python" + languageVersionDefault = "3" + CUDNNVersionDefault = "8" + + aptSourceFilePath = "/etc/apt/sources.list" + pypiIndexFilePath = "/etc/pip.conf" + + pypiConfigTemplate = ` +[global] +index-url=%s +%s + +[install] +src = /tmp +` +) + +var ( + // used inside the container + defaultConfigDir = fileutil.EnvdHomeDir(".config") + starshipConfigPath = fileutil.EnvdHomeDir(".config", "starship.toml") +) diff --git a/pkg/lang/ir/v1/custom.go b/pkg/lang/ir/v0/custom.go similarity index 100% rename from pkg/lang/ir/v1/custom.go rename to pkg/lang/ir/v0/custom.go diff --git a/pkg/lang/ir/v0/editor.go b/pkg/lang/ir/v0/editor.go new file mode 100644 index 000000000..9e95ad688 --- /dev/null +++ b/pkg/lang/ir/v0/editor.go @@ -0,0 +1,120 @@ +// 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 v1 + +import ( + "fmt" + "strconv" + + "github.com/cockroachdb/errors" + "github.com/moby/buildkit/client/llb" + + "github.com/tensorchord/envd/pkg/config" + "github.com/tensorchord/envd/pkg/editor/vscode" + "github.com/tensorchord/envd/pkg/flag" + "github.com/tensorchord/envd/pkg/progress/compileui" + "github.com/tensorchord/envd/pkg/types" + "github.com/tensorchord/envd/pkg/util/fileutil" +) + +func (g generalGraph) compileVSCode() (*llb.State, error) { + if len(g.VSCodePlugins) == 0 { + return nil, nil + } + inputs := []llb.State{} + for _, p := range g.VSCodePlugins { + vscodeClient, err := vscode.NewClient(vscode.MarketplaceVendorOpenVSX) + if err != nil { + return nil, errors.Wrap(err, "failed to create vscode client") + } + g.Writer.LogVSCodePlugin(p, compileui.ActionStart, false) + if cached, err := vscodeClient.DownloadOrCache(p); err != nil { + return nil, err + } else { + g.Writer.LogVSCodePlugin(p, compileui.ActionEnd, cached) + } + ext := llb.Scratch().File(llb.Copy(llb.Local(flag.FlagCacheDir), + vscodeClient.PluginPath(p), + fileutil.EnvdHomeDir(".vscode-server", "extensions", p.String()), + &llb.CopyInfo{ + CreateDestPath: true, + }, llb.WithUIDGID(g.uid, g.gid)), + llb.WithCustomNamef("install vscode plugin %s", p.String())) + inputs = append(inputs, ext) + } + layer := llb.Merge(inputs, llb.WithCustomName("merging plugins for vscode")) + return &layer, nil +} + +func (g *generalGraph) compileJupyter() error { + if g.JupyterConfig == nil { + return nil + } + + g.PyPIPackages = append(g.PyPIPackages, "jupyter") + switch g.Language.Name { + case "python": + return nil + default: + return errors.Newf("Jupyter is not supported in %s yet", g.Language.Name) + } +} + +func (g generalGraph) generateJupyterCommand(workingDir string) []string { + if g.JupyterConfig == nil { + return nil + } + + if g.JupyterConfig.Token == "" { + g.JupyterConfig.Token = "''" + } + + // get from env if not set + if len(workingDir) == 0 { + workingDir = fmt.Sprintf("${%s}", types.EnvdWorkDir) + } + + cmd := []string{ + "python3", "-m", "notebook", + "--ip", "0.0.0.0", "--notebook-dir", workingDir, + "--NotebookApp.token", g.JupyterConfig.Token, + "--port", strconv.Itoa(config.JupyterPortInContainer), + } + + if g.uid == 0 { + cmd = append(cmd, "--allow-root") + } + + return cmd +} + +// nolint:unparam +func (g generalGraph) generateRStudioCommand(workingDir string) []string { + if g.RStudioServerConfig == nil { + return nil + } + + // get from env if not set + // if len(workingDir) == 0 { + // workingDir = fmt.Sprintf("${%s}", types.EnvdWorkDir) + // } + + return []string{ + // TODO(gaocegege): Remove root permission here. + "sudo", + "/usr/lib/rstudio-server/bin/rserver", + // TODO(gaocegege): Support working dir. + } +} diff --git a/pkg/lang/ir/v0/editor_test.go b/pkg/lang/ir/v0/editor_test.go new file mode 100644 index 000000000..dfb41cdc1 --- /dev/null +++ b/pkg/lang/ir/v0/editor_test.go @@ -0,0 +1,97 @@ +// 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 v1 + +import ( + "testing" + + "github.com/tensorchord/envd/pkg/lang/ir" +) + +func TestGenerateCommand(t *testing.T) { + testcases := []struct { + graph generalGraph + dir string + expected []string + }{ + { + graph: generalGraph{ + JupyterConfig: &ir.JupyterConfig{ + Token: "", + Port: 8888, + }, + uid: 1000, + }, + dir: "test", + expected: []string{ + "python3", "-m", "notebook", "--ip", "0.0.0.0", "--notebook-dir", "test", + "--NotebookApp.token", "''", "--port", "8888", + }, + }, + { + graph: generalGraph{ + JupyterConfig: &ir.JupyterConfig{ + Token: "test", + Port: 8888, + }, + uid: 1000, + }, + dir: "test", + expected: []string{ + "python3", "-m", "notebook", "--ip", "0.0.0.0", "--notebook-dir", "test", + "--NotebookApp.token", "test", "--port", "8888", + }, + }, + { + graph: generalGraph{ + JupyterConfig: &ir.JupyterConfig{ + Token: "test", + Port: 8888, + }, + uid: 0, + }, + dir: "test", + expected: []string{ + "python3", "-m", "notebook", "--ip", "0.0.0.0", "--notebook-dir", "test", + "--NotebookApp.token", "test", "--port", "8888", "--allow-root", + }, + }, + { + graph: generalGraph{}, + dir: "test", + expected: []string{}, + }, + } + for _, tc := range testcases { + actual := tc.graph.generateJupyterCommand(tc.dir) + if !equal(actual, tc.expected) { + t.Errorf("failed to generate the command: expected %v, got %v", tc.expected, actual) + } + } +} + +// Equal tells whether a and b contain the same elements. +// A nil argument is equivalent to an empty slice. +func equal(a, b []string) bool { + if len(a) != len(b) { + return false + } + for i, v := range a { + if v != b[i] { + return false + } + } + return true +} diff --git a/pkg/lang/ir/v0/fs.go b/pkg/lang/ir/v0/fs.go new file mode 100644 index 000000000..2a16ddf92 --- /dev/null +++ b/pkg/lang/ir/v0/fs.go @@ -0,0 +1,25 @@ +// 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 v1 + +import ( + "github.com/moby/buildkit/client/llb" +) + +func (g *generalGraph) CompileCacheDir(root llb.State, cacheDir string) llb.State { + g.UserDirectories = append(g.UserDirectories, cacheDir) + run := root.Run(llb.Shlexf("mkdir -p %s", cacheDir), llb.WithCustomName("[internal] create cache dir")) + return run.Root() +} diff --git a/pkg/lang/ir/v0/git.go b/pkg/lang/ir/v0/git.go new file mode 100644 index 000000000..8e456efb0 --- /dev/null +++ b/pkg/lang/ir/v0/git.go @@ -0,0 +1,45 @@ +// 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 v1 + +import ( + "fmt" + + "github.com/moby/buildkit/client/llb" + + "github.com/tensorchord/envd/pkg/util/fileutil" +) + +const ( + templateGitConfig = ` +[user] + email = %s + name = %s +[core] + editor = %s + +` +) + +func (g *generalGraph) compileGit(root llb.State) llb.State { + if g.GitConfig == nil { + return root + } + content := fmt.Sprintf(templateGitConfig, g.GitConfig.Email, g.GitConfig.Name, g.GitConfig.Editor) + installPath := fileutil.EnvdHomeDir(".gitconfig") + gitStage := root.File(llb.Mkfile(installPath, + 0644, []byte(content), llb.WithUIDGID(g.uid, g.gid))) + return gitStage +} diff --git a/pkg/lang/ir/v1/install-conda.sh b/pkg/lang/ir/v0/install-conda.sh similarity index 100% rename from pkg/lang/ir/v1/install-conda.sh rename to pkg/lang/ir/v0/install-conda.sh diff --git a/pkg/lang/ir/v1/install-mamba.sh b/pkg/lang/ir/v0/install-mamba.sh similarity index 100% rename from pkg/lang/ir/v1/install-mamba.sh rename to pkg/lang/ir/v0/install-mamba.sh diff --git a/pkg/lang/ir/v0/interface.go b/pkg/lang/ir/v0/interface.go new file mode 100644 index 000000000..1bee1b3f3 --- /dev/null +++ b/pkg/lang/ir/v0/interface.go @@ -0,0 +1,298 @@ +// 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 v1 + +import ( + "github.com/cockroachdb/errors" + "github.com/opencontainers/go-digest" + + "github.com/tensorchord/envd/pkg/editor/vscode" + "github.com/tensorchord/envd/pkg/lang/ir" + "github.com/tensorchord/envd/pkg/types" +) + +func Base(os, language, image string) error { + l, version, err := parseLanguage(language) + if err != nil { + return err + } + g := DefaultGraph.(*generalGraph) + g.Language = ir.Language{ + Name: l, + Version: version, + } + if len(os) > 0 { + g.OS = os + } + if image != "" { + g.Image = &image + } + return nil +} + +func PyPIPackage(deps []string, requirementsFile string, wheels []string) error { + g := DefaultGraph.(*generalGraph) + + g.PyPIPackages = append(g.PyPIPackages, deps...) + g.PythonWheels = append(g.PythonWheels, wheels...) + + if requirementsFile != "" { + g.RequirementsFile = &requirementsFile + } + + return nil +} + +func RPackage(deps []string) { + g := DefaultGraph.(*generalGraph) + + g.RPackages = append(g.RPackages, deps...) +} + +func JuliaPackage(deps []string) { + g := DefaultGraph.(*generalGraph) + + g.JuliaPackages = append(g.JuliaPackages, deps...) +} + +func SystemPackage(deps []string) { + g := DefaultGraph.(*generalGraph) + + g.SystemPackages = append(g.SystemPackages, deps...) +} + +func GPU(numGPUs int) { + g := DefaultGraph.(*generalGraph) + + g.NumGPUs = numGPUs +} + +func CUDA(version, cudnn string) { + g := DefaultGraph.(*generalGraph) + + g.CUDA = &version + if len(cudnn) > 0 { + g.CUDNN = cudnn + } +} + +func VSCodePlugins(plugins []string) error { + g := DefaultGraph.(*generalGraph) + + for _, p := range plugins { + plugin, err := vscode.ParsePlugin(p) + if err != nil { + return err + } + g.VSCodePlugins = append(g.VSCodePlugins, *plugin) + } + return nil +} + +// UbuntuAPT updates the Ubuntu apt source.list in the image. +func UbuntuAPT(source string) error { + if source == "" { + return errors.New("source is required") + } + g := DefaultGraph.(*generalGraph) + + g.UbuntuAPTSource = &source + return nil +} + +func PyPIIndex(url, extraURL string) error { + if url == "" { + return errors.New("url is required") + } + g := DefaultGraph.(*generalGraph) + + g.PyPIIndexURL = &url + g.PyPIExtraIndexURL = &extraURL + return nil +} + +func CRANMirror(url string) error { + g := DefaultGraph.(*generalGraph) + + g.CRANMirrorURL = &url + return nil +} + +func JuliaPackageServer(url string) error { + g := DefaultGraph.(*generalGraph) + + g.JuliaPackageServer = &url + return nil +} + +func Shell(shell string) error { + g := DefaultGraph.(*generalGraph) + + g.Shell = shell + return nil +} + +func Jupyter(pwd string, port int64) error { + g := DefaultGraph.(*generalGraph) + + g.JupyterConfig = &ir.JupyterConfig{ + Token: pwd, + Port: port, + } + return nil +} + +func RStudioServer() error { + g := DefaultGraph.(*generalGraph) + + g.RStudioServerConfig = &ir.RStudioServerConfig{} + return nil +} + +func Run(commands []string, mount bool) error { + g := DefaultGraph.(*generalGraph) + + g.Exec = append(g.Exec, ir.RunBuildCommand{ + Commands: commands, + MountHost: mount, + }) + return nil +} + +func Git(name, email, editor string) error { + g := DefaultGraph.(*generalGraph) + + g.GitConfig = &ir.GitConfig{ + Name: name, + Email: email, + Editor: editor, + } + return nil +} + +func CondaChannel(channel string, useMamba bool) error { + g := DefaultGraph.(*generalGraph) + + g.CondaConfig.CondaChannel = &channel + g.CondaConfig.UseMicroMamba = useMamba + return nil +} + +func CondaPackage(deps []string, channel []string, envFile string) error { + g := DefaultGraph.(*generalGraph) + + g.CondaConfig.CondaPackages = append( + g.CondaConfig.CondaPackages, deps...) + + g.CondaConfig.CondaEnvFileName = envFile + + if len(channel) != 0 { + g.CondaConfig.AdditionalChannels = append( + g.CondaConfig.AdditionalChannels, channel...) + } + return nil +} + +func Copy(src, dest string) { + g := DefaultGraph.(*generalGraph) + + g.Copy = append(g.Copy, ir.CopyInfo{ + Source: src, + Destination: dest, + }) +} + +func Mount(src, dest string) { + g := DefaultGraph.(*generalGraph) + + g.Mount = append(g.Mount, ir.MountInfo{ + Source: src, + Destination: dest, + }) +} + +func HTTP(url, checksum, filename string) error { + info := ir.HTTPInfo{ + URL: url, + Filename: filename, + } + if len(checksum) > 0 { + d, err := digest.Parse(checksum) + if err != nil { + return err + } + info.Checksum = d + } + g := DefaultGraph.(*generalGraph) + + g.HTTP = append(g.HTTP, info) + return nil +} + +func Entrypoint(args []string) { + g := DefaultGraph.(*generalGraph) + + g.Entrypoint = append(g.Entrypoint, args...) +} + +func RuntimeCommands(commands map[string]string) { + g := DefaultGraph.(*generalGraph) + + for k, v := range commands { + g.RuntimeCommands[k] = v + } +} + +func RuntimeDaemon(commands [][]string) { + g := DefaultGraph.(*generalGraph) + + g.RuntimeDaemon = append(g.RuntimeDaemon, commands...) +} + +func RuntimeExpose(envdPort, hostPort int, serviceName string, listeningAddr string) error { + g := DefaultGraph.(*generalGraph) + + g.RuntimeExpose = append(g.RuntimeExpose, ir.ExposeItem{ + EnvdPort: envdPort, + HostPort: hostPort, + ServiceName: serviceName, + ListeningAddr: listeningAddr, + }) + return nil +} + +func RuntimeEnviron(env map[string]string, path []string) { + g := DefaultGraph.(*generalGraph) + + for k, v := range env { + g.RuntimeEnviron[k] = v + } + g.RuntimeEnvPaths = append(g.RuntimeEnvPaths, path...) +} + +func RuntimeInitScript(commands []string) { + g := DefaultGraph.(*generalGraph) + + g.RuntimeInitScript = append(g.RuntimeInitScript, commands) +} + +func Repo(url, description string) { + g := DefaultGraph.(*generalGraph) + + g.Repo = types.RepoInfo{ + Description: description, + URL: url, + } +} diff --git a/pkg/lang/ir/v0/julia.go b/pkg/lang/ir/v0/julia.go new file mode 100644 index 000000000..c32012e70 --- /dev/null +++ b/pkg/lang/ir/v0/julia.go @@ -0,0 +1,101 @@ +// 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 v1 + +import ( + "fmt" + "strings" + + "github.com/cockroachdb/errors" + "github.com/moby/buildkit/client/llb" + "github.com/sirupsen/logrus" +) + +func (g generalGraph) compileJulia(baseStage llb.State) (llb.State, error) { + if err := g.compileJupyter(); err != nil { + return llb.State{}, errors.Wrap(err, "failed to compile jupyter") + } + + aptStage := g.compileUbuntuAPT(baseStage) + 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")) + + 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")) + + juliaStage := llb.Diff(builtinSystemStage, + g.installJuliaPackages(builtinSystemStage), llb.WithCustomName("install julia 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, juliaStage, *vscodeStage, + }, llb.WithCustomName("[internal] generating the image")) + } else { + merged = llb.Merge([]llb.State{ + builtinSystemStage, systemStage, diffShellStage, + diffSSHStage, juliaStage, + }, llb.WithCustomName("[internal] generating the image")) + } + return merged, nil +} + +func (g generalGraph) installJuliaPackages(root llb.State) llb.State { + if len(g.JuliaPackages) == 0 { + return root + } + + var sb strings.Builder + + sb.WriteString(`/usr/local/julia/bin/julia -e 'using Pkg; Pkg.add([`) + for i, pkg := range g.JuliaPackages { + sb.WriteString(fmt.Sprintf(`"%s"`, pkg)) + if i != len(g.JuliaPackages)-1 { + sb.WriteString(", ") + } + } + + sb.WriteString(`])'`) + + // TODO(gaocegege): Support cache. + cmd := sb.String() + logrus.Debug("install julia packages: ", cmd) + root = llb.User("envd")(root) + if g.JuliaPackageServer != nil { + root = root.AddEnv("JULIA_PKG_SERVER", *g.JuliaPackageServer) + } + root = root.AddEnv("PATH", "/usr/local/julia/bin") + run := root. + Run(llb.Shlex(cmd), llb.WithCustomNamef("install julia packages")) + + return run.Root() +} diff --git a/pkg/lang/ir/v0/python.go b/pkg/lang/ir/v0/python.go new file mode 100644 index 000000000..5dbf3752b --- /dev/null +++ b/pkg/lang/ir/v0/python.go @@ -0,0 +1,226 @@ +// 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 v1 + +import ( + "fmt" + "path/filepath" + "strings" + + "github.com/cockroachdb/errors" + "github.com/moby/buildkit/client/llb" + "github.com/sirupsen/logrus" + + "github.com/tensorchord/envd/pkg/flag" +) + +const ( + pythonVersionDefault = "3.9" +) + +func (g generalGraph) getAppropriatePythonVersion() (string, error) { + if g.Language.Version == nil { + return pythonVersionDefault, nil + } + + version := *g.Language.Version + if version == "3" || version == "" { + return pythonVersionDefault, nil + } + if strings.HasPrefix(version, "3.") { + return version, nil + } + return "", errors.Errorf("python version %s is not supported", version) +} + +func (g generalGraph) compilePython(baseStage llb.State) (llb.State, error) { + if err := g.compileJupyter(); err != nil { + return llb.State{}, errors.Wrap(err, "failed to compile jupyter") + } + aptStage := g.compileUbuntuAPT(baseStage) + systemStage := g.compileSystemPackages(aptStage) + + condaEnvStage, err := g.compileCondaEnvironment(baseStage) + if err != nil { + return llb.State{}, errors.Wrap(err, "failed to compile conda environment") + } + + // Conda affects shell and python, thus we cannot do it in parallel. + shellStage, err := g.compileShell(baseStage) + if err != nil { + return llb.State{}, errors.Wrap(err, "failed to compile shell") + } + condaShellStage := g.compileCondaShell(shellStage) + + diffCondaEnvStage := llb.Diff(baseStage, condaEnvStage, + llb.WithCustomName("[internal] conda python environment")) + diffSystemStage := llb.Diff(baseStage, systemStage, + llb.WithCustomName("[internal] install system packages")) + diffShellStage := llb.Diff(baseStage, condaShellStage, + llb.WithCustomNamef("[internal] configure shell %s", g.Shell)) + prePythonStage := llb.Merge([]llb.State{ + diffSystemStage, + diffCondaEnvStage, + diffShellStage, + baseStage}, llb.WithCustomName("pre-python stage")) + + condaChannelStage := g.compileCondaChannel(prePythonStage) + + condaStage := llb.Diff(prePythonStage, + g.compileCondaPackages(condaChannelStage), + llb.WithCustomName("[internal] install conda packages")) + + pypiMirrorStage := g.compilePyPIIndex(prePythonStage) + + pypiStage := llb.Diff(prePythonStage, + g.compilePyPIPackages(pypiMirrorStage), + llb.WithCustomName("[internal] install PyPI packages")) + + vscodeStage, err := g.compileVSCode() + if err != nil { + return llb.State{}, errors.Wrap(err, "failed to get vscode plugins") + } + sshStage, err := g.copySSHKey(prePythonStage) + if err != nil { + return llb.State{}, errors.Wrap(err, "failed to copy ssh keys") + } + diffSSHStage := llb.Diff(prePythonStage, sshStage, + llb.WithCustomName("[internal] install ssh key")) + + var merged llb.State + if vscodeStage != nil { + merged = llb.Merge([]llb.State{ + prePythonStage, condaStage, pypiStage, + diffSSHStage, *vscodeStage, + }, llb.WithCustomName("[internal] generating the image")) + } else { + merged = llb.Merge([]llb.State{ + prePythonStage, condaStage, + diffSSHStage, pypiStage, + }, llb.WithCustomName("[internal] generating the image")) + } + merged = g.compileAlternative(merged) + return merged, nil +} + +// Set the system default python to envd's python. +func (g generalGraph) compileAlternative(root llb.State) llb.State { + envdPrefix := "/opt/conda/envs/envd/bin" + run := root. + Run(llb.Shlexf("update-alternatives --install /usr/bin/python python %s/python 1", envdPrefix), + llb.WithCustomName("[internal] update alternative python to envd")). + Run(llb.Shlexf("update-alternatives --install /usr/bin/python3 python3 %s/python3 1", envdPrefix), + llb.WithCustomName("[internal] update alternative python3 to envd")). + Run(llb.Shlexf("update-alternatives --install /usr/bin/pip pip %s/pip 1", envdPrefix), + llb.WithCustomName("[internal] update alternative pip to envd")). + Run(llb.Shlexf("update-alternatives --install /usr/bin/pip3 pip3 %s/pip3 1", envdPrefix), + llb.WithCustomName("[internal] update alternative pip3 to envd")) + return run.Root() +} + +func (g generalGraph) compilePyPIPackages(root llb.State) llb.State { + if len(g.PyPIPackages) == 0 && g.RequirementsFile == nil && len(g.PythonWheels) == 0 { + return root + } + + // Create the envd cache directory in the container. see issue #582 + cacheDir := filepath.Join("/", "root", ".cache", "pip") + root = g.CompileCacheDir(root, cacheDir) + + cache := root.File(llb.Mkdir("/cache/pip", 0755, llb.WithParents(true)), + llb.WithCustomName("[internal] setting pip cache mount permissions")) + + if len(g.PyPIPackages) != 0 { + // Compose the package install command. + var sb strings.Builder + // Always use the conda's pip. + sb.WriteString("/opt/conda/envs/envd/bin/python -m pip install") + for _, pkg := range g.PyPIPackages { + sb.WriteString(fmt.Sprintf(" %s", pkg)) + } + + cmd := sb.String() + logrus.WithField("command", cmd). + Debug("Configure pip install statements") + run := root. + Run(llb.Shlex(sb.String()), llb.WithCustomNamef("pip install %s", + strings.Join(g.PyPIPackages, " "))) + // Refer to https://github.com/moby/buildkit/blob/31054718bf775bf32d1376fe1f3611985f837584/frontend/dockerfile/dockerfile2llb/convert_runmount.go#L46 + run.AddMount(cacheDir, cache, + llb.AsPersistentCacheDir(g.CacheID(cacheDir), llb.CacheMountShared), llb.SourcePath("/cache/pip")) + root = run.Root() + } + + if g.RequirementsFile != nil { + // Compose the package install command. + var sb strings.Builder + sb.WriteString("bash -c '") + sb.WriteString("set -euo pipefail\n") + sb.WriteString(fmt.Sprintf("chown -R envd:envd %s\n", g.getWorkingDir())) // Change mount dir permission + envdCmd := strings.Builder{} + envdCmd.WriteString(fmt.Sprintf("cd %s\n", g.getWorkingDir())) + envdCmd.WriteString(fmt.Sprintf("/opt/conda/envs/envd/bin/python -m pip install -r %s\n", *g.RequirementsFile)) + + // Execute the command to write yaml file and conda env using envd user + sb.WriteString(fmt.Sprintf("sudo -i -u envd bash << EOF\n%s\nEOF\n", envdCmd.String())) + sb.WriteString("'") + cmd := sb.String() + + logrus.WithField("command", cmd). + Debug("Configure pip install requirements statements") + root = root.User("root").Dir(g.getWorkingDir()) + run := root. + Run(llb.Shlex(cmd), llb.WithCustomNamef("pip install %s", *g.RequirementsFile)) + run.AddMount(cacheDir, cache, + llb.AsPersistentCacheDir(g.CacheID(cacheDir), llb.CacheMountShared), llb.SourcePath("/cache/pip")) + run.AddMount(g.getWorkingDir(), + llb.Local(flag.FlagBuildContext)) + root = run.Root() + } + + if len(g.PythonWheels) > 0 { + root = root.Dir(g.getWorkingDir()) + cmdTemplate := "/opt/conda/envs/envd/bin/python -m pip install %s" + for _, wheel := range g.PythonWheels { + run := root.Run(llb.Shlex(fmt.Sprintf(cmdTemplate, wheel)), llb.WithCustomNamef("pip install %s", wheel)) + run.AddMount(g.getWorkingDir(), llb.Local(flag.FlagBuildContext), llb.Readonly) + run.AddMount(cacheDir, cache, + llb.AsPersistentCacheDir(g.CacheID(cacheDir), llb.CacheMountShared), llb.SourcePath("/cache/pip")) + root = run.Root() + } + } + return root +} + +func (g generalGraph) compilePyPIIndex(root llb.State) llb.State { + if g.PyPIIndexURL != nil { + logrus.WithField("index", *g.PyPIIndexURL).Debug("using custom PyPI index") + var extraIndex string + if g.PyPIExtraIndexURL != nil { + logrus.WithField("index", *g.PyPIIndexURL).Debug("using extra PyPI index") + extraIndex = "extra-index-url=" + *g.PyPIExtraIndexURL + } + content := fmt.Sprintf(pypiConfigTemplate, *g.PyPIIndexURL, extraIndex) + dir := filepath.Dir(pypiIndexFilePath) + pypiMirror := root. + File(llb.Mkdir(dir, 0755, llb.WithParents(true), llb.WithUIDGID(g.uid, g.gid)), + llb.WithCustomNamef("[internal] setting PyPI index dir %s", dir)). + File(llb.Mkfile(pypiIndexFilePath, + 0644, []byte(content), llb.WithUIDGID(g.uid, g.gid)), + llb.WithCustomNamef("[internal] setting PyPI index file %s", pypiIndexFilePath)) + return pypiMirror + } + return root +} diff --git a/pkg/lang/ir/v0/r.go b/pkg/lang/ir/v0/r.go new file mode 100644 index 000000000..9a64dd212 --- /dev/null +++ b/pkg/lang/ir/v0/r.go @@ -0,0 +1,96 @@ +// 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 v1 + +import ( + "fmt" + "strings" + + "github.com/cockroachdb/errors" + "github.com/moby/buildkit/client/llb" +) + +func (g generalGraph) compileRLang(baseStage llb.State) (llb.State, error) { + if err := g.compileJupyter(); err != nil { + return llb.State{}, errors.Wrap(err, "failed to compile jupyter") + } + aptStage := g.compileUbuntuAPT(baseStage) + 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("[internal] generating the image")) + } else { + merged = llb.Merge([]llb.State{ + builtinSystemStage, systemStage, diffShellStage, + diffSSHStage, rPackageInstallStage, + }, llb.WithCustomName("[internal] generating the image")) + } + return merged, nil +} + +func (g generalGraph) installRPackages(root llb.State) llb.State { + if len(g.RPackages) == 0 { + return root + } + // TODO(terrytangyuan): Support different CRAN mirrors + var sb strings.Builder + mirrorURL := "https://cran.rstudio.com" + if g.CRANMirrorURL != nil { + mirrorURL = *g.CRANMirrorURL + } + sb.WriteString(fmt.Sprintf(`R -e 'options(repos = c(CRAN = "%s")); install.packages(c(`, mirrorURL)) + 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("install R packages")) + return run.Root() +} diff --git a/pkg/lang/ir/v0/shell.go b/pkg/lang/ir/v0/shell.go new file mode 100644 index 000000000..f243b77f5 --- /dev/null +++ b/pkg/lang/ir/v0/shell.go @@ -0,0 +1,119 @@ +// 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 v1 + +import ( + "fmt" + + "github.com/cockroachdb/errors" + "github.com/moby/buildkit/client/llb" + + "github.com/tensorchord/envd/pkg/flag" + "github.com/tensorchord/envd/pkg/progress/compileui" + "github.com/tensorchord/envd/pkg/shell" + "github.com/tensorchord/envd/pkg/util/fileutil" +) + +const ( + starshipConfig = ` +[container] +format = "[$symbol \\[envd\\]]($style)" + +[sudo] +disabled = false +symbol = "sudo " + +[python] +symbol = "Py " + +[status] +format = '[\[$status:$common_meaning$signal_name\]]($style) ' +disabled = false +` +) + +func (g *generalGraph) compileShell(root llb.State) (llb.State, error) { + if g.Shell == shellZSH { + g.RuntimeEnviron["SHELL"] = "/usr/bin/zsh" + return g.compileZSH(root) + } + g.RuntimeEnviron["SHELL"] = "/usr/bin/bash" + return root, nil +} + +func (g *generalGraph) compileCondaShell(root llb.State) llb.State { + var run llb.ExecState + switch g.Shell { + case shellBASH: + run = root.Run( + llb.Shlex( + fmt.Sprintf(`bash -c 'echo "source %s/activate envd" >> %s'`, + condaBinDir, fileutil.EnvdHomeDir(".bashrc"))), + llb.WithCustomName("[internal] add conda environment to bashrc")) + case shellZSH: + run = root.Run( + llb.Shlex(fmt.Sprintf("bash -c \"%s\"", g.condaInitShell(g.Shell))), + llb.WithCustomNamef("[internal] initialize conda %s environment", g.Shell)).Run( + llb.Shlex(fmt.Sprintf(`bash -c 'echo "source %s/activate envd" >> %s'`, condaBinDir, fileutil.EnvdHomeDir(".zshrc"))), + llb.WithCustomName("[internal] add conda environment to zshrc")) + } + return run.Root() +} + +func (g *generalGraph) compilePrompt(root llb.State) llb.State { + // skip this for customized image + if g.Image != nil { + return root + } + // starship config + config := root. + File(llb.Mkdir(defaultConfigDir, 0755, llb.WithParents(true)), + llb.WithCustomName("[internal] creating config dir")). + File(llb.Mkfile(starshipConfigPath, 0644, []byte(starshipConfig), llb.WithUIDGID(g.uid, g.gid)), + llb.WithCustomName("[internal] setting prompt starship config")) + + run := config.Run(llb.Shlex(fmt.Sprintf(`bash -c 'echo "eval \"\$(starship init bash)\"" >> %s'`, fileutil.EnvdHomeDir(".bashrc"))), + llb.WithCustomName("[internal] setting prompt bash config")).Root() + + if g.Shell == shellZSH { + run = run.Run( + llb.Shlex(fmt.Sprintf(`bash -c 'echo "eval \"\$(starship init zsh)\"" >> %s'`, fileutil.EnvdHomeDir(".zshrc"))), + llb.WithCustomName("[internal] setting prompt zsh config")).Root() + } + return run +} + +func (g generalGraph) compileZSH(root llb.State) (llb.State, error) { + installPath := fileutil.EnvdHomeDir("install.sh") + zshrcPath := fileutil.EnvdHomeDir(".zshrc") + ohMyZSHPath := fileutil.EnvdHomeDir(".oh-my-zsh") + m := shell.NewManager() + g.Writer.LogZSH(compileui.ActionStart, false) + if cached, err := m.DownloadOrCache(); err != nil { + return llb.State{}, errors.Wrap(err, "failed to download oh-my-zsh") + } else { + g.Writer.LogZSH(compileui.ActionEnd, cached) + } + zshStage := root. + File(llb.Copy(llb.Local(flag.FlagCacheDir), "oh-my-zsh", ohMyZSHPath, + &llb.CopyInfo{CreateDestPath: true}, llb.WithUIDGID(g.uid, g.gid))). + File(llb.Mkfile(installPath, + 0644, []byte(m.InstallScript()), llb.WithUIDGID(g.uid, g.gid))) + zshrc := zshStage.Run(llb.Shlex(fmt.Sprintf("bash %s", installPath)), + llb.WithCustomName("[internal] install oh-my-zsh")). + File(llb.Mkfile(zshrcPath, + 0644, []byte(m.ZSHRC()), llb.WithUIDGID(g.uid, g.gid))) + return zshrc, nil +} diff --git a/pkg/lang/ir/v0/supervisor.go b/pkg/lang/ir/v0/supervisor.go new file mode 100644 index 000000000..f6b7116da --- /dev/null +++ b/pkg/lang/ir/v0/supervisor.go @@ -0,0 +1,117 @@ +// 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 v1 + +import ( + "fmt" + "path/filepath" + "strings" + + "github.com/cockroachdb/errors" + "github.com/moby/buildkit/client/llb" + + "github.com/tensorchord/envd/pkg/config" + "github.com/tensorchord/envd/pkg/types" +) + +const ( + horustTemplate = ` +name = "%[1]s" +command = """ +%[2]s +""" +stdout = "/var/log/horust/%[1]s_stdout.log" +stderr = "/var/log/horust/%[1]s_stderr.log" +user = "${USER}" +working-directory = "${%[3]s}" +%[4]s + +[environment] +keep-env = true + +[restart] +strategy = "on-failure" +backoff = "1s" +attempts = 2 + +[termination] +wait = "5s" +` +) + +func (g generalGraph) installHorust(root llb.State) llb.State { + horust := root. + File(llb.Copy(llb.Image(types.HorustImage), "/", "/usr/local/bin"), + llb.WithCustomName("[internal] install horust")). + File(llb.Mkdir(types.HorustServiceDir, 0755, llb.WithParents(true)), + llb.WithCustomNamef("[internal] mkdir for horust service: %s", types.HorustServiceDir)). + File(llb.Mkdir(types.HorustLogDir, 0777, llb.WithParents(true)), + llb.WithCustomNamef("[internal] mkdir for horust log: %s", types.HorustLogDir)) + return horust +} + +func (g generalGraph) addNewProcess(root llb.State, name, command string, depends []string) llb.State { + var sb strings.Builder + if len(depends) != 0 { + sb.WriteString("start-after = [") + for _, d := range depends { + sb.WriteString("\"") + sb.WriteString(d) + sb.WriteString("\",") + } + sb.WriteString("]\n") + } + template := fmt.Sprintf(horustTemplate, name, command, types.EnvdWorkDir, sb.String()) + + filename := filepath.Join(types.HorustServiceDir, fmt.Sprintf("%s.toml", name)) + supervisor := root.File(llb.Mkfile(filename, 0644, []byte(template), llb.WithUIDGID(g.uid, g.gid)), llb.WithCustomNamef("[internal] create file %s", filename)) + return supervisor +} + +func (g generalGraph) compileEntrypoint(root llb.State) (llb.State, error) { + if g.Image != nil { + return root, nil + } + if len(g.Entrypoint) > 0 { + return root, errors.New("`config.entrypoint` is only for custom image, maybe you need `runtime.init`") + } + cmd := fmt.Sprintf("/var/envd/bin/envd-sshd --port %d --shell %s", config.SSHPortInContainer, g.Shell) + entrypoint := g.addNewProcess(root, "sshd", cmd, nil) + var deps []string + if g.RuntimeInitScript != nil { + for i, command := range g.RuntimeInitScript { + entrypoint = g.addNewProcess(entrypoint, fmt.Sprintf("init_%d", i), fmt.Sprintf("/bin/bash -c 'set -euo pipefail\n%s'", strings.Join(command, "\n")), nil) + deps = append(deps, fmt.Sprintf("init_%d", i)) + } + } + + if g.RuntimeDaemon != nil { + for i, command := range g.RuntimeDaemon { + entrypoint = g.addNewProcess(entrypoint, fmt.Sprintf("daemon_%d", i), strings.Join(command, " "), deps) + } + } + + if g.JupyterConfig != nil { + jupyterCmd := g.generateJupyterCommand("") + entrypoint = g.addNewProcess(entrypoint, "jupyter", strings.Join(jupyterCmd, " "), deps) + } + + if g.RStudioServerConfig != nil { + rstudioCmd := g.generateRStudioCommand("") + entrypoint = g.addNewProcess(entrypoint, "rstudio", strings.Join(rstudioCmd, " "), deps) + } + + return entrypoint, nil +} diff --git a/pkg/lang/ir/v0/system.go b/pkg/lang/ir/v0/system.go new file mode 100644 index 000000000..6c90819eb --- /dev/null +++ b/pkg/lang/ir/v0/system.go @@ -0,0 +1,289 @@ +// 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 v1 + +import ( + "context" + _ "embed" + "fmt" + "os" + "path/filepath" + "strings" + + "github.com/cockroachdb/errors" + "github.com/moby/buildkit/client/llb" + "github.com/moby/buildkit/client/llb/imagemetaresolver" + "github.com/sirupsen/logrus" + "github.com/spf13/viper" + + "github.com/tensorchord/envd/pkg/config" + "github.com/tensorchord/envd/pkg/flag" + "github.com/tensorchord/envd/pkg/types" + "github.com/tensorchord/envd/pkg/util/fileutil" + "github.com/tensorchord/envd/pkg/version" +) + +func (g generalGraph) compileUbuntuAPT(root llb.State) llb.State { + if g.UbuntuAPTSource != nil { + logrus.WithField("source", *g.UbuntuAPTSource).Debug("using custom APT source") + aptSource := root. + File(llb.Mkdir(filepath.Dir(aptSourceFilePath), 0755, llb.WithParents(true)), + llb.WithCustomName("[internal] setting apt source")). + File(llb.Mkfile(aptSourceFilePath, 0644, []byte(*g.UbuntuAPTSource)), + llb.WithCustomName("[internal] setting apt source")) + return aptSource + } + return root +} + +func (g generalGraph) compileRun(root llb.State) llb.State { + if len(g.Exec) == 0 { + return root + } + + workingDir := g.getWorkingDir() + stage := root.AddEnv("PATH", types.DefaultPathEnvUnix) + for _, execGroup := range g.Exec { + var sb strings.Builder + sb.WriteString("set -euo pipefail\n") + for _, c := range execGroup.Commands { + sb.WriteString(c + "\n") + } + + cmdStr := fmt.Sprintf("/usr/bin/bash -c '%s'", sb.String()) + logrus.WithField("command", cmdStr).Debug("compile run command") + // Mount the build context into the build process. + // TODO(gaocegege): Maybe we should make it readonly, + // but these cases then cannot be supported: + // run(commands=["git clone xx.git"]) + run := stage.Dir(workingDir).Run(llb.Shlex(cmdStr)) + if execGroup.MountHost { + run.AddMount(workingDir, llb.Local(flag.FlagBuildContext)) + } + stage = run.Root() + } + return stage +} + +func (g generalGraph) compileCopy(root llb.State) llb.State { + if len(g.Copy) == 0 { + return root + } + + result := root + // Compose the copy command. + for _, c := range g.Copy { + result = result.File(llb.Copy( + llb.Local(flag.FlagBuildContext), c.Source, c.Destination, + llb.WithUIDGID(g.uid, g.gid))) + } + return result +} + +func (g *generalGraph) compileCUDAPackages(org string) llb.State { + return g.preparePythonBase(llb.Image(fmt.Sprintf( + "docker.io/%s:%s-cudnn%s-devel-%s", + org, *g.CUDA, g.CUDNN, g.OS))) +} + +func (g generalGraph) compileSystemPackages(root llb.State) llb.State { + if len(g.SystemPackages) == 0 { + logrus.Debug("skip the apt since system package is not specified") + return root + } + + // Compose the package install command. + var sb strings.Builder + sb.WriteString("sudo apt-get update && sudo apt-get install -y --no-install-recommends") + + for _, pkg := range g.SystemPackages { + sb.WriteString(fmt.Sprintf(" %s", pkg)) + } + + cacheDir := "/var/cache/apt" + cacheLibDir := "/var/lib/apt" + + run := root.Run(llb.Shlex(fmt.Sprintf("bash -c \"%s\"", sb.String())), + llb.WithCustomNamef("apt-get install %s", + strings.Join(g.SystemPackages, " "))) + run.AddMount(cacheDir, llb.Scratch(), + llb.AsPersistentCacheDir(g.CacheID(cacheDir), llb.CacheMountShared)) + run.AddMount(cacheLibDir, llb.Scratch(), + llb.AsPersistentCacheDir(g.CacheID(cacheLibDir), llb.CacheMountShared)) + return run.Root() +} + +// nolint:unparam +func (g *generalGraph) compileExtraSource(root llb.State) (llb.State, error) { + if len(g.HTTP) == 0 { + return root, nil + } + inputs := []llb.State{} + for _, httpInfo := range g.HTTP { + src := llb.HTTP( + httpInfo.URL, + llb.Checksum(httpInfo.Checksum), + llb.Filename(httpInfo.Filename), + llb.Chown(g.uid, g.gid), + ) + inputs = append(inputs, llb.Scratch().File( + llb.Copy(src, "/", g.getExtraSourceDir(), &llb.CopyInfo{CreateDestPath: true}), + )) + } + inputs = append(inputs, root) + return llb.Merge(inputs, llb.WithCustomName("[internal] build source layers")), nil +} + +func (g *generalGraph) preparePythonBase(root llb.State) llb.State { + for _, env := range types.BaseEnvironment { + root = root.AddEnv(env.Name, env.Value) + } + + // apt packages + var sb strings.Builder + sb.WriteString("apt-get update && apt-get install -y apt-utils && ") + sb.WriteString("apt-get install -y --no-install-recommends --no-install-suggests --fix-missing ") + sb.WriteString(strings.Join(types.BaseAptPackage, " ")) + sb.WriteString("&& rm -rf /var/lib/apt/lists/* ") + // shell prompt + sb.WriteString("&& curl --proto '=https' --tlsv1.2 -sSf https://starship.rs/install.sh | sh -s -- -y") + sb.WriteString("&& locale-gen en_US.UTF-8") + + run := root.Run(llb.Shlex(fmt.Sprintf("bash -c \"%s\"", sb.String())), + llb.WithCustomName("[internal] install built-in packages")) + + return run.Root() +} + +func (g generalGraph) compileSSHD(root llb.State) llb.State { + sshd := root.File(llb.Copy( + llb.Image(types.EnvdSshdImage), "/usr/bin/envd-sshd", "/var/envd/bin/envd-sshd", + &llb.CopyInfo{CreateDestPath: true}), + llb.WithCustomName(fmt.Sprintf("[internal] add envd-sshd from %s", types.EnvdSshdImage))) + return sshd +} + +func (g *generalGraph) compileBase() (llb.State, error) { + 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 + org := viper.GetString(flag.FlagDockerOrganization) + v := version.GetVersionForImageTag() + // Do not update user permission in the base image. + if g.Image != nil { + return g.customBase() + } else if g.CUDA == nil { + switch g.Language.Name { + case "r": + base = llb.Image(fmt.Sprintf("docker.io/%s/r-base:4.2-envd-%s", org, v)) + // 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 + } + case "python": + // TODO(keming) use user input `base(os="")` + base = g.preparePythonBase(llb.Image(types.PythonBaseImage)) + case "julia": + base = llb.Image(fmt.Sprintf( + "docker.io/%s/julia:1.8rc1-ubuntu20.04-envd-%s", org, v)) + } + } else { + base = g.compileCUDAPackages("nvidia/cuda") + } + + // Install conda first. + condaStage, err := g.installConda(base) + if err != nil { + return llb.State{}, errors.Wrap(err, "failed to install conda") + } + supervisor := g.installHorust(condaStage) + sshdStage := g.compileSSHD(supervisor) + source, err := g.compileExtraSource(sshdStage) + if err != nil { + return llb.State{}, errors.Wrap(err, "failed to get extra sources") + } + final := g.compileUserGroup(source) + return final, nil +} + +// customBase get the image and the set the image metadata to graph. +func (g *generalGraph) customBase() (llb.State, error) { + if g.Image == nil { + return llb.State{}, fmt.Errorf("failed to get the image") + } + logrus.WithField("image", *g.Image).Debugf("using custom base image") + + // Fix https://github.com/tensorchord/envd/issues/1147. + // Fetch the image metadata from base image. + base := llb.Image(*g.Image, + llb.WithMetaResolver(imagemetaresolver.Default())) + envs, err := base.Env(context.Background()) + if err != nil { + return llb.State{}, errors.Wrap(err, "failed to get the image metadata") + } + + // Set the environment variables to RuntimeEnviron to keep it in the resulting image. + for _, e := range envs { + kv := strings.Split(e, "=") + g.RuntimeEnviron[kv[0]] = kv[1] + } + return base, nil +} + +func (g generalGraph) copySSHKey(root llb.State) (llb.State, error) { + public := g.PublicKeyPath + bdat, err := os.ReadFile(public) + dat := strings.TrimSuffix(string(bdat), "\n") + if err != nil { + return llb.State{}, errors.Wrap(err, "Cannot read public SSH key") + } + run := root. + File(llb.Mkdir("/var/envd", 0755, llb.WithParents(true), + llb.WithUIDGID(g.uid, g.gid)), + llb.WithCustomName("[internal] create dir for ssh key")). + File(llb.Mkfile(config.ContainerAuthorizedKeysPath, + 0644, []byte(dat+" envd"), llb.WithUIDGID(g.uid, g.gid)), + llb.WithCustomName("[internal] install ssh keys")) + return run, nil +} + +func (g generalGraph) compileMountDir(root llb.State) llb.State { + mount := root + if g.Image == nil { + // create the ENVD_WORKDIR as a placeholder (envd-server may not mount this dir) + workDir := fileutil.EnvdHomeDir(g.EnvironmentName) + mount = root.File(llb.Mkdir(workDir, 0755, llb.WithParents(true), llb.WithUIDGID(g.uid, g.gid)), + llb.WithCustomNamef("[internal] create work dir: %s", workDir)) + } + + for _, m := range g.Mount { + mount = mount.File(llb.Mkdir(m.Destination, 0755, llb.WithParents(true), + llb.WithUIDGID(g.uid, g.gid)), + llb.WithCustomNamef("[internal] create dir for runtime.mount %s", m.Destination), + ) + } + return mount +} diff --git a/pkg/lang/ir/v0/types.go b/pkg/lang/ir/v0/types.go new file mode 100644 index 000000000..9883e374f --- /dev/null +++ b/pkg/lang/ir/v0/types.go @@ -0,0 +1,70 @@ +package v1 + +import ( + "github.com/tensorchord/envd/pkg/editor/vscode" + "github.com/tensorchord/envd/pkg/lang/ir" + "github.com/tensorchord/envd/pkg/progress/compileui" + "github.com/tensorchord/envd/pkg/types" +) + +// A Graph contains the state, +// such as its call stack and thread-local storage. +// TODO(gaocegeg): Refactor it to support order. +type generalGraph struct { + uid int + gid int + + OS string + ir.Language + Image *string + + Shell string + CUDA *string + CUDNN string + NumGPUs int + + UbuntuAPTSource *string + CRANMirrorURL *string + JuliaPackageServer *string + PyPIIndexURL *string + PyPIExtraIndexURL *string + + PublicKeyPath string + + PyPIPackages []string + RequirementsFile *string + PythonWheels []string + RPackages []string + JuliaPackages []string + SystemPackages []string + + VSCodePlugins []vscode.Plugin + UserDirectories []string + RuntimeEnvPaths []string + + Exec []ir.RunBuildCommand + Copy []ir.CopyInfo + Mount []ir.MountInfo + HTTP []ir.HTTPInfo + Entrypoint []string + + Repo types.RepoInfo + + *ir.JupyterConfig + *ir.GitConfig + *ir.CondaConfig + *ir.RStudioServerConfig + + Writer compileui.Writer + // EnvironmentName is the base name of the environment. + // It is the BaseDir(BuildContextDir) + // e.g. mnist, streamlit-mnist + EnvironmentName string + + ir.RuntimeGraph +} + +const ( + shellBASH = "bash" + shellZSH = "zsh" +) diff --git a/pkg/lang/ir/v0/user.go b/pkg/lang/ir/v0/user.go new file mode 100644 index 000000000..4820504a1 --- /dev/null +++ b/pkg/lang/ir/v0/user.go @@ -0,0 +1,73 @@ +// 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 v1 + +import ( + "fmt" + + "github.com/moby/buildkit/client/llb" +) + +// compileUserOwn chown related directories +func (g *generalGraph) compileUserOwn(root llb.State) llb.State { + if g.Image != nil || g.uid == 0 { + g.RuntimeEnviron["USER"] = "root" + return root + } + g.RuntimeEnviron["USER"] = "envd" + if len(g.UserDirectories) == 0 { + return root.User("envd") + } + run := root.Run() + for _, dir := range g.UserDirectories { + run = root.Run(llb.Shlex(fmt.Sprintf("chown -R envd:envd %s", dir)), + llb.WithCustomNamef("[internal] configure user permissions for %s", dir)) + } + return run.Root().User("envd") +} + +// compileUserGroup creates user `envd` +func (g *generalGraph) compileUserGroup(root llb.State) llb.State { + if g.Image != nil { + return root + } + var res llb.ExecState + if g.uid == 0 { + res = root. + Run(llb.Shlex(fmt.Sprintf("groupadd -g %d envd", 1001)), + llb.WithCustomName("[internal] still create group envd for root context")). + Run(llb.Shlex(fmt.Sprintf("useradd -p \"\" -u %d -g envd -s /bin/sh -m envd", 1001)), + llb.WithCustomName("[internal] still create user envd for root context")). + Run(llb.Shlex("usermod -s /bin/sh root"), + llb.WithCustomName("[internal] set root default shell to /bin/sh")). + Run(llb.Shlex("sed -i \"s/envd:x:1001:1001/envd:x:0:0/g\" /etc/passwd"), + llb.WithCustomName("[internal] set envd uid to 0 as root")). + Run(llb.Shlex("sed -i \"s./root./home/envd.g\" /etc/passwd"), + llb.WithCustomName("[internal] set root home dir to /home/envd")). + Run(llb.Shlex("sed -i \"s/envd:x:1001/envd:x:0/g\" /etc/group"), + llb.WithCustomName("[internal] set envd group to 0 as root group")) + } else { + res = root. + Run(llb.Shlex(fmt.Sprintf("groupadd -g %d envd", g.gid)), + llb.WithCustomName("[internal] create user group envd")). + Run(llb.Shlex(fmt.Sprintf("useradd -p \"\" -u %d -g envd -s /bin/sh -m envd", g.uid)), + llb.WithCustomName("[internal] create user envd")). + Run(llb.Shlex("adduser envd sudo"), + llb.WithCustomName("[internal] add user envd to sudoers")). + Run(llb.Shlex(fmt.Sprintf("install -d -o envd -g %d -m 0700 /home/envd/.config /home/envd/.cache", g.gid)), + llb.WithCustomName("[internal] mkdir config and cache dir")) + } + return res.Root() +} diff --git a/pkg/lang/ir/v0/util.go b/pkg/lang/ir/v0/util.go new file mode 100644 index 000000000..2cac34f56 --- /dev/null +++ b/pkg/lang/ir/v0/util.go @@ -0,0 +1,118 @@ +// 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 v1 + +import ( + "bytes" + "crypto/md5" + "encoding/gob" + "encoding/hex" + "os/user" + "regexp" + "strconv" + "strings" + + "github.com/cockroachdb/errors" + "github.com/sirupsen/logrus" + "github.com/spf13/viper" + + "github.com/tensorchord/envd/pkg/flag" + "github.com/tensorchord/envd/pkg/util/fileutil" +) + +func (g generalGraph) getWorkingDir() string { + return fileutil.EnvdHomeDir(g.EnvironmentName) +} + +func (g generalGraph) getExtraSourceDir() string { + return fileutil.EnvdHomeDir("extra_source") +} + +func parseLanguage(l string) (string, *string, error) { + var language, version string + if l == "" { + return "", nil, errors.New("language is required") + } + + // Get version from the string. + 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", "r", "julia": + return language, &version, nil + default: + return "", nil, errors.Newf("language %s is not supported", language) + } +} + +func getUIDGID() (int, int, error) { + owner := viper.GetString(flag.FlagBuildOwner) + if len(owner) > 0 { + logrus.WithField("flag", owner).Info("use owner") + ids := strings.Split(owner, ":") + if len(ids) > 2 { + return 0, 0, errors.Newf("wrong format for owner (uid:gid): %s", owner) + } + uid, err := strconv.Atoi(ids[0]) + if err != nil { + logrus.Info(err) + return 0, 0, errors.Wrap(err, "failed to get uid") + } + // if omit gid, will use the uid as gid + if len(ids) == 1 { + return uid, uid, nil + } + gid, err := strconv.Atoi(ids[1]) + if err != nil { + return 0, 0, errors.Wrap(err, "failed to get gid") + } + return uid, gid, nil + } + user, err := user.Current() + if err != nil { + return 0, 0, errors.Wrap(err, "failed to get uid/gid") + } + // Do not support windows yet. + uid, err := strconv.Atoi(user.Uid) + if err != nil { + return 0, 0, errors.Wrap(err, "failed to get uid") + } + gid, err := strconv.Atoi(user.Gid) + if err != nil { + return 0, 0, errors.Wrap(err, "failed to get gid") + } + return uid, gid, nil +} + +// A stream of gobs is self-describing. Each data item in the stream is preceded by a specification of its type, expressed in terms of a small set of predefined types. Pointers are not transmitted, but the things they point to are transmitted; that is, the values are flattened. +// see https://pkg.go.dev/encoding/gob#hdr-Basics +// we hash the blobs to determined if the graph changed. +func GetDefaultGraphHash() string { + var b bytes.Buffer + err := gob.NewEncoder(&b).Encode(DefaultGraph) + if err != nil { + return "" + } + data := b.Bytes() + hashD := md5.Sum(data) + return hex.EncodeToString(hashD[:]) +} diff --git a/pkg/lang/ir/v0/util_test.go b/pkg/lang/ir/v0/util_test.go new file mode 100644 index 000000000..ee1bb8b3f --- /dev/null +++ b/pkg/lang/ir/v0/util_test.go @@ -0,0 +1,78 @@ +// 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 v1 + +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) + } + } + } + + } +} diff --git a/pkg/lang/ir/v1/compile.go b/pkg/lang/ir/v1/compile.go index 5f4e68333..77367e9ba 100644 --- a/pkg/lang/ir/v1/compile.go +++ b/pkg/lang/ir/v1/compile.go @@ -43,9 +43,8 @@ func NewGraph() ir.Graph { RuntimeEnviron: make(map[string]string), } langVersion := languageVersionDefault - conda := &ir.CondaConfig{} return &generalGraph{ - OS: osDefault, + Image: defaultImage, Language: ir.Language{ Name: languageDefault, Version: &langVersion, @@ -62,13 +61,20 @@ func NewGraph() ir.Graph { UserDirectories: []string{}, RuntimeEnvPaths: []string{types.DefaultPathEnv()}, Shell: shellBASH, - CondaConfig: conda, RuntimeGraph: runtimeGraph, } } var DefaultGraph = NewGraph() +func (g *generalGraph) SetWriter(w compileui.Writer) { + g.Writer = w +} + +func (g generalGraph) GetHTTP() []ir.HTTPInfo { + return g.HTTP +} + func (g generalGraph) GetNumGPUs() int { return g.NumGPUs } @@ -126,6 +132,19 @@ func (g *generalGraph) Compile(ctx context.Context, envName string, pub string) return def, nil } +func (g generalGraph) GetEnviron() []string { + // Add PATH and LC_ALL. + return append(g.EnvString(), + "PATH="+strings.Join(g.RuntimeEnvPaths, ":"), + "LC_ALL=en_US.UTF-8", + "LANG=C.UTF-8", + ) +} + +func (g generalGraph) GetUser() string { + return g.User +} + func (g generalGraph) GPUEnabled() bool { return g.CUDA != nil } @@ -205,8 +224,8 @@ func (g generalGraph) Labels() (map[string]string, error) { func (g generalGraph) ExposedPorts() (map[string]struct{}, error) { ports := make(map[string]struct{}) - // do not expose ports for custom images - if g.Image != nil { + // only expose ports for dev env + if !g.Dev { return ports, nil } @@ -235,25 +254,6 @@ func (g generalGraph) EnvString() []string { return envs } -func (g generalGraph) GetEnviron() []string { - if g.Image != nil { - return g.EnvString() - } - // Add PATH and LC_ALL. - return append(g.EnvString(), - "PATH="+strings.Join(g.RuntimeEnvPaths, ":"), - "LC_ALL=en_US.UTF-8", - ) -} - -func (g *generalGraph) SetWriter(w compileui.Writer) { - g.Writer = w -} - -func (g generalGraph) GetHTTP() []ir.HTTPInfo { - return g.HTTP -} - func (g generalGraph) DefaultCacheImporter() (*string, error) { // The base remote cache should work for all languages. var res string @@ -272,14 +272,14 @@ func (g generalGraph) DefaultCacheImporter() (*string, error) { } func (g *generalGraph) GetEntrypoint(buildContextDir string) ([]string, error) { - if g.Image != nil { + if !g.Dev { return g.Entrypoint, nil } g.RuntimeEnviron[types.EnvdWorkDir] = fileutil.EnvdHomeDir(filepath.Base(buildContextDir)) return []string{"horust"}, nil } -func (g generalGraph) CompileLLB(uid, gid int) (llb.State, error) { +func (g *generalGraph) CompileLLB(uid, gid int) (llb.State, error) { g.uid = uid // TODO(gaocegege): Remove the hack for https://github.com/tensorchord/envd/issues/370 @@ -289,49 +289,73 @@ func (g generalGraph) CompileLLB(uid, gid int) (llb.State, error) { "gid": g.gid, }).Debug("compile LLB") - // TODO(gaocegege): Support more OS and langs. - aptStage, err := g.compileBase() + base, err := g.compileBaseImage() if err != nil { return llb.State{}, errors.Wrap(err, "failed to get the base image") } - var merged llb.State - // Use custom logic when image is specified. - if g.Image != nil { - merged, err = g.compileCustomPython(aptStage) + + // prepare dev env: stable operations should be done here to make it cache friendly + if g.Dev { + dev := g.compileDevPackages(base) + sshd := g.compileSSHD(dev) + horust := g.installHorust(sshd) + userGroup := g.compileUserGroup(horust) + base = userGroup + } + + lang, err := g.compileLanguage(base) + if err != nil { + return llb.State{}, errors.Wrap(err, "failed to compile language") + } + aptMirror := g.compileUbuntuAPT(base) + systemPackages := g.compileSystemPackages(aptMirror) + merge := llb.Merge([]llb.State{ + base, + llb.Diff(base, lang, llb.WithCustomName("[internal] prepare language")), + llb.Diff(base, systemPackages, llb.WithCustomName("[internal] install system packages")), + }, llb.WithCustomName("[internal] language environment and system packages")) + packages := g.compileLanguagePackages(merge) + if err != nil { + return llb.State{}, errors.Wrap(err, "failed to compile language") + } + + source, err := g.compileExtraSource(packages) + if err != nil { + return llb.State{}, errors.Wrap(err, "failed to compile extra source") + } + copy := g.compileCopy(source) + + // dev postprocessing: related to UID, which may not be cached + if g.Dev { + git := g.compileGit(copy) + user := g.compileUserOwn(git) + key, err := g.copySSHKey(user) if err != nil { - return llb.State{}, errors.Wrap(err, "failed to compile custom python image") + return llb.State{}, errors.Wrap(err, "failed to copy ssh key") } - } else { - switch g.Language.Name { - case "r": - merged, err = g.compileRLang(aptStage) - if err != nil { - return llb.State{}, errors.Wrap(err, "failed to compile r language") - } - case "python": - merged, err = g.compilePython(aptStage) - if err != nil { - return llb.State{}, errors.Wrap(err, "failed to compile python") - } - case "julia": - merged, err = g.compileJulia(aptStage) - if err != nil { - return llb.State{}, errors.Wrap(err, "failed to compile julia") - } + shell, err := g.compileShell(key) + if err != nil { + return llb.State{}, errors.Wrap(err, "failed to compile shell") } + prompt := g.compilePrompt(shell) + entrypoint, err := g.compileEntrypoint(prompt) + if err != nil { + return llb.State{}, errors.Wrap(err, "failed to compile entrypoint") + } + vscode, err := g.compileVSCode() + if err != nil { + return llb.State{}, errors.Wrap(err, "failed to compile VSCode extensions") + } + copy = llb.Merge([]llb.State{ + entrypoint, + vscode, + }, llb.WithCustomName("[internal] final dev environment")) } - prompt := g.compilePrompt(merged) - copy := g.compileCopy(prompt) - // TODO(gaocegege): Support order-based exec. + // it's necessary to exec `run`` with the desired user run := g.compileRun(copy) - git := g.compileGit(run) - user := g.compileUserOwn(git) - mount := g.compileMountDir(user) - entrypoint, err := g.compileEntrypoint(mount) - if err != nil { - return llb.State{}, errors.Wrap(err, "failed to compile entrypoint") - } + mount := g.compileMountDir(run) + g.Writer.Finish() - return entrypoint, nil + return mount, nil } diff --git a/pkg/lang/ir/v1/conda.go b/pkg/lang/ir/v1/conda.go index 3f2d6c8b6..7a4ace517 100644 --- a/pkg/lang/ir/v1/conda.go +++ b/pkg/lang/ir/v1/conda.go @@ -28,21 +28,32 @@ import ( ) const ( + builderImage = "curlimages/curl:7.86.0" condaVersionDefault = "py39_4.11.0" - // check the issue https://github.com/mamba-org/mamba/issues/1975 - mambaVersionDefault = "0.25.1" + microMambaImage = "mambaorg/micromamba:1.0.0" condaRootPrefix = "/opt/conda" condaBinDir = "/opt/conda/bin" + condaSourcePath = "/tmp/miniconda.sh" + + mambaRc = ` +channels: + - conda-forge +` + mambaActivate = ` +#!/bin/sh +eval "$(/opt/conda/bin/micromamba shell hook --shell=bash)" || return $? +micromamba activate "$@" +` ) var ( // this file can be used by both conda and mamba // https://mamba.readthedocs.io/en/latest/user_guide/configuration.html#multiple-rc-files condarc = "/opt/conda/.condarc" - //go:embed install-conda.sh + //go:embed get_conda.sh + downloadCondaBash string + //go:embed install_conda.sh installCondaBash string - //go:embed install-mamba.sh - installMambaBash string ) func (g generalGraph) compileCondaChannel(root llb.State) llb.State { @@ -112,8 +123,9 @@ func (g *generalGraph) compileCondaPackages(root llb.State) llb.State { cmd := sb.String() run = root.Dir(g.getWorkingDir()). + AddEnv("MAMBA_ROOT_PREFIX", condaRootPrefix). Run(llb.Shlex(cmd), llb.WithCustomNamef("[internal] %s %s", - cmd, strings.Join(g.CondaPackages, " "))) + cmd, strings.Join(g.CondaConfig.CondaPackages, " "))) run.AddMount(g.getWorkingDir(), llb.Local(flag.FlagBuildContext)) run.AddMount(cacheDir, cacheMount, llb.AsPersistentCacheDir(g.CacheID(cacheDir), llb.CacheMountShared), llb.SourcePath("/cache-conda")) @@ -140,19 +152,50 @@ func (g generalGraph) compileCondaEnvironment(root llb.State) (llb.State, error) } // nolint:unparam -func (g generalGraph) installConda(root llb.State) (llb.State, error) { +func (g *generalGraph) installConda(root llb.State) llb.State { if g.CondaConfig.UseMicroMamba { - run := root.AddEnv("MAMBA_BIN_DIR", condaBinDir). - AddEnv("MAMBA_ROOT_PREFIX", condaRootPrefix). - AddEnv("MAMBA_VERSION", mambaVersionDefault). - Run(llb.Shlex(fmt.Sprintf("bash -c '%s'", installMambaBash)), - llb.WithCustomName("[internal] install micro mamba")) - return run.Root(), nil + return g.installMicroMamba(root) } - run := root.AddEnv("CONDA_VERSION", condaVersionDefault). + return g.installMiniConda(root) +} + +func (g generalGraph) installMiniConda(root llb.State) llb.State { + base := llb.Image(builderImage) + builder := base.AddEnv("CONDA_VERSION", condaVersionDefault). + Run(llb.Shlexf("sh -c '%s'", downloadCondaBash), + llb.WithCustomName("[internal] download conda")).Root() + conda := root. + File(llb.Copy(builder, condaSourcePath, condaSourcePath), + llb.WithCustomName("copy conda from builder")). File(llb.Mkdir(condaRootPrefix, 0755, llb.WithParents(true)), llb.WithCustomName("[internal] create conda directory")). - Run(llb.Shlex(fmt.Sprintf("bash -c '%s'", installCondaBash)), - llb.WithCustomName("[internal] install conda")) - return run.Root(), nil + Run(llb.Shlexf("bash -c '%s'", installCondaBash), + llb.WithCustomName("[internal] install conda")).Root(). + File(llb.Rm(condaSourcePath), llb.WithCustomName("[internal] rm conda source file")) + return conda +} + +func (g *generalGraph) installMicroMamba(root llb.State) llb.State { + g.RuntimeEnviron["MAMBA_ROOT_PREFIX"] = condaRootPrefix + g.RuntimeEnviron["MAMBA_TARGET_PREFIX"] = condaRootPrefix + mamba := root. + AddEnv("MAMBA_ROOT_PREFIX", condaRootPrefix). + AddEnv("MAMBA_TARGET_PREFIX", condaRootPrefix). + File(llb.Mkdir(certPath, 0755, llb.WithParents(true)), + llb.WithCustomName("[internal] mkdir certs")). + File(llb.Copy(llb.Image(microMambaImage), fmt.Sprintf("%s/%s", certPath, "ca-certificates.crt"), certPath), + llb.WithCustomName("[internal] copy cert from mamba")). + File(llb.Mkdir(condaBinDir, 0755, llb.WithParents(true)), + llb.WithCustomName("[internal] create mamba path")). + File(llb.Copy(llb.Image(microMambaImage), "/bin/micromamba", condaBinDir), + llb.WithCustomName("[internal] copy micromamba binary")). + File(llb.Mkfile(fmt.Sprintf("%s/.mambarc", condaRootPrefix), 0644, []byte(mambaRc)), + llb.WithCustomName("[internal] create the mamba rc file")). + File(llb.Mkfile(fmt.Sprintf("%s/activate", condaBinDir), 0755, []byte(mambaActivate)), + llb.WithCustomName("[internal] create the mamba activate file")). + Run(llb.Shlexf("update-alternatives --install /usr/bin/conda conda %s/micromamba 1", condaBinDir), + llb.WithCustomName("[internal] update alternative micromamba to conda")). + Run(llb.Shlexf("bash -c \"%s/micromamba shell init --shell bash\"", condaBinDir), + llb.WithCustomName("[internal] init micromamba for bash")).Root() + return mamba } diff --git a/pkg/lang/ir/v1/consts.go b/pkg/lang/ir/v1/consts.go index 9f154ea66..d6cca6848 100644 --- a/pkg/lang/ir/v1/consts.go +++ b/pkg/lang/ir/v1/consts.go @@ -17,7 +17,7 @@ package v1 import "github.com/tensorchord/envd/pkg/util/fileutil" const ( - osDefault = "ubuntu20.04" + defaultImage = "ubuntu:20.04" languageDefault = "python" languageVersionDefault = "3" CUDNNVersionDefault = "8" diff --git a/pkg/lang/ir/v1/editor.go b/pkg/lang/ir/v1/editor.go index 9e95ad688..3ee0ad3c1 100644 --- a/pkg/lang/ir/v1/editor.go +++ b/pkg/lang/ir/v1/editor.go @@ -29,19 +29,19 @@ import ( "github.com/tensorchord/envd/pkg/util/fileutil" ) -func (g generalGraph) compileVSCode() (*llb.State, error) { +func (g generalGraph) compileVSCode() (llb.State, error) { if len(g.VSCodePlugins) == 0 { - return nil, nil + return llb.Scratch(), nil } inputs := []llb.State{} for _, p := range g.VSCodePlugins { vscodeClient, err := vscode.NewClient(vscode.MarketplaceVendorOpenVSX) if err != nil { - return nil, errors.Wrap(err, "failed to create vscode client") + return llb.State{}, errors.Wrap(err, "failed to create vscode client") } g.Writer.LogVSCodePlugin(p, compileui.ActionStart, false) if cached, err := vscodeClient.DownloadOrCache(p); err != nil { - return nil, err + return llb.State{}, err } else { g.Writer.LogVSCodePlugin(p, compileui.ActionEnd, cached) } @@ -55,9 +55,10 @@ func (g generalGraph) compileVSCode() (*llb.State, error) { inputs = append(inputs, ext) } layer := llb.Merge(inputs, llb.WithCustomName("merging plugins for vscode")) - return &layer, nil + return layer, nil } +// nolint:unused func (g *generalGraph) compileJupyter() error { if g.JupyterConfig == nil { return nil diff --git a/pkg/lang/ir/v1/editor_test.go b/pkg/lang/ir/v1/editor_test.go index dfb41cdc1..11ab187ae 100644 --- a/pkg/lang/ir/v1/editor_test.go +++ b/pkg/lang/ir/v1/editor_test.go @@ -16,19 +16,17 @@ package v1 import ( "testing" - - "github.com/tensorchord/envd/pkg/lang/ir" ) func TestGenerateCommand(t *testing.T) { testcases := []struct { - graph generalGraph + graph Graph dir string expected []string }{ { - graph: generalGraph{ - JupyterConfig: &ir.JupyterConfig{ + graph: Graph{ + JupyterConfig: &JupyterConfig{ Token: "", Port: 8888, }, @@ -41,8 +39,8 @@ func TestGenerateCommand(t *testing.T) { }, }, { - graph: generalGraph{ - JupyterConfig: &ir.JupyterConfig{ + graph: Graph{ + JupyterConfig: &JupyterConfig{ Token: "test", Port: 8888, }, @@ -55,8 +53,8 @@ func TestGenerateCommand(t *testing.T) { }, }, { - graph: generalGraph{ - JupyterConfig: &ir.JupyterConfig{ + graph: Graph{ + JupyterConfig: &JupyterConfig{ Token: "test", Port: 8888, }, @@ -69,7 +67,7 @@ func TestGenerateCommand(t *testing.T) { }, }, { - graph: generalGraph{}, + graph: Graph{}, dir: "test", expected: []string{}, }, diff --git a/pkg/lang/ir/v1/get_conda.sh b/pkg/lang/ir/v1/get_conda.sh new file mode 100644 index 000000000..c3dbbcdd2 --- /dev/null +++ b/pkg/lang/ir/v1/get_conda.sh @@ -0,0 +1,18 @@ +set -euo pipefail && \ +UNAME_M="$(uname -m)" && \ +if [ "${UNAME_M}" = "x86_64" ]; then \ + MINICONDA_URL="https://repo.anaconda.com/miniconda/Miniconda3-${CONDA_VERSION}-Linux-x86_64.sh"; \ + SHA256SUM="4ee9c3aa53329cd7a63b49877c0babb49b19b7e5af29807b793a76bdb1d362b4"; \ +elif [ "${UNAME_M}" = "s390x" ]; then \ + MINICONDA_URL="https://repo.anaconda.com/miniconda/Miniconda3-${CONDA_VERSION}-Linux-s390x.sh"; \ + SHA256SUM="e5e5e89cdcef9332fe632cd25d318cf71f681eef029a24495c713b18e66a8018"; \ +elif [ "${UNAME_M}" = "aarch64" ]; then \ + MINICONDA_URL="https://repo.anaconda.com/miniconda/Miniconda3-${CONDA_VERSION}-Linux-aarch64.sh"; \ + SHA256SUM="00c7127a8a8d3f4b9c2ab3391c661239d5b9a88eafe895fd0f3f2a8d9c0f4556"; \ +elif [ "${UNAME_M}" = "ppc64le" ]; then \ + MINICONDA_URL="https://repo.anaconda.com/miniconda/Miniconda3-${CONDA_VERSION}-Linux-ppc64le.sh"; \ + SHA256SUM="8ee1f8d17ef7c8cb08a85f7d858b1cb55866c06fcf7545b98c3b82e4d0277e66"; \ +fi && \ +wget "${MINICONDA_URL}" -O /tmp/miniconda.sh && \ +echo "${SHA256SUM} /tmp/miniconda.sh" > /tmp/shasum && \ +if [ "${CONDA_VERSION}" != "latest" ]; then sha256sum -c -s /tmp/shasum; fi diff --git a/pkg/lang/ir/v1/install_conda.sh b/pkg/lang/ir/v1/install_conda.sh new file mode 100644 index 000000000..b753d0972 --- /dev/null +++ b/pkg/lang/ir/v1/install_conda.sh @@ -0,0 +1,9 @@ +set -euo pipefail && \ +sh /tmp/miniconda.sh -b -u -p /opt/conda && \ +touch ~/.bashrc && \ +echo ". /opt/conda/etc/profile.d/conda.sh" >> ~/.bashrc && \ +echo "conda activate base" >> ~/.bashrc && \ +echo -e "channels:\n - conda-forge" > /opt/conda/.condarc && \ +find /opt/conda/ -follow -type f -name '*.a' -delete && \ +find /opt/conda/ -follow -type f -name '*.js.map' -delete && \ +/opt/conda/bin/conda clean -afy diff --git a/pkg/lang/ir/v1/interface.go b/pkg/lang/ir/v1/interface.go index 1bee1b3f3..264ce680e 100644 --- a/pkg/lang/ir/v1/interface.go +++ b/pkg/lang/ir/v1/interface.go @@ -15,31 +15,63 @@ package v1 import ( + "strings" + "github.com/cockroachdb/errors" "github.com/opencontainers/go-digest" + "github.com/sirupsen/logrus" "github.com/tensorchord/envd/pkg/editor/vscode" "github.com/tensorchord/envd/pkg/lang/ir" "github.com/tensorchord/envd/pkg/types" ) -func Base(os, language, image string) error { - l, version, err := parseLanguage(language) - if err != nil { - return err +func Base(image string, dev bool) error { + g := DefaultGraph.(*generalGraph) + + if image != "" { + g.Image = image + } + g.Dev = dev + return nil +} + +func Python(version string) error { + if strings.HasPrefix(version, "2") { + logrus.Debugf("envd doesn't support Python2: %s", version) + return errors.New("envd doesn't support this Python version") } g := DefaultGraph.(*generalGraph) + g.Language = ir.Language{ - Name: l, - Version: version, + Name: "python", + Version: &version, } - if len(os) > 0 { - g.OS = os + return nil +} + +func Conda(mamba bool) { + g := DefaultGraph.(*generalGraph) + + g.CondaConfig = &ir.CondaConfig{ + UseMicroMamba: mamba, } - if image != "" { - g.Image = &image +} + +func RLang() { + g := DefaultGraph.(*generalGraph) + + g.Language = ir.Language{ + Name: "r", + } +} + +func Julia() { + g := DefaultGraph.(*generalGraph) + + g.Language = ir.Language{ + Name: "julia", } - return nil } func PyPIPackage(deps []string, requirementsFile string, wheels []string) error { @@ -182,11 +214,13 @@ func Git(name, email, editor string) error { return nil } -func CondaChannel(channel string, useMamba bool) error { +func CondaChannel(channel string) error { g := DefaultGraph.(*generalGraph) + if g.CondaConfig == nil { + return errors.New("cannot config conda when conda is not installed") + } g.CondaConfig.CondaChannel = &channel - g.CondaConfig.UseMicroMamba = useMamba return nil } @@ -224,6 +258,8 @@ func Mount(src, dest string) { } func HTTP(url, checksum, filename string) error { + g := DefaultGraph.(*generalGraph) + info := ir.HTTPInfo{ URL: url, Filename: filename, @@ -235,8 +271,6 @@ func HTTP(url, checksum, filename string) error { } info.Checksum = d } - g := DefaultGraph.(*generalGraph) - g.HTTP = append(g.HTTP, info) return nil } diff --git a/pkg/lang/ir/v1/julia.go b/pkg/lang/ir/v1/julia.go index c32012e70..762231844 100644 --- a/pkg/lang/ir/v1/julia.go +++ b/pkg/lang/ir/v1/julia.go @@ -23,50 +23,8 @@ import ( "github.com/sirupsen/logrus" ) -func (g generalGraph) compileJulia(baseStage llb.State) (llb.State, error) { - if err := g.compileJupyter(); err != nil { - return llb.State{}, errors.Wrap(err, "failed to compile jupyter") - } - - aptStage := g.compileUbuntuAPT(baseStage) - 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")) - - 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")) - - juliaStage := llb.Diff(builtinSystemStage, - g.installJuliaPackages(builtinSystemStage), llb.WithCustomName("install julia 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, juliaStage, *vscodeStage, - }, llb.WithCustomName("[internal] generating the image")) - } else { - merged = llb.Merge([]llb.State{ - builtinSystemStage, systemStage, diffShellStage, - diffSSHStage, juliaStage, - }, llb.WithCustomName("[internal] generating the image")) - } - return merged, nil +func (g generalGraph) installJulia(root llb.State) (llb.State, error) { + return llb.State{}, errors.New("not implemented") } func (g generalGraph) installJuliaPackages(root llb.State) llb.State { @@ -93,7 +51,6 @@ func (g generalGraph) installJuliaPackages(root llb.State) llb.State { if g.JuliaPackageServer != nil { root = root.AddEnv("JULIA_PKG_SERVER", *g.JuliaPackageServer) } - root = root.AddEnv("PATH", "/usr/local/julia/bin") run := root. Run(llb.Shlex(cmd), llb.WithCustomNamef("install julia packages")) diff --git a/pkg/lang/ir/v1/python.go b/pkg/lang/ir/v1/python.go index 5dbf3752b..563b08f56 100644 --- a/pkg/lang/ir/v1/python.go +++ b/pkg/lang/ir/v1/python.go @@ -27,17 +27,51 @@ import ( ) const ( - pythonVersionDefault = "3.9" + PythonVersionDefault = "3.9" + microMambaPathPrefix = "/usr/local/bin" + certPath = "/etc/ssl/certs" ) +func (g *generalGraph) installPython(root llb.State) (llb.State, error) { + if g.CondaConfig == nil { + version, err := g.getAppropriatePythonVersion() + if err != nil { + return llb.State{}, err + } + install := root. + File(llb.Mkdir(certPath, 0755, llb.WithParents(true)), + llb.WithCustomName("[internal] mkdir certs")). + File(llb.Copy(llb.Image(microMambaImage), fmt.Sprintf("%s/%s", certPath, "ca-certificates.crt"), certPath), + llb.WithCustomName("[internal] copy cert from mamba")). + File(llb.Copy(llb.Image(microMambaImage), "/bin/micromamba", microMambaPathPrefix), + llb.WithCustomName("[internal] copy micromamba binary")). + Run(llb.Shlex(fmt.Sprintf("bash -c \"%s/micromamba create -p /opt/conda/envs/envd -c conda-forge python=%s\"", microMambaPathPrefix, version)), + llb.WithCustomNamef("[internal] create envd python=%s", version)). + Run(llb.Shlex(fmt.Sprintf("rm %s/micromamba", microMambaPathPrefix)), + llb.WithCustomName("[internal] rm micromamba binary")).Root() + python := g.compileAlternative(install) + return python, nil + } + + // install Conda to create the env + py := g.installConda(root) + env, err := g.compileCondaEnvironment(py) + if err != nil { + return llb.State{}, err + } + + python := g.compileAlternative(env) + return python, nil +} + func (g generalGraph) getAppropriatePythonVersion() (string, error) { if g.Language.Version == nil { - return pythonVersionDefault, nil + return PythonVersionDefault, nil } version := *g.Language.Version if version == "3" || version == "" { - return pythonVersionDefault, nil + return PythonVersionDefault, nil } if strings.HasPrefix(version, "3.") { return version, nil @@ -45,76 +79,6 @@ func (g generalGraph) getAppropriatePythonVersion() (string, error) { return "", errors.Errorf("python version %s is not supported", version) } -func (g generalGraph) compilePython(baseStage llb.State) (llb.State, error) { - if err := g.compileJupyter(); err != nil { - return llb.State{}, errors.Wrap(err, "failed to compile jupyter") - } - aptStage := g.compileUbuntuAPT(baseStage) - systemStage := g.compileSystemPackages(aptStage) - - condaEnvStage, err := g.compileCondaEnvironment(baseStage) - if err != nil { - return llb.State{}, errors.Wrap(err, "failed to compile conda environment") - } - - // Conda affects shell and python, thus we cannot do it in parallel. - shellStage, err := g.compileShell(baseStage) - if err != nil { - return llb.State{}, errors.Wrap(err, "failed to compile shell") - } - condaShellStage := g.compileCondaShell(shellStage) - - diffCondaEnvStage := llb.Diff(baseStage, condaEnvStage, - llb.WithCustomName("[internal] conda python environment")) - diffSystemStage := llb.Diff(baseStage, systemStage, - llb.WithCustomName("[internal] install system packages")) - diffShellStage := llb.Diff(baseStage, condaShellStage, - llb.WithCustomNamef("[internal] configure shell %s", g.Shell)) - prePythonStage := llb.Merge([]llb.State{ - diffSystemStage, - diffCondaEnvStage, - diffShellStage, - baseStage}, llb.WithCustomName("pre-python stage")) - - condaChannelStage := g.compileCondaChannel(prePythonStage) - - condaStage := llb.Diff(prePythonStage, - g.compileCondaPackages(condaChannelStage), - llb.WithCustomName("[internal] install conda packages")) - - pypiMirrorStage := g.compilePyPIIndex(prePythonStage) - - pypiStage := llb.Diff(prePythonStage, - g.compilePyPIPackages(pypiMirrorStage), - llb.WithCustomName("[internal] install PyPI packages")) - - vscodeStage, err := g.compileVSCode() - if err != nil { - return llb.State{}, errors.Wrap(err, "failed to get vscode plugins") - } - sshStage, err := g.copySSHKey(prePythonStage) - if err != nil { - return llb.State{}, errors.Wrap(err, "failed to copy ssh keys") - } - diffSSHStage := llb.Diff(prePythonStage, sshStage, - llb.WithCustomName("[internal] install ssh key")) - - var merged llb.State - if vscodeStage != nil { - merged = llb.Merge([]llb.State{ - prePythonStage, condaStage, pypiStage, - diffSSHStage, *vscodeStage, - }, llb.WithCustomName("[internal] generating the image")) - } else { - merged = llb.Merge([]llb.State{ - prePythonStage, condaStage, - diffSSHStage, pypiStage, - }, llb.WithCustomName("[internal] generating the image")) - } - merged = g.compileAlternative(merged) - return merged, nil -} - // Set the system default python to envd's python. func (g generalGraph) compileAlternative(root llb.State) llb.State { envdPrefix := "/opt/conda/envs/envd/bin" @@ -146,7 +110,7 @@ func (g generalGraph) compilePyPIPackages(root llb.State) llb.State { // Compose the package install command. var sb strings.Builder // Always use the conda's pip. - sb.WriteString("/opt/conda/envs/envd/bin/python -m pip install") + sb.WriteString("python -m pip install") for _, pkg := range g.PyPIPackages { sb.WriteString(fmt.Sprintf(" %s", pkg)) } @@ -155,7 +119,7 @@ func (g generalGraph) compilePyPIPackages(root llb.State) llb.State { logrus.WithField("command", cmd). Debug("Configure pip install statements") run := root. - Run(llb.Shlex(sb.String()), llb.WithCustomNamef("pip install %s", + Run(llb.Shlex(sb.String()), llb.WithCustomNamef("[internal] pip install %s", strings.Join(g.PyPIPackages, " "))) // Refer to https://github.com/moby/buildkit/blob/31054718bf775bf32d1376fe1f3611985f837584/frontend/dockerfile/dockerfile2llb/convert_runmount.go#L46 run.AddMount(cacheDir, cache, @@ -171,7 +135,7 @@ func (g generalGraph) compilePyPIPackages(root llb.State) llb.State { sb.WriteString(fmt.Sprintf("chown -R envd:envd %s\n", g.getWorkingDir())) // Change mount dir permission envdCmd := strings.Builder{} envdCmd.WriteString(fmt.Sprintf("cd %s\n", g.getWorkingDir())) - envdCmd.WriteString(fmt.Sprintf("/opt/conda/envs/envd/bin/python -m pip install -r %s\n", *g.RequirementsFile)) + envdCmd.WriteString(fmt.Sprintf("python -m pip install -r %s\n", *g.RequirementsFile)) // Execute the command to write yaml file and conda env using envd user sb.WriteString(fmt.Sprintf("sudo -i -u envd bash << EOF\n%s\nEOF\n", envdCmd.String())) @@ -192,7 +156,7 @@ func (g generalGraph) compilePyPIPackages(root llb.State) llb.State { if len(g.PythonWheels) > 0 { root = root.Dir(g.getWorkingDir()) - cmdTemplate := "/opt/conda/envs/envd/bin/python -m pip install %s" + cmdTemplate := "python -m pip install %s" for _, wheel := range g.PythonWheels { run := root.Run(llb.Shlex(fmt.Sprintf(cmdTemplate, wheel)), llb.WithCustomNamef("pip install %s", wheel)) run.AddMount(g.getWorkingDir(), llb.Local(flag.FlagBuildContext), llb.Readonly) @@ -205,22 +169,22 @@ func (g generalGraph) compilePyPIPackages(root llb.State) llb.State { } func (g generalGraph) compilePyPIIndex(root llb.State) llb.State { - if g.PyPIIndexURL != nil { - logrus.WithField("index", *g.PyPIIndexURL).Debug("using custom PyPI index") - var extraIndex string - if g.PyPIExtraIndexURL != nil { - logrus.WithField("index", *g.PyPIIndexURL).Debug("using extra PyPI index") - extraIndex = "extra-index-url=" + *g.PyPIExtraIndexURL - } - content := fmt.Sprintf(pypiConfigTemplate, *g.PyPIIndexURL, extraIndex) - dir := filepath.Dir(pypiIndexFilePath) - pypiMirror := root. - File(llb.Mkdir(dir, 0755, llb.WithParents(true), llb.WithUIDGID(g.uid, g.gid)), - llb.WithCustomNamef("[internal] setting PyPI index dir %s", dir)). - File(llb.Mkfile(pypiIndexFilePath, - 0644, []byte(content), llb.WithUIDGID(g.uid, g.gid)), - llb.WithCustomNamef("[internal] setting PyPI index file %s", pypiIndexFilePath)) - return pypiMirror + if g.PyPIIndexURL == nil { + return root } - return root + logrus.WithField("index", *g.PyPIIndexURL).Debug("using custom PyPI index") + var extraIndex string + if g.PyPIExtraIndexURL != nil { + logrus.WithField("index", *g.PyPIIndexURL).Debug("using extra PyPI index") + extraIndex = "extra-index-url=" + *g.PyPIExtraIndexURL + } + content := fmt.Sprintf(pypiConfigTemplate, *g.PyPIIndexURL, extraIndex) + dir := filepath.Dir(pypiIndexFilePath) + pypiMirror := root. + File(llb.Mkdir(dir, 0755, llb.WithParents(true), llb.WithUIDGID(g.uid, g.gid)), + llb.WithCustomNamef("[internal] setting PyPI index dir %s", dir)). + File(llb.Mkfile(pypiIndexFilePath, + 0644, []byte(content), llb.WithUIDGID(g.uid, g.gid)), + llb.WithCustomNamef("[internal] setting PyPI index file %s", pypiIndexFilePath)) + return pypiMirror } diff --git a/pkg/lang/ir/v1/r.go b/pkg/lang/ir/v1/r.go index 9a64dd212..cd9bfcb1d 100644 --- a/pkg/lang/ir/v1/r.go +++ b/pkg/lang/ir/v1/r.go @@ -22,51 +22,8 @@ import ( "github.com/moby/buildkit/client/llb" ) -func (g generalGraph) compileRLang(baseStage llb.State) (llb.State, error) { - if err := g.compileJupyter(); err != nil { - return llb.State{}, errors.Wrap(err, "failed to compile jupyter") - } - aptStage := g.compileUbuntuAPT(baseStage) - 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("[internal] generating the image")) - } else { - merged = llb.Merge([]llb.State{ - builtinSystemStage, systemStage, diffShellStage, - diffSSHStage, rPackageInstallStage, - }, llb.WithCustomName("[internal] generating the image")) - } - return merged, nil +func (g generalGraph) installRLang(root llb.State) (llb.State, error) { + return llb.State{}, errors.New("not implemented") } func (g generalGraph) installRPackages(root llb.State) llb.State { diff --git a/pkg/lang/ir/v1/shell.go b/pkg/lang/ir/v1/shell.go index f243b77f5..31ba22ee1 100644 --- a/pkg/lang/ir/v1/shell.go +++ b/pkg/lang/ir/v1/shell.go @@ -44,39 +44,40 @@ disabled = false ` ) -func (g *generalGraph) compileShell(root llb.State) (llb.State, error) { +func (g *generalGraph) compileShell(root llb.State) (_ llb.State, err error) { + g.RuntimeEnviron["SHELL"] = "/usr/bin/bash" if g.Shell == shellZSH { g.RuntimeEnviron["SHELL"] = "/usr/bin/zsh" - return g.compileZSH(root) + root, err = g.compileZSH(root) + if err != nil { + return llb.State{}, err + } + } + if g.CondaConfig != nil { + root = g.compileCondaShell(root) } - g.RuntimeEnviron["SHELL"] = "/usr/bin/bash" return root, nil } func (g *generalGraph) compileCondaShell(root llb.State) llb.State { var run llb.ExecState - switch g.Shell { - case shellBASH: - run = root.Run( - llb.Shlex( - fmt.Sprintf(`bash -c 'echo "source %s/activate envd" >> %s'`, - condaBinDir, fileutil.EnvdHomeDir(".bashrc"))), - llb.WithCustomName("[internal] add conda environment to bashrc")) - case shellZSH: - run = root.Run( - llb.Shlex(fmt.Sprintf("bash -c \"%s\"", g.condaInitShell(g.Shell))), - llb.WithCustomNamef("[internal] initialize conda %s environment", g.Shell)).Run( - llb.Shlex(fmt.Sprintf(`bash -c 'echo "source %s/activate envd" >> %s'`, condaBinDir, fileutil.EnvdHomeDir(".zshrc"))), - llb.WithCustomName("[internal] add conda environment to zshrc")) + findDir := fileutil.DefaultHomeDir + if g.Dev { + findDir = fileutil.EnvdHomeDir } + rcPath := findDir(".bashrc") + if g.Shell == shellZSH { + rcPath = findDir(".zshrc") + } + run = root. + Run(llb.Shlexf("bash -c \"%s\"", g.condaInitShell(g.Shell)), + llb.WithCustomNamef("[internal] init conda %s env", g.Shell)). + Run(llb.Shlexf(`bash -c 'echo "source %s/activate envd" >> %s'`, condaBinDir, rcPath), + llb.WithCustomName("[internal] add conda environment to zshrc")) return run.Root() } func (g *generalGraph) compilePrompt(root llb.State) llb.State { - // skip this for customized image - if g.Image != nil { - return root - } // starship config config := root. File(llb.Mkdir(defaultConfigDir, 0755, llb.WithParents(true)), @@ -108,12 +109,10 @@ func (g generalGraph) compileZSH(root llb.State) (llb.State, error) { } zshStage := root. File(llb.Copy(llb.Local(flag.FlagCacheDir), "oh-my-zsh", ohMyZSHPath, - &llb.CopyInfo{CreateDestPath: true}, llb.WithUIDGID(g.uid, g.gid))). - File(llb.Mkfile(installPath, - 0644, []byte(m.InstallScript()), llb.WithUIDGID(g.uid, g.gid))) + &llb.CopyInfo{CreateDestPath: true})). + File(llb.Mkfile(installPath, 0666, []byte(m.InstallScript()))) zshrc := zshStage.Run(llb.Shlex(fmt.Sprintf("bash %s", installPath)), llb.WithCustomName("[internal] install oh-my-zsh")). - File(llb.Mkfile(zshrcPath, - 0644, []byte(m.ZSHRC()), llb.WithUIDGID(g.uid, g.gid))) + File(llb.Mkfile(zshrcPath, 0666, []byte(m.ZSHRC()))) return zshrc, nil } diff --git a/pkg/lang/ir/v1/supervisor.go b/pkg/lang/ir/v1/supervisor.go index f6b7116da..1ce7d87f6 100644 --- a/pkg/lang/ir/v1/supervisor.go +++ b/pkg/lang/ir/v1/supervisor.go @@ -81,9 +81,6 @@ func (g generalGraph) addNewProcess(root llb.State, name, command string, depend } func (g generalGraph) compileEntrypoint(root llb.State) (llb.State, error) { - if g.Image != nil { - return root, nil - } if len(g.Entrypoint) > 0 { return root, errors.New("`config.entrypoint` is only for custom image, maybe you need `runtime.init`") } diff --git a/pkg/lang/ir/v1/system.go b/pkg/lang/ir/v1/system.go index 6c90819eb..ac6c02bb8 100644 --- a/pkg/lang/ir/v1/system.go +++ b/pkg/lang/ir/v1/system.go @@ -26,13 +26,11 @@ import ( "github.com/moby/buildkit/client/llb" "github.com/moby/buildkit/client/llb/imagemetaresolver" "github.com/sirupsen/logrus" - "github.com/spf13/viper" "github.com/tensorchord/envd/pkg/config" "github.com/tensorchord/envd/pkg/flag" "github.com/tensorchord/envd/pkg/types" "github.com/tensorchord/envd/pkg/util/fileutil" - "github.com/tensorchord/envd/pkg/version" ) func (g generalGraph) compileUbuntuAPT(root llb.State) llb.State { @@ -92,12 +90,6 @@ func (g generalGraph) compileCopy(root llb.State) llb.State { return result } -func (g *generalGraph) compileCUDAPackages(org string) llb.State { - return g.preparePythonBase(llb.Image(fmt.Sprintf( - "docker.io/%s:%s-cudnn%s-devel-%s", - org, *g.CUDA, g.CUDNN, g.OS))) -} - func (g generalGraph) compileSystemPackages(root llb.State) llb.State { if len(g.SystemPackages) == 0 { logrus.Debug("skip the apt since system package is not specified") @@ -146,7 +138,44 @@ func (g *generalGraph) compileExtraSource(root llb.State) (llb.State, error) { return llb.Merge(inputs, llb.WithCustomName("[internal] build source layers")), nil } -func (g *generalGraph) preparePythonBase(root llb.State) llb.State { +func (g *generalGraph) compileLanguage(root llb.State) (lang llb.State, err error) { + switch g.Language.Name { + case "python": + lang, err = g.installPython(root) + case "r": + lang, err = g.installRLang(root) + case "julia": + lang, err = g.installJulia(root) + } + + return lang, err +} + +func (g *generalGraph) compileLanguagePackages(root llb.State) (pack llb.State) { + switch g.Language.Name { + case "python": + index := g.compilePyPIIndex(root) + pypi := g.compilePyPIPackages(index) + if g.CondaConfig == nil { + pack = pypi + } else { + channel := g.compileCondaChannel(root) + conda := g.compileCondaPackages(channel) + pack = llb.Merge([]llb.State{ + root, + llb.Diff(root, pypi, llb.WithCustomName("[internal] PyPI packages")), + llb.Diff(root, conda, llb.WithCustomName("[internal] conda packages")), + }, llb.WithCustomName("[internal] Python packages")) + } + case "r": + pack = g.installRPackages(root) + case "julia": + pack = g.installJuliaPackages(root) + } + return pack +} + +func (g *generalGraph) compileDevPackages(root llb.State) llb.State { for _, env := range types.BaseEnvironment { root = root.AddEnv(env.Name, env.Value) } @@ -175,9 +204,14 @@ func (g generalGraph) compileSSHD(root llb.State) llb.State { return sshd } -func (g *generalGraph) compileBase() (llb.State, error) { +func (g *generalGraph) compileBaseImage() (llb.State, error) { + // TODO: find another way to install CUDA + if g.CUDA != nil { + g.Image = GetCUDAImage(g.Image, g.CUDA, g.CUDNN, g.Dev) + } + logger := logrus.WithFields(logrus.Fields{ - "os": g.OS, + "image": g.Image, "language": g.Language.Name, }) if g.Language.Version != nil { @@ -185,61 +219,9 @@ func (g *generalGraph) compileBase() (llb.State, error) { } logger.Debug("compile base image") - var base llb.State - org := viper.GetString(flag.FlagDockerOrganization) - v := version.GetVersionForImageTag() - // Do not update user permission in the base image. - if g.Image != nil { - return g.customBase() - } else if g.CUDA == nil { - switch g.Language.Name { - case "r": - base = llb.Image(fmt.Sprintf("docker.io/%s/r-base:4.2-envd-%s", org, v)) - // 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 - } - case "python": - // TODO(keming) use user input `base(os="")` - base = g.preparePythonBase(llb.Image(types.PythonBaseImage)) - case "julia": - base = llb.Image(fmt.Sprintf( - "docker.io/%s/julia:1.8rc1-ubuntu20.04-envd-%s", org, v)) - } - } else { - base = g.compileCUDAPackages("nvidia/cuda") - } - - // Install conda first. - condaStage, err := g.installConda(base) - if err != nil { - return llb.State{}, errors.Wrap(err, "failed to install conda") - } - supervisor := g.installHorust(condaStage) - sshdStage := g.compileSSHD(supervisor) - source, err := g.compileExtraSource(sshdStage) - if err != nil { - return llb.State{}, errors.Wrap(err, "failed to get extra sources") - } - final := g.compileUserGroup(source) - return final, nil -} - -// customBase get the image and the set the image metadata to graph. -func (g *generalGraph) customBase() (llb.State, error) { - if g.Image == nil { - return llb.State{}, fmt.Errorf("failed to get the image") - } - logrus.WithField("image", *g.Image).Debugf("using custom base image") - // Fix https://github.com/tensorchord/envd/issues/1147. // Fetch the image metadata from base image. - base := llb.Image(*g.Image, - llb.WithMetaResolver(imagemetaresolver.Default())) + base := llb.Image(g.Image, llb.WithMetaResolver(imagemetaresolver.Default())) envs, err := base.Env(context.Background()) if err != nil { return llb.State{}, errors.Wrap(err, "failed to get the image metadata") @@ -250,6 +232,8 @@ func (g *generalGraph) customBase() (llb.State, error) { kv := strings.Split(e, "=") g.RuntimeEnviron[kv[0]] = kv[1] } + // TODO: inherit the USER from base + g.User = "" return base, nil } @@ -272,7 +256,7 @@ func (g generalGraph) copySSHKey(root llb.State) (llb.State, error) { func (g generalGraph) compileMountDir(root llb.State) llb.State { mount := root - if g.Image == nil { + if g.Dev { // create the ENVD_WORKDIR as a placeholder (envd-server may not mount this dir) workDir := fileutil.EnvdHomeDir(g.EnvironmentName) mount = root.File(llb.Mkdir(workDir, 0755, llb.WithParents(true), llb.WithUIDGID(g.uid, g.gid)), diff --git a/pkg/lang/ir/v1/types.go b/pkg/lang/ir/v1/types.go index 9883e374f..407316532 100644 --- a/pkg/lang/ir/v1/types.go +++ b/pkg/lang/ir/v1/types.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 v1 import ( @@ -14,11 +28,12 @@ type generalGraph struct { uid int gid int - OS string ir.Language - Image *string + Image string + User string Shell string + Dev bool CUDA *string CUDNN string NumGPUs int diff --git a/pkg/lang/ir/v1/user.go b/pkg/lang/ir/v1/user.go index 4820504a1..8ea04fce4 100644 --- a/pkg/lang/ir/v1/user.go +++ b/pkg/lang/ir/v1/user.go @@ -22,27 +22,32 @@ import ( // compileUserOwn chown related directories func (g *generalGraph) compileUserOwn(root llb.State) llb.State { - if g.Image != nil || g.uid == 0 { + if g.uid == 0 { g.RuntimeEnviron["USER"] = "root" return root } g.RuntimeEnviron["USER"] = "envd" - if len(g.UserDirectories) == 0 { - return root.User("envd") - } - run := root.Run() + g.User = "envd" for _, dir := range g.UserDirectories { - run = root.Run(llb.Shlex(fmt.Sprintf("chown -R envd:envd %s", dir)), - llb.WithCustomNamef("[internal] configure user permissions for %s", dir)) + root = root.Run(llb.Shlex(fmt.Sprintf("chown -R envd:envd %s", dir)), + llb.WithCustomNamef("[internal] configure user permissions for %s", dir)).Root() } - return run.Root().User("envd") + return root.User("envd") } // compileUserGroup creates user `envd` func (g *generalGraph) compileUserGroup(root llb.State) llb.State { - if g.Image != nil { - return root + if g.Language.Name == "r" { + // 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 + } } + var res llb.ExecState if g.uid == 0 { res = root. @@ -61,13 +66,13 @@ func (g *generalGraph) compileUserGroup(root llb.State) llb.State { } else { res = root. Run(llb.Shlex(fmt.Sprintf("groupadd -g %d envd", g.gid)), - llb.WithCustomName("[internal] create user group envd")). + llb.WithCustomNamef("[internal] create user group envd(g:%d)", g.gid)). Run(llb.Shlex(fmt.Sprintf("useradd -p \"\" -u %d -g envd -s /bin/sh -m envd", g.uid)), - llb.WithCustomName("[internal] create user envd")). - Run(llb.Shlex("adduser envd sudo"), + llb.WithCustomNamef("[internal] create user envd(u:%d)", g.uid)). + Run(llb.Shlex("usermod -a -G sudo envd"), llb.WithCustomName("[internal] add user envd to sudoers")). Run(llb.Shlex(fmt.Sprintf("install -d -o envd -g %d -m 0700 /home/envd/.config /home/envd/.cache", g.gid)), - llb.WithCustomName("[internal] mkdir config and cache dir")) + llb.WithCustomName("[internal] mkdir config and cache")) } return res.Root() } diff --git a/pkg/lang/ir/v1/util.go b/pkg/lang/ir/v1/util.go index 2cac34f56..8465a380d 100644 --- a/pkg/lang/ir/v1/util.go +++ b/pkg/lang/ir/v1/util.go @@ -19,6 +19,7 @@ import ( "crypto/md5" "encoding/gob" "encoding/hex" + "fmt" "os/user" "regexp" "strconv" @@ -116,3 +117,15 @@ func GetDefaultGraphHash() string { hashD := md5.Sum(data) return hex.EncodeToString(hashD[:]) } + +// GetCUDAImage finds the correct CUDA base image +// refer to https://hub.docker.com/r/nvidia/cuda/tags +func GetCUDAImage(image string, cuda *string, cudnn string, dev bool) string { + // TODO: support CUDA 10 + target := "runtime" + if dev { + target = "devel" + } + + return fmt.Sprintf("docker.io/nvidia:%s-cudnn%s-%s-%s", *cuda, cudnn, target, image) +} diff --git a/pkg/util/fileutil/file.go b/pkg/util/fileutil/file.go index 721b78d74..43e635a46 100644 --- a/pkg/util/fileutil/file.go +++ b/pkg/util/fileutil/file.go @@ -182,6 +182,12 @@ func DownloadOrUpdateGitRepo(url string) (path string, err error) { return path, nil } +// EnvdHomeDir returns the envd user path inside the environment func EnvdHomeDir(path ...string) string { return filepath.Join(append([]string{"/", "home", "envd"}, path...)...) } + +// DefaultHomeDir returns the default user path inside the environment +func DefaultHomeDir(path ...string) string { + return filepath.Join(append([]string{"~"}, path...)...) +}