ob-metaflow 2.10.9.1__tar.gz → 2.10.9.2__tar.gz

{ob-metaflow-2.10.9.1/ob_metaflow.egg-info → ob-metaflow-2.10.9.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ob-metaflow
-Version: 2.10.9.1
+Version: 2.10.9.2
 Summary: Metaflow: More Data Science, Less Engineering
 Author: Netflix, Outerbounds & the Metaflow Community
 Author-email: help@outerbounds.co

{ob-metaflow-2.10.9.1 → ob-metaflow-2.10.9.2}/metaflow/metaflow_config.py

@@ -296,6 +296,8 @@ KUBERNETES_CONTAINER_REGISTRY = from_conf(
 )
 # Toggle for trying to fetch EC2 instance metadata
 KUBERNETES_FETCH_EC2_METADATA = from_conf("KUBERNETES_FETCH_EC2_METADATA", False)
+# Default port number to open on the pods
+KUBERNETES_PORT = from_conf("KUBERNETES_PORT", None)
 
 ARGO_WORKFLOWS_KUBERNETES_SECRETS = from_conf("ARGO_WORKFLOWS_KUBERNETES_SECRETS", "")
 ARGO_WORKFLOWS_ENV_VARS_TO_SKIP = from_conf("ARGO_WORKFLOWS_ENV_VARS_TO_SKIP", "")

{ob-metaflow-2.10.9.1 → ob-metaflow-2.10.9.2}/metaflow/plugins/argo/argo_workflows.py

@@ -836,6 +836,11 @@ class ArgoWorkflows(object):
     # Visit every node and yield the uber DAGTemplate(s).
     def _dag_templates(self):
         def _visit(node, exit_node=None, templates=None, dag_tasks=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

{ob-metaflow-2.10.9.1 → ob-metaflow-2.10.9.2}/metaflow/plugins/kubernetes/kubernetes.py

@@ -4,6 +4,7 @@ import os
 import re
 import shlex
 import time
+import copy
 from typing import Dict, List, Optional
 import uuid
 from uuid import uuid4
@@ -174,6 +175,10 @@ class Kubernetes(object):
         persistent_volume_claims=None,
         tolerations=None,
         labels=None,
+        annotations=None,
+        num_parallel=0,
+        attrs={},
+        port=None,
     ):
         if env is None:
             env = {}
@@ -213,6 +218,9 @@ class Kubernetes(object):
                 tmpfs_size=tmpfs_size,
                 tmpfs_path=tmpfs_path,
                 persistent_volume_claims=persistent_volume_claims,
+                num_parallel=num_parallel,
+                attrs=attrs,
+                port=port,
             )
             .environment_variable("METAFLOW_CODE_SHA", code_package_sha)
             .environment_variable("METAFLOW_CODE_URL", code_package_url)
@@ -266,6 +274,7 @@ class Kubernetes(object):
             # see get_datastore_root_from_config in datastore/local.py).
         )
 
+        self.num_parallel = num_parallel
         # Temporary passing of *some* environment variables. Do not rely on this
         # mechanism as it will be removed in the near future
         for k, v in config_values():
@@ -341,7 +350,7 @@ class Kubernetes(object):
             sigmoid = 1.0 / (1.0 + math.exp(-0.01 * secs_since_start + 9.0))
             return 0.5 + sigmoid * 30.0
 
