diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 54bb580..be402e3 100644 --- a/.github/workflows/test.yml +++ b/.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 diff --git a/metaflow_extensions/netflix_ext/plugins/conda/conda.py b/metaflow_extensions/netflix_ext/plugins/conda/conda.py index 07b0710..ec2093a 100644 --- a/metaflow_extensions/netflix_ext/plugins/conda/conda.py +++ b/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]) diff --git a/metaflow_extensions/netflix_ext/plugins/conda/resolvers/conda_resolver.py b/metaflow_extensions/netflix_ext/plugins/conda/resolvers/conda_resolver.py index 077164e..f3f47b0 100644 --- a/metaflow_extensions/netflix_ext/plugins/conda/resolvers/conda_resolver.py +++ b/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"]) diff --git a/metaflow_extensions/netflix_ext/plugins/conda/resolvers/pip_resolver.py b/metaflow_extensions/netflix_ext/plugins/conda/resolvers/pip_resolver.py index 925b690..06244d1 100644 --- a/metaflow_extensions/netflix_ext/plugins/conda/resolvers/pip_resolver.py +++ b/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: diff --git a/metaflow_extensions/netflix_ext/plugins/conda/utils.py b/metaflow_extensions/netflix_ext/plugins/conda/utils.py index 476bdbb..0d3375f 100644 --- a/metaflow_extensions/netflix_ext/plugins/conda/utils.py +++ b/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]: