metaflow 2.14.0__py2.py3-none-any.whl → 2.14.2__py2.py3-none-any.whl

metaflow/plugins/env_escape/stub.py

@@ -276,9 +276,22 @@ class MetaWithConnection(StubMetaClass):
         if len(args) > 0 and id(args[0]) == id(cls.___class_connection___):
             return super(MetaWithConnection, cls).__call__(*args, **kwargs)
         else:
-            return cls.___class_connection___.stub_request(
-                None, OP_INIT, cls.___class_remote_class_name___, *args, **kwargs
-            )
+            if hasattr(cls, "__overriden_init__"):
+                return cls.__overriden_init__(
+                    None,
+                    functools.partial(
+                        cls.___class_connection___.stub_request,
+                        None,
+                        OP_INIT,
+                        cls.___class_remote_class_name___,
+                    ),
+                    *args,
+                    **kwargs
+                )
+            else:
+                return cls.___class_connection___.stub_request(
+                    None, OP_INIT, cls.___class_remote_class_name___, *args, **kwargs
+                )
 
     def __subclasscheck__(cls, subclass):
         subclass_name = "%s.%s" % (subclass.__module__, subclass.__name__)
@@ -381,7 +394,10 @@ def create_class(
             name = name[7:]
             method_type = CLASS_METHOD
         if name in overriden_methods:
-            if method_type == NORMAL_METHOD:
+            if name == "__init__":
+                class_dict["__overriden_init__"] = overriden_methods["__init__"]
+
+            elif method_type == NORMAL_METHOD:
                 class_dict[name] = (
                     lambda override, orig_method: lambda obj, *args, **kwargs: override(
                         obj, functools.partial(orig_method, obj), *args, **kwargs
@@ -412,6 +428,7 @@ def create_class(
             class_dict[name] = _make_method(
                 method_type, connection, class_name, name, doc
             )
+
     # Check for any getattr/setattr overrides
     special_attributes = set(getattr_overrides.keys())
     special_attributes.update(set(setattr_overrides.keys()))

metaflow/plugins/kubernetes/kubernetes_cli.py

@@ -39,7 +39,7 @@ def kubernetes():
     "command inside a Kubernetes pod with the given options. Typically you do not call "
     "this command directly; it is used internally by Metaflow."
 )
-@tracing.cli_entrypoint("kubernetes/step")
+@tracing.cli("kubernetes/step")
 @click.argument("step-name")
 @click.argument("code-package-sha")
 @click.argument("code-package-url")

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

@@ -4,6 +4,7 @@ import time
 
 import requests
 
+from typing import List
 from metaflow.exception import (
     MetaflowException,
     MetaflowInternalError,
@@ -13,6 +14,7 @@ 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
 
 
@@ -318,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,

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:

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()

metaflow/tracing/tracing_modules.py

@@ -13,45 +13,51 @@ from typing import Dict, List, Optional
 from opentelemetry import trace as trace_api, context
 from .span_exporter import get_span_exporter
 
-tracer_provider = None
+tracer_provider: Optional[TracerProvider] = None
 
 
 def init_tracing():
     global tracer_provider
     if tracer_provider is not None:
-        print("Tracing already initialized", file=sys.stderr)
         return
 
     from .propagator import EnvPropagator
 
     set_global_textmap(EnvPropagator(None))
-    span_exporter = get_span_exporter()
 
-    if "METAFLOW_KUBERNETES_POD_NAMESPACE" in os.environ:
-        service_name = "metaflow-kubernetes"
-    elif "AWS_BATCH_JOB_ID" in os.environ:
-        service_name = "metaflow-awsbatch"
-    else:
-        service_name = "metaflow-local"
+    span_exporter = get_span_exporter()
+    if span_exporter is None:
+        return
 
     tracer_provider = TracerProvider(
-        resource=Resource.create({SERVICE_NAME: service_name})
+        resource=Resource.create(
+            {
+                SERVICE_NAME: "metaflow",
+            }
+        )
     )
     trace_api.set_tracer_provider(tracer_provider)
 
     span_processor = BatchSpanProcessor(span_exporter)
     tracer_provider.add_span_processor(span_processor)
 
-    from opentelemetry.instrumentation.requests import RequestsInstrumentor
+    try:
+        from opentelemetry.instrumentation.requests import RequestsInstrumentor
 
-    RequestsInstrumentor().instrument()
+        RequestsInstrumentor().instrument(
+            tracer_provider=tracer_provider,
+        )
+    except ImportError:
+        pass
 
 
 @contextlib.contextmanager
 def post_fork():
     global tracer_provider
+
     tracer_provider = None
     init_tracing()
+
     token = context.attach(extract(os.environ))
     try:
         tracer = trace_api.get_tracer_provider().get_tracer(__name__)
@@ -59,47 +65,27 @@ def post_fork():
             "fork", kind=trace_api.SpanKind.SERVER
         ) as span:
             span.set_attribute("cmd", " ".join(sys.argv))
+            span.set_attribute("pid", str(os.getpid()))
             yield
     finally:
         context.detach(token)
 
 
-def _extract_token_after(tokens: List[str], before_token: str) -> Optional[str]:
-    for i, tok in enumerate(tokens):
-        if i > 0 and tokens[i - 1] == before_token:
-            return tok
-
-
-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):
             global tracer_provider
-
             init_tracing()
 
-            assert tracer_provider is not None  # make type checker happy
+            if tracer_provider is None:
+                return func(*args, **kwargs)
 
             token = context.attach(extract(os.environ))
             try:
                 tracer = trace_api.get_tracer_provider().get_tracer(__name__)
-
-                card_subcommand = _extract_token_after(sys.argv, "card")
-
-                step_name = _extract_token_after(sys.argv, "step")
-                task_id = _extract_token_after(sys.argv, "--task-id")
-                run_id = _extract_token_after(sys.argv, "--run-id")
-                if step_name and task_id and run_id:
-                    better_name = "/".join([run_id, step_name, task_id])
-                elif card_subcommand:
-                    better_name = "card/" + card_subcommand
-                elif "run" in sys.argv:
-                    better_name = "run"
-                else:
-                    better_name = None
-
                 with tracer.start_as_current_span(
-                    better_name or name, kind=trace_api.SpanKind.SERVER
+                    name, kind=trace_api.SpanKind.SERVER
                 ) as span:
                     span.set_attribute("cmd", " ".join(sys.argv))
                     span.set_attribute("pid", str(os.getpid()))
@@ -113,7 +99,7 @@ def cli_entrypoint(name: str):
 
         return wrapper_func
 
-    return cli_entrypoint_wrap
+    return cli_wrap
 
 
 def inject_tracing_vars(env_dict: Dict[str, str]) -> Dict[str, str]:
@@ -122,23 +108,32 @@ def inject_tracing_vars(env_dict: Dict[str, str]) -> Dict[str, str]:
 
 
 def get_trace_id() -> str:
-    return format_trace_id(trace_api.get_current_span().get_span_context().trace_id)
+    try:
+        return format_trace_id(trace_api.get_current_span().get_span_context().trace_id)
+    except Exception:
+        return ""
 
 
 @contextlib.contextmanager
-def traced(name, attrs={}):
+def traced(name: str, attrs: Optional[Dict] = None):
+    if tracer_provider is None:
+        yield
+        return
     tracer = trace_api.get_tracer_provider().get_tracer(__name__)
     with tracer.start_as_current_span(name) as span:
-        for k, v in attrs.items():
-            span.set_attribute(k, v)
+        if attrs:
+            for k, v in attrs.items():
+                span.set_attribute(k, v)
         yield
 
 
 def tracing(func):
     @wraps(func)
     def wrapper_func(*args, **kwargs):
-        tracer = trace_api.get_tracer_provider().get_tracer(func.__module__.__name__)
+        if tracer_provider is None:
+            return func(*args, **kwargs)
 
+        tracer = trace_api.get_tracer_provider().get_tracer(func.__module__)
         with tracer.start_as_current_span(func.__name__):
             return func(*args, **kwargs)
 

metaflow/version.py

@@ -1 +1 @@
-metaflow_version = "2.14.0"
+metaflow_version = "2.14.2"

{metaflow-2.14.0.dist-info → metaflow-2.14.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: metaflow
-Version: 2.14.0
+Version: 2.14.2
 Summary: Metaflow: More Data Science, Less Engineering
 Author: Metaflow Developers
 Author-email: help@metaflow.org
@@ -26,7 +26,7 @@ License-File: LICENSE
 Requires-Dist: requests
 Requires-Dist: boto3
 Provides-Extra: stubs
-Requires-Dist: metaflow-stubs==2.14.0; extra == "stubs"
+Requires-Dist: metaflow-stubs==2.14.2; extra == "stubs"
 Dynamic: author
 Dynamic: author-email
 Dynamic: classifier