fix(ir): apt install and conda env create cache (#962)

Signed-off-by: Keming <[email protected]> Signed-off-by: Keming <[email protected]>
tensorchord · Oct 1, 2022 · c22408c · c22408c
1 parent 006f653
commit c22408c
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 27 deletions.
diff --git a/pkg/lang/ir/conda.go b/pkg/lang/ir/conda.go
@@ -157,15 +157,6 @@ func (g Graph) compileCondaPackages(root llb.State) llb.State {
 func (g Graph) compileCondaEnvironment(root llb.State) (llb.State, error) {
 	root = llb.User("envd")(root)
 
-	cacheDir := filepath.Join(condaRootPrefix, "pkgs")
-	// Create the cache directory to the container. see issue #582
-	root = g.CompileCacheDir(root, cacheDir)
-
-	// 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)),
-		llb.WithCustomName("[internal] setting conda cache mount permissions"))
-
 	// 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"))),
@@ -184,8 +175,6 @@ func (g Graph) compileCondaEnvironment(root llb.State) (llb.State, error) {
 	// Create a conda environment.
 	run = run.Run(llb.Shlex(cmd),
 		llb.WithCustomName("[internal] create conda environment"))
-	run.AddMount(cacheDir, cache, llb.AsPersistentCacheDir(
-		g.CacheID(cacheDir), llb.CacheMountShared), llb.SourcePath("/cache-conda"))
 
 	switch g.Shell {
 	case shellBASH:

diff --git a/pkg/lang/ir/system.go b/pkg/lang/ir/system.go
@@ -149,12 +149,6 @@ func (g *Graph) preparePythonBase(root llb.State) llb.State {
 		root = root.AddEnv(env.Name, env.Value)
 	}
 
-	// envd-sshd
-	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)))
-
 	// apt packages
 	var sb strings.Builder
 	sb.WriteString("apt-get update && apt-get install -y apt-utils && ")
@@ -164,19 +158,20 @@ func (g *Graph) preparePythonBase(root llb.State) llb.State {
 	// shell prompt
 	sb.WriteString("&& curl --proto '=https' --tlsv1.2 -sSf https://starship.rs/install.sh | sh -s -- -y")
 
-	cacheDir := "/var/cache/apt"
-	cacheLibDir := "/var/lib/apt"
-
-	run := sshd.Run(llb.Shlex(fmt.Sprintf("bash -c \"%s\"", sb.String())),
+	run := root.Run(llb.Shlex(fmt.Sprintf("bash -c \"%s\"", sb.String())),
 		llb.WithCustomName("[internal] install system packages"))
-	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()
 }
 
+func (g Graph) 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 *Graph) compileBase() (llb.State, error) {
 	logger := logrus.WithFields(logrus.Fields{
 		"os":       g.OS,
@@ -223,10 +218,12 @@ func (g *Graph) compileBase() (llb.State, error) {
 		return llb.State{}, errors.Wrap(err, "failed to install conda")
 	}
 
+	sshd := g.compileSshd(condaStage)
+
 	// TODO(gaocegege): Refactor user to a separate stage.
 	var res llb.ExecState
 	if g.uid == 0 {
-		res = condaStage.
+		res = sshd.
 			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)),
@@ -242,7 +239,7 @@ func (g *Graph) compileBase() (llb.State, error) {
 			Run(llb.Shlex("chown -R root:root /opt/conda"),
 				llb.WithCustomName("[internal] configure user permissions"))
 	} else {
-		res = condaStage.
+		res = sshd.
 			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)),