metaflow 2.12.34__py2.py3-none-any.whl → 2.12.36__py2.py3-none-any.whl

metaflow/metaflow_config.py

@@ -377,6 +377,8 @@ KUBERNETES_PORT = from_conf("KUBERNETES_PORT", None)
 KUBERNETES_CPU = from_conf("KUBERNETES_CPU", None)
 KUBERNETES_MEMORY = from_conf("KUBERNETES_MEMORY", None)
 KUBERNETES_DISK = from_conf("KUBERNETES_DISK", None)
+# Default kubernetes QoS class
+KUBERNETES_QOS = from_conf("KUBERNETES_QOS", "burstable")
 
 ARGO_WORKFLOWS_KUBERNETES_SECRETS = from_conf("ARGO_WORKFLOWS_KUBERNETES_SECRETS", "")
 ARGO_WORKFLOWS_ENV_VARS_TO_SKIP = from_conf("ARGO_WORKFLOWS_ENV_VARS_TO_SKIP", "")

metaflow/plugins/airflow/airflow.py

@@ -46,6 +46,7 @@ from metaflow.parameters import (
 # TODO: Move chevron to _vendor
 from metaflow.plugins.cards.card_modules import chevron
 from metaflow.plugins.kubernetes.kubernetes import Kubernetes
+from metaflow.plugins.kubernetes.kube_utils import qos_requests_and_limits
 from metaflow.plugins.timeout_decorator import get_run_time_limit_for_task
 from metaflow.util import compress_list, dict_to_cli_options, get_username
 
@@ -428,25 +429,25 @@ class Airflow(object):
             if k8s_deco.attributes["namespace"] is not None
             else "default"
         )
-
+        qos_requests, qos_limits = qos_requests_and_limits(
+            k8s_deco.attributes["qos"],
+            k8s_deco.attributes["cpu"],
+            k8s_deco.attributes["memory"],
+            k8s_deco.attributes["disk"],
+        )
         resources = dict(
-            requests={
-                "cpu": k8s_deco.attributes["cpu"],
-                "memory": "%sM" % str(k8s_deco.attributes["memory"]),
-                "ephemeral-storage": str(k8s_deco.attributes["disk"]),
-            }
+            requests=qos_requests,
+            limits={
+                **qos_limits,
+                **{
+                    "%s.com/gpu".lower()
+                    % k8s_deco.attributes["gpu_vendor"]: str(k8s_deco.attributes["gpu"])
+                    for k in [0]
+                    # Don't set GPU limits if gpu isn't specified.
+                    if k8s_deco.attributes["gpu"] is not None
+                },
+            },
         )
-        if k8s_deco.attributes["gpu"] is not None:
-            resources.update(
-                dict(
-                    limits={
-                        "%s.com/gpu".lower()
-                        % k8s_deco.attributes["gpu_vendor"]: str(
-                            k8s_deco.attributes["gpu"]
-                        )
-                    }
-                )
-            )
 
         annotations = {
             "metaflow/production_token": self.production_token,

metaflow/plugins/argo/argo_workflows.py

@@ -53,6 +53,7 @@ from metaflow.metaflow_config import (
 from metaflow.metaflow_config_funcs import config_values
 from metaflow.mflog import BASH_SAVE_LOGS, bash_capture_logs, export_mflog_env_vars
 from metaflow.parameters import deploy_time_eval
+from metaflow.plugins.kubernetes.kube_utils import qos_requests_and_limits
 from metaflow.plugins.kubernetes.kubernetes import (
     parse_kube_keyvalue_list,
     validate_kube_labels,
@@ -1842,6 +1843,13 @@ class ArgoWorkflows(object):
             if tmpfs_enabled and tmpfs_tempdir:
                 env["METAFLOW_TEMPDIR"] = tmpfs_path
 
+            qos_requests, qos_limits = qos_requests_and_limits(
+                resources["qos"],
+                resources["cpu"],
+                resources["memory"],
+                resources["disk"],
+            )
+
             # Create a ContainerTemplate for this node. Ideally, we would have
             # liked to inline this ContainerTemplate and avoid scanning the workflow
             # twice, but due to issues with variable substitution, we will have to
@@ -1905,6 +1913,7 @@ class ArgoWorkflows(object):
                     persistent_volume_claims=resources["persistent_volume_claims"],
                     shared_memory=shared_memory,
                     port=port,
+                    qos=resources["qos"],
                 )
 
                 for k, v in env.items():
@@ -2090,17 +2099,17 @@ class ArgoWorkflows(object):
                                 image=resources["image"],
                                 image_pull_policy=resources["image_pull_policy"],
                                 resources=kubernetes_sdk.V1ResourceRequirements(
-                                    requests={
-                                        "cpu": str(resources["cpu"]),
-                                        "memory": "%sM" % str(resources["memory"]),
-                                        "ephemeral-storage": "%sM"
-                                        % str(resources["disk"]),
-                                    },
+                                    requests=qos_requests,
                                     limits={
-                                        "%s.com/gpu".lower()
-                                        % resources["gpu_vendor"]: str(resources["gpu"])
-                                        for k in [0]
-                                        if resources["gpu"] is not None
+                                        **qos_limits,
+                                        **{
+                                            "%s.com/gpu".lower()
+                                            % resources["gpu_vendor"]: str(
+                                                resources["gpu"]
+                                            )
+                                            for k in [0]
+                                            if resources["gpu"] is not None
+                                        },
                                     },
                                 ),
                                 # Configure secrets
@@ -2337,7 +2346,7 @@ class ArgoWorkflows(object):
                                 "memory": "500Mi",
                             },
                         ),
