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

metaflow_extensions/outerbounds/plugins/fast_bakery/baker.py

@@ -0,0 +1,110 @@
+import threading
+import time
+import sys
+from typing import Dict, Optional, Any, Callable
+from functools import partial
+from metaflow.exception import MetaflowException
+from metaflow.metaflow_config import FAST_BAKERY_URL
+
+from .fast_bakery import FastBakery, FastBakeryApiResponse, FastBakeryException
+from .docker_environment import cache_request
+
+BAKERY_METAFILE = ".imagebakery-cache"
+
+
+class BakerException(MetaflowException):
+    headline = "Ran into an error while baking image"
+
+    def __init__(self, msg):
+        super(BakerException, self).__init__(msg)
+
+
+def bake_image(
+    cache_file_path: str,
+    ref: Optional[str] = None,
+    python: Optional[str] = None,
+    pypi_packages: Optional[Dict[str, str]] = None,
+    conda_packages: Optional[Dict[str, str]] = None,
+    base_image: Optional[str] = None,
+    logger: Optional[Callable[[str], Any]] = None,
+) -> FastBakeryApiResponse:
+    """
+    Bakes a Docker image with the specified dependencies.
+
+    Args:
+        cache_file_path: Path to the cache file
+        ref: Reference identifier for this bake (for logging purposes)
+        python: Python version to use
+        pypi_packages: Dictionary of PyPI packages and versions
+        conda_packages: Dictionary of Conda packages and versions
+        base_image: Base Docker image to use
+        logger: Optional logger function to output progress
+
+    Returns:
+        FastBakeryApiResponse: The response from the bakery service
+
+    Raises:
+        BakerException: If the baking process fails
+    """
+    # Default logger if none provided
+    if logger is None:
+        logger = partial(print, file=sys.stderr)
+
+    # Thread lock for logging
+    logger_lock = threading.Lock()
+    images_baked = 0
+
+    @cache_request(cache_file_path)
+    def _cached_bake(
+        ref=None,
+        python=None,
+        pypi_packages=None,
+        conda_packages=None,
+        base_image=None,
+    ):
+        try:
+            bakery = FastBakery(url=FAST_BAKERY_URL)
+            bakery._reset_payload()
+            bakery.python_version(python)
+            bakery.pypi_packages(pypi_packages)
+            bakery.conda_packages(conda_packages)
+            bakery.base_image(base_image)
+            # bakery.ignore_cache()
+
+            with logger_lock:
+                logger(f"🍳 Baking [{ref}] ...")
+                logger(f"     🐍 Python: {python}")
+
+                if pypi_packages:
+                    logger(f"     📦 PyPI packages:")
+                    for package, version in pypi_packages.items():
+                        logger(f"        🔧 {package}: {version}")
+
+                if conda_packages:
+                    logger(f"     📦 Conda packages:")
+                    for package, version in conda_packages.items():
+                        logger(f"        🔧 {package}: {version}")
+
+                logger(f"     🏗️  Base image: {base_image}")
+
+            start_time = time.time()
+            res = bakery.bake()
+            # TODO: Get actual bake time from bakery
+            bake_time = time.time() - start_time
+
+            with logger_lock:
+                logger(f"🏁 Baked [{ref}] in {bake_time:.2f} seconds!")
+            nonlocal images_baked
+            images_baked += 1
+            return res
+        except FastBakeryException as ex:
+            raise BakerException(f"Bake [{ref}] failed: {str(ex)}")
+
+    # Call the cached bake function with the provided parameters
+    return _cached_bake(
+        ref=ref,
+        python=python,
+        pypi_packages=pypi_packages,
+        conda_packages=conda_packages,
+        base_image=base_image,
+    )

metaflow_extensions/outerbounds/plugins/fast_bakery/docker_environment.py

@@ -351,12 +351,16 @@ class DockerEnvironment(MetaflowEnvironment):
         config.append("--disable=F0401")
         return config
 
