ob-metaflow 2.12.30.2__py2.py3-none-any.whl → 2.13.6.1__py2.py3-none-any.whl

metaflow/plugins/pypi/micromamba.py

@@ -1,12 +1,15 @@
+import functools
 import json
 import os
+import re
 import subprocess
 import tempfile
+import time
 
 from metaflow.exception import MetaflowException
 from metaflow.util import which
 
-from .utils import conda_platform
+from .utils import MICROMAMBA_MIRROR_URL, MICROMAMBA_URL, conda_platform
 
 
 class MicromambaException(MetaflowException):
@@ -19,8 +22,13 @@ class MicromambaException(MetaflowException):
         super(MicromambaException, self).__init__(msg)
 
 
+GLIBC_VERSION = os.environ.get("CONDA_OVERRIDE_GLIBC", "2.38")
+
+_double_equal_match = re.compile("==(?=[<=>!~])")
+
+
 class Micromamba(object):
-    def __init__(self):
+    def __init__(self, logger=None):
         # micromamba is a tiny version of the mamba package manager and comes with
         # metaflow specific performance enhancements.
 
@@ -33,6 +41,12 @@ class Micromamba(object):
             os.path.expanduser(_home),
             "micromamba",
         )
+
+        if logger:
+            self.logger = logger
+        else:
+            self.logger = lambda *args, **kwargs: None  # No-op logger if not provided
+
         self.bin = (
             which(os.environ.get("METAFLOW_PATH_TO_MICROMAMBA") or "micromamba")
             or which("./micromamba")  # to support remote execution
@@ -70,6 +84,9 @@ class Micromamba(object):
                 "MAMBA_ADD_PIP_AS_PYTHON_DEPENDENCY": "true",
                 "CONDA_SUBDIR": platform,
                 # "CONDA_UNSATISFIABLE_HINTS_CHECK_DEPTH": "0" # https://github.com/conda/conda/issues/9862
+                # Add a default glibc version for linux-64 environments (ignored for other platforms)
+                # TODO: Make the version configurable
+                "CONDA_OVERRIDE_GLIBC": GLIBC_VERSION,
             }
             cmd = [
                 "create",
@@ -78,6 +95,7 @@ class Micromamba(object):
                 "--dry-run",
                 "--no-extra-safety-checks",
                 "--repodata-ttl=86400",
+                "--safety-checks=disabled",
                 "--retry-clean-cache",
                 "--prefix=%s/prefix" % tmp_dir,
             ]
@@ -86,15 +104,17 @@ class Micromamba(object):
                 cmd.append("--channel=%s" % channel)
 
             for package, version in packages.items():
-                cmd.append("%s==%s" % (package, version))
+                version_string = "%s==%s" % (package, version)
+                cmd.append(_double_equal_match.sub("", version_string))
             if python:
                 cmd.append("python==%s" % python)
             # TODO: Ensure a human readable message is returned when the environment
             #       can't be resolved for any and all reasons.
-            return [
+            solved_packages = [
                 {k: v for k, v in item.items() if k in ["url"]}
                 for item in self._call(cmd, env)["actions"]["LINK"]
             ]
+            return solved_packages
 
     def download(self, id_, packages, python, platform):
         # Unfortunately all the packages need to be catalogued in package cache
@@ -103,8 +123,6 @@ class Micromamba(object):
         # Micromamba is painfully slow in determining if many packages are infact
         # already cached. As a perf heuristic, we check if the environment already
         # exists to short circuit package downloads.
-        if self.path_to_environment(id_, platform):
-            return
 
         prefix = "{env_dirs}/{keyword}/{platform}/{id}".format(
             env_dirs=self.info()["envs_dirs"][0],
@@ -113,13 +131,18 @@ class Micromamba(object):
             id=id_,
         )
 
-        # Another forced perf heuristic to skip cross-platform downloads.
+        # cheap check
         if os.path.exists(f"{prefix}/fake.done"):
             return
 
+        # somewhat expensive check
+        if self.path_to_environment(id_, platform):
+            return
+
         with tempfile.TemporaryDirectory() as tmp_dir:
             env = {
                 "CONDA_SUBDIR": platform,
+                "CONDA_OVERRIDE_GLIBC": GLIBC_VERSION,
             }
             cmd = [
                 "create",
@@ -159,6 +182,7 @@ class Micromamba(object):
             # use hardlinks when possible, otherwise copy files
             # disabled for now since it adds to environment creation latencies
             "CONDA_ALLOW_SOFTLINKS": "0",
+            "CONDA_OVERRIDE_GLIBC": GLIBC_VERSION,
         }
         cmd = [
             "create",
@@ -174,6 +198,7 @@ class Micromamba(object):
             cmd.append("{url}".format(**package))
         self._call(cmd, env)
 
+    @functools.lru_cache(maxsize=None)
     def info(self):
         return self._call(["config", "list", "-a"])
 
@@ -198,18 +223,24 @@ class Micromamba(object):
         }
         directories = self.info()["pkgs_dirs"]
         # search all package caches for packages
-        metadata = {
-            url: os.path.join(d, file)
+
+        file_to_path = {}
+        for d in directories:
+            if os.path.isdir(d):
+                try:
+                    with os.scandir(d) as entries:
+                        for entry in entries:
+                            if entry.is_file():
+                                # Prefer the first occurrence if the file exists in multiple directories
+                                file_to_path.setdefault(entry.name, entry.path)
+                except OSError:
+                    continue
+        ret = {
+            # set package tarball local paths to None if package tarballs are missing
+            url: file_to_path.get(file)
             for url, file in packages_to_filenames.items()
-            for d in directories
-            if os.path.isdir(d)
-            and file in os.listdir(d)
-            and os.path.isfile(os.path.join(d, file))
         }
-        # set package tarball local paths to None if package tarballs are missing
-        for url in packages_to_filenames:
-            metadata.setdefault(url, None)
-        return metadata
+        return ret
 
     def interpreter(self, id_):
         return os.path.join(self.path_to_environment(id_), "bin/python")
@@ -296,7 +327,7 @@ class Micromamba(object):
                         stderr="\n".join(err),
                     )
                 )
