metaflow 2.13.10__py2.py3-none-any.whl → 2.14.1__py2.py3-none-any.whl

metaflow/plugins/metadata_providers/local.py

@@ -2,10 +2,12 @@ import collections
 import glob
 import json
 import os
+import re
 import random
 import tempfile
 import time
 from collections import namedtuple
+from typing import List
 
 from metaflow.exception import MetaflowInternalError, MetaflowTaggingError
 from metaflow.metadata_provider.metadata import ObjectOrder
@@ -202,6 +204,70 @@ class LocalMetadataProvider(MetadataProvider):
             "Tagging failed due to too many conflicting updates from other processes"
         )
 
+    @classmethod
+    def filter_tasks_by_metadata(
+        cls,
+        flow_name: str,
+        run_id: str,
+        step_name: str,
+        field_name: str,
+        pattern: str,
+    ) -> List[str]:
+        """
+        Filter tasks by metadata field and pattern, returning task pathspecs that match criteria.
+
+        Parameters
+        ----------
+        flow_name : str
+            Identifier for the flow
+        run_id : str
+            Identifier for the run
+        step_name : str
+            Name of the step to query tasks from
+        field_name : str
+            Name of metadata field to query
+        pattern : str
+            Pattern to match in metadata field value
+
+        Returns
+        -------
+        List[str]
+            List of task pathspecs that match the query criteria
+        """
+        tasks = cls.get_object("step", "task", {}, None, flow_name, run_id, step_name)
+        if not tasks:
+            return []
+
+        regex = re.compile(pattern)
+        matching_task_pathspecs = []
+
+        for task in tasks:
+            task_id = task.get("task_id")
+            if not task_id:
+                continue
+
+            if pattern == ".*":
+                # If the pattern is ".*", we can match all tasks without reading metadata
+                matching_task_pathspecs.append(
+                    f"{flow_name}/{run_id}/{step_name}/{task_id}"
+                )
+                continue
+
+            metadata = cls.get_object(
+                "task", "metadata", {}, None, flow_name, run_id, step_name, task_id
+            )
+
+            if any(
+                meta.get("field_name") == field_name
+                and regex.match(meta.get("value", ""))
+                for meta in metadata
+            ):
+                matching_task_pathspecs.append(
+                    f"{flow_name}/{run_id}/{step_name}/{task_id}"
+                )
+
+        return matching_task_pathspecs
+
     @classmethod
     def _get_object_internal(
         cls, obj_type, obj_order, sub_type, sub_order, filters, attempt, *args

metaflow/plugins/metadata_providers/service.py

@@ -1,23 +1,20 @@
 import os
 import random
+import time
 
 import requests
-import time
 
+from typing import List
 from metaflow.exception import (
     MetaflowException,
-    MetaflowTaggingError,
     MetaflowInternalError,
-)
-from metaflow.metaflow_config import (
-    SERVICE_RETRY_COUNT,
-    SERVICE_HEADERS,
-    SERVICE_URL,
+    MetaflowTaggingError,
 )
 from metaflow.metadata_provider import MetadataProvider
 from metaflow.metadata_provider.heartbeat import HB_URL_KEY
+from metaflow.metaflow_config import SERVICE_HEADERS, SERVICE_RETRY_COUNT, SERVICE_URL
 from metaflow.sidecar import Message, MessageTypes, Sidecar
-
+from urllib.parse import urlencode
 from metaflow.util import version_parse
 
 
@@ -39,6 +36,23 @@ class ServiceException(MetaflowException):
 class ServiceMetadataProvider(MetadataProvider):
     TYPE = "service"
 
+    _session = requests.Session()
+    _session.mount(
+        "http://",
+        requests.adapters.HTTPAdapter(
+            pool_connections=20,
+            pool_maxsize=20,
+            max_retries=0,  # Handle retries explicitly
+            pool_block=False,
+        ),
+    )
+    _session.mount(
+        "https://",
+        requests.adapters.HTTPAdapter(
+            pool_connections=20, pool_maxsize=20, max_retries=0, pool_block=False
+        ),
+    )
+
     _supports_attempt_gets = None
     _supports_tag_mutation = None
 
@@ -59,7 +73,9 @@ class ServiceMetadataProvider(MetadataProvider):
     def compute_info(cls, val):
         v = val.rstrip("/")
         try:
-            resp = requests.get(os.path.join(v, "ping"), headers=SERVICE_HEADERS.copy())
+            resp = cls._session.get(
+                os.path.join(v, "ping"), headers=SERVICE_HEADERS.copy()
+            )
             resp.raise_for_status()
         except:  # noqa E722
             raise ValueError("Metaflow service [%s] unreachable." % v)
@@ -304,6 +320,55 @@ class ServiceMetadataProvider(MetadataProvider):
             self._register_system_metadata(run_id, step_name, task["task_id"], attempt)
         return task["task_id"], did_create
 
+    @classmethod
+    def filter_tasks_by_metadata(
+        cls,
+        flow_name: str,
+        run_id: str,
+        step_name: str,
+        field_name: str,
+        pattern: str,
+    ) -> List[str]:
+        """
+        Filter tasks by metadata field and pattern, returning task pathspecs that match criteria.
+
+        Parameters
+        ----------
+        flow_name : str
+            Flow name, that the run belongs to.
+        run_id: str
+            Run id, together with flow_id, that identifies the specific Run whose tasks to query
+        step_name: str
+            Step name to query tasks from
+        field_name: str
+            Metadata field name to query
+        pattern: str
+            Pattern to match in metadata field value
+
+        Returns
+        -------
+        List[str]
+            List of task pathspecs that satisfy the query
+        """
+        query_params = {
+            "metadata_field_name": field_name,
+            "pattern": pattern,
+            "step_name": step_name,
+        }
+        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")
+        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."
+                ) from e
+        return resp
+
     @staticmethod
     def _obj_path(
         flow_name,
@@ -412,27 +477,27 @@ class ServiceMetadataProvider(MetadataProvider):
                 if method == "GET":
                     if monitor:
                         with monitor.measure("metaflow.service_metadata.get"):
-                            resp = requests.get(url, headers=SERVICE_HEADERS.copy())
+                            resp = cls._session.get(url, headers=SERVICE_HEADERS.copy())
                     else:
-                        resp = requests.get(url, headers=SERVICE_HEADERS.copy())
+                        resp = cls._session.get(url, headers=SERVICE_HEADERS.copy())
                 elif method == "POST":
                     if monitor:
                         with monitor.measure("metaflow.service_metadata.post"):
-                            resp = requests.post(
+                            resp = cls._session.post(
                                 url, headers=SERVICE_HEADERS.copy(), json=data
                             )
                     else:
-                        resp = requests.post(
+                        resp = cls._session.post(
                             url, headers=SERVICE_HEADERS.copy(), json=data
                         )
                 elif method == "PATCH":
                     if monitor:
                         with monitor.measure("metaflow.service_metadata.patch"):
-                            resp = requests.patch(
+                            resp = cls._session.patch(
                                 url, headers=SERVICE_HEADERS.copy(), json=data
                             )
                     else:
-                        resp = requests.patch(
+                        resp = cls._session.patch(
                             url, headers=SERVICE_HEADERS.copy(), json=data
                         )
                 else:
@@ -475,7 +540,6 @@ class ServiceMetadataProvider(MetadataProvider):
                         resp.text,
                     )
             time.sleep(2**i)
-
         if resp:
             raise ServiceException(
                 "Metadata request (%s) failed (code %s): %s"
@@ -499,9 +563,9 @@ class ServiceMetadataProvider(MetadataProvider):
             try:
                 if monitor:
                     with monitor.measure("metaflow.service_metadata.get"):
-                        resp = requests.get(url, headers=SERVICE_HEADERS.copy())
+                        resp = cls._session.get(url, headers=SERVICE_HEADERS.copy())
                 else:
-                    resp = requests.get(url, headers=SERVICE_HEADERS.copy())
+                    resp = cls._session.get(url, headers=SERVICE_HEADERS.copy())
             except:
                 if monitor:
                     with monitor.count("metaflow.service_metadata.failed_request"):

metaflow/plugins/pypi/bootstrap.py

@@ -10,7 +10,7 @@ import tarfile
 import time
 from urllib.error import URLError
 from urllib.request import urlopen
-from metaflow.metaflow_config import DATASTORE_LOCAL_DIR, CONDA_USE_FAST_INIT
+from metaflow.metaflow_config import DATASTORE_LOCAL_DIR
 from metaflow.plugins import DATASTORES
 from metaflow.plugins.pypi.utils import MICROMAMBA_MIRROR_URL, MICROMAMBA_URL
 from metaflow.util import which
@@ -329,8 +329,6 @@ if __name__ == "__main__":
 
     @timer
     def fast_setup_environment(architecture, storage, env, prefix, pkgs_dir):
-        install_fast_initializer(architecture)
-
         # Get package urls
         conda_pkgs = env["conda"]
         pypi_pkgs = env.get("pypi", [])
@@ -381,7 +379,9 @@ if __name__ == "__main__":
         with open(os.path.join(manifest_dir, MAGIC_FILE)) as f:
             env = json.load(f)[id_][architecture]
 
-        if CONDA_USE_FAST_INIT:
+        if datastore_type == "s3":
+            # TODO: Remove this once fast-initializer is ready for all datastores
+            install_fast_initializer(architecture)
             fast_setup_environment(architecture, storage, env, prefix, pkgs_dir)
         else:
             setup_environment(

metaflow/runner/click_api.py

@@ -97,12 +97,13 @@ def _method_sanity_check(
             check_type(supplied_v, annotations[supplied_k])
         except TypeCheckError:
             raise TypeError(
-                "Invalid type for '%s' (%s), expected: '%s', default is '%s'"
+                "Invalid type for '%s' (%s), expected: '%s', default is '%s' but found '%s'"
                 % (
                     supplied_k,
                     type(supplied_k),
                     annotations[supplied_k],
                     defaults[supplied_k],
+                    str(supplied_v),
                 )
             )
 
@@ -218,7 +219,7 @@ 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,
         ),
-        annotation,
+        Optional[TTuple[annotation]] if is_vararg else annotation,
     )
 
 
@@ -392,7 +393,9 @@ class MetaflowAPI(object):
                 options = params.pop("options", {})
 
                 for _, v in args.items():
-                    if isinstance(v, list):
+                    if v is None:
+                        continue
+                    if isinstance(v, (list, tuple)):
                         for i in v:
                             components.append(i)
                     else:
@@ -527,12 +530,6 @@ def extract_all_params(cmd_obj: Union[click.Command, click.Group]):
             )
             arg_parameters[each_param.name] = each_param
         elif isinstance(each_param, click.Option):
-            if each_param.hidden:
-                # Skip hidden options because users should not be setting those.
-                # These are typically internal only options (used by the Runner in part
-                # for example to pass state files or configs to pass local-config-file).
-                continue
-
             opt_params_sigs[each_param.name], annotations[each_param.name] = (
                 get_inspect_param_obj(each_param, inspect.Parameter.KEYWORD_ONLY)
             )

metaflow/sidecar/sidecar_worker.py

@@ -49,7 +49,7 @@ def process_messages(worker_type, worker):
 
 
 @click.command(help="Initialize workers")
-@tracing.cli_entrypoint("sidecar")
+@tracing.cli("sidecar")
 @click.argument("worker-type")
 def main(worker_type):
     sidecar_type = SIDECARS.get(worker_type)

metaflow/task.py

@@ -493,6 +493,25 @@ class MetaflowTask(object):
                     )
                 )
 
+            # Add runtime dag information to the metadata of the task
+            foreach_execution_path = ",".join(
+                [
+                    "{}:{}".format(foreach_frame.step, foreach_frame.index)
+                    for foreach_frame in foreach_stack
+                ]
+            )
+            if foreach_execution_path:
+                metadata.extend(
+                    [
+                        MetaDatum(
+                            field="foreach-execution-path",
+                            value=foreach_execution_path,
+                            type="foreach-execution-path",
+                            tags=metadata_tags,
+                        ),
+                    ]
+                )
+
         self.metadata.register_metadata(
             run_id,
             step_name,
@@ -559,6 +578,7 @@ class MetaflowTask(object):
                 self.flow._success = False
                 self.flow._task_ok = None
                 self.flow._exception = None
+
                 # Note: All internal flow attributes (ie: non-user artifacts)
                 # should either be set prior to running the user code or listed in
                 # FlowSpec._EPHEMERAL to allow for proper merging/importing of
@@ -616,7 +636,6 @@ class MetaflowTask(object):
                                 "graph_info": self.flow._graph_info,
                             }
                         )
-
                 for deco in decorators:
                     deco.task_pre_step(
                         step_name,
@@ -728,7 +747,7 @@ class MetaflowTask(object):
                                 value=attempt_ok,
                                 type="internal_attempt_status",
                                 tags=["attempt_id:{0}".format(retry_count)],
-                            )
+                            ),
                         ],
                     )
 

