metaflow 2.12.28__py2.py3-none-any.whl → 2.12.29__py2.py3-none-any.whl

metaflow/datastore/task_datastore.py

@@ -10,7 +10,7 @@ from types import MethodType, FunctionType
 
 from .. import metaflow_config
 from ..exception import MetaflowInternalError
-from ..metadata import DataArtifact, MetaDatum
+from ..metadata_provider import DataArtifact, MetaDatum
 from ..parameters import Parameter
 from ..util import Path, is_stringish, to_fileobj
 

metaflow/extension_support/plugins.py

@@ -178,6 +178,7 @@ _plugin_categories = {
     "environment": lambda x: x.TYPE,
     "metadata_provider": lambda x: x.TYPE,
     "datastore": lambda x: x.TYPE,
+    "dataclient": lambda x: x.TYPE,
     "secrets_provider": lambda x: x.TYPE,
     "gcp_client_provider": lambda x: x.name,
     "deployer_impl_provider": lambda x: x.TYPE,

metaflow/flowspec.py

@@ -64,7 +64,7 @@ class ParallelUBF(UnboundedForeachInput):
         return item or 0  # item is None for the control task, but it is also split 0
 
 
-class _FlowSpecMeta(type):
+class FlowSpecMeta(type):
     def __new__(cls, name, bases, dct):
         f = super().__new__(cls, name, bases, dct)
         # This makes sure to give _flow_decorators to each
@@ -75,7 +75,7 @@ class _FlowSpecMeta(type):
         return f
 
 
-class FlowSpec(metaclass=_FlowSpecMeta):
+class FlowSpec(metaclass=FlowSpecMeta):
     """
     Main class from which all Flows should inherit.
 

metaflow/includefile.py

@@ -1,6 +1,7 @@
 from collections import namedtuple
 import gzip
 
+import importlib
 import io
 import json
 import os
@@ -17,6 +18,8 @@ from .parameters import (
     Parameter,
     ParameterContext,
 )
+
+from .plugins import DATACLIENTS
 from .util import get_username
 
 import functools
@@ -47,16 +50,7 @@ _DelayedExecContext = namedtuple(
 
 
 # From here on out, this is the IncludeFile implementation.
-from metaflow.plugins.datatools import Local, S3
-from metaflow.plugins.azure.includefile_support import Azure
-from metaflow.plugins.gcp.includefile_support import GS
-
-DATACLIENTS = {
-    "local": Local,
-    "s3": S3,
-    "azure": Azure,
-    "gs": GS,
-}
+_dict_dataclients = {d.TYPE: d for d in DATACLIENTS}
 
 
 class IncludedFile(object):
@@ -167,7 +161,7 @@ class FilePathClass(click.ParamType):
                 "IncludeFile using a direct reference to a file in cloud storage is no "
                 "longer supported. Contact the Metaflow team if you need this supported"
             )
-            # if DATACLIENTS.get(path[:prefix_pos]) is None:
+            # if _dict_dataclients.get(path[:prefix_pos]) is None:
             #     self.fail(
             #         "IncludeFile: no handler for external file of type '%s' "
             #         "(given path is '%s')" % (path[:prefix_pos], path)
@@ -187,7 +181,7 @@ class FilePathClass(click.ParamType):
                     pass
             except OSError:
                 self.fail("IncludeFile: could not open file '%s' for reading" % path)
-            handler = DATACLIENTS.get(ctx.ds_type)
+            handler = _dict_dataclients.get(ctx.ds_type)
             if handler is None:
                 self.fail(
                     "IncludeFile: no data-client for datastore of type '%s'"
@@ -213,7 +207,7 @@ class FilePathClass(click.ParamType):
                         ctx.path,
                         ctx.is_text,
                         ctx.encoding,
-                        DATACLIENTS[ctx.handler_type],
+                        _dict_dataclients[ctx.handler_type],
                         ctx.echo,
                     )
                 )
@@ -425,7 +419,7 @@ class UploaderV1:
         if prefix_pos < 0:
             raise MetaflowException("Malformed URL: '%s'" % url)
         prefix = url[:prefix_pos]
-        handler = DATACLIENTS.get(prefix)
+        handler = _dict_dataclients.get(prefix)
         if handler is None:
             raise MetaflowException("Could not find data client for '%s'" % prefix)
         return handler

metaflow/metaflow_config.py

@@ -43,6 +43,10 @@ DEFAULT_GCP_CLIENT_PROVIDER = from_conf("DEFAULT_GCP_CLIENT_PROVIDER", "gcp-defa
 DEFAULT_SECRETS_BACKEND_TYPE = from_conf("DEFAULT_SECRETS_BACKEND_TYPE")
 DEFAULT_SECRETS_ROLE = from_conf("DEFAULT_SECRETS_ROLE")
 
+DEFAULT_FROM_DEPLOYMENT_IMPL = from_conf(
+    "DEFAULT_FROM_DEPLOYMENT_IMPL", "argo-workflows"
+)
+
 ###
 # User configuration
 ###

metaflow/metaflow_current.py

@@ -30,7 +30,7 @@ class Current(object):
             raise ex
 
         self.__class__.graph = property(
-            fget=lambda _: _raise(RuntimeError("Graph is not available"))
+            fget=lambda self: _raise(RuntimeError("Graph is not available"))
         )
 
     def _set_env(

metaflow/parameters.py

@@ -438,3 +438,6 @@ def add_custom_parameters(deploy_mode=False):
         return cmd
 
     return wrapper
+
+
+JSONType = JSONTypeClass()

metaflow/plugins/__init__.py

@@ -73,8 +73,8 @@ ENVIRONMENTS_DESC = [
 
 # Add metadata providers here
 METADATA_PROVIDERS_DESC = [
-    ("service", ".metadata.service.ServiceMetadataProvider"),
-    ("local", ".metadata.local.LocalMetadataProvider"),
+    ("service", ".metadata_providers.service.ServiceMetadataProvider"),
+    ("local", ".metadata_providers.local.LocalMetadataProvider"),
 ]
 
 # Add datastore here
@@ -85,13 +85,21 @@ DATASTORES_DESC = [
     ("gs", ".datastores.gs_storage.GSStorage"),
 ]
 
+# Dataclients are used for IncludeFile
+DATACLIENTS_DESC = [
+    ("local", ".datatools.Local"),
+    ("s3", ".datatools.S3"),
+    ("azure", ".azure.includefile_support.Azure"),
+    ("gs", ".gcp.includefile_support.GS"),
+]
+
 # Add non monitoring/logging sidecars here
 SIDECARS_DESC = [
     (
         "save_logs_periodically",
         "..mflog.save_logs_periodically.SaveLogsPeriodicallySidecar",
     ),
-    ("heartbeat", "metaflow.metadata.heartbeat.MetadataHeartBeat"),
+    ("heartbeat", "metaflow.metadata_provider.heartbeat.MetadataHeartBeat"),
 ]
 
 # Add logging sidecars here
@@ -161,6 +169,7 @@ FLOW_DECORATORS = resolve_plugins("flow_decorator")
 ENVIRONMENTS = resolve_plugins("environment")
 METADATA_PROVIDERS = resolve_plugins("metadata_provider")
 DATASTORES = resolve_plugins("datastore")
+DATACLIENTS = resolve_plugins("dataclient")
 SIDECARS = resolve_plugins("sidecar")
 LOGGING_SIDECARS = resolve_plugins("logging_sidecar")
 MONITOR_SIDECARS = resolve_plugins("monitor_sidecar")

metaflow/plugins/airflow/airflow_cli.py

@@ -389,6 +389,11 @@ def _validate_workflow(flow, graph, flow_datastore, metadata, workflow_timeout):
                 "Step *%s* is marked for execution on AWS Batch with Airflow which isn't currently supported."
                 % node.name
             )
+        if any([d.name == "slurm" for d in node.decorators]):
+            raise NotSupportedException(
+                "Step *%s* is marked for execution on Slurm with Airflow which isn't currently supported."
+                % node.name
+            )
     SUPPORTED_DATASTORES = ("azure", "s3", "gs")
     if flow_datastore.TYPE not in SUPPORTED_DATASTORES:
         raise AirflowException(

metaflow/plugins/airflow/airflow_decorator.py

@@ -1,7 +1,7 @@
 import json
 import os
 from metaflow.decorators import StepDecorator
-from metaflow.metadata import MetaDatum
+from metaflow.metadata_provider import MetaDatum
 
 from .airflow_utils import (
     TASK_ID_XCOM_KEY,

metaflow/plugins/argo/argo_workflows_decorator.py

@@ -6,7 +6,7 @@ import time
 from metaflow import current
 from metaflow.decorators import StepDecorator
 from metaflow.events import Trigger
-from metaflow.metadata import MetaDatum
+from metaflow.metadata_provider import MetaDatum
 from metaflow.metaflow_config import ARGO_EVENTS_WEBHOOK_URL
 from metaflow.graph import DAGNode, FlowGraph
 from metaflow.flowspec import FlowSpec

metaflow/plugins/argo/argo_workflows_deployer.py

@@ -1,292 +1,106 @@
-import sys
-import tempfile
-from typing import Optional, ClassVar
+from typing import Any, ClassVar, Dict, Optional, TYPE_CHECKING, Type
 
-from metaflow.plugins.argo.argo_workflows import ArgoWorkflows
-from metaflow.runner.deployer import (
-    DeployerImpl,
-    DeployedFlow,
-    TriggeredRun,
-    get_lower_level_group,
-    handle_timeout,
-)
+from metaflow.runner.deployer_impl import DeployerImpl
 
-
-def suspend(instance: TriggeredRun, **kwargs):
-    """
-    Suspend the running workflow.
-
-    Parameters
-    ----------
-    **kwargs : Any
-        Additional arguments to pass to the suspend command.
-
-    Returns
-    -------
-    bool
-        True if the command was successful, False otherwise.
-    """
-    _, run_id = instance.pathspec.split("/")
-
-    # every subclass needs to have `self.deployer_kwargs`
-    command = get_lower_level_group(
-        instance.deployer.api,
-        instance.deployer.top_level_kwargs,
-        instance.deployer.TYPE,
-        instance.deployer.deployer_kwargs,
-    ).suspend(run_id=run_id, **kwargs)
-
-    pid = instance.deployer.spm.run_command(
-        [sys.executable, *command],
-        env=instance.deployer.env_vars,
-        cwd=instance.deployer.cwd,
-        show_output=instance.deployer.show_output,
-    )
-
-    command_obj = instance.deployer.spm.get(pid)
-    return command_obj.process.returncode == 0
-
-
-def unsuspend(instance: TriggeredRun, **kwargs):
-    """
-    Unsuspend the suspended workflow.
-
-    Parameters
-    ----------
-    **kwargs : Any
-        Additional arguments to pass to the unsuspend command.
-
-    Returns
-    -------
-    bool
-        True if the command was successful, False otherwise.
-    """
-    _, run_id = instance.pathspec.split("/")
-
-    # every subclass needs to have `self.deployer_kwargs`
-    command = get_lower_level_group(
-        instance.deployer.api,
-        instance.deployer.top_level_kwargs,
-        instance.deployer.TYPE,
-        instance.deployer.deployer_kwargs,
-    ).unsuspend(run_id=run_id, **kwargs)
-
-    pid = instance.deployer.spm.run_command(
-        [sys.executable, *command],
-        env=instance.deployer.env_vars,
-        cwd=instance.deployer.cwd,
-        show_output=instance.deployer.show_output,
-    )
-
-    command_obj = instance.deployer.spm.get(pid)
-    return command_obj.process.returncode == 0
-
-
-def terminate(instance: TriggeredRun, **kwargs):
-    """
-    Terminate the running workflow.
-
-    Parameters
-    ----------
-    **kwargs : Any
-        Additional arguments to pass to the terminate command.
-
-    Returns
-    -------
-    bool
-        True if the command was successful, False otherwise.
-    """
-    _, run_id = instance.pathspec.split("/")
-
-    # every subclass needs to have `self.deployer_kwargs`
-    command = get_lower_level_group(
-        instance.deployer.api,
-        instance.deployer.top_level_kwargs,
-        instance.deployer.TYPE,
-        instance.deployer.deployer_kwargs,
-    ).terminate(run_id=run_id, **kwargs)
-
-    pid = instance.deployer.spm.run_command(
-        [sys.executable, *command],
-        env=instance.deployer.env_vars,
-        cwd=instance.deployer.cwd,
-        show_output=instance.deployer.show_output,
-    )
-
-    command_obj = instance.deployer.spm.get(pid)
-    return command_obj.process.returncode == 0
-
-
-def status(instance: TriggeredRun):
-    """
-    Get the status of the triggered run.
-
-    Returns
-    -------
-    str, optional
-        The status of the workflow considering the run object, or None if the status could not be retrieved.
-    """
-    from metaflow.plugins.argo.argo_workflows_cli import (
-        get_status_considering_run_object,
-    )
-
-    flow_name, run_id = instance.pathspec.split("/")
-    name = run_id[5:]
-    status = ArgoWorkflows.get_workflow_status(flow_name, name)
-    if status is not None:
-        return get_status_considering_run_object(status, instance.run)
-    return None
-
-
-def production_token(instance: DeployedFlow):
-    """
-    Get the production token for the deployed flow.
-
-    Returns
-    -------
-    str, optional
-        The production token, None if it cannot be retrieved.
-    """
-    try:
-        _, production_token = ArgoWorkflows.get_existing_deployment(
-            instance.deployer.name
-        )
-        return production_token
-    except TypeError:
-        return None
-
-
-def delete(instance: DeployedFlow, **kwargs):
-    """
-    Delete the deployed flow.
-
-    Parameters
-    ----------
-    **kwargs : Any
-        Additional arguments to pass to the delete command.
-
-    Returns
-    -------
-    bool
-        True if the command was successful, False otherwise.
-    """
-    command = get_lower_level_group(
-        instance.deployer.api,
-        instance.deployer.top_level_kwargs,
-        instance.deployer.TYPE,
-        instance.deployer.deployer_kwargs,
-    ).delete(**kwargs)
-
-    pid = instance.deployer.spm.run_command(
-        [sys.executable, *command],
-        env=instance.deployer.env_vars,
-        cwd=instance.deployer.cwd,
-        show_output=instance.deployer.show_output,
-    )
-
-    command_obj = instance.deployer.spm.get(pid)
-    return command_obj.process.returncode == 0
-
-
-def trigger(instance: DeployedFlow, **kwargs):
-    """
-    Trigger a new run for the deployed flow.
-
-    Parameters
-    ----------
-    **kwargs : Any
-        Additional arguments to pass to the trigger command, `Parameters` in particular
-
-    Returns
-    -------
-    ArgoWorkflowsTriggeredRun
-        The triggered run instance.
-
-    Raises
-    ------
-    Exception
-        If there is an error during the trigger process.
-    """
-    with tempfile.TemporaryDirectory() as temp_dir:
-        tfp_runner_attribute = tempfile.NamedTemporaryFile(dir=temp_dir, delete=False)
-
-        # every subclass needs to have `self.deployer_kwargs`
-        command = get_lower_level_group(
-            instance.deployer.api,
-            instance.deployer.top_level_kwargs,
-            instance.deployer.TYPE,
-            instance.deployer.deployer_kwargs,
-        ).trigger(deployer_attribute_file=tfp_runner_attribute.name, **kwargs)
-
-        pid = instance.deployer.spm.run_command(
-            [sys.executable, *command],
-            env=instance.deployer.env_vars,
-            cwd=instance.deployer.cwd,
-            show_output=instance.deployer.show_output,
-        )
-
-        command_obj = instance.deployer.spm.get(pid)
-        content = handle_timeout(
-            tfp_runner_attribute, command_obj, instance.deployer.file_read_timeout
-        )
-
-        if command_obj.process.returncode == 0:
-            triggered_run = TriggeredRun(deployer=instance.deployer, content=content)
-            triggered_run._enrich_object(
-                {
-                    "status": property(status),
-                    "terminate": terminate,
-                    "suspend": suspend,
-                    "unsuspend": unsuspend,
-                }
-            )
-            return triggered_run
-
-    raise Exception(
-        "Error triggering %s on %s for %s"
-        % (instance.deployer.name, instance.deployer.TYPE, instance.deployer.flow_file)
-    )
+if TYPE_CHECKING:
+    import metaflow.plugins.argo.argo_workflows_deployer_objects
 
 
 class ArgoWorkflowsDeployer(DeployerImpl):
     """
     Deployer implementation for Argo Workflows.
 
-    Attributes
+    Parameters
     ----------
-    TYPE : ClassVar[Optional[str]]
-        The type of the deployer, which is "argo-workflows".
+    name : str, optional, default None
+        Argo workflow name. The flow name is used instead if this option is not specified.
     """
 
     TYPE: ClassVar[Optional[str]] = "argo-workflows"
 
-    def __init__(self, deployer_kwargs, **kwargs):
+    def __init__(self, deployer_kwargs: Dict[str, str], **kwargs):
         """
         Initialize the ArgoWorkflowsDeployer.
 
         Parameters
         ----------
-        deployer_kwargs : dict
+        deployer_kwargs : Dict[str, str]
             The deployer-specific keyword arguments.
         **kwargs : Any
             Additional arguments to pass to the superclass constructor.
         """
-        self.deployer_kwargs = deployer_kwargs
+        self._deployer_kwargs = deployer_kwargs
         super().__init__(**kwargs)
 
-    def _enrich_deployed_flow(self, deployed_flow: DeployedFlow):
+    @property
+    def deployer_kwargs(self) -> Dict[str, Any]:
+        return self._deployer_kwargs
+
+    @staticmethod
+    def deployed_flow_type() -> (
+        Type[
+            "metaflow.plugins.argo.argo_workflows_deployer_objects.ArgoWorkflowsDeployedFlow"
+        ]
+    ):
+        from .argo_workflows_deployer_objects import ArgoWorkflowsDeployedFlow
+
+        return ArgoWorkflowsDeployedFlow
+
+    def create(
+        self, **kwargs
+    ) -> "metaflow.plugins.argo.argo_workflows_deployer_objects.ArgoWorkflowsDeployedFlow":
         """
-        Enrich the DeployedFlow object with additional properties and methods.
+        Create a new ArgoWorkflow deployment.
 
         Parameters
         ----------
-        deployed_flow : DeployedFlow
-            The deployed flow object to enrich.
+        authorize : str, optional, default None
+            Authorize using this production token. Required when re-deploying an existing flow
+            for the first time. The token is cached in METAFLOW_HOME.
+        generate_new_token : bool, optional, default False
+            Generate a new production token for this flow. Moves the production flow to a new namespace.
+        given_token : str, optional, default None
+            Use the given production token for this flow. Moves the production flow to the given namespace.
+        tags : List[str], optional, default None
+            Annotate all objects produced by Argo Workflows runs with these tags.
+        user_namespace : str, optional, default None
+            Change the namespace from the default (production token) to the given tag.
+        only_json : bool, optional, default False
+            Only print out JSON sent to Argo Workflows without deploying anything.
+        max_workers : int, optional, default 100
+            Maximum number of parallel processes.
+        workflow_timeout : int, optional, default None
+            Workflow timeout in seconds.
+        workflow_priority : int, optional, default None
+            Workflow priority as an integer. Higher priority workflows are processed first
+            if Argo Workflows controller is configured to process limited parallel workflows.
+        auto_emit_argo_events : bool, optional, default True
+            Auto emits Argo Events when the run completes successfully.
+        notify_on_error : bool, optional, default False
+            Notify if the workflow fails.
+        notify_on_success : bool, optional, default False
+            Notify if the workflow succeeds.
+        notify_slack_webhook_url : str, optional, default ''
+            Slack incoming webhook url for workflow success/failure notifications.
+        notify_pager_duty_integration_key : str, optional, default ''
+            PagerDuty Events API V2 Integration key for workflow success/failure notifications.
+        enable_heartbeat_daemon : bool, optional, default False
+            Use a daemon container to broadcast heartbeats.
+        deployer_attribute_file : str, optional, default None
+            Write the workflow name to the specified file. Used internally for Metaflow's Deployer API.
+        enable_error_msg_capture : bool, optional, default True
+            Capture stack trace of first failed task in exit hook.
+
+        Returns
+        -------
+        ArgoWorkflowsDeployedFlow
+            The Flow deployed to Argo Workflows.
         """
-        deployed_flow._enrich_object(
-            {
-                "production_token": property(production_token),
-                "trigger": trigger,
-                "delete": delete,
-            }
-        )
+
+        # Prevent circular import
+        from .argo_workflows_deployer_objects import ArgoWorkflowsDeployedFlow
+
+        return self._create(ArgoWorkflowsDeployedFlow, **kwargs)
+
+
+_addl_stubgen_modules = ["metaflow.plugins.argo.argo_workflows_deployer_objects"]