metaflow-stubs 2.12.8__py2.py3-none-any.whl → 2.12.9__py2.py3-none-any.whl

metaflow-stubs/plugins/argo/argo_workflows.pyi

@@ -1,17 +1,18 @@
 ##################################################################################
 #                       Auto-generated Metaflow stub file                        #
-# MF version: 2.12.8                                                             #
-# Generated on 2024-07-16T15:51:55.417818                                        #
+# MF version: 2.12.9                                                             #
+# Generated on 2024-07-25T18:20:16.325569                                        #
 ##################################################################################
 
 from __future__ import annotations
 
 import typing
 if typing.TYPE_CHECKING:
-    import metaflow._vendor.click.types
     import metaflow.parameters
-    import metaflow.exception
     import metaflow.metaflow_current
+    import metaflow._vendor.click.types
+    import metaflow.graph
+    import metaflow.exception
 
 inf: float
 
@@ -103,6 +104,10 @@ SERVICE_INTERNAL_URL: None
 
 UI_URL: None
 
+UBF_CONTROL: str
+
+UBF_TASK: str
+
 BASH_SAVE_LOGS: str
 
 def deploy_time_eval(value):
@@ -123,6 +128,32 @@ def validate_kube_labels(labels: typing.Optional[typing.Dict[str, typing.Optiona
     """
     ...
 
+class KubernetesArgoJobSet(object, metaclass=type):
+    def __init__(self, kubernetes_sdk, name = None, namespace = None, **kwargs):
+        ...
+    @property
+    def jobset_control_addr(self):
+        ...
+    @property
+    def worker(self):
+        ...
+    @property
+    def control(self):
+        ...
+    def environment_variable_from_selector(self, name, label_value):
+        ...
+    def environment_variables_from_selectors(self, env_dict):
+        ...
+    def environment_variable(self, name, value):
+        ...
+    def label(self, name, value):
+        ...
+    def annotation(self, name, value):
+        ...
+    def dump(self):
+        ...
+    ...
+
 class ArgoClient(object, metaclass=type):
     def __init__(self, namespace = None):
         ...
@@ -176,7 +207,7 @@ class ArgoWorkflowsSchedulingException(metaflow.exception.MetaflowException, met
     ...
 
 class ArgoWorkflows(object, metaclass=type):
-    def __init__(self, name, graph, flow, code_package_sha, code_package_url, production_token, metadata, flow_datastore, environment, event_logger, monitor, tags = None, namespace = None, username = None, max_workers = None, workflow_timeout = None, workflow_priority = None, auto_emit_argo_events = False, notify_on_error = False, notify_on_success = False, notify_slack_webhook_url = None, notify_pager_duty_integration_key = None):
+    def __init__(self, name, graph: metaflow.graph.FlowGraph, flow, code_package_sha, code_package_url, production_token, metadata, flow_datastore, environment, event_logger, monitor, tags = None, namespace = None, username = None, max_workers = None, workflow_timeout = None, workflow_priority = None, auto_emit_argo_events = False, notify_on_error = False, notify_on_success = False, notify_slack_webhook_url = None, notify_pager_duty_integration_key = None):
         ...
     def __str__(self):
         ...
@@ -359,6 +390,8 @@ class Template(object, metaclass=type):
         ...
     def to_json(self):
         ...
+    def resource(self, action, manifest, success_criteria, failure_criteria):
+        ...
     def __str__(self):
         ...
     ...

metaflow-stubs/plugins/argo/argo_workflows_cli.pyi

@@ -1,22 +1,281 @@
 ##################################################################################
 #                       Auto-generated Metaflow stub file                        #
-# MF version: 2.12.8                                                             #
-# Generated on 2024-07-16T15:51:55.420043                                        #
+# MF version: 2.12.9                                                             #
+# Generated on 2024-07-25T18:20:16.329989                                        #
 ##################################################################################
 
 from __future__ import annotations
 
 import typing
 if typing.TYPE_CHECKING:
-    import metaflow.decorators
+    import datetime
+    import metaflow.client.core
     import metaflow.parameters
-    import metaflow.exception
     import metaflow.metaflow_current
+    import metaflow.decorators
+    import metaflow.graph
+    import metaflow.events
+    import metaflow.exception
+
+class Run(metaflow.client.core.MetaflowObject, metaclass=type):
+    def steps(self, *tags: str) -> typing.Iterator[metaflow.client.core.Step]:
+        """
+        [Legacy function - do not use]
+        
+        Returns an iterator over all `Step` objects in the step. This is an alias
+        to iterating the object itself, i.e.
+        ```
+        list(Run(...)) == list(Run(...).steps())
+        ```
+        
+        Parameters
+        ----------
+        tags : str
+            No op (legacy functionality)
+        
+        Yields
+        ------
+        Step
+            `Step` objects in this run.
+        """
+        ...
+    @property
+    def code(self) -> typing.Optional[metaflow.client.core.MetaflowCode]:
+        """
+        Returns the MetaflowCode object for this run, if present.
+        Code is packed if atleast one `Step` runs remotely, else None is returned.
+        
+        Returns
+        -------
+        MetaflowCode, optional
+            Code package for this run
+        """
+        ...
+    @property
+    def data(self) -> typing.Optional[metaflow.client.core.MetaflowData]:
+        """
+        Returns a container of data artifacts produced by this run.
+        
+        You can access data produced by this run as follows:
+        ```
+        print(run.data.my_var)
+        ```
+        This is a shorthand for `run['end'].task.data`. If the 'end' step has not yet
+        executed, returns None.
+        
+        Returns
+        -------
+        MetaflowData, optional
+            Container of all artifacts produced by this task
+        """
+        ...
+    @property
+    def successful(self) -> bool:
+        """
+        Indicates whether or not the run completed successfully.
+        
+        A run is successful if its 'end' step is successful.
+        
+        Returns
+        -------
+        bool
+            True if the run completed successfully and False otherwise
+        """
+        ...
+    @property
+    def finished(self) -> bool:
+        """
+        Indicates whether or not the run completed.
+        
+        A run completed if its 'end' step completed.
+        
+        Returns
+        -------
+        bool
+            True if the run completed and False otherwise
+        """
+        ...
+    @property
+    def finished_at(self) -> typing.Optional[datetime.datetime]:
+        """
+        Returns the datetime object of when the run finished (successfully or not).
+        
+        The completion time of a run is the same as the completion time of its 'end' step.
+        If the 'end' step has not completed, returns None.
+        
+        Returns
+        -------
+        datetime, optional
+            Datetime of when the run finished
+        """
+        ...
+    @property
+    def end_task(self) -> typing.Optional[metaflow.client.core.Task]:
+        """
+        Returns the Task corresponding to the 'end' step.
+        
+        This returns None if the end step does not yet exist.
+        
+        Returns
+        -------
+        Task, optional
+            The 'end' task
+        """
+        ...
+    def add_tag(self, tag: str):
+        """
+        Add a tag to this `Run`.
+        
+        Note that if the tag is already a system tag, it is not added as a user tag,
+        and no error is thrown.
+        
+        Parameters
+        ----------
+        tag : str
+            Tag to add.
+        """
+        ...
+    def add_tags(self, tags: typing.Iterable[str]):
+        """
+        Add one or more tags to this `Run`.
+        
+        Note that if any tag is already a system tag, it is not added as a user tag
+        and no error is thrown.
+        
+        Parameters
+        ----------
+        tags : Iterable[str]
+            Tags to add.
+        """
+        ...
+    def remove_tag(self, tag: str):
+        """
+        Remove one tag from this `Run`.
+        
+        Removing a system tag is an error. Removing a non-existent
+        user tag is a no-op.
+        
+        Parameters
+        ----------
+        tag : str
+            Tag to remove.
+        """
+        ...
+    def remove_tags(self, tags: typing.Iterable[str]):
+        """
+        Remove one or more tags to this `Run`.
+        
+        Removing a system tag will result in an error. Removing a non-existent
+        user tag is a no-op.
+        
+        Parameters
+        ----------
+        tags : Iterable[str]
+            Tags to remove.
+        """
+        ...
+    def replace_tag(self, tag_to_remove: str, tag_to_add: str):
+        """
+        Remove a tag and add a tag atomically. Removal is done first.
+        The rules for `Run.add_tag` and `Run.remove_tag` also apply here.
+        
+        Parameters
+        ----------
+        tag_to_remove : str
+            Tag to remove.
+        tag_to_add : str
+            Tag to add.
+        """
+        ...
+    def replace_tags(self, tags_to_remove: typing.Iterable[str], tags_to_add: typing.Iterable[str]):
+        """
+        Remove and add tags atomically; the removal is done first.
+        The rules for `Run.add_tag` and `Run.remove_tag` also apply here.
+        
+        Parameters
+        ----------
+        tags_to_remove : Iterable[str]
+            Tags to remove.
+        tags_to_add : Iterable[str]
+            Tags to add.
+        """
+        ...
+    def __iter__(self) -> typing.Iterator[metaflow.client.core.Step]:
+        """
+        Iterate over all children Step of this Run
+        
+        Yields
+        ------
+        Step
+            A Step in this Run
+        """
+        ...
+    def __getitem__(self, name: str) -> metaflow.client.core.Step:
+        """
+        Returns the Step object with the step name 'name'
+        
+        Parameters
+        ----------
+        name : str
+            Step name
+        
+        Returns
+        -------
+        Step
+            Step for this step name in this Run
+        
+        Raises
+        ------
+        KeyError
+            If the name does not identify a valid Step object
+        """
+        ...
+    def __getstate__(self):
+        ...
+    def __setstate__(self, state):
+        ...
+    @property
+    def trigger(self) -> typing.Optional[metaflow.events.Trigger]:
+        """
+        Returns a container of events that triggered this run.
+        
+        This returns None if the run was not triggered by any events.
+        
+        Returns
+        -------
+        Trigger, optional
+            Container of triggering events
+        """
+        ...
+    ...
 
 JSONType: metaflow.parameters.JSONTypeClass
 
 current: metaflow.metaflow_current.Current
 
+def get_metadata() -> str:
+    """
+    Returns the current Metadata provider.
+    
+    If this is not set explicitly using `metadata`, the default value is
+    determined through the Metaflow configuration. You can use this call to
+    check that your configuration is set up properly.
+    
+    If multiple configuration profiles are present, this call returns the one
+    selected through the `METAFLOW_PROFILE` environment variable.
+    
+    Returns
+    -------
+    str
+        Information about the Metadata provider currently selected. This information typically
+        returns provider specific information (like URL for remote providers or local paths for
+        local providers).
+    """
+    ...
+
+class MetaflowNotFound(metaflow.exception.MetaflowException, metaclass=type):
+    ...
+
 class MetaflowException(Exception, metaclass=type):
     def __init__(self, msg = "", lineno = None):
         ...
@@ -80,7 +339,7 @@ def validate_tags(tags, existing_tags = None):
     ...
 
 class ArgoWorkflows(object, metaclass=type):
-    def __init__(self, name, graph, flow, code_package_sha, code_package_url, production_token, metadata, flow_datastore, environment, event_logger, monitor, tags = None, namespace = None, username = None, max_workers = None, workflow_timeout = None, workflow_priority = None, auto_emit_argo_events = False, notify_on_error = False, notify_on_success = False, notify_slack_webhook_url = None, notify_pager_duty_integration_key = None):
+    def __init__(self, name, graph: metaflow.graph.FlowGraph, flow, code_package_sha, code_package_url, production_token, metadata, flow_datastore, environment, event_logger, monitor, tags = None, namespace = None, username = None, max_workers = None, workflow_timeout = None, workflow_priority = None, auto_emit_argo_events = False, notify_on_error = False, notify_on_success = False, notify_slack_webhook_url = None, notify_pager_duty_integration_key = None):
         ...
     def __str__(self):
         ...
@@ -159,6 +418,12 @@ def validate_token(name, token_prefix, authorize, instructions_fn = None):
     """
     ...
 
+def get_run_object(pathspec: str):
+    ...
+
+def get_status_considering_run_object(status, run_obj):
+    ...
+
 def validate_run_id(workflow_name, token_prefix, authorize, run_id, instructions_fn = None):
     """
     Validates that a run_id adheres to the Argo Workflows naming rules, and

metaflow-stubs/plugins/argo/argo_workflows_decorator.pyi

@@ -1,7 +1,7 @@
 ##################################################################################
 #                       Auto-generated Metaflow stub file                        #
-# MF version: 2.12.8                                                             #
-# Generated on 2024-07-16T15:51:55.411940                                        #
+# MF version: 2.12.9                                                             #
+# Generated on 2024-07-25T18:20:16.319363                                        #
 ##################################################################################
 
 from __future__ import annotations
@@ -9,8 +9,12 @@ from __future__ import annotations
 import typing
 if typing.TYPE_CHECKING:
     import metaflow
+    import metaflow.flowspec
+    import typing
     import metaflow.metaflow_current
+    import metaflow.datastore.inputs
     import metaflow.decorators
+    import metaflow.graph
     import metaflow.events
 
 current: metaflow.metaflow_current.Current
@@ -109,6 +113,227 @@ class MetaDatum(tuple, metaclass=type):
 
 ARGO_EVENTS_WEBHOOK_URL: None
 
+class FlowSpec(object, metaclass=metaflow.flowspec._FlowSpecMeta):
+    def __init__(self, use_cli = True):
+        """
+        Construct a FlowSpec
+        
+        Parameters
+        ----------
+        use_cli : bool, default True
+            Set to True if the flow is invoked from __main__ or the command line
+        """
+        ...
+    @property
+    def script_name(self) -> str:
+        """
+        [Legacy function - do not use. Use `current` instead]
+        
+        Returns the name of the script containing the flow
+        
+        Returns
+        -------
+        str
+            A string containing the name of the script
+        """
+        ...
+    def __iter__(self):
+        """
+        [Legacy function - do not use]
+        
+        Iterate over all steps in the Flow
+        
+        Returns
+        -------
+        Iterator[graph.DAGNode]
+            Iterator over the steps in the flow
+        """
+        ...
+    def __getattr__(self, name: str):
+        ...
+    def cmd(self, cmdline, input = {}, output = []):
+        """
+        [Legacy function - do not use]
+        """
+        ...
+    @property
+    def index(self) -> typing.Optional[int]:
+        """
+        The index of this foreach branch.
+        
+        In a foreach step, multiple instances of this step (tasks) will be executed,
+        one for each element in the foreach. This property returns the zero based index
+        of the current task. If this is not a foreach step, this returns None.
+        
+        If you need to know the indices of the parent tasks in a nested foreach, use
+        `FlowSpec.foreach_stack`.
+        
+        Returns
+        -------
+        int, optional
+            Index of the task in a foreach step.
+        """
+        ...
+    @property
+    def input(self) -> typing.Optional[typing.Any]:
+        """
+        The value of the foreach artifact in this foreach branch.
+        
+        In a foreach step, multiple instances of this step (tasks) will be executed,
+        one for each element in the foreach. This property returns the element passed
+        to the current task. If this is not a foreach step, this returns None.
+        
+        If you need to know the values of the parent tasks in a nested foreach, use
+        `FlowSpec.foreach_stack`.
+        
+        Returns
+        -------
+        object, optional
+            Input passed to the foreach task.
+        """
+        ...
+    def foreach_stack(self) -> typing.Optional[typing.List[typing.Tuple[int, int, typing.Any]]]:
+        """
+        Returns the current stack of foreach indexes and values for the current step.
+        
+        Use this information to understand what data is being processed in the current
+        foreach branch. For example, considering the following code:
+        ```
+        @step
+        def root(self):
+            self.split_1 = ['a', 'b', 'c']
+            self.next(self.nest_1, foreach='split_1')
+        
+        @step
+        def nest_1(self):
+            self.split_2 = ['d', 'e', 'f', 'g']
+            self.next(self.nest_2, foreach='split_2'):
+        
+        @step
+        def nest_2(self):
+            foo = self.foreach_stack()
+        ```
+        
+        `foo` will take the following values in the various tasks for nest_2:
+        ```
+            [(0, 3, 'a'), (0, 4, 'd')]
+            [(0, 3, 'a'), (1, 4, 'e')]
+            ...
+            [(0, 3, 'a'), (3, 4, 'g')]
+            [(1, 3, 'b'), (0, 4, 'd')]
+            ...
+        ```
+        where each tuple corresponds to:
+        
+        - The index of the task for that level of the loop.
+        - The number of splits for that level of the loop.
+        - The value for that level of the loop.
+        
+        Note that the last tuple returned in a task corresponds to:
+        
+        - 1st element: value returned by `self.index`.
+        - 3rd element: value returned by `self.input`.
+        
+        Returns
+        -------
+        List[Tuple[int, int, Any]]
+            An array describing the current stack of foreach steps.
+        """
+        ...
+    def merge_artifacts(self, inputs: metaflow.datastore.inputs.Inputs, exclude: typing.Optional[typing.List[str]] = None, include: typing.Optional[typing.List[str]] = None):
+        """
+        Helper function for merging artifacts in a join step.
+        
+        This function takes all the artifacts coming from the branches of a
+        join point and assigns them to self in the calling step. Only artifacts
+        not set in the current step are considered. If, for a given artifact, different
+        values are present on the incoming edges, an error will be thrown and the artifacts
+        that conflict will be reported.
+        
+        As a few examples, in the simple graph: A splitting into B and C and joining in D:
+        ```
+        A:
+          self.x = 5
+          self.y = 6
+        B:
+          self.b_var = 1
+          self.x = from_b
+        C:
+          self.x = from_c
+        
+        D:
+          merge_artifacts(inputs)
+        ```
+        In D, the following artifacts are set:
+          - `y` (value: 6), `b_var` (value: 1)
+          - if `from_b` and `from_c` are the same, `x` will be accessible and have value `from_b`
+          - if `from_b` and `from_c` are different, an error will be thrown. To prevent this error,
+            you need to manually set `self.x` in D to a merged value (for example the max) prior to
+            calling `merge_artifacts`.
+        
+        Parameters
+        ----------
+        inputs : Inputs
+            Incoming steps to the join point.
+        exclude : List[str], optional, default None
+            If specified, do not consider merging artifacts with a name in `exclude`.
+            Cannot specify if `include` is also specified.
+        include : List[str], optional, default None
+            If specified, only merge artifacts specified. Cannot specify if `exclude` is
+            also specified.
+        
+        Raises
+        ------
+        MetaflowException
+            This exception is thrown if this is not called in a join step.
+        UnhandledInMergeArtifactsException
+            This exception is thrown in case of unresolved conflicts.
+        MissingInMergeArtifactsException
+            This exception is thrown in case an artifact specified in `include` cannot
+            be found.
+        """
+        ...
+    def next(self, *dsts: typing.Callable[..., None], **kwargs):
+        """
+        Indicates the next step to execute after this step has completed.
+        
+        This statement should appear as the last statement of each step, except
+        the end step.
+        
+        There are several valid formats to specify the next step:
+        
+        - Straight-line connection: `self.next(self.next_step)` where `next_step` is a method in
+          the current class decorated with the `@step` decorator.
+        
+        - Static fan-out connection: `self.next(self.step1, self.step2, ...)` where `stepX` are
+          methods in the current class decorated with the `@step` decorator.
+        
+        - Foreach branch:
+          ```
+          self.next(self.foreach_step, foreach='foreach_iterator')
+          ```
+          In this situation, `foreach_step` is a method in the current class decorated with the
+          `@step` decorator and `foreach_iterator` is a variable name in the current class that
+          evaluates to an iterator. A task will be launched for each value in the iterator and
+          each task will execute the code specified by the step `foreach_step`.
+        
+        Parameters
+        ----------
+        dsts : Callable[..., None]
+            One or more methods annotated with `@step`.
+        
+        Raises
+        ------
+        InvalidNextException
+            Raised if the format of the arguments does not match one of the ones given above.
+        """
+        ...
+    def __str__(self):
+        ...
+    def __getstate__(self):
+        ...
+    ...
+
 class ArgoEvent(object, metaclass=type):
     def __init__(self, name, url = None, payload = None, access_token = None):
         ...
@@ -163,7 +388,7 @@ class ArgoEvent(object, metaclass=type):
 class ArgoWorkflowsInternalDecorator(metaflow.decorators.StepDecorator, metaclass=type):
     def task_pre_step(self, step_name, task_datastore, metadata, run_id, task_id, flow, graph, retry_count, max_user_code_retries, ubf_context, inputs):
         ...
-    def task_finished(self, step_name, flow, graph, is_task_ok, retry_count, max_user_code_retries):
+    def task_finished(self, step_name, flow: metaflow.flowspec.FlowSpec, graph: metaflow.graph.FlowGraph, is_task_ok, retry_count, max_user_code_retries):
         ...
     ...