Repository: tensorchord/envd (language: python)

Issue:
feat: make it cache friendly for pip/conda/apt
## Description

Currently, all the required pip/conda/apt packages will be joined together in one installation command. These are not cache-friendly. BTW, some packages may change the environment which causes a cache miss (https://github.com/tensorchord/envd/issues/1252).

We should make each install command a separate process, just like the changes to `run` in https://github.com/tensorchord/envd/pull/1044.

---
<!-- Issue Author: Don't delete this message to encourage other users to support your issue! -->
**Message from the maintainers**:

Love this enhancement proposal? Give it a 👍. We prioritise the proposals with the most 👍.



Discussion / hints:

Fix patch:
diff --git a/pkg/lang/ir/v0/compile.go b/pkg/lang/ir/v0/compile.go
index d3f157295..5ca4bdb79 100644
--- a/pkg/lang/ir/v0/compile.go
+++ b/pkg/lang/ir/v0/compile.go
@@ -54,7 +54,7 @@ func NewGraph() ir.Graph {
 		CUDNN:   CUDNNVersionDefault,
 		NumGPUs: 0,
 
-		PyPIPackages:    []string{},
+		PyPIPackages:    [][]string{},
 		RPackages:       []string{},
 		JuliaPackages:   []string{},
 		SystemPackages:  []string{},
diff --git a/pkg/lang/ir/v0/custom.go b/pkg/lang/ir/v0/custom.go
index 899f3a58f..d6b9d61e3 100644
--- a/pkg/lang/ir/v0/custom.go
+++ b/pkg/lang/ir/v0/custom.go
@@ -37,26 +37,20 @@ func (g generalGraph) compileCustomPyPIPackages(root llb.State) llb.State {
 	}
 
 	cacheDir := "/home/root/.cache"
-
-	// Compose the package install command.
-	var sb strings.Builder
-	// Always use the conda's pip.
-	sb.WriteString(`pip install`)
-	for _, pkg := range g.PyPIPackages {
-		sb.WriteString(fmt.Sprintf(" %s", pkg))
-	}
-
-	cmd := sb.String()
 	// 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.WithCustomName("[internal] settings pip cache mount permissions"))
-	run := root.
-		Run(llb.Shlex(cmd), llb.WithCustomNamef("pip install %s",
-			strings.Join(g.PyPIPackages, " ")))
-	run.AddMount(cacheDir, cache,
-		llb.AsPersistentCacheDir(g.CacheID(cacheDir), llb.CacheMountShared),
-		llb.SourcePath("/cache"))
-	return run.Root()
+
+	for _, packages := range g.PyPIPackages {
+		cmd := fmt.Sprintf("pip install %s", strings.Join(packages, " "))
+		run := root.
+			Run(llb.Shlex(cmd), llb.WithCustomNamef("pip install %s",
+				strings.Join(packages, " ")))
+		run.AddMount(cacheDir, cache,
+			llb.AsPersistentCacheDir(g.CacheID(cacheDir), llb.CacheMountShared), llb.SourcePath("/cache"))
+		root = run.Root()
+	}
+	return root
 }
 
 func (g generalGraph) compileCustomSystemPackages(root llb.State) llb.State {
diff --git a/pkg/lang/ir/v0/editor.go b/pkg/lang/ir/v0/editor.go
index 641d06d35..4d7bd68ad 100644
--- a/pkg/lang/ir/v0/editor.go
+++ b/pkg/lang/ir/v0/editor.go
@@ -63,7 +63,7 @@ func (g *generalGraph) compileJupyter() error {
 		return nil
 	}
 
-	g.PyPIPackages = append(g.PyPIPackages, "jupyter")
+	g.PyPIPackages = append(g.PyPIPackages, []string{"jupyter"})
 	switch g.Language.Name {
 	case "python":
 		return nil
diff --git a/pkg/lang/ir/v0/interface.go b/pkg/lang/ir/v0/interface.go
index d064cb1ec..49e3209b8 100644
--- a/pkg/lang/ir/v0/interface.go
+++ b/pkg/lang/ir/v0/interface.go
@@ -45,7 +45,9 @@ func Base(os, language, image string) error {
 func PyPIPackage(deps []string, requirementsFile string, wheels []string) error {
 	g := DefaultGraph.(*generalGraph)
 
-	g.PyPIPackages = append(g.PyPIPackages, deps...)
+	if len(deps) > 0 {
+		g.PyPIPackages = append(g.PyPIPackages, deps)
+	}
 	g.PythonWheels = append(g.PythonWheels, wheels...)
 
 	if requirementsFile != "" {
diff --git a/pkg/lang/ir/v0/python.go b/pkg/lang/ir/v0/python.go
index 2d4e32ac6..1aeed261b 100644
--- a/pkg/lang/ir/v0/python.go
+++ b/pkg/lang/ir/v0/python.go
@@ -143,24 +143,16 @@ func (g generalGraph) compilePyPIPackages(root llb.State) llb.State {
 		llb.WithCustomName("[internal] setting pip cache mount permissions"))
 
 	if len(g.PyPIPackages) != 0 {
-		// Compose the package install command.
-		var sb strings.Builder
-		// Always use the conda's pip.
-		sb.WriteString("/opt/conda/envs/envd/bin/python -m pip install")
-		for _, pkg := range g.PyPIPackages {
-			sb.WriteString(fmt.Sprintf(" %s", pkg))
+		for _, packages := range g.PyPIPackages {
+			cmd := fmt.Sprintf("/opt/conda/envs/envd/bin/python -m pip install %s", strings.Join(packages, " "))
+			logrus.WithField("command", cmd).Debug("Configure pip install statements")
+			run := root.
+				Run(llb.Shlex(cmd), llb.WithCustomNamef("pip install %s", strings.Join(packages, " ")))
+			// Refer to https://github.com/moby/buildkit/blob/31054718bf775bf32d1376fe1f3611985f837584/frontend/dockerfile/dockerfile2llb/convert_runmount.go#L46
+			run.AddMount(cacheDir, cache,
+				llb.AsPersistentCacheDir(g.CacheID(cacheDir), llb.CacheMountShared), llb.SourcePath("/cache/pip"))
+			root = run.Root()
 		}
-
-		cmd := sb.String()
-		logrus.WithField("command", cmd).
-			Debug("Configure pip install statements")
-		run := root.
-			Run(llb.Shlex(sb.String()), llb.WithCustomNamef("pip install %s",
-				strings.Join(g.PyPIPackages, " ")))
-		// Refer to https://github.com/moby/buildkit/blob/31054718bf775bf32d1376fe1f3611985f837584/frontend/dockerfile/dockerfile2llb/convert_runmount.go#L46
-		run.AddMount(cacheDir, cache,
-			llb.AsPersistentCacheDir(g.CacheID(cacheDir), llb.CacheMountShared), llb.SourcePath("/cache/pip"))
-		root = run.Root()
 	}
 
 	if g.RequirementsFile != nil {
diff --git a/pkg/lang/ir/v0/types.go b/pkg/lang/ir/v0/types.go
index 0f3944acd..eee98fd5d 100644
--- a/pkg/lang/ir/v0/types.go
+++ b/pkg/lang/ir/v0/types.go
@@ -45,7 +45,7 @@ type generalGraph struct {
 
 	PublicKeyPath string
 
-	PyPIPackages     []string
+	PyPIPackages     [][]string
 	RequirementsFile *string
 	PythonWheels     []string
 	RPackages        []string
diff --git a/pkg/lang/ir/v1/compile.go b/pkg/lang/ir/v1/compile.go
index 014b46172..a1adc4381 100644
--- a/pkg/lang/ir/v1/compile.go
+++ b/pkg/lang/ir/v1/compile.go
@@ -49,7 +49,7 @@ func NewGraph() ir.Graph {
 		CUDNN:    CUDNNVersionDefault,
 		NumGPUs:  0,
 
-		PyPIPackages:    []string{},
+		PyPIPackages:    [][]string{},
 		RPackages:       []string{},
 		JuliaPackages:   []string{},
 		SystemPackages:  []string{},
diff --git a/pkg/lang/ir/v1/editor.go b/pkg/lang/ir/v1/editor.go
index 3ee0ad3c1..a92e87ba9 100644
--- a/pkg/lang/ir/v1/editor.go
+++ b/pkg/lang/ir/v1/editor.go
@@ -64,7 +64,7 @@ func (g *generalGraph) compileJupyter() error {
 		return nil
 	}
 
-	g.PyPIPackages = append(g.PyPIPackages, "jupyter")
+	g.PyPIPackages = append(g.PyPIPackages, []string{"jupyter"})
 	switch g.Language.Name {
 	case "python":
 		return nil
diff --git a/pkg/lang/ir/v1/interface.go b/pkg/lang/ir/v1/interface.go
index 264ce680e..da2ee9008 100644
--- a/pkg/lang/ir/v1/interface.go
+++ b/pkg/lang/ir/v1/interface.go
@@ -77,7 +77,9 @@ func Julia() {
 func PyPIPackage(deps []string, requirementsFile string, wheels []string) error {
 	g := DefaultGraph.(*generalGraph)
 
-	g.PyPIPackages = append(g.PyPIPackages, deps...)
+	if len(deps) > 0 {
+		g.PyPIPackages = append(g.PyPIPackages, deps)
+	}
 	g.PythonWheels = append(g.PythonWheels, wheels...)
 
 	if requirementsFile != "" {
diff --git a/pkg/lang/ir/v1/python.go b/pkg/lang/ir/v1/python.go
index 563b08f56..5c6f6f8de 100644
--- a/pkg/lang/ir/v1/python.go
+++ b/pkg/lang/ir/v1/python.go
@@ -103,28 +103,21 @@ func (g generalGraph) compilePyPIPackages(root llb.State) llb.State {
 	cacheDir := filepath.Join("/", "root", ".cache", "pip")
 	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/pip", 0755, llb.WithParents(true)),
 		llb.WithCustomName("[internal] setting pip cache mount permissions"))
 
 	if len(g.PyPIPackages) != 0 {
-		// Compose the package install command.
-		var sb strings.Builder
-		// Always use the conda's pip.
-		sb.WriteString("python -m pip install")
-		for _, pkg := range g.PyPIPackages {
-			sb.WriteString(fmt.Sprintf(" %s", pkg))
+		for _, packages := range g.PyPIPackages {
+			command := fmt.Sprintf("python -m pip install %s", strings.Join(packages, " "))
+			logrus.WithField("command", command).Debug("Configure pip install statements")
+			run := root.
+				Run(llb.Shlex(command), llb.WithCustomNamef("[internal] pip install %s",
+					strings.Join(packages, " ")))
+			run.AddMount(cacheDir, cache,
+				llb.AsPersistentCacheDir(g.CacheID(cacheDir), llb.CacheMountShared), llb.SourcePath("/cache/pip"))
+			root = run.Root()
 		}
-
-		cmd := sb.String()
-		logrus.WithField("command", cmd).
-			Debug("Configure pip install statements")
-		run := root.
-			Run(llb.Shlex(sb.String()), llb.WithCustomNamef("[internal] pip install %s",
-				strings.Join(g.PyPIPackages, " ")))
-		// Refer to https://github.com/moby/buildkit/blob/31054718bf775bf32d1376fe1f3611985f837584/frontend/dockerfile/dockerfile2llb/convert_runmount.go#L46
-		run.AddMount(cacheDir, cache,
-			llb.AsPersistentCacheDir(g.CacheID(cacheDir), llb.CacheMountShared), llb.SourcePath("/cache/pip"))
-		root = run.Root()
 	}
 
 	if g.RequirementsFile != nil {
diff --git a/pkg/lang/ir/v1/types.go b/pkg/lang/ir/v1/types.go
index 407316532..f2341ff67 100644
--- a/pkg/lang/ir/v1/types.go
+++ b/pkg/lang/ir/v1/types.go
@@ -46,7 +46,7 @@ type generalGraph struct {
 
 	PublicKeyPath string
 
-	PyPIPackages     []string
+	PyPIPackages     [][]string
 	RequirementsFile *string
 	PythonWheels     []string
 	RPackages        []string
diff --git a/pkg/types/envd.go b/pkg/types/envd.go
index 4f4ac19a9..d05b5a383 100644
--- a/pkg/types/envd.go
+++ b/pkg/types/envd.go
@@ -316,9 +316,11 @@ func NewDependencyFromLabels(label map[string]string) (*Dependency, error) {
 		}
 		packages := []string{}
 
-		for i, pkg := range lst {
-			if !strings.HasPrefix(pkg, "-") && (i == 0 || !strings.HasPrefix(lst[i-1], "-")) {
-				packages = append(packages, pkg)
+		for _, pkgs := range lst {
+			for i, pkg := range pkgs {
+				if !strings.HasPrefix(pkg, "-") && (i == 0 || !strings.HasPrefix(pkgs[i-1], "-")) {
+					packages = append(packages, pkg)
+				}
 			}
 		}
 		dep.PyPIPackages = packages
@@ -332,8 +334,8 @@ func parseAPTPackages(lst string) ([]string, error) {
 	return pkgs, err
 }
 
-func parsePyPICommands(lst string) ([]string, error) {
-	var pkgs []string
+func parsePyPICommands(lst string) ([][]string, error) {
+	var pkgs [][]string
 	err := json.Unmarshal([]byte(lst), &pkgs)
 	return pkgs, err
 }


Test patch:
diff --git a/pkg/types/envd_test.go b/pkg/types/envd_test.go
index 8190250d4..e8395cf43 100644
--- a/pkg/types/envd_test.go
+++ b/pkg/types/envd_test.go
@@ -25,7 +25,7 @@ var _ = g.Describe("pypi label", func() {
 	g.When("pypi command contains slashes", func() {
 		g.It("should ignore such commands", func() {
 			labels := make(map[string]string)
-			s, _ := json.Marshal([]string{"numpy", "torch", "--extra-url", "https://download.pytorch.org/whl/torch_stable.html"})
+			s, _ := json.Marshal([][]string{{"numpy", "torch", "--extra-url", "https://download.pytorch.org/whl/torch_stable.html"}})
 			labels[ImageLabelPyPI] = string(s)
 			d, err := NewDependencyFromLabels(labels)
 			Expect(err).NotTo(HaveOccurred())