Compatibility with mamba 2.0; added more test and more robust retries

.github/workflows/test.yml

@@ -6,22 +6,23 @@ on:
       - main
 
 jobs:
-
   pre-commit:
     runs-on: ubuntu-latest
     steps:
-    - uses: actions/checkout@692973e3d937129bcbf40652eb9f2f61becf3332 # v4.1.7
-    - uses: actions/setup-python@39cd14951b08e74b54015e9e001cdefcf80e669f # v5.1.1
-    - uses: pre-commit/action@2c7b3805fd2a0fd8c1884dcaebf91fc102a13ecd # v3.0.1
+      - uses: actions/checkout@692973e3d937129bcbf40652eb9f2f61becf3332 # v4.1.7
+      - uses: actions/setup-python@39cd14951b08e74b54015e9e001cdefcf80e669f # v5.1.1
+      - uses: pre-commit/action@2c7b3805fd2a0fd8c1884dcaebf91fc102a13ecd # v3.0.1
 
   test:
     runs-on: ${{ matrix.os }}
     strategy:
       fail-fast: false
       matrix:
-        os: [ ubuntu-latest, macos-latest ]
-        python-version: [ "3.8", "3.10", "3.12" ]
-        resolver: [ mamba, conda, micromamba ]
+        os: [ubuntu-latest, macos-latest]
+        python-version: ["3.8", "3.10", "3.12"]
+        resolver: [mamba, conda, micromamba]
+        micromamba-version: ["1.5.10-0", "latest"]
+        mamba-version: ["1.5.10", "~2"]
     env:
       METAFLOW_CONDA_DEPENDENCY_RESOLVER: ${{ matrix.resolver }}
       METAFLOW_CONDA_TEST: 1
@@ -31,18 +32,19 @@ jobs:
 
       - uses: mamba-org/setup-micromamba@f8b8a1e23a26f60a44c853292711bacfd3eac822 # v1.9.0
         with:
-          micromamba-version: 1.5.10-0
+          micromamba-version: ${{ matrix.micromamba-version }}
           environment-file: dev-env.yml
           init-shell: bash
           create-args: >-
             python=${{ matrix.python-version }}
+            mamba=${{ matrix.mamba-version }}
 
       - name: install nflx-extension
         shell: bash -eo pipefail -l {0}
         run: |
           which pip
           pip install -e . --force-reinstall -U
-      
+
       - name: install bash
         if: runner.os == 'macOS'
         run: brew install bash

metaflow_extensions/netflix_ext/plugins/conda/conda.py

@@ -36,6 +36,7 @@
 from shutil import which
 
 from requests.auth import AuthBase
+from urllib3 import Retry
 
 from metaflow.plugins.datastores.local_storage import LocalStorage
 from metaflow.datastore.datastore_storage import DataStoreStorage
@@ -124,6 +125,9 @@ def _modified_logger(*args: Any, **kwargs: Any):
         self._mode = mode
         self._bins = None  # type: Optional[Dict[str, Optional[str]]]
         self._conda_executable_type = None  # type: Optional[str]
+        # True when using micromamba or mamba 2.0+ which doesn't wrap
+        # conda anymore
+        self.is_non_conda_exec = False  # type: bool
 
         self._have_micromamba_server = False  # type: bool
         self._micromamba_server_port = None  # type: Optional[int]
