ob-metaflow 2.15.0.1__py2.py3-none-any.whl → 2.15.5.1__py2.py3-none-any.whl

metaflow/plugins/kubernetes/kubernetes_decorator.py

@@ -29,6 +29,7 @@ from metaflow.metaflow_config import (
     KUBERNETES_SHARED_MEMORY,
     KUBERNETES_TOLERATIONS,
     KUBERNETES_QOS,
+    KUBERNETES_CONDA_ARCH,
 )
 from metaflow.plugins.resources_decorator import ResourcesDecorator
 from metaflow.plugins.timeout_decorator import get_run_time_limit_for_task
@@ -160,7 +161,7 @@ class KubernetesDecorator(StepDecorator):
 
     # Conda environment support
     supports_conda_environment = True
-    target_platform = "linux-64"
+    target_platform = KUBERNETES_CONDA_ARCH or "linux-64"
 
     def init(self):
         super(KubernetesDecorator, self).init()

metaflow/plugins/kubernetes/kubernetes_jobsets.py

@@ -319,6 +319,8 @@ class RunningJobSet(object):
     def kill(self):
         plural = "jobsets"
         client = self._client.get()
+        if not (self.is_running or self.is_waiting):
+            return
         try:
             # Killing the control pod will trigger the jobset to mark everything as failed.
             # Since jobsets have a successPolicy set to `All` which ensures that everything has

metaflow/plugins/metadata_providers/service.py

