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

metaflow/plugins/aws/step_functions/step_functions_deployer.py

@@ -1,253 +1,94 @@
-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("/")
+from typing import Any, ClassVar, Dict, Optional, TYPE_CHECKING, Type
 
-    # 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)
+from metaflow.runner.deployer_impl import DeployerImpl
 
-    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, instance.deployer.file_read_timeout
-        )
-
-        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)
-    )
+if TYPE_CHECKING:
+    import metaflow.plugins.aws.step_functions.step_functions_deployer_objects
 
 
 class StepFunctionsDeployer(DeployerImpl):
     """
     Deployer implementation for AWS Step Functions.
 
-    Attributes
+    Parameters
     ----------
-    TYPE : ClassVar[Optional[str]]
-        The type of the deployer, which is "step-functions".
+    name : str, optional, default None
+        State Machine name. The flow name is used instead if this option is not specified.
     """
 
     TYPE: ClassVar[Optional[str]] = "step-functions"
 
-    def __init__(self, deployer_kwargs, **kwargs):
+    def __init__(self, deployer_kwargs: Dict[str, str], **kwargs):
         """
         Initialize the StepFunctionsDeployer.
 
         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.aws.step_functions.step_functions_deployer_objects.StepFunctionsDeployedFlow"
+        ]
+    ):
+        from .step_functions_deployer_objects import StepFunctionsDeployedFlow
+
+        return StepFunctionsDeployedFlow
+
+    def create(
+        self, **kwargs
+    ) -> "metaflow.plugins.aws.step_functions.step_functions_deployer_objects.StepFunctionsDeployedFlow":
         """
