feat(lang): Fix pip cache (#243)

* fix(gitignore): Add demo

Signed-off-by: Ce Gao <cegao@tensorchord.ai>

* feat(lang): Fix cache

Signed-off-by: Ce Gao <cegao@tensorchord.ai>

.gitignore

@@ -37,3 +37,4 @@ debug-bin/
 cover.html
 
 dist/
+.demo/

pkg/lang/ir/cache.go

@@ -14,12 +14,20 @@
 
 package ir
 
-import "fmt"
+import (
+	"fmt"
+
+	"github.com/sirupsen/logrus"
+)
 
 func (g Graph) CacheID(filename string) string {
-	gpu := g.CUDA == nil || g.CUDNN == nil
+	gpu := g.CUDA != nil || g.CUDNN != nil
+	var cacheID string
 	if gpu {
-		return fmt.Sprintf("/%s-gpu-%s", g.CachePrefix, filename)
+		cacheID = fmt.Sprintf("/%s-gpu-%s", g.CachePrefix, filename)
+	} else {
+		cacheID = fmt.Sprintf("/%s-cpu-%s", g.CachePrefix, filename)
 	}
-	return fmt.Sprintf("/%s-cpu-%s", g.CachePrefix, filename)
+	logrus.Debugf("apt/pypi calculated cacheID: %s", cacheID)
+	return cacheID
 }

pkg/lang/ir/python.go

@@ -28,27 +28,25 @@ func (g Graph) compilePyPIPackages(root llb.State) llb.State {
 		return root
 	}
 
-	cacheDir := "/home/envd/.cache/pip"
+	cacheDir := "/home/envd/.cache"
 
 	// Compose the package install command.
 	var sb strings.Builder
-	// Wait until https://github.com/moby/buildkit/commit/31054718bf775bf32d1376fe1f3611985f837584 is released in v0.10.4
-	sb.WriteString("sudo chown -R 1000:1000 ")
-	sb.WriteString(filepath.Dir(cacheDir))
-	sb.WriteString("&& pip install --no-warn-script-location")
+	sb.WriteString("pip install --no-warn-script-location")
 	for _, pkg := range g.PyPIPackages {
 		sb.WriteString(fmt.Sprintf(" %s", pkg))
 	}
 
 	cmd := sb.String()
-	// Wait until https://github.com/moby/buildkit/commit/31054718bf775bf32d1376fe1f3611985f837584 is released in v0.10.4
-	// run := root.File(llb.Mkdir(cacheDir,
-	// 	0755, llb.WithParents(true), llb.WithUIDGID(1000, 1000)), llb.WithCustomName("[internal] settings pip cache mount permissions")).
+	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(1000, 1000)), llb.WithCustomName("[internal] settings pip cache mount permissions"))
 	run := root.
-		Run(llb.Shlex(fmt.Sprintf(`sh -c "%s"`, cmd)), llb.WithCustomNamef("pip install %s",
+		Run(llb.Shlex(cmd), llb.WithCustomNamef("pip install %s",
 			strings.Join(g.PyPIPackages, " ")))
-	run.AddMount(cacheDir, llb.Scratch(),
-		llb.AsPersistentCacheDir(g.CacheID(cacheDir), llb.CacheMountShared))
+	run.AddMount(cacheDir, cache,
+		llb.AsPersistentCacheDir(g.CacheID(cacheDir), llb.CacheMountShared), llb.SourcePath("/cache"))
 	return run.Root()
 }
 
@@ -61,12 +59,12 @@ func (g Graph) compilePyPIIndex(root llb.State) llb.State {
 			extraIndex = "extra-index-url=" + *g.PyPIExtraIndexURL
 		}
 		content := fmt.Sprintf(pypiConfigTemplate, *g.PyPIIndexURL, extraIndex)
-		pypiMirror := llb.Scratch().
+		pypiMirror := root.
 			File(llb.Mkdir(filepath.Dir(pypiIndexFilePath),
 				0755, llb.WithParents(true), llb.WithUIDGID(defaultUID, defaultGID))).
 			File(llb.Mkfile(pypiIndexFilePath,
 				0644, []byte(content), llb.WithUIDGID(defaultUID, defaultGID)))
-		return llb.Merge([]llb.State{root, pypiMirror}, llb.WithCustomName("add PyPI mirror"))
+		return pypiMirror
 	}
 	return root
 }

pkg/lang/ir/system.go

@@ -73,7 +73,7 @@ func (g Graph) compileSystemPackages(root llb.State) llb.State {
 	cacheLibDir := "/var/lib/apt"
 
 	run := root.Run(llb.Shlex(fmt.Sprintf("bash -c \"%s\"", sb.String())),
-		llb.WithCustomNamef("(user-defined packages) apt-get install %s",
+		llb.WithCustomNamef("apt-get install %s",
 			strings.Join(g.SystemPackages, " ")))
 	run.AddMount(cacheDir, llb.Scratch(),
 		llb.AsPersistentCacheDir(g.CacheID(cacheDir), llb.CacheMountShared))
@@ -93,7 +93,7 @@ func (g *Graph) compileBase() llb.State {
 	res := base.
 		Run(llb.Shlex("groupadd -g 1000 envd"), llb.WithCustomName("create user group envd")).
 		Run(llb.Shlex("useradd -p \"\" -u 1000 -g envd -s /bin/sh -m envd"), llb.WithCustomName("create user envd")).
-		Run(llb.Shlex("adduser envd sudo"), llb.WithCustomName("add user envd to sudoers"))
+		Run(llb.Shlex("adduser envd sudo"), llb.WithCustomName("add user envd to sudoers")).Run(llb.Shlex("chown -R envd:envd /usr/local/lib"))
 	return llb.User("envd")(res.Root())
 }