metaflow/tracing/__init__.py

@@ -20,15 +20,15 @@ def post_fork():
     yield
 
 
-def cli_entrypoint(name: str):
-    def cli_entrypoint_wrap(func):
+def cli(name: str):
+    def cli_wrap(func):
         @wraps(func)
         def wrapper_func(*args, **kwargs):
             return func(*args, **kwargs)
 
         return wrapper_func
 
-    return cli_entrypoint_wrap
+    return cli_wrap
 
 
 def inject_tracing_vars(env_dict: Dict[str, str]) -> Dict[str, str]:
@@ -40,7 +40,9 @@ def get_trace_id() -> str:
 
 
 @contextlib.contextmanager
-def traced(name, attrs={}):
+def traced(name, attrs=None):
+    if attrs is None:
+        attrs = {}
     yield
 
 
@@ -54,17 +56,15 @@ def tracing(func):
 
 if not DISABLE_TRACING and (CONSOLE_TRACE_ENABLED or OTEL_ENDPOINT or ZIPKIN_ENDPOINT):
     try:
-        # Overrides No-Op implementations if a specific provider is configured.
         from .tracing_modules import (
             init_tracing,
             post_fork,
-            cli_entrypoint,
+            cli,
             inject_tracing_vars,
             get_trace_id,
             traced,
             tracing,
         )
