bug: fix conda install with env file (#837)

* fix conda install env file Signed-off-by: Jinjing.Zhou <allenzhou@tensorchord.ai> * fix Signed-off-by: Jinjing.Zhou <allenzhou@tensorchord.ai> * fix Signed-off-by: Jinjing.Zhou <allenzhou@tensorchord.ai> * revert change Signed-off-by: Jinjing.Zhou <allenzhou@tensorchord.ai> * lint Signed-off-by: Jinjing.Zhou <allenzhou@tensorchord.ai> * add root condarc Signed-off-by: Jinjing.Zhou <allenzhou@tensorchord.ai> * fix Signed-off-by: Jinjing.Zhou <allenzhou@tensorchord.ai> * fix Signed-off-by: Jinjing.Zhou <allenzhou@tensorchord.ai> * fix Signed-off-by: Jinjing.Zhou <allenzhou@tensorchord.ai> * fix Signed-off-by: Jinjing.Zhou <allenzhou@tensorchord.ai> Signed-off-by: Jinjing.Zhou <allenzhou@tensorchord.ai>
tensorchord · Sep 8, 2022 · 1e5e24d · 1e5e24d
1 parent ecb9e26
commit 1e5e24d
Show file tree

Hide file tree

Showing 11 changed files with 111 additions and 103 deletions.
diff --git a/e2e/language/testdata/python/conda/env.yaml b/e2e/language/testdata/python/conda/env.yaml
@@ -1,6 +1,9 @@
+name: imarobot # should ignore this name
 channels:
-  - pytorch
+  - conda-forge
   - defaults
 dependencies:
-  - pytorch=1.11.0
-  - cpuonly
+  - babel
+  - pip:
+    - via
+    - -e .
diff --git a/e2e/language/testdata/python/conda/setup.py b/e2e/language/testdata/python/conda/setup.py
@@ -0,0 +1,11 @@
+from setuptools import setup, find_packages
+
+setup(
+    name='envd_e2e_conda_test',
+    version='1.0.0',
+    url='https://github.com/mypackage.git',
+    author='Author Name',
+    author_email='author@gmail.com',
+    description='Description of my package',
+    packages=find_packages(),
+)
diff --git a/e2e/language/testdata/python/requirements/requirements.txt b/e2e/language/testdata/python/requirements/requirements.txt
@@ -1 +1,2 @@
 via
+-e .
diff --git a/e2e/language/testdata/python/requirements/setup.py b/e2e/language/testdata/python/requirements/setup.py
@@ -0,0 +1,11 @@
+from setuptools import setup, find_packages
+
+setup(
+    name='envd_e2e_conda_test',
+    version='1.0.0',
+    url='https://github.com/mypackage.git',
+    author='Author Name',
+    author_email='author@gmail.com',
+    description='Description of my package',
+    packages=find_packages(),
+)
diff --git a/pkg/lang/frontend/starlark/install/install.go b/pkg/lang/frontend/starlark/install/install.go
@@ -15,13 +15,11 @@
 package install
 
 import (
-	"path/filepath"
-
+	"github.com/cockroachdb/errors"
 	"github.com/sirupsen/logrus"
 	"go.starlark.net/starlark"
 	"go.starlark.net/starlarkstruct"
 
-	"github.com/tensorchord/envd/pkg/lang/frontend/starlark/builtin"
 	"github.com/tensorchord/envd/pkg/lang/ir"
 )
 
@@ -208,17 +206,16 @@ func ruleFuncConda(thread *starlark.Thread, _ *starlark.Builtin,
 			channelList = append(channelList, channel.Index(i).(starlark.String).GoString())
 		}
 	}
