refact: conda/mamba create/update env, fix user permissions (#933)

* inii mamba Signed-off-by: Keming <kemingyang@tensorchord.ai> * mamba init bash err Signed-off-by: Keming <kemingyang@tensorchord.ai> * change default to conda Signed-off-by: Keming <kemingyang@tensorchord.ai> * fix mamba activate shell Signed-off-by: Keming <kemingyang@tensorchord.ai> * try to use create -f Signed-off-by: Keming <kemingyang@tensorchord.ai> * make conda&mamba happy Signed-off-by: Keming <kemingyang@tensorchord.ai> * fix conda packages Signed-off-by: Keming <kemingyang@tensorchord.ai> * do not create user envd or chown for custom images Signed-off-by: Keming <kemingyang@tensorchord.ai> * Update pkg/lang/ir/conda.go Co-authored-by: Ce Gao <ce.gao@outlook.com> Signed-off-by: Keming <kemingy94@gmail.com> * chore: use conda-forge as default miniconda channel (#929) Signed-off-by: Keming <kemingyang@tensorchord.ai> Signed-off-by: Keming <kemingyang@tensorchord.ai> * fix *Graph Signed-off-by: Keming <kemingyang@tensorchord.ai> * test cache Signed-off-by: Keming <kemingyang@tensorchord.ai> * fix lint Signed-off-by: Keming <kemingyang@tensorchord.ai> * fix permission Signed-off-by: Keming <kemingyang@tensorchord.ai> * fix tests Signed-off-by: Keming <kemingyang@tensorchord.ai> * reduce chown Signed-off-by: Keming <kemingyang@tensorchord.ai> * no need to chown conda Signed-off-by: Keming <kemingyang@tensorchord.ai> Signed-off-by: Keming <kemingyang@tensorchord.ai> Signed-off-by: Keming <kemingy94@gmail.com> Co-authored-by: Ce Gao <ce.gao@outlook.com>
tensorchord · Oct 4, 2022 · bf993e2 · bf993e2
1 parent 0e79fb9
commit bf993e2
Show file tree

Hide file tree

Showing 13 changed files with 134 additions and 144 deletions.
diff --git a/pkg/lang/frontend/starlark/install/install.go b/pkg/lang/frontend/starlark/install/install.go
@@ -193,7 +193,6 @@ func ruleFuncConda(thread *starlark.Thread, _ *starlark.Builtin,
 
 	envFileStr := envFile.GoString()
 	if envFileStr != "" {
-
 		if (len(nameList) != 0) || (len(channelList) != 0) {
 			return nil, errors.New("env_file and name/channel are mutually exclusive")
 		}

diff --git a/pkg/lang/ir/compile.go b/pkg/lang/ir/compile.go
@@ -41,6 +41,7 @@ func NewGraph() *Graph {
 		RuntimeEnviron:  make(map[string]string),
 	}
 	langVersion := languageVersionDefault
+	conda := &CondaConfig{}
 	return &Graph{
 		OS: osDefault,
 		Language: Language{
@@ -51,13 +52,15 @@ func NewGraph() *Graph {
 		CUDNN:   "8", // default version
 		NumGPUs: -1,
 
-		PyPIPackages:   []string{},
-		RPackages:      []string{},
-		JuliaPackages:  []string{},
-		SystemPackages: []string{},
-		Exec:           []string{},
-		Shell:          shellBASH,
-		RuntimeGraph:   runtimeGraph,
+		PyPIPackages:    []string{},
+		RPackages:       []string{},
+		JuliaPackages:   []string{},
+		SystemPackages:  []string{},
+		Exec:            []string{},
+		UserDirectories: []string{},
+		Shell:           shellBASH,
+		CondaConfig:     conda,
+		RuntimeGraph:    runtimeGraph,
 	}
 }
 
@@ -293,10 +296,8 @@ func (g Graph) Compile(uid, gid int) (llb.State, error) {
 	copy := g.compileCopy(prompt)
 	// TODO(gaocegege): Support order-based exec.
 	run := g.compileRun(copy)
-	finalStage, err := g.compileGit(run)
-	if err != nil {
-		return llb.State{}, errors.Wrap(err, "failed to compile git")
-	}
+	git := g.compileGit(run)
+	finalStage := g.compileUserOwn(git)
 	g.Writer.Finish()
 	return finalStage, nil
 }
diff --git a/pkg/lang/ir/conda.go b/pkg/lang/ir/conda.go
@@ -30,6 +30,8 @@ import (
 
 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"
 )
@@ -44,19 +46,8 @@ var (
 	installMambaBash string
 )
 
-func (g Graph) CondaEnabled() bool {
-	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 {
-	if g.CondaConfig != nil && g.CondaConfig.CondaChannel != nil {
+	if g.CondaConfig.CondaChannel != nil {
 		logrus.WithField("conda-channel", *g.CondaChannel).Debug("using custom conda channel")
 		stage := root.
 			File(llb.Mkfile(condarc,
@@ -66,69 +57,46 @@ func (g Graph) compileCondaChannel(root llb.State) llb.State {
 	return root
 }
 
-func (g Graph) microMambaEnabled() bool {
-	if g.CondaConfig != nil && g.CondaConfig.UseMicroMamba {
-		return true
-	}
-	return false
-}
-
 func (g Graph) condaCommandPath() string {
-	if g.microMambaEnabled() {
+	if g.CondaConfig.UseMicroMamba {
 		return filepath.Join(condaBinDir, "micromamba")
 	}
 	return filepath.Join(condaBinDir, "conda")
 }
 
 func (g Graph) condaInitShell(shell string) string {
 	path := g.condaCommandPath()
-	if g.microMambaEnabled() {
+	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 Graph) compileCondaPackages(root llb.State) llb.State {
-	if !g.CondaEnabled() {
+func (g Graph) 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 *Graph) 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
-	cache := root.File(llb.Mkdir("/cache-conda",
-		0755, llb.WithParents(true), llb.WithUIDGID(g.uid, g.gid)),
+	cacheMount := root.File(llb.Mkdir("/cache-conda", 0755, llb.WithParents(true)),
 		llb.WithCustomName("[internal] setting conda cache mount permissions"))
 
-	root = llb.User("envd")(root)
 	// Compose the package install command.
 	var sb strings.Builder
 	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("%s env update -n envd --file %s\n", g.condaCommandPath(), 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))
 
+	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()))
@@ -138,25 +106,22 @@ func (g Graph) compileCondaPackages(root llb.State) llb.State {
 				sb.WriteString(fmt.Sprintf(" -c %s", channel))
 			}
 		}
-
 		for _, pkg := range g.CondaConfig.CondaPackages {
 			sb.WriteString(fmt.Sprintf(" %s", pkg))
 		}
-
-		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"))
 	}
+
+	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 Graph) compileCondaEnvironment(root llb.State) (llb.State, error) {
-	root = llb.User("envd")(root)
-
 	// 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"))),
@@ -167,14 +132,10 @@ func (g Graph) compileCondaEnvironment(root llb.State) (llb.State, error) {
 	if err != nil {
 		return llb.State{}, errors.Wrap(err, "failed to get python version")
 	}
-
-	cmd := fmt.Sprintf(
-		"bash -c \"%s create -n envd python=%s\"",
-		g.condaCommandPath(), pythonVersion)
-
 	// Create a conda environment.
-	run = run.Run(llb.Shlex(cmd),
-		llb.WithCustomName("[internal] create conda environment"))
+	cmd := fmt.Sprintf("bash -c \"%s create -n envd python=%s\"", g.condaCommandPath(), pythonVersion)
+	run = run.Dir(g.getWorkingDir()).Run(llb.Shlex(cmd),
+		llb.WithCustomNamef("[internal] create conda environment: %s", cmd))
 
 	switch g.Shell {
 	case shellBASH:
@@ -194,9 +155,10 @@ func (g Graph) compileCondaEnvironment(root llb.State) (llb.State, error) {
 }
 
 func (g Graph) installConda(root llb.State) (llb.State, error) {
-	if g.microMambaEnabled() {
+	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

diff --git a/pkg/lang/ir/fs.go b/pkg/lang/ir/fs.go
@@ -17,8 +17,8 @@ import (
 	"github.com/moby/buildkit/client/llb"
 )
 
-func (g Graph) CompileCacheDir(root llb.State, cacheDir string) llb.State {
-	root = llb.User("envd")(root)
+func (g *Graph) 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/git.go b/pkg/lang/ir/git.go
@@ -33,13 +33,13 @@ const (
 `
 )
 
-func (g *Graph) compileGit(root llb.State) (llb.State, error) {
+func (g *Graph) compileGit(root llb.State) llb.State {
 	if g.GitConfig == nil {
-		return root, 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, nil
+	return gitStage
 }
diff --git a/pkg/lang/ir/install-mamba.sh b/pkg/lang/ir/install-mamba.sh
@@ -8,7 +8,8 @@ else \
     ARCH="64"; \
 fi && \
 mkdir -p ${MAMBA_BIN_DIR} && \
-curl -Ls https://micro.mamba.pm/api/micromamba/linux-${ARCH}/latest | tar -xvj -C ${MAMBA_BIN_DIR} --strip-components=1 bin/micromamba && \
+curl -Ls https://micro.mamba.pm/api/micromamba/linux-${ARCH}/${MAMBA_VERSION} | tar -xvj -C ${MAMBA_BIN_DIR} --strip-components=1 bin/micromamba && \
+chown $(id -u):$(id -g) ${MAMBA_BIN_DIR}/micromamba
 ln -s ${MAMBA_BIN_DIR}/micromamba ${MAMBA_BIN_DIR}/conda && \
 echo -e "channels:\n  - conda-forge" > ${MAMBA_ROOT_PREFIX}/.mambarc
 echo -e "#!/bin/sh\n\. ${MAMBA_ROOT_PREFIX}/etc/profile.d/micromamba.sh || return \$?\nmicromamba activate \"\$@\"" > ${MAMBA_BIN_DIR}/activate
diff --git a/pkg/lang/ir/interface.go b/pkg/lang/ir/interface.go
@@ -145,19 +145,12 @@ func Git(name, email, editor string) error {
 }
 
 func CondaChannel(channel string, useMamba bool) error {
-	if !DefaultGraph.CondaEnabled() {
-		DefaultGraph.CondaConfig = &CondaConfig{}
-	}
-
 	DefaultGraph.CondaConfig.CondaChannel = &channel
 	DefaultGraph.CondaConfig.UseMicroMamba = useMamba
 	return nil
 }
 
 func CondaPackage(deps []string, channel []string, envFile string) error {
-	if !DefaultGraph.CondaEnabled() {
-		DefaultGraph.CondaConfig = &CondaConfig{}
-	}
 	DefaultGraph.CondaConfig.CondaPackages = append(
 		DefaultGraph.CondaConfig.CondaPackages, deps...)
 

diff --git a/pkg/lang/ir/python.go b/pkg/lang/ir/python.go
@@ -24,7 +24,6 @@ import (
 	"github.com/sirupsen/logrus"
 
 	"github.com/tensorchord/envd/pkg/flag"
-	"github.com/tensorchord/envd/pkg/util/fileutil"
 )
 
 const (
@@ -105,7 +104,6 @@ func (g Graph) compilePython(aptStage llb.State) (llb.State, error) {
 
 // Set the system default python to envd's python.
 func (g Graph) compileAlternative(root llb.State) llb.State {
-	root = llb.User("root")(root)
 	envdPrefix := "/opt/conda/envs/envd/bin"
 	run := root.
 		Run(llb.Shlexf("update-alternatives --install /usr/bin/python python %s/python 1", envdPrefix), llb.WithCustomName("update alternative python to envd")).
@@ -120,12 +118,12 @@ func (g Graph) compilePyPIPackages(root llb.State) llb.State {
 		return root
 	}
 
-	cacheDir := fileutil.EnvdHomeDir(".cache")
-	// Create the cache directory to the container. see issue #582
+	// 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",
-		0755, llb.WithParents(true), llb.WithUIDGID(g.uid, g.gid)), llb.WithCustomName("[internal] setting pip cache mount permissions"))
+	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.
@@ -139,13 +137,12 @@ func (g Graph) compilePyPIPackages(root llb.State) llb.State {
 		cmd := sb.String()
 		logrus.WithField("command", cmd).
 			Debug("Configure pip install statements")
-		root = llb.User("envd")(root)
 		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"))
+			llb.AsPersistentCacheDir(g.CacheID(cacheDir), llb.CacheMountShared), llb.SourcePath("/cache/pip"))
 		root = run.Root()
 	}
 
@@ -170,7 +167,7 @@ func (g Graph) compilePyPIPackages(root llb.State) llb.State {
 		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"))
+			llb.AsPersistentCacheDir(g.CacheID(cacheDir), llb.CacheMountShared), llb.SourcePath("/cache/pip"))
 		run.AddMount(g.getWorkingDir(),
 			llb.Local(flag.FlagBuildContext))
 		root = run.Root()
@@ -183,7 +180,7 @@ func (g Graph) compilePyPIPackages(root llb.State) llb.State {
 			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"))
+				llb.AsPersistentCacheDir(g.CacheID(cacheDir), llb.CacheMountShared), llb.SourcePath("/cache/pip"))
 			root = run.Root()
 		}
 	}

diff --git a/pkg/lang/ir/shell.go b/pkg/lang/ir/shell.go
@@ -61,15 +61,15 @@ func (g *Graph) compilePrompt(root llb.State) llb.State {
 		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 config"))
+			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 config")).Root()
+		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 config")).Root()
+			llb.WithCustomName("[internal] setting prompt zsh config")).Root()
 	}
 	return run
 }