ob-metaflow-extensions 1.1.151__py2.py3-none-any.whl → 1.4.33__py2.py3-none-any.whl

metaflow_extensions/outerbounds/plugins/fast_bakery/baker.py

@@ -0,0 +1,110 @@
+import threading
+import time
+import sys
+from typing import Dict, Optional, Any, Callable
+from functools import partial
+from metaflow.exception import MetaflowException
+from metaflow.metaflow_config import FAST_BAKERY_URL
+
+from .fast_bakery import FastBakery, FastBakeryApiResponse, FastBakeryException
+from .docker_environment import cache_request
+
+BAKERY_METAFILE = ".imagebakery-cache"
+
+
+class BakerException(MetaflowException):
+    headline = "Ran into an error while baking image"
+
+    def __init__(self, msg):
+        super(BakerException, self).__init__(msg)
+
+
+def bake_image(
+    cache_file_path: str,
+    ref: Optional[str] = None,
+    python: Optional[str] = None,
+    pypi_packages: Optional[Dict[str, str]] = None,
+    conda_packages: Optional[Dict[str, str]] = None,
+    base_image: Optional[str] = None,
+    logger: Optional[Callable[[str], Any]] = None,
+) -> FastBakeryApiResponse:
+    """
+    Bakes a Docker image with the specified dependencies.
+
+    Args:
+        cache_file_path: Path to the cache file
+        ref: Reference identifier for this bake (for logging purposes)
+        python: Python version to use
+        pypi_packages: Dictionary of PyPI packages and versions
+        conda_packages: Dictionary of Conda packages and versions
+        base_image: Base Docker image to use
+        logger: Optional logger function to output progress
+
+    Returns:
+        FastBakeryApiResponse: The response from the bakery service
+
+    Raises:
+        BakerException: If the baking process fails
+    """
+    # Default logger if none provided
+    if logger is None:
+        logger = partial(print, file=sys.stderr)
+
+    # Thread lock for logging
+    logger_lock = threading.Lock()
+    images_baked = 0
+
+    @cache_request(cache_file_path)
+    def _cached_bake(
+        ref=None,
+        python=None,
+        pypi_packages=None,
+        conda_packages=None,
+        base_image=None,
+    ):
+        try:
+            bakery = FastBakery(url=FAST_BAKERY_URL)
+            bakery._reset_payload()
+            bakery.python_version(python)
+            bakery.pypi_packages(pypi_packages)
+            bakery.conda_packages(conda_packages)
+            bakery.base_image(base_image)
+            # bakery.ignore_cache()
+
+            with logger_lock:
+                logger(f"🍳 Baking [{ref}] ...")
+                logger(f"     🐍 Python: {python}")
+
+                if pypi_packages:
+                    logger(f"     📦 PyPI packages:")
+                    for package, version in pypi_packages.items():
+                        logger(f"        🔧 {package}: {version}")
+
+                if conda_packages:
+                    logger(f"     📦 Conda packages:")
+                    for package, version in conda_packages.items():
+                        logger(f"        🔧 {package}: {version}")
+
+                logger(f"     🏗️  Base image: {base_image}")
+
+            start_time = time.time()
+            res = bakery.bake()
+            # TODO: Get actual bake time from bakery
+            bake_time = time.time() - start_time
+
+            with logger_lock:
+                logger(f"🏁 Baked [{ref}] in {bake_time:.2f} seconds!")
+            nonlocal images_baked
+            images_baked += 1
+            return res
+        except FastBakeryException as ex:
+            raise BakerException(f"Bake [{ref}] failed: {str(ex)}")
+
+    # Call the cached bake function with the provided parameters
+    return _cached_bake(
+        ref=ref,
+        python=python,
+        pypi_packages=pypi_packages,
+        conda_packages=conda_packages,
+        base_image=base_image,
+    )

metaflow_extensions/outerbounds/plugins/fast_bakery/docker_environment.py

@@ -90,6 +90,7 @@ class DockerEnvironmentException(MetaflowException):
 class DockerEnvironment(MetaflowEnvironment):
     TYPE = "fast-bakery"
     _filecache = None
+    _force_rebuild = False
 
     def __init__(self, flow):
         self.skipped_steps = set()
@@ -178,12 +179,20 @@ class DockerEnvironment(MetaflowEnvironment):
 
         if self.skipped_steps:
             self.delegate = CondaEnvironment(self.flow)
+            self.delegate._force_rebuild = self._force_rebuild
             self.delegate.set_local_root(self.local_root)
             self.delegate.validate_environment(echo, self.datastore_type)
             self.delegate.init_environment(echo, self.skipped_steps)
 
     def _bake(self, steps) -> Dict[str, FastBakeryApiResponse]:
         metafile_path = get_fastbakery_metafile_path(self.local_root, self.flow.name)
