ob-metaflow-extensions 1.2.8__tar.gz → 1.2.9rc1__tar.gz

{ob_metaflow_extensions-1.2.8 → ob_metaflow_extensions-1.2.9rc1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ob_metaflow_extensions
-Version: 1.2.8
+Version: 1.2.9rc1
 Summary: Outerbounds Platform Extensions for Metaflow
 Author: Outerbounds, Inc.
 License: Commercial

{ob_metaflow_extensions-1.2.8 → ob_metaflow_extensions-1.2.9rc1}/metaflow_extensions/outerbounds/plugins/aws/assume_role_decorator.py

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

ob_metaflow_extensions-1.2.9rc1/metaflow_extensions/outerbounds/plugins/checkpoint_datastores/coreweave.py

@@ -0,0 +1,70 @@
+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
+import os
+
+
+class coreweave_checkpoints(_ExternalCheckpointFlowDeco):
+
+    """
+
+    This decorator is used for setting the coreweave object store as the artifact store for checkpoints/models created by the flow.
+
+    Parameters
+    ----------
+    secrets: list
+        A list of secrets to be added to the step. These secrets should contain any secrets that are required globally and the secret
+        for the coreweave object store. The secret should contain the following keys:
+        - COREWEAVE_ACCESS_KEY
+        - COREWEAVE_SECRET_KEY
+
+    bucket_path: str
+        The path to the bucket to store the checkpoints/models.
+
+    Usage
+    -----
+    ```python
+    from metaflow import checkpoint, step, FlowSpec, coreweave_checkpoints
+
+    @coreweave_checkpoints(secrets=[], bucket_path=None)
+    class MyFlow(FlowSpec):
+        @checkpoint
+        @step
+        def start(self):
+            # Saves the checkpoint in the coreweave object store
+            current.checkpoint.save("./foo.txt")
+
+        @step
+        def end(self):
+            pass
+    ```
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def init(self, *args, **kwargs):
+        super().init(*args, **kwargs)
+        self.coreweave_endpoint_url = f"https://cwobject.com"
+
+    def pre_mutate(self, mutable_flow: MutableFlow) -> None:
+        from metaflow import (
+            with_artifact_store,
+        )
+
+        def _coreweave_config():
+            return {
+                "root": self.bucket_path,
+                "client_params": {
+                    "aws_access_key_id": os.environ.get("COREWEAVE_ACCESS_KEY"),
+                    "aws_secret_access_key": os.environ.get("COREWEAVE_SECRET_KEY"),
+                    "endpoint_url": self.coreweave_endpoint_url,
+                    "config": dict(s3={"addressing_style": "virtual"}),
+                },
+            }
+        mutable_flow.add_decorator(
+            with_artifact_store,
+            deco_kwargs=dict(type="coreweave", config=_coreweave_config),
+        )
+        self._swap_secrets(mutable_flow)

ob_metaflow_extensions-1.2.9rc1/metaflow_extensions/outerbounds/plugins/checkpoint_datastores/external_chckpt.py

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

ob_metaflow_extensions-1.2.9rc1/metaflow_extensions/outerbounds/plugins/checkpoint_datastores/nebius.py

@@ -0,0 +1,73 @@
+from metaflow.user_decorators.mutable_flow import MutableFlow
+from .external_chckpt import _ExternalCheckpointFlowDeco
+import os
+
+NEBIUS_ENDPOINT_URL = "https://storage.eu-north1.nebius.cloud:443"
+
+
+class nebius_checkpoints(_ExternalCheckpointFlowDeco):
+
+    """
+
+    This decorator is used for setting the nebius's S3 compatible object store as the artifact store for
+    checkpoints/models created by the flow.
+
+    Parameters
+    ----------
+    secrets: list
+        A list of secrets to be added to the step. These secrets should contain any secrets that are required globally and the secret
+        for the nebius object store. The secret should contain the following keys:
+        - NEBIUS_ACCESS_KEY
+        - NEBIUS_SECRET_KEY
+
+    bucket_path: str
+        The path to the bucket to store the checkpoints/models.
+
+    endpoint_url: str
+        The endpoint url for the nebius object store. Defaults to `https://storage.eu-north1.nebius.cloud:443`
+
+    Usage
+    -----
+    ```python
+    from metaflow import checkpoint, step, FlowSpec, nebius_checkpoints
+
+    @nebius_checkpoints(secrets=[], bucket_path=None)
+    class MyFlow(FlowSpec):
+        @checkpoint
+        @step
+        def start(self):
+            # Saves the checkpoint in the nebius object store
+            current.checkpoint.save("./foo.txt")
+
+        @step
+        def end(self):
+            pass
+    ```
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def init(self, *args, **kwargs):
+        super().init(*args, **kwargs)
+        self.nebius_endpoint_url = kwargs.get("endpoint_url", NEBIUS_ENDPOINT_URL)
+
+    def pre_mutate(self, mutable_flow: MutableFlow) -> None:
+        from metaflow import (
+            with_artifact_store,
+        )
+
+        def _nebius_config():
+            return {
+                "root": self.bucket_path,
+                "client_params": {
+                    "aws_access_key_id": os.environ.get("NEBIUS_ACCESS_KEY"),
+                    "aws_secret_access_key": os.environ.get("NEBIUS_SECRET_KEY"),
+                    "endpoint_url": self.nebius_endpoint_url,
+                },
+            }
+
+        mutable_flow.add_decorator(
+            with_artifact_store, deco_kwargs=dict(type="s3", config=_nebius_config)
+        )
+        self._swap_secrets(mutable_flow)

{ob_metaflow_extensions-1.2.8 → ob_metaflow_extensions-1.2.9rc1}/metaflow_extensions/outerbounds/plugins/fast_bakery/docker_environment.py

@@ -351,12 +351,16 @@ class DockerEnvironment(MetaflowEnvironment):
         config.append("--disable=F0401")
         return config
 
-    def get_package_commands(self, codepackage_url, datastore_type):
+    def get_package_commands(
+        self, codepackage_url, datastore_type, code_package_metadata=None
+    ):
         # 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)
+        ] + super().get_package_commands(
+            codepackage_url, datastore_type, code_package_metadata=code_package_metadata
+        )
 
     def bootstrap_commands(self, step_name, datastore_type):
         if step_name in self.skipped_steps:

