feat: Suport conda env

examples/conda/build.envd

@@ -0,0 +1,27 @@
+def build():
+    config.conda_channel(channel="""
+channels:
+  - defaults
+show_channel_urls: true
+default_channels:
+  - https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main
+  - https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/r
+  - https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/msys2
+  - https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/free
+custom_channels:
+  conda-forge: https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud
+  msys2: https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud
+  bioconda: https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud
+  menpo: https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud
+  pytorch: https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud
+  pytorch-lts: https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud
+  simpleitk: https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud
+""")
+    install.conda_packages(name = [
+        "pytorch=0.4.1", "cuda90"
+    ], channel= ["pytorch"])
+    base(os="ubuntu20.04", language="python3.5")
+    install.python_packages(name = [
+        "flask"
+    ])
+    install.cuda(version="11.6", cudnn="8")

pkg/app/destroy.go

@@ -71,7 +71,7 @@ func destroy(clicontext *cli.Context) error {
 	}
 	if ctrName, err := dockerClient.Destroy(clicontext.Context, ctrName); err != nil {
 		return errors.Wrapf(err, "failed to destroy the environment: %s", ctrName)
-	} else if name != "" {
+	} else if ctrName != "" {
 		logrus.Infof("%s is destroyed", ctrName)
 	}
 

pkg/editor/jupyter/util.go

