ob-metaflow-extensions 1.3.0__py2.py3-none-any.whl → 1.3.1__py2.py3-none-any.whl

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 = [

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

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)

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(

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,

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:

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.

metaflow_extensions/outerbounds/plugins/s3_proxy/__init__.py

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

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"]

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"

metaflow_extensions/outerbounds/plugins/s3_proxy/s3_proxy_api.py

@@ -0,0 +1,93 @@
+import json
+import time
+from typing import Dict, Optional
+
+from .exceptions import S3ProxyConfigException, S3ProxyApiException
+
+
+class S3ProxyConfigResponse:
+    def __init__(self, data: Dict):
+        self.bucket_name = data.get("bucket_name")
+        self.endpoint_url = data.get("endpoint_url")
+        self.access_key_id = data.get("access_key_id")
+        self.secret_access_key = data.get("secret_access_key")
+        self.region = data.get("region")
+
+
+class S3ProxyApiClient:
+    def __init__(self):
+        self.perimeter, self.integrations_url = self._get_api_configs()
+
+    def _get_api_configs(self):
+        from metaflow_extensions.outerbounds.remote_config import init_config
+        from os import environ
+
+        conf = init_config()
+        perimeter = conf.get("OBP_PERIMETER") or environ.get("OBP_PERIMETER", "")
+        integrations_url = conf.get("OBP_INTEGRATIONS_URL") or environ.get(
+            "OBP_INTEGRATIONS_URL", ""
+        )
+
+        if not perimeter:
+            raise S3ProxyConfigException(
+                "No perimeter set. Please run `outerbounds configure` command."
+            )
+
+        if not integrations_url:
+            raise S3ProxyConfigException(
+                "No integrations URL set. Please contact your Outerbounds support team."
+            )
+
+        return perimeter, integrations_url
+
+    def fetch_s3_proxy_config(
+        self, integration_name: Optional[str] = None
+    ) -> S3ProxyConfigResponse:
+        url = f"{self.integrations_url}/s3proxy"
+
+        payload = {"perimeter_name": self.perimeter}
+        if integration_name:
+            payload["integration_name"] = integration_name
+
+        headers = {"Content-Type": "application/json"}
+
+        try:
+            from metaflow.metaflow_config import SERVICE_HEADERS
+
+            headers.update(SERVICE_HEADERS or {})
+        except ImportError:
+            pass
+
+        response = self._make_request(url, headers, payload)
+        return S3ProxyConfigResponse(response)
+
+    def _make_request(self, url: str, headers: Dict, payload: Dict) -> Dict:
+        from metaflow_extensions.outerbounds.plugins.secrets.secrets import (
+            _api_server_get,
+        )
+
+        retryable_status_codes = [409]
+        json_payload = json.dumps(payload)
+
+        for attempt in range(3):
+            response = _api_server_get(
+                url, data=json_payload, headers=headers, conn_error_retries=5
+            )
+
+            if response.status_code not in retryable_status_codes:
+                break
+
+            if attempt < 2:
+                time.sleep(0.5 * (attempt + 1))
+
+        if response.status_code != 200:
+            error_msg = f"API request failed with status {response.status_code}"
+            try:
+                error_data = response.json()
+                if "message" in error_data:
+                    error_msg = error_data["message"]
+            except:
+                pass
+            raise S3ProxyApiException(error_msg)
+
+        return response.json()

metaflow_extensions/outerbounds/plugins/s3_proxy/s3_proxy_decorator.py

@@ -0,0 +1,106 @@
+import functools
+from typing import Optional
+
+from metaflow import current
+from metaflow.decorators import StepDecorator
+
+from .s3_proxy_manager import S3ProxyManager
+from .exceptions import S3ProxyException
+from .constants import S3_PROXY_WRITE_MODES
+
+
+class S3ProxyDecorator(StepDecorator):
+    """
+    S3 Proxy decorator for routing S3 requests through a local proxy service.
+
+    Parameters
+    ----------
+    integration_name : str, optional
+        Name of the S3 proxy integration. If not specified, will use the only
+        available S3 proxy integration in the namespace (fails if multiple exist).
+    write_mode : str, optional
+        The desired behavior during write operations to target (origin) S3 bucket.
+        allowed options are:
+            "origin-and-cache" -> write to both the target S3 bucket and local object
+                storage
+            "origin" -> only write to the target S3 bucket
+            "cache" -> only write to the object storage service used for caching
+    debug : bool, optional
+        Enable debug logging for proxy operations.
+    """
+
+    name = "s3_proxy"
+    defaults = {
+        "integration_name": None,
+        "write_mode": None,
+        "debug": False,
+    }
+
+    def step_init(self, flow, graph, step, decos, environment, flow_datastore, logger):
+        write_mode = self.attributes["write_mode"]
+        if write_mode and write_mode not in S3_PROXY_WRITE_MODES:
+            raise S3ProxyException(
+                f"unexpected write_mode specified: {write_mode}. Allowed values are: {','.join(S3_PROXY_WRITE_MODES)}."
+            )
+
+        self.manager = S3ProxyManager(
+            integration_name=self.attributes["integration_name"],
+            write_mode=self.attributes["write_mode"],
+            debug=self.attributes["debug"],
+        )
+
+        current._update_env({"s3_proxy": self.manager})
+
+    def task_pre_step(
+        self,
+        step_name,
+        task_datastore,
+        metadata,
+        run_id,
+        task_id,
+        flow,
+        graph,
+        retry_count,
+        max_user_code_retries,
+        ubf_context,
+        inputs,
+    ):
+        """Setup S3 proxy before step execution"""
+        self.manager.setup_proxy()
+
+    def task_finished(
+        self, step_name, flow, graph, is_task_ok, retry_count, max_retries
+    ):
+        """Cleanup S3 proxy after step execution"""
+        if self.manager:
+            self.manager.cleanup()
+
+
+class NebiusS3ProxyDecorator(S3ProxyDecorator):
+    """
+    Nebius-specific S3 Proxy decorator for routing S3 requests through a local proxy service.
+    It exists to make it easier for users to know that this decorator should only be used with
+    a Neo Cloud like Nebius.
+    """
+
+    name = "nebius_s3_proxy"
+    defaults = {
+        "integration_name": None,
+        "write_mode": None,
+        "debug": False,
+    }
+
+
+class CoreWeaveS3ProxyDecorator(S3ProxyDecorator):
+    """
+    CoreWeave-specific S3 Proxy decorator for routing S3 requests through a local proxy service.
+    It exists to make it easier for users to know that this decorator should only be used with
+    a Neo Cloud like CoreWeave.
+    """
+
+    name = "coreweave_s3_proxy"
+    defaults = {
+        "integration_name": None,
+        "write_mode": None,
+        "debug": False,
+    }

metaflow_extensions/outerbounds/plugins/s3_proxy/s3_proxy_manager.py

@@ -0,0 +1,222 @@
+import os
+import json
+import gzip
+import time
+import threading
+import subprocess
+from pathlib import Path
+from typing import Optional
+
+import requests
+
+from .constants import (
+    S3_PROXY_BINARY_URLS,
+    DEFAULT_PROXY_PORT,
+    DEFAULT_PROXY_HOST,
+)
+from metaflow.metaflow_config import AWS_SECRETS_MANAGER_DEFAULT_REGION
+from .s3_proxy_api import S3ProxyApiClient
+from .exceptions import S3ProxyException
+
+
+class S3ProxyManager:
+    def __init__(
+        self,
+        integration_name: Optional[str] = None,
+        write_mode: Optional[str] = None,
+        debug: bool = False,
+    ):
+        self.integration_name = integration_name
+        self.write_mode = write_mode
+        self.debug = debug
+        self.process = None
+        self.binary_path = None
+        self.config_path = None
+        self.api_client = S3ProxyApiClient()
+        self.proxy_config = None
+
+    def setup_proxy(self) -> bool:
+        try:
+            if self._is_running_in_kubernetes():
+                config_data = self.api_client.fetch_s3_proxy_config(
+                    self.integration_name
+                )
+                self.binary_path = self._download_binary()
+                self.config_path = self._write_config_file(config_data)
+                self.process = self._start_proxy_process()
+                self._setup_proxy_config(config_data)
+                return True
+
+            print(
+                "[@s3_proxy] skipping s3-proxy set up because metaflow has not detected a Kubernetes environment"
+            )
+            return False
+        except Exception as e:
+            if self.debug:
+                print(f"[@s3_proxy] Setup failed: {e}")
+            self.cleanup()
+            raise
+
+    def _is_running_in_kubernetes(self) -> bool:
+        """Check if running inside a Kubernetes pod by checking for Kubernetes service account token."""
+        return (
+            os.path.exists("/var/run/secrets/kubernetes.io/serviceaccount/token")
+            and os.environ.get("KUBERNETES_SERVICE_HOST") is not None
+        )
+
+    def _download_binary(self) -> str:
+        binary_path = Path("/tmp/s3-proxy")
+        if binary_path.exists():
+            if self.debug:
+                print("[@s3_proxy] Binary already exists, skipping download")
+            return str(binary_path.absolute())
+
+        try:
+            if self.debug:
+                print("[@s3_proxy] Downloading binary...")
+
+            from platform import machine
+
+            arch = machine()
+            if arch not in S3_PROXY_BINARY_URLS:
+                raise S3ProxyException(
+                    f"unsupported platform architecture: {arch}. Please reach out to your Outerbounds Support team for more help."
+                )
+
+            response = requests.get(S3_PROXY_BINARY_URLS[arch], stream=True, timeout=60)
+            response.raise_for_status()
+
+            with open(binary_path, "wb") as f:
+                with gzip.GzipFile(fileobj=response.raw) as gz:
+                    f.write(gz.read())
+
+            binary_path.chmod(0o755)
+
+            if self.debug:
+                print("[@s3_proxy] Binary downloaded successfully")
+
+            return str(binary_path.absolute())
+
+        except Exception as e:
+            if self.debug:
+                print(f"[@s3_proxy] Binary download failed: {e}")
+            raise S3ProxyException(f"Failed to download S3 proxy binary: {e}")
+
+    def _write_config_file(self, config_data) -> str:
+        config_path = Path("/tmp/s3-proxy-config.json")
+
+        proxy_config = {
+            "bucketName": config_data.bucket_name,
+            "endpointUrl": config_data.endpoint_url,
+            "accessKeyId": config_data.access_key_id,
+            "accessKeySecret": config_data.secret_access_key,
+            "region": config_data.region,
+        }
+
+        config_path.write_text(json.dumps(proxy_config, indent=2))
+
+        if self.debug:
+            print(f"[@s3_proxy] Config written to {config_path}")
+
+        return str(config_path.absolute())
+
+    def _start_proxy_process(self) -> subprocess.Popen:
+        cmd = [self.binary_path, "--bucket-config", self.config_path, "serve"]
+
+        if self.debug:
+            print(f"[@s3_proxy] Starting proxy: {' '.join(cmd)}")
+
+        process = subprocess.Popen(
+            cmd,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.STDOUT,  # Redirect stderr to stdout
+            text=True,
+            start_new_session=True,
+        )
+
+        self._setup_log_streaming(process)
+
+        time.sleep(3)
+
+        if process.poll() is None:
+            if self.debug:
+                print(f"[@s3_proxy] Proxy started successfully (pid: {process.pid})")
+
+            return process
+        else:
+            stdout_data, stderr_data = process.communicate()
+            if self.debug:
+                print(f"[@s3_proxy] Proxy failed to start - output: {stdout_data}")
+            raise S3ProxyException(f"S3 proxy failed to start: {stdout_data}")
+
+    def _setup_log_streaming(self, process: subprocess.Popen):
+        def stream_logs():
+            try:
+                # Read stdout line by line (stderr is redirected to stdout)
+                while True:
+                    line = process.stdout.readline()
+                    if not line:
+                        # Process has ended
+                        break
+                    line = line.strip()
+                    if line and self.debug:
+                        print(f"[@s3_proxy] {line}")
+
+            except Exception as e:
+                if self.debug:
+                    print(f"[@s3_proxy] Log streaming error: {e}")
+
+        log_thread = threading.Thread(target=stream_logs, daemon=True)
+        log_thread.start()
+
+    def _setup_proxy_config(self, config_data):
+        from metaflow_extensions.outerbounds.toplevel.global_aliases_for_metaflow_package import (
+            set_s3_proxy_config,
+        )
+        from metaflow.metaflow_config import AWS_SECRETS_MANAGER_DEFAULT_REGION
+
+        region = os.environ.get(
+            "METAFLOW_AWS_SECRETS_MANAGER_DEFAULT_REGION",
+            AWS_SECRETS_MANAGER_DEFAULT_REGION,
+        )
+
+        proxy_config = {
+            "endpoint_url": f"http://{DEFAULT_PROXY_HOST}:{DEFAULT_PROXY_PORT}",
+            "region": region,
+            "bucket_name": config_data.bucket_name,
+            "active": True,
+        }
+
+        if self.write_mode:
+            proxy_config["write_mode"] = self.write_mode
+
+        set_s3_proxy_config(proxy_config)
+        self.proxy_config = proxy_config
+
+        if self.debug:
+            print("[@s3_proxy] Global S3 proxy configuration activated")
+
+    def cleanup(self):
+        try:
+            from metaflow_extensions.outerbounds.toplevel.global_aliases_for_metaflow_package import (
+                clear_s3_proxy_config,
+            )
+
+            clear_s3_proxy_config()
+
+            if self.process and self.process.poll() is None:
+                self.process.terminate()
+                self.process.wait(timeout=5)
+                if self.debug:
+                    print("[@s3_proxy] Proxy process stopped")
+
+                from os import remove
+
+                remove(self.config_path)
+                remove(self.binary_path)
+
+        except Exception as e:
+            if self.debug:
+                print(f"[@s3_proxy] Cleanup error: {e}")
+        finally:
+            self.proxy_config = None

metaflow_extensions/outerbounds/toplevel/global_aliases_for_metaflow_package.py

@@ -5,6 +5,27 @@
 __version__ = "v1"
 __mf_extensions__ = "ob"
 
+from metaflow_extensions.outerbounds.toplevel.s3_proxy import (
+    get_aws_client_with_s3_proxy,
+    get_S3_with_s3_proxy,
+)
+
+_S3_PROXY_CONFIG = None
+
+
+def set_s3_proxy_config(config):
+    global _S3_PROXY_CONFIG
+    _S3_PROXY_CONFIG = config
+
+
+def clear_s3_proxy_config():
+    global _S3_PROXY_CONFIG
+    _S3_PROXY_CONFIG = None
+
+
+def get_s3_proxy_config():
+    return _S3_PROXY_CONFIG
+
 
 # Must match the signature of metaflow.plugins.aws.aws_client.get_aws_client
 # This function is called by the "userland" code inside tasks. Metaflow internals
@@ -34,7 +55,12 @@ def get_aws_client(
         if decorator_role_arn:
             role_arn = decorator_role_arn
 
-    return metaflow.plugins.aws.aws_client.get_aws_client(
+    if module == "s3" and _S3_PROXY_CONFIG is not None:
+        return get_aws_client_with_s3_proxy(
+            module, with_error, role_arn, session_vars, client_params, _S3_PROXY_CONFIG
+        )
+
+    client = metaflow.plugins.aws.aws_client.get_aws_client(
         module,
         with_error=with_error,
         role_arn=role_arn or USE_CSPR_ROLE_ARN_IF_SET,
@@ -42,6 +68,8 @@ def get_aws_client(
         client_params=client_params,
     )
 
+    return client
+
 
 # This should match the signature of metaflow.plugins.datatools.s3.S3.
 #
@@ -68,6 +96,10 @@ def S3(*args, **kwargs):
         else:
             kwargs["role"] = USE_CSPR_ROLE_ARN_IF_SET
 
+    # Check if S3 proxy is active using module variable (like CSPR)
+    if _S3_PROXY_CONFIG is not None:
+        return get_S3_with_s3_proxy(_S3_PROXY_CONFIG, *args, **kwargs)
+
     return metaflow.plugins.datatools.s3.S3(*args, **kwargs)
 
 

metaflow_extensions/outerbounds/toplevel/s3_proxy.py

@@ -0,0 +1,88 @@
+from metaflow_extensions.outerbounds.plugins import USE_CSPR_ROLE_ARN_IF_SET
+from metaflow.metaflow_config import AWS_SECRETS_MANAGER_DEFAULT_REGION
+from metaflow_extensions.outerbounds.plugins.s3_proxy.constants import (
+    DEFAULT_PROXY_HOST,
+    DEFAULT_PROXY_PORT,
+)
+
+
+def get_aws_client_with_s3_proxy(
+    module,
+    with_error=False,
+    role_arn=None,
+    session_vars=None,
+    client_params=None,
+    s3_config=None,
+):
+    if not client_params:
+        client_params = {}
+
+    client_params["region_name"] = client_params.get(
+        "region_name", s3_config.get("region")
+    )
+    client_params["endpoint_url"] = s3_config.get(
+        "endpoint_url", f"http://{DEFAULT_PROXY_HOST}:{DEFAULT_PROXY_PORT}"
+    )
+
+    import metaflow.plugins.aws.aws_client
+
+    client = metaflow.plugins.aws.aws_client.get_aws_client(
+        module,
+        with_error=with_error,
+        role_arn=role_arn or USE_CSPR_ROLE_ARN_IF_SET,
+        session_vars=session_vars,
+        client_params=client_params,
+    )
+
+    def override_s3_proxy_host_header(request, **kwargs):
+        region = kwargs["region_name"]
+        request.headers["Host"] = f"s3.{region}.amazonaws.com"
+        if "x-ob-write-to" not in request.headers and "write_mode" in s3_config:
+            request.headers["x-ob-write-to"] = s3_config.get("write_mode")
+
+    client.meta.events.register("before-sign", override_s3_proxy_host_header)
+
+    return client
+
+
+def get_S3_with_s3_proxy(s3_config, *args, **kwargs):
+    if "region_name" not in kwargs:
+        kwargs["region_name"] = s3_config.get(
+            "region", AWS_SECRETS_MANAGER_DEFAULT_REGION
+        )
+
+    kwargs["endpoint_url"] = s3_config.get(
+        "endpoint_url", f"http://{DEFAULT_PROXY_HOST}:{DEFAULT_PROXY_PORT}"
+    )
+
+    import metaflow.plugins.datatools.s3
+
+    mf_s3 = metaflow.plugins.datatools.s3.S3(*args, **kwargs)
+
+    # Override reset_client to ensure proxy endpoint is preserved
+    original_reset_client = mf_s3._s3_client.reset_client
+
+    def proxy_reset_client():
+        original_reset_client()
+        import boto3
+
+        proxy_client = boto3.client(
+            "s3",
+            region_name=kwargs.get("region_name", s3_config.get("region")),
+            endpoint_url=s3_config.get("endpoint_url"),
+        )
+        mf_s3._s3_client._s3_client = proxy_client
+
+    mf_s3._s3_client.reset_client = proxy_reset_client
+    mf_s3._s3_client.reset_client()
+
+    def override_s3_proxy_host_header(request, **kwargs):
+        region = kwargs["region_name"]
+        request.headers["Host"] = f"s3.{region}.amazonaws.com"
+        if "x-ob-write-to" not in request.headers and "write_mode" in s3_config:
+            request.headers["x-ob-write-to"] = s3_config.get("write_mode")
+
+    mf_s3._s3_client._s3_client.meta.events.register(
+        "before-sign", override_s3_proxy_host_header
+    )
+    return mf_s3

{ob_metaflow_extensions-1.3.0.dist-info → ob_metaflow_extensions-1.3.1.dist-info}/METADATA

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

{ob_metaflow_extensions-1.3.0.dist-info → ob_metaflow_extensions-1.3.1.dist-info}/RECORD

@@ -1,7 +1,7 @@
 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=c4zZnULrBOU2bmleU_D7Xr5H3kTuTgDhtZIKUAJiNqw,13795
+metaflow_extensions/outerbounds/plugins/__init__.py,sha256=ISFPm4cn8IvfZnM6Q4LqDalycr02v6gMsCPo9mRNSwk,14028
 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
@@ -12,18 +12,18 @@ metaflow_extensions/outerbounds/plugins/apps/consts.py,sha256=iHsyqbUg9k-rgswCs1
 metaflow_extensions/outerbounds/plugins/apps/deploy_decorator.py,sha256=VkmiMdNYHhNdt-Qm9AVv7aE2LWFsIFEc16YcOYjwF6Q,8568
 metaflow_extensions/outerbounds/plugins/apps/supervisord_utils.py,sha256=GQoN2gyPClcpR9cLldJmbCfqXnoAHxp8xUnY7vzaYtY,9026
 metaflow_extensions/outerbounds/plugins/apps/core/__init__.py,sha256=c6uCgKlgEkTmM9BVdAO-m3vZvUpK2KW_AZZ2236now4,237
-metaflow_extensions/outerbounds/plugins/apps/core/_state_machine.py,sha256=8z4MnPeat3Vm8ekLIMJj8vVsvUAQ__i8daf-9UItUIQ,19926
-metaflow_extensions/outerbounds/plugins/apps/core/app_cli.py,sha256=9YyvOQzPNlpxA2K9AZ4jYpfDWpLSp66u_NotGGE5DHg,42155
+metaflow_extensions/outerbounds/plugins/apps/core/_state_machine.py,sha256=q2in_xmHsFJ1KwgQhVbg_DuJM5I2XdQHhbMGKlJpdcs,19957
+metaflow_extensions/outerbounds/plugins/apps/core/app_cli.py,sha256=V0Ki_VwjVIyIa2sgXPC7miOPLYWLrsHvzMpTfQypU2U,42169
 metaflow_extensions/outerbounds/plugins/apps/core/app_config.py,sha256=PHt-HdNfTHIuhY-eB5vkRMp1RKQNWJ4DKdgZWyYgUuc,4167
 metaflow_extensions/outerbounds/plugins/apps/core/artifacts.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-metaflow_extensions/outerbounds/plugins/apps/core/capsule.py,sha256=9_LrpZcDg3w0FZvs9h_6Mmiy5YTYLCLnEu8wnNau24E,34445
+metaflow_extensions/outerbounds/plugins/apps/core/capsule.py,sha256=tswwFa-c3P0CQtomQV6d--eP4hDQI4SjbbUccvwkmGo,34314
 metaflow_extensions/outerbounds/plugins/apps/core/click_importer.py,sha256=kgoPQmK_-8PSSTc3QMSaynCLQ5VWTkKFOC69FPURyXA,998
 metaflow_extensions/outerbounds/plugins/apps/core/config_schema.yaml,sha256=iTThO2vNQrFWe9nYfjiOcMf6FOQ6vU_1ZhXhUAr0L24,8142
-metaflow_extensions/outerbounds/plugins/apps/core/dependencies.py,sha256=HDPj7rDARcsKeex5GwH0IP8rOXMH6YdOufgXDknP1S8,4006
+metaflow_extensions/outerbounds/plugins/apps/core/dependencies.py,sha256=JlWT9f27yzZeJPlqTQk134WDfQgOdyxC5iaw3pLlhqY,4006
 metaflow_extensions/outerbounds/plugins/apps/core/deployer.py,sha256=VkYe8mK_VOr-bAiR2RohhKeLf8Z3gHZw7RoRBSCu2FA,9765
 metaflow_extensions/outerbounds/plugins/apps/core/perimeters.py,sha256=jeFGAUnFQkBFiOMp_Ls7Ofb80Qogh509suam5sMucYU,3030
-metaflow_extensions/outerbounds/plugins/apps/core/secrets.py,sha256=aWzcAayQEJghQgFP_qp9w6jyvan_hoL4_ceqZ0ZjLd4,6126
-metaflow_extensions/outerbounds/plugins/apps/core/utils.py,sha256=RLO6p25Fzq4olLFtQmfSl9LT0NPDfUosxPrsjO9sczo,7897
+metaflow_extensions/outerbounds/plugins/apps/core/secrets.py,sha256=sgDiAmpSC8Y5xjlaOEp79F6m0S3x4RONf_vJ5PUAfu8,6127
+metaflow_extensions/outerbounds/plugins/apps/core/utils.py,sha256=2M2zU8DhbAlJee8P0xKXINAku81PcUylS3sVCSb0TUs,7896
 metaflow_extensions/outerbounds/plugins/apps/core/validations.py,sha256=Inr9AJDe-L3PMMMxcJPH1zulh9_SynqITb2BzGseLh4,471
 metaflow_extensions/outerbounds/plugins/apps/core/_vendor/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 metaflow_extensions/outerbounds/plugins/apps/core/_vendor/spinner/__init__.py,sha256=gTm0NdQGTRxmghye1CYkhRtodji0MezsqAWs9OrLLRc,102
@@ -86,6 +86,12 @@ metaflow_extensions/outerbounds/plugins/ollama/status_card.py,sha256=F5e4McDl28l
 metaflow_extensions/outerbounds/plugins/profilers/deco_injector.py,sha256=oI_C3c64XBm7n88FILqHwn-Nnc5DeT_68I67lM9rXaI,2434
 metaflow_extensions/outerbounds/plugins/profilers/gpu_profile_decorator.py,sha256=gDHQ2sMIp4NuZSzUspbSd8RGdFAoO5mgZAyFcZ2a51Y,2619
 metaflow_extensions/outerbounds/plugins/profilers/simple_card_decorator.py,sha256=4W9tLGCmkFx-4XYLa1xF6qMiaWOBYYFx_RclZDKej30,3259
+metaflow_extensions/outerbounds/plugins/s3_proxy/__init__.py,sha256=9Kw86B331pQJAzkfBMPIDoPrJsW0LVRHXBYikbcc2xk,204
+metaflow_extensions/outerbounds/plugins/s3_proxy/constants.py,sha256=Jjsd3cuo3IMi8rcKsUJx2PK188hMhFNyPTNKCFKfAQI,319
+metaflow_extensions/outerbounds/plugins/s3_proxy/exceptions.py,sha256=IkPqDvSeYQukNeu0aIVCmfQWTvUHsTs-qv7nvry2KjM,305
+metaflow_extensions/outerbounds/plugins/s3_proxy/s3_proxy_api.py,sha256=WjpprW0tCICLOihFywEtgJbCnx-OFmwuT_hR27ACl2A,3007
+metaflow_extensions/outerbounds/plugins/s3_proxy/s3_proxy_decorator.py,sha256=yqQryVjaUQ6Aq_SMI8IRHXwzPokkznHncLDpLSEcQeM,3285
+metaflow_extensions/outerbounds/plugins/s3_proxy/s3_proxy_manager.py,sha256=qmSypJXY-l7P2sI4mO6y-Rut5vGL2m1TjvGIXHUi6vs,7379
 metaflow_extensions/outerbounds/plugins/secrets/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 metaflow_extensions/outerbounds/plugins/secrets/secrets.py,sha256=3s98hO_twKkM22tKyDdcUjGQNfYpSXW_jLKISV9ju_U,8433
 metaflow_extensions/outerbounds/plugins/snowflake/__init__.py,sha256=RG4ixt3jwqcK1_tt0QxLcUbNmf7wWAMnZhBx-ZMGgLk,114
@@ -108,8 +114,9 @@ metaflow_extensions/outerbounds/plugins/vllm/vllm_manager.py,sha256=sp_TX2SrImJG
 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=xVCIvKeEzMwu2vfsWYqnq4aetMGmPBjpzOXrZfBr5iI,3036
+metaflow_extensions/outerbounds/toplevel/global_aliases_for_metaflow_package.py,sha256=StTRMBHjuxfxe-wQs8ikoAZc4xnhlceY0R4avaJ1Ps8,3823
 metaflow_extensions/outerbounds/toplevel/ob_internal.py,sha256=DXCaAtLzlE-bFIiVWEv-iV2JKIWsoSGaUeH4jIQZ9gs,193
+metaflow_extensions/outerbounds/toplevel/s3_proxy.py,sha256=zdqG7Z12cGuoYYCi2P4kqC3WsgL3xfdJGIb7ejecHH4,2862
 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
@@ -117,7 +124,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.3.0.dist-info/METADATA,sha256=AqSHKfzK42NBfzMZAHKk3eqN2WWVkxbQQNp_EQA_EnA,518
-ob_metaflow_extensions-1.3.0.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
-ob_metaflow_extensions-1.3.0.dist-info/top_level.txt,sha256=NwG0ukwjygtanDETyp_BUdtYtqIA_lOjzFFh1TsnxvI,20
-ob_metaflow_extensions-1.3.0.dist-info/RECORD,,
+ob_metaflow_extensions-1.3.1.dist-info/METADATA,sha256=HnIBZEkQl7rStKXqEPoFTDNrG5O6a_n8iYmSdnQT_q4,518
+ob_metaflow_extensions-1.3.1.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
+ob_metaflow_extensions-1.3.1.dist-info/top_level.txt,sha256=NwG0ukwjygtanDETyp_BUdtYtqIA_lOjzFFh1TsnxvI,20
+ob_metaflow_extensions-1.3.1.dist-info/RECORD,,