-                    )
+                    ).to_dict()
                 )
             ),
             Template("capture-error-hook-fn-preflight").steps(
@@ -2688,7 +2697,7 @@ class ArgoWorkflows(object):
                             },
                         ),
                     )
-                )
+                ).to_dict()
             )
         )
 
@@ -2858,7 +2867,7 @@ class ArgoWorkflows(object):
                                         "memory": "250Mi",
                                     },
                                 ),
-                            )
+                            ).to_dict()
                         )
                     )
                     .service_account_name(ARGO_EVENTS_SERVICE_ACCOUNT)

metaflow/plugins/kubernetes/kube_utils.py

@@ -23,3 +23,32 @@ def parse_cli_options(flow_name, run_id, user, my_runs, echo):
             raise CommandException("A previous run id was not found. Specify --run-id.")
 
     return flow_name, run_id, user
+
+
+def qos_requests_and_limits(qos: str, cpu: int, memory: int, storage: int):
+    "return resource requests and limits for the kubernetes pod based on the given QoS Class"
+    # case insensitive matching for QoS class
+    qos = qos.lower()
+    # Determine the requests and limits to define chosen QoS class
+    qos_limits = {}
+    qos_requests = {}
+    if qos == "guaranteed":
+        # Guaranteed - has both cpu/memory limits. requests not required, as these will be inferred.
+        qos_limits = {
+            "cpu": str(cpu),
+            "memory": "%sM" % str(memory),
+            "ephemeral-storage": "%sM" % str(storage),
+        }
+        # NOTE: Even though Kubernetes will produce matching requests for the specified limits, this happens late in the lifecycle.
+        # We specify them explicitly here to make some K8S tooling happy, in case they rely on .resources.requests being present at time of submitting the job.
+        qos_requests = qos_limits
+    else:
+        # Burstable - not Guaranteed, and has a memory/cpu limit or request
+        qos_requests = {
+            "cpu": str(cpu),
+            "memory": "%sM" % str(memory),
+            "ephemeral-storage": "%sM" % str(storage),
+        }
+    # TODO: Add support for BestEffort once there is a use case for it.
+    # BestEffort - no limit or requests for cpu/memory
+    return qos_requests, qos_limits

metaflow/plugins/kubernetes/kubernetes.py

@@ -196,6 +196,7 @@ class Kubernetes(object):
         shared_memory=None,
         port=None,
         num_parallel=None,
+        qos=None,
     ):
         name = "js-%s" % str(uuid4())[:6]
         jobset = (
@@ -228,6 +229,7 @@ class Kubernetes(object):
                 shared_memory=shared_memory,
                 port=port,
                 num_parallel=num_parallel,
+                qos=qos,
             )
             .environment_variable("METAFLOW_CODE_SHA", code_package_sha)
             .environment_variable("METAFLOW_CODE_URL", code_package_url)
@@ -488,6 +490,7 @@ class Kubernetes(object):
         shared_memory=None,
         port=None,
         name_pattern=None,
+        qos=None,
     ):
         if env is None:
             env = {}
@@ -528,6 +531,7 @@ class Kubernetes(object):
                 persistent_volume_claims=persistent_volume_claims,
                 shared_memory=shared_memory,
                 port=port,
+                qos=qos,
             )
             .environment_variable("METAFLOW_CODE_SHA", code_package_sha)
             .environment_variable("METAFLOW_CODE_URL", code_package_url)

metaflow/plugins/kubernetes/kubernetes_cli.py

@@ -126,6 +126,12 @@ def kubernetes():
     type=int,
     help="Number of parallel nodes to run as a multi-node job.",
 )
+@click.option(
+    "--qos",
+    default=None,
+    type=str,
+    help="Quality of Service class for the Kubernetes pod",
+)
 @click.pass_context
 def step(
     ctx,
@@ -154,6 +160,7 @@ def step(
     shared_memory=None,
     port=None,
     num_parallel=None,
+    qos=None,
     **kwargs
 ):
     def echo(msg, stream="stderr", job_id=None, **kwargs):
@@ -294,6 +301,7 @@ def step(
                 shared_memory=shared_memory,
                 port=port,
                 num_parallel=num_parallel,
+                qos=qos,
             )
     except Exception as e:
         traceback.print_exc(chain=False)

metaflow/plugins/kubernetes/kubernetes_decorator.py

@@ -26,6 +26,7 @@ from metaflow.metaflow_config import (
     KUBERNETES_SERVICE_ACCOUNT,
     KUBERNETES_SHARED_MEMORY,
     KUBERNETES_TOLERATIONS,
+    KUBERNETES_QOS,
 )
 from metaflow.plugins.resources_decorator import ResourcesDecorator
 from metaflow.plugins.timeout_decorator import get_run_time_limit_for_task
@@ -41,6 +42,8 @@ except NameError:
     unicode = str
     basestring = str
 
+SUPPORTED_KUBERNETES_QOS_CLASSES = ["Guaranteed", "Burstable"]
+
 
 class KubernetesDecorator(StepDecorator):
     """
@@ -109,6 +112,8 @@ class KubernetesDecorator(StepDecorator):
     hostname_resolution_timeout: int, default 10 * 60
         Timeout in seconds for the workers tasks in the gang scheduled cluster to resolve the hostname of control task.
         Only applicable when @parallel is used.
+    qos: str, default: Burstable
+        Quality of Service class to assign to the pod. Supported values are: Guaranteed, Burstable, BestEffort
     """
 
     name = "kubernetes"
@@ -136,6 +141,7 @@ class KubernetesDecorator(StepDecorator):
         "compute_pool": None,
         "executable": None,
         "hostname_resolution_timeout": 10 * 60,
+        "qos": KUBERNETES_QOS,
     }
     package_url = None
     package_sha = None
