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...)...)
+}