@@ -27,7 +27,7 @@ func GenerateCommand(g ir.Graph, notebookDir string) []string {
 
 	var cmd []string
 	// Use python in conda env.
-	if len(g.CondaPackages) != 0 {
+	if g.CondaEnabled() {
 		cmd = append(cmd, "/opt/conda/bin/python3")
 	} else {
 		cmd = append(cmd, "python3")

pkg/lang/frontend/starlark/install/install.go

@@ -158,10 +158,10 @@ func ruleFuncVSCode(thread *starlark.Thread, _ *starlark.Builtin,
 
 func ruleFuncConda(thread *starlark.Thread, _ *starlark.Builtin,
 	args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
-	var name *starlark.List
+	var name, channel *starlark.List
 
 	if err := starlark.UnpackArgs(ruleConda,
-		args, kwargs, "name", &name); err != nil {
+		args, kwargs, "name", &name, "channel?", &channel); err != nil {
 		return nil, err
 	}
 
@@ -172,8 +172,15 @@ func ruleFuncConda(thread *starlark.Thread, _ *starlark.Builtin,
 		}
 	}
 
-	logger.Debugf("rule `%s` is invoked, name=%v", ruleConda, nameList)
-	ir.CondaPackage(nameList)
+	channelList := []string{}
+	if channel != nil {
+		for i := 0; i < channel.Len(); i++ {
+			channelList = append(channelList, channel.Index(i).(starlark.String).GoString())
+		}
+	}
+
+	logger.Debugf("rule `%s` is invoked, name=%v, channel=%v", ruleConda, nameList, channelList)
+	ir.CondaPackage(nameList, channelList)
 
 	return starlark.None, nil
 }

pkg/lang/frontend/starlark/universe/universe.go

@@ -52,7 +52,10 @@ func ruleFuncBase(thread *starlark.Thread, _ *starlark.Builtin,
 
 	logger.Debugf("rule `%s` is invoked, os=%s, language=%s", ruleBase,
 		osStr, langStr)
-	ir.Base(osStr, langStr)
+	err := ir.Base(osStr, langStr)
+	if err != nil {
+		return starlark.None, err
+	}
 
 	return starlark.None, nil
 }

pkg/lang/ir/compile.go

@@ -28,10 +28,12 @@ import (
 
 func NewGraph() *Graph {
 	return &Graph{
-		OS:       osDefault,
-		Language: languageDefault,
-		CUDA:     nil,
-		CUDNN:    nil,
+		OS: osDefault,
+		Language: Language{
+			Name: languageDefault,
+		},
+		CUDA:  nil,
+		CUDNN: nil,
 
 		PyPIPackages:   []string{},
 		RPackages:      []string{},
@@ -109,7 +111,7 @@ func (g Graph) Compile() (llb.State, error) {
 	base := g.compileBase()
 	aptStage := g.compileUbuntuAPT(base)
 	var merged llb.State
-	if g.Language == "r" {
+	if g.Language.Name == "r" {
 		// TODO(terrytangyuan): Support RStudio local server
 		rPackageInstallStage := llb.Diff(aptStage, g.installRPackages(aptStage), llb.WithCustomName("install R packages"))
 		merged = llb.Merge([]llb.State{
@@ -133,12 +135,15 @@ func (g Graph) Compile() (llb.State, error) {
 		if err != nil {
 			return llb.State{}, errors.Wrap(err, "failed to compile shell")
 		}
+
+		condaEnvStage := g.setCondaENV(shellStage)
+
 		condaStage := llb.Diff(builtinSystemStage,
-			g.compileCondaPackages(shellStage),
-			llb.WithCustomName("install PyPI packages"))
+			g.compileCondaPackages(condaEnvStage),
+			llb.WithCustomName("install conda packages"))
 
-		pypiStage := llb.Diff(builtinSystemStage,
-			g.compilePyPIPackages(builtinSystemStage),
+		pypiStage := llb.Diff(condaEnvStage,
+			g.compilePyPIPackages(condaEnvStage),
 			llb.WithCustomName("install PyPI packages"))
 		systemStage := llb.Diff(builtinSystemStage, g.compileSystemPackages(builtinSystemStage),
 			llb.WithCustomName("install system packages"))

pkg/lang/ir/conda.go

@@ -26,8 +26,12 @@ const (
 	condarc = "/home/envd/.condarc"
 )
 
+func (g Graph) CondaEnabled() bool {
+	return g.CondaConfig != nil
+}
+
 func (g Graph) compileCondaChannel(root llb.State) llb.State {
-	if g.CondaChannel != nil {
+	if g.CondaConfig != nil && g.CondaConfig.CondaChannel != nil {
 		logrus.WithField("conda-channel", *g.CondaChannel).Debug("using custom connda channel")
 		stage := root.
 			File(llb.Mkfile(condarc,
@@ -38,16 +42,25 @@ func (g Graph) compileCondaChannel(root llb.State) llb.State {
 }
 
 func (g Graph) compileCondaPackages(root llb.State) llb.State {
-	if len(g.CondaPackages) == 0 {
+	if !g.CondaEnabled() || len(g.CondaConfig.CondaPackages) == 0 {
 		return root
 	}
 
 	cacheDir := "/opt/conda/pkgs"
 
 	// Compose the package install command.
 	var sb strings.Builder
-	sb.WriteString("/opt/conda/bin/conda install")
-	for _, pkg := range g.CondaPackages {
+	if len(g.CondaConfig.AdditionalChannels) == 0 {
+		sb.WriteString("/opt/conda/bin/conda install -n envd")
+
+	} else {
+		sb.WriteString("/opt/conda/bin/conda install -n envd -c")
+		for _, channel := range g.CondaConfig.AdditionalChannels {
+			sb.WriteString(fmt.Sprintf(" %s", channel))
+		}
+	}
+
+	for _, pkg := range g.CondaConfig.CondaPackages {
 		sb.WriteString(fmt.Sprintf(" %s", pkg))
 	}
 
@@ -62,16 +75,38 @@ func (g Graph) compileCondaPackages(root llb.State) llb.State {
 			strings.Join(g.CondaPackages, " ")))
 	run.AddMount(cacheDir, cache,
 		llb.AsPersistentCacheDir(g.CacheID(cacheDir), llb.CacheMountShared), llb.SourcePath("/cache"))
-	return g.setCondaENV(run.Root())
+	return run.Root()
 }
 
 func (g Graph) setCondaENV(root llb.State) llb.State {
+	if !g.CondaEnabled() {
+		return root
+	}
+
 	root = llb.User("envd")(root)
 	// Always init bash since we will use it to create jupyter notebook service.
-	run := root.Run(llb.Shlex("/opt/conda/bin/conda init bash"), llb.WithCustomName("[internal] initialize conda bash environment"))
-	if g.Shell != shellBASH {
-		run = run.Run(llb.Shlex(fmt.Sprintf("/opt/conda/bin/conda init %s", g.Shell)),
-			llb.WithCustomNamef("[internal] initialize conda %s environment", g.Shell))
+	run := root.Run(llb.Shlex("bash -c \"/opt/conda/bin/conda init bash\""), llb.WithCustomName("[internal] initialize conda bash environment"))
+
+	pythonVersion := "3.9"
+	if g.Language.Version != nil {
+		pythonVersion = *g.Language.Version
+	}
+	cmd := fmt.Sprintf(
+		"bash -c \"/opt/conda/bin/conda create -n envd python=%s\"", pythonVersion)
+	// Create a conda environment.
+	run = run.Run(llb.Shlex(cmd), llb.WithCustomName("[internal] create conda environment"))
+
+	switch g.Shell {
+	case shellBASH:
+		run = run.Run(
+			llb.Shlex(`bash -c 'echo "source /opt/conda/bin/activate envd" >> /home/envd/.bashrc'`),
+			llb.WithCustomName("[internal] add conda environment to bashrc"))
+	case shellZSH:
+		run = run.Run(
+			llb.Shlex(fmt.Sprintf("bash -c \"/opt/conda/bin/conda init %s\"", g.Shell)),
+			llb.WithCustomNamef("[internal] initialize conda %s environment", g.Shell)).Run(
+			llb.Shlex(`bash -c 'echo "source /opt/conda/bin/activate envd" >> /home/envd/.zshrc'`),
+			llb.WithCustomName("[internal] add conda environment to zshrc"))
 	}
 	return run.Root()
 }

pkg/lang/ir/consts.go

@@ -16,7 +16,7 @@ package ir
 
 const (
 	osDefault         = "ubuntu20.04"
-	languageDefault   = "python3.8"
+	languageDefault   = "python"
 	pypiIndexModeAuto = "auto"
 
 	aptSourceFilePath = "/etc/apt/sources.list"

pkg/lang/ir/interface.go

@@ -20,9 +20,17 @@ import (
 	"github.com/tensorchord/envd/pkg/editor/vscode"
 )
 
-func Base(os, language string) {
-	DefaultGraph.Language = language
+func Base(os, language string) error {
+	l, version, err := parseLanguage(language)
+	if err != nil {
+		return err
+	}
+	DefaultGraph.Language = Language{
+		Name:    l,
+		Version: version,
+	}
 	DefaultGraph.OS = os
+	return nil
 }
 
 func PyPIPackage(deps []string) {
@@ -116,10 +124,23 @@ func CondaChannel(channel string) error {
 		return errors.New("channel is required")
 	}
 
-	DefaultGraph.CondaChannel = &channel
+	if !DefaultGraph.CondaEnabled() {
+		DefaultGraph.CondaConfig = &CondaConfig{}
+	}
+
+	DefaultGraph.CondaConfig.CondaChannel = &channel
 	return nil
 }
 
-func CondaPackage(deps []string) {
-	DefaultGraph.CondaPackages = append(DefaultGraph.CondaPackages, deps...)
+func CondaPackage(deps []string, channel []string) {
+	if !DefaultGraph.CondaEnabled() {
+		DefaultGraph.CondaConfig = &CondaConfig{}
+	}
+	DefaultGraph.CondaConfig.CondaPackages = append(
+		DefaultGraph.CondaConfig.CondaPackages, deps...)
+
+	if len(channel) != 0 {
+		DefaultGraph.CondaConfig.AdditionalChannels = append(
+			DefaultGraph.CondaConfig.AdditionalChannels, channel...)
+	}
 }

pkg/lang/ir/python.go

@@ -32,8 +32,8 @@ func (g Graph) compilePyPIPackages(root llb.State) llb.State {
 
 	// Compose the package install command.
 	var sb strings.Builder
-	if len(g.CondaPackages) != 0 {
-		sb.WriteString("/opt/conda/bin/pip install --no-warn-script-location")
+	if g.CondaEnabled() {
+		sb.WriteString("/opt/conda/bin/conda run -n envd pip install")
 	} else {
 		sb.WriteString("pip install --no-warn-script-location")
 	}

pkg/lang/ir/system.go

@@ -86,10 +86,19 @@ func (g Graph) compileSystemPackages(root llb.State) llb.State {
 }
 
 func (g *Graph) compileBase() llb.State {
+	logger := logrus.WithFields(logrus.Fields{
+		"os":       g.OS,
+		"language": g.Language.Name,
+	})
+	if g.Language.Version != nil {
+		logger = logger.WithField("version", *g.Language.Version)
+	}
+	logger.Debug("compile base image")
+
 	var base llb.State
 	var groupID string = "1000"
 	if g.CUDA == nil && g.CUDNN == nil {
-		if g.Language == "r" {
+		if g.Language.Name == "r" {
 			base = llb.Image("docker.io/r-base:4.2.0")
 			// r-base image already has GID 1000.
 			groupID = "1001"

pkg/lang/ir/types.go

@@ -22,34 +22,45 @@ import (
 // A Graph contains the state,
 // such as its call stack and thread-local storage.
 type Graph struct {
-	OS       string
-	Language string
-	Shell    string
-	CUDA     *string
-	CUDNN    *string
+	OS string
+	Language
+
+	Shell string
+	CUDA  *string
+	CUDNN *string
 
 	UbuntuAPTSource   *string
 	PyPIIndexURL      *string
 	PyPIExtraIndexURL *string
-	CondaChannel      *string
 
 	PublicKeyPath string
 
 	PyPIPackages   []string
 	RPackages      []string
-	CondaPackages  []string
 	SystemPackages []string
 
 	VSCodePlugins []vscode.Plugin
 
 	Exec []string
 	*JupyterConfig
 	*GitConfig
+	*CondaConfig
 
 	Writer      compileui.Writer
 	CachePrefix string
 }
 
+type Language struct {
+	Name    string
+	Version *string
+}
+
+type CondaConfig struct {
+	CondaPackages      []string
+	AdditionalChannels []string
+	CondaChannel       *string
+}
+
 type GitConfig struct {
 	Name   string
 	Email  string