Added mamba support

.github/workflows/test.yml

@@ -58,9 +58,15 @@ jobs:
         with:
           path: ${{ env.BAZEL_CACHE_DIR }}
           key: ${{ runner.os }}-bazel-${{ env.CACHE_NUMBER }}
-      - name: Run miniconda tests
-        working-directory: tests/miniconda
+      - name: Test miniconda with conda
+        working-directory: tests/miniconda/conda
         run: bazel test --test_output=all --test_arg=--verbose --test_arg=-rA ...
-      - name: Run miniforge tests
-        working-directory: tests/miniforge
+      - name: Test miniconda with mamba
+        working-directory: tests/miniconda/mamba
+        run: bazel test --test_output=all --test_arg=--verbose --test_arg=-rA ...
+      - name: Test miniforge with conda
+        working-directory: tests/miniforge/conda
+        run: bazel test --test_output=all --test_arg=--verbose --test_arg=-rA ...
+      - name: Test miniforge with mamba
+        working-directory: tests/miniforge/mamba
         run: bazel test --test_output=all --test_arg=--verbose --test_arg=-rA ...

conda.bzl

@@ -110,9 +110,21 @@ def _install_conda(rctx, installer):
         fail("Failure installing conda.\nstdout: {}\nstderr: {}".format(result.stdout, result.stderr))
     return "{}/condabin/conda{}".format(rctx.attr.conda_dir, CONDA_EXT_MAP[os])
 
+def _install_mamba(rctx, executable):
+    rctx.report_progress("Installing mamba")
+    mamba_with_version = "mamba={}".format(rctx.attr.mamba_version)
+
+    # `-n base` is necessary so that mamba is installed to the conda in the bazel cache.
+    # If we omit `-n base`, then mamba can get installed to the user's environment i.e. ~/anaconda3 which breaks
+    # the hermetic nature of the build.
+    args = [rctx.path(executable), "install", "-n", "base", "-c", "conda-forge", mamba_with_version, "-y"]
+    result = rctx.execute(args, quiet = rctx.attr.quiet, working_directory = rctx.attr.conda_dir, timeout = rctx.attr.timeout)
+    if result.return_code:
+        fail("Failure when installing mamba.\nstdout: {}\nstderr: {}".format(result.stdout, result.stderr))
+
 # use conda to update itself
 def _update_conda(rctx, executable):
-    conda_with_version = "conda={}".format(rctx.attr.version)
+    conda_with_version = "conda={}".format(rctx.attr.conda_version)
     args = [rctx.path(executable), "install", conda_with_version, "-y"]
 
     # update conda itself
@@ -130,22 +142,32 @@ def _create_conda_build_file(rctx, executable):
 
 def _load_conda_impl(rctx):
     installer = _download_conda(rctx)
