ob-metaflow-extensions 1.1.86__py2.py3-none-any.whl → 1.1.88__py2.py3-none-any.whl

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:

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:

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.dist-info → ob_metaflow_extensions-1.1.88.dist-info}/METADATA

@@ -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.dist-info → ob_metaflow_extensions-1.1.88.dist-info}/RECORD

@@ -1,12 +1,12 @@
 metaflow_extensions/outerbounds/__init__.py,sha256=TRGvIUMjkfneWtYUFSWoubu_Kf2ekAL4WLbV3IxOj9k,499
 metaflow_extensions/outerbounds/remote_config.py,sha256=Zpfpjgz68_ZgxlXezjzlsDLo4840rkWuZgwDB_5H57U,4059
 metaflow_extensions/outerbounds/config/__init__.py,sha256=JsQGRuGFz28fQWjUvxUgR8EKBLGRdLUIk_buPLJplJY,1225
-metaflow_extensions/outerbounds/plugins/__init__.py,sha256=g07Xj7YifwLfTTa94cyf3kLebHezeLWNG9HoGnpDyMo,12519
+metaflow_extensions/outerbounds/plugins/__init__.py,sha256=RNN9wxnIeNHO9S9lI-sgVINbAcO0C-29lYNd4BqNs-M,12508
 metaflow_extensions/outerbounds/plugins/auth_server.py,sha256=1v2GBqoMBxp5E7Lejz139w-jxJtPnLDvvHXP0HhEIHI,2361
 metaflow_extensions/outerbounds/plugins/perimeters.py,sha256=QXh3SFP7GQbS-RAIxUOPbhPzQ7KDFVxZkTdKqFKgXjI,2697
 metaflow_extensions/outerbounds/plugins/fast_bakery/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-metaflow_extensions/outerbounds/plugins/fast_bakery/docker_environment.py,sha256=P8AOdQYz6T0D4Uj0_hUF91MF0FkRHBU9P7Wgihouk20,9852
-metaflow_extensions/outerbounds/plugins/fast_bakery/fast_bakery.py,sha256=7ppPwnwb_b2bRdZnmxnZXsnDlysLe2Ah4o0vxJdrP6Q,4621
+metaflow_extensions/outerbounds/plugins/fast_bakery/docker_environment.py,sha256=9iokzw--nCpk-omGXDSB-oct7-57rCJbq62HKvmt2NA,12126
+metaflow_extensions/outerbounds/plugins/fast_bakery/fast_bakery.py,sha256=6OQdij49tjBwQuQ8Kdu16kfLBzKLGM4_f80KQm6WurQ,5069
 metaflow_extensions/outerbounds/plugins/fast_bakery/fast_bakery_cli.py,sha256=kqFyu2bJSnc9_9aYfBpz5xK6L6luWFZK_NMuh8f1eVk,1494
 metaflow_extensions/outerbounds/plugins/fast_bakery/fast_bakery_decorator.py,sha256=EZDbyrfZ7fgcU-P9dMS_hpCxsdDeUE0K5VU3uNM4aW4,1506
 metaflow_extensions/outerbounds/plugins/kubernetes/__init__.py,sha256=5zG8gShSj8m7rgF4xgWBZFuY3GDP5n1T0ktjRpGJLHA,69
@@ -33,7 +33,7 @@ metaflow_extensions/outerbounds/toplevel/global_aliases_for_metaflow_package.py,
 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
-ob_metaflow_extensions-1.1.86.dist-info/METADATA,sha256=LKi6v-1FMvAWE2u9eW2PKoVBi1xl6Wor1TjhwYVzPNk,520
-ob_metaflow_extensions-1.1.86.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
-ob_metaflow_extensions-1.1.86.dist-info/top_level.txt,sha256=NwG0ukwjygtanDETyp_BUdtYtqIA_lOjzFFh1TsnxvI,20
-ob_metaflow_extensions-1.1.86.dist-info/RECORD,,
+ob_metaflow_extensions-1.1.88.dist-info/METADATA,sha256=zgizoMrEL9ZXm9vj3WTY0K6y0DbQuoBkBuxiZGeXkCc,520
+ob_metaflow_extensions-1.1.88.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
+ob_metaflow_extensions-1.1.88.dist-info/top_level.txt,sha256=NwG0ukwjygtanDETyp_BUdtYtqIA_lOjzFFh1TsnxvI,20
+ob_metaflow_extensions-1.1.88.dist-info/RECORD,,