feat: support micromamba as an alternative to miniconda (#891)

* inii mamba

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

* mamba init bash err

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

* change default to conda

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

* fix mamba activate shell

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

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

pkg/lang/frontend/starlark/config/config.go

@@ -177,18 +177,17 @@ func ruleFuncRStudioServer(thread *starlark.Thread, _ *starlark.Builtin,
 
 func ruleFuncCondaChannel(thread *starlark.Thread, _ *starlark.Builtin,
  args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
- var channel starlark.String
+ var channel string
+ var useMamba bool
 
  if err := starlark.UnpackArgs(ruleCondaChannel, args, kwargs,
- "channel?", &channel); err != nil {
+ "channel?", &channel, "use_mamba?", &useMamba); err != nil {
  return nil, err
  }
 
- channelStr := channel.GoString()
-
- logger.Debugf("rule `%s` is invoked, channel=%s",
- ruleCondaChannel, channelStr)
- if err := ir.CondaChannel(channelStr); err != nil {
+ logger.Debugf("rule `%s` is invoked, channel=%s, use_mamba=%t\n",
+ ruleCondaChannel, channel, useMamba)
+ if err := ir.CondaChannel(channel, useMamba); err != nil {
  return nil, err
  }
 

pkg/lang/frontend/starlark/universe/universe.go

@@ -44,21 +44,18 @@ func RegisterBuildContext(buildContextDir string) {
 
 func ruleFuncBase(thread *starlark.Thread, _ *starlark.Builtin,
  args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
- var os, language, image starlark.String
+ var os, language, image string
+ var useConda bool
 
  if err := starlark.UnpackArgs(ruleBase, args, kwargs,
- "os?", &os, "language?", &language, "image?", &image); err != nil {
+ "os?", &os, "language?", &language, "image?", &image, "use_conda?", &useConda); err != nil {
  return nil, err
  }
 
- osStr := os.GoString()
- langStr := language.GoString()
- imageStr := image.GoString()
+ logger.Debugf("rule `%s` is invoked, os=%s, language=%s, image=%s\n",
+ ruleBase, os, language, image)
 
- logger.Debugf("rule `%s` is invoked, os=%s, language=%s, image=%s",
- ruleBase, osStr, langStr, imageStr)
-
- err := ir.Base(osStr, langStr, imageStr)
+ err := ir.Base(os, language, image)
  return starlark.None, err
 }
 

pkg/lang/ir/conda.go

@@ -17,6 +17,7 @@ package ir
 import (
  _ "embed"
  "fmt"
+ "path/filepath"
  "strings"
 
  "github.com/cockroachdb/errors"
@@ -29,12 +30,18 @@ import (
 
 const (
  condaVersionDefault = "py39_4.11.0"
+ condaRootPrefix = "/opt/conda"
+ condaBinDir = "/opt/conda/bin"
 )
 
 var (
+ // this file can be used by both conda and mamba
+ // https://mamba.readthedocs.io/en/latest/user_guide/configuration.html#multiple-rc-files
  condarc = fileutil.EnvdHomeDir(".condarc")
  //go:embed install-conda.sh
  installCondaBash string
+ //go:embed install-mamba.sh
+ installMambaBash string
 )
 
 func (g Graph) CondaEnabled() bool {
@@ -50,7 +57,7 @@ func (g Graph) CondaEnabled() bool {
 
 func (g Graph) compileCondaChannel(root llb.State) llb.State {
  if g.CondaConfig != nil && g.CondaConfig.CondaChannel != nil {
- logrus.WithField("conda-channel", *g.CondaChannel).Debug("using custom connda channel")
+ logrus.WithField("conda-channel", *g.CondaChannel).Debug("using custom conda channel")
  stage := root.
  File(llb.Mkfile(condarc,
  0644, []byte(*g.CondaChannel), llb.WithUIDGID(g.uid, g.gid)), llb.WithCustomName("[internal] setting conda channel"))
@@ -59,13 +66,35 @@ func (g Graph) compileCondaChannel(root llb.State) llb.State {
  return root
 }
 
+func (g Graph) microMambaEnabled() bool {
+ if g.CondaConfig != nil && g.CondaConfig.UseMicroMamba {
+ return true
+ }
+ return false
+}
+
+func (g Graph) condaCommandPath() string {
+ if g.microMambaEnabled() {
+ return filepath.Join(condaBinDir, "micromamba")
+ }
+ return filepath.Join(condaBinDir, "conda")
+}
+
+func (g Graph) condaInitShell(shell string) string {
+ path := g.condaCommandPath()
+ if g.microMambaEnabled() {
+ return fmt.Sprintf("%s shell init -p %s -s %s", path, condaRootPrefix, shell)
+ }
+ return fmt.Sprintf("%s init %s", path, shell)
+}
+
 func (g Graph) compileCondaPackages(root llb.State) llb.State {
  if !g.CondaEnabled() {
  logrus.Debug("Conda packages not enabled")
  return root
  }
 
- cacheDir := "/opt/conda/pkgs"
+ cacheDir := filepath.Join(condaRootPrefix, "pkgs")
  // Refer to https://github.com/moby/buildkit/blob/31054718bf775bf32d1376fe1f3611985f837584/frontend/dockerfile/dockerfile2llb/convert_runmount.go#L46
  cache := root.File(llb.Mkdir("/cache-conda",
  0755, llb.WithParents(true), llb.WithUIDGID(g.uid, g.gid)),
@@ -82,7 +111,7 @@ func (g Graph) compileCondaPackages(root llb.State) llb.State {
  sb.WriteString(fmt.Sprintf("chown -R envd:envd %s\n", g.getWorkingDir())) // Change mount dir permission
  envdCmd := strings.Builder{}
  envdCmd.WriteString(fmt.Sprintf("cd %s\n", g.getWorkingDir()))
- envdCmd.WriteString(fmt.Sprintf("/opt/conda/bin/conda env update -n envd --file %s\n", g.CondaConfig.CondaEnvFileName))
+ envdCmd.WriteString(fmt.Sprintf("%s env update -n envd --file %s\n", g.condaCommandPath(), g.CondaConfig.CondaEnvFileName))
 
  // Execute the command to write yaml file and conda env using envd user
  sb.WriteString(fmt.Sprintf("sudo -i -u envd bash << EOF\nset -euo pipefail\n%s\nEOF\n", envdCmd.String()))
@@ -102,9 +131,9 @@ func (g Graph) compileCondaPackages(root llb.State) llb.State {
 
  } else {
  if len(g.CondaConfig.AdditionalChannels) == 0 {
- sb.WriteString("/opt/conda/bin/conda install -n envd")
+ sb.WriteString(fmt.Sprintf("%s install -n envd", g.condaCommandPath()))
  } else {
- sb.WriteString("/opt/conda/bin/conda install -n envd")
+ sb.WriteString(fmt.Sprintf("%s install -n envd", g.condaCommandPath()))
  for _, channel := range g.CondaConfig.AdditionalChannels {
  sb.WriteString(fmt.Sprintf(" -c %s", channel))
  }
@@ -128,7 +157,7 @@ func (g Graph) compileCondaPackages(root llb.State) llb.State {
 func (g Graph) compileCondaEnvironment(root llb.State) (llb.State, error) {
  root = llb.User("envd")(root)
 
- cacheDir := "/opt/conda/pkgs"
+ cacheDir := filepath.Join(condaRootPrefix, "pkgs")
  // Create the cache directory to the container. see issue #582
  root = g.CompileCacheDir(root, cacheDir)
 
@@ -138,16 +167,19 @@ func (g Graph) compileCondaEnvironment(root llb.State) (llb.State, error) {
  llb.WithCustomName("[internal] setting conda cache mount permissions"))
 
  // Always init bash since we will use it to create jupyter notebook service.
- run := root.Run(llb.Shlex("bash -c \"/opt/conda/bin/conda init bash\""), llb.WithCustomName("[internal] initialize conda bash environment"))
+ run := root.Run(
+ llb.Shlex(fmt.Sprintf("bash -c \"%s\"", g.condaInitShell("bash"))),
+ llb.WithCustomName("[internal] initialize conda bash environment"),
+ )
 
  pythonVersion, err := g.getAppropriatePythonVersion()
  if err != nil {
  return llb.State{}, errors.Wrap(err, "failed to get python version")
  }
 
  cmd := fmt.Sprintf(
- "bash -c \"/opt/conda/bin/conda create -n envd python=%s\"",
- pythonVersion)
+ "bash -c \"%s create -n envd python=%s\"",
+ g.condaCommandPath(), pythonVersion)
 
  // Create a conda environment.
  run = run.Run(llb.Shlex(cmd),
@@ -158,21 +190,30 @@ func (g Graph) compileCondaEnvironment(root llb.State) (llb.State, error) {
  switch g.Shell {
  case shellBASH:
  run = run.Run(
- llb.Shlex(fmt.Sprintf(`bash -c 'echo "source /opt/conda/bin/activate envd" >> %s'`, fileutil.EnvdHomeDir(".bashrc"))),
+ llb.Shlex(
+ fmt.Sprintf(`bash -c 'echo "source %s/activate envd" >> %s'`,
+ condaBinDir, fileutil.EnvdHomeDir(".bashrc"))),
  llb.WithCustomName("[internal] add conda environment to bashrc"))
  case shellZSH:
  run = run.Run(
- llb.Shlex(fmt.Sprintf("bash -c \"/opt/conda/bin/conda init %s\"", g.Shell)),
+ llb.Shlex(fmt.Sprintf("bash -c \"%s\"", g.condaInitShell(g.Shell))),
  llb.WithCustomNamef("[internal] initialize conda %s environment", g.Shell)).Run(
- llb.Shlex(fmt.Sprintf(`bash -c 'echo "source /opt/conda/bin/activate envd" >> %s'`, fileutil.EnvdHomeDir(".zshrc"))),
+ llb.Shlex(fmt.Sprintf(`bash -c 'echo "source %s/activate envd" >> %s'`, condaBinDir, fileutil.EnvdHomeDir(".zshrc"))),
  llb.WithCustomName("[internal] add conda environment to zshrc"))
  }
  return run.Root(), nil
 }
 
 func (g Graph) installConda(root llb.State) (llb.State, error) {
+ if g.microMambaEnabled() {
+ run := root.AddEnv("MAMBA_BIN_DIR", condaBinDir).
+ AddEnv("MAMBA_ROOT_PREFIX", condaRootPrefix).
+ Run(llb.Shlex(fmt.Sprintf("bash -c '%s'", installMambaBash)),
+ llb.WithCustomName("[internal] install micro mamba"))
+ return run.Root(), nil
+ }
  run := root.AddEnv("CONDA_VERSION", condaVersionDefault).
- File(llb.Mkdir("/opt/conda", 0755, llb.WithParents(true)),
+ File(llb.Mkdir(condaRootPrefix, 0755, llb.WithParents(true)),
  llb.WithCustomName("[internal] create conda directory")).
  Run(llb.Shlex(fmt.Sprintf("bash -c '%s'", installCondaBash)),
  llb.WithCustomName("[internal] install conda"))

pkg/lang/ir/install-conda.sh

@@ -1,4 +1,4 @@
-set -x && \
+set -euo pipefail && \
 UNAME_M="$(uname -m)" && \
 if [ "${UNAME_M}" = "x86_64" ]; then \
  MINICONDA_URL="https://repo.anaconda.com/miniconda/Miniconda3-${CONDA_VERSION}-Linux-x86_64.sh"; \

pkg/lang/ir/install-mamba.sh

@@ -0,0 +1,14 @@
+set -euo pipefail && \
+ARCH="$(uname -m)" && \
+if [[ "${ARCH}" == "aarch64" ]]; then \
+ ARCH="aarch64"; \
+elif [[ "${ARCH}" == "ppc64le" ]]; then \
+ ARCH="ppc64le"; \
+else \
+ ARCH="64"; \
+fi && \
+mkdir -p ${MAMBA_BIN_DIR} && \
+curl -Ls https://micro.mamba.pm/api/micromamba/linux-${ARCH}/latest | tar -xvj -C ${MAMBA_BIN_DIR} --strip-components=1 bin/micromamba && \
+ln -s ${MAMBA_BIN_DIR}/micromamba ${MAMBA_BIN_DIR}/conda && \
+echo -e "channels:\n - conda-forge" > ${MAMBA_ROOT_PREFIX}/.mambarc
+echo -e "#!/bin/sh\n\. ${MAMBA_ROOT_PREFIX}/etc/profile.d/micromamba.sh || return \$?\nmicromamba activate \"\$@\"" > ${MAMBA_BIN_DIR}/activate

pkg/lang/ir/interface.go

@@ -143,16 +143,13 @@ func Git(name, email, editor string) error {
  return nil
 }
 
-func CondaChannel(channel string) error {
- if channel == "" {
- return errors.New("channel is required")
- }
-
+func CondaChannel(channel string, useMamba bool) error {
  if !DefaultGraph.CondaEnabled() {
  DefaultGraph.CondaConfig = &CondaConfig{}
  }
 
  DefaultGraph.CondaConfig.CondaChannel = &channel
+ DefaultGraph.CondaConfig.UseMicroMamba = useMamba
  return nil
 }
 

pkg/lang/ir/types.go

@@ -111,6 +111,7 @@ type CondaConfig struct {
  AdditionalChannels []string
  CondaChannel *string
  CondaEnvFileName string
+ UseMicroMamba bool
 }
 
 type GitConfig struct {