ob-metaflow-extensions 1.1.90__tar.gz → 1.1.92__tar.gz

{ob-metaflow-extensions-1.1.90 → ob-metaflow-extensions-1.1.92}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ob-metaflow-extensions
-Version: 1.1.90
+Version: 1.1.92
 Summary: Outerbounds Platform Extensions for Metaflow
 Author: Outerbounds, Inc.
 License: Commercial

{ob-metaflow-extensions-1.1.90 → ob-metaflow-extensions-1.1.92}/metaflow_extensions/outerbounds/plugins/fast_bakery/docker_environment.py

@@ -116,15 +116,6 @@ class DockerEnvironment(MetaflowEnvironment):
 
         self.datastore = [d for d in DATASTORES if d.TYPE == self.datastore_type][0]
 
-        # Mixing @pypi/@conda in a single step is not supported yet
-        for step in self.flow:
-            if sum(1 for deco in step.decorators if _is_env_deco(deco)) > 1:
-                raise MetaflowException(
-                    "Mixing and matching PyPI packages and Conda packages within a\n"
-                    "step is not yet supported. Use one of @pypi or @conda only for the *%s* step."
-                    % step.name
-                )
-
     def init_environment(self, echo):
         self.skipped_steps = {
             step.name for step in self.flow if not _step_executes_remotely(step)
@@ -132,35 +123,45 @@ class DockerEnvironment(MetaflowEnvironment):
         # Attach environment decorator as needed. This is done on a step-by-step basis
         # as we require a conda decorator for fallback steps, but prefer pypi for the baked ones.
         for step in self.flow:
-            if not _step_has_environment_deco(step):
-                if step.name in self.skipped_steps:
-                    # Conda fallback requires a conda decorator as the default for a step
-                    decorators._attach_decorators_to_step(step, ["conda"])
-                else:
+            if step.name in self.skipped_steps:
+                # Conda fallback requires a conda decorator as the default for a step
+                decorators._attach_decorators_to_step(step, ["conda"])
+            else:
+                if not _step_has_environment_deco(step):
                     # We default to PyPI for steps that are going to be baked.
                     decorators._attach_decorators_to_step(step, ["pypi"])
-                # Initialize the decorator we attached.
-                # This is crucial for the conda decorator to work properly in the fallback environment
-                for deco in step.decorators:
-                    if _is_env_deco(deco):
-                        deco.step_init(
-                            self.flow,
-                            None,  # not passing graph as it is not available, and not required by conda/pypi decorators
-                            step.name,
-                            step.decorators,
-                            self,
-                            self.datastore,
-                            echo,
-                        )
+            # Initialize the decorator we attached.
+            # This is crucial for the conda decorator to work properly in the fallback environment
+            for deco in step.decorators:
+                if _is_env_deco(deco):
+                    deco.step_init(
+                        self.flow,
+                        None,  # not passing graph as it is not available, and not required by conda/pypi decorators
+                        step.name,
+                        step.decorators,
+                        self,
+                        self.datastore,
+                        echo,
+                    )
+
+            # Mixing @pypi/@conda in a single step is not supported yet
+            if sum(1 for deco in step.decorators if _is_env_deco(deco)) > 1:
+                raise MetaflowException(
+                    "Mixing and matching PyPI packages and Conda packages within a\n"
+                    "step is not yet supported. Use one of @pypi or @conda only for the *%s* step."
+                    % step.name
+                )
 
         steps_to_bake = [
-            step for step in self.flow if step.name not in self.skipped_steps
+            step
+            for step in self.flow
+            if step.name not in self.skipped_steps and not self.is_disabled(step)
         ]
         if steps_to_bake:
             self.logger("🚀 Baking container image(s) ...")
             start_time = time.time()
             self.results = self._bake(steps_to_bake)
-            for step in self.flow:
+            for step in steps_to_bake:
                 for d in step.decorators:
                     if _is_remote_deco(d):
                         d.attributes["image"] = self.results[step.name].container_image

ob-metaflow-extensions-1.1.92/metaflow_extensions/outerbounds/plugins/kubernetes/kubernetes_client.py

@@ -0,0 +1,157 @@
+from concurrent.futures import ThreadPoolExecutor
+import os
+import sys
+import time
+
+from metaflow.exception import MetaflowException
+from metaflow.metaflow_config import KUBERNETES_NAMESPACE
+
+
+CLIENT_REFRESH_INTERVAL_SECONDS = 300
+
+
+class KubernetesClientException(MetaflowException):
+    headline = "Kubernetes client error"
+
+
+class KubernetesClient(object):
+    def __init__(self):
+        try:
+            # Kubernetes is a soft dependency.
+            from kubernetes import client, config
+        except (NameError, ImportError):
+            raise KubernetesClientException(
+                "Could not import module 'kubernetes'.\n\nInstall Kubernetes "
+                "Python package (https://pypi.org/project/kubernetes/) first.\n"
+                "You can install the module by executing - "
+                "%s -m pip install kubernetes\n"
+                "or equivalent through your favorite Python package manager."
+                % sys.executable
+            )
+        self._refresh_client()
+        self._namespace = KUBERNETES_NAMESPACE
+
+    def _refresh_client(self):
+        from metaflow_extensions.outerbounds.plugins.auth_server import get_token
+        from kubernetes import client
+
+        config = client.Configuration()
+        token_info = get_token("/generate/k8s")
+        config.host = token_info["endpoint"]
+        config.api_key["authorization"] = "Bearer " + token_info["token"]
+        config.verify_ssl = False  # TODO: FIX THIS
+        client.Configuration.set_default(config)
+        self._client = client
+        self._client_refresh_timestamp = time.time()
+
+    def get(self):
+        if (
+            time.time() - self._client_refresh_timestamp
+            > CLIENT_REFRESH_INTERVAL_SECONDS
+        ):
+            self._refresh_client()
+
+        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:
+            executor.map(_kill_pod, list(pods))
+
+    def job(self, **kwargs):
+        from metaflow.plugins.kubernetes.kubernetes_job import KubernetesJob
+
+        return KubernetesJob(self, **kwargs)
+
+    def jobset(self, **kwargs):
+        from metaflow.plugins.kubernetes.kubernetes_job import KubernetesJobSet
+
+        return KubernetesJobSet(self, **kwargs)

{ob-metaflow-extensions-1.1.90 → ob-metaflow-extensions-1.1.92}/metaflow_extensions/outerbounds/plugins/nvcf/heartbeat_store.py

@@ -8,26 +8,27 @@ from datetime import datetime, timezone
 from metaflow.exception import MetaflowException
 
 
-def kill_process_and_descendants(pid, termination_timeout=5):
-    try:
-        subprocess.check_call(["pkill", "-TERM", "-P", str(pid)])
-        subprocess.check_call(["kill", "-TERM", str(pid)])
-    except subprocess.CalledProcessError:
-        pass
+def kill_process_and_descendants(pid, termination_timeout=1, iterations=20, delay=0.5):
+    for i in range(iterations):
+        try:
+            subprocess.check_call(["pkill", "-TERM", "-P", str(pid)])
+            subprocess.check_call(["kill", "-TERM", str(pid)])
+
+            time.sleep(termination_timeout)
 
-    time.sleep(termination_timeout)
+            subprocess.check_call(["pkill", "-KILL", "-P", str(pid)])
+            subprocess.check_call(["kill", "-KILL", str(pid)])
+        except subprocess.CalledProcessError:
+            pass
 
-    try:
-        subprocess.check_call(["pkill", "-KILL", "-P", str(pid)])
-        subprocess.check_call(["kill", "-KILL", str(pid)])
-    except subprocess.CalledProcessError:
-        pass
+        # Don't delay after the last iteration
+        if i < iterations - 1:
+            time.sleep(delay)
 
 
 class HeartbeatStore(object):
     def __init__(
         self,
-        heartbeat_prefix,
         main_pid=None,
         storage_backend=None,
         emit_frequency=30,
@@ -35,7 +36,6 @@ class HeartbeatStore(object):
         monitor_frequency=15,
         max_missed_heartbeats=3,
     ) -> None:
-        self.heartbeat_prefix = heartbeat_prefix
         self.main_pid = main_pid
         self.storage_backend = storage_backend
         self.emit_frequency = emit_frequency
@@ -44,8 +44,8 @@ class HeartbeatStore(object):
         self.max_missed_heartbeats = max_missed_heartbeats
         self.missed_heartbeats = 0
 
-    def emit_heartbeat(self, folder_name=None):
-        heartbeat_key = f"{self.heartbeat_prefix}/heartbeat"
+    def emit_heartbeat(self, heartbeat_prefix: str, folder_name=None):
+        heartbeat_key = f"{heartbeat_prefix}/heartbeat"
         if folder_name:
             heartbeat_key = f"{folder_name}/{heartbeat_key}"
 
@@ -63,8 +63,8 @@ class HeartbeatStore(object):
 
             time.sleep(self.emit_frequency)
 
-    def emit_tombstone(self, folder_name=None):
-        tombstone_key = f"{self.heartbeat_prefix}/tombstone"
+    def emit_tombstone(self, tombstone_prefix: str, folder_name=None):
+        tombstone_key = f"{tombstone_prefix}/tombstone"
         if folder_name:
             tombstone_key = f"{folder_name}/{tombstone_key}"
 
@@ -113,12 +113,12 @@ class HeartbeatStore(object):
             return False
         return True
 
-    def monitor(self, folder_name=None):
-        heartbeat_key = f"{self.heartbeat_prefix}/heartbeat"
+    def monitor(self, heartbeat_prefix: str, tombstone_prefix: str, folder_name=None):
+        heartbeat_key = f"{heartbeat_prefix}/heartbeat"
         if folder_name:
             heartbeat_key = f"{folder_name}/{heartbeat_key}"
 
-        tombstone_key = f"{self.heartbeat_prefix}/tombstone"
+        tombstone_key = f"{tombstone_prefix}/tombstone"
         if folder_name:
             tombstone_key = f"{folder_name}/{tombstone_key}"
 
@@ -162,12 +162,17 @@ if __name__ == "__main__":
     storage = datastores[0](datastore_sysroot)
 
     heartbeat_prefix = f"{os.getenv('MF_PATHSPEC')}/{os.getenv('MF_ATTEMPT')}"
+    flow_name, run_id, _, _ = os.getenv("MF_PATHSPEC").split("/")
+    tombstone_prefix = f"{flow_name}/{run_id}"
 
     store = HeartbeatStore(
-        heartbeat_prefix=heartbeat_prefix,
         main_pid=int(main_pid),
         storage_backend=storage,
         max_missed_heartbeats=int(NVIDIA_HEARTBEAT_THRESHOLD),
     )
 
-    store.monitor(folder_name=folder_name)
+    store.monitor(
+        heartbeat_prefix=heartbeat_prefix,
+        tombstone_prefix=tombstone_prefix,
+        folder_name=folder_name,
+    )

{ob-metaflow-extensions-1.1.90 → ob-metaflow-extensions-1.1.92}/metaflow_extensions/outerbounds/plugins/nvcf/nvcf.py

@@ -185,13 +185,17 @@ class Job(object):
         )
 
         store = HeartbeatStore(
-            heartbeat_prefix=heartbeat_prefix,
             main_pid=None,
             storage_backend=backend,
         )
 
         self.heartbeat_thread = threading.Thread(
-            target=store.emit_heartbeat, args=("nvcf_heartbeats",), daemon=True
+            target=store.emit_heartbeat,
+            args=(
+                heartbeat_prefix,
+                "nvcf_heartbeats",
+            ),
+            daemon=True,
         )
         self.heartbeat_thread.start()
 