-            except (TypeError, ValueError) as ve:
+            except (TypeError, ValueError):
                 pass
             raise MicromambaException(
                 msg.format(
@@ -312,23 +343,37 @@ def _install_micromamba(installation_location):
     # Unfortunately no 32bit binaries are available for micromamba, which ideally
     # shouldn't be much of a problem in today's world.
     platform = conda_platform()
-    try:
-        subprocess.Popen(f"mkdir -p {installation_location}", shell=True).wait()
-        # https://mamba.readthedocs.io/en/latest/micromamba-installation.html#manual-installation
-        # requires bzip2
-        result = subprocess.Popen(
-            f"curl -Ls https://micro.mamba.pm/api/micromamba/{platform}/1.5.7 | tar -xvj -C {installation_location} bin/micromamba",
-            shell=True,
-            stderr=subprocess.PIPE,
-            stdout=subprocess.PIPE,
-        )
-        _, err = result.communicate()
-        if result.returncode != 0:
-            raise MicromambaException(
-                f"Micromamba installation '{result.args}' failed:\n{err.decode()}"
-            )
+    url = MICROMAMBA_URL.format(platform=platform, version="1.5.7")
+    mirror_url = MICROMAMBA_MIRROR_URL.format(platform=platform, version="1.5.7")
+    os.makedirs(installation_location, exist_ok=True)
 
-    except subprocess.CalledProcessError as e:
-        raise MicromambaException(
-            "Micromamba installation failed:\n{}".format(e.stderr.decode())
-        )
+    def _download_and_extract(url):
+        max_retries = 3
+        for attempt in range(max_retries):
+            try:
+                # https://mamba.readthedocs.io/en/latest/micromamba-installation.html#manual-installation
+                # requires bzip2
+                result = subprocess.Popen(
+                    f"curl -Ls {url} | tar -xvj -C {installation_location} bin/micromamba",
+                    shell=True,
+                    stderr=subprocess.PIPE,
+                    stdout=subprocess.PIPE,
+                )
+                _, err = result.communicate()
+                if result.returncode != 0:
+                    raise MicromambaException(
+                        f"Micromamba installation '{result.args}' failed:\n{err.decode()}"
+                    )
+            except subprocess.CalledProcessError as e:
+                if attempt == max_retries - 1:
+                    raise MicromambaException(
+                        "Micromamba installation failed:\n{}".format(e.stderr.decode())
+                    )
+                time.sleep(2**attempt)
+
+    try:
+        # prioritize downloading from mirror
+        _download_and_extract(mirror_url)
+    except Exception:
+        # download from official source as a fallback
+        _download_and_extract(url)

metaflow/plugins/pypi/pip.py

@@ -50,10 +50,14 @@ INSTALLATION_MARKER = "{prefix}/.pip/id"
 
 
 class Pip(object):
-    def __init__(self, micromamba=None):
+    def __init__(self, micromamba=None, logger=None):
         # pip is assumed to be installed inside a conda environment managed by
         # micromamba. pip commands are executed using `micromamba run --prefix`
-        self.micromamba = micromamba or Micromamba()
+        self.micromamba = micromamba or Micromamba(logger)
+        if logger:
+            self.logger = logger
+        else:
+            self.logger = lambda *args, **kwargs: None  # No-op logger if not provided
 
     def solve(self, id_, packages, python, platform):
         prefix = self.micromamba.path_to_environment(id_)
@@ -102,9 +106,8 @@ class Pip(object):
             except PipPackageNotFound as ex:
                 # pretty print package errors
                 raise PipException(
-                    "Could not find a binary distribution for %s \n"
-                    "for the platform %s\n\n"
-                    "Note that ***@pypi*** does not currently support source distributions"
+                    "Unable to find a binary distribution compatible with %s for %s.\n\n"
+                    "Note: ***@pypi*** does not currently support source distributions"
                     % (ex.package_spec, platform)
                 )
 
@@ -123,7 +126,7 @@ class Pip(object):
                         **res,
                         subdir_str=(
                             "#subdirectory=%s" % subdirectory if subdirectory else ""
-                        )
+                        ),
                     )
                     # used to deduplicate the storage location in case wheel does not
                     # build with enough unique identifiers.

