ob-metaflow-extensions 1.1.146__py2.py3-none-any.whl → 1.1.147__py2.py3-none-any.whl

metaflow_extensions/outerbounds/plugins/checkpoint_datastores/__init__.py

@@ -0,0 +1,2 @@
+from .nebius import nebius_checkpoints
+from .coreweave import coreweave_checkpoints

metaflow_extensions/outerbounds/plugins/checkpoint_datastores/coreweave.py

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

metaflow_extensions/outerbounds/plugins/checkpoint_datastores/nebius.py

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

metaflow_extensions/outerbounds/plugins/nvcf/nvcf_decorator.py

@@ -198,7 +198,12 @@ class NvcfDecorator(StepDecorator):
             meta["nvcf-nspectid"] = os.environ.get("NVCF_NSPECTID")
 
             entries = [
-                MetaDatum(field=k, value=v, type=k, tags=[])
+                MetaDatum(
+                    field=k,
+                    value=v,
+                    type=k,
+                    tags=["attempt_id:{0}".format(retry_count)],
+                )
                 for k, v in meta.items()
                 if v is not None
             ]

metaflow_extensions/outerbounds/toplevel/global_aliases_for_metaflow_package.py

@@ -52,3 +52,4 @@ def S3(*args, **kwargs):
 
 from .. import profilers
 from ..plugins.snowflake import Snowflake
+from ..plugins.checkpoint_datastores import nebius_checkpoints, coreweave_checkpoints

{ob_metaflow_extensions-1.1.146.dist-info → ob_metaflow_extensions-1.1.147.dist-info}/METADATA

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

{ob_metaflow_extensions-1.1.146.dist-info → ob_metaflow_extensions-1.1.147.dist-info}/RECORD

@@ -13,6 +13,9 @@ metaflow_extensions/outerbounds/plugins/card_utilities/__init__.py,sha256=47DEQp
 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=_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/docker_environment.py,sha256=Tl520HdBteg-aDOM7mnnJJpdDCZc49BmFFmLUc_vTi8,15018
 metaflow_extensions/outerbounds/plugins/fast_bakery/fast_bakery.py,sha256=PE81ZB54OAMXkMGSB7JqgvgMg7N9kvoVclrWL-6jc2U,5626
@@ -30,7 +33,7 @@ metaflow_extensions/outerbounds/plugins/nvcf/exceptions.py,sha256=-Pm9cOWUzpv94T
 metaflow_extensions/outerbounds/plugins/nvcf/heartbeat_store.py,sha256=pOWwm8LFQBbtku0zNBBwCyXxLK8U-hhC4naQcmU69nE,6217
 metaflow_extensions/outerbounds/plugins/nvcf/nvcf.py,sha256=3ZFdYItVpFWnHMOeyV1nslUyelfvX5rknh2d2IWxVws,15591
 metaflow_extensions/outerbounds/plugins/nvcf/nvcf_cli.py,sha256=3D-r5XO88Yh2k1EAZFJTe_PwdbhWp5qXflG8AgE4ZIU,9500
-metaflow_extensions/outerbounds/plugins/nvcf/nvcf_decorator.py,sha256=pxxNxW4bW3kbB6ybRam657GyKIhvIkMuidB94iFfCD8,9116
+metaflow_extensions/outerbounds/plugins/nvcf/nvcf_decorator.py,sha256=pSWKaPyGXBEfxG35QA1FQljio8ADjwf-DnPgEsqXoUM,9251
 metaflow_extensions/outerbounds/plugins/nvcf/utils.py,sha256=DxWSCayfa95e0HJkWacey1s1nxoTpaunGhrb_0Ayv28,133
 metaflow_extensions/outerbounds/plugins/ollama/__init__.py,sha256=HEsI5U4ckQby7K2NsGBOdizhPY3WWqXSnXx_IHL7_No,2307
 metaflow_extensions/outerbounds/plugins/ollama/ollama.py,sha256=KlP8_EmnUoi8-PidyU0IDuENYxKjQaHFC33yGsvaeic,13320
@@ -52,13 +55,13 @@ metaflow_extensions/outerbounds/plugins/tensorboard/__init__.py,sha256=9lUM4Cqi5
 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=2CNRACc2EO-d_OtN7Ob-DlZAhD1ryYKhltB-4dVndAQ,1923
+metaflow_extensions/outerbounds/toplevel/global_aliases_for_metaflow_package.py,sha256=oCRFfwoLODDju2-6JU2h74x6O-GCyO7PRRKTZsj-73k,2009
 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
 metaflow_extensions/outerbounds/toplevel/plugins/ollama/__init__.py,sha256=GRSz2zwqkvlmFS6bcfYD_CX6CMko9DHQokMaH1iBshA,47
 metaflow_extensions/outerbounds/toplevel/plugins/snowflake/__init__.py,sha256=LptpH-ziXHrednMYUjIaosS1SXD3sOtF_9_eRqd8SJw,50
-ob_metaflow_extensions-1.1.146.dist-info/METADATA,sha256=Vtp-8yJaPeg8Hh33WXQYAerLoLlx5Ldya5M2U1lp5z0,520
-ob_metaflow_extensions-1.1.146.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
-ob_metaflow_extensions-1.1.146.dist-info/top_level.txt,sha256=NwG0ukwjygtanDETyp_BUdtYtqIA_lOjzFFh1TsnxvI,20
-ob_metaflow_extensions-1.1.146.dist-info/RECORD,,
+ob_metaflow_extensions-1.1.147.dist-info/METADATA,sha256=KjFRs8JFmUIUzP-ghZUKxAGefc9i6tpf4GXDpJrSbSE,520
+ob_metaflow_extensions-1.1.147.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
+ob_metaflow_extensions-1.1.147.dist-info/top_level.txt,sha256=NwG0ukwjygtanDETyp_BUdtYtqIA_lOjzFFh1TsnxvI,20
+ob_metaflow_extensions-1.1.147.dist-info/RECORD,,