-        def wait_for_launch(job):
+        def wait_for_launch(job, child_jobs):
             status = job.status
             echo(
                 "Task is starting (%s)..." % status,
@@ -351,11 +360,38 @@ class Kubernetes(object):
             t = time.time()
             start_time = time.time()
             while job.is_waiting:
-                new_status = job.status
-                if status != new_status or (time.time() - t) > 30:
-                    status = new_status
+                # new_status = job.status
+                if status != job.status or (time.time() - t) > 30:
+                    if not child_jobs:
+                        child_statuses = ""
+                    else:
+                        status_keys = set(
+                            [child_job.status for child_job in child_jobs]
+                        )
+                        status_counts = [
+                            (
+                                status,
+                                len(
+                                    [
+                                        child_job.status == status
+                                        for child_job in child_jobs
+                                    ]
+                                ),
+                            )
+                            for status in status_keys
+                        ]
+                        child_statuses = " (parallel node status: [{}])".format(
+                            ", ".join(
+                                [
+                                    "{}:{}".format(status, num)
+                                    for (status, num) in sorted(status_counts)
+                                ]
+                            )
+                        )
+
+                    status = job.status
                     echo(
-                        "Task is starting (%s)..." % status,
+                        "Task is starting (status %s)... %s" % (status, child_statuses),
                         "stderr",
                         job_id=job.id,
                     )
@@ -367,8 +403,9 @@ class Kubernetes(object):
         stdout_tail = get_log_tailer(stdout_location, self._datastore.TYPE)
         stderr_tail = get_log_tailer(stderr_location, self._datastore.TYPE)
 
+        child_jobs = []
         # 1) Loop until the job has started
-        wait_for_launch(self._job)
+        wait_for_launch(self._job, child_jobs)
 
         # 2) Tail logs until the job has finished
         tail_logs(

{ob-metaflow-2.10.9.1 → ob-metaflow-2.10.9.2}/metaflow/plugins/kubernetes/kubernetes_cli.py

@@ -107,6 +107,26 @@ def kubernetes():
     type=JSONTypeClass(),
     multiple=False,
 )
+@click.option(
+    "--labels",
+    default=None,
+    type=JSONTypeClass(),
+    multiple=False,
+)
+@click.option(
+    "--annotations",
+    default=None,
+    type=JSONTypeClass(),
+    multiple=False,
+)
+@click.option("--ubf-context", default=None, type=click.Choice([None, "ubf_control"]))
+@click.option(
+    "--num-parallel",
+    default=0,
+    type=int,
+    help="Number of parallel nodes to run as a multi-node job.",
+)
+@click.option("--port", default=None, help="port number")
 @click.pass_context
 def step(
     ctx,
@@ -132,6 +152,10 @@ def step(
     run_time_limit=None,
     persistent_volume_claims=None,
     tolerations=None,
+    labels=None,
+    annotations=None,
+    num_parallel=None,
+    port=None,
     **kwargs
 ):
     def echo(msg, stream="stderr", job_id=None, **kwargs):
@@ -177,11 +201,17 @@ def step(
         )
         time.sleep(minutes_between_retries * 60)
 
+    step_args = " ".join(util.dict_to_cli_options(kwargs))
+    num_parallel = num_parallel or 0
+    if num_parallel and num_parallel > 1:
+        # For multinode, we need to add a placeholder that can be mutated by the caller
+        step_args += " [multinode-args]"
+
     step_cli = "{entrypoint} {top_args} step {step} {step_args}".format(
         entrypoint="%s -u %s" % (executable, os.path.basename(sys.argv[0])),
         top_args=" ".join(util.dict_to_cli_options(ctx.parent.parent.params)),
         step=step_name,
-        step_args=" ".join(util.dict_to_cli_options(kwargs)),
+        step_args=step_args,
     )
 
     # Set log tailing.
@@ -207,6 +237,10 @@ def step(
                 ),
             )
 
+    attrs = {
+        "metaflow.task_id": kwargs["task_id"],
+        "requires_passwordless_ssh": any([getattr(deco, "requires_passwordless_ssh", False) for deco in node.decorators]),
+    }
     try:
         kubernetes = Kubernetes(
             datastore=ctx.obj.flow_datastore,
@@ -245,6 +279,11 @@ def step(
                 env=env,
                 persistent_volume_claims=persistent_volume_claims,
                 tolerations=tolerations,
+                labels=labels,
+                annotations=annotations,
+                num_parallel=num_parallel,
+                port=port,
+                attrs=attrs,
             )
     except Exception as e:
         traceback.print_exc(chain=False)

{ob-metaflow-2.10.9.1 → ob-metaflow-2.10.9.2}/metaflow/plugins/kubernetes/kubernetes_decorator.py

@@ -2,6 +2,7 @@ import json
 import os
 import platform
 import sys
+import time
 
 from metaflow import current
 from metaflow.decorators import StepDecorator
@@ -20,10 +21,12 @@ from metaflow.metaflow_config import (
     KUBERNETES_PERSISTENT_VOLUME_CLAIMS,
     KUBERNETES_TOLERATIONS,
     KUBERNETES_SERVICE_ACCOUNT,
+    KUBERNETES_PORT,
 )
 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
@@ -88,6 +91,8 @@ class KubernetesDecorator(StepDecorator):
     persistent_volume_claims: Dict[str, str], optional
         A map (dictionary) of persistent volumes to be mounted to the pod for this step. The map is from persistent
         volumes to the path to which the volume is to be mounted, e.g., `{'pvc-name': '/path/to/mount/on'}`.
+    port: int, optional
+        Number of the port to specify in the Kubernetes job object
     """
 
     name = "kubernetes"
@@ -110,6 +115,7 @@ class KubernetesDecorator(StepDecorator):
         "tmpfs_size": None,
         "tmpfs_path": "/metaflow_temp",
         "persistent_volume_claims": None,  # e.g., {"pvc-name": "/mnt/vol", "another-pvc": "/mnt/vol2"}
+        "port": None,
     }
     package_url = None
     package_sha = None
@@ -195,6 +201,8 @@ class KubernetesDecorator(StepDecorator):
             if not self.attributes["tmpfs_size"]:
                 # default tmpfs behavior - https://man7.org/linux/man-pages/man5/tmpfs.5.html
                 self.attributes["tmpfs_size"] = int(self.attributes["memory"]) // 2
+        if not self.attributes["port"]:
+            self.attributes["port"] = KUBERNETES_PORT
 
     # Refer https://github.com/Netflix/metaflow/blob/master/docs/lifecycle.png
     def step_init(self, flow, graph, step, decos, environment, flow_datastore, logger):
@@ -216,12 +224,6 @@ class KubernetesDecorator(StepDecorator):
                 "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."
-                )
-
         # 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:
@@ -432,6 +434,27 @@ 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)
+
+        if num_parallel >= 1:
+            _setup_multinode_environment()
+
     def task_finished(
         self, step_name, flow, graph, is_task_ok, retry_count, max_retries
     ):
@@ -459,9 +482,53 @@ 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:
             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"]
+    os.environ["MF_PARALLEL_NODE_INDEX"] = os.environ["RANK"]