{ob_metaflow_extensions-1.2.8 → ob_metaflow_extensions-1.2.9rc1}/ob_metaflow_extensions.egg-info/PKG-INFO

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

{ob_metaflow_extensions-1.2.8 → ob_metaflow_extensions-1.2.9rc1}/ob_metaflow_extensions.egg-info/SOURCES.txt

@@ -51,6 +51,7 @@ metaflow_extensions/outerbounds/plugins/card_utilities/extra_components.py
 metaflow_extensions/outerbounds/plugins/card_utilities/injector.py
 metaflow_extensions/outerbounds/plugins/checkpoint_datastores/__init__.py
 metaflow_extensions/outerbounds/plugins/checkpoint_datastores/coreweave.py
+metaflow_extensions/outerbounds/plugins/checkpoint_datastores/external_chckpt.py
 metaflow_extensions/outerbounds/plugins/checkpoint_datastores/nebius.py
 metaflow_extensions/outerbounds/plugins/fast_bakery/__init__.py
 metaflow_extensions/outerbounds/plugins/fast_bakery/baker.py

ob_metaflow_extensions-1.2.9rc1/ob_metaflow_extensions.egg-info/requires.txt

@@ -0,0 +1,3 @@
+boto3
+kubernetes
+ob-metaflow==2.16.5.1

{ob_metaflow_extensions-1.2.8 → ob_metaflow_extensions-1.2.9rc1}/setup.py

@@ -2,7 +2,7 @@ from setuptools import setup, find_namespace_packages
 from pathlib import Path
 
 
-version = "1.2.8"
+version = "1.2.9rc1"
 this_directory = Path(__file__).parent
 long_description = (this_directory / "README.md").read_text()
 
@@ -19,6 +19,6 @@ setup(
     install_requires=[
         "boto3",
         "kubernetes",
-        "ob-metaflow == 2.15.21.4",
+        "ob-metaflow == 2.16.5.1",
     ],
 )