-        Enrich the DeployedFlow object with additional properties and methods.
+        Create a new AWS Step Functions State Machine 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 AWS Step Functions 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 AWS Step Functions without deploying anything.
+        max_workers : int, optional, default 100
+            Maximum number of parallel processes.
+        workflow_timeout : int, optional, default None
+            Workflow timeout in seconds.
+        log_execution_history : bool, optional, default False
+            Log AWS Step Functions execution history to AWS CloudWatch Logs log group.
+        use_distributed_map : bool, optional, default False
+            Use AWS Step Functions Distributed Map instead of Inline Map for defining foreach
+            tasks in Amazon State Language.
+        deployer_attribute_file : str, optional, default None
+            Write the workflow name to the specified file. Used internally for Metaflow's Deployer API.
+
+        Returns
+        -------
+        StepFunctionsDeployedFlow
+            The Flow deployed to AWS Step Functions.
         """
-        deployed_flow._enrich_object(
-            {
-                "production_token": property(production_token),
-                "trigger": trigger,
-                "delete": delete,
-                "list_runs": list_runs,
-            }
-        )
+        from .step_functions_deployer_objects import StepFunctionsDeployedFlow
+
+        return self._create(StepFunctionsDeployedFlow, **kwargs)
+
+
+_addl_stubgen_modules = [
+    "metaflow.plugins.aws.step_functions.step_functions_deployer_objects"
+]

metaflow/plugins/aws/step_functions/step_functions_deployer_objects.py

@@ -0,0 +1,236 @@
+import sys
+import json
+import tempfile
+from typing import ClassVar, Optional, List
+
+from metaflow.plugins.aws.step_functions.step_functions import StepFunctions
+from metaflow.runner.deployer import DeployedFlow, TriggeredRun
+
+from metaflow.runner.utils import get_lower_level_group, handle_timeout
+
+
+class StepFunctionsTriggeredRun(TriggeredRun):
+    """
+    A class representing a triggered AWS Step Functions state machine execution.
+    """
+
+    def terminate(self, **kwargs) -> bool:
+        """
+        Terminate the running state machine execution.
+
+        Parameters
+        ----------
+        authorize : str, optional, default None
+            Authorize the termination with a production token.
+
+        Returns
+        -------
+        bool
+            True if the command was successful, False otherwise.
+        """
+        _, run_id = self.pathspec.split("/")
+
+        # every subclass needs to have `self.deployer_kwargs`
+        command = get_lower_level_group(
+            self.deployer.api,
+            self.deployer.top_level_kwargs,
+            self.deployer.TYPE,
+            self.deployer.deployer_kwargs,
+        ).terminate(run_id=run_id, **kwargs)
+
+        pid = self.deployer.spm.run_command(
+            [sys.executable, *command],
+            env=self.deployer.env_vars,
+            cwd=self.deployer.cwd,
+            show_output=self.deployer.show_output,
+        )
+
+        command_obj = self.deployer.spm.get(pid)
+        return command_obj.process.returncode == 0
+
+
+class StepFunctionsDeployedFlow(DeployedFlow):
+    """
+    A class representing a deployed AWS Step Functions state machine.
+    """
+
+    TYPE: ClassVar[Optional[str]] = "step-functions"
+
+    @classmethod
+    def from_deployment(cls, identifier: str, metadata: Optional[str] = None):
+        """
+        This method is not currently implemented for Step Functions.
+
+        Raises
+        ------
+        NotImplementedError
+            This method is not implemented for Step Functions.
+        """
+        raise NotImplementedError(
+            "from_deployment is not implemented for StepFunctions"
+        )
+
+    @property
+    def production_token(self: DeployedFlow) -> Optional[str]:
+        """
+        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(
+                self.deployer.name
+            )
+            return production_token
+        except TypeError:
+            return None
+
+    def list_runs(
+        self, states: Optional[List[str]] = None
+    ) -> List[StepFunctionsTriggeredRun]:
+        """
+        List runs of the deployed flow.
+
+        Parameters
+        ----------
+        states : List[str], optional, default None
+            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[StepFunctionsTriggeredRun]
+            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(self.deployer.name, states)
+
+        for e in executions:
+            run_id = "sfn-%s" % e["name"]
+            tr = StepFunctionsTriggeredRun(
+                deployer=self.deployer,
+                content=json.dumps(
+                    {
+                        "metadata": self.deployer.metadata,
+                        "pathspec": "/".join((self.deployer.flow_name, run_id)),
+                        "name": run_id,
+                    }
+                ),
+            )
+            triggered_runs.append(tr)
+
+        return triggered_runs
+
+    def delete(self, **kwargs) -> bool:
+        """
+        Delete the deployed state machine.
+
+        Parameters
+        ----------
+        authorize : str, optional, default None
+            Authorize the deletion with a production token.
+
+        Returns
+        -------
+        bool
+            True if the command was successful, False otherwise.
+        """
+        command = get_lower_level_group(
+            self.deployer.api,
+            self.deployer.top_level_kwargs,
+            self.deployer.TYPE,
+            self.deployer.deployer_kwargs,
+        ).delete(**kwargs)
+
+        pid = self.deployer.spm.run_command(
+            [sys.executable, *command],
+            env=self.deployer.env_vars,
+            cwd=self.deployer.cwd,
+            show_output=self.deployer.show_output,
+        )
+
+        command_obj = self.deployer.spm.get(pid)
+        return command_obj.process.returncode == 0
+
+    def trigger(self, **kwargs) -> StepFunctionsTriggeredRun:
+        """
+        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(
+                self.deployer.api,
+                self.deployer.top_level_kwargs,
+                self.deployer.TYPE,
+                self.deployer.deployer_kwargs,
+            ).trigger(deployer_attribute_file=tfp_runner_attribute.name, **kwargs)
+
+            pid = self.deployer.spm.run_command(
+                [sys.executable, *command],
+                env=self.deployer.env_vars,
+                cwd=self.deployer.cwd,
+                show_output=self.deployer.show_output,
+            )
+
+            command_obj = self.deployer.spm.get(pid)
+            content = handle_timeout(
+                tfp_runner_attribute, command_obj, self.deployer.file_read_timeout
+            )
+
+            if command_obj.process.returncode == 0:
+                return StepFunctionsTriggeredRun(
+                    deployer=self.deployer, content=content
+                )
+
+        raise Exception(
+            "Error triggering %s on %s for %s"
+            % (
+                self.deployer.name,
+                self.deployer.TYPE,
+                self.deployer.flow_file,
+            )
+        )

metaflow/plugins/azure/includefile_support.py

@@ -8,6 +8,8 @@ from metaflow.exception import MetaflowException, MetaflowInternalError
 
 
 class Azure(object):
+    TYPE = "azure"
+
     @classmethod
     def get_root_from_config(cls, echo, create_on_absent=True):
         from metaflow.metaflow_config import DATATOOLS_AZUREROOT

metaflow/plugins/cards/card_cli.py

@@ -1,5 +1,6 @@
 from metaflow.client import Task
