ob-metaflow-extensions 1.1.91__py2.py3-none-any.whl → 1.1.93__py2.py3-none-any.whl

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:
+            # Mixing @pypi/@conda in a single step is not supported yet.
+            # We validate this before attaching any new ones as the OSS Conda environment requires an implicit conda decorator for pypi environments which would fail the validation.
+            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
+                )
+            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,
+                    )
 
         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

metaflow_extensions/outerbounds/plugins/kubernetes/kubernetes_client.py

@@ -1,8 +1,10 @@
+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
@@ -27,6 +29,7 @@ class KubernetesClient(object):
                 % sys.executable
             )
         self._refresh_client()
+        self._namespace = KUBERNETES_NAMESPACE
 
     def _refresh_client(self):
         from metaflow_extensions.outerbounds.plugins.auth_server import get_token
@@ -50,6 +53,99 @@ 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:
+            executor.map(_kill_pod, list(pods))
+
     def job(self, **kwargs):
         from metaflow.plugins.kubernetes.kubernetes_job import KubernetesJob
 

{ob_metaflow_extensions-1.1.91.dist-info → ob_metaflow_extensions-1.1.93.dist-info}/METADATA

@@ -1,13 +1,13 @@
 Metadata-Version: 2.1
 Name: ob-metaflow-extensions
-Version: 1.1.91
+Version: 1.1.93
 Summary: Outerbounds Platform Extensions for Metaflow
 Author: Outerbounds, Inc.
 License: Commercial
 Description-Content-Type: text/markdown
 Requires-Dist: boto3
 Requires-Dist: kubernetes
-Requires-Dist: ob-metaflow (==2.12.19.1)
+Requires-Dist: ob-metaflow (==2.12.20.1)
 
 # Outerbounds platform package
 

{ob_metaflow_extensions-1.1.91.dist-info → ob_metaflow_extensions-1.1.93.dist-info}/RECORD

@@ -5,12 +5,12 @@ metaflow_extensions/outerbounds/plugins/__init__.py,sha256=WwvFcN5kserbPwhrE4hXp
 metaflow_extensions/outerbounds/plugins/auth_server.py,sha256=1v2GBqoMBxp5E7Lejz139w-jxJtPnLDvvHXP0HhEIHI,2361
 metaflow_extensions/outerbounds/plugins/perimeters.py,sha256=QXh3SFP7GQbS-RAIxUOPbhPzQ7KDFVxZkTdKqFKgXjI,2697
 metaflow_extensions/outerbounds/plugins/fast_bakery/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-metaflow_extensions/outerbounds/plugins/fast_bakery/docker_environment.py,sha256=gFeXt8nLv7SzZ2zEoAKMpdxAePN4jZclSsmQAm6hB4w,13846
+metaflow_extensions/outerbounds/plugins/fast_bakery/docker_environment.py,sha256=jpEBbdRLq7826LqUKIHqo3cYei2CNeVAxggw6vFmYso,13996
 metaflow_extensions/outerbounds/plugins/fast_bakery/fast_bakery.py,sha256=MAPRQsfqeEkL1LXqgwPrUJOzZ3kY3C00QjdDgQ7wdIg,5160
 metaflow_extensions/outerbounds/plugins/fast_bakery/fast_bakery_cli.py,sha256=kqFyu2bJSnc9_9aYfBpz5xK6L6luWFZK_NMuh8f1eVk,1494
 metaflow_extensions/outerbounds/plugins/fast_bakery/fast_bakery_decorator.py,sha256=EZDbyrfZ7fgcU-P9dMS_hpCxsdDeUE0K5VU3uNM4aW4,1506
 metaflow_extensions/outerbounds/plugins/kubernetes/__init__.py,sha256=5zG8gShSj8m7rgF4xgWBZFuY3GDP5n1T0ktjRpGJLHA,69
-metaflow_extensions/outerbounds/plugins/kubernetes/kubernetes_client.py,sha256=gj6Iaz26bGbZm3aQuNS18Mqh_80iJp5PgFwFSlJRcn8,1968
+metaflow_extensions/outerbounds/plugins/kubernetes/kubernetes_client.py,sha256=Vbchj4D9HbMBQy7Te2QbFnQqEkfABeGxrWrvBWy5TTY,5858
 metaflow_extensions/outerbounds/plugins/nim/__init__.py,sha256=GVnvSTjqYVj5oG2yh8KJFt7iZ33cEadDD5HbdmC9hJ0,1457
 metaflow_extensions/outerbounds/plugins/nim/nim_manager.py,sha256=SWieODDxtIaeZwdMYtObDi57Kjyfw2DUuE6pJtU750w,9206
 metaflow_extensions/outerbounds/plugins/nvcf/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -33,7 +33,7 @@ metaflow_extensions/outerbounds/toplevel/global_aliases_for_metaflow_package.py,
 metaflow_extensions/outerbounds/toplevel/plugins/azure/__init__.py,sha256=WUuhz2YQfI4fz7nIcipwwWq781eaoHEk7n4GAn1npDg,63
 metaflow_extensions/outerbounds/toplevel/plugins/gcp/__init__.py,sha256=BbZiaH3uILlEZ6ntBLKeNyqn3If8nIXZFq_Apd7Dhco,70
 metaflow_extensions/outerbounds/toplevel/plugins/kubernetes/__init__.py,sha256=5zG8gShSj8m7rgF4xgWBZFuY3GDP5n1T0ktjRpGJLHA,69
-ob_metaflow_extensions-1.1.91.dist-info/METADATA,sha256=rTZDEBQQ5bWEZnqLoralSTjginYiJ-Q81f054qg3vLs,520
-ob_metaflow_extensions-1.1.91.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
-ob_metaflow_extensions-1.1.91.dist-info/top_level.txt,sha256=NwG0ukwjygtanDETyp_BUdtYtqIA_lOjzFFh1TsnxvI,20
-ob_metaflow_extensions-1.1.91.dist-info/RECORD,,
+ob_metaflow_extensions-1.1.93.dist-info/METADATA,sha256=yWB9qUR6LTg61tnxvkMqcdGw3SRZjcSnSIhn6J6u5aQ,520
+ob_metaflow_extensions-1.1.93.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
+ob_metaflow_extensions-1.1.93.dist-info/top_level.txt,sha256=NwG0ukwjygtanDETyp_BUdtYtqIA_lOjzFFh1TsnxvI,20
+ob_metaflow_extensions-1.1.93.dist-info/RECORD,,