{ob-metaflow-extensions-1.1.90 → ob-metaflow-extensions-1.1.92}/metaflow_extensions/outerbounds/plugins/nvcf/nvcf_cli.py

@@ -56,23 +56,21 @@ def list(ctx, run_id):
     flow_name = ctx.obj.flow.name
     run_obj = Run(pathspec=f"{flow_name}/{run_id}", _namespace_check=False)
     running_invocations = []
+
     for each_step in run_obj:
-        if (
-            not each_step.task.finished
-            and "nvcf-function-id" in each_step.task.metadata_dict
-        ):
-            task_pathspec = each_step.task.pathspec
-            attempt = each_step.task.metadata_dict.get("attempt")
-            flow_name, run_id, step_name, task_id = task_pathspec.split("/")
-            running_invocations.append(
-                f"Flow Name: {flow_name}, Run ID: {run_id}, Step Name: {step_name}, Task ID: {task_id}, Retry Count: {attempt}"
-            )
+        for each_task in each_step:
+            if not each_task.finished and "nvcf-function-id" in each_task.metadata_dict:
+
+                task_pathspec = each_task.pathspec
+                attempt = each_task.metadata_dict.get("attempt")
+                flow_name, run_id, step_name, task_id = task_pathspec.split("/")
+                running_invocations.append(
+                    f"Flow Name: {flow_name}, Run ID: {run_id}, Step Name: {step_name}, Task ID: {task_id}, Retry Count: {attempt}"
+                )
 
     if running_invocations:
         for each_invocation in running_invocations:
             ctx.obj.echo(each_invocation)
