feat: add mamba backend (#88)

* feat: add mamba backend

* tests: add mamba tests

* fix tests

* feat: make mamba default conda resolver

* Update src/isolate/backends/conda.py

Co-authored-by: Batuhan Taskaya <[email protected]>

* PR comments

* PR comments

* Fix tests

---------

Co-authored-by: Batuhan Taskaya <[email protected]>

.github/workflows/test.yml

@@ -28,13 +28,16 @@ jobs:
       with:
         python-version: ${{ matrix.python }}
 
-    - name: Install conda
-      uses: s-weigand/setup-conda@v1
+    - name: Install mamba
+      uses: mamba-org/setup-micromamba@v1
       with:
-        activate-conda: false
-        python-version: ${{ matrix.python }}
-        # PyJokes is available on conda-forge
-        conda-channels: anaconda, conda-forge
+        create-args: >-
+          python=${{ matrix.python }}
+        condarc: |
+          channels:
+            - anaconda
+            - conda-forge
+            - pytorch
 
     - uses: actions/checkout@v3
       with:

src/isolate/backends/conda.py

@@ -22,10 +22,12 @@
 from isolate.connections import PythonIPC
 from isolate.logs import LogLevel
 
-# Specify the path where the conda binary might reside in (or
-# mamba, if it is the preferred one).
+# Specify paths where conda and mamba binaries might reside
 _CONDA_COMMAND = os.environ.get("CONDA_EXE", "conda")
+_MAMBA_COMMAND = os.environ.get("MAMBA_EXE", "micromamba")
 _ISOLATE_CONDA_HOME = os.getenv("ISOLATE_CONDA_HOME")
+_ISOLATE_MAMBA_HOME = os.getenv("ISOLATE_MAMBA_HOME")
+_ISOLATE_DEFAULT_RESOLVER = os.getenv("ISOLATE_DEFAULT_RESOLVER", "mamba")
 
 # Conda accepts the following version specifiers: =, ==, >=, <=, >, <, !=
 _POSSIBLE_CONDA_VERSION_IDENTIFIERS = (
@@ -43,6 +45,9 @@ class CondaEnvironment(BaseEnvironment[Path]):
     environment_definition: Dict[str, Any] = field(default_factory=dict)
     python_version: Optional[str] = None
     tags: List[str] = field(default_factory=list)
+    _exec_home: Optional[str] = _ISOLATE_MAMBA_HOME
+    _exec_command: Optional[str] = _MAMBA_COMMAND
+
 
     @classmethod
     def from_config(
@@ -52,7 +57,15 @@ def from_config(
     ) -> BaseEnvironment:
         processing_config = copy.deepcopy(config)
         processing_config.setdefault("python_version", active_python())
-
+        resolver = processing_config.pop("resolver", _ISOLATE_DEFAULT_RESOLVER)
+        if resolver == "conda":
+            _exec_home = _ISOLATE_CONDA_HOME
+            _exec_command = _CONDA_COMMAND
+        elif resolver == "mamba":
+            _exec_home = _ISOLATE_MAMBA_HOME
+            _exec_command = _MAMBA_COMMAND
+        else:
+            raise Exception(f"Conda resolver of type {resolver} is not supported")
         if "env_dict" in processing_config:
             definition = processing_config.pop("env_dict")
         elif "env_yml_str" in processing_config:
@@ -100,6 +113,8 @@ def from_config(
 
         environment = cls(
             environment_definition=definition,
+            _exec_home=_exec_home,
+            _exec_command=_exec_command,
             **processing_config,
         )
         environment.apply_settings(settings)
@@ -110,6 +125,7 @@ def key(self) -> str:
         return sha256_digest_of(
             repr(self.environment_definition),
             self.python_version,
+            self._exec_command,
             *sorted(self.tags),
         )
 
@@ -126,9 +142,7 @@ def create(self, *, force: bool = False) -> Path:
                 tf.flush()
 
                 try:
-                    self._run_conda(
-                        "env", "create", "--force", "--prefix", env_path, "-f", tf.name
-                    )
+                    self._run_create(str(env_path), tf.name)
                 except subprocess.SubprocessError as exc:
                     raise EnvironmentCreationError(
                         f"Failure during 'conda create': {exc}"
@@ -144,16 +158,19 @@ def destroy(self, connection_key: Path) -> None:
             if not connection_key.exists():
                 return
 
-            self._run_conda(
-                "remove",
-                "--yes",
-                "--all",
-                "--prefix",
-                connection_key,
-            )
+            self._run_destroy(str(connection_key))
+
+    def _run_create(self, env_path: str, env_name: str) -> None:
+        if self._exec_command == "conda":
+            self._run_conda("env", "create", "--force", "--prefix", env_path, "-f", env_name)
+        else:
+            self._run_conda("env", "create", "--prefix", env_path, "-f", env_name)
+
+    def _run_destroy(self, connection_key: str) -> None:
+        self._run_conda("remove","--yes","--all","--prefix", connection_key)
 
     def _run_conda(self, *args: Any) -> None:
-        conda_executable = _get_conda_executable()
+        conda_executable = _get_executable(self._exec_command, self._exec_home)
         with logged_io(partial(self.log, level=LogLevel.INFO)) as (stdout, stderr):
             subprocess.check_call(
                 [conda_executable, *args],
@@ -170,9 +187,9 @@ def open_connection(self, connection_key: Path) -> PythonIPC:
 
 
 @functools.lru_cache(1)
-def _get_conda_executable() -> Path:
-    for path in [_ISOLATE_CONDA_HOME, None]:
-        conda_path = shutil.which(_CONDA_COMMAND, path=path)
+def _get_executable(command: str, home: str | None = None) -> Path:
+    for path in [home, None]:
+        conda_path = shutil.which(command, path=path)
         if conda_path is not None:
             return Path(conda_path)
     else:

tests/test_backends.py

@@ -13,7 +13,7 @@
 import isolate
 from isolate.backends import BaseEnvironment, EnvironmentCreationError
 from isolate.backends.common import sha256_digest_of
-from isolate.backends.conda import CondaEnvironment, _get_conda_executable
+from isolate.backends.conda import CondaEnvironment, _get_executable
 from isolate.backends.local import LocalPythonEnvironment
 from isolate.backends.pyenv import PyenvEnvironment, _get_pyenv_executable
 from isolate.backends.remote import IsolateServer
@@ -375,17 +375,17 @@ def test_tags_in_key(self, tmp_path, monkeypatch):
         assert tagged_environment.key == tagged_environment_2.key, "Tag order should not matter"
 
 
-# Since conda is an external dependency, we'll skip tests using it
+# Since mamba is an external dependency, we'll skip tests using it
 # if it is not installed.
 try:
-    _get_conda_executable()
+    _get_executable("micromamba")
 except FileNotFoundError:
-    IS_CONDA_AVAILABLE = False
+    IS_MAMBA_AVAILABLE = False
 else:
-    IS_CONDA_AVAILABLE = True
+    IS_MAMBA_AVAILABLE = True
 
 
-@pytest.mark.skipif(not IS_CONDA_AVAILABLE, reason="Conda is not available")
+@pytest.mark.skipif(not IS_MAMBA_AVAILABLE, reason="Mamba is not available")
 class TestConda(GenericEnvironmentTests):
 
     backend_cls = CondaEnvironment

tools/Dockerfile

@@ -22,6 +22,15 @@ RUN wget --quiet https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86
     /bin/bash ~/miniconda.sh -b -p /opt/conda
 #### END CONDA ####
 
+#### MAMBA ####
+ENV ISOLATE_MAMBA_HOME=/opt/mamba/bin
+ENV CONDA_DIR /opt/mamba
+RUN mkdir -p /opt/mamba/bin
+RUN curl -Ls https://micro.mamba.pm/api/micromamba/linux-64/latest | tar -xvj -C /opt/mamba/bin/ --strip-components=1 bin/micromamba
+RUN /opt/mamba/bin/micromamba config append channels conda-forge
+RUN /opt/mamba/bin/micromamba config append channels pytorch
+#### END MAMBA ####
+
 RUN pip install --upgrade pip virtualenv wheel poetry-core
 
 # Since system-level debian does not comply with
@@ -36,7 +45,7 @@ COPY tools/requirements.txt /tmp/requirements.txt
 RUN pip install -r /tmp/requirements.txt
 
 COPY . /isolate
-RUN pip install /isolate[server]
+RUN pip install /isolate[server,build]
 
 ENV ISOLATE_INHERIT_FROM_LOCAL=1
 ENV AGENT_REQUIREMENTS_TXT=/isolate/tools/agent_requirements.txt