ob-metaflow-extensions 1.4.14__tar.gz → 1.4.30__tar.gz

{ob_metaflow_extensions-1.4.14 → ob_metaflow_extensions-1.4.30}/PKG-INFO

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

{ob_metaflow_extensions-1.4.14 → ob_metaflow_extensions-1.4.30}/metaflow_extensions/outerbounds/plugins/apps/core/app_config.py

@@ -45,10 +45,11 @@ def _try_loading_yaml(file):
 class AuthType:
     BROWSER = "Browser"
     API = "API"
+    BROWSER_AND_API = "BrowserAndApi"
 
     @classmethod
     def enums(cls):
-        return [cls.BROWSER, cls.API]
+        return [cls.BROWSER, cls.API, cls.BROWSER_AND_API]
 
     @classproperty
     def default(cls):

{ob_metaflow_extensions-1.4.14 → ob_metaflow_extensions-1.4.30}/metaflow_extensions/outerbounds/plugins/apps/core/capsule.py

@@ -678,7 +678,7 @@ class CapsuleDeployer:
         auth_type = self._app_config.get_state("auth", {}).get("type", AuthType.default)
         if auth_type == AuthType.BROWSER:
             return "App"
-        elif auth_type == AuthType.API:
+        elif auth_type == AuthType.API or auth_type == AuthType.BROWSER_AND_API:
             return "Endpoint"
         else:
             raise TODOException(f"Unknown auth type: {auth_type}")

{ob_metaflow_extensions-1.4.14 → ob_metaflow_extensions-1.4.30}/metaflow_extensions/outerbounds/plugins/apps/core/code_package/code_packager.py

@@ -381,6 +381,7 @@ class CodePackager:
                     "google-auth",
                     "simple-gcp-object-downloader",
                     "google-cloud-secret-manager",
+                    "packaging",
                 ],
                 "local": [],
             }

{ob_metaflow_extensions-1.4.14 → ob_metaflow_extensions-1.4.30}/metaflow_extensions/outerbounds/plugins/apps/core/config/typed_configs.py

@@ -97,6 +97,7 @@ class TypedCoreConfig:
         branch: Optional[str] = None,
         models: Optional[list] = None,
         data: Optional[list] = None,
+        generate_static_url: Optional[bool] = None,
         **kwargs
     ) -> None:
         self._kwargs = {
@@ -122,6 +123,7 @@ class TypedCoreConfig:
             "branch": branch,
             "models": models,
             "data": data,
+            "generate_static_url": generate_static_url,
         }
         # Add any additional kwargs
         self._kwargs.update(kwargs)

{ob_metaflow_extensions-1.4.14 → ob_metaflow_extensions-1.4.30}/metaflow_extensions/outerbounds/plugins/apps/core/config/unified_config.py

@@ -36,10 +36,11 @@ from .config_utils import (
 class AuthType:
     BROWSER = "Browser"
     API = "API"
+    BROWSER_AND_API = "BrowserAndApi"
 
     @classmethod
     def choices(cls):
-        return [cls.BROWSER, cls.API]
+        return [cls.BROWSER, cls.API, cls.BROWSER_AND_API]
 
 
 class UnitParser:
@@ -966,6 +967,17 @@ How to read this schema:
         is_experimental=True,
         example=[{"asset_id": "data-789", "asset_instance_id": "instance-101"}],
     )
+    generate_static_url = ConfigField(
+        cli_meta=CLIOption(
+            name="generate_static_url",
+            cli_option_str="--generate-static-url",
+            help="Generate a static URL for the app based on its name.",
+            is_flag=True,
+        ),
+        field_type=bool,
+        default=False,
+        help="Generate a static URL for the app based on its name.",
+    )
     # ------- /Experimental -------------
 
     def to_dict(self):

{ob_metaflow_extensions-1.4.14 → ob_metaflow_extensions-1.4.30}/metaflow_extensions/outerbounds/plugins/apps/core/config_schema.yaml

@@ -149,6 +149,7 @@ properties:
         enum:
         - Browser
         - API
+        - BrowserAndApi
         example: Browser
         mutation_behavior: union
         cli_option: --auth-type
@@ -328,3 +329,9 @@ properties:
       asset_instance_id: instance-101
     experimental: true
     mutation_behavior: union
+  generate_static_url:
+    description: Generate a static URL for the app based on its name.
+    type: boolean
+    default: false
+    mutation_behavior: union
+    cli_option: --generate-static-url

{ob_metaflow_extensions-1.4.14 → ob_metaflow_extensions-1.4.30}/metaflow_extensions/outerbounds/plugins/apps/core/experimental/__init__.py

@@ -21,6 +21,10 @@ def capsule_input_overrides(app_config: "AppConfig", capsule_input: dict):
     if persistence is not None and persistence != "none":
         capsule_input["persistence"] = persistence
 
+    capsule_input["generateStaticUrl"] = app_config.get_state(
+        "generate_static_url", False
+    )
+
     model_asset_conf = app_config.get_state("models", None)
     data_asset_conf = app_config.get_state("data", None)
     code_info = _code_info(app_config)

ob_metaflow_extensions-1.4.30/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.14 → ob_metaflow_extensions-1.4.30}/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.30/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.30/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,
+    }