ob-metaflow 2.11.13.1__py2.py3-none-any.whl → 2.19.7.1rc0__py2.py3-none-any.whl

metaflow/plugins/pypi/micromamba.py

@@ -1,12 +1,18 @@
+import functools
 import json
 import os
+import re
+import shutil
 import subprocess
 import tempfile
+import time
 
+from metaflow.debug import debug
 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
+from threading import Lock
 
 
 class MicromambaException(MetaflowException):
@@ -19,8 +25,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, force_rebuild=False):
         # micromamba is a tiny version of the mamba package manager and comes with
         # metaflow specific performance enhancements.
 
@@ -29,26 +40,53 @@ class Micromamba(object):
             _home = os.environ.get("METAFLOW_TOKEN_HOME")
         else:
             _home = os.environ.get("METAFLOW_HOME", "~/.metaflowconfig")
-        _path_to_hidden_micromamba = os.path.join(
+        self._path_to_hidden_micromamba = os.path.join(
             os.path.expanduser(_home),
             "micromamba",
         )
-        self.bin = (
+
+        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
             or which("./bin/micromamba")
-            or which(os.path.join(_path_to_hidden_micromamba, "bin/micromamba"))
+            or which(os.path.join(self._path_to_hidden_micromamba, "bin/micromamba"))
         )
-        if self.bin is None:
+
+        # We keep a mutex as environments are resolved in parallel,
+        # which causes a race condition in case micromamba needs to be installed first.
+        self.install_mutex = Lock()
+
+        self.force_rebuild = force_rebuild
+
+    @property
+    def bin(self) -> str:
+        "Defer installing Micromamba until when the binary path is actually requested"
+        if self._bin is not None:
+            return self._bin
+        with self.install_mutex:
+            # another check as micromamba might have been installed when the mutex is released.
+            if self._bin is not None:
+                return self._bin
+
             # Install Micromamba on the fly.
             # TODO: Make this optional at some point.
-            _install_micromamba(_path_to_hidden_micromamba)
-            self.bin = which(os.path.join(_path_to_hidden_micromamba, "bin/micromamba"))
+            debug.conda_exec("No Micromamba binary found. Installing micromamba")
+            _install_micromamba(self._path_to_hidden_micromamba)
+            self._bin = which(
+                os.path.join(self._path_to_hidden_micromamba, "bin/micromamba")
+            )
+
+            if self._bin is None:
+                msg = "No installation for *Micromamba* found.\n"
+                msg += "Visit https://mamba.readthedocs.io/en/latest/micromamba-installation.html for installation instructions."
+                raise MetaflowException(msg)
 
-        if self.bin is None:
-            msg = "No installation for *Micromamba* found.\n"
-            msg += "Visit https://mamba.readthedocs.io/en/latest/micromamba-installation.html for installation instructions."
-            raise MetaflowException(msg)
+        return self._bin
 
     def solve(self, id_, packages, python, platform):
         # Performance enhancements
@@ -65,11 +103,15 @@ class Micromamba(object):
         #    environment
         # 4. Multiple solves can progress at the same time while relying on the same
         #    index
+        debug.conda_exec("Solving packages for conda environment %s" % id_)
         with tempfile.TemporaryDirectory() as tmp_dir:
             env = {
                 "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 +120,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 +129,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 +148,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],
@@ -112,14 +155,25 @@ class Micromamba(object):
             keyword="metaflow",  # indicates metaflow generated environment
             id=id_,
         )
+        # If we are forcing a rebuild of the environment, we make sure to remove existing files beforehand.
+        # This is to ensure that no irrelevant packages get bundled relative to the resolved environment.
+        # NOTE: download always happens before create, so we want to do the cleanup here instead.
+        if self.force_rebuild:
+            shutil.rmtree(self.path_to_environment(id_, platform), ignore_errors=True)
 
-        # 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
+
+        debug.conda_exec("Downloading packages for conda environment %s" % id_)
         with tempfile.TemporaryDirectory() as tmp_dir:
             env = {
                 "CONDA_SUBDIR": platform,
+                "CONDA_OVERRIDE_GLIBC": GLIBC_VERSION,
             }
             cmd = [
                 "create",
@@ -148,6 +202,7 @@ class Micromamba(object):
         if platform != self.platform() or self.path_to_environment(id_, platform):
             return
 
+        debug.conda_exec("Creating local Conda environment %s" % id_)
         prefix = "{env_dirs}/{keyword}/{platform}/{id}".format(
             env_dirs=self.info()["envs_dirs"][0],
             platform=platform,
@@ -159,6 +214,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 +230,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 +255,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")
@@ -253,7 +316,40 @@ class Micromamba(object):
             try:
                 output = json.loads(e.output)
                 err = []
+                v_pkgs = ["__cuda", "__glibc"]
                 for error in output.get("solver_problems", []):
+                    # raise a specific error message for virtual package related errors
+                    match = next((p for p in v_pkgs if p in error), None)
+                    if match is not None:
+                        vpkg_name = match[2:]
+                        # try to strip version from error msg which are of the format:
+                        # nothing provides <__vpkg> >=2.17,<3.0.a0 needed by <pkg_name>
+                        try:
+                            vpkg_version = error[
+                                len("nothing provides %s " % match) : error.index(
+                                    " needed by"
+                                )
+                            ]
+                        except ValueError:
+                            vpkg_version = None
+                        raise MicromambaException(
+                            "{msg}\n\n"
+                            "*Please set the environment variable CONDA_OVERRIDE_{var} to a specific version{version} of {name}.*\n\n"
+                            "Here is an example of supplying environment variables through the command line\n"
+                            "CONDA_OVERRIDE_{var}=<{name}-version> python flow.py <args>".format(
+                                msg=msg.format(
+                                    cmd=" ".join(e.cmd),
+                                    code=e.returncode,
+                                    output=e.output.decode(),
+                                    stderr=error,
+                                ),
+                                var=vpkg_name.upper(),
+                                version=(
+                                    "" if not vpkg_version else f" ({vpkg_version})"
+                                ),
+                                name=vpkg_name,
+                            )
+                        )
                     err.append(error)
                 raise MicromambaException(
                     msg.format(
@@ -263,7 +359,7 @@ class Micromamba(object):
                         stderr="\n".join(err),
                     )
                 )
-            except (TypeError, ValueError) as ve:
+            except (TypeError, ValueError):
                 pass
             raise MicromambaException(
                 msg.format(
@@ -279,23 +375,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}/latest | 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/parsers.py

@@ -0,0 +1,268 @@
+# this file can be overridden by extensions as is (e.g. metaflow-nflx-extensions)
+from metaflow.exception import MetaflowException
+
+
+class ParserValueError(MetaflowException):
+    headline = "Value error"
+
+
+def requirements_txt_parser(content: str):
+    """
+    Parse non-comment lines from a requirements.txt file as strictly valid
+    PEP 508 requirements.
+
+    Recognizes direct references (e.g. "my_lib @ git+https://..."), extras
+    (e.g. "requests[security]"), and version specifiers (e.g. "==2.0"). If
+    the package name is "python", its specifier is stored in the "python"
+    key instead of "packages".
+
+    Parameters
+    ----------
+    content : str
+        Contents of a requirements.txt file.
+
+    Returns
+    -------
+    dict
+        A dictionary with two keys:
+            - "packages": dict(str -> str)
+              Mapping from package name (plus optional extras/references) to a
+              version specifier string.
+            - "python": str or None
+              The Python version constraints if present, otherwise None.
+
+    Raises
+    ------
+     ParserValueError
+        If a requirement line is invalid PEP 508 or if environment markers are
+        detected, or if multiple Python constraints are specified.
+    """
+    import re
+    from metaflow._vendor.packaging.requirements import Requirement, InvalidRequirement
+
+    parsed = {"packages": {}, "python": None}
+
+    inline_comment_pattern = re.compile(r"\s+#.*$")
+    for line in content.splitlines():
+        line = line.strip()
+
+        # support Rye lockfiles by skipping lines not compliant with requirements
+        if line == "-e file:.":
+            continue
+
+        if not line or line.startswith("#"):
+            continue
+
+        line = inline_comment_pattern.sub("", line).strip()
+        if not line:
+            continue
+
+        try:
+            req = Requirement(line)
+        except InvalidRequirement:
+            raise ParserValueError(f"Not a valid PEP 508 requirement: '{line}'")
+
+        if req.marker is not None:
+            raise ParserValueError(
+                "Environment markers (e.g. 'platform_system==\"Linux\"') "
+                f"are not supported for line: '{line}'"
+            )
+
+        dep_key = req.name
+        if req.extras:
+            dep_key += f"[{','.join(req.extras)}]"
+        if req.url:
+            dep_key += f"@{req.url}"
+
+        dep_spec = str(req.specifier).lstrip(" =")
+
+        if req.name.lower() == "python":
+            if parsed["python"] is not None and dep_spec:
+                raise ParserValueError(
+                    f"Multiple Python version specs not allowed: '{line}'"
+                )
+            parsed["python"] = dep_spec or None
+        else:
+            parsed["packages"][dep_key] = dep_spec
+
+    return parsed
+
+
+def pyproject_toml_parser(content: str):
+    """
+    Parse a pyproject.toml file per PEP 621.
+
+    Reads the 'requires-python' and 'dependencies' fields from the "[project]" section.
+    Each dependency line must be a valid PEP 508 requirement. If the package name is
+    "python", its specifier is stored in the "python" key instead of "packages".
+
+    Parameters
+    ----------
+    content : str
+        Contents of a pyproject.toml file.
+
+    Returns
+    -------
+    dict
+        A dictionary with two keys:
+            - "packages": dict(str -> str)
+              Mapping from package name (plus optional extras/references) to a
+              version specifier string.
+            - "python": str or None
+              The Python version constraints if present, otherwise None.
+
+    Raises
+    ------
+    RuntimeError
+        If no TOML library (tomllib in Python 3.11+ or tomli in earlier versions) is found.
+     ParserValueError
+        If a dependency is not valid PEP 508, if environment markers are used, or if
+        multiple Python constraints are specified.
+    """
+    try:
+        import tomllib as toml  # Python 3.11+
+    except ImportError:
+        try:
+            import tomli as toml  # Python < 3.11 (requires "tomli" package)
+        except ImportError:
+            raise RuntimeError(
+                "Could not import a TOML library. For Python <3.11, please install 'tomli'."
+            )
+    from metaflow._vendor.packaging.requirements import Requirement, InvalidRequirement
+
+    data = toml.loads(content)
+
+    project = data.get("project", {})
+    requirements = project.get("dependencies", [])
+    requires_python = project.get("requires-python")
+
+    parsed = {"packages": {}, "python": None}
+
+    if requires_python is not None:
+        # If present, store verbatim; note that PEP 621 does not necessarily
+        # require "python" to be a dependency in the usual sense.
+        # Example: "requires-python" = ">=3.7,<4"
+        parsed["python"] = requires_python.lstrip("=").strip()
+
+    for dep_line in requirements:
+        dep_line_stripped = dep_line.strip()
+        try:
+            req = Requirement(dep_line_stripped)
+        except InvalidRequirement:
+            raise ParserValueError(
+                f"Not a valid PEP 508 requirement: '{dep_line_stripped}'"
+            )
+
+        if req.marker is not None:
+            raise ParserValueError(
+                f"Environment markers not supported for line: '{dep_line_stripped}'"
+            )
+
+        dep_key = req.name
+        if req.extras:
+            dep_key += f"[{','.join(req.extras)}]"
+        if req.url:
+            dep_key += f"@{req.url}"
+
+        dep_spec = str(req.specifier).lstrip("=")
+
+        if req.name.lower() == "python":
+            if parsed["python"] is not None and dep_spec:
+                raise ParserValueError(
+                    f"Multiple Python version specs not allowed: '{dep_line_stripped}'"
+                )
+            parsed["python"] = dep_spec or None
+        else:
+            parsed["packages"][dep_key] = dep_spec
+
+    return parsed
+
+
+def conda_environment_yml_parser(content: str):
+    """
+    Parse a minimal environment.yml file under strict assumptions.
+
+    The file must contain a 'dependencies:' line, after which each dependency line
+    appears with a '- ' prefix. Python can appear as 'python=3.9', etc.; other
+    packages as 'numpy=1.21.2' or simply 'numpy'. Non-compliant lines raise  ParserValueError.
+
+    Parameters
+    ----------
+    content : str
+        Contents of a environment.yml file.
+
+    Returns
+    -------
+    dict
+        A dictionary with keys:
+        {
+            "packages": dict(str -> str),
+            "python": str or None
+        }
+
+    Raises
+    ------
+     ParserValueError
+        If the file has malformed lines or unsupported sections.
+    """
+    import re
+
+    packages = {}
+    python_version = None
+
+    inside_dependencies = False
+
+    # Basic pattern for lines like "numpy=1.21.2"
+    # Group 1: package name
+    # Group 2: optional operator + version (could be "=1.21.2", "==1.21.2", etc.)
+    line_regex = re.compile(r"^([A-Za-z0-9_\-\.]+)(\s*[=<>!~].+\s*)?$")
+    inline_comment_pattern = re.compile(r"\s+#.*$")
+
+    for line in content.splitlines():
+        line = line.strip()
+        if not line or line.startswith("#"):
+            continue
+
+        line = inline_comment_pattern.sub("", line).strip()
+        if not line:
+            continue
+
+        if line.lower().startswith("dependencies:"):
+            inside_dependencies = True
+            continue
+
+        if inside_dependencies and not line.startswith("-"):
+            inside_dependencies = False
+            continue
+
+        if not inside_dependencies:
+            continue
+
+        dep_line = line.lstrip("-").strip()
+        if dep_line.endswith(":"):
+            raise ParserValueError(
+                f"Unsupported subsection '{dep_line}' in environment.yml."
+            )
+
+        match = line_regex.match(dep_line)
+        if not match:
+            raise ParserValueError(
+                f"Line '{dep_line}' is not a valid conda package specifier."
+            )
+
+        pkg_name, pkg_version_part = match.groups()
+        version_spec = pkg_version_part.strip() if pkg_version_part else ""
+
+        if version_spec.startswith("="):
+            version_spec = version_spec.lstrip("=").strip()
+
+        if pkg_name.lower() == "python":
+            if python_version is not None and version_spec:
+                raise ParserValueError(
+                    f"Multiple Python version specs detected: '{dep_line}'"
+                )
+            python_version = version_spec
+        else:
+            packages[pkg_name] = version_spec
+
+    return {"packages": packages, "python": python_version}