[@parallel on Kubernetes] support for Jobsets

metaflow/metaflow_config.py

@@ -331,6 +331,9 @@
 ARGO_WORKFLOWS_KUBERNETES_SECRETS = from_conf("ARGO_WORKFLOWS_KUBERNETES_SECRETS", "")
 ARGO_WORKFLOWS_ENV_VARS_TO_SKIP = from_conf("ARGO_WORKFLOWS_ENV_VARS_TO_SKIP", "")
 
+KUBERNETES_JOBSET_GROUP = from_conf("KUBERNETES_JOBSET_GROUP", "jobset.x-k8s.io")
+KUBERNETES_JOBSET_VERSION = from_conf("KUBERNETES_JOBSET_VERSION", "v1alpha2")
+
 ##
 # Argo Events Configuration
 ##

metaflow/plugins/argo/argo_workflows.py

@@ -839,6 +839,11 @@ def _dag_templates(self):
         def _visit(
             node, exit_node=None, templates=None, dag_tasks=None, parent_foreach=None
         ):
+            if node.parallel_foreach:
+                raise ArgoWorkflowsException(
+                    "Deploying flows with @parallel decorator(s) "
+                    "as Argo Workflows is not supported currently."
+                )
             # Every for-each node results in a separate subDAG and an equivalent
             # DAGTemplate rooted at the child of the for-each node. Each DAGTemplate
             # has a unique name - the top-level DAGTemplate is named as the name of

metaflow/plugins/kubernetes/kubernetes.py

@@ -3,9 +3,9 @@
 import os
 import re
 import shlex
+import copy
 import time
 from typing import Dict, List, Optional
-import uuid
 from uuid import uuid4
 
 from metaflow import current, util
@@ -67,6 +67,12 @@ class KubernetesKilledException(MetaflowException):
     headline = "Kubernetes Batch job killed"
 
 
+def _extract_labels_and_annotations_from_job_spec(job_spec):
+    annotations = job_spec.template.metadata.annotations
+    labels = job_spec.template.metadata.labels
+    return copy.copy(annotations), copy.copy(labels)
+
+
 class Kubernetes(object):
     def __init__(
         self,
@@ -141,9 +147,64 @@ def _command(
         return shlex.split('bash -c "%s"' % cmd_str)
 
     def launch_job(self, **kwargs):
-        self._job = self.create_job(**kwargs).execute()
+        if (
+            "num_parallel" in kwargs
+            and kwargs["num_parallel"]
+            and int(kwargs["num_parallel"]) > 0
+        ):
+            job = self.create_job_object(**kwargs)
+            spec = job.create_job_spec()
+            # `kwargs["step_cli"]` is setting `ubf_context` as control to ALL pods.
+            # This will be modified by the KubernetesJobSet object
+            annotations, labels = _extract_labels_and_annotations_from_job_spec(spec)
+            self._job = self.create_jobset(
+                job_spec=spec,
+                run_id=kwargs["run_id"],
+                step_name=kwargs["step_name"],
+                task_id=kwargs["task_id"],
+                namespace=kwargs["namespace"],
+                env=kwargs["env"],
+                num_parallel=kwargs["num_parallel"],
+                port=kwargs["port"],
+                annotations=annotations,
+                labels=labels,
+            ).execute()
+        else:
+            kwargs["name_pattern"] = "t-{uid}-".format(uid=str(uuid4())[:8])
+            self._job = self.create_job_object(**kwargs).create().execute()
+
+    def create_jobset(
+        self,
+        job_spec=None,
+        run_id=None,
+        step_name=None,
+        task_id=None,
+        namespace=None,
+        env=None,
+        num_parallel=None,
+        port=None,
+        annotations=None,
+        labels=None,
+    ):
+        if env is None:
+            env = {}
 
-    def create_job(
+        _prefix = str(uuid4())[:6]
+        js = KubernetesClient().jobset(
+            name="js-%s" % _prefix,
+            run_id=run_id,
+            task_id=task_id,
+            step_name=step_name,
+            namespace=namespace,
+            labels=self._get_labels(labels),
+            annotations=annotations,
+            num_parallel=num_parallel,
+            job_spec=job_spec,
+            port=port,
+        )
+        return js
+
+    def create_job_object(
         self,
         flow_name,
         run_id,
@@ -177,14 +238,15 @@ def create_job(
         labels=None,
         shared_memory=None,
         port=None,
+        name_pattern=None,
+        num_parallel=None,
     ):
         if env is None:
             env = {}
-
         job = (
             KubernetesClient()
             .job(
-                generate_name="t-{uid}-".format(uid=str(uuid4())[:8]),
+                generate_name=name_pattern,
                 namespace=namespace,
                 service_account=service_account,
                 secrets=secrets,
@@ -218,6 +280,7 @@ def create_job(
                 persistent_volume_claims=persistent_volume_claims,
                 shared_memory=shared_memory,
                 port=port,
+                num_parallel=num_parallel,
             )
             .environment_variable("METAFLOW_CODE_SHA", code_package_sha)
             .environment_variable("METAFLOW_CODE_URL", code_package_url)
@@ -336,6 +399,9 @@ def create_job(
             .label("app.kubernetes.io/part-of", "metaflow")
         )
 
+        return job
+
+    def create_k8sjob(self, job):
         return job.create()
 
     def wait(self, stdout_location, stderr_location, echo=None):
@@ -370,7 +436,7 @@ def wait_for_launch(job):
                     t = time.time()
                 time.sleep(update_delay(time.time() - start_time))
 
-        prefix = b"[%s] " % util.to_bytes(self._job.id)
+        _make_prefix = lambda: b"[%s] " % util.to_bytes(self._job.id)
 
         stdout_tail = get_log_tailer(stdout_location, self._datastore.TYPE)
         stderr_tail = get_log_tailer(stderr_location, self._datastore.TYPE)
@@ -380,7 +446,7 @@ def wait_for_launch(job):
 
         # 2) Tail logs until the job has finished
         tail_logs(
-            prefix=prefix,
+            prefix=_make_prefix(),
             stdout_tail=stdout_tail,
             stderr_tail=stderr_tail,
             echo=echo,
@@ -396,7 +462,6 @@ def wait_for_launch(job):
         #        exists prior to calling S3Tail and note the user about
         #        truncated logs if it doesn't.
         # TODO : For hard crashes, we can fetch logs from the pod.
-
         if self._job.has_failed:
             exit_code, reason = self._job.reason
             msg = next(

metaflow/plugins/kubernetes/kubernetes_cli.py

@@ -7,11 +7,17 @@
 from metaflow._vendor import click
 from metaflow.exception import METAFLOW_EXIT_DISALLOW_RETRY, CommandException
 from metaflow.metadata.util import sync_local_metadata_from_datastore
+from metaflow.unbounded_foreach import UBF_CONTROL, UBF_TASK
 from metaflow.metaflow_config import DATASTORE_LOCAL_DIR, KUBERNETES_LABELS
 from metaflow.mflog import TASK_LOG_SOURCE
 import metaflow.tracing as tracing
 
-from .kubernetes import Kubernetes, KubernetesKilledException, parse_kube_keyvalue_list
+from .kubernetes import (
+    Kubernetes,
+    KubernetesKilledException,
+    parse_kube_keyvalue_list,
+    KubernetesException,
+)
 from .kubernetes_decorator import KubernetesDecorator
 
 
@@ -109,6 +115,15 @@ def kubernetes():
 )
 @click.option("--shared-memory", default=None, help="Size of shared memory in MiB")
 @click.option("--port", default=None, help="Port number to expose from the container")
+@click.option(
+    "--ubf-context", default=None, type=click.Choice([None, UBF_CONTROL, UBF_TASK])
+)
+@click.option(
+    "--num-parallel",
+    default=None,
+    type=int,
+    help="Number of parallel nodes to run as a multi-node job.",
+)
 @click.pass_context
 def step(
     ctx,
@@ -136,6 +151,7 @@ def step(
     tolerations=None,
     shared_memory=None,
     port=None,
+    num_parallel=None,
     **kwargs
 ):
     def echo(msg, stream="stderr", job_id=None, **kwargs):
@@ -167,6 +183,12 @@ def echo(msg, stream="stderr", job_id=None, **kwargs):
         kwargs["input_paths"] = "".join("${%s}" % s for s in split_vars.keys())
         env.update(split_vars)
 
+    if num_parallel is not None and num_parallel <= 1:
+        raise KubernetesException(
+            "Using @parallel with `num_parallel` <= 1 is not supported with Kubernetes. "
+            "Please set the value of `num_parallel` to be greater than 1."
+        )
+
     # Set retry policy.
     retry_count = int(kwargs.get("retry_count", 0))
     retry_deco = [deco for deco in node.decorators if deco.name == "retry"]
@@ -251,6 +273,7 @@ def _sync_metadata():
                 tolerations=tolerations,
                 shared_memory=shared_memory,
                 port=port,
+                num_parallel=num_parallel,
             )
     except Exception as e:
         traceback.print_exc(chain=False)

metaflow/plugins/kubernetes/kubernetes_client.py

@@ -4,7 +4,7 @@
 
 from metaflow.exception import MetaflowException
 
-from .kubernetes_job import KubernetesJob
+from .kubernetes_job import KubernetesJob, KubernetesJobSet
 
 
 CLIENT_REFRESH_INTERVAL_SECONDS = 300
@@ -61,5 +61,8 @@ def get(self):
 
         return self._client
 
+    def jobset(self, **kwargs):
+        return KubernetesJobSet(self, **kwargs)
+
     def job(self, **kwargs):
         return KubernetesJob(self, **kwargs)

metaflow/plugins/kubernetes/kubernetes_decorator.py

@@ -32,6 +32,8 @@
 
 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
 
 try:
     unicode
@@ -239,11 +241,15 @@ def step_init(self, flow, graph, step, decos, environment, flow_datastore, logge
                 "Kubernetes. Please use one or the other.".format(step=step)
             )
 
-        for deco in decos:
-            if getattr(deco, "IS_PARALLEL", False):
-                raise KubernetesException(
-                    "@kubernetes does not support parallel execution currently."
+        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)
@@ -421,6 +427,10 @@ def task_pre_step(
                 "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
@@ -453,6 +463,24 @@ def task_pre_step(
             self._save_logs_sidecar = Sidecar("save_logs_periodically")
             self._save_logs_sidecar.start()
 
+        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 and num_parallel > 1:
+            _setup_multinode_environment()
+
     def task_finished(
         self, step_name, flow, graph, is_task_ok, retry_count, max_retries
     ):
@@ -486,3 +514,20 @@ def _save_package_once(cls, flow_datastore, package):
             cls.package_url, cls.package_sha = flow_datastore.save_data(
                 [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"]
+    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"
+        )