ob-metaflow-extensions 1.1.86__tar.gz → 1.1.88__tar.gz

{ob-metaflow-extensions-1.1.86 → ob-metaflow-extensions-1.1.88}/PKG-INFO

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

{ob-metaflow-extensions-1.1.86 → ob-metaflow-extensions-1.1.88}/metaflow_extensions/outerbounds/plugins/__init__.py

@@ -104,7 +104,7 @@ def get_boto3_session(role_arn=None, session_vars=None):
             tmp_aws_config_file = f.name
         os.rename(tmp_aws_config_file, aws_config_file)
         os.environ["AWS_CONFIG_FILE"] = aws_config_file
-        os.environ["AWS_DEFAULT_PROFILE"] = "cspr"
+        os.environ["AWS_PROFILE"] = "cspr"
     else:
         os.environ["AWS_WEB_IDENTITY_TOKEN_FILE"] = token_file
         os.environ["AWS_ROLE_ARN"] = token_info["role_arn"]
@@ -122,7 +122,7 @@ def get_boto3_session(role_arn=None, session_vars=None):
             #  AWS_CONFIG_FILE environment variable above.
             if role_arn == USE_CSPR_ROLE_ARN_IF_SET:
                 # Otherwise start from the default profile, assuming CSPR role
-                session = boto3.session.Session(profile_name="default")
+                session = boto3.session.Session(profile_name="cspr")
             else:
                 session = boto3.session.Session(profile_name="task")
         else:

{ob-metaflow-extensions-1.1.86 → ob-metaflow-extensions-1.1.88}/metaflow_extensions/outerbounds/plugins/fast_bakery/docker_environment.py

@@ -1,30 +1,30 @@
 import hashlib
 import json
 import os
-
+import threading
+import time
+import uuid
 from concurrent.futures import ThreadPoolExecutor
 from typing import Dict
+
 from metaflow.exception import MetaflowException
-from metaflow.metaflow_config import (
-    FAST_BAKERY_URL,
-    get_pinned_conda_libs,
-)
+from metaflow.metaflow_config import FAST_BAKERY_URL, get_pinned_conda_libs
 from metaflow.metaflow_environment import MetaflowEnvironment
-from metaflow.plugins.pypi.conda_environment import CondaEnvironment
-from .fast_bakery import FastBakery, FastBakeryApiResponse, FastBakeryException
 from metaflow.plugins.aws.batch.batch_decorator import BatchDecorator
 from metaflow.plugins.kubernetes.kubernetes_decorator import KubernetesDecorator
 from metaflow.plugins.pypi.conda_decorator import CondaStepDecorator
+from metaflow.plugins.pypi.conda_environment import CondaEnvironment
 from metaflow.plugins.pypi.pypi_decorator import PyPIStepDecorator
 
+from .fast_bakery import FastBakery, FastBakeryApiResponse, FastBakeryException
+
 BAKERY_METAFILE = ".imagebakery-cache"
 
+import fcntl
 import json
 import os
-import fcntl
-from functools import wraps
 from concurrent.futures import ThreadPoolExecutor
-
+from functools import wraps
 
 # TODO - ensure that both @conda/@pypi are not assigned to the same step
 
@@ -36,6 +36,9 @@ def cache_request(cache_file):
             call_args = kwargs.copy()
             call_args.update(zip(func.__code__.co_varnames, args))
             call_args.pop("self", None)
+            call_args.pop("ref", None)
+            # invalidate cache when moving from one deployment to another
+            call_args.update({"fast_bakery_url": FAST_BAKERY_URL})
             cache_key = hashlib.md5(
                 json.dumps(call_args, sort_keys=True).encode("utf-8")
             ).hexdigest()
@@ -79,7 +82,7 @@ def cache_request(cache_file):
 
 
 class DockerEnvironmentException(MetaflowException):
-    headline = "Ran into an error while setting up the environment"
+    headline = "Ran into an error while baking image"
 
     def __init__(self, msg):
         super(DockerEnvironmentException, self).__init__(msg)
@@ -93,8 +96,8 @@ class DockerEnvironment(MetaflowEnvironment):
         self.skipped_steps = set()
         self.flow = flow
 
-        self.bakery = FastBakery(url=FAST_BAKERY_URL)
         self.results = {}
+        self.images_baked = 0
 
     def set_local_root(self, local_root):
         self.local_root = local_root
@@ -102,15 +105,31 @@ class DockerEnvironment(MetaflowEnvironment):
     def decospecs(self):
         return ("conda", "fast_bakery_internal") + super().decospecs()
 
-    def validate_environment(self, echo, datastore_type):
+    def validate_environment(self, logger, datastore_type):
         self.datastore_type = datastore_type
-        self.echo = echo
+        self.logger = logger
 
         # Avoiding circular imports.
         from metaflow.plugins import DATASTORES
 
         self.datastore = [d for d in DATASTORES if d.TYPE == self.datastore_type][0]
 
+        # Mixing @pypi/@conda in a single step is not supported yet
+        for step in self.flow:
+            if (
+                sum(
+                    1
+                    for deco in step.decorators
+                    if isinstance(deco, (PyPIStepDecorator, CondaStepDecorator))
+                )
+                > 1
+            ):
+                raise MetaflowException(
+                    "Mixing and matching PyPI packages and Conda packages within a\n"
+                    "step is not yet supported. Use one of @pypi or @conda only for the *%s* step."
+                    % step.name
+                )
+
     def init_environment(self, echo):
         self.skipped_steps = {
             step.name
@@ -125,14 +144,21 @@ class DockerEnvironment(MetaflowEnvironment):
             step for step in self.flow if step.name not in self.skipped_steps
         ]
         if steps_to_bake:
