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

metaflow_extensions/outerbounds/plugins/s3_proxy/binary_caller.py

@@ -0,0 +1,132 @@
+import sys
+import os
+import subprocess
+from metaflow.mflog.mflog import decorate
+from metaflow.mflog import TASK_LOG_SOURCE
+from typing import Union, TextIO, BinaryIO, Callable, Optional
+from queue import Queue, Empty
+from concurrent.futures import ThreadPoolExecutor
+import subprocess
+
+
+def enqueue_output(file, queue):
+    for line in iter(file.readline, ""):
+        queue.put(line)
+    file.close()
+
+
+def read_popen_pipes(p: subprocess.Popen):
+
+    with ThreadPoolExecutor(2) as pool:
+        q_stdout, q_stderr = Queue(), Queue()
+        pool.submit(enqueue_output, p.stdout, q_stdout)
+        pool.submit(enqueue_output, p.stderr, q_stderr)
+        while True:
+
+            if p.poll() is not None and q_stdout.empty() and q_stderr.empty():
+                break
+
+            out_line = err_line = ""
+
+            try:
+                out_line = q_stdout.get_nowait()
+            except Empty:
+                pass
+            try:
+                err_line = q_stderr.get_nowait()
+            except Empty:
+                pass
+
+            yield (out_line, err_line)
+
+
+class LogBroadcaster:
+    def __init__(
+        self,
+        process: subprocess.Popen,
+    ):
+        self._process = process
+        self._file_descriptors_and_parsers = []
+
+    def add_channel(
+        self, file_path: str, parser: Optional[Callable[[str], str]] = None
+    ):
+        self._file_descriptors_and_parsers.append((open(file_path, "a"), parser))
+
+    def _broadcast_lines(self, out_line: str, err_line: str):
+        for file_descriptor, parser in self._file_descriptors_and_parsers:
+            if out_line != "":
+                if parser:
+                    out_line = parser(out_line)
+                print(out_line, file=file_descriptor, end="", flush=True)
+            if err_line != "":
+                if parser:
+                    err_line = parser(err_line)
+                print(err_line, file=file_descriptor, end="", flush=True)
+
+    def publish_line(self, out_line: str, err_line: str):
+        self._broadcast_lines(out_line, err_line)
+
+    def broadcast_logs_to_files(self):
+        for out_line, err_line in read_popen_pipes(self._process):
+            self._broadcast_lines(out_line, err_line)
+
+        self._process.wait()
+
+        for file_descriptor, _ in self._file_descriptors_and_parsers:
+            file_descriptor.close()
+
+
+def run_with_mflog_capture(command, debug=False):
+    """
+    Run a subprocess with proper mflog integration for stdout/stderr capture.
+    This mimics what bash_capture_logs does but in Python.
+    """
+    # Get the log file paths from environment variables
+    stdout_path = os.environ.get("MFLOG_STDOUT")
+    stderr_path = os.environ.get("MFLOG_STDERR")
+
+    if not stdout_path or not stderr_path:
+        # Fallback to regular subprocess if mflog env vars aren't set
+        return subprocess.run(command, check=True, shell=True)
+
+    pipe = subprocess.PIPE if debug else subprocess.DEVNULL
+    # Start the subprocess with pipes
+    process = subprocess.Popen(
+        command,
+        shell=True,
+        stdout=pipe,
+        stderr=pipe,
+        text=False,  # Use bytes for proper mflog handling
+        bufsize=0,  # Unbuffered for real-time logging
+    )
+
+    broadcaster = LogBroadcaster(process)
+
+    broadcaster.add_channel(
+        stderr_path, lambda line: decorate(TASK_LOG_SOURCE, line).decode("utf-8")
+    )
+    broadcaster.publish_line(f"[S3 PROXY] Starting Fast S3 Proxy.....\n", "")
+    broadcaster.broadcast_logs_to_files()
+
+    # Check the return code and raise if non-zero
+    if process.returncode != 0:
+        raise subprocess.CalledProcessError(process.returncode, command)
+
+    return process.returncode
+
+
+if __name__ == "__main__":
+    s3_proxy_binary_path = os.environ.get("S3_PROXY_BINARY_COMMAND")
+    s3_proxy_debug = bool(os.environ.get("S3_PROXY_BINARY_DEBUG", False))
+    if not s3_proxy_binary_path:
+        print("S3_PROXY_BINARY_COMMAND environment variable not set")
+        sys.exit(1)
+
+    try:
+        run_with_mflog_capture(s3_proxy_binary_path, debug=s3_proxy_debug)
+    except subprocess.CalledProcessError as e:
+        sys.exit(e.returncode)
+    except Exception as e:
+        print(f"Error running S3 proxy binary: {e}", file=sys.stderr)
+        sys.exit(1)

metaflow_extensions/outerbounds/plugins/s3_proxy/constants.py

@@ -0,0 +1,11 @@
+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",
+]

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/proxy_bootstrap.py

@@ -0,0 +1,59 @@
+from .s3_proxy_manager import S3ProxyManager
+from metaflow._vendor import click
+from metaflow import JSONType
+import json
+
+
+@click.group()
+def cli():
+    pass
+
+
+@cli.command()
+@click.option(
+    "--integration-name", type=str, help="The integration name", required=True
+)
+@click.option("--write-mode", type=str, help="The write mode")
+@click.option("--debug", type=bool, help="The debug mode", default=False)
+@click.option(
+    "--uc-proxy-cfg-write-path",
+    type=str,
+    help="The path to write the user code proxy config",
+    required=True,
+)
+@click.option(
+    "--proxy-status-write-path",
+    type=str,
+    help="The path to write the proxy status",
+    required=True,
+)
+def bootstrap(
+    integration_name,
+    write_mode,
+    debug,
+    uc_proxy_cfg_write_path,
+    proxy_status_write_path,
+):
+    manager = S3ProxyManager(
+        integration_name=integration_name,
+        write_mode=write_mode,
+        debug=debug,
+    )
+    user_code_proxy_config, proxy_pid, config_path, binary_path = manager.setup_proxy()
+    with open(uc_proxy_cfg_write_path, "w") as f:
+        f.write(json.dumps(user_code_proxy_config))
+    with open(proxy_status_write_path, "w") as f:
+        f.write(
+            json.dumps(
+                {
+                    "proxy_pid": proxy_pid,
+                    "config_path": config_path,
+                    "binary_path": binary_path,
+                }
+            )
+        )
+
+
+if __name__ == "__main__":
+    print("[@s3_proxy] Jumpstarting the proxy....")
+    cli()

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,
+    }