+        if self._force_rebuild:
+            # clear the metafile if force rebuilding, effectively skipping the cache.
+            try:
+                os.remove(metafile_path)
+            except Exception:
+                pass
+
         logger_lock = threading.Lock()
 
         @cache_request(metafile_path)
@@ -201,7 +210,8 @@ class DockerEnvironment(MetaflowEnvironment):
                 bakery.pypi_packages(pypi_packages)
                 bakery.conda_packages(conda_packages)
                 bakery.base_image(base_image)
-                # bakery.ignore_cache()
+                if self._force_rebuild:
+                    bakery.ignore_cache()
 
                 with logger_lock:
                     self.logger(f"🍳 Baking [{ref}] ...")
@@ -341,12 +351,16 @@ class DockerEnvironment(MetaflowEnvironment):
         config.append("--disable=F0401")
         return config
 
-    def get_package_commands(self, codepackage_url, datastore_type):
+    def get_package_commands(
+        self, codepackage_url, datastore_type, code_package_metadata=None
+    ):
         # we must set the skip install flag at this stage in order to skip package downloads,
         # doing so in bootstrap_commands is too late in the lifecycle.
         return [
             "export METAFLOW_SKIP_INSTALL_DEPENDENCIES=$FASTBAKERY_IMAGE",
-        ] + super().get_package_commands(codepackage_url, datastore_type)
+        ] + super().get_package_commands(
+            codepackage_url, datastore_type, code_package_metadata=code_package_metadata
+        )
 
     def bootstrap_commands(self, step_name, datastore_type):
         if step_name in self.skipped_steps:

metaflow_extensions/outerbounds/plugins/fast_bakery/fast_bakery.py

@@ -122,6 +122,7 @@ class FastBakery:
             "responseMaxAgeSeconds": 0,
             "layerMaxAgeSeconds": 0,
             "baseImageMaxAgeSeconds": 0,
+            "overwriteExistingLayers": True,  # Used primarily to rewrite possibly corrupted layers.
         }
         return self
 

metaflow_extensions/outerbounds/plugins/kubernetes/kubernetes_client.py

@@ -5,6 +5,7 @@ import time
 
 from metaflow.exception import MetaflowException
 from metaflow.metaflow_config import KUBERNETES_NAMESPACE
+from .pod_killer import PodKiller
 
 
 CLIENT_REFRESH_INTERVAL_SECONDS = 300
@@ -105,50 +106,23 @@ class KubernetesClient(object):
         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)
+        # Create PodKiller instance
+        killer = PodKiller(self._client, echo, self._namespace)
+
+        # Process all matching jobs and jobsets based on their outcomes
+        (
+            job_jobset_results,
+            num_jobs,
+            num_jobsets,
+        ) = killer.process_matching_jobs_and_jobsets(flow_name, run_id, user)
+
+        if job_jobset_results:
+            successful_operations = sum(1 for result in job_jobset_results if result)
+            echo(
+                f"Found and processed {num_jobs} jobs and {num_jobsets} jobsets, {successful_operations} operations successful\n"
+            )
+        else:
+            echo("No matching jobs or jobsets found for run *%s*" % run_id)
 
     def job(self, **kwargs):
         from metaflow.plugins.kubernetes.kubernetes_job import KubernetesJob

metaflow_extensions/outerbounds/plugins/kubernetes/pod_killer.py