-
     except ImportError as e:
         # We keep the errors silent by default so that having tracing environment variables present
         # does not affect users with no need for tracing.

metaflow/tracing/span_exporter.py

@@ -3,60 +3,53 @@ from metaflow.metaflow_config import (
     OTEL_ENDPOINT,
     ZIPKIN_ENDPOINT,
     CONSOLE_TRACE_ENABLED,
+    SERVICE_AUTH_KEY,
+    SERVICE_HEADERS,
 )
 
 
 def get_span_exporter():
-    if OTEL_ENDPOINT:
-        return set_otel_exporter()
+    exporter_map = {
+        OTEL_ENDPOINT: _create_otel_exporter,
+        ZIPKIN_ENDPOINT: _create_zipkin_exporter,
+        CONSOLE_TRACE_ENABLED: _create_console_exporter,
+    }
 
-    elif ZIPKIN_ENDPOINT:
-        return set_zipkin_exporter()
+    for config, create_exporter in exporter_map.items():
+        if config:
+            return create_exporter()
 
-    elif CONSOLE_TRACE_ENABLED:
-        return set_console_exporter()
-    else:
-        print("WARNING: endpoints not set up for Opentelemetry", file=sys.stderr)
-        return
+    print("WARNING: endpoints not set up for OpenTelemetry", file=sys.stderr)
+    return None
 
 
-def set_otel_exporter():
+def _create_otel_exporter():
     from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
 
