metaflow 2.12.8__py2.py3-none-any.whl → 2.12.10__py2.py3-none-any.whl

metaflow/plugins/aws/step_functions/step_functions_cli.py

@@ -4,6 +4,7 @@ import re
 from hashlib import sha1
 
 from metaflow import JSONType, current, decorators, parameters
+from metaflow.client.core import get_metadata
 from metaflow._vendor import click
 from metaflow.exception import MetaflowException, MetaflowInternalError
 from metaflow.metaflow_config import (
@@ -130,6 +131,14 @@ def step_functions(obj, name=None):
     help="Use AWS Step Functions Distributed Map instead of Inline Map for "
     "defining foreach tasks in Amazon State Language.",
 )
+@click.option(
+    "--deployer-attribute-file",
+    default=None,
+    show_default=True,
+    type=str,
+    help="Write the workflow name to the file specified. Used internally for Metaflow's Deployer API.",
+    hidden=True,
+)
 @click.pass_obj
 def create(
     obj,
@@ -143,9 +152,21 @@ def create(
     workflow_timeout=None,
     log_execution_history=False,
     use_distributed_map=False,
+    deployer_attribute_file=None,
 ):
     validate_tags(tags)
 
+    if deployer_attribute_file:
+        with open(deployer_attribute_file, "w") as f:
+            json.dump(
+                {
+                    "name": obj.state_machine_name,
+                    "flow_name": obj.flow.name,
+                    "metadata": get_metadata(),
+                },
+                f,
+            )
+
     obj.echo(
         "Deploying *%s* to AWS Step Functions..." % obj.state_machine_name, bold=True
     )
@@ -231,8 +252,10 @@ def check_metadata_service_version(obj):
 
 
 def resolve_state_machine_name(obj, name):
-    def attach_prefix(name):
-        if SFN_STATE_MACHINE_PREFIX is not None:
+    def attach_prefix(name: str):
+        if SFN_STATE_MACHINE_PREFIX is not None and (
+            not name.startswith(SFN_STATE_MACHINE_PREFIX)
+        ):
             return SFN_STATE_MACHINE_PREFIX + "_" + name
         return name
 
@@ -440,8 +463,16 @@ def resolve_token(
     type=str,
     help="Write the ID of this run to the file specified.",
 )
+@click.option(
+    "--deployer-attribute-file",
+    default=None,
+    show_default=True,
+    type=str,
+    help="Write the metadata and pathspec of this run to the file specified.\nUsed internally for Metaflow's Deployer API.",
+    hidden=True,
+)
 @click.pass_obj
-def trigger(obj, run_id_file=None, **kwargs):
+def trigger(obj, run_id_file=None, deployer_attribute_file=None, **kwargs):
     def _convert_value(param):
         # Swap `-` with `_` in parameter name to match click's behavior
         val = kwargs.get(param.name.replace("-", "_").lower())
@@ -466,6 +497,17 @@ def trigger(obj, run_id_file=None, **kwargs):
         with open(run_id_file, "w") as f:
             f.write(str(run_id))
 
+    if deployer_attribute_file:
+        with open(deployer_attribute_file, "w") as f:
+            json.dump(
+                {
+                    "name": obj.state_machine_name,
+                    "metadata": get_metadata(),
+                    "pathspec": "/".join((obj.flow.name, run_id)),
+                },
+                f,
+            )
+
     obj.echo(
         "Workflow *{name}* triggered on AWS Step Functions "
         "(run-id *{run_id}*).".format(name=obj.state_machine_name, run_id=run_id),

metaflow/plugins/aws/step_functions/step_functions_deployer.py

@@ -0,0 +1,251 @@
+import sys
+import json
+import tempfile
+from typing import Optional, ClassVar, List
+
+from metaflow.plugins.aws.step_functions.step_functions import StepFunctions
+from metaflow.runner.deployer import (
+    DeployerImpl,
+    DeployedFlow,
+    TriggeredRun,
+    get_lower_level_group,
+    handle_timeout,
+)
+
+
+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 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 = StepFunctions.get_existing_deployment(
+            instance.deployer.name
+        )
+        return production_token
+    except TypeError:
+        return None
+
+
+def list_runs(instance: DeployedFlow, states: Optional[List[str]] = None):
+    """
+    List runs of the deployed flow.
+
+    Parameters
+    ----------
+    states : Optional[List[str]], optional
+        A list of states to filter the runs by. Allowed values are:
+        RUNNING, SUCCEEDED, FAILED, TIMED_OUT, ABORTED.
+        If not provided, all states will be considered.
+
+    Returns
+    -------
+    List[TriggeredRun]
+        A list of TriggeredRun objects representing the runs of the deployed flow.
+
+    Raises
+    ------
+    ValueError
+        If any of the provided states are invalid or if there are duplicate states.
+    """
+    VALID_STATES = {"RUNNING", "SUCCEEDED", "FAILED", "TIMED_OUT", "ABORTED"}
+
+    if states is None:
+        states = []
+
+    unique_states = set(states)
+    if not unique_states.issubset(VALID_STATES):
+        invalid_states = unique_states - VALID_STATES
+        raise ValueError(
+            f"Invalid states found: {invalid_states}. Valid states are: {VALID_STATES}"
+        )
+
+    if len(states) != len(unique_states):
+        raise ValueError("Duplicate states are not allowed")
+
+    triggered_runs = []
+    executions = StepFunctions.list(instance.deployer.name, states)
+
+    for e in executions:
+        run_id = "sfn-%s" % e["name"]
+        tr = TriggeredRun(
+            deployer=instance.deployer,
+            content=json.dumps(
+                {
+                    "metadata": instance.deployer.metadata,
+                    "pathspec": "/".join((instance.deployer.flow_name, run_id)),
+                    "name": run_id,
+                }
+            ),
+        )
+        tr._enrich_object({"terminate": terminate})
+        triggered_runs.append(tr)
+
+    return triggered_runs
+
+
+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
+    -------
+    StepFunctionsTriggeredRun
+        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)
+
+        if command_obj.process.returncode == 0:
+            triggered_run = TriggeredRun(deployer=instance.deployer, content=content)
+            triggered_run._enrich_object({"terminate": terminate})
+            return triggered_run
+
+    raise Exception(
+        "Error triggering %s on %s for %s"
+        % (instance.deployer.name, instance.deployer.TYPE, instance.deployer.flow_file)
+    )
+
+
+class StepFunctionsDeployer(DeployerImpl):
+    """
+    Deployer implementation for AWS Step Functions.
+
+    Attributes
+    ----------
+    TYPE : ClassVar[Optional[str]]
+        The type of the deployer, which is "step-functions".
+    """
+
+    TYPE: ClassVar[Optional[str]] = "step-functions"
+
+    def __init__(self, deployer_kwargs, **kwargs):
+        """
+        Initialize the StepFunctionsDeployer.
+
+        Parameters
+        ----------
+        deployer_kwargs : dict
+            The deployer-specific keyword arguments.
+        **kwargs : Any
+            Additional arguments to pass to the superclass constructor.
+        """
+        self.deployer_kwargs = deployer_kwargs
+        super().__init__(**kwargs)
+
+    def _enrich_deployed_flow(self, deployed_flow: DeployedFlow):
+        """
+        Enrich the DeployedFlow object with additional properties and methods.
+
+        Parameters
+        ----------
+        deployed_flow : DeployedFlow
+            The deployed flow object to enrich.
+        """
+        deployed_flow._enrich_object(
+            {
+                "production_token": property(production_token),
+                "trigger": trigger,
+                "delete": delete,
+                "list_runs": list_runs,
+            }
+        )

metaflow/plugins/cards/card_cli.py

@@ -752,7 +752,7 @@ def create(
         return _card.render(
             task,
             stack_trace=stack_trace,
-        ).replace(mf_card.RELOAD_POLICY_TOKEN, token)
+        ).replace(_card.RELOAD_POLICY_TOKEN, token)
 
     if error_stack_trace is not None and mode != "refresh":
         rendered_content = _render_error_card(error_stack_trace)