-    else:
-        ctx.obj.echo("No running @nvidia invocations for Run ID: %s" % run_id)
 
 
 @nvidia.command(help="Kill steps / tasks running as an nvidia job.")
@@ -88,29 +86,17 @@ def kill(ctx, run_id):
         HeartbeatStore,
     )
 
-    flow_name = ctx.obj.flow.name
-    run_obj = Run(pathspec=f"{flow_name}/{run_id}", _namespace_check=False)
-
-    for each_step in run_obj:
-        if (
-            not each_step.task.finished
-            and "nvcf-function-id" in each_step.task.metadata_dict
-        ):
-            task_pathspec = each_step.task.pathspec
-            attempt = each_step.task.metadata_dict.get("attempt")
-            heartbeat_prefix = "{task_pathspec}/{attempt}".format(
-                task_pathspec=task_pathspec, attempt=attempt
-            )
+    datastore_root = ctx.obj.datastore_impl.datastore_root
+    store = HeartbeatStore(
+        main_pid=None,
+        storage_backend=ctx.obj.datastore_impl(datastore_root),
+    )
 
-            datastore_root = ctx.obj.datastore_impl.datastore_root
-            store = HeartbeatStore(
-                heartbeat_prefix=heartbeat_prefix,
-                main_pid=None,
-                storage_backend=ctx.obj.datastore_impl(datastore_root),
-            )
-            store.emit_tombstone(folder_name="nvcf_heartbeats")
-        else:
-            ctx.obj.echo("No running @nvidia invocations for Run ID: %s" % run_id)
+    flow_name = ctx.obj.flow.name
+    tombstone_prefix = f"{flow_name}/{run_id}"
+    store.emit_tombstone(
+        tombstone_prefix=tombstone_prefix, folder_name="nvcf_heartbeats"
+    )
 
 
 @nvidia.command(

{ob-metaflow-extensions-1.1.90 → ob-metaflow-extensions-1.1.92}/metaflow_extensions/outerbounds/plugins/snowpark/snowpark.py

@@ -31,7 +31,7 @@ from metaflow.mflog import (
 )
 
 from .snowpark_client import SnowparkClient
-from .snowpark_exceptions import SnowparkException
+from .snowpark_exceptions import SnowparkException, SnowparkKilledException
 from .snowpark_job import SnowparkJob
 
 # Redirect structured logs to $PWD/.logs/
@@ -291,7 +291,7 @@ class Snowpark(object):
         else:
             if self.job.is_running:
                 # Kill the job if it is still running by throwing an exception.
-                raise SnowparkException("Task failed!")
+                raise SnowparkKilledException("Task failed!")
             echo(
                 "Task finished with message '%s'." % self.job.message,
                 "stderr",

{ob-metaflow-extensions-1.1.90 → ob-metaflow-extensions-1.1.92}/ob_metaflow_extensions.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ob-metaflow-extensions
-Version: 1.1.90
+Version: 1.1.92
 Summary: Outerbounds Platform Extensions for Metaflow
 Author: Outerbounds, Inc.
 License: Commercial

ob-metaflow-extensions-1.1.92/ob_metaflow_extensions.egg-info/requires.txt

@@ -0,0 +1,3 @@
+boto3
+kubernetes
+ob-metaflow==2.12.20.1

{ob-metaflow-extensions-1.1.90 → ob-metaflow-extensions-1.1.92}/setup.py

@@ -2,7 +2,7 @@ from setuptools import setup, find_namespace_packages
 from pathlib import Path
 
 
-version = "1.1.90"
+version = "1.1.92"
 this_directory = Path(__file__).parent
 long_description = (this_directory / "README.md").read_text()
 
@@ -18,6 +18,6 @@ setup(
     install_requires=[
         "boto3",
         "kubernetes",
-        "ob-metaflow == 2.12.19.1",
+        "ob-metaflow == 2.12.20.1",
     ],
 )

ob-metaflow-extensions-1.1.90/metaflow_extensions/outerbounds/plugins/kubernetes/kubernetes_client.py

@@ -1,61 +0,0 @@
-import os
-import sys
-import time
-
-from metaflow.exception import MetaflowException
-
-
-CLIENT_REFRESH_INTERVAL_SECONDS = 300
-
-
-class KubernetesClientException(MetaflowException):
-    headline = "Kubernetes client error"
-
-
-class KubernetesClient(object):
-    def __init__(self):
-        try:
-            # Kubernetes is a soft dependency.
-            from kubernetes import client, config
-        except (NameError, ImportError):
-            raise KubernetesClientException(
-                "Could not import module 'kubernetes'.\n\nInstall Kubernetes "
-                "Python package (https://pypi.org/project/kubernetes/) first.\n"
-                "You can install the module by executing - "
-                "%s -m pip install kubernetes\n"
-                "or equivalent through your favorite Python package manager."
-                % sys.executable
-            )
-        self._refresh_client()
-
-    def _refresh_client(self):
-        from metaflow_extensions.outerbounds.plugins.auth_server import get_token
-        from kubernetes import client
-
-        config = client.Configuration()
-        token_info = get_token("/generate/k8s")
-        config.host = token_info["endpoint"]
-        config.api_key["authorization"] = "Bearer " + token_info["token"]
-        config.verify_ssl = False  # TODO: FIX THIS
-        client.Configuration.set_default(config)
-        self._client = client
-        self._client_refresh_timestamp = time.time()
-
-    def get(self):
-        if (
-            time.time() - self._client_refresh_timestamp
-            > CLIENT_REFRESH_INTERVAL_SECONDS
-        ):
-            self._refresh_client()
-
-        return self._client
-
-    def job(self, **kwargs):
-        from metaflow.plugins.kubernetes.kubernetes_job import KubernetesJob
-
-        return KubernetesJob(self, **kwargs)
-
-    def jobset(self, **kwargs):
-        from metaflow.plugins.kubernetes.kubernetes_job import KubernetesJobSet
-
-        return KubernetesJobSet(self, **kwargs)

ob-metaflow-extensions-1.1.90/ob_metaflow_extensions.egg-info/requires.txt

@@ -1,3 +0,0 @@
-boto3
-kubernetes
-ob-metaflow==2.12.19.1