ob-metaflow-extensions 1.1.130__py2.py3-none-any.whl → 1.5.1__py2.py3-none-any.whl

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,250 @@
+import os
+import functools
+import json
+import signal
+import time
+from typing import Dict, List, Optional, Tuple
+
+from metaflow import current
+from metaflow.decorators import StepDecorator
+from .exceptions import S3ProxyException
+from .constants import S3_PROXY_WRITE_MODES
+from collections import namedtuple
+
+S3ProxyBinaryConfig = namedtuple(
+    "S3ProxyBinaryConfig", ["integration_name", "write_mode", "debug"]
+)
+
+
+def monkey_patch_environment(
+    environment, step_name_and_deco_attrs: Dict[str, S3ProxyBinaryConfig]
+):
+    wrapping_func = environment.bootstrap_commands
+
+    @functools.wraps(wrapping_func)
+    def wrapper(step_name, ds_type, *args):
+        base_boostrap_cmd = wrapping_func(step_name, ds_type, *args)
+        additional_cmd = []
+
+        if step_name in step_name_and_deco_attrs:
+            integration_name = step_name_and_deco_attrs[step_name].integration_name
+            write_mode = step_name_and_deco_attrs[step_name].write_mode
+            debug = step_name_and_deco_attrs[step_name].debug
+            additional_cmd = [
+                "echo 'Setting up the S3 proxy.'",
+                f"python -m metaflow_extensions.outerbounds.plugins.s3_proxy.proxy_bootstrap bootstrap --integration-name {integration_name} --write-mode {write_mode} --debug {debug} --uc-proxy-cfg-write-path ./.uc_proxy_cfg_file --proxy-status-write-path ./.proxy_status_file",
+                "export METAFLOW_S3_PROXY_USER_CODE_CONFIG=$(cat ./.uc_proxy_cfg_file)",
+                "export METAFLOW_S3_PROXY_STATUS=$(cat ./.proxy_status_file)",
+                "export METAFLOW_S3_PROXY_SETUP_SUCCESS=True",
+                "flush_mflogs",
+            ]
+        return base_boostrap_cmd + additional_cmd
+
+    environment.bootstrap_commands = wrapper
+
+
+class S3ProxyDecorator(StepDecorator):
+    """
+    Set up an S3 proxy that caches objects in an external, S3‑compatible bucket
+    for S3 read and write requests.
+
+    This decorator requires an integration in the Outerbounds platform that
+    points to an external bucket. It affects S3 operations performed via
+    Metaflow's `get_aws_client` and `S3` within a `@step`.
+
+    Read operations
+    ---------------
+    All read operations pass through the proxy. If an object does not already
+    exist in the external bucket, it is cached there. For example, if code reads
+    from buckets `FOO` and `BAR` using the `S3` interface, objects from both
+    buckets are cached in the external bucket.
+
+    During task execution, all S3‑related read requests are routed through the
+    proxy:
+        - If the object is present in the external object store, the proxy
+          streams it directly from there without accessing the requested origin
+          bucket.
+        - If the object is not present in the external storage, the proxy
+          fetches it from the requested bucket, caches it in the external
+          storage, and streams the response from the origin bucket.
+
+    Warning
+    -------
+    All READ operations (e.g., GetObject, HeadObject) pass through the external
+    bucket regardless of the bucket specified in user code. Even
+    `S3(run=self)` and `S3(s3root="mybucketfoo")` requests go through the
+    external bucket cache.
+
+    Write operations
+    ----------------
+    Write behavior is controlled by the `write_mode` parameter, which determines
+    whether writes also persist objects in the cache.
+
+    `write_mode` values:
+        - `origin-and-cache`: objects are written both to the cache and to their
+          intended origin bucket.
+        - `origin`: objects are written only to their intended origin bucket.
+
+    Parameters
+    ----------
+    integration_name : str, optional
+        [Outerbounds integration name](https://docs.outerbounds.com/outerbounds/configuring-secrets/#integrations-view)
+        that holds the configuration for the external, S3‑compatible object
+        storage bucket. If not specified, the only available S3 proxy
+        integration in the namespace is used (fails if multiple exist).
+    write_mode : str, optional
+        Controls whether writes also go to the external bucket.
+            - `origin` (default)
+            - `origin-and-cache`
+    debug : bool, optional
+        Enables debug logging for proxy operations.
+    """
+
+    name = "s3_proxy"
+    defaults = {
+        "integration_name": None,
+        "write_mode": None,
+        "debug": False,
+    }
+
+    _environment_patched = False
+
+    _proxy_status = None
+
+    @classmethod
+    def patch_environment(cls, flow, environment):
+        """
+        We need to patch the environment boostrap command so that
+        we can launch the s3 proxy before the step code execution.
+        We also want to ensure that we are running the proxy bootstrap
+        only for the steps that have the decorator set. This is why we pass down all
+        the step names that will change the boostrap commands.
+        """
+        if cls._environment_patched:
+            return
+
+        steps_with_s3_proxy = [
+            step
+            for step in flow
+            if any(deco.name == "s3_proxy" for deco in step.decorators)
+        ]
+        if len(steps_with_s3_proxy) == 0:  # weird but y not?
+            return
+
+        step_names_and_deco_attrs = {}
+        for s in steps_with_s3_proxy:
+            _decos = [x for x in s.decorators if x.name == "s3_proxy"]
+            deco = _decos[0]
+            step_names_and_deco_attrs[s.name] = S3ProxyBinaryConfig(
+                integration_name=deco.attributes["integration_name"],
+                write_mode=deco.attributes["write_mode"],
+                debug=deco.attributes["debug"],
+            )
+
+        monkey_patch_environment(environment, step_names_and_deco_attrs)
+        cls._environment_patched = True
+
+    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.patch_environment(flow, environment)
+        if (
+            os.environ.get("METAFLOW_S3_PROXY_USER_CODE_CONFIG")
+            and os.environ.get("METAFLOW_S3_PROXY_STATUS")
+            and self.attributes["debug"]
+        ):
+            print("[@s3_proxy] S3 Proxy detected. Debug mode is enabled.")
+
+        if os.environ.get("METAFLOW_S3_PROXY_STATUS"):
+            proxy_status = json.loads(os.environ.get("METAFLOW_S3_PROXY_STATUS"))
+            self._proxy_status = proxy_status
+
+    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"""
+        pass
+
+    def task_finished(
+        self, step_name, flow, graph, is_task_ok, retry_count, max_retries
+    ):
+        if not self._proxy_status:
+            return
+
+        status = self._proxy_status
+        proxy_pid = status.get("proxy_pid")
+        config_path = status.get("config_path")
+        binary_path = status.get("binary_path")
+
+        # 1) Stop processes: try to terminate the process group for clean child shutdown
+        if proxy_pid:
+            try:
+                pgid = os.getpgid(proxy_pid)
+                os.killpg(pgid, signal.SIGTERM)
+                time.sleep(1)
+            except Exception:
+                # Fall back to killing the pid directly if pgid is unavailable
+                try:
+                    os.kill(proxy_pid, signal.SIGTERM)
+                except Exception:
+                    pass
+
+        # 2) Clear files based on status
+        for path in (config_path, binary_path):
+            try:
+                if path and os.path.exists(path):
+                    os.remove(path)
+            except Exception:
+                pass
+
+
+class NebiusS3ProxyDecorator(S3ProxyDecorator):
+
+    __doc__ = (
+        """
+    `@nebius_s3_proxy` is a 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. The underlying mechanics of the decorator is the same as the `@s3_proxy`:\n
+    """
+        + S3ProxyDecorator.__doc__
+    )
+
+    name = "nebius_s3_proxy"
+    defaults = {
+        "integration_name": None,
+        "write_mode": None,
+        "debug": False,
+    }
+
+
+class CoreWeaveS3ProxyDecorator(S3ProxyDecorator):
+    __doc__ = (
+        """
+    `@coreweave_s3_proxy` is a 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. The underlying mechanics of the decorator is the same as the `@s3_proxy`:\n
+    """
+        + S3ProxyDecorator.__doc__
+    )
+
+    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,225 @@
+import os
+import json
+import gzip
+import sys
+import time
+import threading
+import subprocess
+from pathlib import Path
+from typing import Optional, Tuple
+
+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) -> Tuple[dict, int, str, str]:
+        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)
+                # In the new world where the binary is being called
+                # before even the metaflow code exection starts,
+                # so this implies a few important things:
+                # 1, We start the actual proxy process via another python file that safely ships logs to mflog.
+                # 2. We passback the right values to the metaflow step process via env vars.
+                # 3. Metaflow step code relies on env vars to decide if clients need to have s3 proxy in them.
+                self.process = self._start_proxy_process()
+
+                user_code_proxy_config = self._setup_proxy_config(config_data)
+
+                return_tuple = (
+                    user_code_proxy_config,  # this is the config that will be used within the metaflow `step` code.
+                    self.process.pid,  # This is the pid of the process that will jumpstart, monitor and ship logs to MFLOG for the proxy process
+                    self.config_path,  # This is the path to the config that is derived from the integration. It contains the actual bucket path and name where external objects are stored.
+                    self.binary_path,  # This is the path to the binary for the proxy.
+                )
+                # We return a tuple because these values need to be passed down to the metaflow step process where
+                # it will handle thier removal gracefully after the step is finished.
+                return return_tuple
+
+            print(
+                "[@s3_proxy] skipping s3-proxy set up because metaflow has not detected a Kubernetes environment"
+            )
+            raise S3ProxyException(
+                "S3 proxy setup failed because metaflow has not detected a Kubernetes environment"
+            )
+        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:
+        # This command will jump start a process that will then call the proxy binary
+        # The reason we do something like this is because we need to run all of this before
+        # even the `step` command is called. So we need a python process that will ship the logs
+        # of the proxy process to MFLOG instead of setting print statements. We need this process
+        # to run independently since the S3ProxyManager gets called in the boostrap_proxy which will
+        # exit after jump starting the proxy process.
+        cmd = [self.binary_path, "--bucket-config", self.config_path, "serve"]
+        _env = os.environ.copy()
+        _env["S3_PROXY_BINARY_COMMAND"] = " ".join(cmd)
+        if self.debug:
+            _env["S3_PROXY_BINARY_DEBUG"] = "True"
+        _cmd = [
+            sys.executable,
+            "-m",
+            "metaflow_extensions.outerbounds.plugins.s3_proxy.binary_caller",
+        ]
+        devnull = subprocess.DEVNULL
+        process = subprocess.Popen(
+            _cmd,
+            stdout=devnull,
+            stderr=devnull,
+            text=True,
+            start_new_session=True,
+            env=_env,
+        )
+        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_proxy_config(self, config_data):
+        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
+
+        self.proxy_config = proxy_config
+        return proxy_config
+
+    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/plugins/secrets/secrets.py

@@ -4,13 +4,48 @@ from typing import Dict
 import base64
 import json
 import requests
+import random
 import time
+import sys
 
 
 class OuterboundsSecretsException(Exception):
     pass
 
 
+def _api_server_get(*args, conn_error_retries=2, **kwargs):
+    """
+    There are two categories of errors that we need to handle when dealing with any API server.
+    1. HTTP errors. These are are errors that are returned from the API server.
+        - 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.
+
+    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.
+    Generally this is a safe enough number of retries after which we can assume that something is really broken. Until then,
+    there can be intermittent issues that would resolve themselves if we retry gracefully.
+    """
+    _num_retries = 0
+    noise = random.uniform(-0.5, 0.5)
+    while _num_retries < conn_error_retries:
+        try:
+            return requests.get(*args, **kwargs)
+        except requests.exceptions.ConnectionError:
+            if _num_retries <= conn_error_retries - 1:
+                # Exponential backoff with 2^(_num_retries+1) seconds
+                time.sleep((2 ** (_num_retries + 1)) + noise)
+                _num_retries += 1
+            else:
+                print(
+                    "[@secrets] Failed to connect to the API server. ",
+                    file=sys.stderr,
+                )
+                raise
+
+
 class OuterboundsSecretsApiResponse:
     def __init__(self, response):
         self.response = response
@@ -133,7 +168,9 @@ class OuterboundsSecretsProvider(SecretsProvider):
         retryable_status_codes = [409]
         json_payload = json.dumps(payload)
         for attempt in range(2):  # 0 = initial attempt, 1-2 = retries
-            response = requests.get(url, data=json_payload, headers=request_headers)
+            response = _api_server_get(
+                url, data=json_payload, headers=request_headers, conn_error_retries=5
+            )
             if response.status_code not in retryable_status_codes:
                 break
 
@@ -141,7 +178,6 @@ class OuterboundsSecretsProvider(SecretsProvider):
                 sleep_time = 0.5 * (attempt + 1)
                 time.sleep(sleep_time)
 
-        response = requests.get(url, data=json_payload, headers=request_headers)
         self._handle_error_response(response)
         return OuterboundsSecretsApiResponse(response.json())