ob_metaflow_extensions-1.2.8/metaflow_extensions/outerbounds/plugins/checkpoint_datastores/coreweave.py

@@ -1,139 +0,0 @@
-from metaflow.user_configs.config_decorators import (
-    MutableFlow,
-    MutableStep,
-    CustomFlowDecorator,
-)
-import os
-
-
-class coreweave_checkpoints(CustomFlowDecorator):
-
-    """
-
-    This decorator is used for setting the coreweave object store as the artifact store for checkpoints/models created by the flow.
-
-    Parameters
-    ----------
-    secrets: list
-        A list of secrets to be added to the step. These secrets should contain any secrets that are required globally and the secret
-        for the coreweave object store. The secret should contain the following keys:
-        - COREWEAVE_ACCESS_KEY
-        - COREWEAVE_SECRET_KEY
-
-    bucket_path: str
-        The path to the bucket to store the checkpoints/models.
-
-    Usage
-    -----
-    ```python
-    from metaflow import checkpoint, step, FlowSpec, coreweave_checkpoints
-
-    @coreweave_checkpoints(secrets=[], bucket_path=None)
-    class MyFlow(FlowSpec):
-        @checkpoint
-        @step
-        def start(self):
-            # Saves the checkpoint in the coreweave object store
-            current.checkpoint.save("./foo.txt")
-
-        @step
-        def end(self):
-            pass
-    ```
-    """
-
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-
-    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"
-            )
-
-        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 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,
-                "client_params": {
-                    "aws_access_key_id": os.environ.get("COREWEAVE_ACCESS_KEY"),
-                    "aws_secret_access_key": os.environ.get("COREWEAVE_SECRET_KEY"),
-                    "endpoint_url": self.coreweave_endpoint_url,
-                    "config": dict(s3={"addressing_style": "virtual"}),
-                },
-            }
-
-        mutable_flow.add_decorator(
-            with_artifact_store,
-            type="coreweave",
-            config=_coreweave_config,
-        )
-
-        for step_name, step in mutable_flow.steps:
-            _add_secrets(step)

ob_metaflow_extensions-1.2.8/metaflow_extensions/outerbounds/plugins/checkpoint_datastores/nebius.py

@@ -1,144 +0,0 @@
-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(CustomFlowDecorator):
-
-    """
-
-    This decorator is used for setting the nebius's S3 compatible object store as the artifact store for
-    checkpoints/models created by the flow.
-
-    Parameters
-    ----------
-    secrets: list
-        A list of secrets to be added to the step. These secrets should contain any secrets that are required globally and the secret
-        for the nebius object store. The secret should contain the following keys:
-        - NEBIUS_ACCESS_KEY
-        - NEBIUS_SECRET_KEY
-
-    bucket_path: str
-        The path to the bucket to store the checkpoints/models.
-
-    endpoint_url: str
-        The endpoint url for the nebius object store. Defaults to `https://storage.eu-north1.nebius.cloud:443`
-
-    Usage
-    -----
-    ```python
-    from metaflow import checkpoint, step, FlowSpec, nebius_checkpoints
-
-    @nebius_checkpoints(secrets=[], bucket_path=None)
-    class MyFlow(FlowSpec):
-        @checkpoint
-        @step
-        def start(self):
-            # Saves the checkpoint in the nebius object store
-            current.checkpoint.save("./foo.txt")
-
-        @step
-        def end(self):
-            pass
-    ```
-    """
-
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-
-    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"
-            )
-
-        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 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,
-                "client_params": {
-                    "aws_access_key_id": os.environ.get("NEBIUS_ACCESS_KEY"),
-                    "aws_secret_access_key": os.environ.get("NEBIUS_SECRET_KEY"),
-                    "endpoint_url": self.nebius_endpoint_url,
-                },
-            }
-
-        mutable_flow.add_decorator(
-            with_artifact_store,
-            type="s3",
-            config=_nebius_config,
-        )
-
-        for step_name, step in mutable_flow.steps:
-            _add_secrets(step)

ob_metaflow_extensions-1.2.8/ob_metaflow_extensions.egg-info/requires.txt

@@ -1,3 +0,0 @@
-boto3
-kubernetes
-ob-metaflow==2.15.21.4