-	var path *string = nil
 	envFileStr := envFile.GoString()
 	if envFileStr != "" {
-		buildContextDir := starlark.Universe[builtin.BuildContextDir]
-		buildContextDirStr := buildContextDir.(starlark.String).GoString()
-		buf := filepath.Join(buildContextDirStr, envFileStr)
-		path = &buf
+
+		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, path); err != nil {
+	if err := ir.CondaPackage(nameList, channelList, envFileStr); err != nil {
 		return starlark.None, err
 	}
 

diff --git a/pkg/lang/ir/conda.go b/pkg/lang/ir/conda.go
@@ -22,6 +22,8 @@ import (
 	"github.com/cockroachdb/errors"
 	"github.com/moby/buildkit/client/llb"
 	"github.com/sirupsen/logrus"
+
+	"github.com/tensorchord/envd/pkg/flag"
 )
 
 const (
@@ -35,7 +37,14 @@ var (
 )
 
 func (g Graph) CondaEnabled() bool {
-	return g.CondaConfig != nil
+	if g.CondaConfig == nil {
+		return false
+	} else {
+		if g.CondaConfig.CondaPackages == nil && len(g.CondaConfig.CondaEnvFileName) == 0 {
+			return false
+		}
+	}
+	return true
 }
 
 func (g Graph) compileCondaChannel(root llb.State) llb.State {
@@ -50,7 +59,8 @@ func (g Graph) compileCondaChannel(root llb.State) llb.State {
 }
 
 func (g Graph) compileCondaPackages(root llb.State) llb.State {
-	if !g.CondaEnabled() || len(g.CondaConfig.CondaPackages) == 0 {
+	if !g.CondaEnabled() {
+		logrus.Debug("Conda packages not enabled")
 		return root
 	}
 
@@ -60,30 +70,57 @@ func (g Graph) compileCondaPackages(root llb.State) llb.State {
 		0755, llb.WithParents(true), llb.WithUIDGID(g.uid, g.gid)),
 		llb.WithCustomName("[internal] setting conda cache mount permissions"))
 
+	root = llb.User("envd")(root)
 	// Compose the package install command.
 	var sb strings.Builder
-	if len(g.CondaConfig.AdditionalChannels) == 0 {
-		sb.WriteString("/opt/conda/bin/conda install -n envd")
+	var run llb.ExecState
+	if len(g.CondaConfig.CondaEnvFileName) != 0 {
+		logrus.Debugf("using custom conda environment file content: %s", g.CondaConfig.CondaEnvFileName)
+		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/bin/conda env update -n envd --file %s\n", g.CondaConfig.CondaEnvFileName))
+
+		// Execute the command to write yaml file and conda env using envd user
+		sb.WriteString(fmt.Sprintf("sudo -i -u envd bash << EOF\nset -euo pipefail\n%s\nEOF\n", envdCmd.String()))
+		sb.WriteString("'")
+		cmd := sb.String()
+
+		run = root.User("root").
+			Dir(g.getWorkingDir()).
+			Run(llb.Shlex(cmd),
+				llb.WithCustomNamef("conda install from file %s", g.CondaConfig.CondaEnvFileName))
+
+		run.AddMount(cacheDir, cache,
+			llb.AsPersistentCacheDir(g.CacheID(cacheDir), llb.CacheMountShared),
+			llb.SourcePath("/cache-conda"))
+		run.AddMount(g.getWorkingDir(),
+			llb.Local(flag.FlagBuildContext))
 
 	} else {
-		sb.WriteString("/opt/conda/bin/conda install -n envd")
-		for _, channel := range g.CondaConfig.AdditionalChannels {
-			sb.WriteString(fmt.Sprintf(" -c %s", channel))
+		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")
+			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))
-	}
+		for _, pkg := range g.CondaConfig.CondaPackages {
+			sb.WriteString(fmt.Sprintf(" %s", pkg))
+		}
 
-	cmd := sb.String()
-	root = llb.User("envd")(root)
+		cmd := sb.String()
 
-	run := root.
-		Run(llb.Shlex(cmd), llb.WithCustomNamef("conda install %s",
-			strings.Join(g.CondaPackages, " ")))
-	run.AddMount(cacheDir, cache,
-		llb.AsPersistentCacheDir(g.CacheID(cacheDir), llb.CacheMountShared), llb.SourcePath("/cache-conda"))
+		run = root.
+			Run(llb.Shlex(cmd), llb.WithCustomNamef("conda install %s",
+				strings.Join(g.CondaPackages, " ")))
+		run.AddMount(cacheDir, cache,
+			llb.AsPersistentCacheDir(g.CacheID(cacheDir), llb.CacheMountShared), llb.SourcePath("/cache-conda"))
+	}
 	return run.Root()
 }
 

diff --git a/pkg/lang/ir/consts.go b/pkg/lang/ir/consts.go
@@ -30,5 +30,8 @@ const (
 [global]
 index-url=%s
 %s
+
+[install]
+src = /tmp
 `
 )
diff --git a/pkg/lang/ir/interface.go b/pkg/lang/ir/interface.go
@@ -18,7 +18,6 @@ import (
 	"github.com/cockroachdb/errors"
 
 	"github.com/tensorchord/envd/pkg/editor/vscode"
-	"github.com/tensorchord/envd/pkg/lang/ir/parser"
 )
 
 func Base(os, language, image string) error {
@@ -156,23 +155,15 @@ func CondaChannel(channel string) error {
 	return nil
 }
 
-func CondaPackage(deps []string, channel []string, envFile *string) error {
+func CondaPackage(deps []string, channel []string, envFile string) error {
 	if !DefaultGraph.CondaEnabled() {
 		DefaultGraph.CondaConfig = &CondaConfig{}
 	}
 	DefaultGraph.CondaConfig.CondaPackages = append(
 		DefaultGraph.CondaConfig.CondaPackages, deps...)
 
-	if envFile != nil {
-		parsed, err := parser.ParseCondaEnvYaml(*envFile)
-		if err != nil {
-			return err
-		}
-		DefaultGraph.CondaConfig.CondaPackages = append(DefaultGraph.CondaConfig.CondaPackages, parsed.CondaPackages...)
-		DefaultGraph.PyPIPackages = append(DefaultGraph.PyPIPackages, parsed.PipPackages...)
-		DefaultGraph.CondaConfig.AdditionalChannels = append(
-			DefaultGraph.CondaConfig.AdditionalChannels, parsed.Channels...)
-	}
+	DefaultGraph.CondaConfig.CondaEnvFileName = envFile
+
 	if len(channel) != 0 {
 		DefaultGraph.CondaConfig.AdditionalChannels = append(
 			DefaultGraph.CondaConfig.AdditionalChannels, channel...)

diff --git a/pkg/lang/ir/parser/conda.go b/pkg/lang/ir/parser/conda.go
diff --git a/pkg/lang/ir/python.go b/pkg/lang/ir/python.go
@@ -139,7 +139,7 @@ func (g Graph) compilePyPIPackages(root llb.State) llb.State {
 		logrus.WithField("command", cmd).
 			Debug("Configure pip install statements")
 		root = llb.User("envd")(root)
-		run := root.Dir("/tmp").
+		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
@@ -151,18 +151,27 @@ func (g Graph) compilePyPIPackages(root llb.State) llb.State {
 	if g.RequirementsFile != nil {
 		// Compose the package install command.
 		var sb strings.Builder
-		sb.WriteString("/opt/conda/envs/envd/bin/python -m pip install -r ")
-		sb.WriteString(*g.RequirementsFile)
+		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.Dir(g.getWorkingDir())
+		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"))
 		run.AddMount(g.getWorkingDir(),
-			llb.Local(flag.FlagBuildContext), llb.Readonly)
+			llb.Local(flag.FlagBuildContext))
 		root = run.Root()
 	}
 

diff --git a/pkg/lang/ir/types.go b/pkg/lang/ir/types.go
@@ -101,6 +101,7 @@ type CondaConfig struct {
 	CondaPackages      []string
 	AdditionalChannels []string
 	CondaChannel       *string
+	CondaEnvFileName   string
 }
 
 type GitConfig struct {