ob-metaflow 2.11.13.1__py2.py3-none-any.whl → 2.19.7.1rc0__py2.py3-none-any.whl

metaflow/plugins/metadata_providers/spin.py

@@ -0,0 +1,16 @@
+from metaflow.plugins.metadata_providers.local import LocalMetadataProvider
+from metaflow.metaflow_config import DATASTORE_SPIN_LOCAL_DIR
+
+
+class SpinMetadataProvider(LocalMetadataProvider):
+    TYPE = "spin"
+    DATASTORE_DIR = DATASTORE_SPIN_LOCAL_DIR  # ".metaflow_spin"
+
+    @classmethod
+    def _get_storage_class(cls):
+        from metaflow.plugins.datastores.spin_storage import SpinStorage
+
+        return SpinStorage
+
+    def version(self):
+        return "spin"

metaflow/plugins/package_cli.py

@@ -9,49 +9,61 @@ def cli():
 
 
 @cli.group(help="Commands related to code packages.")
+@click.option(
+    "--timeout", default=60, help="Timeout for package operations in seconds."
+)
 @click.pass_obj
-def package(obj):
+def package(obj, timeout):
     # Prepare the package before any of the sub-commands are invoked.
+    # We explicitly will *not* upload it to the datastore.
     obj.package = MetaflowPackage(
-        obj.flow, obj.environment, obj.echo, obj.package_suffixes
+        obj.flow,
+        obj.environment,
+        obj.echo,
+        suffixes=obj.package_suffixes,
+        flow_datastore=None,
     )
+    obj.package_op_timeout = timeout
 
 
-@package.command(help="Output information about the current code package.")
+@package.command(help="Output information about the code package.")
 @click.pass_obj
 def info(obj):
-    obj.echo("Status of the current working directory:", fg="magenta", bold=False)
-    obj.echo_always(
-        "Hash: *%s*" % sha1(obj.package.blob).hexdigest(),
-        highlight="green",
-        highlight_bold=False,
-    )
-    obj.echo_always(
-        "Package size: *%d* KB" % (len(obj.package.blob) / 1024),
-        highlight="green",
-        highlight_bold=False,
-    )
-    num = sum(1 for _ in obj.package.path_tuples())
-    obj.echo_always(
-        "Number of files: *%d*" % num, highlight="green", highlight_bold=False
-    )
+    obj.echo_always(obj.package.show())
 
 
-@package.command(help="List files included in the code package.")
+@package.command(help="List all files included in the code package.")
+@click.option(
+    "--archive/--no-archive",
+    default=False,
+    help="If True, lists the file paths as present in the code package archive; "
+    "otherwise, lists the files on your filesystem included in the code package",
+    show_default=True,
+)
 @click.pass_obj
-def list(obj):
+def list(obj, archive=False):
+    _ = obj.package.blob_with_timeout(timeout=obj.package_op_timeout)
+    # We now have all the information about the blob
     obj.echo(
-        "Files included in the code package " "(change with --package-suffixes):",
+        "Files included in the code package (change with --package-suffixes):",
         fg="magenta",
         bold=False,
     )
-    obj.echo_always("\n".join(path for path, _ in obj.package.path_tuples()))
+    if archive:
+        obj.echo_always("\n".join(path for _, path in obj.package.path_tuples()))
+    else:
+        obj.echo_always("\n".join(path for path, _ in obj.package.path_tuples()))
 
 
-@package.command(help="Save the current code package in a tar file")
+@package.command(help="Save the current code package to a file.")
 @click.argument("path")
 @click.pass_obj
 def save(obj, path):
     with open(path, "wb") as f:
         f.write(obj.package.blob)
-    obj.echo("Code package saved in *%s*." % path, fg="magenta", bold=False)
+    obj.echo(
+        "Code package saved in *%s* with metadata: %s"
+        % (path, obj.package.package_metadata),
+        fg="magenta",
+        bold=False,
+    )

metaflow/plugins/parallel_decorator.py

@@ -1,30 +1,145 @@
+from collections import namedtuple
 from metaflow.decorators import StepDecorator
