feat(lang): Support pip cache with uid (#198)

Signed-off-by: Ce Gao <cegao@tensorchord.ai>
tensorchord · May 25, 2022 · 6ed91e5 · 6ed91e5
1 parent a589700
commit 6ed91e5
Show file tree

Hide file tree

Showing 5 changed files with 27 additions and 14 deletions.
diff --git a/pkg/builder/builder.go b/pkg/builder/builder.go
@@ -165,6 +165,7 @@ func (b generalBuilder) build(ctx context.Context, def *llb.Definition, pw progr
 				flag.FlagContextDir: b.buildContextDir,
 				flag.FlagCacheDir:   home.GetManager().CacheDir(),
 			},
+			// TODO(gaocegege): Use llb.WithProxy to implement it.
 			FrontendAttrs: map[string]string{
 				"build-arg:HTTPS_PROXY": os.Getenv("HTTPS_PROXY"),
 			},

diff --git a/pkg/lang/ir/consts.go b/pkg/lang/ir/consts.go
@@ -26,5 +26,6 @@ const (
 [global]
 index-url=%s`
 
-	defaultUID = "1000"
+	defaultUID = 1000
+	defaultGID = 1000
 )
diff --git a/pkg/lang/ir/editor.go b/pkg/lang/ir/editor.go
@@ -44,7 +44,7 @@ func (g Graph) compileVSCode() (*llb.State, error) {
 			"/home/envd/.vscode-server/extensions/"+p.String(),
 			&llb.CopyInfo{
 				CreateDestPath: true,
-			}, llb.WithUser(defaultUID)),
+			}, llb.WithUIDGID(defaultUID, defaultGID)),
 			llb.WithCustomNamef("install vscode plugin %s", p.String()))
 		inputs = append(inputs, ext)
 	}

diff --git a/pkg/lang/ir/python.go b/pkg/lang/ir/python.go
@@ -28,18 +28,25 @@ func (g Graph) compilePyPIPackages(root llb.State) llb.State {
 		return root
 	}
 
+	cacheDir := "/home/envd/.cache/pip"
+
 	// Compose the package install command.
 	var sb strings.Builder
-	// TODO(gaocegege): Support per-user config to keep the mirror.
-	sb.WriteString("pip install")
+	// 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")
 	for _, pkg := range g.PyPIPackages {
 		sb.WriteString(fmt.Sprintf(" %s", pkg))
 	}
 
-	cacheDir := "/home/envd/.cache/pip"
 	cmd := sb.String()
-	run := root.Run(llb.Shlex(cmd), llb.WithCustomNamef("pip install %s",
-		strings.Join(g.PyPIPackages, " ")))
+	// 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")).
+	run := root.
+		Run(llb.Shlex(fmt.Sprintf(`sh -c "%s"`, cmd)), llb.WithCustomNamef("pip install %s",
+			strings.Join(g.PyPIPackages, " ")))
 	run.AddMount(cacheDir, llb.Scratch(),
 		llb.AsPersistentCacheDir("/"+cacheDir, llb.CacheMountShared))
 	return run.Root()
@@ -49,10 +56,12 @@ func (g Graph) compilePyPIMirror(root llb.State) llb.State {
 	if g.PyPIMirror != nil {
 		logrus.WithField("mirror", *g.PyPIMirror).Debug("using custom PyPI mirror")
 		content := fmt.Sprintf(pypiConfigTemplate, *g.PyPIMirror)
-		aptSource := llb.Scratch().
-			File(llb.Mkdir(filepath.Dir(pypiMirrorFilePath), 0755, llb.WithParents(true))).
-			File(llb.Mkfile(pypiMirrorFilePath, 0644, []byte(content)))
-		return llb.Merge([]llb.State{root, aptSource}, llb.WithCustomName("add PyPI mirror"))
+		pypiMirror := llb.Scratch().
+			File(llb.Mkdir(filepath.Dir(pypiMirrorFilePath),
+				0755, llb.WithParents(true), llb.WithUIDGID(defaultUID, defaultGID))).
+			File(llb.Mkfile(pypiMirrorFilePath,
+				0644, []byte(content), llb.WithUIDGID(defaultUID, defaultGID)))
+		return llb.Merge([]llb.State{root, pypiMirror}, llb.WithCustomName("add PyPI mirror"))
 	}
 	return root
 }
diff --git a/pkg/lang/ir/shell.go b/pkg/lang/ir/shell.go
@@ -44,10 +44,12 @@ func (g Graph) compileZSH(root llb.State) (llb.State, error) {
 	}
 	zshStage := root.
 		File(llb.Copy(llb.Local(flag.FlagCacheDir), "oh-my-zsh", ohMyZSHPath,
-			&llb.CopyInfo{CreateDestPath: true}, llb.WithUser(defaultUID))).
-		File(llb.Mkfile(installPath, 0644, []byte(m.InstallScript()), llb.WithUser(defaultUID)))
+			&llb.CopyInfo{CreateDestPath: true}, llb.WithUIDGID(defaultUID, defaultGID))).
+		File(llb.Mkfile(installPath,
+			0644, []byte(m.InstallScript()), llb.WithUIDGID(defaultUID, defaultGID)))
 	run := zshStage.Run(llb.Shlex(fmt.Sprintf("bash %s", installPath)),
 		llb.WithCustomName("install oh-my-zsh")).
-		File(llb.Mkfile(zshrcPath, 0644, []byte(m.ZSHRC()), llb.WithUser(defaultUID)))
+		File(llb.Mkfile(zshrcPath,
+			0644, []byte(m.ZSHRC()), llb.WithUIDGID(defaultUID, defaultGID)))
 	return run, nil
 }