ob-metaflow-extensions 1.1.168rc5__py2.py3-none-any.whl → 1.1.169__py2.py3-none-any.whl

metaflow_extensions/outerbounds/plugins/__init__.py

@@ -52,10 +52,6 @@ def get_boto3_session(role_arn=None, session_vars=None):
 
     token_info = get_token("/generate/aws")
 
-    # Check if the assume_role decorator has set a role ARN via environment variable
-    # This takes precedence over CSPR role
-    decorator_role_arn = os.environ.get(OBP_ASSUME_ROLE_ARN_ENV_VAR)
-
     # Write token to a file. The file name is derived from the user name
     # so it works with multiple users on the same machine.
     #
@@ -76,13 +72,18 @@ def get_boto3_session(role_arn=None, session_vars=None):
     if token_info.get("cspr_role_arn"):
         cspr_role = token_info["cspr_role_arn"]
 
-    # If assume_role decorator is used, prioritize it over CSPR role
-    effective_role = decorator_role_arn or cspr_role
-    if effective_role:
-        # If we have either a decorator role or CSPR role, set up AWS config
+    # Check if the assume_role decorator has set a CSPR ARN via environment variable
+    # This takes precedence over CSPR role that comes from the token_info response
+    decorator_role_arn = os.environ.get(OBP_ASSUME_ROLE_ARN_ENV_VAR)
+    if decorator_role_arn:
+        cspr_role = decorator_role_arn
+
+    if cspr_role:
+        # If CSPR role is set, we set it as the default role to assume
+        # for the AWS SDK. We do this by writing an AWS config file
         # with two profiles. One to get credentials for the task role
         # in exchange for the OIDC token, and second to assume the
-        # effective role using the task role credentials.
+        # CSPR role using the task role credentials.
         import configparser
         from io import StringIO
 
@@ -94,9 +95,9 @@ def get_boto3_session(role_arn=None, session_vars=None):
             "web_identity_token_file": token_file,
         }
 
-        # Effective role profile (decorator role or CSPR role)
-        aws_config["profile effective"] = {
-            "role_arn": effective_role,
+        # CSPR role profile (default)
+        aws_config["profile cspr"] = {
+            "role_arn": cspr_role,
             "source_profile": "task",
         }
 
@@ -112,7 +113,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_PROFILE"] = "effective"
+        os.environ["AWS_PROFILE"] = "cspr"
     else:
         os.environ["AWS_WEB_IDENTITY_TOKEN_FILE"] = token_file
         os.environ["AWS_ROLE_ARN"] = token_info["role_arn"]
@@ -124,29 +125,21 @@ def get_boto3_session(role_arn=None, session_vars=None):
     if token_info.get("region"):
         os.environ["AWS_DEFAULT_REGION"] = token_info["region"]
 
-    if effective_role:
-        if role_arn == decorator_role_arn or role_arn == USE_CSPR_ROLE_ARN_IF_SET:
-            # We have either a decorator role or CSPR role, use the effective profile
-            # The generated AWS config will be used here since we set the
-            # AWS_CONFIG_FILE environment variable above.
-            session = boto3.session.Session(profile_name="effective")
+    if cspr_role:
+        #  The generated AWS config will be used here since we set the
+        #  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="cspr")
         else:
             session = boto3.session.Session(profile_name="task")
     else:
-        # No decorator role or CSPR role, use default session
         # Not using AWS config, just AWS_WEB_IDENTITY_TOKEN_FILE + AWS_ROLE_ARN
         session = boto3.session.Session()
 
-    if (
-        role_arn
-        and role_arn != USE_CSPR_ROLE_ARN_IF_SET
-        and role_arn != decorator_role_arn
-    ):
-        # If the user provided a role_arn that's different from the decorator role,
-        # we assume that role using the current session credentials.
-        # This works for both cases:
-        # 1. No decorator role: Task role -> Secrets role
-        # 2. With decorator role: Decorator role -> Secrets role
+    if role_arn and role_arn != USE_CSPR_ROLE_ARN_IF_SET:
+        # If the user provided a role_arn, we assume that role
+        # using the task role credentials. CSPR role is not used.
         fetcher = botocore.credentials.AssumeRoleCredentialFetcher(
             client_creator=session._session.create_client,
             source_credentials=session._session.get_credentials(),
@@ -162,8 +155,8 @@ def get_boto3_session(role_arn=None, session_vars=None):
     else:
         # If the user didn't provide a role_arn, or if the role_arn
         # is set to USE_CSPR_ROLE_ARN_IF_SET, we return the default
-        # session which would use the effective role (decorator or CSPR) if set,
-        # or the task role otherwise.
+        # session which would use the CSPR role if it is set on the
+        # server, and the task role otherwise.
         return session
 
 
@@ -331,7 +324,6 @@ CLIS_DESC = [
     ("nvct", ".nvct.nvct_cli.cli"),
     ("fast-bakery", ".fast_bakery.fast_bakery_cli.cli"),
     ("snowpark", ".snowpark.snowpark_cli.cli"),
-    ("app", ".apps.app_cli.cli"),
 ]
 STEP_DECORATORS_DESC = [
     ("nvidia", ".nvcf.nvcf_decorator.NvcfDecorator"),

metaflow_extensions/outerbounds/plugins/kubernetes/pod_killer.py

@@ -39,7 +39,7 @@ def derive_job_outcome(job_status: "V1JobStatus"):
 
     # This means that the job has neither finished or succedded.
     if job_status.active:
-        return JobOutcomes.KILL
+        return JobOutcomes.DELETE
 
     # This means that the job is not active. Had started. There is not succedded/fail.
     # This is a weird state. Better to just kill the job
@@ -47,7 +47,7 @@ def derive_job_outcome(job_status: "V1JobStatus"):
 
 
 class PodKiller:
-    def __init__(self, kubernetes_client, echo_func, namespace):
+    def __init__(self, kubernetes_client, echo_func, namespace, progress_bar=None):
         self.client = kubernetes_client
         self.echo = echo_func
         self.api_instance = self.client.CoreV1Api()
@@ -55,6 +55,7 @@ class PodKiller:
         self._namespace = namespace
         self.jobset_api = None
         self.jobset_api = self.client.CustomObjectsApi()
+        self.progress_bar = progress_bar
 
     def _delete_jobset(self, owner_ref, namespace):
         """Delete a JobSet if it's the owner of a job."""
@@ -147,20 +148,31 @@ class PodKiller:
 
     def _find_matching_jobs(self, flow_name, run_id=None, user=None):
         """Find jobs that match the flow_name, run_id, and user criteria using similar logic to _find_active_pods"""
+
+        def paginated_job_finder(namespace):
+            continue_token = None
+            while True:
+                response = self.job_api.list_namespaced_job(
+                    namespace=namespace, limit=100, _continue=continue_token
+                )
+                yield response.items
+                continue_token = response.metadata._continue
+                if not continue_token:
+                    break
+
         try:
-            jobs = self.job_api.list_namespaced_job(namespace=self._namespace)
             matching_jobs = []
-            for _job in jobs.items:
-                job = _job.to_dict()
-                _match = self._metaflow_matching_spec(
-                    run_id=run_id,
-                    user=user,
-                    flow_name=flow_name,
-                    annotations=job.get("metadata", {}).get("annotations", {}),
-                    labels=job.get("metadata", {}).get("labels", {}),
-                )
-                if _match:
-                    matching_jobs.append(_job)
+            for _jobs in paginated_job_finder(self._namespace):
+                for job in _jobs:
+                    _match = self._metaflow_matching_spec(
+                        run_id=run_id,
+                        user=user,
+                        flow_name=flow_name,
+                        annotations=job.metadata.annotations,
+                        labels=job.metadata.labels,
+                    )
+                    if _match:
+                        matching_jobs.append(job)
             return matching_jobs
         except Exception as e:
             self.echo(f"Error finding jobs: {str(e)}\n")
@@ -171,25 +183,38 @@ class PodKiller:
         if not self.jobset_api:
             return []
 
+        def paginated_jobset_finder(namespace):
+            continue_token = None
+            responses = []
+            while True:
+                response = self.jobset_api.list_namespaced_custom_object(
+                    group="jobset.x-k8s.io",
+                    version="v1alpha2",
+                    namespace=namespace,
+                    plural="jobsets",
+                    limit=100,
+                    **({"_continue": continue_token} if continue_token else {}),
+                )
+                continue_token = response.get("metadata", {}).get("continue", None)
+                responses.append(response)
+                if not continue_token:
+                    break
+            return responses
+
         try:
-            jobsets = self.jobset_api.list_namespaced_custom_object(
-                group="jobset.x-k8s.io",
-                version="v1alpha2",
-                namespace=self._namespace,
-                plural="jobsets",
-            )
             matching_jobsets = []
 
-            for jobset in jobsets.get("items", []):
-                _match = self._metaflow_matching_spec(
-                    run_id=run_id,
-                    user=user,
-                    flow_name=flow_name,
-                    annotations=jobset.get("metadata", {}).get("annotations", {}),
-                    labels=jobset.get("metadata", {}).get("labels", {}),
-                )
-                if _match:
-                    matching_jobsets.append(jobset)
+            for jobset_response in paginated_jobset_finder(self._namespace):
+                for jobset in jobset_response.get("items", []):
+                    _match = self._metaflow_matching_spec(
+                        run_id=run_id,
+                        user=user,
+                        flow_name=flow_name,
+                        annotations=jobset.get("metadata", {}).get("annotations", {}),
+                        labels=jobset.get("metadata", {}).get("labels", {}),
+                    )
+                    if _match:
+                        matching_jobsets.append(jobset)
 
             return matching_jobsets
         except Exception as e:
@@ -270,9 +295,20 @@ class PodKiller:
             self.echo(f"Unknown outcome {outcome} for JobSet {jobset_name}")
             return False
 
+    def extract_matching_jobs_and_jobsets(self, flow_name, run_id, user):
+        """Extract matching jobs and jobsets based on the flow_name, run_id, and user criteria"""
+        jobs = self._find_matching_jobs(flow_name, run_id, user)
+        jobsets = self._find_matching_jobsets(flow_name, run_id, user)
+        return [(j, derive_job_outcome(j.status)) for j in jobs], [
+            (j, derive_jobset_outcome(j.get("status", {}))) for j in jobsets
+        ]
+
     def process_matching_jobs_and_jobsets(self, flow_name, run_id, user):
         """Process all matching jobs and jobsets based on their derived outcomes"""
         results = []
+        progress_update = lambda x: x
+        if self.progress_bar:
+            progress_update = lambda x: self.progress_bar.update(1, x)
 
         # Process matching jobs
         _jobs, _jobsets = [], []
@@ -282,6 +318,7 @@ class PodKiller:
             result = self._handle_job_outcome(job, outcome)
             # results.append(result)
             if result is not None:
+                progress_update("💀 Killing Job %s" % job.metadata.name)
                 results.append(result)
                 _jobs.append(result)
 
@@ -292,7 +329,46 @@ class PodKiller:
             outcome = derive_jobset_outcome(jobset_status)
             result = self._handle_jobset_outcome(jobset, outcome)
             if result is not None:
+                progress_update(
+                    "💀 Deleting JobSet %s"
+                    % jobset.get("metadata", {}).get("name", "unknown")
+                )
                 results.append(result)
                 _jobsets.append(result)
 
         return results, len(_jobs), len(_jobsets)
+
+    def process_matching_jobs_and_jobsets_force_all(self, flow_name, run_id, user):
+        """Force process ALL matching jobs and jobsets regardless of their status/outcome"""
+        results = []
+        progress_update = lambda x: x
+        if self.progress_bar:
+            progress_update = lambda x: self.progress_bar.update(1, x)
+
+        # Process matching jobs - FORCE DELETE ALL
+        _jobs, _jobsets = [], []
+        jobs = self._find_matching_jobs(flow_name, run_id, user)
+        for job in jobs:
+            # Force DELETE outcome regardless of actual status
+            result = self._handle_job_outcome(job, JobOutcomes.DELETE)
+            progress_update("🔥 FORCE Deleting Job %s" % job.metadata.name)
+            results.append(
+                result if result is not None else True
+            )  # Treat None as success for force mode
+            _jobs.append(result if result is not None else True)
+
+        # Process matching jobsets - FORCE DELETE ALL
+        jobsets = self._find_matching_jobsets(flow_name, run_id, user)
+        for jobset in jobsets:
+            # Force DELETE outcome regardless of actual status
+            result = self._handle_jobset_outcome(jobset, JobOutcomes.DELETE)
+            progress_update(
+                "🔥 FORCE Deleting JobSet %s"
+                % jobset.get("metadata", {}).get("name", "unknown")
+            )
+            results.append(
+                result if result is not None else True
+            )  # Treat None as success for force mode
+            _jobsets.append(result if result is not None else True)
+
+        return results, len(_jobs), len(_jobsets)

metaflow_extensions/outerbounds/toplevel/global_aliases_for_metaflow_package.py

@@ -75,4 +75,3 @@ from .. import profilers
 from ..plugins.snowflake import Snowflake
 from ..plugins.checkpoint_datastores import nebius_checkpoints, coreweave_checkpoints
 from ..plugins.aws import assume_role
-from . import ob_internal

metaflow_extensions/outerbounds/toplevel/ob_internal.py

@@ -1 +1 @@
-from ..plugins.fast_bakery.baker import bake_image
+from ..plugins.kubernetes.pod_killer import PodKiller

{ob_metaflow_extensions-1.1.168rc5.dist-info → ob_metaflow_extensions-1.1.169.dist-info}/METADATA

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

{ob_metaflow_extensions-1.1.168rc5.dist-info → ob_metaflow_extensions-1.1.169.dist-info}/RECORD

@@ -1,11 +1,10 @@
 metaflow_extensions/outerbounds/__init__.py,sha256=Gb8u06s9ClQsA_vzxmkCzuMnigPy7kKcDnLfb7eB-64,514
 metaflow_extensions/outerbounds/remote_config.py,sha256=pEFJuKDYs98eoB_-ryPjVi9b_c4gpHMdBHE14ltoxIU,4672
 metaflow_extensions/outerbounds/config/__init__.py,sha256=JsQGRuGFz28fQWjUvxUgR8EKBLGRdLUIk_buPLJplJY,1225
-metaflow_extensions/outerbounds/plugins/__init__.py,sha256=8NoHkOVUbd7YWK1S7GjAe7FzB_dTj-uPIZfla7-t-jE,14129
+metaflow_extensions/outerbounds/plugins/__init__.py,sha256=qaGCEa_QFWgGURABv-ss0TYbC1RJTLlckogTA5tc3Bk,13713
 metaflow_extensions/outerbounds/plugins/auth_server.py,sha256=_Q9_2EL0Xy77bCRphkwT1aSu8gQXRDOH-Z-RxTUO8N4,2202
 metaflow_extensions/outerbounds/plugins/perimeters.py,sha256=QXh3SFP7GQbS-RAIxUOPbhPzQ7KDFVxZkTdKqFKgXjI,2697
 metaflow_extensions/outerbounds/plugins/apps/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-metaflow_extensions/outerbounds/plugins/apps/app_cli.py,sha256=erfKCC6zKuwax0ye9j-tqIkAVdCg7MJfbRGhMhViSzU,575
 metaflow_extensions/outerbounds/plugins/apps/app_utils.py,sha256=sw9whU17lAzlD2K2kEDNjlk1Ib-2xE2UNhJkmzD8Qv8,8543
 metaflow_extensions/outerbounds/plugins/apps/consts.py,sha256=iHsyqbUg9k-rgswCs1Jxf5QZIxR1V-peCDRjgr9kdBM,177
 metaflow_extensions/outerbounds/plugins/apps/deploy_decorator.py,sha256=VkmiMdNYHhNdt-Qm9AVv7aE2LWFsIFEc16YcOYjwF6Q,8568
@@ -21,14 +20,13 @@ 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
 metaflow_extensions/outerbounds/plugins/fast_bakery/fast_bakery_decorator.py,sha256=MXSIp05-jvt8Q2uGaLKjtuM_ToLeRLxhtMbfHc9Kcko,1515
 metaflow_extensions/outerbounds/plugins/kubernetes/__init__.py,sha256=5zG8gShSj8m7rgF4xgWBZFuY3GDP5n1T0ktjRpGJLHA,69
 metaflow_extensions/outerbounds/plugins/kubernetes/kubernetes_client.py,sha256=sjBhQ4aa-i1UkKsJyTswdDLYOBAFIvHRco4r7wfs9Tc,5003
-metaflow_extensions/outerbounds/plugins/kubernetes/pod_killer.py,sha256=R483_Qg8eRifx5Z-3iHSQsFL6hUDmNTVLBlPWudCsQA,11404
+metaflow_extensions/outerbounds/plugins/kubernetes/pod_killer.py,sha256=3b9y6gYEqWkN-_-GOqcAV2pOJOeCjpVmfW-nOMCJ4Z0,14924
 metaflow_extensions/outerbounds/plugins/nim/card.py,sha256=dXOJvsZed5NyYyxYLPDvtwg9z_X4azL9HTJGYaiNriY,4690
 metaflow_extensions/outerbounds/plugins/nim/nim_decorator.py,sha256=50YVvC7mcZYlPluM0Wq1UtufhzlQb-RxzZkTOJJ3LkM,3439
 metaflow_extensions/outerbounds/plugins/nim/nim_manager.py,sha256=y8U71106KJtrC6nlhsNnzX9Xkv3RnyZ1KEpRFwqZZFk,13686
@@ -77,8 +75,8 @@ metaflow_extensions/outerbounds/plugins/vllm/vllm_manager.py,sha256=E6SCBoanNNsF
 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=_fKWv_-O1k5Nk5A1q05Ioh-PSsFXGL-jiAt7zfl8pIE,2999
-metaflow_extensions/outerbounds/toplevel/ob_internal.py,sha256=53xM6d_UYT3uGFFA59UzxN23H5QMO5_F39pALpmGy04,51
+metaflow_extensions/outerbounds/toplevel/global_aliases_for_metaflow_package.py,sha256=FS0ZKQJsJKlw9PgtLqVV4kH9o7_qwYcIVtyu2Kqwa_U,2973
+metaflow_extensions/outerbounds/toplevel/ob_internal.py,sha256=RtF7t9EIyGjWi36jAHbRaLo_MNOxF5Jo8X9QWzOr0co,54
 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
@@ -86,7 +84,7 @@ metaflow_extensions/outerbounds/toplevel/plugins/ollama/__init__.py,sha256=GRSz2
 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
 metaflow_extensions/outerbounds/toplevel/plugins/vllm/__init__.py,sha256=ekcgD3KVydf-a0xMI60P4uy6ePkSEoFHiGnDq1JM940,45
-ob_metaflow_extensions-1.1.168rc5.dist-info/METADATA,sha256=1a5w7uL3QYY5Oif4eylHkXQQO9cfaKqWZ59YGtXBWgo,524
-ob_metaflow_extensions-1.1.168rc5.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
-ob_metaflow_extensions-1.1.168rc5.dist-info/top_level.txt,sha256=NwG0ukwjygtanDETyp_BUdtYtqIA_lOjzFFh1TsnxvI,20
-ob_metaflow_extensions-1.1.168rc5.dist-info/RECORD,,
+ob_metaflow_extensions-1.1.169.dist-info/METADATA,sha256=x6w5D_39YQbFxSSjlZMRqvkO0M6eW4jzZzq-_5w3d4g,521
+ob_metaflow_extensions-1.1.169.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
+ob_metaflow_extensions-1.1.169.dist-info/top_level.txt,sha256=NwG0ukwjygtanDETyp_BUdtYtqIA_lOjzFFh1TsnxvI,20
+ob_metaflow_extensions-1.1.169.dist-info/RECORD,,

metaflow_extensions/outerbounds/plugins/apps/app_cli.py

@@ -1,26 +0,0 @@
-from metaflow._vendor import click
-
-OUTERBOUNDS_APP_CLI_AVAILABLE = True
-try:
-    import outerbounds.apps.app_cli as ob_apps_cli
-except ImportError:
-    OUTERBOUNDS_APP_CLI_AVAILABLE = False
-
-
-if not OUTERBOUNDS_APP_CLI_AVAILABLE:
-
-    @click.group()
-    def _cli():
-        pass
-
-    @_cli.group(help="Dummy Group to append to CLI for Safety")
-    def app():
-        pass
-
-    @app.command(help="Dummy Command to append to CLI for Safety")
-    def cannot_deploy():
-        raise Exception("Outerbounds App CLI not available")
-
-    cli = _cli
-else:
-    cli = ob_apps_cli.cli

metaflow_extensions/outerbounds/plugins/fast_bakery/baker.py

@@ -1,110 +0,0 @@
-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,
-    )