Merge pull request #1062 from wolfv/micromamba

bootstrap base env with micromamba

Author: minrk

repo2docker/buildpacks/conda/__init__.py

@@ -55,6 +55,8 @@ def get_build_env(self):
             ("NPM_DIR", "${APP_BASE}/npm"),
             ("NPM_CONFIG_GLOBALCONFIG", "${NPM_DIR}/npmrc"),
             ("NB_ENVIRONMENT_FILE", self._nb_environment_file),
+            ("MAMBA_ROOT_PREFIX", "${CONDA_DIR}"),
+            ("MAMBA_EXE", "/tmp/bin/micromamba"),
         ]
         if self._nb_requirements_file:
             env.append(("NB_REQUIREMENTS_FILE", self._nb_requirements_file))
@@ -98,7 +100,7 @@ def get_build_scripts(self):
 
         All scripts here should be independent of contents of the repository.
 
-        This sets up through `install-miniforge.bash` (found in this directory):
+        This sets up through `install-base-env.bash` (found in this directory):
 
         - a directory for the conda environment and its ownership by the
           notebook user
@@ -115,8 +117,8 @@ def get_build_scripts(self):
                 "root",
                 r"""
                 TIMEFORMAT='time: %3R' \
-                bash -c 'time /tmp/install-miniforge.bash' && \
-                rm -rf /tmp/install-miniforge.bash /tmp/env
+                bash -c 'time /tmp/install-base-env.bash' && \
+                rm -rf /tmp/install-base-env.bash /tmp/env
                 """,
             ),
             (
@@ -146,7 +148,7 @@ def get_build_script_files(self):
 
         """
         files = {
-            "conda/install-miniforge.bash": "/tmp/install-miniforge.bash",
+            "conda/install-base-env.bash": "/tmp/install-base-env.bash",
             "conda/activate-conda.sh": "/etc/profile.d/activate-conda.sh",
         }
         py_version = self.python_version
@@ -337,9 +339,9 @@ def get_env_scripts(self):
                     "${NB_USER}",
                     r"""
                 TIMEFORMAT='time: %3R' \
-                bash -c 'time mamba env update -p {0} -f "{1}" && \
-                time mamba clean --all -f -y && \
-                mamba list -p {0} \
+                bash -c 'time ${{MAMBA_EXE}} install -p {0} -f "{1}" && \
+                time ${{MAMBA_EXE}} clean --all -y && \
+                ${{MAMBA_EXE}} list -p {0} \
                 '
                 """.format(
                         env_prefix, environment_yml
@@ -356,9 +358,9 @@ def get_env_scripts(self):
                 (
                     "${NB_USER}",
                     r"""
-                mamba install -p {0} r-base{1} r-irkernel=1.2 r-devtools -y && \
-                mamba clean --all -f -y && \
-                mamba list -p {0}
+                ${{MAMBA_EXE}} install -p {0} r-base{1} r-irkernel=1.2 r-devtools -y && \
+                ${{MAMBA_EXE}} clean --all -y && \
+                ${{MAMBA_EXE}} list -p {0}
                 """.format(
                         env_prefix, r_pin
                     ),

repo2docker/buildpacks/conda/activate-conda.sh

@@ -1,11 +1,13 @@
 # enable conda and activate the notebook environment
-CONDA_PROFILE="${CONDA_DIR}/etc/profile.d/conda.sh"
+/tmp/bin/micromamba shell init -s bash -p ${CONDA_DIR}
+export MAMBA_ROOT_PREFIX="/srv/conda"
+CONDA_PROFILE="${CONDA_DIR}/etc/profile.d/mamba.sh"
 test -f $CONDA_PROFILE && . $CONDA_PROFILE
 if [[ "${KERNEL_PYTHON_PREFIX}" != "${NB_PYTHON_PREFIX}" ]]; then
     # if the kernel is a separate env, stack them
     # so both are on PATH, notebook first
-    conda activate ${KERNEL_PYTHON_PREFIX}
-    conda activate --stack ${NB_PYTHON_PREFIX}
+    micromamba activate ${KERNEL_PYTHON_PREFIX}
+    micromamba activate --stack ${NB_PYTHON_PREFIX}
 
     # even though it's second on $PATH
     # make sure CONDA_DEFAULT_ENV is the *kernel* env
@@ -14,5 +16,5 @@ if [[ "${KERNEL_PYTHON_PREFIX}" != "${NB_PYTHON_PREFIX}" ]]; then
     # which only contains UI when the two are different
     export CONDA_DEFAULT_ENV="${KERNEL_PYTHON_PREFIX}"
 else
-    conda activate ${NB_PYTHON_PREFIX}
+    micromamba activate ${NB_PYTHON_PREFIX}
 fi

repo2docker/buildpacks/conda/install-base-env.bash

@@ -0,0 +1,87 @@
+#!/bin/bash
+# This downloads and installs a pinned version of micromamba
+# and sets up the base environment
+set -ex
+
+cd $(dirname $0)
+
+export MAMBA_VERSION=0.19.1
+export CONDA_VERSION=4.11.0
+
+URL="https://anaconda.org/conda-forge/micromamba/${MAMBA_VERSION}/download/linux-64/micromamba-${MAMBA_VERSION}-0.tar.bz2"
+
+# make sure we don't do anything funky with user's $HOME
+# since this is run as root
+unset HOME
+mkdir -p ${CONDA_DIR}
+
+time wget -qO- ${URL} | tar -xvj bin/micromamba
+chmod 0755 /tmp/bin/micromamba
+
+export MAMBA_ROOT_PREFIX=${CONDA_DIR}
+export MAMBA_EXE="/tmp/bin/micromamba"
+
+eval "$(/tmp/bin/micromamba shell hook -p ${CONDA_DIR} -s posix)"
+
+micromamba activate
+
+export PATH="${PWD}/bin:$PATH"
+
+cat <<EOT >> ${CONDA_DIR}/.condarc
+channels:
+  - conda-forge
+  - defaults
+auto_update_conda: false
+show_channel_urls: true
+update_dependencies: false
+# channel_priority: flexible
+EOT
+
+micromamba install conda=${CONDA_VERSION} mamba=${MAMBA_VERSION} -y
+
+echo "installing notebook env:"
+cat "${NB_ENVIRONMENT_FILE}"
+
+
+time micromamba create -p ${NB_PYTHON_PREFIX} -f "${NB_ENVIRONMENT_FILE}"
+
+if [[ ! -z "${NB_REQUIREMENTS_FILE:-}" ]]; then
+    echo "installing pip requirements"
+    cat "${NB_REQUIREMENTS_FILE}"
+    ${NB_PYTHON_PREFIX}/bin/python -mpip install --no-cache --no-deps -r "${NB_REQUIREMENTS_FILE}"
+fi
+# empty conda history file,
+# which seems to result in some effective pinning of packages in the initial env,
+# which we don't intend.
+# this file must not be *removed*, however
+echo '' > ${NB_PYTHON_PREFIX}/conda-meta/history
+
+if [[ ! -z "${KERNEL_ENVIRONMENT_FILE:-}" ]]; then
+    # install kernel env and register kernelspec
+    echo "installing kernel env:"
+    cat "${KERNEL_ENVIRONMENT_FILE}"
+    time micromamba create -p ${KERNEL_PYTHON_PREFIX} -f "${KERNEL_ENVIRONMENT_FILE}"
+
+    if [[ ! -z "${KERNEL_REQUIREMENTS_FILE:-}" ]]; then
+        echo "installing pip requirements for kernel"
+        cat "${KERNEL_REQUIREMENTS_FILE}"
+        ${KERNEL_PYTHON_PREFIX}/bin/python -mpip install --no-cache --no-deps -r "${KERNEL_REQUIREMENTS_FILE}"
+    fi
+
+    ${KERNEL_PYTHON_PREFIX}/bin/ipython kernel install --prefix "${NB_PYTHON_PREFIX}"
+    echo '' > ${KERNEL_PYTHON_PREFIX}/conda-meta/history
+    micromamba list -p ${KERNEL_PYTHON_PREFIX}
+fi
+
+# Clean things out!
+time micromamba clean --all -y
+
+# Remove the pip cache created as part of installing micromamba
+rm -rf /root/.cache
+
+chown -R $NB_USER:$NB_USER ${CONDA_DIR}
+
+micromamba list -p ${NB_PYTHON_PREFIX}
+
+# Set NPM config
+${NB_PYTHON_PREFIX}/bin/npm config --global set prefix ${NPM_DIR}

repo2docker/buildpacks/conda/install-miniforge.bash

@@ -4,7 +4,18 @@ from subprocess import check_output
 
 assert sys.version_info[:2] == (3, 5), sys.version
 
-out = check_output(["conda", "--version"]).decode("utf8").strip()
-assert out == "conda 4.9.2", out
+out = check_output(["/tmp/bin/micromamba", "--version"]).decode("utf8").strip()
+assert (
+    out
+    == """micromamba: 0.19.1
+libmamba: 0.19.1"""
+), out
+
+out = check_output(["mamba", "--version"]).decode("utf8").strip()
+assert (
+    out
+    == """mamba 0.19.1
+conda 4.11.0"""
+), out
 
 import numpy

tests/conda/binder-dir/verify

@@ -14,13 +14,17 @@ assert sorted(specs) == ["python2", "python3"], specs.keys()
 import json
 from subprocess import check_output
 
-envs = json.loads(check_output(["conda", "env", "list", "--json"]).decode("utf8"))
+envs = json.loads(
+    check_output(["/tmp/bin/micromamba", "env", "list", "--json"]).decode("utf8")
+)
 assert envs == {
     "envs": ["/srv/conda", "/srv/conda/envs/kernel", "/srv/conda/envs/notebook"]
 }, envs
 
 pkgs = json.loads(
-    check_output(["conda", "list", "-n", "kernel", "--json"]).decode("utf8")
+    check_output(["/tmp/bin/micromamba", "list", "-n", "kernel", "--json"]).decode(
+        "utf8"
+    )
 )
 pkg_names = [pkg["name"] for pkg in pkgs]
 assert "ipykernel" in pkg_names, pkg_names

tests/conda/simple-py2/verify

@@ -1,3 +1,4 @@
 dependencies:
+  - pip
   - pip:
       - pypi-pkg-test