chroe: make python package after apt package (#1048)

* make python package after apt package Signed-off-by: Jinjing.Zhou <allenzhou@tensorchord.ai> * fix Signed-off-by: Jinjing.Zhou <allenzhou@tensorchord.ai> * tweak Signed-off-by: Jinjing.Zhou <allenzhou@tensorchord.ai> Signed-off-by: Jinjing.Zhou <allenzhou@tensorchord.ai>
tensorchord · Oct 19, 2022 · b601304 · b601304
1 parent cf78dd1
commit b601304
Showing 1 changed file with 10 additions and 6 deletions.
diff --git a/pkg/lang/ir/python.go b/pkg/lang/ir/python.go
@@ -66,20 +66,24 @@ func (g Graph) compilePython(aptStage llb.State) (llb.State, error) {
 		return llb.State{}, errors.Wrap(err, "failed to compile shell")
 	}
 
+	systemStage := g.compileSystemPackages(shellStage)
+
 	condaEnvStage, err := g.compileCondaEnvironment(shellStage)
 	if err != nil {
 		return llb.State{}, errors.Wrap(err, "failed to compile conda environment")
 	}
 
+	diffCondaEnvStage := llb.Diff(shellStage, condaEnvStage)
+	diffSystemStage := llb.Diff(shellStage, systemStage)
+	prePythonStage := llb.Merge([]llb.State{diffSystemStage, diffCondaEnvStage, shellStage}, llb.WithCustomName("pre-python stage"))
+
 	condaStage := llb.Diff(builtinSystemStage,
-		g.compileCondaPackages(condaEnvStage),
+		g.compileCondaPackages(prePythonStage),
 		llb.WithCustomName("install conda packages"))
 
 	pypiStage := llb.Diff(condaEnvStage,
-		g.compilePyPIPackages(condaEnvStage),
+		g.compilePyPIPackages(prePythonStage),
 		llb.WithCustomName("install PyPI packages"))
-	systemStage := llb.Diff(builtinSystemStage, g.compileSystemPackages(builtinSystemStage),
-		llb.WithCustomName("install system packages"))
 
 	vscodeStage, err := g.compileVSCode()
 	if err != nil {
@@ -89,12 +93,12 @@ func (g Graph) compilePython(aptStage llb.State) (llb.State, error) {
 	var merged llb.State
 	if vscodeStage != nil {
 		merged = llb.Merge([]llb.State{
-			builtinSystemStage, systemStage, condaStage,
+			builtinSystemStage, condaStage,
 			diffSSHStage, pypiStage, *vscodeStage,
 		}, llb.WithCustomName("merging all components into one"))
 	} else {
 		merged = llb.Merge([]llb.State{
-			builtinSystemStage, systemStage, condaStage,
+			builtinSystemStage, condaStage,
 			diffSSHStage, pypiStage,
 		}, llb.WithCustomName("merging all components into one"))
 	}