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

metaflow/plugins/kubernetes/kubernetes_cli.py

@@ -3,16 +3,25 @@ import sys
 import time
 import traceback
 
+from metaflow.plugins.kubernetes.kube_utils import (
+    parse_cli_options,
+    parse_kube_keyvalue_list,
+)
+from metaflow.plugins.kubernetes.kubernetes_client import KubernetesClient
+import metaflow.tracing as tracing
 from metaflow import JSONTypeClass, util
 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.metaflow_config import DATASTORE_LOCAL_DIR, KUBERNETES_LABELS
+from metaflow.exception import METAFLOW_EXIT_DISALLOW_RETRY, MetaflowException
+from metaflow.metadata_provider.util import sync_local_metadata_from_datastore
+from metaflow.metaflow_config import DATASTORE_LOCAL_DIR
 from metaflow.mflog import TASK_LOG_SOURCE
-import metaflow.tracing as tracing
+from metaflow.unbounded_foreach import UBF_CONTROL, UBF_TASK
 
-from .kubernetes import Kubernetes, KubernetesKilledException, parse_kube_keyvalue_list
-from .kubernetes_decorator import KubernetesDecorator
+from .kubernetes import (
+    Kubernetes,
+    KubernetesException,
+    KubernetesKilledException,
+)
 
 
 @click.group()
@@ -25,13 +34,14 @@ def kubernetes():
     pass
 
 
-@tracing.cli_entrypoint("kubernetes/step")
 @kubernetes.command(
     help="Execute a single task on Kubernetes. This command calls the top-level step "
     "command inside a Kubernetes pod with the given options. Typically you do not call "
     "this command directly; it is used internally by Metaflow."
 )
+@tracing.cli("kubernetes/step")
 @click.argument("step-name")
+@click.argument("code-package-metadata")
 @click.argument("code-package-sha")
 @click.argument("code-package-url")
 @click.option(
@@ -44,6 +54,12 @@ def kubernetes():
     default=None,
     help="Optional Docker Image Pull Policy for Kubernetes pod.",
 )
+@click.option(
+    "--image-pull-secrets",
+    default=None,
+    type=JSONTypeClass(),
+    multiple=False,
+)
 @click.option(
     "--service-account",
     help="IRSA requirement for Kubernetes pod.",
@@ -107,6 +123,23 @@ def kubernetes():
     type=JSONTypeClass(),
     multiple=False,
 )
+@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.option(
+    "--qos",
+    default=None,
+    type=str,
+    help="Quality of Service class for the Kubernetes pod",
+)
 @click.option(
     "--labels",
     default=None,
@@ -119,24 +152,23 @@ def kubernetes():
     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.",
+    "--security-context",
+    default=None,
+    type=JSONTypeClass(),
+    multiple=False,
 )
-@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.pass_context
 def step(
     ctx,
     step_name,
+    code_package_metadata,
     code_package_sha,
     code_package_url,
     executable=None,
     image=None,
     image_pull_policy=None,
+    image_pull_secrets=None,
     service_account=None,
     secrets=None,
     node_selector=None,
@@ -153,11 +185,13 @@ def step(
     run_time_limit=None,
     persistent_volume_claims=None,
     tolerations=None,
-    labels=None,
-    annotations=None,
-    num_parallel=None,
     shared_memory=None,
     port=None,
+    num_parallel=None,
+    qos=None,
+    labels=None,
+    annotations=None,
+    security_context=None,
     **kwargs
 ):
     def echo(msg, stream="stderr", job_id=None, **kwargs):
@@ -172,7 +206,7 @@ def step(
     executable = ctx.obj.environment.executable(step_name, executable)
 
     # Set environment
-    env = {}
+    env = {"METAFLOW_FLOW_FILENAME": os.path.basename(sys.argv[0])}
     env_deco = [deco for deco in node.decorators if deco.name == "environment"]
     if env_deco:
         env = env_deco[0].attributes["vars"]
@@ -189,6 +223,12 @@ def step(
         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"]
@@ -203,25 +243,37 @@ 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]"
+    # Explicitly Remove `ubf_context` from `kwargs` so that it's not passed as a commandline option
+    # If an underlying step command is executing a vanilla Kubernetes job, then it should never need
+    # to know about the UBF context.
+    # If it is a jobset which is executing a multi-node job, then the UBF context is set based on the
+    # `ubf_context` parameter passed to the jobset.
+    kwargs.pop("ubf_context", None)
+    # `task_id` is also need to be removed from `kwargs` as it needs to be dynamically
+    # set in the downstream code IF num_parallel is > 1
+    task_id = kwargs["task_id"]
+    if num_parallel:
+        kwargs.pop("task_id")
 
     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=step_args,
+        step_args=" ".join(util.dict_to_cli_options(kwargs)),
     )
+    # Since it is a parallel step there are some parts of the step_cli that need to be modified
+    # based on the type of worker in the JobSet. This is why we will create a placeholder string
+    # in the template which will be replaced based on the type of worker.
+
+    if num_parallel:
+        step_cli = "%s {METAFLOW_PARALLEL_STEP_CLI_OPTIONS_TEMPLATE}" % step_cli
 
     # Set log tailing.
     ds = ctx.obj.flow_datastore.get_task_datastore(
         mode="w",
         run_id=kwargs["run_id"],
         step_name=step_name,
-        task_id=kwargs["task_id"],
+        task_id=task_id,
         attempt=int(retry_count),
     )
     stdout_location = ds.get_log_location(TASK_LOG_SOURCE, "stdout")
@@ -235,14 +287,10 @@ def step(
             sync_local_metadata_from_datastore(
                 DATASTORE_LOCAL_DIR,
                 ctx.obj.flow_datastore.get_task_datastore(
-                    kwargs["run_id"], step_name, kwargs["task_id"]
+                    kwargs["run_id"], step_name, task_id
                 ),
             )
 
-    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,
@@ -255,15 +303,17 @@ def step(
                 flow_name=ctx.obj.flow.name,
                 run_id=kwargs["run_id"],
                 step_name=step_name,
-                task_id=kwargs["task_id"],
+                task_id=task_id,
                 attempt=str(retry_count),
                 user=util.get_username(),
+                code_package_metadata=code_package_metadata,
                 code_package_sha=code_package_sha,
                 code_package_url=code_package_url,
                 code_package_ds=ctx.obj.flow_datastore.TYPE,
                 step_cli=step_cli,
                 docker_image=image,
                 docker_image_pull_policy=image_pull_policy,
+                image_pull_secrets=image_pull_secrets,
                 service_account=service_account,
                 secrets=secrets,
                 node_selector=node_selector,
@@ -281,14 +331,15 @@ def step(
                 env=env,
                 persistent_volume_claims=persistent_volume_claims,
                 tolerations=tolerations,
-                labels=labels,
-                annotations=annotations,
-                num_parallel=num_parallel,
                 shared_memory=shared_memory,
                 port=port,
-                attrs=attrs,
+                num_parallel=num_parallel,
+                qos=qos,
+                labels=labels,
+                annotations=annotations,
+                security_context=security_context,
             )
-    except Exception as e:
+    except Exception:
         traceback.print_exc(chain=False)
         _sync_metadata()
         sys.exit(METAFLOW_EXIT_DISALLOW_RETRY)
@@ -300,3 +351,84 @@ def step(
         sys.exit(METAFLOW_EXIT_DISALLOW_RETRY)
     finally:
         _sync_metadata()
+
+
+@kubernetes.command(help="List unfinished Kubernetes tasks of this flow.")
+@click.option(
+    "--my-runs",
+    default=False,
+    is_flag=True,
+    help="List all my unfinished tasks.",
+)
+@click.option("--user", default=None, help="List unfinished tasks for the given user.")
+@click.option(
+    "--run-id",
+    default=None,
+    help="List unfinished tasks corresponding to the run id.",
+)
+@click.pass_obj
+def list(obj, run_id, user, my_runs):
+    flow_name, run_id, user = parse_cli_options(
+        obj.flow.name, run_id, user, my_runs, obj.echo
+    )
+    kube_client = KubernetesClient()
+    pods = kube_client.list(obj.flow.name, run_id, user)
+
+    def format_timestamp(timestamp=None):
+        if timestamp is None:
+            return "-"
+        return timestamp.strftime("%Y-%m-%d %H:%M:%S")
+
+    for pod in pods:
+        obj.echo(
+            "Run: *{run_id}* "
+            "Pod: *{pod_id}* "
+            "Started At: {startedAt} "
+            "Status: *{status}*".format(
+                run_id=pod.metadata.annotations.get(
+                    "metaflow/run_id",
+                    pod.metadata.labels.get("workflows.argoproj.io/workflow"),
+                ),
+                pod_id=pod.metadata.name,
+                startedAt=format_timestamp(pod.status.start_time),
+                status=pod.status.phase,
+            )
+        )
+
+    if not pods:
+        obj.echo("No active Kubernetes pods found.")
+
+
+@kubernetes.command(
+    help="Terminate unfinished Kubernetes tasks of this flow. Killed pods may result in newer attempts when using @retry."
+)
+@click.option(
+    "--my-runs",
+    default=False,
+    is_flag=True,
+    help="Kill all my unfinished tasks.",
+)
+@click.option(
+    "--user",
+    default=None,
+    help="Terminate unfinished tasks for the given user.",
+)
+@click.option(
+    "--run-id",
+    default=None,
+    help="Terminate unfinished tasks corresponding to the run id.",
+)
+@click.pass_obj
+def kill(obj, run_id, user, my_runs):
+    flow_name, run_id, user = parse_cli_options(
+        obj.flow.name, run_id, user, my_runs, obj.echo
+    )
+
+    if run_id is not None and run_id.startswith("argo-") or user == "argo-workflows":
+        raise MetaflowException(
+            "Killing pods launched by Argo Workflows is not supported. "
+            "Use *argo-workflows terminate* instead."
+        )
+
+    kube_client = KubernetesClient()
+    kube_client.kill_pods(flow_name, run_id, user, obj.echo)

metaflow/plugins/kubernetes/kubernetes_client.py

@@ -1,11 +1,12 @@
+from concurrent.futures import ThreadPoolExecutor
 import os
 import sys
 import time
 
 from metaflow.exception import MetaflowException
+from metaflow.metaflow_config import KUBERNETES_NAMESPACE
 
-from .kubernetes_job import KubernetesJob
-
+from .kubernetes_job import KubernetesJob, KubernetesJobSet
 
 CLIENT_REFRESH_INTERVAL_SECONDS = 300
 
@@ -29,6 +30,7 @@ class KubernetesClient(object):
                 % sys.executable
             )
         self._refresh_client()
+        self._namespace = KUBERNETES_NAMESPACE
 
     def _refresh_client(self):
         from kubernetes import client, config
@@ -61,5 +63,105 @@ class KubernetesClient(object):
 
         return self._client
 
+    def _find_active_pods(self, flow_name, run_id=None, user=None):
+        def _request(_continue=None):
+            # handle paginated responses
+            return self._client.CoreV1Api().list_namespaced_pod(
+                namespace=self._namespace,
+                # limited selector support for K8S api. We want to cover multiple statuses: Running / Pending / Unknown
+                field_selector="status.phase!=Succeeded,status.phase!=Failed",
+                limit=1000,
+                _continue=_continue,
+            )
+
+        results = _request()
+
+        if run_id is not None:
+            # handle argo prefixes in run_id
+            run_id = run_id[run_id.startswith("argo-") and len("argo-") :]
+
+        while results.metadata._continue or results.items:
+            for pod in results.items:
+                match = (
+                    # arbitrary pods might have no annotations at all.
+                    pod.metadata.annotations
+                    and pod.metadata.labels
+                    and (
+                        run_id is None
+                        or (pod.metadata.annotations.get("metaflow/run_id") == run_id)
+                        # we want to also match pods launched by argo-workflows
+                        or (
+                            pod.metadata.labels.get("workflows.argoproj.io/workflow")
+                            == run_id
+                        )
+                    )
+                    and (
+                        user is None
+                        or pod.metadata.annotations.get("metaflow/user") == user
+                    )
+                    and (
+                        pod.metadata.annotations.get("metaflow/flow_name") == flow_name
+                    )
+                )
+                if match:
+                    yield pod
+            if not results.metadata._continue:
+                break
+            results = _request(results.metadata._continue)
+
+    def list(self, flow_name, run_id, user):
+        results = self._find_active_pods(flow_name, run_id, user)
+
+        return list(results)
+
+    def kill_pods(self, flow_name, run_id, user, echo):
+        from kubernetes.stream import stream
+
+        api_instance = self._client.CoreV1Api()
+        job_api = self._client.BatchV1Api()
+        pods = self._find_active_pods(flow_name, run_id, user)
+
+        def _kill_pod(pod):
+            echo("Killing Kubernetes pod %s\n" % pod.metadata.name)
+            try:
+                stream(
+                    api_instance.connect_get_namespaced_pod_exec,
+                    name=pod.metadata.name,
+                    namespace=pod.metadata.namespace,
+                    command=[
+                        "/bin/sh",
+                        "-c",
+                        "/sbin/killall5",
+                    ],
+                    stderr=True,
+                    stdin=False,
+                    stdout=True,
+                    tty=False,
+                )
+            except Exception:
+                # best effort kill for pod can fail.
+                try:
+                    job_name = pod.metadata.labels.get("job-name", None)
+                    if job_name is None:
+                        raise Exception("Could not determine job name")
+
+                    job_api.patch_namespaced_job(
+                        name=job_name,
+                        namespace=pod.metadata.namespace,
+                        field_manager="metaflow",
+                        body={"spec": {"parallelism": 0}},
+                    )
+                except Exception as e:
+                    echo("failed to kill pod %s - %s" % (pod.metadata.name, str(e)))
+
+        with ThreadPoolExecutor() as executor:
+            operated_pods = list(executor.map(_kill_pod, pods))
+
+            if not operated_pods:
+                echo("No active Kubernetes pods found for run *%s*" % run_id)
+
+    def jobset(self, **kwargs):
+        return KubernetesJobSet(self, **kwargs)
+
     def job(self, **kwargs):
         return KubernetesJob(self, **kwargs)