fix: v0 bashrc file (#1307)

* fix: v0 bashrc file Signed-off-by: Keming <[email protected]> * reorder py llb Signed-off-by: Keming <[email protected]> * fix go emb Signed-off-by: Keming <[email protected]> Signed-off-by: Keming <[email protected]>
tensorchord · Dec 15, 2022 · 31eec50 · 31eec50
1 parent 1fec54e
commit 31eec50
Show file tree

Hide file tree

Showing 4 changed files with 14 additions and 21 deletions.
diff --git a/envd/api.go b/envd/api.go
@@ -23,7 +23,7 @@ import (
 	"path/filepath"
 )
 
-//go:embed api/*.py api/**/*.py
+//go:embed api/**/*.py
 var api embed.FS
 
 const autoGeneratedHeader = `### AUTO-GENERATED ###

diff --git a/pkg/lang/ir/v0/python.go b/pkg/lang/ir/v0/python.go
@@ -62,13 +62,12 @@ func (g generalGraph) compilePython(baseStage llb.State) (llb.State, error) {
 	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,
+	diffShellStage := llb.Diff(baseStage, shellStage,
 		llb.WithCustomNamef("[internal] configure shell %s", g.Shell))
 	prePythonStage := llb.Merge([]llb.State{
 		diffSystemStage,
@@ -112,7 +111,8 @@ func (g generalGraph) compilePython(baseStage llb.State) (llb.State, error) {
 		}, llb.WithCustomName("[internal] generating the image"))
 	}
 	merged = g.compileAlternative(merged)
-	return merged, nil
+	condaShell := g.compileCondaShell(merged)
+	return condaShell, nil
 }
 
 // Set the system default python to envd's python.

diff --git a/pkg/lang/ir/v0/shell.go b/pkg/lang/ir/v0/shell.go
@@ -67,21 +67,15 @@ func (g *generalGraph) compileShell(root llb.State) (llb.State, error) {
 }
 
 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"))
+	rcPath := fileutil.EnvdHomeDir(".bashrc")
+	if g.Shell == shellZSH {
+		rcPath = fileutil.EnvdHomeDir(".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.WithCustomNamef("[internal] add conda environment to %s", rcPath))
 	return run.Root()
 }
 

diff --git a/pkg/lang/ir/v1/shell.go b/pkg/lang/ir/v1/shell.go
@@ -73,7 +73,6 @@ func (g *generalGraph) compileShell(root llb.State) (_ llb.State, err error) {
 }
 
 func (g *generalGraph) compileCondaShell(root llb.State) llb.State {
-	var run llb.ExecState
 	findDir := fileutil.DefaultHomeDir
 	if g.Dev {
 		findDir = fileutil.EnvdHomeDir
@@ -82,11 +81,11 @@ func (g *generalGraph) compileCondaShell(root llb.State) llb.State {
 	if g.Shell == shellZSH {
 		rcPath = findDir(".zshrc")
 	}
-	run = root.
+	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"))
+			llb.WithCustomNamef("[internal] add conda environment to %s", rcPath))
 	return run.Root()
 }