@@ -0,0 +1,374 @@
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from kubernetes.client.models.v1_job import V1Job
+    from kubernetes.client.models.v1_job_status import V1JobStatus
+
+
+def _is_jobset_child(job: "V1Job"):
+    if job.metadata.owner_references:
+        for owner_ref in job.metadata.owner_references:
+            if owner_ref.kind == "JobSet":
+                return owner_ref
+    return None
+
+
+class JobOutcomes:
+    KILL = "kill"
+    DELETE = "delete"
+    LEAVE_UNCHANGED = "leave_unchanged"
+
+
+def derive_jobset_outcome(jobset_status):
+    return (
+        JobOutcomes.LEAVE_UNCHANGED
+        if jobset_status.get("terminalState", None)
+        else JobOutcomes.DELETE
+    )
+
+
+def derive_job_outcome(job_status: "V1JobStatus"):
+    if job_status.start_time is None:
+        # If the job has not started even then just wipe it!
+        return JobOutcomes.DELETE
+    if job_status.succeeded or job_status.failed:
+        return JobOutcomes.LEAVE_UNCHANGED
+
+    if job_status.completion_time is not None:
+        return JobOutcomes.LEAVE_UNCHANGED
+
+    # This means that the job has neither finished or succedded.
+    if job_status.active:
+        return JobOutcomes.DELETE
+
+    # This means that the job is not active. Had started. There is not succedded/fail.
+    # This is a weird state. Better to just kill the job
+    return JobOutcomes.DELETE
+
+
+class PodKiller:
+    def __init__(self, kubernetes_client, echo_func, namespace, progress_bar=None):
+        self.client = kubernetes_client
+        self.echo = echo_func
+        self.api_instance = self.client.CoreV1Api()
+        self.job_api = self.client.BatchV1Api()
+        self._namespace = namespace
+        self.jobset_api = None
+        self.jobset_api = self.client.CustomObjectsApi()
+        self.progress_bar = progress_bar
+
+    def _delete_jobset(self, owner_ref, namespace):
+        """Delete a JobSet if it's the owner of a job."""
+        if not self.jobset_api:
+            self.echo("JobSet API not available, cannot delete JobSet\n")
+            return False
+
+        try:
+            jobset_name = owner_ref.name
+            self.echo(f"Deleting JobSet: {jobset_name}\n")
+
+            self.jobset_api.delete_namespaced_custom_object(
+                group="jobset.x-k8s.io",
+                version="v1alpha2",
+                namespace=namespace,
+                plural="jobsets",
+                name=jobset_name,
+            )
+            return True
+        except Exception as e:
+            self.echo(f"Failed to delete JobSet {owner_ref.name}: {str(e)}\n")
+            return False
+
+    def _delete_job(self, job_name, namespace):
+        """Delete a Batch Job and check for JobSet owner reference."""
+        try:
+            # First get the job to check for owner references
+            job = self.job_api.read_namespaced_job(name=job_name, namespace=namespace)
+            # Check for JobSet owner reference
+            jobset_ref = _is_jobset_child(job)
+            if jobset_ref:
+                if self._delete_jobset(jobset_ref, namespace):
+                    return True
+
+            # If no JobSet owner or JobSet deletion failed, delete the job
+            self.echo(f"Deleting Batch Job: {job_name}")
+            self.job_api.delete_namespaced_job(
+                name=job_name, namespace=namespace, propagation_policy="Background"
+            )
+            return True
+
+        except Exception as e:
+            self.echo(f"Failed to delete job {job_name}: {str(e)}")
+            return False
+
+    def _kill_pod_process(self, pod):
+        """Attempt to kill processes inside a pod."""
+        from kubernetes.stream import stream
+
+        try:
+            stream(
+                self.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,
+            )
+            return True
+        except Exception as e:
+            self.echo(
+                f"Failed to kill processes in pod {pod.metadata.name}: {str(e)}\n"
+            )
+            return False
+
+    @staticmethod
+    def _metaflow_matching_spec(run_id, user, flow_name, annotations, labels):
+        # Handle argo prefixes in run_id like in _find_active_pods
+        _argo_run_id = None
+        if run_id is not None:
+            _argo_run_id = run_id[run_id.startswith("argo-") and len("argo-") :]
+        return (
+            annotations
+            and (
+                run_id is None
+                or (annotations.get("metaflow/run_id") == run_id)
+                # we want to also match jobsets launched by argo-workflows
+                # This line has no real value since the We already avoid any
+                # argo-workflows related terminations.
+                or (
+                    labels.get("workflows.argoproj.io/workflow") is not None
+                    and labels.get("workflows.argoproj.io/workflow") == _argo_run_id
+                )
+            )
+            and (user is None or annotations.get("metaflow/user") == user)
+            and (annotations.get("metaflow/flow_name") == flow_name)
+        )
+
+    def _find_matching_jobs(self, flow_name, run_id=None, user=None):
+        """Find jobs that match the flow_name, run_id, and user criteria using similar logic to _find_active_pods"""
+
+        def paginated_job_finder(namespace):
+            continue_token = None
+            while True:
+                response = self.job_api.list_namespaced_job(
+                    namespace=namespace, limit=100, _continue=continue_token
+                )
+                yield response.items
+                continue_token = response.metadata._continue
+                if not continue_token:
+                    break
+
+        try:
+            matching_jobs = []
+            for _jobs in paginated_job_finder(self._namespace):
+                for job in _jobs:
+                    _match = self._metaflow_matching_spec(
+                        run_id=run_id,
+                        user=user,
+                        flow_name=flow_name,
+                        annotations=job.metadata.annotations,
+                        labels=job.metadata.labels,
+                    )
+                    if _match:
+                        matching_jobs.append(job)
+            return matching_jobs
+        except Exception as e:
+            self.echo(f"Error finding jobs: {str(e)}\n")
+            return []
+
+    def _find_matching_jobsets(self, flow_name, run_id=None, user=None):
+        """Find jobsets that match the flow_name, run_id, and user criteria using similar logic to _find_active_pods"""
+        if not self.jobset_api:
+            return []
+
+        def paginated_jobset_finder(namespace):
+            continue_token = None
+            responses = []
+            while True:
+                response = self.jobset_api.list_namespaced_custom_object(
+                    group="jobset.x-k8s.io",
+                    version="v1alpha2",
+                    namespace=namespace,
+                    plural="jobsets",
+                    limit=100,
+                    **({"_continue": continue_token} if continue_token else {}),
+                )
+                continue_token = response.get("metadata", {}).get("continue", None)
+                responses.append(response)
+                if not continue_token:
+                    break
+            return responses
+
+        try:
+            matching_jobsets = []
+
+            for jobset_response in paginated_jobset_finder(self._namespace):
+                for jobset in jobset_response.get("items", []):
+                    _match = self._metaflow_matching_spec(
+                        run_id=run_id,
+                        user=user,
+                        flow_name=flow_name,
+                        annotations=jobset.get("metadata", {}).get("annotations", {}),
+                        labels=jobset.get("metadata", {}).get("labels", {}),
+                    )
+                    if _match:
+                        matching_jobsets.append(jobset)
+
+            return matching_jobsets
+        except Exception as e:
+            self.echo(f"Error finding jobsets: {str(e)}\n")
+            return []
+
+    def _kill_pods_for_job(self, job):
+        """Find and kill pods associated with a specific job"""
+        job_name = job.metadata.name
+        namespace = job.metadata.namespace
+
+        try:
+            # Find pods with the job-name label matching this job
+            pods = self.api_instance.list_namespaced_pod(
+                namespace=namespace, label_selector=f"job-name={job_name}"
+            )
+
+            killed_pods = 0
+            for pod in pods.items:
+                if pod.status.phase in ["Running"]:
+                    self.echo(
+                        f"Killing processes in pod {pod.metadata.name} for job {job_name}"
+                    )
+                    if self._kill_pod_process(pod):
+                        killed_pods += 1
+
+            return killed_pods > 0
+        except Exception as e:
+            self.echo(f"Failed to find/kill pods for job {job_name}: {str(e)}")
+            return False
+
+    def _handle_job_outcome(self, job, outcome):
+        """Handle a job based on the derived outcome"""
+        job_name = job.metadata.name
+        namespace = job.metadata.namespace
+
+        if outcome == JobOutcomes.LEAVE_UNCHANGED:
+            # self.echo(f"Job {job_name} is in terminal state, leaving unchanged")
+            return None
+        elif outcome == JobOutcomes.DELETE:
+            self.echo(f"Deleting Job {job_name}")
+            return self._delete_job(job_name, namespace)
+        elif outcome == JobOutcomes.KILL:
+            self.echo(f"Killing Job {job_name}")
+            # First try to kill the pod processes
+            pods_killed = self._kill_pods_for_job(job)
+            if pods_killed > 0:
+                return True
+            # Worst case if we are not able to delete any pod, then delete the Job.
+            return self._delete_job(job_name, namespace)
+        else:
+            self.echo(f"Unknown outcome {outcome} for job {job_name}\n")
+            return False
+
+    def _handle_jobset_outcome(self, jobset, outcome):
+        """Handle a jobset based on the derived outcome"""
+        jobset_name = jobset.get("metadata", {}).get("name", "unknown")
+        namespace = jobset.get("metadata", {}).get("namespace", self._namespace)
+
+        if outcome == JobOutcomes.LEAVE_UNCHANGED:
+            # self.echo(f"JobSet {jobset_name} is in terminal state, leaving unchanged")
+            return None
+        elif outcome == JobOutcomes.DELETE:
+            self.echo(f"Deleting JobSet {jobset_name}")
+            try:
+                self.jobset_api.delete_namespaced_custom_object(
+                    group="jobset.x-k8s.io",
+                    version="v1alpha2",
+                    namespace=namespace,
+                    plural="jobsets",
+                    name=jobset_name,
+                )
+                return True
+            except Exception as e:
+                self.echo(f"Failed to delete JobSet {jobset_name}: {str(e)}")
+                return False
+        else:
+            self.echo(f"Unknown outcome {outcome} for JobSet {jobset_name}")
+            return False
+
+    def extract_matching_jobs_and_jobsets(self, flow_name, run_id, user):
+        """Extract matching jobs and jobsets based on the flow_name, run_id, and user criteria"""
+        jobs = self._find_matching_jobs(flow_name, run_id, user)
+        jobsets = self._find_matching_jobsets(flow_name, run_id, user)
+        return [(j, derive_job_outcome(j.status)) for j in jobs], [
+            (j, derive_jobset_outcome(j.get("status", {}))) for j in jobsets
+        ]
+
+    def process_matching_jobs_and_jobsets(self, flow_name, run_id, user):
+        """Process all matching jobs and jobsets based on their derived outcomes"""
+        results = []
+        progress_update = lambda x: x
+        if self.progress_bar:
+            progress_update = lambda x: self.progress_bar.update(1, x)
+
+        # Process matching jobs
+        _jobs, _jobsets = [], []
+        jobs = self._find_matching_jobs(flow_name, run_id, user)
+        for job in jobs:
+            outcome = derive_job_outcome(job.status)
+            result = self._handle_job_outcome(job, outcome)
+            # results.append(result)
+            if result is not None:
+                progress_update("💀 Killing Job %s" % job.metadata.name)
+                results.append(result)
+                _jobs.append(result)
+
+        # Process matching jobsets
+        jobsets = self._find_matching_jobsets(flow_name, run_id, user)
+        for jobset in jobsets:
+            jobset_status = jobset.get("status", {})
+            outcome = derive_jobset_outcome(jobset_status)
+            result = self._handle_jobset_outcome(jobset, outcome)
+            if result is not None:
+                progress_update(
+                    "💀 Deleting JobSet %s"
+                    % jobset.get("metadata", {}).get("name", "unknown")
+                )
+                results.append(result)
+                _jobsets.append(result)
+
+        return results, len(_jobs), len(_jobsets)
+
+    def process_matching_jobs_and_jobsets_force_all(self, flow_name, run_id, user):
+        """Force process ALL matching jobs and jobsets regardless of their status/outcome"""
+        results = []
+        progress_update = lambda x: x
+        if self.progress_bar:
+            progress_update = lambda x: self.progress_bar.update(1, x)
+
+        # Process matching jobs - FORCE DELETE ALL
+        _jobs, _jobsets = [], []
+        jobs = self._find_matching_jobs(flow_name, run_id, user)
+        for job in jobs:
+            # Force DELETE outcome regardless of actual status
+            result = self._handle_job_outcome(job, JobOutcomes.DELETE)
+            progress_update("🔥 FORCE Deleting Job %s" % job.metadata.name)
+            results.append(
+                result if result is not None else True
+            )  # Treat None as success for force mode
+            _jobs.append(result if result is not None else True)
+
+        # Process matching jobsets - FORCE DELETE ALL
+        jobsets = self._find_matching_jobsets(flow_name, run_id, user)
+        for jobset in jobsets:
+            # Force DELETE outcome regardless of actual status
+            result = self._handle_jobset_outcome(jobset, JobOutcomes.DELETE)
+            progress_update(
+                "🔥 FORCE Deleting JobSet %s"
+                % jobset.get("metadata", {}).get("name", "unknown")
+            )
+            results.append(
+                result if result is not None else True
+            )  # Treat None as success for force mode
+            _jobsets.append(result if result is not None else True)
+
+        return results, len(_jobs), len(_jobsets)