-from metaflow.unbounded_foreach import UBF_CONTROL
+from metaflow.unbounded_foreach import UBF_CONTROL, CONTROL_TASK_TAG
 from metaflow.exception import MetaflowException
+from metaflow.metadata_provider import MetaDatum
+from metaflow.metaflow_current import current, Parallel
 import os
 import sys
 
 
 class ParallelDecorator(StepDecorator):
+    """
+    MF Add To Current
+    -----------------
+    parallel -> metaflow.metaflow_current.Parallel
+        Returns a namedtuple with relevant information about the parallel task.
+
+        @@ Returns
+        -------
+        Parallel
+            `namedtuple` with the following fields:
+                - main_ip (`str`)
+                    The IP address of the control task.
+                - num_nodes (`int`)
+                    The total number of tasks created by @parallel
+                - node_index (`int`)
+                    The index of the current task in all the @parallel tasks.
+                - control_task_id (`Optional[str]`)
+                    The task ID of the control task. Available to all tasks.
+
+    is_parallel -> bool
+        True if the current step is a @parallel step.
+    """
+
     name = "parallel"
     defaults = {}
     IS_PARALLEL = True
 
-    def __init__(self, attributes=None, statically_defined=False):
-        super(ParallelDecorator, self).__init__(attributes, statically_defined)
+    def __init__(self, attributes=None, statically_defined=False, inserted_by=None):
+        super(ParallelDecorator, self).__init__(
+            attributes, statically_defined, inserted_by
+        )
 
     def runtime_step_cli(
         self, cli_args, retry_count, max_user_code_retries, ubf_context
     ):
-
         if ubf_context == UBF_CONTROL:
             num_parallel = cli_args.task.ubf_iter.num_parallel
             cli_args.command_options["num-parallel"] = str(num_parallel)
+            if os.environ.get("METAFLOW_RUNTIME_ENVIRONMENT", "local") == "local":
+                cli_args.command_options["split_index"] = "0"
 
     def step_init(
         self, flow, graph, step_name, decorators, environment, flow_datastore, logger
     ):
+        # TODO: This can be supported in the future, but for the time being we disable the transition as it leads to
+        # a UBF exception during runtime when the actual parallel-join step is conditional (switching between different join implementations from the @parallel step).
+        if graph[step_name].type == "split-switch":
+            raise MetaflowException(
+                "A @parallel step can not be a conditional switch step. Please add a join step after *%s*"
+                % step_name
+            )
         self.environment = environment
