ob-metaflow-extensions 1.4.17__tar.gz → 1.4.19__tar.gz

{ob_metaflow_extensions-1.4.17 → ob_metaflow_extensions-1.4.19}/PKG-INFO

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

{ob_metaflow_extensions-1.4.17 → ob_metaflow_extensions-1.4.19}/metaflow_extensions/outerbounds/plugins/fast_bakery/docker_environment.py

@@ -254,7 +254,8 @@ class DockerEnvironment(MetaflowEnvironment):
                 (d for d in step.decorators if _is_env_deco(d)),
                 None,
             )
-            python = next(
+            python = None
+            conda_set_python = next(
                 (
                     d.attributes["python"]
                     for d in step.decorators
@@ -262,17 +263,25 @@ class DockerEnvironment(MetaflowEnvironment):
                 ),
                 None,
             )
-            pypi_deco = next(
-                (d for d in step.decorators if isinstance(d, PyPIStepDecorator)), None
+            pypi_set_python = next(
+                (
+                    d.attributes["python"]
+                    for d in step.decorators
+                    if isinstance(d, PyPIStepDecorator)
+                ),
+                None,
             )
-            # if pypi decorator is set and user has specified a python version, we must create a new environment.
-            # otherwise rely on the base environment
-            if pypi_deco is not None:
-                python = (
-                    pypi_deco.attributes["python"]
-                    if pypi_deco.is_attribute_user_defined("python")
-                    else None
-                )
+            # Since the introduction of custom decorators and new packaging support in Metaflow
+            # @pypi/@pypi_base etc can be set at runtime and can also be overriden at runtime.
+            # So at time of execution here, we can assume the python version set in the decorator land
+            # is the one that has been finally resolved. So we just coalesce the conda_set_python and pypi_set_python.
+            def coalesce(*args):
+                for arg in args:
+                    if arg is not None:
+                        return arg
+                return None
+
+            python = coalesce(conda_set_python, pypi_set_python)
 
             packages = get_pinned_conda_libs(python, self.datastore_type)
             packages.update(dependencies.attributes["packages"] if dependencies else {})

ob_metaflow_extensions-1.4.19/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)

{ob_metaflow_extensions-1.4.17 → ob_metaflow_extensions-1.4.19}/metaflow_extensions/outerbounds/plugins/s3_proxy/constants.py

@@ -5,4 +5,7 @@ S3_PROXY_BINARY_URLS = {
 
 DEFAULT_PROXY_PORT = 8081
 DEFAULT_PROXY_HOST = "localhost"
-S3_PROXY_WRITE_MODES = ["origin-and-cache", "origin", "cache"]
+S3_PROXY_WRITE_MODES = [
+    "origin-and-cache",
+    "origin",
+]

ob_metaflow_extensions-1.4.19/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()

ob_metaflow_extensions-1.4.19/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,
+    }

{ob_metaflow_extensions-1.4.17 → ob_metaflow_extensions-1.4.19}/metaflow_extensions/outerbounds/plugins/s3_proxy/s3_proxy_manager.py

@@ -1,11 +1,12 @@
 import os
 import json
 import gzip
+import sys
 import time
 import threading
 import subprocess
 from pathlib import Path
-from typing import Optional
+from typing import Optional, Tuple
 
 import requests
 
@@ -35,7 +36,7 @@ class S3ProxyManager:
         self.api_client = S3ProxyApiClient()
         self.proxy_config = None
 
-    def setup_proxy(self) -> bool:
+    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(
@@ -43,14 +44,32 @@ class S3ProxyManager:
                 )
                 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()
-                self._setup_proxy_config(config_data)
-                return True
+
+                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"
             )
-            return False
+            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}")
@@ -121,21 +140,31 @@ class S3ProxyManager:
         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:
-            print(f"[@s3_proxy] Starting proxy: {' '.join(cmd)}")
-
+            _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=subprocess.PIPE,
-            stderr=subprocess.STDOUT,  # Redirect stderr to stdout
+            _cmd,
+            stdout=devnull,
+            stderr=devnull,
             text=True,
             start_new_session=True,
+            env=_env,
         )
-
-        self._setup_log_streaming(process)
-
         time.sleep(3)
 
         if process.poll() is None:
@@ -149,30 +178,7 @@ class S3ProxyManager:
                 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(
@@ -190,11 +196,8 @@ class S3ProxyManager:
         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")
+        return proxy_config
 
     def cleanup(self):
         try: