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

metaflow_extensions/outerbounds/plugins/checkpoint_datastores/external_chckpt.py

@@ -0,0 +1,85 @@
+from metaflow.user_decorators.user_flow_decorator import FlowMutator
+from metaflow.user_decorators.mutable_flow import MutableFlow
+from metaflow.user_decorators.mutable_step import MutableStep
+import os
+
+
+class _ExternalCheckpointFlowDeco(FlowMutator):
+    def init(self, *args, **kwargs):
+        self.bucket_path = kwargs.get("bucket_path", None)
+
+        self.secrets = kwargs.get("secrets", [])
+        if self.bucket_path is None:
+            raise ValueError(
+                "`bucket_path` keyword argument is required for the coreweave_datastore"
+            )
+        if not self.bucket_path.startswith("s3://"):
+            raise ValueError(
+                "`bucket_path` must start with `s3://` for the coreweave_datastore"
+            )
+        if self.secrets is None:
+            raise ValueError(
+                "`secrets` keyword argument is required for the coreweave_datastore"
+            )
+
+    def _swap_secrets(self, mutable_flow: MutableFlow) -> None:
+        from metaflow import (
+            checkpoint,
+            model,
+            huggingface_hub,
+            secrets,
+            with_artifact_store,
+        )
+
+        def _add_secrets(step: MutableStep) -> None:
+            decos_to_add = []
+            swapping_decos = {
+                "huggingface_hub": huggingface_hub,
+                "model": model,
+                "checkpoint": checkpoint,
+            }
+            already_has_secrets = False
+            secrets_present_in_deco = []
+            for d in step.decorator_specs:
+                name, _, _, deco_kwargs = d
+                if name in swapping_decos:
+                    decos_to_add.append((name, deco_kwargs))
+                elif name == "secrets":
+                    already_has_secrets = True
+                    secrets_present_in_deco.extend(deco_kwargs["sources"])
+
+            # If the step aleady has secrets then take all the sources in
+            # the secrets and add the addtional secrets to the existing secrets
+            secrets_to_add = self.secrets
+            if already_has_secrets:
+                secrets_to_add.extend(secrets_present_in_deco)
+
+            secrets_to_add = list(set(secrets_to_add))
+
+            if len(decos_to_add) == 0:
+                if already_has_secrets:
+                    step.remove_decorator("secrets")
+
+                step.add_decorator(
+                    secrets,
+                    deco_kwargs=dict(
+                        sources=secrets_to_add,
+                    ),
+                )
+                return
+
+            for d, _ in decos_to_add:
+                step.remove_decorator(d)
+
+            step.add_decorator(
+                secrets,
+                deco_kwargs=dict(
+                    sources=secrets_to_add,
+                ),
+            )
+            for d, attrs in decos_to_add:
+                _deco_to_add = swapping_decos[d]
+                step.add_decorator(_deco_to_add, deco_kwargs=attrs)
+
+        for step_name, step in mutable_flow.steps:
+            _add_secrets(step)

metaflow_extensions/outerbounds/plugins/checkpoint_datastores/nebius.py

@@ -1,14 +1,11 @@
-from metaflow.user_configs.config_decorators import (
-    MutableFlow,
-    MutableStep,
-    CustomFlowDecorator,
-)
+from metaflow.user_decorators.mutable_flow import MutableFlow
+from .external_chckpt import _ExternalCheckpointFlowDeco
 import os
 
 NEBIUS_ENDPOINT_URL = "https://storage.eu-north1.nebius.cloud:443"
 
 
-class nebius_checkpoints(CustomFlowDecorator):
+class nebius_checkpoints(_ExternalCheckpointFlowDeco):
 
     """
 
@@ -52,78 +49,14 @@ class nebius_checkpoints(CustomFlowDecorator):
         super().__init__(*args, **kwargs)
 
     def init(self, *args, **kwargs):
-        self.bucket_path = kwargs.get("bucket_path", None)
-
-        self.secrets = kwargs.get("secrets", [])
-        if self.bucket_path is None:
-            raise ValueError(
-                "`bucket_path` keyword argument is required for the coreweave_datastore"
-            )
-        if not self.bucket_path.startswith("s3://"):
-            raise ValueError(
-                "`bucket_path` must start with `s3://` for the coreweave_datastore"
-            )
-
+        super().init(*args, **kwargs)
         self.nebius_endpoint_url = kwargs.get("endpoint_url", NEBIUS_ENDPOINT_URL)
-        if self.secrets is None:
-            raise ValueError(
-                "`secrets` keyword argument is required for the coreweave_datastore"
-            )
 
-    def evaluate(self, mutable_flow: MutableFlow) -> None:
+    def pre_mutate(self, mutable_flow: MutableFlow) -> None:
         from metaflow import (
-            checkpoint,
-            model,
-            huggingface_hub,
-            secrets,
             with_artifact_store,
         )
 
-        def _add_secrets(step: MutableStep) -> None:
-            decos_to_add = []
-            swapping_decos = {
-                "huggingface_hub": huggingface_hub,
-                "model": model,
-                "checkpoint": checkpoint,
-            }
-            already_has_secrets = False
-            secrets_present_in_deco = []
-            for d in step.decorators:
-                if d.name in swapping_decos:
-                    decos_to_add.append((d.name, d.attributes))
-                elif d.name == "secrets":
-                    already_has_secrets = True
-                    secrets_present_in_deco.extend(d.attributes["sources"])
-
-            # If the step aleady has secrets then take all the sources in
-            # the secrets and add the addtional secrets to the existing secrets
-            secrets_to_add = self.secrets
-            if already_has_secrets:
-                secrets_to_add.extend(secrets_present_in_deco)
-
-            secrets_to_add = list(set(secrets_to_add))
-
-            if len(decos_to_add) == 0:
-                if already_has_secrets:
-                    step.remove_decorator("secrets")
-
-                step.add_decorator(
-                    secrets,
-                    sources=secrets_to_add,
-                )
-                return
-
-            for d, _ in decos_to_add:
-                step.remove_decorator(d)
-
-            step.add_decorator(
-                secrets,
-                sources=secrets_to_add,
-            )
-            for d, attrs in decos_to_add:
-                _deco_to_add = swapping_decos[d]
-                step.add_decorator(_deco_to_add, **attrs)
-
         def _nebius_config():
             return {
                 "root": self.bucket_path,
@@ -135,10 +68,6 @@ class nebius_checkpoints(CustomFlowDecorator):
             }
 
         mutable_flow.add_decorator(
-            with_artifact_store,
-            type="s3",
-            config=_nebius_config,
+            with_artifact_store, deco_kwargs=dict(type="s3", config=_nebius_config)
         )
-
-        for step_name, step in mutable_flow.steps:
-            _add_secrets(step)
+        self._swap_secrets(mutable_flow)

metaflow_extensions/outerbounds/plugins/fast_bakery/baker.py

@@ -0,0 +1,119 @@
+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,
+    fast_bakery_url: Optional[str] = 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
+        fast_bakery_url: Optional FB URL
+
+    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)
+
+    if all([fast_bakery_url is None and FAST_BAKERY_URL is None]):
+        raise BakerException(
+            "Image Bakery endpoint missing. METAFLOW_FAST_BAKERY_URL environment/configuration variable not found."
+        )
+
+    fast_bakery_url = fast_bakery_url or FAST_BAKERY_URL
+
+    # 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