metaflow/plugins/pypi/pypi_decorator.py

@@ -25,9 +25,10 @@ class PyPIStepDecorator(StepDecorator):
     defaults = {"packages": {}, "python": None, "disabled": None}  # wheels
 
     def __init__(self, attributes=None, statically_defined=False):
-        self._user_defined_attributes = (
-            attributes.copy() if attributes is not None else {}
+        self._attributes_with_user_values = (
+            set(attributes.keys()) if attributes is not None else set()
         )
+
         super().__init__(attributes, statically_defined)
 
     def step_init(self, flow, graph, step, decos, environment, flow_datastore, logger):
@@ -42,10 +43,9 @@ class PyPIStepDecorator(StepDecorator):
         if "pypi_base" in self.flow._flow_decorators:
             pypi_base = self.flow._flow_decorators["pypi_base"][0]
             super_attributes = pypi_base.attributes
-            self._user_defined_attributes = {
-                **self._user_defined_attributes,
-                **pypi_base._user_defined_attributes,
-            }
+            self._attributes_with_user_values.update(
+                pypi_base._attributes_with_user_values
+            )
             self.attributes["packages"] = {
                 **super_attributes["packages"],
                 **self.attributes["packages"],
@@ -106,7 +106,7 @@ class PyPIStepDecorator(StepDecorator):
         environment.set_local_root(LocalStorage.get_datastore_root_from_config(logger))
 
     def is_attribute_user_defined(self, name):
-        return name in self._user_defined_attributes
+        return name in self._attributes_with_user_values
 
 
 class PyPIFlowDecorator(FlowDecorator):
@@ -129,9 +129,10 @@ class PyPIFlowDecorator(FlowDecorator):
     defaults = {"packages": {}, "python": None, "disabled": None}
 
     def __init__(self, attributes=None, statically_defined=False):
-        self._user_defined_attributes = (
-            attributes.copy() if attributes is not None else {}
+        self._attributes_with_user_values = (
+            set(attributes.keys()) if attributes is not None else set()
         )
+
         super().__init__(attributes, statically_defined)
 
     def flow_init(
@@ -140,6 +141,7 @@ class PyPIFlowDecorator(FlowDecorator):
         from metaflow import decorators
 
         decorators._attach_decorators(flow, ["pypi"])
+        decorators._init(flow)
 
         # @pypi uses a conda environment to create a virtual environment.
         # The conda environment can be created through micromamba.

metaflow/plugins/pypi/utils.py

@@ -1,4 +1,3 @@
-import os
 import platform
 import sys
 
@@ -17,10 +16,13 @@ else:
     from metaflow._vendor.packaging import tags
     from metaflow._vendor.packaging.utils import parse_wheel_filename
 
-from urllib.parse import unquote, urlparse
+from urllib.parse import unquote
 
 from metaflow.exception import MetaflowException
 
+MICROMAMBA_URL = "https://micro.mamba.pm/api/micromamba/{platform}/{version}"
+MICROMAMBA_MIRROR_URL = "https://micromamba.outerbounds.sh/{platform}/{version}.tar.bz2"
+
 
 def conda_platform():
     # Returns the conda platform for the Python interpreter

metaflow/plugins/timeout_decorator.py

@@ -37,8 +37,8 @@ class TimeoutDecorator(StepDecorator):
     name = "timeout"
     defaults = {"seconds": 0, "minutes": 0, "hours": 0}
 
-    def __init__(self, *args, **kwargs):
-        super(TimeoutDecorator, self).__init__(*args, **kwargs)
+    def init(self):
+        super().init()
         # Initialize secs in __init__ so other decorators could safely use this
         # value without worrying about decorator order.
         # Convert values in attributes to type:int since they can be type:str