feat: upgrade to micromamba2 (#1971)

Signed-off-by: Keming <kemingyang@tensorchord.ai>

Author: kemingy

pkg/lang/ir/v1/conda.go

@@ -30,7 +30,7 @@ import (
 const (
 	builderImage        = "curlimages/curl:8.11.1"
 	condaVersionDefault = "py311_24.11.1-0"
-	microMambaImage     = "mambaorg/micromamba:1.0.0"
+	microMambaImage     = "mambaorg/micromamba:2.0.6"
 	condaRootPrefix     = "/opt/conda"
 	condaBinDir         = "/opt/conda/bin"
 	condaSourcePath     = "/tmp/miniconda.sh"
@@ -86,7 +86,7 @@ func (g generalGraph) condaCommandPath() string {
 func (g generalGraph) condaInitShell(shell string) string {
 	path := g.condaCommandPath()
 	if g.CondaConfig.UseMicroMamba {
-		return fmt.Sprintf("%s shell init -p %s -s %s", path, condaRootPrefix, shell)
+		return fmt.Sprintf("%s shell init -s %s %s", path, shell, condaRootPrefix)
 	}
 	return fmt.Sprintf("%s init %s", path, shell)
 }

pkg/lang/ir/v1/python.go

@@ -91,14 +91,13 @@ func (g generalGraph) getAppropriatePythonVersion() (string, error) {
 func (g generalGraph) compileAlternative(root llb.State) llb.State {
 	envdPrefix := "/opt/conda/envs/envd/bin"
 	run := root.
-		Run(llb.Shlexf("update-alternatives --install /usr/bin/python python %s/python 1", envdPrefix),
-			llb.WithCustomName("[internal] update alternative python to envd")).
-		Run(llb.Shlexf("update-alternatives --install /usr/bin/python3 python3 %s/python3 1", envdPrefix),
-			llb.WithCustomName("[internal] update alternative python3 to envd")).
-		Run(llb.Shlexf("update-alternatives --install /usr/bin/pip pip %s/pip 1", envdPrefix),
-			llb.WithCustomName("[internal] update alternative pip to envd")).
-		Run(llb.Shlexf("update-alternatives --install /usr/bin/pip3 pip3 %s/pip3 1", envdPrefix),
-			llb.WithCustomName("[internal] update alternative pip3 to envd"))
+		Run(llb.Shlexf(`sh -c "
+			update-alternatives --install /usr/bin/python python %[1]s/python 1 &&
+			update-alternatives --install /usr/bin/python3 python3 %[1]s/python3 1 &&
+			update-alternatives --install /usr/bin/pip pip %[1]s/pip 1 &&
+			update-alternatives --install /usr/bin/pip3 pip3 %[1]s/pip3 1
+			"`, envdPrefix),
+			llb.WithCustomName("[internal] update alternative python/python3/pip/pip3 to envd"))
 	return run.Root()
 }