ob-metaflow-extensions 1.1.170__py2.py3-none-any.whl → 1.4.35__py2.py3-none-any.whl

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/snowpark/snowpark_client.py

@@ -27,9 +27,12 @@ class SnowparkClient(object):
         except (NameError, ImportError, ModuleNotFoundError):
             raise SnowflakeException(
                 "Could not import module 'snowflake'.\n\nInstall Snowflake "
-                "Python package (https://pypi.org/project/snowflake/) first.\n"
-                "You can install the module by executing - "
-                "%s -m pip install snowflake\n"
+                "Python packages first:\n"
+                "  snowflake==1.8.0\n"
+                "  snowflake-connector-python==3.18.0\n"
+                "  snowflake-snowpark-python==1.40.0\n\n"
+                "You can install them by executing:\n"
+                "%s -m pip install snowflake==1.8.0 snowflake-connector-python==3.18.0 snowflake-snowpark-python==1.40.0\n"
                 "or equivalent through your favorite Python package manager."
                 % sys.executable
             )

metaflow_extensions/outerbounds/plugins/snowpark/snowpark_decorator.py

@@ -42,10 +42,13 @@ class Snowflake(object):
             return session
         except (NameError, ImportError, ModuleNotFoundError):
             raise SnowflakeException(
-                "Could not import module 'snowflake'.\n\nInstall Snowflake "
-                "Python package (https://pypi.org/project/snowflake/) first.\n"
-                "You can install the module by using the @pypi decorator - "
-                "Eg: @pypi(packages={'snowflake': '0.11.0'})\n"
+                "Could not import module 'snowflake'.\n\n"
+                "Install required Snowflake packages using the @pypi decorator:\n"
+                "@pypi(packages={\n"
+                "    'snowflake': '1.8.0',\n"
+                "    'snowflake-connector-python': '3.18.0',\n"
+                "    'snowflake-snowpark-python': '1.40.0'\n"
+                "})\n"
             )
 
 
@@ -143,9 +146,12 @@ class SnowparkDecorator(StepDecorator):
         except (NameError, ImportError, ModuleNotFoundError):
             raise SnowflakeException(
                 "Could not import module 'snowflake'.\n\nInstall Snowflake "
-                "Python package (https://pypi.org/project/snowflake/) first.\n"
-                "You can install the module by executing - "
-                "%s -m pip install snowflake\n"
+                "Python packages first:\n"
+                "  snowflake==1.8.0\n"
+                "  snowflake-connector-python==3.18.0\n"
+                "  snowflake-snowpark-python==1.40.0\n\n"
+                "You can install them by executing:\n"
+                "%s -m pip install snowflake==1.8.0 snowflake-connector-python==3.18.0 snowflake-snowpark-python==1.40.0\n"
                 "or equivalent through your favorite Python package manager."
                 % sys.executable
             )

metaflow_extensions/outerbounds/plugins/snowpark/snowpark_job.py

@@ -199,11 +199,17 @@ class RunningJob(object):
 
     @property
     def status(self):
-        return self.status_obj()[0].get("status")
+        status_list = self.status_obj()
+        if not status_list:
+            return "UNKNOWN"
+        return status_list[0].get("status", "UNKNOWN")
 
     @property
     def message(self):
-        return self.status_obj()[0].get("message")
+        status_list = self.status_obj()
+        if not status_list:
+            return None
+        return status_list[0].get("message")
 
     @property
     def is_waiting(self):

metaflow_extensions/outerbounds/plugins/torchtune/__init__.py

@@ -13,12 +13,14 @@ class TorchTune:
     def __init__(
         self,
         use_multi_node_config: bool = False,
+        config_overrides: Optional[Dict] = None,
     ):
         """
         Initialize the Tune launcher.
 
         :param use_multi_node_config: If True, attempt to build a distributed configuration
                                       from current.torch.torchrun_args.
+        :param config_overrides: Optional dictionary of config overrides for tune run.
         """
         self.multi_node_config = {}
         if use_multi_node_config:
@@ -37,6 +39,8 @@ class TorchTune:
                     "num_processes": current.torch.torchrun_args["nproc_per_node"]
                     * current.torch.torchrun_args["nnodes"],
                 }
+                if config_overrides:
+                    self.multi_node_config.update(config_overrides)
                 print(
                     f"[Metaflow Tune] Discovered multi-node config for torchrun: {self.multi_node_config}"
                 )