ob-metaflow-extensions 1.2.10__tar.gz → 1.3.1__tar.gz

{ob_metaflow_extensions-1.2.10 → ob_metaflow_extensions-1.3.1}/PKG-INFO

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

{ob_metaflow_extensions-1.2.10 → ob_metaflow_extensions-1.3.1}/metaflow_extensions/outerbounds/plugins/__init__.py

@@ -340,6 +340,9 @@ STEP_DECORATORS_DESC = [
     ("ollama", ".ollama.OllamaDecorator"),
     ("vllm", ".vllm.VLLMDecorator"),
     ("app_deploy", ".apps.app_deploy_decorator.AppDeployDecorator"),
+    ("s3_proxy", ".s3_proxy.s3_proxy_decorator.S3ProxyDecorator"),
+    ("nebius_s3_proxy", ".s3_proxy.s3_proxy_decorator.NebiusS3ProxyDecorator"),
+    ("coreweave_s3_proxy", ".s3_proxy.s3_proxy_decorator.CoreWeaveS3ProxyDecorator"),
 ]
 
 TOGGLE_STEP_DECORATOR = [

{ob_metaflow_extensions-1.2.10 → ob_metaflow_extensions-1.3.1}/metaflow_extensions/outerbounds/plugins/apps/core/_state_machine.py

@@ -199,29 +199,29 @@ class DEPLOYMENT_READY_CONDITIONS:
     This allows users or platform designers to configure the criteria for deployment readiness.
 
     Why do we need deployment readiness conditions?
-    - Deployments might be taking place from a CI/CD esq environment, In these setups, the downstream build triggers might be depending on a specific criteria for deployment completion. Having readiness conditions allows the CI/CD systems to get a signal of when the deployment is ready.
+        - Deployments might be taking place from a CI/CD-esque environment, In these setups, the downstream build triggers might be depending on a specific criteria for deployment completion. Having readiness conditions allows the CI/CD systems to get a signal of when the deployment is ready.
     - Users might be calling the deployment API under different conditions:
         - Some users might want a cluster of workers ready before serving traffic while others might want just one worker ready to start serving traffic.
 
     Some readiness conditions include:
-    1) [at_least_one_running] Atleast min(min_replicas, 1) workers of the current deployment instance's version have started running.
+            1) [at_least_one_running] At least min(min_replicas, 1) workers of the current deployment instance's version have started running.
         - Usecase: Some endpoints may be deployed ephemerally and are considered ready when at least one instance is running; additional instances are for load management.
-    2) [all_running] Atleast min_replicas number of workers are running for the deployment to be considered ready.
+        2) [all_running] At least min_replicas number of workers are running for the deployment to be considered ready.
         - Usecase: Operators may require that all replicas are available before traffic is routed. Needed when inference endpoints maybe under some SLA or require a larger load
-    3) [fully_finished] Atleast min_replicas number of workers are running for the deployment and there are no pending or crashlooping workers from previous versions lying around.
+        3) [fully_finished] At least min_replicas number of workers are running for the deployment and there are no pending or crashlooping workers from previous versions lying around.
         - Usecase: Ensuring endpoint is fully available and no other versions are running or endpoint has been fully scaled down.
     4) [async] The deployment will be assumed ready as soon as the server responds with a 200.
         - Usecase: Operators may only care that the URL is minted for the deployment or the deployment eventually scales down to 0.
     """
 
-    # `ATLEAST_ONE_RUNNING` implies that atleast one worker of the current deployment instance's version has started running.
+    # `ATLEAST_ONE_RUNNING` implies that at least one worker of the current deployment instance's version has started running.
     ATLEAST_ONE_RUNNING = "at_least_one_running"
 
     # `ALL_RUNNING` implies that all workers of the current deployment instance's version have started running (i.e. all workers aligning to the minimum number of replicas).
     # It doesn't imply that all the workers relating to other deployments have been torn down.
     ALL_RUNNING = "all_running"
 
-    # `FULLY_FINISHED` implies Atleast min_replicas number of workers are running for the deployment and there are no pending or crashlooping workers from previous versions lying around.
+    # `FULLY_FINISHED` implies at least min_replicas number of workers are running for the deployment and there are no pending or crashlooping workers from previous versions lying around.
     FULLY_FINISHED = "fully_finished"
 
     # `ASYNC` implies that the deployment will be assumed ready after the URL is minted and the worker statuses are not checked.
@@ -447,9 +447,9 @@ def _capsule_worker_semantic_status(
     crashlooping_workers = _make_version_dict(workers, "CrashLoopBackOff")
 
     # current_status (formulated basis):
-    # - atleast one pods are pending for `_end_state_capsule_version`
-    # - atleast one pod is in Running state for `_end_state_capsule_version` (maybe terminal) [Might require heath-check thing here]
-    # - alteast one pod is crashlooping for `_end_state_capsule_version` (maybe terminal)
+    # - at least one pods are pending for `_end_state_capsule_version`
+    # - at least one pod is in Running state for `_end_state_capsule_version` (maybe terminal) [Might require health-check thing here]
+    # - at least one pod is crashlooping for `_end_state_capsule_version` (maybe terminal)
     # - all pods are running for `_end_state_capsule_version` that match the minimum number of replicas
     # - all pods are running for `_end_state_capsule_version` that match the maximum number of replicas and no other pods of older versions are running
     # - no pods relating to `_end_state_capsule_version` are pending/running/crashlooping

{ob_metaflow_extensions-1.2.10 → ob_metaflow_extensions-1.3.1}/metaflow_extensions/outerbounds/plugins/apps/core/app_cli.py

@@ -239,7 +239,7 @@ def _bake_image(app_config: AppConfig, cache_dir: str, logger):
         baking_status.resolved_image,
     )
     app_config.set_state("python_path", baking_status.python_path)
-    logger("🐳 Using The Docker Image : %s" % app_config.get_state("image"))
+    logger("🐳 Using the docker image : %s" % app_config.get_state("image"))
 
 
 def print_table(data, headers):
@@ -374,7 +374,7 @@ def _package_necessary_things(app_config: AppConfig, logger):
     #      or is it relative to where the caller command is sitting. Ideally it should work
     #      like Kustomizations where its relative to where the yaml file sits for simplicity
     #      of understanding relationships between config files. Ideally users can pass the src_path
-    #      from the command line and that will aliviate any need to package any other directories for
+    #      from the command line and that will alleviate any need to package any other directories for
     #
 
     package_dir = app_config.get_state("packaging_directory")
@@ -395,7 +395,7 @@ def _package_necessary_things(app_config: AppConfig, logger):
     )
     app_config.set_state("code_package_url", package_url)
     app_config.set_state("code_package_key", package_key)
-    logger("💾 Code Package Saved to : %s" % app_config.get_state("code_package_url"))
+    logger("💾 Code package saved to : %s" % app_config.get_state("code_package_url"))
 
 
 @app.command(help="Deploy an app to the Outerbounds Platform.")
@@ -465,7 +465,7 @@ def deploy(
 
         app_config.set_state("packaging_directory", packaging_directory)
         logger(
-            "📦 Packaging Directory : %s" % app_config.get_state("packaging_directory"),
+            "📦 Packaging directory : %s" % app_config.get_state("packaging_directory"),
         )
 
         if app_config.get("no_deps", False):
@@ -611,7 +611,7 @@ def deploy(
                     )
                 raise AppConfigError(message)
             capsule_logger(
-                f"🚀 {'' if not force_upgrade else 'Force'} Upgrading {capsule.capsule_type.lower()} `{capsule.name}`....",
+                f"🚀 {'Upgrading' if not force_upgrade else 'Force upgrading'} {capsule.capsule_type.lower()} `{capsule.name}`....",
                 color=ColorTheme.INFO_COLOR,
                 system_msg=True,
             )
@@ -632,7 +632,7 @@ def deploy(
             capsule_spinner.stop()
 
         logger(
-            f"💊 {capsule.capsule_type} {app_config.config['name']} ({capsule.identifier}) deployed! {capsule.capsule_type} exposed on the URL: {capsule.url}",
+            f"💊 {capsule.capsule_type} {app_config.config['name']} ({capsule.identifier}) deployed! {capsule.capsule_type} available on the URL: {capsule.url}",
             color=ColorTheme.INFO_COLOR,
             system_msg=True,
         )
@@ -761,7 +761,7 @@ def list(ctx, project, branch, name, tags, format, auth_type):
 def delete(ctx, name, cap_id, project, branch, tags, auto_approve):
 
     """Delete an app/apps from the Outerbounds Platform."""
-    # Atleast one of the args need to be provided
+    # At least one of the args need to be provided
     if not any(
         [
             name is not None,
@@ -772,7 +772,7 @@ def delete(ctx, name, cap_id, project, branch, tags, auto_approve):
         ]
     ):
         raise AppConfigError(
-            "Atleast one of the options need to be provided. You can use --name, --id, --project, --branch, --tag"
+            "At least one of the options need to be provided. You can use --name, --id, --project, --branch, --tag"
         )
 
     capsule_api = CapsuleApi(ctx.obj.api_url, ctx.obj.perimeter)

{ob_metaflow_extensions-1.2.10 → ob_metaflow_extensions-1.3.1}/metaflow_extensions/outerbounds/plugins/apps/core/capsule.py

@@ -44,24 +44,24 @@ class CapsuleStateMachine:
     - Happy Path:
         - First time Create :
             - wait for status.updateInProgress to be set to False
-                - (interleved) Poll the worker endpoints to check their status
+                - (interleaved) Poll the worker endpoints to check their status
                     - showcase how many workers are coming up if things are on the cli side.
                 - If the user has set some flag like `--dont-wait-to-fully-finish` then we check the `status.currentlyServedVersion` to see if even one replica is ready to
                 serve traffic.
             - once the status.updateInProgress is set to False, it means that the replicas are ready
         - Upgrade:
             - wait for status.updateInProgress to be set to False
-                - (interleved) Poll the worker endpoints to check their status and signal the user the number replicas coming up
+                - (interleaved) Poll the worker endpoints to check their status and signal the user the number replicas coming up
                 - If the user has set some flag like `--dont-wait-to-fully-finish` then we check the `status.currentlyServedVersion` to see if even one replica is ready to
                 serve traffic.
     - Unhappy Path:
         - First time Create :
             - wait for status.updateInProgress to be set to False,
-                - (interleved) Poll the workers to check their status.
+                - (interleaved) Poll the workers to check their status.
                     - If the worker pertaining the current deployment instance version is crashlooping then crash the deployment process with the error messages and logs.
         - Upgrade:
             - wait for status.updateInProgress to be set to False,
-                - (interleved) Poll the workers to check their status.
+                - (interleaved) Poll the workers to check their status.
                     - If the worker pertaining the current deployment instance version is crashlooping then crash the deployment process with the error messages and logs.
 
     """
@@ -210,9 +210,9 @@ class CapsuleInput:
     def construct_exec_command(cls, commands: List[str]):
         commands = ["set -eEuo pipefail"] + commands
         command_string = "\n".join(commands)
-        # First constuct a base64 encoded string of the quoted command
+        # First construct a base64 encoded string of the quoted command
         # One of the reasons we don't directly pass the command string to the backend with a `\n` join
-        # is because the backend controller doesnt play nice when the command can be a multi-line string.
+        # is because the backend controller doesn't play nice when the command can be a multi-line string.
         # So we encode it to a base64 string and then decode it back to a command string at runtime to provide to
         # `bash -c`. The ideal thing to have done is to run "bash -c {shlex.quote(command_string)}" and call it a day
         # but the backend controller yields the following error:
@@ -682,7 +682,7 @@ class CapsuleDeployer:
         """
         - `capsule_response.version` contains the version of the object present in the database
         - `current_deployment_instance_version` contains the version of the object that was deployed by this instance of the deployer.
-        In the situtation that the versions of the objects become a mismatch then it means that current deployment process is not giving the user the
+        In the situation that the versions of the objects become a mismatch then it means that current deployment process is not giving the user the
         output that they desire.
         """
         if capsule_response.get("version", None) != current_deployment_instance_version:
@@ -783,24 +783,22 @@ class CapsuleDeployer:
             )
             if capsule_ready or failure_condition_satisfied:
                 logger(
-                    "💊 %s deployment status: %s | worker states: [success :%s | failure :%s ] "
+                    "💊 %s deployment status: %s "
                     % (
                         self.capsule_type.title(),
                         "in progress"
                         if state_machine.update_in_progress
                         else "completed",
-                        capsule_ready,
-                        failure_condition_satisfied,
                     )
                 )
                 _further_readiness_check_failed = False
                 if further_check_worker_readiness:
                     # HACK : monitor the workers for N seconds to make sure they are healthy
-                    # this is a hack. Ideally we should implment a healtcheck as a first class citizen
+                    # this is a hack. Ideally we should implement a healthcheck as a first class citizen
                     # but it will take some time to do that so in the meanwhile a timeout set on the cli
                     # side will be really helpful.
                     logger(
-                        "💊 running last minute readiness check for %s..."
+                        "💊 Running last minute readiness check for %s..."
                         % self.identifier
                     )
                     _further_readiness_check_failed = self._monitor_worker_readiness(

{ob_metaflow_extensions-1.2.10 → ob_metaflow_extensions-1.3.1}/metaflow_extensions/outerbounds/plugins/apps/core/dependencies.py

@@ -95,7 +95,7 @@ def bake_deployment_image(
     pinned_conda_libs = get_pinned_conda_libs(python_version, DEFAULT_DATASTORE)
     pypi_packages.update(pinned_conda_libs)
     _reference = app_config.get("name", "default")
-    # `image` cannot be None. If by change it is none, FB will fart.
+    # `image` cannot be None. If by chance it is none, FB will fart.
     fb_response = bake_image(
         cache_file_path=cache_file_path,
         pypi_packages=pypi_packages,

{ob_metaflow_extensions-1.2.10 → ob_metaflow_extensions-1.3.1}/metaflow_extensions/outerbounds/plugins/apps/core/secrets.py

@@ -116,7 +116,7 @@ class SecretRetriever:
 
         if not perimeter:
             raise OuterboundsSecretsException(
-                "No perimeter set. Please make sure to run `outerbounds configure <...>` command which can be found on the Ourebounds UI or reach out to your Outerbounds support team."
+                "No perimeter set. Please make sure to run `outerbounds configure <...>` command which can be found on the Outerbounds UI or reach out to your Outerbounds support team."
             )
 
         if not integrations_url:

{ob_metaflow_extensions-1.2.10 → ob_metaflow_extensions-1.3.1}/metaflow_extensions/outerbounds/plugins/apps/core/utils.py

@@ -185,8 +185,8 @@ def safe_requests_wrapper(
         - How to handle retries for this case will be application specific.
     2. Errors when the API server may not be reachable (DNS resolution / network issues)
         - In this scenario, we know that something external to the API server is going wrong causing the issue.
-        - Failing pre-maturely in the case might not be the best course of action since critical user jobs might crash on intermittent issues.
-        - So in this case, we can just planely retry the request.
+        - Failing prematurely in the case might not be the best course of action since critical user jobs might crash on intermittent issues.
+        - So in this case, we can just plainly retry the request.
 
     This function handles the second case. It's a simple wrapper to handle the retry logic for connection errors.
     If this function is provided a `conn_error_retries` of 5, then the last retry will have waited 32 seconds.

{ob_metaflow_extensions-1.2.10 → ob_metaflow_extensions-1.3.1}/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.3.1/metaflow_extensions/outerbounds/plugins/checkpoint_datastores/coreweave.py

@@ -0,0 +1,71 @@
+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.3.1/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.3.1/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.10 → ob_metaflow_extensions-1.3.1}/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.3.1/metaflow_extensions/outerbounds/plugins/s3_proxy/__init__.py

@@ -0,0 +1,7 @@
+from .s3_proxy_decorator import (
+    S3ProxyDecorator,
+    NebiusS3ProxyDecorator,
+    CoreWeaveS3ProxyDecorator,
+)
+
+__all__ = ["S3ProxyDecorator", "NebiusS3ProxyDecorator", "CoreWeaveS3ProxyDecorator"]

ob_metaflow_extensions-1.3.1/metaflow_extensions/outerbounds/plugins/s3_proxy/constants.py

@@ -0,0 +1,8 @@
+S3_PROXY_BINARY_URLS = {
+    "aarch64": "https://fast-s3-proxy.outerbounds.sh/linux-arm64/s3-proxy-0.1.1.gz",
+    "x86_64": "https://fast-s3-proxy.outerbounds.sh/linux-amd64/s3-proxy-0.1.1.gz",
+}
+
+DEFAULT_PROXY_PORT = 8081
+DEFAULT_PROXY_HOST = "localhost"
+S3_PROXY_WRITE_MODES = ["origin-and-cache", "origin", "cache"]

ob_metaflow_extensions-1.3.1/metaflow_extensions/outerbounds/plugins/s3_proxy/exceptions.py

@@ -0,0 +1,13 @@
+from metaflow.exception import MetaflowException
+
+
+class S3ProxyException(MetaflowException):
+    headline = "S3 Proxy Error"
+
+
+class S3ProxyConfigException(S3ProxyException):
+    headline = "S3 Proxy Configuration Error"
+
+
+class S3ProxyApiException(S3ProxyException):
+    headline = "S3 Proxy API Error"