+        # Previously, the `parallel` property was a hardcoded, static property within `current`.
+        # Whenever `current.parallel` was called, it returned a named tuple with values coming from
+        # environment variables, loaded dynamically at runtime.
+        # Now, many of these environment variables are set by compute-related decorators in `task_pre_step`.
+        # This necessitates ensuring the correct ordering of the `parallel` and compute decorators if we want to
+        # statically set the namedtuple via `current._update_env` in `task_pre_step`. Hence we avoid using
+        # `current._update_env` since:
+        # - it will set a static named tuple, resolving environment variables only once (at the time of calling `current._update_env`).
+        # - we cannot guarantee the order of calling the decorator's `task_pre_step` (calling `current._update_env` may not set
+        #   the named tuple with the correct values).
+        # Therefore, we explicitly set the property in `step_init` to ensure the property can resolve the appropriate values in the named tuple
+        # when accessed at runtime.
+        setattr(
+            current.__class__,
+            "parallel",
+            property(
+                fget=lambda _: Parallel(
+                    main_ip=os.environ.get("MF_PARALLEL_MAIN_IP", "127.0.0.1"),
+                    num_nodes=int(os.environ.get("MF_PARALLEL_NUM_NODES", "1")),
+                    node_index=int(os.environ.get("MF_PARALLEL_NODE_INDEX", "0")),
+                    control_task_id=os.environ.get("MF_PARALLEL_CONTROL_TASK_ID", None),
+                )
+            ),
+        )
+
+    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,
+    ):
+        from metaflow import current
+
+        # Set `is_parallel` to `True` in `current` just like we
+        # with `is_production` in the project decorator.
+        current._update_env(
+            {
+                "is_parallel": True,
+            }
+        )
+
+        self.input_paths = [obj.pathspec for obj in inputs]
+        task_metadata_list = [
+            MetaDatum(
+                field="parallel-world-size",
+                value=flow._parallel_ubf_iter.num_parallel,
+                type="parallel-world-size",
+                tags=["attempt_id:{0}".format(0)],
+            )
+        ]
+        if ubf_context == UBF_CONTROL:
+            # A Task's tags are now those of its ancestral Run, so we are not able
+            # to rely on a task's tags to indicate the presence of a control task
+            # so, on top of adding the tags above, we also add a task metadata
+            # entry indicating that this is a "control task".
+            #
+            # Here we will also add a task metadata entry to indicate "control
+            # task". Within the metaflow repo, the only dependency of such a
+            # "control task" indicator is in the integration test suite (see
+            # Step.control_tasks() in client API).
+            task_metadata_list += [
+                MetaDatum(
+                    field="internal_task_type",
+                    value=CONTROL_TASK_TAG,
+                    type="internal_task_type",
+                    tags=["attempt_id:{0}".format(0)],
+                )
+            ]
+            flow._control_task_is_mapper_zero = True
+
+        metadata.register_metadata(run_id, step_name, task_id, task_metadata_list)
 
     def task_decorate(
         self, step_func, flow, graph, retry_count, max_user_code_retries, ubf_context
@@ -47,6 +162,7 @@ class ParallelDecorator(StepDecorator):
                 env_to_use,
                 _step_func_with_setup,
                 retry_count,
+                ",".join(self.input_paths),
             )
         else:
             return _step_func_with_setup
@@ -56,7 +172,9 @@ class ParallelDecorator(StepDecorator):
         pass
 
 
-def _local_multinode_control_task_step_func(flow, env_to_use, step_func, retry_count):
+def _local_multinode_control_task_step_func(
+    flow, env_to_use, step_func, retry_count, input_paths
+):
     """
     Used as multinode UBF control task when run in local mode.
     """