@@ -259,6 +265,17 @@ class KubernetesDecorator(StepDecorator):
         self.step = step
         self.flow_datastore = flow_datastore
 
+        if (
+            self.attributes["qos"] is not None
+            # case insensitive matching.
+            and self.attributes["qos"].lower()
+            not in [c.lower() for c in SUPPORTED_KUBERNETES_QOS_CLASSES]
+        ):
+            raise MetaflowException(
+                "*%s* is not a valid Kubernetes QoS class. Choose one of the following: %s"
+                % (self.attributes["qos"], ", ".join(SUPPORTED_KUBERNETES_QOS_CLASSES))
+            )
+
         if any([deco.name == "batch" for deco in decos]):
             raise MetaflowException(
                 "Step *{step}* is marked for execution both on AWS Batch and "

metaflow/plugins/kubernetes/kubernetes_job.py

@@ -15,6 +15,8 @@ from .kubernetes_jobsets import (
     KubernetesJobSet,
 )  # We need this import for Kubernetes Client.
 
+from .kube_utils import qos_requests_and_limits
+
 
 class KubernetesJobException(MetaflowException):
     headline = "Kubernetes job error"
@@ -74,6 +76,13 @@ class KubernetesJob(object):
             if self._kwargs["shared_memory"]
             else None
         )
