ob-metaflow-extensions 1.1.161__py2.py3-none-any.whl → 1.1.161rc1__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/toplevel/global_aliases_for_metaflow_package.py

@@ -53,3 +53,4 @@ def S3(*args, **kwargs):
 from .. import profilers
 from ..plugins.snowflake import Snowflake
 from ..plugins.checkpoint_datastores import nebius_checkpoints, coreweave_checkpoints
+from . import ob_internal

metaflow_extensions/outerbounds/toplevel/ob_internal.py

@@ -0,0 +1 @@
+from ..plugins.fast_bakery.baker import bake_image

{ob_metaflow_extensions-1.1.161.dist-info → ob_metaflow_extensions-1.1.161rc1.dist-info}/METADATA

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

{ob_metaflow_extensions-1.1.161.dist-info → ob_metaflow_extensions-1.1.161rc1.dist-info}/RECORD

@@ -17,6 +17,7 @@ metaflow_extensions/outerbounds/plugins/checkpoint_datastores/__init__.py,sha256
 metaflow_extensions/outerbounds/plugins/checkpoint_datastores/coreweave.py,sha256=_WzoOROFjoFa8TzsMNFp-r_1Zz7NUp-5ljn_kKlczXA,4534
 metaflow_extensions/outerbounds/plugins/checkpoint_datastores/nebius.py,sha256=zgqDLFewCeF5jqh-hUNKmC_OAjld09ln0bb8Lkeqapc,4659
 metaflow_extensions/outerbounds/plugins/fast_bakery/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+metaflow_extensions/outerbounds/plugins/fast_bakery/baker.py,sha256=ShE5omFBr83wkvEhL_ptRFvDNMs6wefg4BjaafQjTcM,3602
 metaflow_extensions/outerbounds/plugins/fast_bakery/docker_environment.py,sha256=Tl520HdBteg-aDOM7mnnJJpdDCZc49BmFFmLUc_vTi8,15018
 metaflow_extensions/outerbounds/plugins/fast_bakery/fast_bakery.py,sha256=PE81ZB54OAMXkMGSB7JqgvgMg7N9kvoVclrWL-6jc2U,5626
 metaflow_extensions/outerbounds/plugins/fast_bakery/fast_bakery_cli.py,sha256=kqFyu2bJSnc9_9aYfBpz5xK6L6luWFZK_NMuh8f1eVk,1494
@@ -66,14 +67,15 @@ metaflow_extensions/outerbounds/plugins/torchtune/__init__.py,sha256=TOXNeyhcgd8
 metaflow_extensions/outerbounds/profilers/__init__.py,sha256=wa_jhnCBr82TBxoS0e8b6_6sLyZX0fdHicuGJZNTqKw,29
 metaflow_extensions/outerbounds/profilers/gpu.py,sha256=3Er8uKQzfm_082uadg4yn_D4Y-iSCgzUfFmguYxZsz4,27485
 metaflow_extensions/outerbounds/toplevel/__init__.py,sha256=qWUJSv_r5hXJ7jV_On4nEasKIfUCm6_UjkjXWA_A1Ts,90
-metaflow_extensions/outerbounds/toplevel/global_aliases_for_metaflow_package.py,sha256=oCRFfwoLODDju2-6JU2h74x6O-GCyO7PRRKTZsj-73k,2009
+metaflow_extensions/outerbounds/toplevel/global_aliases_for_metaflow_package.py,sha256=505fLAY4NkBCGtjpXlE0RgaKVpJ0jOUq1-Fq-EKzPew,2035
+metaflow_extensions/outerbounds/toplevel/ob_internal.py,sha256=53xM6d_UYT3uGFFA59UzxN23H5QMO5_F39pALpmGy04,51
 metaflow_extensions/outerbounds/toplevel/plugins/azure/__init__.py,sha256=WUuhz2YQfI4fz7nIcipwwWq781eaoHEk7n4GAn1npDg,63
 metaflow_extensions/outerbounds/toplevel/plugins/gcp/__init__.py,sha256=BbZiaH3uILlEZ6ntBLKeNyqn3If8nIXZFq_Apd7Dhco,70
 metaflow_extensions/outerbounds/toplevel/plugins/kubernetes/__init__.py,sha256=5zG8gShSj8m7rgF4xgWBZFuY3GDP5n1T0ktjRpGJLHA,69
 metaflow_extensions/outerbounds/toplevel/plugins/ollama/__init__.py,sha256=GRSz2zwqkvlmFS6bcfYD_CX6CMko9DHQokMaH1iBshA,47
 metaflow_extensions/outerbounds/toplevel/plugins/snowflake/__init__.py,sha256=LptpH-ziXHrednMYUjIaosS1SXD3sOtF_9_eRqd8SJw,50
 metaflow_extensions/outerbounds/toplevel/plugins/torchtune/__init__.py,sha256=uTVkdSk3xZ7hEKYfdlyVteWj5KeDwaM1hU9WT-_YKfI,50
-ob_metaflow_extensions-1.1.161.dist-info/METADATA,sha256=5iL8FqYFqx1U_GR1YUAzqeV1NrR44jnr_vFSUOmjJEo,521
-ob_metaflow_extensions-1.1.161.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
-ob_metaflow_extensions-1.1.161.dist-info/top_level.txt,sha256=NwG0ukwjygtanDETyp_BUdtYtqIA_lOjzFFh1TsnxvI,20
-ob_metaflow_extensions-1.1.161.dist-info/RECORD,,
+ob_metaflow_extensions-1.1.161rc1.dist-info/METADATA,sha256=QZ_Y0zbT95Qr57WQ4XnvhZUSXoOQy3x2Pd-KT73mnxo,524
+ob_metaflow_extensions-1.1.161rc1.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
+ob_metaflow_extensions-1.1.161rc1.dist-info/top_level.txt,sha256=NwG0ukwjygtanDETyp_BUdtYtqIA_lOjzFFh1TsnxvI,20
+ob_metaflow_extensions-1.1.161rc1.dist-info/RECORD,,