@@ -76,14 +194,12 @@ def _local_multinode_control_task_step_func(flow, env_to_use, step_func, retry_c
     num_parallel = foreach_iter.num_parallel
     os.environ["MF_PARALLEL_NUM_NODES"] = str(num_parallel)
     os.environ["MF_PARALLEL_MAIN_IP"] = "127.0.0.1"
+    os.environ["MF_PARALLEL_CONTROL_TASK_ID"] = str(current.task_id)
 
     run_id = current.run_id
     step_name = current.step_name
     control_task_id = current.task_id
-
-    (_, split_step_name, split_task_id) = control_task_id.split("-")[1:]
     # UBF handling for multinode case
-    top_task_id = control_task_id.replace("control-", "")  # chop "-0"
     mapper_task_ids = [control_task_id]
     # If we are running inside Conda, we use the base executable FIRST;
     # the conda environment will then be used when runtime_step_cli is
@@ -93,12 +209,13 @@ def _local_multinode_control_task_step_func(flow, env_to_use, step_func, retry_c
     script = sys.argv[0]
 
     # start workers
+    # TODO: Logs for worker processes are assigned to control process as of today, which
+    #       should be fixed at some point
     subprocesses = []
     for node_index in range(1, num_parallel):
-        task_id = "%s_node_%d" % (top_task_id, node_index)
+        task_id = "%s_node_%d" % (control_task_id, node_index)
         mapper_task_ids.append(task_id)
         os.environ["MF_PARALLEL_NODE_INDEX"] = str(node_index)
-        input_paths = "%s/%s/%s" % (run_id, split_step_name, split_task_id)
         # Override specific `step` kwargs.
         kwargs = cli_args.step_kwargs
         kwargs["split_index"] = str(node_index)
@@ -109,6 +226,7 @@ def _local_multinode_control_task_step_func(flow, env_to_use, step_func, retry_c
         kwargs["retry_count"] = str(retry_count)
 
         cmd = cli_args.step_command(executable, script, step_name, step_kwargs=kwargs)
+
         p = subprocess.Popen(cmd)
         subprocesses.append(p)
 
@@ -116,7 +234,6 @@ def _local_multinode_control_task_step_func(flow, env_to_use, step_func, retry_c
         "%s/%s/%s" % (run_id, step_name, mapper_task_id)
         for mapper_task_id in mapper_task_ids
     ]
-    flow._control_task_is_mapper_zero = True
 
     # run the step function ourselves
     os.environ["MF_PARALLEL_NODE_INDEX"] = "0"

metaflow/plugins/parsers.py

@@ -0,0 +1,16 @@
+from metaflow._vendor import yaml
+
+
+def yaml_parser(content: str) -> dict:
+    """
+    Parse YAML content to a dictionary.
+
+    Parameters
+    ----------
+    content : str
+
+    Returns
+    -------
+    dict
+    """
+    return yaml.safe_load(content)

metaflow/plugins/project_decorator.py

@@ -26,6 +26,24 @@ class ProjectDecorator(FlowDecorator):
         projects that use the same production scheduler. The name may
         contain only lowercase alphanumeric characters and underscores.
 
+    branch : Optional[str], default None
+        The branch to use. If not specified, the branch is set to
+        `user.<username>` unless `production` is set to `True`. This can
+        also be set on the command line using `--branch` as a top-level option.
+        It is an error to specify `branch` in the decorator and on the command line.
+
+    production : bool, default False
+        Whether or not the branch is the production branch. This can also be set on the
+        command line using `--production` as a top-level option. It is an error to specify
+        `production` in the decorator and on the command line.
+        The project branch name will be:
+          - if `branch` is specified:
+            - if `production` is True: `prod.<branch>`
+            - if `production` is False: `test.<branch>`
+          - if `branch` is not specified:
+            - if `production` is True: `prod`
+            - if `production` is False: `user.<username>`
+
     MF Add To Current
     -----------------
     project_name -> str
@@ -72,7 +90,6 @@ class ProjectDecorator(FlowDecorator):
     """
 
     name = "project"
-    defaults = {"name": None}
 
     options = {
         "production": dict(
@@ -91,19 +108,48 @@ class ProjectDecorator(FlowDecorator):
         ),
     }
 
+    defaults = {"name": None, **{k: v["default"] for k, v in options.items()}}
+
     def flow_init(
         self, flow, graph, environment, flow_datastore, metadata, logger, echo, options
     ):
         self._option_values = options
         project_name = self.attributes.get("name")
+        for op in options:
+            if (
+                op in self._user_defined_attributes
+                and options[op] != self.defaults[op]
+                and self.attributes[op] != options[op]
+            ):
+                # Exception if:
+                #  - the user provides a value in the attributes field
+                #  - AND the user provided a value in the command line (non default)
+                #  - AND the values are different
+                # Note that this won't raise an error if the user provided the default
+                # value in the command line and provided one in attribute but although
+                # slightly inconsistent, it is not incorrect.
+                raise MetaflowException(
+                    "You cannot pass %s as both a command-line argument and an attribute "
+                    "of the @project decorator." % op
+                )
+        if "branch" in self._user_defined_attributes:
+            project_branch = self.attributes["branch"]
+        else:
+            project_branch = options["branch"]
+
+        if "production" in self._user_defined_attributes:
+            project_production = self.attributes["production"]
+        else:
+            project_production = options["production"]
+
         project_flow_name, branch_name = format_name(
             flow.name,
             project_name,
-            options["production"],
-            options["branch"],
+            project_production,
+            project_branch,
             get_username(),
         )
-        is_user_branch = options["branch"] is None and not options["production"]
+        is_user_branch = project_branch is None and not project_production
         echo(
             "Project: *%s*, Branch: *%s*" % (project_name, branch_name),
             fg="magenta",
@@ -114,7 +160,7 @@ class ProjectDecorator(FlowDecorator):
                 "project_name": project_name,
                 "branch_name": branch_name,
                 "is_user_branch": is_user_branch,
-                "is_production": options["production"],
+                "is_production": project_production,
                 "project_flow_name": project_flow_name,
             }
         )