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

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

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

{ob-metaflow-extensions-1.1.91 → 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.91 → 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.91
+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.91 → 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.91"
+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.91/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.91/ob_metaflow_extensions.egg-info/requires.txt

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