ob-metaflow-extensions 1.1.45rc3__py2.py3-none-any.whl → 1.5.1__py2.py3-none-any.whl

metaflow_extensions/outerbounds/plugins/card_utilities/injector.py

@@ -0,0 +1,70 @@
+from metaflow.exception import MetaflowException
+from collections import defaultdict
+
+
+class CardDecoratorInjector:
+    """
+    Mixin Useful for injecting @card decorators from other first class Metaflow decorators.
+    """
+
+    _first_time_init = defaultdict(dict)
+
+    @classmethod
+    def _get_first_time_init_cached_value(cls, step_name, card_id):
+        return cls._first_time_init.get(step_name, {}).get(card_id, None)
+
+    @classmethod
+    def _set_first_time_init_cached_value(cls, step_name, card_id, value):
+        cls._first_time_init[step_name][card_id] = value
+
+    def _card_deco_already_attached(self, step, card_id):
+        for decorator in step.decorators:
+            if decorator.name == "card":
+                if decorator.attributes["id"] and card_id == decorator.attributes["id"]:
+                    return True
+        return False
+
+    def _get_step(self, flow, step_name):
+        for step in flow:
+            if step.name == step_name:
+                return step
+        return None
+
+    def _first_time_init_check(self, step_dag_node, card_id):
+        """ """
+        return not self._card_deco_already_attached(step_dag_node, card_id)
+
+    def attach_card_decorator(
+        self,
+        flow,
+        step_name,
+        card_id,
+        card_type,
+        refresh_interval=5,
+    ):
+        """
+        This method is called `step_init` in your StepDecorator code since
+        this class is used as a Mixin
+        """
+        from metaflow import decorators as _decorators
+
+        if not all([card_id, card_type]):
+            raise MetaflowException(
+                "`INJECTED_CARD_ID` and `INJECTED_CARD_TYPE` must be set in the `CardDecoratorInjector` Mixin"
+            )
+
+        step_dag_node = self._get_step(flow, step_name)
+        if (
+            self._get_first_time_init_cached_value(step_name, card_id) is None
+        ):  # First check class level setting.
+            if self._first_time_init_check(step_dag_node, card_id):
+                self._set_first_time_init_cached_value(step_name, card_id, True)
+                _decorators._attach_decorators_to_step(
+                    step_dag_node,
+                    [
+                        "card:type=%s,id=%s,refresh_interval=%s"
+                        % (card_type, card_id, str(refresh_interval))
+                    ],
+                )
+            else:
+                self._set_first_time_init_cached_value(step_name, card_id, False)

metaflow_extensions/outerbounds/plugins/checkpoint_datastores/__init__.py

@@ -0,0 +1,2 @@
+from .nebius import nebius_checkpoints
+from .coreweave import coreweave_checkpoints

metaflow_extensions/outerbounds/plugins/checkpoint_datastores/coreweave.py

@@ -0,0 +1,71 @@
+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
+from .external_chckpt import _ExternalCheckpointFlowDeco
+import os
+
+
+class coreweave_checkpoints(_ExternalCheckpointFlowDeco):
+
+    """
+
+    This decorator is used for setting the coreweave object store as the artifact store for checkpoints/models created by the flow.
+
+    Parameters
+    ----------
+    secrets: list
+        A list of secrets to be added to the step. These secrets should contain any secrets that are required globally and the secret
+        for the coreweave object store. The secret should contain the following keys:
+        - COREWEAVE_ACCESS_KEY
+        - COREWEAVE_SECRET_KEY
+
+    bucket_path: str
+        The path to the bucket to store the checkpoints/models.
+
+    Usage
+    -----
+    ```python
+    from metaflow import checkpoint, step, FlowSpec, coreweave_checkpoints
+
+    @coreweave_checkpoints(secrets=[], bucket_path=None)
+    class MyFlow(FlowSpec):
+        @checkpoint
+        @step
+        def start(self):
+            # Saves the checkpoint in the coreweave object store
+            current.checkpoint.save("./foo.txt")
+
+        @step
+        def end(self):
+            pass
+    ```
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def init(self, *args, **kwargs):
+        super().init(*args, **kwargs)
+        self.coreweave_endpoint_url = f"https://cwobject.com"
+
+    def pre_mutate(self, mutable_flow: MutableFlow) -> None:
+        from metaflow import (
+            with_artifact_store,
+        )
+
+        def _coreweave_config():
+            return {
+                "root": self.bucket_path,
+                "client_params": {
+                    "aws_access_key_id": os.environ.get("COREWEAVE_ACCESS_KEY"),
+                    "aws_secret_access_key": os.environ.get("COREWEAVE_SECRET_KEY"),
+                    "endpoint_url": self.coreweave_endpoint_url,
+                    "config": dict(s3={"addressing_style": "virtual"}),
+                },
+            }
+
+        mutable_flow.add_decorator(
+            with_artifact_store,
+            deco_kwargs=dict(type="coreweave", config=_coreweave_config),
+        )
+        self._swap_secrets(mutable_flow)

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

@@ -0,0 +1,73 @@
+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(_ExternalCheckpointFlowDeco):
+
+    """
+
+    This decorator is used for setting the nebius's S3 compatible object store as the artifact store for
+    checkpoints/models created by the flow.
+
+    Parameters
+    ----------
+    secrets: list
+        A list of secrets to be added to the step. These secrets should contain any secrets that are required globally and the secret
+        for the nebius object store. The secret should contain the following keys:
+        - NEBIUS_ACCESS_KEY
+        - NEBIUS_SECRET_KEY
+
+    bucket_path: str
+        The path to the bucket to store the checkpoints/models.
+
+    endpoint_url: str
+        The endpoint url for the nebius object store. Defaults to `https://storage.eu-north1.nebius.cloud:443`
+
+    Usage
+    -----
+    ```python
+    from metaflow import checkpoint, step, FlowSpec, nebius_checkpoints
+
+    @nebius_checkpoints(secrets=[], bucket_path=None)
+    class MyFlow(FlowSpec):
+        @checkpoint
+        @step
+        def start(self):
+            # Saves the checkpoint in the nebius object store
+            current.checkpoint.save("./foo.txt")
+
+        @step
+        def end(self):
+            pass
+    ```
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def init(self, *args, **kwargs):
+        super().init(*args, **kwargs)
+        self.nebius_endpoint_url = kwargs.get("endpoint_url", NEBIUS_ENDPOINT_URL)
+
+    def pre_mutate(self, mutable_flow: MutableFlow) -> None:
+        from metaflow import (
+            with_artifact_store,
+        )
+
+        def _nebius_config():
+            return {
+                "root": self.bucket_path,
+                "client_params": {
+                    "aws_access_key_id": os.environ.get("NEBIUS_ACCESS_KEY"),
+                    "aws_secret_access_key": os.environ.get("NEBIUS_SECRET_KEY"),
+                    "endpoint_url": self.nebius_endpoint_url,
+                },
+            }
+
+        mutable_flow.add_decorator(
+            with_artifact_store, deco_kwargs=dict(type="s3", config=_nebius_config)
+        )
+        self._swap_secrets(mutable_flow)

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,
+    )