@@ -271,10 +275,7 @@ def call_conda(
             if (
                 args
                 and args[0] not in ("package", "info")
-                and (
-                    self._conda_executable_type == "micromamba"
-                    or binary == "micromamba"
-                )
+                and (self.is_non_conda_exec or binary == "micromamba")
             ):
                 args.extend(["-r", self.root_prefix, "--json"])
             debug.conda_exec("Conda call: %s" % str([self._bins[binary]] + args))
@@ -1634,7 +1635,7 @@ def _micromamba_transmute(src_file: str, dst_file: str, dst_format: str):
                         a = requests.adapters.HTTPAdapter(
                             pool_connections=executor._max_workers,
                             pool_maxsize=executor._max_workers,
-                            max_retries=3,
+                            max_retries=Retry(total=5, backoff_factor=0.1),
                         )
                         s.mount("https://", a)
                         download_results = [
@@ -1921,6 +1922,7 @@ def _ensure_remote_conda(self):
             self._bins = {"conda": self._ensure_micromamba()}
             self._bins["micromamba"] = self._bins["conda"]
             self._conda_executable_type = "micromamba"
+            self.is_non_conda_exec = True
 
     def _install_remote_conda(self):
         # We download the installer and return a path to it
@@ -1971,6 +1973,7 @@ def _install_remote_conda(self):
         os.sync()
         self._bins = {"conda": final_path, "micromamba": final_path}
         self._conda_executable_type = "micromamba"
+        self.is_non_conda_exec = True
 
     def _validate_conda_installation(self) -> Optional[Exception]:
         # If this is installed in CONDA_LOCAL_PATH look for special marker
@@ -2042,6 +2045,16 @@ def _validate_conda_installation(self) -> Optional[Exception]:
                 return InvalidEnvironmentException(
                     self._install_message_for_resolver("micromamba")
                 )
+            else:
+                self.is_non_conda_exec = True
+        elif "mamba version" in self._info_no_lock:
+            # Mamba 2.0+ has mamba version but no conda version
+            if parse_version(self._info_no_lock) < parse_version("2.0.0"):
+                return InvalidEnvironmentException(
+                    self._install_message_for_resolver("mamba")
+                )
+            else:
+                self.is_non_conda_exec = True
         else:
             if parse_version(self._info_no_lock["conda_version"]) < parse_version(
                 "4.14.0"
@@ -2108,12 +2121,9 @@ def _check_match(dir_name: str) -> Optional[EnvID]:
                     self._remove(os.path.basename(dir_name))
             return None
 
-        if (
-            self._conda_executable_type == "micromamba"
-            or CONDA_LOCAL_PATH is not None
-            or CONDA_TEST
-        ):
-            # Micromamba does not record created environments so we look around for them
+        if self.is_non_conda_exec or CONDA_LOCAL_PATH is not None or CONDA_TEST:
+            # Micromamba (or Mamba 2.0+) does not record created environments so we look
+            # around for them
             # in the root env directory. We also do this if had a local installation
             # because we don't want to look around at other environments created outside
             # of that local installation. Finally, we also do this in test mode for
@@ -2273,7 +2283,7 @@ def _info(self) -> Dict[str, Any]:
     def _info_no_lock(self) -> Dict[str, Any]:
         if self._cached_info is None:
             self._cached_info = json.loads(self.call_conda(["info", "--json"]))
-            if self._conda_executable_type == "micromamba":
+            if "root_prefix" not in self._cached_info:  # Micromamba and Mamba 2.0+
                 self._cached_info["root_prefix"] = self._cached_info["base environment"]
                 self._cached_info["envs_dirs"] = self._cached_info["envs directories"]
                 self._cached_info["pkgs_dirs"] = self._cached_info["package cache"]
@@ -2423,13 +2433,12 @@ def _create(self, env: ResolvedEnvironment, env_name: str) -> str:
                 "--offline",
                 "--no-deps",
             ]
-            if self._conda_executable_type == "micromamba":
-                # micromamba seems to have a bug when compiling .py files. In some
+            if self.is_non_conda_exec:
+                # Micromamba (some version) seems to have a bug when compiling .py files. In some
                 # circumstances, it just hangs forever. We avoid this by not compiling
                 # any file and letting things get compiled lazily. This may have the
                 # added benefit of a faster environment creation.
-                # This option is only available for micromamba so we don't add it
-                # for anything else. This should cover all remote installations though.
+                # This works with Micromamba and Mamba 2.0+.
                 args.append("--no-pyc")
             args.extend(
                 [
@@ -2467,7 +2476,7 @@ def _create(self, env: ResolvedEnvironment, env_name: str) -> str:
                         "--no-deps",
                         "--no-input",
                     ]
-                    if self._conda_executable_type == "micromamba":
+                    if self.is_non_conda_exec:
                         # Be consistent with what we install with micromamba
                         arg_list.append("--no-compile")
                     arg_list.extend(["-r", pypi_list.name])

metaflow_extensions/netflix_ext/plugins/conda/resolvers/conda_resolver.py

@@ -13,7 +13,12 @@
     PackageSpecification,
     ResolvedEnvironment,
 )
-from ..utils import CondaException, channel_or_url, parse_explicit_url_conda
+from ..utils import (
+    CondaException,
+    channel_or_url,
+    clean_up_double_equal,
+    parse_explicit_url_conda,
+)
 from . import Resolver
 
 
@@ -40,7 +45,9 @@ def resolve(
                     % ", ".join([p.package_name for p in local_packages])
                 )
         sys_overrides = {k: v for d in deps.get("sys", []) for k, v in [d.split("==")]}
-        real_deps = list(chain(deps.get("conda", []), deps.get("npconda", [])))
+        real_deps = clean_up_double_equal(
+            chain(deps.get("conda", []), deps.get("npconda", []))
+        )
         packages = []  # type: List[PackageSpecification]
         with tempfile.TemporaryDirectory() as mamba_dir:
             args = [
@@ -86,7 +93,8 @@ def resolve(
         #  - actions:
         #    - FETCH: List of objects to fetch -- this is where we get hash and URL
         #    - LINK: Packages to actually install (in that order)
-        # On micromamba, we can just use the LINK blob since it has all information we need
+        # On micromamba (or Mamba 2+), we can just use the LINK blob since it has all
+        # information we need
         if not conda_result["success"]:
             print(
                 "Pretty-printed Conda create result:\n%s" % conda_result,
@@ -96,7 +104,7 @@ def resolve(
                 "Could not resolve environment -- see above pretty-printed error."
             )
 
-        if self._conda.conda_executable_type == "micromamba":
+        if self._conda.is_non_conda_exec:
             for lnk in conda_result["actions"]["LINK"]:
                 parse_result = parse_explicit_url_conda(
                     "%s#%s" % (lnk["url"], lnk["md5"])

metaflow_extensions/netflix_ext/plugins/conda/resolvers/pip_resolver.py

@@ -24,6 +24,7 @@
 from ..utils import (
     CondaException,
     arch_id,
+    clean_up_double_equal,
     correct_splitext,
     get_glibc_version,
     parse_explicit_path_pypi,
@@ -254,6 +255,7 @@ def resolve(
 
             # Unfortunately, pip doesn't like things like ==<= so we need to strip
             # the ==
+            args.extend(clean_up_double_equal(real_deps))
             for d in real_deps:
                 splits = d.split("==", 1)
                 if len(splits) == 1:

metaflow_extensions/netflix_ext/plugins/conda/utils.py

@@ -18,6 +18,7 @@
     Any,
     Dict,
     FrozenSet,
+    Iterable,
     List,
     Mapping,
     NamedTuple,
@@ -100,6 +101,8 @@ class AliasType(Enum):
     CONDA_FORMATS = _ALL_CONDA_FORMATS  # type: Tuple[str, ...]
 FAKEURL_PATHCOMPONENT = "_fake"
 
+_double_equal_match = re.compile("==(?=[<=>!~])")
+
 
 class CondaException(MetaflowException):
     headline = "Conda ran into an error while setting up environment."
@@ -472,6 +475,10 @@ def split_into_dict(deps: List[str]) -> Dict[str, str]:
     return result
 
 
+def clean_up_double_equal(deps: Iterable[str]) -> List[str]:
+    return [_double_equal_match.sub("", d) for d in deps]
+
+
 def merge_dep_dicts(
     d1: Dict[str, str], d2: Dict[str, str], only_last_deps: bool = False
 ) -> Dict[str, str]: