From 71fb1ce154ba7335de6d40e38f63d65012706c86 Mon Sep 17 00:00:00 2001 From: Keming Date: Thu, 15 Sep 2022 13:34:52 +0800 Subject: [PATCH] feat: support micromamba as an alternative to miniconda (#891) * inii mamba Signed-off-by: Keming * mamba init bash err Signed-off-by: Keming * change default to conda Signed-off-by: Keming * fix mamba activate shell Signed-off-by: Keming Signed-off-by: Keming --- pkg/lang/frontend/starlark/config/config.go | 13 ++-- .../frontend/starlark/universe/universe.go | 15 ++--- pkg/lang/ir/conda.go | 67 +++++++++++++++---- pkg/lang/ir/install-conda.sh | 2 +- pkg/lang/ir/install-mamba.sh | 14 ++++ pkg/lang/ir/interface.go | 7 +- pkg/lang/ir/types.go | 1 + 7 files changed, 84 insertions(+), 35 deletions(-) create mode 100644 pkg/lang/ir/install-mamba.sh diff --git a/pkg/lang/frontend/starlark/config/config.go b/pkg/lang/frontend/starlark/config/config.go index 7e78d2538..da4f096b8 100644 --- a/pkg/lang/frontend/starlark/config/config.go +++ b/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 } diff --git a/pkg/lang/frontend/starlark/universe/universe.go b/pkg/lang/frontend/starlark/universe/universe.go index 71b75dede..76912a91d 100644 --- a/pkg/lang/frontend/starlark/universe/universe.go +++ b/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 } diff --git a/pkg/lang/ir/conda.go b/pkg/lang/ir/conda.go index 88628da60..ad418ffa9 100644 --- a/pkg/lang/ir/conda.go +++ b/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,7 +167,10 @@ 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 { @@ -146,8 +178,8 @@ func (g Graph) compileCondaEnvironment(root llb.State) (llb.State, error) { } 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")) diff --git a/pkg/lang/ir/install-conda.sh b/pkg/lang/ir/install-conda.sh index 1a04c5553..9673cea34 100644 --- a/pkg/lang/ir/install-conda.sh +++ b/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"; \ diff --git a/pkg/lang/ir/install-mamba.sh b/pkg/lang/ir/install-mamba.sh new file mode 100644 index 000000000..461829ffd --- /dev/null +++ b/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 diff --git a/pkg/lang/ir/interface.go b/pkg/lang/ir/interface.go index e0caea24d..9045f14de 100644 --- a/pkg/lang/ir/interface.go +++ b/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 } diff --git a/pkg/lang/ir/types.go b/pkg/lang/ir/types.go index 73f4ddcdc..bda42c6cc 100644 --- a/pkg/lang/ir/types.go +++ b/pkg/lang/ir/types.go @@ -111,6 +111,7 @@ type CondaConfig struct { AdditionalChannels []string CondaChannel *string CondaEnvFileName string + UseMicroMamba bool } type GitConfig struct {