-    from metaflow.metaflow_config import (
-        SERVICE_AUTH_KEY,
-        SERVICE_HEADERS,
-    )
+    if not any([SERVICE_AUTH_KEY, SERVICE_HEADERS]):
+        print("WARNING: no auth settings for OpenTelemetry", file=sys.stderr)
+        return None
+
+    config = {
+        "endpoint": OTEL_ENDPOINT,
+        "timeout": 1,
+    }
 
     if SERVICE_AUTH_KEY:
-        span_exporter = OTLPSpanExporter(
-            endpoint=OTEL_ENDPOINT,
-            headers={"x-api-key": SERVICE_AUTH_KEY},
-            timeout=1,
-        )
+        config["headers"] = {"x-api-key": SERVICE_AUTH_KEY}
     elif SERVICE_HEADERS:
-        span_exporter = OTLPSpanExporter(
-            endpoint=OTEL_ENDPOINT,
-            headers=SERVICE_HEADERS,
-            timeout=1,
-        )
-    else:
-        print("WARNING: no auth settings for Opentelemetry", file=sys.stderr)
-        return
-    return span_exporter
-
-
-def set_zipkin_exporter():
+        config["headers"] = SERVICE_HEADERS
+
+    return OTLPSpanExporter(**config)
+
+
+def _create_zipkin_exporter():
     from opentelemetry.exporter.zipkin.proto.http import ZipkinExporter
 
-    span_exporter = ZipkinExporter(
-        endpoint=ZIPKIN_ENDPOINT,
-    )
-    return span_exporter
+    return ZipkinExporter(endpoint=ZIPKIN_ENDPOINT)
 
 
-def set_console_exporter():
+def _create_console_exporter():
     from opentelemetry.sdk.trace.export import ConsoleSpanExporter
 
-    span_exporter = ConsoleSpanExporter()
-    return span_exporter
+    return ConsoleSpanExporter()