-            echo("Baking container image(s) ...")
-            self.results = self._bake(steps_to_bake, echo)
+            self.logger("🚀 Baking container image(s) ...")
+            start_time = time.time()
+            self.results = self._bake(steps_to_bake)
             for step in self.flow:
                 for d in step.decorators:
                     if isinstance(d, (BatchDecorator, KubernetesDecorator)):
                         d.attributes["image"] = self.results[step.name].container_image
                         d.attributes["executable"] = self.results[step.name].python_path
-            echo("Container image(s) baked!")
+            if self.images_baked > 0:
+                bake_time = time.time() - start_time
+                self.logger(
+                    f"🎉 All container image(s) baked in {bake_time:.2f} seconds!"
+                )
+            else:
+                self.logger("🎉 All container image(s) baked!")
 
         if self.skipped_steps:
             self.delegate = CondaEnvironment(self.flow)
@@ -140,29 +166,54 @@ class DockerEnvironment(MetaflowEnvironment):
             self.delegate.validate_environment(echo, self.datastore_type)
             self.delegate.init_environment(echo, self.skipped_steps)
 
-    def _bake(self, steps, echo) -> Dict[str, FastBakeryApiResponse]:
+    def _bake(self, steps) -> Dict[str, FastBakeryApiResponse]:
         metafile_path = get_fastbakery_metafile_path(self.local_root, self.flow.name)
+        logger_lock = threading.Lock()
 
         @cache_request(metafile_path)
         def _cached_bake(
-            python=None, pypi_packages=None, conda_packages=None, base_image=None
+            ref=None,
+            python=None,
+            pypi_packages=None,
+            conda_packages=None,
+            base_image=None,
         ):
-            self.bakery._reset_payload()
-            self.bakery.python_version(python)
-            self.bakery.pypi_packages(pypi_packages)
-            self.bakery.conda_packages(conda_packages)
-            self.bakery.base_image(base_image)
-            # self.bakery.ignore_cache()
+            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:
+                self.logger(f"🍳 Baking [{ref}] ...")
+                self.logger(f"     🐍 Python: {python}")
+
+                if pypi_packages:
+                    self.logger(f"     📦 PyPI packages:")
+                    for package, version in pypi_packages.items():
+                        self.logger(f"        🔧 {package}: {version}")
+
+                if conda_packages:
+                    self.logger(f"     📦 Conda packages:")
+                    for package, version in conda_packages.items():
+                        self.logger(f"        🔧 {package}: {version}")
+
+                self.logger(f"     🏗️  Base image: {base_image}")
+
+            start_time = time.time()
             try:
-                res = self.bakery.bake()
-                if res.baking_stats:
-                    echo(
-                        "baked image in: %s milliseconds"
-                        % res.baking_stats.solver_stats.duration_ms
-                    )
+                res = bakery.bake()
+                # TODO: Get actual bake time from bakery
+                bake_time = time.time() - start_time
+
+                with logger_lock:
+                    self.logger(f"🏁 Baked [{ref}] in {bake_time:.2f} seconds!")
+                self.images_baked += 1
                 return res
             except FastBakeryException as ex:
-                raise DockerEnvironmentException(str(ex))
+                raise DockerEnvironmentException(f"Bake [{ref}] failed: {str(ex)}")
 
         def prepare_step(step):
             base_image = next(
@@ -216,10 +267,15 @@ class DockerEnvironment(MetaflowEnvironment):
             }
 
         with ThreadPoolExecutor() as executor:
-            return {
-                step.name: _cached_bake(**args)
-                for step, args in zip(steps, executor.map(prepare_step, steps))
-            }
+            prepared_args = list(executor.map(prepare_step, steps))
+            for i, args in enumerate(prepared_args, 1):
+                args["ref"] = f"#{i:02d}"
+            futures = [executor.submit(_cached_bake, **args) for args in prepared_args]
+            results = {}
+            for step, future in zip(steps, futures):
+                results[step.name] = future.result()
+
+            return results
 
     def executable(self, step_name, default=None):
         if step_name in self.skipped_steps:

{ob-metaflow-extensions-1.1.86 → ob-metaflow-extensions-1.1.88}/metaflow_extensions/outerbounds/plugins/fast_bakery/fast_bakery.py

@@ -1,5 +1,6 @@
 from typing import Dict, Optional
 import requests
+import time
 
 
 class FastBakeryException(Exception):
@@ -140,6 +141,19 @@ class FastBakery:
             headers = {**self.headers, **(SERVICE_HEADERS or {})}
         except ImportError:
             headers = self.headers
+
+        retryable_status_codes = [409]
+
+        for attempt in range(2):  # 0 = initial attempt, 1-2 = retries
+            response = requests.post(self.url, json=payload, headers=headers)
+
+            if response.status_code not in retryable_status_codes:
+                break
+
+            if attempt < 2:  # Don't sleep after the last attempt
+                sleep_time = 0.5 * (attempt + 1)
+                time.sleep(sleep_time)
+
         response = requests.post(self.url, json=payload, headers=headers)
         self._handle_error_response(response)
         return FastBakeryApiResponse(response.json())

{ob-metaflow-extensions-1.1.86 → ob-metaflow-extensions-1.1.88}/ob_metaflow_extensions.egg-info/PKG-INFO

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

{ob-metaflow-extensions-1.1.86 → ob-metaflow-extensions-1.1.88}/setup.py

@@ -2,7 +2,7 @@ from setuptools import setup, find_namespace_packages
 from pathlib import Path
 
 
-version = "1.1.86"
+version = "1.1.88"
 this_directory = Path(__file__).parent
 long_description = (this_directory / "README.md").read_text()