metaflow_extensions/outerbounds/plugins/nim/card.py

@@ -1,8 +1,7 @@
-import sqlite3
 from metaflow.cards import Markdown, Table
 from metaflow.metaflow_current import current
 
-from .utilities import get_storage_path
+from .utils import get_storage_path
 from ..card_utilities.async_cards import CardRefresher
 from ..card_utilities.extra_components import BarPlot, ViolinPlot
 
@@ -17,9 +16,7 @@ class NimMetricsRefresher(CardRefresher):
         self._file_name = get_storage_path(current.task_id)
 
     def sqlite_fetch_func(self, conn):
-        cursor = conn.cursor()
         try:
-            conn = sqlite3.connect(self._file_name)
             cursor = conn.cursor()
             cursor.execute(
                 "SELECT error, success, status_code, prompt_tokens, completion_tokens, e2e_time, model FROM metrics"
@@ -85,7 +82,6 @@ class NimMetricsRefresher(CardRefresher):
         current_card.refresh()
 
     def on_error(self, current_card, error_message):
-
         if isinstance(error_message, FileNotFoundError):
             return
 
@@ -99,7 +95,6 @@ class NimMetricsRefresher(CardRefresher):
             current_card.refresh()
 
     def update_only_components(self, current_card, data_object):
-
         # update request success data
         self._metrics_charts["request_success"].spec["data"][0]["values"] = [
             {