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

metaflow_extensions/outerbounds/plugins/apps/core/utils.py

@@ -0,0 +1,233 @@
+import random
+import time
+import sys
+import json
+import requests
+from typing import Optional
+
+# This click import is not used to construct any ob
+# package cli. Its used only for printing stuff.
+# So we can use the static metaflow._vendor import path
+from metaflow._vendor import click
+from .app_config import CAPSULE_DEBUG
+import sys
+import threading
+import time
+import logging
+import itertools
+from typing import Union, Callable, Any, List
+
+from ._vendor.spinner import (
+    Spinners,
+)
+
+
+class MultiStepSpinner:
+    """
+    A spinner that supports multi-step progress and configurable alignment.
+
+    Parameters
+    ----------
+    spinner : Spinners
+        Which spinner frames/interval to use.
+    text : str
+        Static text to display beside the spinner.
+    color : str, optional
+        Click color name.
+    align : {'left','right'}
+        Whether to render the spinner to the left (default) or right of the text.
+    """
+
+    def __init__(
+        self,
+        spinner: Spinners = Spinners.dots,
+        text: str = "",
+        color: Optional[str] = None,
+        align: str = "right",
+        file=sys.stdout,
+    ):
+        cfg = spinner.value
+        self.frames = cfg["frames"]
+        self.interval = float(cfg["interval"]) / 1000.0  # type: ignore
+        self.text = text
+        self.color = color
+        if align not in ("left", "right"):
+            raise ValueError("align must be 'left' or 'right'")
+        self.align = align
+        self._write_file = file
+        # precompute clear length: max frame width + space + text length
+        max_frame = max(self.frames, key=lambda x: len(x))  # type: ignore
+        self.clear_len = len(self.main_text) + len(max_frame) + 1
+
+        self._stop_evt = threading.Event()
+        self._pause_evt = threading.Event()
+        self._thread = None
+        self._write_lock = threading.Lock()
+
+    @property
+    def main_text(self):
+        # if self.text is a callable then call it
+        if callable(self.text):
+            return self.text()
+        return self.text
+
+    def _spin(self):
+        for frame in itertools.cycle(self.frames):
+            if self._stop_evt.is_set():
+                break
+            if self._pause_evt.is_set():
+                time.sleep(0.05)
+                continue
+
+            # ---- Core logging critical section ----
+            with self._write_lock:
+                symbol = click.style(frame, fg=self.color) if self.color else frame
+                if self.align == "left":
+                    msg = f"{symbol} {self.main_text}"
+                else:
+                    msg = f"{self.main_text} {symbol}"
+
+                click.echo(msg, nl=False, file=self._write_file)
+                click.echo("\r", nl=False, file=self._write_file)
+                self._write_file.flush()
+            # ---- End of critical section ----
+            time.sleep(self.interval)
+        # clear the line when done
+        self._clear_line()
+
+    def _clear_line(self):
+        with self._write_lock:
+            click.echo(" " * self.clear_len, nl=False, file=self._write_file)
+            click.echo("\r", nl=False, file=self._write_file)
+            self._write_file.flush()
+
+    def start(self):
+        if self._thread and self._thread.is_alive():
+            return
+        self._stop_evt.clear()
+        self._pause_evt.clear()
+        self._thread = threading.Thread(target=self._spin, daemon=True)
+        self._thread.start()
+
+    def stop(self):
+        self._stop_evt.set()
+        if self._thread:
+            self._thread.join()
+
+    def log(self, *messages: str):
+        """Pause the spinner, emit a ✔ + message, then resume."""
+        self._pause_evt.set()
+        self._clear_line()
+        # ---- Core logging critical section ----
+        with self._write_lock:
+            self._write_file.flush()
+            for message in messages:
+                click.echo(f"{message}", file=self._write_file, nl=True)
+            self._write_file.flush()
+        # ---- End of critical section ----
+        self._pause_evt.clear()
+
+    def __enter__(self):
+        self.start()
+        return self
+
+    def __exit__(self, exc_type, exc, tb):
+        self.stop()
+
+
+class SpinnerLogHandler(logging.Handler):
+    def __init__(self, spinner: MultiStepSpinner, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.spinner = spinner
+
+    def emit(self, record):
+        msg = self.format(record)
+        self.spinner.log(msg)
+
+
+class MaximumRetriesExceeded(Exception):
+    def __init__(self, url, method, status_code, text):
+        self.url = url
+        self.method = method
+        self.status_code = status_code
+        self.text = text
+
+    def __str__(self):
+        return f"Maximum retries exceeded for {self.url}[{self.method}] {self.status_code} {self.text}"
+
+
+class TODOException(Exception):
+    pass
+
+
+requests_funcs = [
+    requests.get,
+    requests.post,
+    requests.put,
+    requests.delete,
+    requests.patch,
+    requests.head,
+    requests.options,
+]
+
+
+def safe_requests_wrapper(
+    requests_module_fn,
+    *args,
+    conn_error_retries=2,
+    retryable_status_codes=[409],
+    logger_fn=None,
+    **kwargs,
+):
+    """
+    There are two categories of errors that we need to handle when dealing with any API server.
+    1. HTTP errors. These are are errors that are returned from the API server.
+        - How to handle retries for this case will be application specific.
+    2. Errors when the API server may not be reachable (DNS resolution / network issues)
+        - In this scenario, we know that something external to the API server is going wrong causing the issue.
+        - Failing prematurely in the case might not be the best course of action since critical user jobs might crash on intermittent issues.
+        - So in this case, we can just plainly retry the request.
+
+    This function handles the second case. It's a simple wrapper to handle the retry logic for connection errors.
+    If this function is provided a `conn_error_retries` of 5, then the last retry will have waited 32 seconds.
+    Generally this is a safe enough number of retries after which we can assume that something is really broken. Until then,
+    there can be intermittent issues that would resolve themselves if we retry gracefully.
+    """
+    if requests_module_fn not in requests_funcs:
+        raise ValueError(
+            f"safe_requests_wrapper doesn't support {requests_module_fn.__name__}. You can only use the following functions: {requests_funcs}"
+        )
+
+    _num_retries = 0
+    noise = random.uniform(-0.5, 0.5)
+    response = None
+    while _num_retries < conn_error_retries:
+        try:
+            response = requests_module_fn(*args, **kwargs)
+            if response.status_code not in retryable_status_codes:
+                return response
+            if CAPSULE_DEBUG:
+                if logger_fn:
+                    logger_fn(
+                        f"[outerbounds-debug] safe_requests_wrapper: {response.url}[{requests_module_fn.__name__}] {response.status_code} {response.text}",
+                    )
+                else:
+                    print(
+                        f"[outerbounds-debug] safe_requests_wrapper: {response.url}[{requests_module_fn.__name__}] {response.status_code} {response.text}",
+                        file=sys.stderr,
+                    )
+            _num_retries += 1
+            time.sleep((2 ** (_num_retries + 1)) + noise)
+        except requests.exceptions.ConnectionError:
+            if _num_retries <= conn_error_retries - 1:
+                # Exponential backoff with 2^(_num_retries+1) seconds
+                time.sleep((2 ** (_num_retries + 1)) + noise)
+                _num_retries += 1
+            else:
+                raise
+    raise MaximumRetriesExceeded(
+        response.url,
+        requests_module_fn.__name__,
+        response.status_code,
+        response.text,
+    )

metaflow_extensions/outerbounds/plugins/apps/core/validations.py

@@ -0,0 +1,17 @@
+import os
+from typing import List
+from .app_config import AppConfig, AppConfigError
+from .secrets import SecretRetriever, SecretNotFound
+
+
+def secrets_validator(secrets: List[str]):
+    secret_retriever = SecretRetriever()
+    for secret in secrets:
+        try:
+            secret_retriever.get_secret_as_dict(secret)
+        except SecretNotFound:
+            raise Exception(f"Secret named `{secret}` not found")
+
+
+def run_validations(app_config: AppConfig):
+    pass

metaflow_extensions/outerbounds/plugins/aws/__init__.py

@@ -0,0 +1,4 @@
+from .assume_role_decorator import assume_role
+from .assume_role import OBP_ASSUME_ROLE_ARN_ENV_VAR
+
+__mf_promote_submodules__ = ["plugins.aws"]

metaflow_extensions/outerbounds/plugins/aws/assume_role.py

@@ -0,0 +1,3 @@
+# Environment variable name used by the assume_role decorator to pass
+# the role ARN to AWS client functions
+OBP_ASSUME_ROLE_ARN_ENV_VAR = "OBP_ASSUME_ROLE_ARN"

metaflow_extensions/outerbounds/plugins/aws/assume_role_decorator.py

@@ -0,0 +1,78 @@
+from metaflow.user_decorators.mutable_flow import MutableFlow
+from metaflow.user_decorators.mutable_step import MutableStep
+from metaflow.user_decorators.user_flow_decorator import FlowMutator
+from .assume_role import OBP_ASSUME_ROLE_ARN_ENV_VAR
+
+
+class assume_role(FlowMutator):
+    """
+    Flow-level decorator for assuming AWS IAM roles.
+
+    When applied to a flow, all steps in the flow will automatically use the specified IAM role-arn
+    as their source principal.
+
+    Usage:
+    ------
+    @assume_role(role_arn="arn:aws:iam::123456789012:role/my-iam-role")
+    class MyFlow(FlowSpec):
+        @step
+        def start(self):
+            import boto3
+            client = boto3.client("dynamodb")  # Automatically uses the role in the flow decorator
+            self.next(self.end)
+
+        @step
+        def end(self):
+            from metaflow import get_aws_client
+            client = get_aws_client("dynamodb")  # Automatically uses the role in the flow decorator
+    """
+
+    def init(self, *args, **kwargs):
+        self.role_arn = kwargs.get("role_arn", None)
+
+        if self.role_arn is None:
+            raise ValueError(
+                "`role_arn` keyword argument is required for the assume_role decorator"
+            )
+
+        if not self.role_arn.startswith("arn:aws:iam::"):
+            raise ValueError(
+                "`role_arn` must be a valid AWS IAM role ARN starting with 'arn:aws:iam::'"
+            )
+
+    def pre_mutate(self, mutable_flow: MutableFlow) -> None:
+        """
+        This method is called by Metaflow to apply the decorator to the flow.
+        It sets up environment variables that will be used by the AWS client
+        to automatically assume the specified role.
+        """
+        # Import environment decorator at runtime to avoid circular imports
+        from metaflow import environment
+
+        def _swap_environment_variables(step: MutableStep, role_arn: str) -> None:
+            _step_has_env_set = True
+            _env_kwargs = {OBP_ASSUME_ROLE_ARN_ENV_VAR: role_arn}
+            for d in step.decorator_specs:
+                name, _, _, deco_kwargs = d
+                if name == "environment":
+                    _env_kwargs.update(deco_kwargs["vars"])
+                    _step_has_env_set = True
+
+            if _step_has_env_set:
+                # remove the environment decorator
+                step.remove_decorator("environment")
+
+            # add the environment decorator
+            step.add_decorator(
+                environment,
+                deco_kwargs=dict(vars=_env_kwargs),
+            )
+
+        # Set the role ARN as an environment variable that will be picked up
+        # by the get_aws_client function
+        def _setup_role_assumption(step: MutableStep) -> None:
+            _swap_environment_variables(step, self.role_arn)
+
+        # Apply the role assumption setup to all steps in the flow
+        for _, step in mutable_flow.steps:
+            _setup_role_assumption(step)

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