-from metaflow import JSONType, namespace
+from metaflow.parameters import JSONTypeClass
+from metaflow import namespace
 from metaflow.util import resolve_identity
 from metaflow.exception import (
     CommandException,
@@ -551,7 +552,7 @@ def update_card(mf_card, mode, task, data, timeout_value=None):
     "--options",
     default=None,
     show_default=True,
-    type=JSONType,
+    type=JSONTypeClass(),
     help="arguments of the card being created.",
 )
 @click.option(

metaflow/plugins/cards/card_modules/components.py

@@ -712,15 +712,15 @@ class ProgressBar(UserComponent):
 
     Parameters
     ----------
-    max : int
+    max : int, default 100
         The maximum value of the progress bar.
-    label : str, optional
+    label : str, optional, default None
         Optional label for the progress bar.
-    value : int, optional
+    value : int, default 0
         Optional initial value of the progress bar.
-    unit : str, optional
+    unit : str, optional, default None
         Optional unit for the progress bar.
-    metadata : str, optional
+    metadata : str, optional, default None
         Optional additional information to show on the progress bar.
     """
 
@@ -731,10 +731,10 @@ class ProgressBar(UserComponent):
     def __init__(
         self,
         max: int = 100,
-        label: str = None,
+        label: Optional[str] = None,
         value: int = 0,
-        unit: str = None,
-        metadata: str = None,
+        unit: Optional[str] = None,
+        metadata: Optional[str] = None,
     ):
         self._label = label
         self._max = max
@@ -742,7 +742,7 @@ class ProgressBar(UserComponent):
         self._unit = unit
         self._metadata = metadata
 
-    def update(self, new_value: int, metadata: str = None):
+    def update(self, new_value: int, metadata: Optional[str] = None):
         self._value = new_value
         if metadata is not None:
             self._metadata = metadata

metaflow/plugins/cards/card_server.py

@@ -20,14 +20,9 @@ except ImportError:
 from .card_client import CardContainer
 from .exception import CardNotPresentException
 from .card_resolver import resolve_paths_from_task
-from metaflow.metaflow_config import DATASTORE_LOCAL_DIR
 from metaflow import namespace
-from metaflow.exception import (
-    CommandException,
-    MetaflowNotFound,
-    MetaflowNamespaceMismatch,
-)
-
+from metaflow.exception import MetaflowNotFound
+from metaflow.plugins.datastores.local_storage import LocalStorage
 
 VIEWER_PATH = os.path.join(
     os.path.dirname(os.path.abspath(__file__)), "card_viewer", "viewer.html"
@@ -50,18 +45,48 @@ class RunWatcher(Thread):
     def __init__(self, flow_name, connection: Connection):
         super().__init__()
 
-        self._watch_file = os.path.join(
-            os.getcwd(), DATASTORE_LOCAL_DIR, flow_name, "latest_run"
-        )
-        self._current_run_id = self.get_run_id()
         self.daemon = True
         self._connection = connection
+        self._flow_name = flow_name
+
+        self._watch_file = self._initialize_watch_file()
+        if self._watch_file is None:
+            _ClickLogger(
+                "Warning: Could not initialize watch file location.", fg="yellow"
+            )
+
+        self._current_run_id = self.get_run_id()
+
+    def _initialize_watch_file(self):
+        local_root = LocalStorage.datastore_root
+        if local_root is None:
+            local_root = LocalStorage.get_datastore_root_from_config(
+                lambda _: None, create_on_absent=False
+            )
+
+        return (
+            os.path.join(local_root, self._flow_name, "latest_run")
+            if local_root
+            else None
+        )
 
     def get_run_id(self):
-        if not os.path.exists(self._watch_file):
+        # Try to reinitialize watch file if needed
+        if not self._watch_file:
+            self._watch_file = self._initialize_watch_file()
+
+        # Early return if watch file is still None or doesn't exist
+        if not (self._watch_file and os.path.exists(self._watch_file)):
+            return None
+
+        try:
+            with open(self._watch_file, "r") as f:
+                return f.read().strip()
+        except (IOError, OSError) as e:
+            _ClickLogger(
+                "Warning: Could not read run ID from watch file: %s" % e, fg="yellow"
+            )
             return None
-        with open(self._watch_file, "r") as f:
-            return f.read().strip()
 
     def watch(self):
         while True: