ob-metaflow-extensions 1.2.9rc0__py2.py3-none-any.whl → 1.2.10__py2.py3-none-any.whl

metaflow_extensions/outerbounds/plugins/aws/assume_role_decorator.py

@@ -1,10 +1,12 @@
-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 metaflow.user_configs.config_decorators import (
+    MutableFlow,
+    MutableStep,
+    CustomFlowDecorator,
+)
 from .assume_role import OBP_ASSUME_ROLE_ARN_ENV_VAR
 
 
-class assume_role(FlowMutator):
+class assume_role(CustomFlowDecorator):
     """
     Flow-level decorator for assuming AWS IAM roles.
 
@@ -40,7 +42,7 @@ class assume_role(FlowMutator):
                 "`role_arn` must be a valid AWS IAM role ARN starting with 'arn:aws:iam::'"
             )
 
-    def pre_mutate(self, mutable_flow: MutableFlow) -> None:
+    def evaluate(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
@@ -49,29 +51,14 @@ class assume_role(FlowMutator):
         # 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)
+            # We'll inject the role assumption by adding an environment decorator
+            # The role will be available through an environment variable
+            step.add_decorator(
+                environment, vars={OBP_ASSUME_ROLE_ARN_ENV_VAR: self.role_arn}
+            )
 
         # Apply the role assumption setup to all steps in the flow
         for _, step in mutable_flow.steps:

metaflow_extensions/outerbounds/plugins/checkpoint_datastores/coreweave.py

@@ -1,11 +1,12 @@
-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
+from metaflow.user_configs.config_decorators import (
+    MutableFlow,
+    MutableStep,
+    CustomFlowDecorator,
+)
 import os
 
 
-class coreweave_checkpoints(_ExternalCheckpointFlowDeco):
+class coreweave_checkpoints(CustomFlowDecorator):
 
     """
 
@@ -45,14 +46,78 @@ class coreweave_checkpoints(_ExternalCheckpointFlowDeco):
         super().__init__(*args, **kwargs)
 
     def init(self, *args, **kwargs):
-        super().init(*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"
+            )
+
         self.coreweave_endpoint_url = f"https://cwobject.com"
+        if self.secrets is None:
+            raise ValueError(
+                "`secrets` keyword argument is required for the coreweave_datastore"
+            )
 
-    def pre_mutate(self, mutable_flow: MutableFlow) -> None:
+    def evaluate(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.decorators:
+                if d.name in swapping_decos:
+                    decos_to_add.append((d.name, d.attributes))
+                elif d.name == "secrets":
+                    already_has_secrets = True
+                    secrets_present_in_deco.extend(d.attributes["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,
+                    sources=secrets_to_add,
+                )
+                return
+
+            for d, _ in decos_to_add:
+                step.remove_decorator(d)
+
+            step.add_decorator(
+                secrets,
+                sources=secrets_to_add,
+            )
+            for d, attrs in decos_to_add:
+                _deco_to_add = swapping_decos[d]
+                step.add_decorator(_deco_to_add, **attrs)
+
         def _coreweave_config():
             return {
                 "root": self.bucket_path,
@@ -63,8 +128,12 @@ class coreweave_checkpoints(_ExternalCheckpointFlowDeco):
                     "config": dict(s3={"addressing_style": "virtual"}),
                 },
             }
+
         mutable_flow.add_decorator(
             with_artifact_store,
-            deco_kwargs=dict(type="coreweave", config=_coreweave_config),
+            type="coreweave",
+            config=_coreweave_config,
         )
-        self._swap_secrets(mutable_flow)
+
+        for step_name, step in mutable_flow.steps:
+            _add_secrets(step)

metaflow_extensions/outerbounds/plugins/checkpoint_datastores/nebius.py

@@ -1,11 +1,14 @@
-from metaflow.user_decorators.mutable_flow import MutableFlow
-from .external_chckpt import _ExternalCheckpointFlowDeco
+from metaflow.user_configs.config_decorators import (
+    MutableFlow,
+    MutableStep,
+    CustomFlowDecorator,
+)
 import os
 
 NEBIUS_ENDPOINT_URL = "https://storage.eu-north1.nebius.cloud:443"
 
 
-class nebius_checkpoints(_ExternalCheckpointFlowDeco):
+class nebius_checkpoints(CustomFlowDecorator):
 
     """
 
@@ -49,14 +52,78 @@ class nebius_checkpoints(_ExternalCheckpointFlowDeco):
         super().__init__(*args, **kwargs)
 
     def init(self, *args, **kwargs):
-        super().init(*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"
+            )
+
         self.nebius_endpoint_url = kwargs.get("endpoint_url", NEBIUS_ENDPOINT_URL)
+        if self.secrets is None:
+            raise ValueError(
+                "`secrets` keyword argument is required for the coreweave_datastore"
+            )
 
-    def pre_mutate(self, mutable_flow: MutableFlow) -> None:
+    def evaluate(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.decorators:
+                if d.name in swapping_decos:
+                    decos_to_add.append((d.name, d.attributes))
+                elif d.name == "secrets":
+                    already_has_secrets = True
+                    secrets_present_in_deco.extend(d.attributes["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,
+                    sources=secrets_to_add,
+                )
+                return
+
+            for d, _ in decos_to_add:
+                step.remove_decorator(d)
+
+            step.add_decorator(
+                secrets,
+                sources=secrets_to_add,
+            )
+            for d, attrs in decos_to_add:
+                _deco_to_add = swapping_decos[d]
+                step.add_decorator(_deco_to_add, **attrs)
+
         def _nebius_config():
             return {
                 "root": self.bucket_path,
@@ -68,6 +135,10 @@ class nebius_checkpoints(_ExternalCheckpointFlowDeco):
             }
 
         mutable_flow.add_decorator(
-            with_artifact_store, deco_kwargs=dict(type="s3", config=_nebius_config)
+            with_artifact_store,
+            type="s3",
+            config=_nebius_config,
         )
-        self._swap_secrets(mutable_flow)
+
+        for step_name, step in mutable_flow.steps:
+            _add_secrets(step)

metaflow_extensions/outerbounds/plugins/fast_bakery/docker_environment.py

@@ -351,16 +351,12 @@ class DockerEnvironment(MetaflowEnvironment):
         config.append("--disable=F0401")
         return config
 
-    def get_package_commands(
-        self, codepackage_url, datastore_type, code_package_metadata=None
-    ):
+    def get_package_commands(self, codepackage_url, datastore_type):
         # we must set the skip install flag at this stage in order to skip package downloads,
         # doing so in bootstrap_commands is too late in the lifecycle.
         return [
             "export METAFLOW_SKIP_INSTALL_DEPENDENCIES=$FASTBAKERY_IMAGE",
-        ] + super().get_package_commands(
-            codepackage_url, datastore_type, code_package_metadata=code_package_metadata
-        )
+        ] + super().get_package_commands(codepackage_url, datastore_type)
 
     def bootstrap_commands(self, step_name, datastore_type):
         if step_name in self.skipped_steps:

metaflow_extensions/outerbounds/plugins/fast_bakery/fast_bakery.py

@@ -122,6 +122,7 @@ class FastBakery:
             "responseMaxAgeSeconds": 0,
             "layerMaxAgeSeconds": 0,
             "baseImageMaxAgeSeconds": 0,
+            "overwriteExistingLayers": True,  # Used primarily to rewrite possibly corrupted layers.
         }
         return self
 

{ob_metaflow_extensions-1.2.9rc0.dist-info → ob_metaflow_extensions-1.2.10.dist-info}/METADATA

@@ -1,13 +1,13 @@
 Metadata-Version: 2.1
 Name: ob-metaflow-extensions
-Version: 1.2.9rc0
+Version: 1.2.10
 Summary: Outerbounds Platform Extensions for Metaflow
 Author: Outerbounds, Inc.
 License: Commercial
 Description-Content-Type: text/markdown
 Requires-Dist: boto3
 Requires-Dist: kubernetes
-Requires-Dist: ob-metaflow (==2.16.6.5)
+Requires-Dist: ob-metaflow (==2.15.21.5)
 
 # Outerbounds platform package
 

{ob_metaflow_extensions-1.2.9rc0.dist-info → ob_metaflow_extensions-1.2.10.dist-info}/RECORD

@@ -41,19 +41,18 @@ metaflow_extensions/outerbounds/plugins/apps/core/config/unified_config.py,sha25
 metaflow_extensions/outerbounds/plugins/apps/core/experimental/__init__.py,sha256=rd4qGTkHndKYfJmoAKZWiY0KK4j5BK6RBrtle-it1Mg,2746
 metaflow_extensions/outerbounds/plugins/aws/__init__.py,sha256=VBGdjNKeFLXGZuqh4jVk8cFtO1AWof73a6k_cnbAOYA,145
 metaflow_extensions/outerbounds/plugins/aws/assume_role.py,sha256=mBewNlnSYsR2rFXFkX-DUH6ku01h2yOcMcLHoCL7eyI,161
-metaflow_extensions/outerbounds/plugins/aws/assume_role_decorator.py,sha256=MrExIdwnQ4ZVKbrWqEQAIuZO4lWEy1rE0ihGUx7GMJA,2976
+metaflow_extensions/outerbounds/plugins/aws/assume_role_decorator.py,sha256=b3te8-RJxsj5cqXbQTFxxeIfnqTERAFInzLgXLVLyYA,2391
 metaflow_extensions/outerbounds/plugins/card_utilities/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 metaflow_extensions/outerbounds/plugins/card_utilities/async_cards.py,sha256=6rQhtZXK5DenXPfCRS1ul0jvLJYd48jrJAlnodID21w,4347
 metaflow_extensions/outerbounds/plugins/card_utilities/extra_components.py,sha256=D8rgCaSc7PuLD0MHJjqsjN0g0PQMN1H-ySOJqi5uIOE,19111
 metaflow_extensions/outerbounds/plugins/card_utilities/injector.py,sha256=UQtHrviFs7Z5LBh0nbaIv8aEnHnyXFjs9TiGCjtcIWc,2452
 metaflow_extensions/outerbounds/plugins/checkpoint_datastores/__init__.py,sha256=H96Bdv2XKTewKIU587sjNbU538SKMZOYWQ2qCp1JaNU,84
-metaflow_extensions/outerbounds/plugins/checkpoint_datastores/coreweave.py,sha256=r_Vgy1YexaGmM3ROl2EqzNpvZjerxwaIKV5YwlO8Q-g,2339
-metaflow_extensions/outerbounds/plugins/checkpoint_datastores/external_chckpt.py,sha256=BRigyuyYYv5H1CxfifFJPPvgMdQeCNIT8Hm0j1xo-YA,3025
-metaflow_extensions/outerbounds/plugins/checkpoint_datastores/nebius.py,sha256=xFGdSxlHcwZwgTSKj8efEv6ydCWYqRGrreGqXhf02ko,2321
+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=vKN3Waz0SgbOALVTRcf_8MCGhxdRueRFZJ7JmWbRJRM,15487
-metaflow_extensions/outerbounds/plugins/fast_bakery/fast_bakery.py,sha256=PE81ZB54OAMXkMGSB7JqgvgMg7N9kvoVclrWL-6jc2U,5626
+metaflow_extensions/outerbounds/plugins/fast_bakery/docker_environment.py,sha256=nmp_INGIAiWyrhyJ71BH38eRLu1xCIEEKejmXNQ6RlA,15378
+metaflow_extensions/outerbounds/plugins/fast_bakery/fast_bakery.py,sha256=9iqIBowKk95vUYvnMrFNWmmSuK-whKnqJwcCmEcRjTU,5727
 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
@@ -117,7 +116,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.2.9rc0.dist-info/METADATA,sha256=cVYNeLe2LkENesJ85BCscHX_5r8HY5etza4fM602xqE,521
-ob_metaflow_extensions-1.2.9rc0.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
-ob_metaflow_extensions-1.2.9rc0.dist-info/top_level.txt,sha256=NwG0ukwjygtanDETyp_BUdtYtqIA_lOjzFFh1TsnxvI,20
-ob_metaflow_extensions-1.2.9rc0.dist-info/RECORD,,
+ob_metaflow_extensions-1.2.10.dist-info/METADATA,sha256=5BoeRaxkfZYm-hBikHcUdYdH4dF-PWkG5MHnJ4fosT4,520
+ob_metaflow_extensions-1.2.10.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
+ob_metaflow_extensions-1.2.10.dist-info/top_level.txt,sha256=NwG0ukwjygtanDETyp_BUdtYtqIA_lOjzFFh1TsnxvI,20
+ob_metaflow_extensions-1.2.10.dist-info/RECORD,,

metaflow_extensions/outerbounds/plugins/checkpoint_datastores/external_chckpt.py

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