-    def get_package_commands(self, codepackage_url, datastore_type):
+    def get_package_commands(
+        self, codepackage_url, datastore_type, code_package_metadata=None
+    ):
         # we must set the skip install flag at this stage in order to skip package downloads,
         # doing so in bootstrap_commands is too late in the lifecycle.
         return [
             "export METAFLOW_SKIP_INSTALL_DEPENDENCIES=$FASTBAKERY_IMAGE",
-        ] + super().get_package_commands(codepackage_url, datastore_type)
+        ] + super().get_package_commands(
+            codepackage_url, datastore_type, code_package_metadata=code_package_metadata
+        )
 
     def bootstrap_commands(self, step_name, datastore_type):
         if step_name in self.skipped_steps:

metaflow_extensions/outerbounds/plugins/fast_bakery/fast_bakery.py

@@ -122,6 +122,7 @@ class FastBakery:
             "responseMaxAgeSeconds": 0,
             "layerMaxAgeSeconds": 0,
             "baseImageMaxAgeSeconds": 0,
+            "overwriteExistingLayers": True,  # Used primarily to rewrite possibly corrupted layers.
         }
         return self
 

metaflow_extensions/outerbounds/plugins/nvct/nvct_decorator.py

@@ -46,23 +46,23 @@ SUPPORTABLE_GPU_TYPES = {
     "H100": [
         {
             "n_gpus": 1,
-            "instance_type": "GCP.GPU.H100_1x",
-            "backend": "gcp-asia-se-1a",
+            "instance_type": "OCI.GPU.H100_1x",
+            "backend": "nvcf-dgxc-k8s-oci-nrt-prd8",
         },
         {
             "n_gpus": 2,
-            "instance_type": "GCP.GPU.H100_2x",
-            "backend": "gcp-asia-se-1a",
+            "instance_type": "OCI.GPU.H100_2x",
+            "backend": "nvcf-dgxc-k8s-oci-nrt-prd8",
         },
         {
             "n_gpus": 4,
-            "instance_type": "GCP.GPU.H100_4x",
-            "backend": "gcp-asia-se-1a",
+            "instance_type": "OCI.GPU.H100_4x",
+            "backend": "nvcf-dgxc-k8s-oci-nrt-prd8",
         },
         {
             "n_gpus": 8,
-            "instance_type": "GCP.GPU.H100_8x",
-            "backend": "gcp-asia-se-1a",
+            "instance_type": "OCI.GPU.H100_8x",
+            "backend": "nvcf-dgxc-k8s-oci-nrt-prd8",
         },
     ],
     "NEBIUS_H100": [

metaflow_extensions/outerbounds/plugins/optuna/__init__.py

@@ -0,0 +1,48 @@
+import os
+import json
+
+__mf_promote_submodules__ = ["plugins.optuna"]
+
+
+def auth():
+    from metaflow.metaflow_config_funcs import init_config
+
+    conf = init_config()
+    if conf:
+        headers = {"x-api-key": conf["METAFLOW_SERVICE_AUTH_KEY"]}
+    else:
+        headers = json.loads(os.environ["METAFLOW_SERVICE_HEADERS"])
+    return headers
+
+
+def get_deployment_db_access_endpoint(name: str):
+    from ..apps.core.perimeters import PerimeterExtractor
+    from ..apps.core.capsule import CapsuleApi
+
+    perimeter, cap_url = PerimeterExtractor.during_metaflow_execution()
+    deployment = CapsuleApi(cap_url, perimeter).get_by_name(name)
+    if not deployment:
+        raise Exception(f"No app deployment found with name `{name}`")
+
+    if (
+        "status" in deployment
+        and "accessInfo" in deployment["status"]
+        and "extraAccessUrls" in deployment["status"]["accessInfo"]
+    ):
+        for extra_url in deployment["status"]["accessInfo"]["extraAccessUrls"]:
+            if extra_url["name"] == "in_cluster_db_access":
+                db_url = extra_url["url"].replace("http://", "")
+                return db_url
+
+    raise Exception(f"No db access endpoint found for deployment `{name}`")
+
+
+def get_db_url(app_name: str):
+    """
+    Example usage:
+        >>> from metaflow.plugins.optuna import get_db_url
+        >>> s = optuna.create_study(..., storage=get_db_url("optuna-dashboard"))
+    """
+    mf_token = auth()["x-api-key"]
+    app_url = get_deployment_db_access_endpoint(app_name)
+    return f"postgresql://userspace_default:{mf_token}@{app_url}/userspace_default?sslmode=disable"

metaflow_extensions/outerbounds/plugins/profilers/simple_card_decorator.py

@@ -0,0 +1,96 @@
+from datetime import datetime
+from metaflow.decorators import StepDecorator
+from ..card_utilities.injector import CardDecoratorInjector
+
+
+class DynamicCardAppendDecorator(StepDecorator):
+    """
+    A simple decorator that demonstrates using CardDecoratorInjector
+    to inject a card and render simple markdown content.
+    """
+
+    name = "test_append_card"
+
+    defaults = {
+        "title": "Simple Card",
+        "message": "Hello from DynamicCardAppendDecorator!",
+        "show_timestamp": True,
+        "refresh_interval": 5,
+    }
+
+    CARD_ID = "simple_card"
+
+    def step_init(
+        self, flow, graph, step_name, decorators, environment, flow_datastore, logger
+    ):
+        """Initialize the decorator and inject the card."""
+        self.deco_injector = CardDecoratorInjector()
+        self.deco_injector.attach_card_decorator(
+            flow,
+            step_name,
+            self.CARD_ID,
+            "blank",
+            refresh_interval=self.attributes["refresh_interval"],
+        )
+
+    def task_decorate(
+        self, step_func, flow, graph, retry_count, max_user_code_retries, ubf_context
+    ):
+        """Decorate the step function to add card content."""
+        from metaflow import current
+        from metaflow.cards import Markdown
+
+        # Create the card content
+        title = self.attributes["title"]
+        message = self.attributes["message"]
+        show_timestamp = self.attributes["show_timestamp"]
+
+        # Add title to the card
+        current.card[self.CARD_ID].append(Markdown(f"# {title}"))
+
+        # Add message to the card
+        current.card[self.CARD_ID].append(Markdown(f"**Message:** {message}"))
+
+        # Add timestamp if requested
+        if show_timestamp:
+            timestamp = datetime.now().astimezone().strftime("%Y-%m-%d %H:%M:%S %z")
+            current.card[self.CARD_ID].append(Markdown(f"**Created at:** {timestamp}"))
+
+        # Add step information
+        current.card[self.CARD_ID].append(Markdown(f"**Step:** `{current.pathspec}`"))
+
+        # Add a simple divider
+        current.card[self.CARD_ID].append(Markdown("---"))
+
+        # Add some dynamic content that shows this is working
+        current.card[self.CARD_ID].append(
+            Markdown("**Status:** Card successfully injected! 🎉")
+        )
+
+        def wrapped_step_func():
+            """Execute the original step function."""
+            try:
+                # Before execution
+                current.card[self.CARD_ID].append(
+                    Markdown("**Execution:** Step started...")
+                )
+                current.card[self.CARD_ID].refresh()
+
+                # Execute the original step
+                step_func()
+
+                # After execution
+                current.card[self.CARD_ID].append(
+                    Markdown("**Execution:** Step completed successfully! ✅")
+                )
+                current.card[self.CARD_ID].refresh()
+
+            except Exception as e:
+                # Handle errors
+                current.card[self.CARD_ID].append(
+                    Markdown(f"**Error:** Step failed with error: `{str(e)}` ❌")
+                )
+                current.card[self.CARD_ID].refresh()
+                raise
+
+        return wrapped_step_func

metaflow_extensions/outerbounds/plugins/s3_proxy/__init__.py

@@ -0,0 +1,7 @@
+from .s3_proxy_decorator import (
+    S3ProxyDecorator,
+    NebiusS3ProxyDecorator,
+    CoreWeaveS3ProxyDecorator,
+)
+
+__all__ = ["S3ProxyDecorator", "NebiusS3ProxyDecorator", "CoreWeaveS3ProxyDecorator"]

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()