ob-metaflow 2.11.15.2__py2.py3-none-any.whl → 2.11.16.1__py2.py3-none-any.whl

metaflow/plugins/kubernetes/kubernetes_decorator.py

@@ -2,7 +2,6 @@ import json
 import os
 import platform
 import sys
-import time
 
 from metaflow import current
 from metaflow.decorators import StepDecorator
@@ -21,7 +20,6 @@ from metaflow.metaflow_config import (
     KUBERNETES_PERSISTENT_VOLUME_CLAIMS,
     KUBERNETES_TOLERATIONS,
     KUBERNETES_SERVICE_ACCOUNT,
-    KUBERNETES_PORT,
     KUBERNETES_SHARED_MEMORY,
     KUBERNETES_PORT,
     KUBERNETES_CPU,
@@ -31,10 +29,11 @@ from metaflow.metaflow_config import (
 from metaflow.plugins.resources_decorator import ResourcesDecorator
 from metaflow.plugins.timeout_decorator import get_run_time_limit_for_task
 from metaflow.sidecar import Sidecar
-from metaflow.unbounded_foreach import UBF_CONTROL
 
 from ..aws.aws_utils import get_docker_registry, get_ec2_instance_metadata
 from .kubernetes import KubernetesException, parse_kube_keyvalue_list
+from metaflow.unbounded_foreach import UBF_CONTROL
+from .kubernetes_jobsets import TaskIdConstructor
 
 from metaflow.metaflow_config import MAX_MEMORY_PER_TASK, MAX_CPU_PER_TASK
 
@@ -244,6 +243,16 @@ class KubernetesDecorator(StepDecorator):
                 "Kubernetes. Please use one or the other.".format(step=step)
             )
 
+        if any([deco.name == "parallel" for deco in decos]) and any(
+            [deco.name == "catch" for deco in decos]
+        ):
+            raise MetaflowException(
+                "Step *{step}* contains a @parallel decorator "
+                "with the @catch decorator. @catch is not supported with @parallel on Kubernetes.".format(
+                    step=step
+                )
+            )
+
         # Set run time limit for the Kubernetes job.
         self.run_time_limit = get_run_time_limit_for_task(decos)
         if self.run_time_limit < 60:
@@ -327,7 +336,7 @@ class KubernetesDecorator(StepDecorator):
 
         if self.attributes["shared_memory"]:
             if not (
-                isinstance(self.attributes["shared_memory"], (int, unicode, basestring))
+                isinstance(self.attributes["shared_memory"], int)
                 and int(self.attributes["shared_memory"]) > 0
             ):
                 raise KubernetesException(
@@ -433,6 +442,10 @@ class KubernetesDecorator(StepDecorator):
                 "METAFLOW_KUBERNETES_SERVICE_ACCOUNT_NAME"
             ]
             meta["kubernetes-node-ip"] = os.environ["METAFLOW_KUBERNETES_NODE_IP"]
+            if os.environ.get("METAFLOW_KUBERNETES_JOBSET_NAME"):
+                meta["kubernetes-jobset-name"] = os.environ[
+                    "METAFLOW_KUBERNETES_JOBSET_NAME"
+                ]
 
             # TODO (savin): Introduce equivalent support for Microsoft Azure and
             #               Google Cloud Platform
@@ -465,25 +478,22 @@ class KubernetesDecorator(StepDecorator):
             self._save_logs_sidecar = Sidecar("save_logs_periodically")
             self._save_logs_sidecar.start()
 
-        num_parallel = int(os.environ.get("WORLD_SIZE", 0))
-        if num_parallel >= 1:
-            if ubf_context == UBF_CONTROL:
-                control_task_id = current.task_id
-                top_task_id = control_task_id.replace("control-", "")
-                mapper_task_ids = [control_task_id] + [
-                    "%s-node-%d" % (top_task_id, node_idx)
-                    for node_idx in range(1, num_parallel)
-                ]
-                flow._control_mapper_tasks = [
-                    "%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
-            else:
-                worker_job_rank = int(os.environ["RANK"])
-                os.environ["RANK"] = str(worker_job_rank + 1)
+        num_parallel = None
+        if hasattr(flow, "_parallel_ubf_iter"):
+            num_parallel = flow._parallel_ubf_iter.num_parallel
+
+        if num_parallel and num_parallel >= 1 and ubf_context == UBF_CONTROL:
+            control_task_id, worker_task_ids = TaskIdConstructor.join_step_task_ids(
+                num_parallel
+            )
+            mapper_task_ids = [control_task_id] + worker_task_ids
+            flow._control_mapper_tasks = [
+                "%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
 
-        if num_parallel >= 1:
+        if num_parallel and num_parallel > 1:
             _setup_multinode_environment()
 
     def task_finished(
@@ -513,44 +523,6 @@ class KubernetesDecorator(StepDecorator):
             # Best effort kill
             pass
 
-        if is_task_ok and len(getattr(flow, "_control_mapper_tasks", [])) > 1:
-            self._wait_for_mapper_tasks(flow, step_name)
-
-    def _wait_for_mapper_tasks(self, flow, step_name):
-        """
-        When launching multinode task with UBF, need to wait for the secondary
-        tasks to finish cleanly and produce their output before exiting the
-        main task. Otherwise, the main task finishing will cause secondary nodes
-        to terminate immediately, and possibly prematurely.
-        """
-        from metaflow import Step  # avoid circular dependency
-
-        TIMEOUT = 600
-        last_completion_timeout = time.time() + TIMEOUT
-        print("Waiting for batch secondary tasks to finish")
-        while last_completion_timeout > time.time():
-            time.sleep(2)
-            try:
-                step_path = "%s/%s/%s" % (flow.name, current.run_id, step_name)
-                tasks = [task for task in Step(step_path)]
-                if len(tasks) == len(flow._control_mapper_tasks):
-                    if all(
-                        task.finished_at is not None for task in tasks
-                    ):  # for some reason task.finished fails
-                        return True
-                else:
-                    print(
-                        "Waiting for all parallel tasks to finish. Finished: {}/{}".format(
-                            len(tasks),
-                            len(flow._control_mapper_tasks),
-                        )
-                    )
-            except Exception as e:
-                pass
-        raise Exception(
-            "Batch secondary workers did not finish in %s seconds" % TIMEOUT
-        )
-
     @classmethod
     def _save_package_once(cls, flow_datastore, package):
         if cls.package_url is None:
@@ -558,8 +530,19 @@ class KubernetesDecorator(StepDecorator):
                 [package.blob], len_hint=1
             )[0]
 
+
 def _setup_multinode_environment():
     import socket
+
     os.environ["MF_PARALLEL_MAIN_IP"] = socket.gethostbyname(os.environ["MASTER_ADDR"])
     os.environ["MF_PARALLEL_NUM_NODES"] = os.environ["WORLD_SIZE"]
-    os.environ["MF_PARALLEL_NODE_INDEX"] = os.environ["RANK"]
+    if os.environ.get("CONTROL_INDEX") is not None:
+        os.environ["MF_PARALLEL_NODE_INDEX"] = str(0)
+    elif os.environ.get("WORKER_REPLICA_INDEX") is not None:
+        os.environ["MF_PARALLEL_NODE_INDEX"] = str(
+            int(os.environ["WORKER_REPLICA_INDEX"]) + 1
+        )
+    else:
+        raise MetaflowException(
+            "Jobset related ENV vars called $CONTROL_INDEX or $WORKER_REPLICA_INDEX not found"
+        )