ob-metaflow 2.14.0.2rc0__py2.py3-none-any.whl → 2.14.2.1__py2.py3-none-any.whl

metaflow/plugins/env_escape/server.py

@@ -264,6 +264,7 @@ class Server(object):
     def serve(self, path=None, port=None):
         # Open up a connection
         if path is not None:
+            # Keep the print line to facilitate debugging
             # print("SERVER: Starting at %s" % path)
             sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
             __try_op__("bind", sock.bind, BIND_RETRY, path)
@@ -503,6 +504,12 @@ class Server(object):
         class_type = self._known_classes.get(class_name)
         if class_type is None:
             raise ValueError("Unknown class %s" % class_name)
+        # Check if __init__ is overridden
+        override_mapping = self._overrides.get(class_type)
+        if override_mapping:
+            override_func = override_mapping.get("__init__")
+            if override_func:
+                return override_func(None, class_type, *args, **kwargs)
         return class_type(*args, **kwargs)
 
     def _handle_subclasscheck(self, target, class_name, otherclass_name, reverse=False):

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.py

@@ -194,7 +194,6 @@ class Kubernetes(object):
         port=None,
         num_parallel=None,
         qos=None,
-        extended_resources=None,
     ):
         name = "js-%s" % str(uuid4())[:6]
         jobset = (
@@ -228,7 +227,6 @@ class Kubernetes(object):
                 port=port,
                 num_parallel=num_parallel,
                 qos=qos,
-                extended_resources=extended_resources,
             )
             .environment_variable("METAFLOW_CODE_SHA", code_package_sha)
             .environment_variable("METAFLOW_CODE_URL", code_package_url)
@@ -506,7 +504,6 @@ class Kubernetes(object):
         name_pattern=None,
         qos=None,
         annotations=None,
-        extended_resources=None,
     ):
         if env is None:
             env = {}
@@ -549,7 +546,6 @@ class Kubernetes(object):
                 shared_memory=shared_memory,
                 port=port,
                 qos=qos,
-                extended_resources=extended_resources,
             )
             .environment_variable("METAFLOW_CODE_SHA", code_package_sha)
             .environment_variable("METAFLOW_CODE_URL", code_package_url)

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")
@@ -145,12 +145,6 @@ def kubernetes():
     type=JSONTypeClass(),
     multiple=False,
 )
-@click.option(
-    "--extended-resources",
-    default=None,
-    type=JSONTypeClass(),
-    multiple=False,
-)
 @click.pass_context
 def step(
     ctx,
@@ -182,7 +176,6 @@ def step(
     qos=None,
     labels=None,
     annotations=None,
-    extended_resources=None,
     **kwargs
 ):
     def echo(msg, stream="stderr", job_id=None, **kwargs):
@@ -326,7 +319,6 @@ def step(
                 qos=qos,
                 labels=labels,
                 annotations=annotations,
-                extended_resources=extended_resources,
             )
     except Exception:
         traceback.print_exc(chain=False)

metaflow/plugins/kubernetes/kubernetes_decorator.py

@@ -123,9 +123,6 @@ class KubernetesDecorator(StepDecorator):
         Only applicable when @parallel is used.
     qos: str, default: Burstable
         Quality of Service class to assign to the pod. Supported values are: Guaranteed, Burstable, BestEffort
-    extended_resources: Dict[str, str], optional, default None
-        Extended resources to be requested for the pod.
-        https://kubernetes.io/docs/tasks/administer-cluster/extended-resource-node/
     """
 
     name = "kubernetes"
@@ -156,7 +153,6 @@ class KubernetesDecorator(StepDecorator):
         "executable": None,
         "hostname_resolution_timeout": 10 * 60,
         "qos": KUBERNETES_QOS,
-        "extended_resources": {},
     }
     package_url = None
     package_sha = None
@@ -492,7 +488,6 @@ class KubernetesDecorator(StepDecorator):
                     "persistent_volume_claims",
                     "labels",
                     "annotations",
-                    "extended_resources",
                 ]:
                     cli_args.command_options[k] = json.dumps(v)
                 else:

metaflow/plugins/kubernetes/kubernetes_job.py

@@ -89,11 +89,10 @@ class KubernetesJob(object):
 
         additional_obp_configs = {
             "OBP_PERIMETER": initial_configs["OBP_PERIMETER"],
-            "OBP_INTEGRATIONS_URL": initial_configs["OBP_INTEGRATIONS_URL"],
+            "OBP_INTEGRATIONS_URL": initial_configs[
+                "OBP_INTEGRATIONS_URL"
+            ],
         }
-        extended_resources = self._kwargs.get("extended_resources", {})
-        qos_requests = {**qos_requests, **extended_resources}
-        qos_limits = {**qos_limits, **extended_resources}
 
         return client.V1JobSpec(
             # Retries are handled by Metaflow when it is responsible for
@@ -119,7 +118,6 @@ class KubernetesJob(object):
                     containers=[
                         client.V1Container(
                             command=self._kwargs["command"],
-                            security_context=client.V1SecurityContext(privileged=True),
                             termination_message_policy="FallbackToLogsOnError",
                             ports=(
                                 []

metaflow/plugins/kubernetes/kubernetes_jobsets.py

@@ -560,11 +560,6 @@ class JobSetSpec(object):
             self._kwargs["memory"],
             self._kwargs["disk"],
         )
-
-        extended_resources = self._kwargs.get("extended_resources", {})
-        qos_requests = {**qos_requests, **extended_resources}
-        qos_limits = {**qos_limits, **extended_resources}
-
         return dict(
             name=self.name,
             template=client.api_client.ApiClient().sanitize_for_serialization(
@@ -616,9 +611,6 @@ class JobSetSpec(object):
                                                 )
                                             ]
                                         ),
-                                        security_context=client.V1SecurityContext(
-                                            privileged=True
-                                        ),
                                         env=[
                                             client.V1EnvVar(name=k, value=str(v))
                                             for k, v in self._kwargs.get(

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,6 +3,8 @@ from metaflow.metaflow_config import (
     OTEL_ENDPOINT,
     ZIPKIN_ENDPOINT,
     CONSOLE_TRACE_ENABLED,
+    SERVICE_AUTH_KEY,
+    SERVICE_HEADERS,
 )
 
 if OTEL_ENDPOINT:
@@ -11,56 +13,47 @@ if OTEL_ENDPOINT:
 
 
 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": 15,
+    }
 
     if SERVICE_AUTH_KEY:
-        span_exporter = OTLPSpanExporter(
-            endpoint=OTEL_ENDPOINT,
-            headers={"x-api-key": SERVICE_AUTH_KEY},
-            timeout=15,
-        )
+        config["headers"] = {"x-api-key": SERVICE_AUTH_KEY}
     elif SERVICE_HEADERS:
-        span_exporter = OTLPSpanExporter(
-            endpoint=OTEL_ENDPOINT,
-            headers=SERVICE_HEADERS,
-            timeout=15,
-        )
-    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()