@@ -350,23 +350,32 @@ class ServiceMetadataProvider(MetadataProvider):
         List[str]
             List of task pathspecs that satisfy the query
         """
-        query_params = {
-            "metadata_field_name": field_name,
-            "pattern": pattern,
-            "step_name": step_name,
-        }
+        query_params = {}
+
+        if pattern == ".*":
+            # we do not need to filter tasks at all if pattern allows 'any'
+            query_params = {}
+        else:
+            if field_name:
+                query_params["metadata_field_name"] = field_name
+            if pattern:
+                query_params["pattern"] = pattern
+
         url = ServiceMetadataProvider._obj_path(flow_name, run_id, step_name)
         url = f"{url}/filtered_tasks?{urlencode(query_params)}"
+
         try:
-            resp = cls._request(None, url, "GET")
+            resp, _ = cls._request(None, url, "GET")
         except Exception as e:
             if e.http_code == 404:
                 # filter_tasks_by_metadata endpoint does not exist in the version of metadata service
                 # deployed currently. Raise a more informative error message.
                 raise MetaflowInternalError(
                     "The version of metadata service deployed currently does not support filtering tasks by metadata. "
-                    "Upgrade Metadata service to version 2.15 or greater to use this feature."
+                    "Upgrade Metadata service to version 2.5.0 or greater to use this feature."
                 ) from e
+            # Other unknown exception
+            raise e
         return resp
 
     @staticmethod

metaflow/plugins/pypi/bootstrap.py

@@ -8,6 +8,7 @@ import subprocess
 import sys
 import tarfile
 import time
+import platform
 from urllib.error import URLError
 from urllib.request import urlopen
 from metaflow.metaflow_config import DATASTORE_LOCAL_DIR, CONDA_USE_FAST_INIT
@@ -36,29 +37,6 @@ def timer(func):
 
 
 if __name__ == "__main__":
-    # TODO: Detect architecture on the fly when dealing with arm architectures.
-    # ARCH=$(uname -m)
-    # OS=$(uname)
-
-    # if [[ "$OS" == "Linux" ]]; then
-    #     PLATFORM="linux"
-    #     if [[ "$ARCH" == "aarch64" ]]; then
-    #         ARCH="aarch64";
-    #     elif [[ $ARCH == "ppc64le" ]]; then
-    #         ARCH="ppc64le";
-    #     else
-    #         ARCH="64";
-    #     fi
-    # fi
-
-    # if [[ "$OS" == "Darwin" ]]; then
-    #     PLATFORM="osx";
-    #     if [[ "$ARCH" == "arm64" ]]; then
-    #         ARCH="arm64";
-    #     else
-    #         ARCH="64"
-    #     fi
-    # fi
 
     def run_cmd(cmd, stdin_str=None):
         result = subprocess.run(
@@ -350,12 +328,25 @@ if __name__ == "__main__":
         cmd = f"fast-initializer --prefix {prefix} --packages-dir {pkgs_dir}"
         run_cmd(cmd, all_package_urls)
 
-    if len(sys.argv) != 5:
-        print("Usage: bootstrap.py <flow_name> <id> <datastore_type> <architecture>")
+    if len(sys.argv) != 4:
+        print("Usage: bootstrap.py <flow_name> <id> <datastore_type>")
         sys.exit(1)
 
     try:
-        _, flow_name, id_, datastore_type, architecture = sys.argv
+        _, flow_name, id_, datastore_type = sys.argv
+
+        system = platform.system().lower()
+        arch_machine = platform.machine().lower()
+
+        if system == "darwin" and arch_machine == "arm64":
+            architecture = "osx-arm64"
+        elif system == "darwin":
+            architecture = "osx-64"
+        elif system == "linux" and arch_machine == "aarch64":
+            architecture = "linux-aarch64"
+        else:
+            # default fallback
+            architecture = "linux-64"
 
         prefix = os.path.join(os.getcwd(), architecture, id_)
         pkgs_dir = os.path.join(os.getcwd(), ".pkgs")

metaflow/plugins/pypi/conda_environment.py

@@ -190,7 +190,6 @@ class CondaEnvironment(MetaflowEnvironment):
         #  4. Start PyPI solves in parallel after each conda environment is created
         #  5. Download PyPI packages sequentially
         #  6. Create and cache PyPI environments in parallel
-
         with ThreadPoolExecutor() as executor:
             # Start all conda solves in parallel
             conda_futures = [
@@ -213,14 +212,14 @@ class CondaEnvironment(MetaflowEnvironment):
 
                 # Queue PyPI solve to start after conda create
                 if result[0] in pypi_envs:
+                    # solve pypi envs uniquely
+                    pypi_env = pypi_envs.pop(result[0])
 
                     def pypi_solve(env):
                         create_future.result()  # Wait for conda create
                         return solve(*env, "pypi")
 
-                    pypi_futures.append(
-                        executor.submit(pypi_solve, pypi_envs[result[0]])
-                    )
+                    pypi_futures.append(executor.submit(pypi_solve, pypi_env))
 
             # Process PyPI results sequentially for downloads
             for solve_future in pypi_futures:
@@ -242,7 +241,7 @@ class CondaEnvironment(MetaflowEnvironment):
         if id_:
             # bootstrap.py is responsible for ensuring the validity of this executable.
             # -s is important! Can otherwise leak packages to other environments.
-            return os.path.join("linux-64", id_, "bin/python -s")
+            return os.path.join("$MF_ARCH", id_, "bin/python -s")
         else:
             # for @conda/@pypi(disabled=True).
             return super().executable(step_name, default)
@@ -315,7 +314,6 @@ class CondaEnvironment(MetaflowEnvironment):
         # 5. All resolved packages (Conda or PyPI) are cached
         # 6. PyPI packages are only installed for local platform
 
-        # Resolve `linux-64` Conda environments if @batch or @kubernetes are in play
         target_platform = conda_platform()
         for decorator in step.decorators:
             # NOTE: Keep the list of supported decorator names for backward compatibility purposes.
@@ -329,7 +327,6 @@ class CondaEnvironment(MetaflowEnvironment):
                 "snowpark",
                 "slurm",
             ]:
-                # TODO: Support arm architectures
                 target_platform = getattr(decorator, "target_platform", "linux-64")
                 break
 
@@ -427,15 +424,18 @@ class CondaEnvironment(MetaflowEnvironment):
         if id_:
             return [
                 "echo 'Bootstrapping virtual environment...'",
+                "flush_mflogs",
                 # We have to prevent the tracing module from loading,
                 # as the bootstrapping process uses the internal S3 client which would fail to import tracing
                 # due to the required dependencies being bundled into the conda environment,
                 # which is yet to be initialized at this point.
-                'DISABLE_TRACING=True python -m metaflow.plugins.pypi.bootstrap "%s" %s "%s" linux-64'
+                'DISABLE_TRACING=True python -m metaflow.plugins.pypi.bootstrap "%s" %s "%s"'
                 % (self.flow.name, id_, self.datastore_type),
                 "echo 'Environment bootstrapped.'",
+                "flush_mflogs",
                 # To avoid having to install micromamba in the PATH in micromamba.py, we add it to the PATH here.
                 "export PATH=$PATH:$(pwd)/micromamba/bin",
+                "export MF_ARCH=$(case $(uname)/$(uname -m) in Darwin/arm64)echo osx-arm64;;Darwin/*)echo osx-64;;Linux/aarch64)echo linux-aarch64;;*)echo linux-64;;esac)",
             ]
         else:
             # for @conda/@pypi(disabled=True).

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}

metaflow/plugins/pypi/utils.py

@@ -72,6 +72,24 @@ def pip_tags(python_version, mamba_platform):
             )
         ]
         platforms.append("linux_x86_64")
+    elif mamba_platform == "linux-aarch64":
+        platforms = [
+            "manylinux%s_aarch64" % s
+            for s in (
+                "2014",
+                "_2_17",
+                "_2_18",
+                "_2_19",
+                "_2_20",
+                "_2_21",
+                "_2_23",
+                "_2_24",
+                "_2_25",
+                "_2_26",
+                "_2_27",
+            )
+        ]
+        platforms.append("linux_aarch64")
     elif mamba_platform == "osx-64":
         platforms = tags.mac_platforms(arch="x86_64")
     elif mamba_platform == "osx-arm64":

metaflow/runner/click_api.py

@@ -219,7 +219,11 @@ def get_inspect_param_obj(p: Union[click.Argument, click.Option], kind: str):
             default=inspect.Parameter.empty if is_vararg else p.default,
             annotation=annotation,
         ),
-        Optional[TTuple[annotation]] if is_vararg else annotation,
+        (
+            Optional[Union[TTuple[annotation], List[annotation]]]
+            if is_vararg
+            else annotation
+        ),
     )
 
 

metaflow/runner/deployer.py

@@ -106,7 +106,7 @@ class TriggeredRun(object):
         self.pathspec = content_json.get("pathspec")
         self.name = content_json.get("name")
 
-    def wait_for_run(self, timeout: Optional[int] = None):
+    def wait_for_run(self, check_interval: int = 5, timeout: Optional[int] = None):
         """
         Wait for the `run` property to become available.
 
@@ -115,6 +115,8 @@ class TriggeredRun(object):
 
         Parameters
         ----------
+        check_interval: int, default: 5
+            Frequency of checking for the `run` to become available, in seconds.
         timeout : int, optional, default None
             Maximum time to wait for the `run` to become available, in seconds. If
             None, wait indefinitely.
@@ -125,7 +127,6 @@ class TriggeredRun(object):
             If the `run` is not available within the specified timeout.
         """
         start_time = time.time()
-        check_interval = 5
         while True:
             if self.run is not None:
                 return self.run

metaflow/version.py

@@ -1 +1 @@
-metaflow_version = "2.15.0.1"
+metaflow_version = "2.15.5.1"