+        qos_requests, qos_limits = qos_requests_and_limits(
+            self._kwargs["qos"],
+            self._kwargs["cpu"],
+            self._kwargs["memory"],
+            self._kwargs["disk"],
+        )
+
         return client.V1JobSpec(
             # Retries are handled by Metaflow when it is responsible for
             # executing the flow. The responsibility is moved to Kubernetes
@@ -154,20 +163,18 @@ class KubernetesJob(object):
                             image_pull_policy=self._kwargs["image_pull_policy"],
                             name=self._kwargs["step_name"].replace("_", "-"),
                             resources=client.V1ResourceRequirements(
-                                requests={
-                                    "cpu": str(self._kwargs["cpu"]),
-                                    "memory": "%sM" % str(self._kwargs["memory"]),
-                                    "ephemeral-storage": "%sM"
-                                    % str(self._kwargs["disk"]),
-                                },
+                                requests=qos_requests,
                                 limits={
-                                    "%s.com/gpu".lower()
-                                    % self._kwargs["gpu_vendor"]: str(
-                                        self._kwargs["gpu"]
-                                    )
-                                    for k in [0]
-                                    # Don't set GPU limits if gpu isn't specified.
-                                    if self._kwargs["gpu"] is not None
+                                    **qos_limits,
+                                    **{
+                                        "%s.com/gpu".lower()
+                                        % self._kwargs["gpu_vendor"]: str(
+                                            self._kwargs["gpu"]
+                                        )
+                                        for k in [0]
+                                        # Don't set GPU limits if gpu isn't specified.
+                                        if self._kwargs["gpu"] is not None
+                                    },
                                 },
                             ),
                             volume_mounts=(

metaflow/plugins/kubernetes/kubernetes_jobsets.py

@@ -9,6 +9,8 @@ from metaflow.metaflow_config import KUBERNETES_JOBSET_GROUP, KUBERNETES_JOBSET_
 from metaflow.tracing import inject_tracing_vars
 from metaflow.metaflow_config import KUBERNETES_SECRETS
 
+from .kube_utils import qos_requests_and_limits
+
 
 class KubernetesJobsetException(MetaflowException):
     headline = "Kubernetes jobset error"
@@ -554,7 +556,12 @@ class JobSetSpec(object):
             if self._kwargs["shared_memory"]
             else None
         )
-
+        qos_requests, qos_limits = qos_requests_and_limits(
+            self._kwargs["qos"],
+            self._kwargs["cpu"],
+            self._kwargs["memory"],
+            self._kwargs["disk"],
+        )
         return dict(
             name=self.name,
             template=client.api_client.ApiClient().sanitize_for_serialization(
@@ -653,21 +660,18 @@ class JobSetSpec(object):
                                             "_", "-"
                                         ),
                                         resources=client.V1ResourceRequirements(
-                                            requests={
-                                                "cpu": str(self._kwargs["cpu"]),
-                                                "memory": "%sM"
-                                                % str(self._kwargs["memory"]),
-                                                "ephemeral-storage": "%sM"
-                                                % str(self._kwargs["disk"]),
-                                            },
+                                            requests=qos_requests,
                                             limits={
-                                                "%s.com/gpu".lower()
-                                                % self._kwargs["gpu_vendor"]: str(
-                                                    self._kwargs["gpu"]
-                                                )
-                                                for k in [0]
-                                                # Don't set GPU limits if gpu isn't specified.
-                                                if self._kwargs["gpu"] is not None
+                                                **qos_limits,
+                                                **{
+                                                    "%s.com/gpu".lower()
+                                                    % self._kwargs["gpu_vendor"]: str(
+                                                        self._kwargs["gpu"]
+                                                    )
+                                                    for k in [0]
+                                                    # Don't set GPU limits if gpu isn't specified.
+                                                    if self._kwargs["gpu"] is not None
+                                                },
                                             },
                                         ),
                                         volume_mounts=(

metaflow/plugins/parallel_decorator.py

@@ -45,6 +45,8 @@ class ParallelDecorator(StepDecorator):
         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
@@ -126,6 +128,8 @@ class ParallelDecorator(StepDecorator):
                     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(
@@ -221,7 +225,6 @@ def _local_multinode_control_task_step_func(
         "%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/runtime.py

@@ -9,6 +9,7 @@ from __future__ import print_function
 import os
 import sys
 import fcntl
+import re
 import time
 import subprocess
 from datetime import datetime
@@ -49,7 +50,13 @@ PROGRESS_INTERVAL = 300  # s
 # The following is a list of the (data) artifacts used by the runtime while
 # executing a flow. These are prefetched during the resume operation by
 # leveraging the TaskDataStoreSet.
-PREFETCH_DATA_ARTIFACTS = ["_foreach_stack", "_task_ok", "_transition"]
+PREFETCH_DATA_ARTIFACTS = [
+    "_foreach_stack",
+    "_task_ok",
+    "_transition",
+    "_control_mapper_tasks",
+    "_control_task_is_mapper_zero",
+]
 RESUME_POLL_SECONDS = 60
 
 # Runtime must use logsource=RUNTIME_LOG_SOURCE for all loglines that it
@@ -269,6 +276,8 @@ class NativeRuntime(object):
         step_name,
         task_id,
         pathspec_index,
+        cloned_task_pathspec_index,
+        finished_tuple,
         ubf_context,
         generate_task_obj,
         verbose=False,
@@ -281,8 +290,13 @@ class NativeRuntime(object):
                     task.ubf_context = ubf_context
                 new_task_id = task.task_id
                 self._cloned_tasks.append(task)
-                self._cloned_task_index.add(task.task_index)
-
+                self._cloned_task_index.add(cloned_task_pathspec_index)
+                task_pathspec = "{}/{}/{}".format(self._run_id, step_name, new_task_id)
+            else:
+                task_pathspec = "{}/{}/{}".format(self._run_id, step_name, new_task_id)
+                Task.clone_pathspec_mapping[task_pathspec] = "{}/{}/{}".format(
+                    self._clone_run_id, step_name, task_id
+                )
             if verbose:
                 self._logger(
                     "Cloning task from {}/{}/{}/{} to {}/{}/{}/{}".format(
@@ -308,6 +322,8 @@ class NativeRuntime(object):
                 self._metadata,
                 origin_ds_set=self._origin_ds_set,
             )
+            self._finished[(step_name, finished_tuple)] = task_pathspec
+            self._is_cloned[task_pathspec] = True
         except Exception as e:
             self._logger(
                 "Cloning {}/{}/{}/{} failed with error: {}".format(
@@ -323,7 +339,8 @@ class NativeRuntime(object):
 
         inputs = []
 
-        ubf_mapper_tasks_to_clone = []
+        ubf_mapper_tasks_to_clone = set()
+        ubf_control_tasks = set()
         # We only clone ubf mapper tasks if the control task is complete.
         # Here we need to check which control tasks are complete, and then get the corresponding
         # mapper tasks.
@@ -331,13 +348,25 @@ class NativeRuntime(object):
             _, step_name, task_id = task_ds.pathspec.split("/")
             pathspec_index = task_ds.pathspec_index
             if task_ds["_task_ok"] and step_name != "_parameters":
-                # Only control task can have _control_mapper_tasks. We then store the corresponding mapepr task pathspecs.
+                # Control task contains "_control_mapper_tasks" but, in the case of
+                # @parallel decorator, the control task is also a mapper task so we
+                # need to distinguish this using _control_task_is_mapper_zero
                 control_mapper_tasks = (
                     []
                     if "_control_mapper_tasks" not in task_ds
                     else task_ds["_control_mapper_tasks"]
                 )
-                ubf_mapper_tasks_to_clone.extend(control_mapper_tasks)
+                if control_mapper_tasks:
+                    if task_ds.get("_control_task_is_mapper_zero", False):
+                        # Strip out the control task of list of mapper tasks
+                        ubf_control_tasks.add(control_mapper_tasks[0])
+                        ubf_mapper_tasks_to_clone.update(control_mapper_tasks[1:])
+                    else:
+                        ubf_mapper_tasks_to_clone.update(control_mapper_tasks)
+                        # Since we only add mapper tasks here, if we are not in the list
+                        # we are a control task
+                        if task_ds.pathspec not in ubf_mapper_tasks_to_clone:
+                            ubf_control_tasks.add(task_ds.pathspec)
 
         for task_ds in self._origin_ds_set:
             _, step_name, task_id = task_ds.pathspec.split("/")
@@ -350,33 +379,54 @@ class NativeRuntime(object):
             ):
                 # "_unbounded_foreach" is a special flag to indicate that the transition is an unbounded foreach.
                 # Both parent and splitted children tasks will have this flag set. The splitted control/mapper tasks
-                # have no "foreach_param" because UBF is always followed by a join step.
+                # are not foreach types because UBF is always followed by a join step.
                 is_ubf_task = (
                     "_unbounded_foreach" in task_ds and task_ds["_unbounded_foreach"]
-                ) and (self._graph[step_name].foreach_param is None)
+                ) and (self._graph[step_name].type != "foreach")
 
-                # Only the control task has "_control_mapper_tasks" artifact.
-                is_ubf_control_task = (
-                    is_ubf_task
-                    and ("_control_mapper_tasks" in task_ds)
-                    and task_ds["_control_mapper_tasks"]
-                )
-                is_ubf_mapper_tasks = is_ubf_task and (not is_ubf_control_task)
-                if is_ubf_mapper_tasks and (
+                is_ubf_control_task = task_ds.pathspec in ubf_control_tasks
+
+                is_ubf_mapper_task = is_ubf_task and (not is_ubf_control_task)
+
+                if is_ubf_mapper_task and (
                     task_ds.pathspec not in ubf_mapper_tasks_to_clone
                 ):
-                    # Skip copying UBF mapper tasks if control tasks is incomplete.
+                    # Skip copying UBF mapper tasks if control task is incomplete.
                     continue
 
                 ubf_context = None
                 if is_ubf_task:
-                    ubf_context = "ubf_test" if is_ubf_mapper_tasks else "ubf_control"
+                    ubf_context = "ubf_test" if is_ubf_mapper_task else "ubf_control"
+
+                finished_tuple = tuple(
+                    [s._replace(value=0) for s in task_ds.get("_foreach_stack", ())]
+                )
+                cloned_task_pathspec_index = pathspec_index.split("/")[1]
+                if task_ds.get("_control_task_is_mapper_zero", False):
+                    # Replace None with index 0 for control task as it is part of the
+                    # UBF (as a mapper as well)
+                    finished_tuple = finished_tuple[:-1] + (
+                        finished_tuple[-1]._replace(index=0),
+                    )
+                    # We need this reverse override though because when we check
+                    # if a task has been cloned in _queue_push, the index will be None
+                    # because the _control_task_is_mapper_zero is set in the control
+                    # task *itself* and *not* in the one that is launching the UBF nest.
+                    # This means that _translate_index will use None.
+                    cloned_task_pathspec_index = re.sub(
+                        r"(\[(?:\d+, ?)*)0\]",
+                        lambda m: (m.group(1) or "[") + "None]",
+                        cloned_task_pathspec_index,
+                    )
+
                 inputs.append(
                     (
                         step_name,
                         task_id,
                         pathspec_index,
-                        is_ubf_mapper_tasks,
+                        cloned_task_pathspec_index,
+                        finished_tuple,
+                        is_ubf_mapper_task,
                         ubf_context,
                     )
                 )
@@ -388,15 +438,19 @@ class NativeRuntime(object):
                     step_name,
                     task_id,
                     pathspec_index,
+                    cloned_task_pathspec_index,
+                    finished_tuple,
                     ubf_context=ubf_context,
-                    generate_task_obj=generate_task_obj and (not is_ubf_mapper_tasks),
+                    generate_task_obj=generate_task_obj and (not is_ubf_mapper_task),
                     verbose=verbose,
                 )
                 for (
                     step_name,
                     task_id,
                     pathspec_index,
-                    is_ubf_mapper_tasks,
+                    cloned_task_pathspec_index,
+                    finished_tuple,
+                    is_ubf_mapper_task,
                     ubf_context,
                 ) in inputs
             ]
@@ -546,7 +600,6 @@ class NativeRuntime(object):
     # Given the current task information (task_index), the type of transition,
     # and the split index, return the new task index.
     def _translate_index(self, task, next_step, type, split_index=None):
-        import re
 
         match = re.match(r"^(.+)\[(.*)\]$", task.task_index)
         if match:
@@ -640,15 +693,18 @@ class NativeRuntime(object):
                 # If the control task is cloned, all mapper tasks should have been cloned
                 # as well, so we no longer need to handle cloning of mapper tasks in runtime.
 
-                # Update _finished since these tasks were successfully
-                # run elsewhere so that join will be unblocked.
-                _, foreach_stack = task.finished_id
-                top = foreach_stack[-1]
-                bottom = list(foreach_stack[:-1])
-                for i in range(num_splits):
-                    s = tuple(bottom + [top._replace(index=i)])
-                    self._finished[(task.step, s)] = mapper_tasks[i]
-                    self._is_cloned[mapper_tasks[i]] = False
+                # Update _finished if we are not cloned. If we were cloned, we already
+                # updated _finished with the new tasks. Note that the *value* of mapper
+                # tasks is incorrect and contains the pathspec of the *cloned* run
+                # but we don't use it for anything. We could look to clean it up though
+                if not task.is_cloned:
+                    _, foreach_stack = task.finished_id
+                    top = foreach_stack[-1]
+                    bottom = list(foreach_stack[:-1])
+                    for i in range(num_splits):
+                        s = tuple(bottom + [top._replace(index=i)])
+                        self._finished[(task.step, s)] = mapper_tasks[i]
+                        self._is_cloned[mapper_tasks[i]] = False
 
             # Find and check status of control task and retrieve its pathspec
             # for retrieving unbounded foreach cardinality.
@@ -1080,7 +1136,7 @@ class Task(object):
                     # To avoid the edge case where the resume leader is selected but has not
                     # yet written the _resume_leader metadata, we will wait for a few seconds.
                     # We will wait for resume leader for at most 3 times.
-                    for resume_leader_wait_retry in range(3):
+                    for _ in range(3):
                         if ds.has_metadata("_resume_leader", add_attempt=False):
                             resume_leader = ds.load_metadata(
                                 ["_resume_leader"], add_attempt=False

metaflow/version.py

@@ -1 +1 @@
-metaflow_version = "2.12.34"
+metaflow_version = "2.12.36"

{metaflow-2.12.34.dist-info → metaflow-2.12.36.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: metaflow
-Version: 2.12.34
+Version: 2.12.36
 Summary: Metaflow: More Data Science, Less Engineering
 Author: Metaflow Developers
 Author-email: help@metaflow.org
@@ -26,7 +26,7 @@ License-File: LICENSE
 Requires-Dist: requests
 Requires-Dist: boto3
 Provides-Extra: stubs
-Requires-Dist: metaflow-stubs==2.12.34; extra == "stubs"
+Requires-Dist: metaflow-stubs==2.12.36; extra == "stubs"
 
 ![Metaflow_Logo_Horizontal_FullColor_Ribbon_Dark_RGB](https://user-images.githubusercontent.com/763451/89453116-96a57e00-d713-11ea-9fa6-82b29d4d6eff.png)
 

{metaflow-2.12.34.dist-info → metaflow-2.12.36.dist-info}/RECORD

@@ -16,7 +16,7 @@ metaflow/includefile.py,sha256=rDJnxF0U7vD3cz9hhPkKlW_KS3ToaXnlOjhjNZ__Rx4,19628
 metaflow/info_file.py,sha256=wtf2_F0M6dgiUu74AFImM8lfy5RrUw5Yj7Rgs2swKRY,686
 metaflow/integrations.py,sha256=LlsaoePRg03DjENnmLxZDYto3NwWc9z_PtU6nJxLldg,1480
 metaflow/lint.py,sha256=5rj1MlpluxyPTSINjtMoJ7viotyNzfjtBJSAihlAwMU,10870
-metaflow/metaflow_config.py,sha256=gsE6LxvZUc68HXw3Uf7Olv8Oaq5AY17nwhKvF0oXF_8,22950
+metaflow/metaflow_config.py,sha256=bA6myygTf5WAEhtPiCUPQIYKP136Ozieg7mTHLkTfps,23039
 metaflow/metaflow_config_funcs.py,sha256=5GlvoafV6SxykwfL8D12WXSfwjBN_NsyuKE_Q3gjGVE,6738
 metaflow/metaflow_current.py,sha256=pfkXmkyHeMJhxIs6HBJNBEaBDpcl5kz9Wx5mW6F_3qo,7164
 metaflow/metaflow_environment.py,sha256=rojFyGdyY56sN1HaEb1-0XX53Q3XPNnl0SaH-8xXZ8w,7987
@@ -29,14 +29,14 @@ metaflow/parameters.py,sha256=pzjG0ssuVHPyYQqWE86dS3yYChEqbT90rOUcRD4wNog,16079
 metaflow/procpoll.py,sha256=U2tE4iK_Mwj2WDyVTx_Uglh6xZ-jixQOo4wrM9OOhxg,2859
 metaflow/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 metaflow/pylint_wrapper.py,sha256=zzBY9YaSUZOGH-ypDKAv2B_7XcoyMZj-zCoCrmYqNRc,2865
-metaflow/runtime.py,sha256=RXA8-YOeGnBl7IJtoMIl4T1zEOmvFsTKv5rBxnnnDEY,68884
+metaflow/runtime.py,sha256=wWoJILHG2JGEPnQqNI07vHT6bfJuu-pxcu5kH07787M,71841
 metaflow/tagging_util.py,sha256=ctyf0Q1gBi0RyZX6J0e9DQGNkNHblV_CITfy66axXB4,2346
 metaflow/task.py,sha256=xVVLWy8NH16OlLu2VoOb1OfiFzcOVVCdQldlmb1Zb_w,29691
 metaflow/tuple_util.py,sha256=_G5YIEhuugwJ_f6rrZoelMFak3DqAR2tt_5CapS1XTY,830
 metaflow/unbounded_foreach.py,sha256=p184WMbrMJ3xKYHwewj27ZhRUsSj_kw1jlye5gA9xJk,387
 metaflow/util.py,sha256=w7oylILPaNAjtM8MR8dfUazTVBArV_CKPpqGs4HnowM,13785
 metaflow/vendor.py,sha256=FchtA9tH22JM-eEtJ2c9FpUdMn8sSb1VHuQS56EcdZk,5139
-metaflow/version.py,sha256=1QFMy_kk8f34f8BuD9G4-ZA0vcxYQo1O9WhpDuY8MXQ,29
+metaflow/version.py,sha256=DYRO3aeKT1fdQdfo72nEW7XP8dt5jMsZSc1x4ewMnmU,29
 metaflow/_vendor/__init__.py,sha256=y_CiwUD3l4eAKvTVDZeqgVujMy31cAM1qjAB-HfI-9s,353
 metaflow/_vendor/typing_extensions.py,sha256=0nUs5p1A_UrZigrAVBoOEM6TxU37zzPDUtiij1ZwpNc,110417
 metaflow/_vendor/zipp.py,sha256=ajztOH-9I7KA_4wqDYygtHa6xUBVZgFpmZ8FE74HHHI,8425
@@ -151,7 +151,7 @@ metaflow/plugins/environment_decorator.py,sha256=6m9j2B77d-Ja_l_9CTJ__0O6aB2a8Qt
 metaflow/plugins/events_decorator.py,sha256=8YSapp_sT3UzNrb6cYBJ19-wmX_CKow6OOJN8kNVnpg,26456
 metaflow/plugins/logs_cli.py,sha256=77W5UNagU2mOKSMMvrQxQmBLRzvmjK-c8dWxd-Ygbqs,11410
 metaflow/plugins/package_cli.py,sha256=-J6D4cupHfWSZ4GEFo2yy9Je9oL3owRWm5pEJwaiqd4,1649
-metaflow/plugins/parallel_decorator.py,sha256=GIjZZVTqkvtnMuGE8RNtObX6CAJavZTxttqRujGmnGs,8973
+metaflow/plugins/parallel_decorator.py,sha256=GR6LKIW7_S7AoU50Ar2_0nndVtO2epdn3LuthE0vKMQ,9127
 metaflow/plugins/project_decorator.py,sha256=eJOe0Ea7CbUCReEhR_XQvRkhV6jyRqDxM72oZI7EMCk,5336
 metaflow/plugins/resources_decorator.py,sha256=AtoOwg4mHYHYthg-CAfbfam-QiT0ViuDLDoukoDvF6Q,1347
 metaflow/plugins/retry_decorator.py,sha256=tz_2Tq6GLg3vjDBZp0KKVTk3ADlCvqaWTSf7blmFdUw,1548
@@ -160,7 +160,7 @@ metaflow/plugins/tag_cli.py,sha256=10039-0DUF0cmhudoDNrRGLWq8tCGQJ7tBsQAGAmkBQ,1
 metaflow/plugins/test_unbounded_foreach_decorator.py,sha256=33p5aCWnyk9MT5DmXcm4Q_Qnwfd4y4xvVTEfeqs4by0,5957
 metaflow/plugins/timeout_decorator.py,sha256=R-X8rKeMqd-xhfJFqskWb6ZpmZt2JB14U1BZJSRriwM,3648
 metaflow/plugins/airflow/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-metaflow/plugins/airflow/airflow.py,sha256=V9a96L5Cb1Yg5CvcbX6WDa7J8lQh5AjIvnVQod20_7k,32160
+metaflow/plugins/airflow/airflow.py,sha256=GDaKLdzzySttJfhl_OiYjkK_ubIaRKR4YgcLaKRCQLk,32293
 metaflow/plugins/airflow/airflow_cli.py,sha256=2JzrGUY9mPNyXRO6fJbgdtiOZXsEUv4Fn2MSPcoPszU,14692
 metaflow/plugins/airflow/airflow_decorator.py,sha256=IWT6M9gga8t65FR4Wi7pIZvOupk3hE75B5NRg9tMEps,1781
 metaflow/plugins/airflow/airflow_utils.py,sha256=dvRllfQeOWfDUseFnOocIGaL3gRI_A7cEHnC1w01vfk,28905
@@ -175,7 +175,7 @@ metaflow/plugins/airflow/sensors/s3_sensor.py,sha256=iDReG-7FKnumrtQg-HY6cCUAAqN
 metaflow/plugins/argo/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 metaflow/plugins/argo/argo_client.py,sha256=Z_A1TO9yw4Y-a8VAlwrFS0BwunWzXpbtik-j_xjcuHE,16303
 metaflow/plugins/argo/argo_events.py,sha256=_C1KWztVqgi3zuH57pInaE9OzABc2NnncC-zdwOMZ-w,5909
-metaflow/plugins/argo/argo_workflows.py,sha256=d-d_xvjUeb5t11kaHZ98yIrGAqYg__FhxNcQTY8LVoA,173881
+metaflow/plugins/argo/argo_workflows.py,sha256=G1yu-bBGpOyoZKRlgKdJLxNueRCv25sPuDM_gbQkBJw,174184
 metaflow/plugins/argo/argo_workflows_cli.py,sha256=NdLwzfBcTsR72qLycZBesR4Pwv48o3Z_v6OfYrZuVEY,36721
 metaflow/plugins/argo/argo_workflows_decorator.py,sha256=QdM1rK9gM-lDhyZldK8WqvFqJDvfJ7i3JPR5Uzaq2as,7887
 metaflow/plugins/argo/argo_workflows_deployer.py,sha256=6kHxEnYXJwzNCM9swI8-0AckxtPWqwhZLerYkX8fxUM,4444
@@ -281,13 +281,13 @@ metaflow/plugins/gcp/gs_tail.py,sha256=qz0QZKT-5LvL8qgZZK2yyMOwuEnx1YOz-pTSAUmwv
 metaflow/plugins/gcp/gs_utils.py,sha256=ZmIGFse1qYyvAVrwga23PQUzF6dXEDLLsZ2F-YRmvow,2030
 metaflow/plugins/gcp/includefile_support.py,sha256=OQO0IVWv4ObboL0VqEZwcDOyj9ORLdur66JToxQ84vU,3887
 metaflow/plugins/kubernetes/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-metaflow/plugins/kubernetes/kube_utils.py,sha256=fYDlvqi8jYPsWijDwT6Z2qhQswyFqv7tiwtic_I80Vg,749
-metaflow/plugins/kubernetes/kubernetes.py,sha256=_cq_4N8l40cP0kifYMDAL8Te0DnIKUhVCqee3covwdM,31642
-metaflow/plugins/kubernetes/kubernetes_cli.py,sha256=TAYOKTQegYxex5piasLc53kNQPLzxuO8FzHXxkdjPjY,13472
+metaflow/plugins/kubernetes/kube_utils.py,sha256=CbJRMn-sQyGqG-hKMBBjA6xmw15_DyQmhU8TxNyWqcQ,2124
+metaflow/plugins/kubernetes/kubernetes.py,sha256=FrIL2wRUzy4bJr6pNz3I-tNFH-OJWHJcrarJsBKRPLE,31728
+metaflow/plugins/kubernetes/kubernetes_cli.py,sha256=tvxwgBojuiezIUdum95t9fHzDs50q-a-gkPyIApWnCY,13633
 metaflow/plugins/kubernetes/kubernetes_client.py,sha256=tuvXP-QKpdeSmzVolB2R_TaacOr5DIb0j642eKcjsiM,6491
-metaflow/plugins/kubernetes/kubernetes_decorator.py,sha256=Gq3lGKA8SPh3pHDbP_FCkUQPMRrIxvbcmw6Jly5PhEY,27846
-metaflow/plugins/kubernetes/kubernetes_job.py,sha256=Cfkee8LbXC17jSXWoeNdomQRvF_8YSeXNg1gvxm6E_M,31806
-metaflow/plugins/kubernetes/kubernetes_jobsets.py,sha256=iehUEKv2KogyJKnp5jejdGP8R-TtF2aX9Wx1WpjKLvM,42030
+metaflow/plugins/kubernetes/kubernetes_decorator.py,sha256=wHwm5pQvhjI9Dcj6yBdKZGprFGrSEgd88-TZew6vgNs,28569
+metaflow/plugins/kubernetes/kubernetes_job.py,sha256=CoDzG0eEcJezfMTmgYJ4Ea9G_o5INYm0w1DvjGwJT2A,31916
+metaflow/plugins/kubernetes/kubernetes_jobsets.py,sha256=0SGOfStlh6orXVpF3s6Mu26OFR36eh4dj-sFYwg3HaA,42066
 metaflow/plugins/metadata_providers/__init__.py,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
 metaflow/plugins/metadata_providers/local.py,sha256=9UAxe9caN6kU1lkSlIoJbRGgTqsMa62cBTnyMwhqiaA,22446
 metaflow/plugins/metadata_providers/service.py,sha256=NKZfFMamx6upP6aFRJfXlfYIhySgFNzz6kbp1yPD7LA,20222
@@ -348,9 +348,9 @@ metaflow/tutorials/07-worldview/README.md,sha256=5vQTrFqulJ7rWN6r20dhot9lI2sVj9W
 metaflow/tutorials/07-worldview/worldview.ipynb,sha256=ztPZPI9BXxvW1QdS2Tfe7LBuVzvFvv0AToDnsDJhLdE,2237
 metaflow/tutorials/08-autopilot/README.md,sha256=GnePFp_q76jPs991lMUqfIIh5zSorIeWznyiUxzeUVE,1039
 metaflow/tutorials/08-autopilot/autopilot.ipynb,sha256=DQoJlILV7Mq9vfPBGW-QV_kNhWPjS5n6SJLqePjFYLY,3191
-metaflow-2.12.34.dist-info/LICENSE,sha256=nl_Lt5v9VvJ-5lWJDT4ddKAG-VZ-2IaLmbzpgYDz2hU,11343
-metaflow-2.12.34.dist-info/METADATA,sha256=nyJPMX2oVYbAC9lY9KuNPMngDomft-j5DJoeHGAnJuY,5907
-metaflow-2.12.34.dist-info/WHEEL,sha256=pxeNX5JdtCe58PUSYP9upmc7jdRPgvT0Gm9kb1SHlVw,109
-metaflow-2.12.34.dist-info/entry_points.txt,sha256=IKwTN1T3I5eJL3uo_vnkyxVffcgnRdFbKwlghZfn27k,57
-metaflow-2.12.34.dist-info/top_level.txt,sha256=v1pDHoWaSaKeuc5fKTRSfsXCKSdW1zvNVmvA-i0if3o,9
-metaflow-2.12.34.dist-info/RECORD,,
+metaflow-2.12.36.dist-info/LICENSE,sha256=nl_Lt5v9VvJ-5lWJDT4ddKAG-VZ-2IaLmbzpgYDz2hU,11343
+metaflow-2.12.36.dist-info/METADATA,sha256=QSuhOl0_WqTjlMSRhMgAt6sxTsOzHzEyA1F5fzwXe_w,5907
+metaflow-2.12.36.dist-info/WHEEL,sha256=pxeNX5JdtCe58PUSYP9upmc7jdRPgvT0Gm9kb1SHlVw,109
+metaflow-2.12.36.dist-info/entry_points.txt,sha256=IKwTN1T3I5eJL3uo_vnkyxVffcgnRdFbKwlghZfn27k,57
+metaflow-2.12.36.dist-info/top_level.txt,sha256=v1pDHoWaSaKeuc5fKTRSfsXCKSdW1zvNVmvA-i0if3o,9
+metaflow-2.12.36.dist-info/RECORD,,