-    executable = _install_conda(rctx, installer)
-    _update_conda(rctx, executable)
-    _create_conda_build_file(rctx, executable)
+    conda_executable = _install_conda(rctx, installer)
+    _update_conda(rctx, conda_executable)
+    if rctx.attr.install_mamba:
+        _install_mamba(rctx, conda_executable)
+    _create_conda_build_file(rctx, conda_executable)
 
 load_conda_rule = repository_rule(
     _load_conda_impl,
     attrs = {
         "conda_dir": attr.string(mandatory = True),
-        "version": attr.string(
+        "conda_version": attr.string(
             mandatory = True,
             doc = "Conda version to install",
         ),
         "installer": attr.string(
             default = "miniconda",
             doc = 'Installer to use, either "miniconda" or "miniforge". Note that miniconda and miniforge have different OS/arch support.',
         ),
+        "install_mamba": attr.bool(
+            default = False,
+            doc = "False if mamba should not be installed",
+        ),
+        "mamba_version": attr.string(
+            mandatory = True,
+            doc = "Mamba version to install",
+        ),
         "quiet": attr.bool(
             default = True,
             doc = "False if conda output should be shown",

defs.bzl

@@ -9,14 +9,16 @@ DEFAULT_CONDA_VERSION = "4.10.3"
 DEFAULT_ENV_NAME = "my_env"
 DEFAULT_TOOLCHAIN_REPO_NAME = "conda_tools"
 DEFAULT_TOOLCHAIN_NAME = "python_toolchain"
+DEFAULT_MAMBA_VERSION = "0.17.0"
 
 # download and install conda
-def load_conda(version = DEFAULT_CONDA_VERSION, **kwargs):
+def load_conda(conda_version = DEFAULT_CONDA_VERSION, mamba_version = DEFAULT_MAMBA_VERSION, **kwargs):
     maybe(
         load_conda_rule,
         CONDA_REPO_NAME,
         conda_dir = CONDA_DIR,
-        version = version,
+        conda_version = conda_version,
+        mamba_version = mamba_version,
         **kwargs
     )
 

docs/docs/usage/api.md

@@ -8,12 +8,14 @@
 
     **Parameters:**
 
-    | Name          | Description                                                                                | Default                |
-    | ------------- | ------------------------------------------------------------------------------------------ | ---------------------- |
-    | `installer`   | Which `conda` installer to download, either `miniconda` or `miniforge`                     | `miniconda`            |
-    | `version`     | Version of `conda` to download                                                             | `4.10.3`               |
-    | `quiet`       | `True` if `conda` output should be hidden                                                  | `True`                 |
-    | `timeout`     | How many seconds each execute action can take                                              | `3600`                 |
+    | Name            | Description                                                                                | Default                |
+    | --------------- | ------------------------------------------------------------------------------------------ | ---------------------- |
+    | `conda_version` | Version of `conda` to download                                                             | `4.10.3`               |
+    | `installer`     | Which `conda` installer to download, either `miniconda` or `miniforge`                     | `miniconda`            |
+    | `install_mamba` | Whether to install `mamba`, which is a faster drop-in replacement for `conda`              | `False`                |
+    | `mamba_version` | Version of `conda` to install                                                              | `0.17.0`               |
+    | `quiet`         | `True` if `conda` output should be hidden                                                  | `True`                 |
+    | `timeout`       | How many seconds each execute action can take                                              | `3600`                 |
 
 ## `conda_create`
 
@@ -23,13 +25,14 @@
 
     **Parameters:**
 
-    | Name          | Description                                                                                | Default                |
-    | ------------- | ------------------------------------------------------------------------------------------ | ---------------------- |
-    | `environment` | label pointing to environment configuration file (typically named `environment.yml`)       |                        |
-    | `name`        | Name of the environment                                                                    | `my_env`               |
-    | `quiet`       | `True` if `conda` output should be hidden                                                  | `True`                 |
-    | `timeout`     | How many seconds each execute action can take                                              | `3600`                 |
-    | `clean`       | `True` if `conda` cache should be cleaned (less space taken, but slower subsequent builds) | `False`                |
+    | Name          | Description                                                                                                      | Default                |
+    | ------------- | ---------------------------------------------------------------------------------------------------------------- | ---------------------- |
+    | `environment` | label pointing to environment configuration file (typically named `environment.yml`)                             |                        |
+    | `name`        | Name of the environment                                                                                          | `my_env`               |
+    | `quiet`       | `True` if `conda` output should be hidden                                                                        | `True`                 |
+    | `timeout`     | How many seconds each execute action can take                                                                    | `3600`                 |
+    | `clean`       | `True` if `conda` cache should be cleaned (less space taken, but slower subsequent builds)                       | `False`                |
+    | `use_mamba`   | Whether to use `mamba` to create the `conda` environment. If this is `True`, `install_mamba` must also be `True` | `False`                |
 
 ## `register_toolchain`
 

docs/docs/usage/example.md

@@ -33,6 +33,7 @@ load_conda(
     installer = "miniconda",  # "miniconda" or "miniforge", defaults to "miniconda"
     quiet = False,  # use True to hide conda output
     version = "4.10.3",  # optional, defaults to 4.10.3
+    install_mamba = False, # use True to use mamba, which a faster drop-in replacement for conda
 )
 
 conda_create(
@@ -41,6 +42,7 @@ conda_create(
     clean = False,  # use True if you want to clean conda cache (less space taken, but slower subsequent builds)
     environment = "@//:environment.yml",  # label pointing to environment.yml file
     quiet = False,  # use True to hide conda output
+    use_mamba = False, # use True to use mamba, which a faster drop-in replacement for conda
 )
 
 register_toolchain(

env.bzl

@@ -13,28 +13,52 @@ py_runtime(
 )
 """
 
+def _label_from_condabin(rctx, exe_name):
+    return Label("@{}//:{}/condabin/{}{}".format(
+        rctx.attr.conda_repo,
+        rctx.attr.conda_dir,
+        exe_name,
+        CONDA_EXT_MAP[get_os(rctx)],
+    ))
+
+def _user_chosen_executable(rctx):
+    if rctx.attr.use_mamba:
+        return _label_from_condabin(rctx, "mamba")
+    return _label_from_condabin(rctx, "conda")
+
 # clean conda caches and unused packages
 def _clean(rctx, executable):
     rctx.report_progress("Cleaning up")
 
-    args = [rctx.path(executable), "clean", "-a", "-y"]
+    args = [executable, "clean", "-a", "-y"]
 
     result = rctx.execute(args, quiet = rctx.attr.quiet, timeout = rctx.attr.timeout)
     if result.return_code:
         fail("Failure cleaning up.\nstdout: {}\nstderr: {}".format(result.stdout, result.stderr))
 
+def _create_empty_environment(rctx, executable, env_name):
+    rctx.report_progress("Creating empty conda environment, to be populated afterwards")
+    args = [executable, "create", "-y", "-p", "./{}".format(env_name)]
+    result = rctx.execute(args, quiet = rctx.attr.quiet, timeout = rctx.attr.timeout)
+    if result.return_code:
+        fail("Failure creating empty environment.\nstdout: {}\nstderr: {}".format(result.stdout, result.stderr))
+
+def _update_environment(rctx, executable, env_name, env_file):
+    rctx.report_progress("Updating empty conda environment to populate it")
+    args = [executable, "env", "update", "-f", env_file, "-p", "./{}".format(env_name)]
+    result = rctx.execute(args, quiet = rctx.attr.quiet, timeout = rctx.attr.timeout)
+    if result.return_code:
+        fail("Failure updating environment.\nstdout: {}\nstderr: {}".format(result.stdout, result.stderr))
+
 # create new local conda environment from file
 def _create_environment(rctx, executable, env_name):
     rctx.report_progress("Creating conda environment")
 
     # path to env file as string
-    env_file = str(rctx.path(rctx.attr.environment))
-
-    args = [rctx.path(executable), "env", "create", "-f", env_file, "-p", "./{}".format(env_name)]
+    env_file = rctx.path(rctx.attr.environment)
 
-    result = rctx.execute(args, quiet = rctx.attr.quiet, timeout = rctx.attr.timeout)
-    if result.return_code:
-        fail("Failure creating environment.\nstdout: {}\nstderr: {}".format(result.stdout, result.stderr))
+    _create_empty_environment(rctx, executable, env_name)
+    _update_environment(rctx, executable, env_name, env_file)
 
 # check if python2 or python3 has been installed
 def _get_py_major(rctx, env_path, interpreter_path):
@@ -59,8 +83,7 @@ def _create_env_build_file(rctx, env_name):
     )
 
 def _conda_create_impl(rctx):
-    conda_label = Label("@{}//:{}/condabin/conda{}".format(rctx.attr.conda_repo, rctx.attr.conda_dir, CONDA_EXT_MAP[get_os(rctx)]))
-    executable = str(rctx.path(conda_label))
+    executable = _user_chosen_executable(rctx)
     env_name = rctx.name
     _create_environment(rctx, executable, env_name)
     if rctx.attr.clean:
@@ -81,6 +104,10 @@ conda_create_rule = repository_rule(
             default = True,
             doc = "False if conda output should be shown",
         ),
+        "use_mamba": attr.bool(
+            default = False,
+            doc = "True if mamba should be used",
+        ),
         "timeout": attr.int(
             default = EXECUTE_TIMEOUT,
             doc = "Timeout in seconds for each execute action",

example/WORKSPACE

@@ -62,8 +62,10 @@ load("@rules_conda//:defs.bzl", "conda_create", "load_conda", "register_toolchai
 
 # download and install conda
 load_conda(
+    conda_version = "4.10.3",  # optional, defaults to 4.10.3
+    install_mamba = False,  # use True to install mamba, which a faster drop-in replacement for conda
+    mamba_version = "0.17.0",  # optional, defaults to 0.17.0
     quiet = False,  # use True to hide conda output
-    version = "4.10.3",  # optional, defaults to 4.10.3
 )
 
 # create environment with python2
@@ -73,6 +75,7 @@ conda_create(
     clean = False,  # use True if you want to clean conda cache (less space taken, but slower subsequent builds)
     environment = "@//third_party/conda:py2_environment.yml",  # label pointing to environment.yml file
     quiet = False,  # use True to hide conda output
+    use_mamba = False,  # use True to use mamba, which a faster drop-in replacement for conda
 )
 
 # create environment with python3
@@ -82,6 +85,7 @@ conda_create(
     clean = False,  # use True if you want to clean conda cache (less space taken, but slower subsequent builds)
     environment = "@//third_party/conda:py3_environment.yml",  # label pointing to environment.yml file
     quiet = False,  # use True to hide conda output
+    use_mamba = False,  # use True to use mamba, which a faster drop-in replacement for conda
 )
 
 # register pythons from environment as toolchain

tests/miniconda/WORKSPACE → tests/miniconda/conda/WORKSPACE

@@ -25,7 +25,7 @@ http_archive(
 
 local_repository(
     name = "rules_conda",
-    path = "../../",
+    path = "../../../",
 )
 
 load("@rules_conda//:defs.bzl", "conda_create", "load_conda", "register_toolchain")

tests/miniconda/mamba/.bazeliskrc

@@ -0,0 +1 @@
+USE_BAZEL_VERSION=4.2.1

tests/miniconda/mamba/WORKSPACE

@@ -0,0 +1,49 @@
+load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
+
+RULES_PYTHON_NAME = "rules_python"
+
+RULES_PYTHON_TAG = "0.4.0"
+
+RULES_PYTHON_SHA = "954aa89b491be4a083304a2cb838019c8b8c3720a7abb9c4cb81ac7a24230cea"
+
+RULES_PYTHON_REPO = "bazelbuild"
+
+RULES_PYTHON_ARCHIVE = "tar.gz"
+
+RULES_PYTHON_URL = "https://github.com/{repo}/{name}/releases/download/{tag}/{name}-{tag}.{archive}".format(
+    name = RULES_PYTHON_NAME,
+    archive = RULES_PYTHON_ARCHIVE,
+    repo = RULES_PYTHON_REPO,
+    tag = RULES_PYTHON_TAG,
+)
+
+http_archive(
+    name = RULES_PYTHON_NAME,
+    sha256 = RULES_PYTHON_SHA,
+    url = RULES_PYTHON_URL,
+)
+
+local_repository(
+    name = "rules_conda",
+    path = "../../../",
+)
+
+load("@rules_conda//:defs.bzl", "conda_create", "load_conda", "register_toolchain")
+
+load_conda(
+    install_mamba = True,
+    quiet = False,
+)
+
+conda_create(
+    name = "test_env",
+    timeout = 600,
+    clean = False,
+    environment = "@//:environment.yml",
+    quiet = False,
+    use_mamba = True,
+)
+
+register_toolchain(
+    py3_env = "test_env",
+)

tests/miniforge/conda/BUILD

@@ -0,0 +1,7 @@
+load("@rules_python//python:defs.bzl", "py_test")
+
+py_test(
+    name = "test",
+    size = "small",
+    srcs = ["test.py"],
+)

tests/miniforge/WORKSPACE → tests/miniforge/conda/WORKSPACE

@@ -25,7 +25,7 @@ http_archive(
 
 local_repository(
     name = "rules_conda",
-    path = "../../",
+    path = "../../../",
 )
 
 load("@rules_conda//:defs.bzl", "conda_create", "load_conda", "register_toolchain")

tests/miniforge/conda/environment.yml

@@ -0,0 +1,6 @@
+channels:
+  - default
+  - conda-forge
+dependencies:
+  - python==3.9.7
+  - pytest==6.2.4

tests/miniforge/conda/test.py

@@ -0,0 +1,11 @@
+import sys
+
+import pytest
+
+
+def test_correct_pytest_version():
+    assert pytest.__version__ == "6.2.4"
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__] + sys.argv[1:]))

tests/miniforge/mamba/BUILD

@@ -0,0 +1,7 @@
+load("@rules_python//python:defs.bzl", "py_test")
+
+py_test(
+    name = "test",
+    size = "small",
+    srcs = ["test.py"],
+)