feat(lang): Add default conda pkg cache (#705)

Signed-off-by: Ce Gao <cegao@tensorchord.ai>
tensorchord · Aug 2, 2022 · a1e0395 · a1e0395
1 parent c515f3c
commit a1e0395
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 9 deletions.
diff --git a/pkg/lang/ir/conda.go b/pkg/lang/ir/conda.go
@@ -48,8 +48,10 @@ func (g Graph) compileCondaPackages(root llb.State) llb.State {
 	}
 
 	cacheDir := "/opt/conda/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"))
 
 	// Compose the package install command.
 	var sb strings.Builder
@@ -69,20 +71,27 @@ func (g Graph) compileCondaPackages(root llb.State) llb.State {
 
 	cmd := sb.String()
 	root = llb.User("envd")(root)
-	// Refer to https://github.com/moby/buildkit/blob/31054718bf775bf32d1376fe1f3611985f837584/frontend/dockerfile/dockerfile2llb/convert_runmount.go#L46
-	cache := root.File(llb.Mkdir("/cache",
-		0755, llb.WithParents(true), llb.WithUIDGID(g.uid, g.gid)),
-		llb.WithCustomName("[internal] setting conda cache mount permissions"))
+
 	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"))
+		llb.AsPersistentCacheDir(g.CacheID(cacheDir), llb.CacheMountShared), llb.SourcePath("/cache-conda"))
 	return run.Root()
 }
 
-func (g Graph) setCondaENV(root llb.State) llb.State {
+func (g Graph) compileCondaEnvironment(root llb.State) llb.State {
 	root = llb.User("envd")(root)
+
+	cacheDir := "/opt/conda/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("bash -c \"/opt/conda/bin/conda init bash\""), llb.WithCustomName("[internal] initialize conda bash environment"))
 
@@ -92,8 +101,11 @@ func (g Graph) setCondaENV(root llb.State) llb.State {
 	}
 	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"))
+	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/python.go b/pkg/lang/ir/python.go
@@ -56,7 +56,7 @@ func (g Graph) compilePython(aptStage llb.State) (llb.State, error) {
 		return llb.State{}, errors.Wrap(err, "failed to compile shell")
 	}
 
-	condaEnvStage := g.setCondaENV(shellStage)
+	condaEnvStage := g.compileCondaEnvironment(shellStage)
 
 	condaStage := llb.Diff(builtinSystemStage,
 		g.compileCondaPackages(condaEnvStage),