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

metaflow_extensions/outerbounds/plugins/ollama/status_card.py

@@ -0,0 +1,292 @@
+from metaflow.cards import Markdown, Table, VegaChart
+from metaflow.metaflow_current import current
+from datetime import datetime
+import threading
+import time
+
+from ..card_utilities.async_cards import CardRefresher
+
+
+class OllamaStatusCard(CardRefresher):
+    """
+    Real-time status card for Ollama system monitoring.
+    Shows circuit breaker state, server health, model status, and recent events.
+    """
+
+    CARD_ID = "ollama_status"
+
+    def __init__(self, refresh_interval=10):
+        self.refresh_interval = refresh_interval
+        self.status_data = {
+            "circuit_breaker": {
+                "state": "CLOSED",
+                "failure_count": 0,
+                "last_failure_time": None,
+                "last_open_time": None,
+            },
+            "server": {
+                "status": "Starting",
+                "uptime_start": None,
+                "restart_attempts": 0,
+                "last_health_check": None,
+                "health_status": "Unknown",
+            },
+            "models": {},  # model_name -> {status, pull_time, load_time, etc}
+            "performance": {
+                "install_time": None,
+                "server_startup_time": None,
+                "total_initialization_time": None,
+            },
+            "versions": {
+                "ollama_system": "Detecting...",
+                "ollama_python": "Detecting...",
+            },
+            "cache": {
+                "policy": "auto",
+                "model_status": {},  # model_name -> cache status
+            },
+            "events": [],  # Recent events log
+        }
+        self._lock = threading.Lock()
+        self._already_rendered = False
+
+    def update_status(self, category, data):
+        """Thread-safe method to update status data"""
+        with self._lock:
+            if category in self.status_data:
+                self.status_data[category].update(data)
+
+    def add_event(self, event_type, message, timestamp=None):
+        """Add an event to the timeline"""
+        if timestamp is None:
+            timestamp = datetime.now()
+
+        with self._lock:
+            self.status_data["events"].insert(
+                0,
+                {
+                    "type": event_type,  # 'info', 'warning', 'error', 'success'
+                    "message": message,
+                    "timestamp": timestamp,
+                },
+            )
+            # Keep only last 10 events
+            self.status_data["events"] = self.status_data["events"][:10]
+
+    def get_circuit_breaker_emoji(self, state):
+        """Get status emoji for circuit breaker state"""
+        emoji_map = {"CLOSED": "🟢", "OPEN": "🔴", "HALF_OPEN": "🟡"}
+        return emoji_map.get(state, "⚪")
+
+    def get_uptime_string(self, start_time):
+        """Calculate uptime string"""
+        if not start_time:
+            return "Not started"
+
+        uptime = datetime.now() - start_time
+        hours, remainder = divmod(int(uptime.total_seconds()), 3600)
+        minutes, seconds = divmod(remainder, 60)
+
+        if hours > 0:
+            return f"{hours}h {minutes}m {seconds}s"
+        elif minutes > 0:
+            return f"{minutes}m {seconds}s"
+        else:
+            return f"{seconds}s"
+
+    def on_startup(self, current_card):
+        """Initialize the card when monitoring starts"""
+        current_card.append(Markdown("# 🦙 `@ollama` Status Dashboard"))
+        current_card.append(Markdown("_Initializing Ollama system..._"))
+        current_card.refresh()
+
+    def render_card_fresh(self, current_card, data):
+        """Render the complete card with all status information"""
+        self._already_rendered = True
+        current_card.clear()
+
+        # Header with version information
+        current_card.append(Markdown("# 🦙 `@ollama` Status Dashboard"))
+
+        # Version information in header
+        versions = data.get("versions", {})
+        system_version = versions.get("ollama_system", "Unknown")
+        python_version = versions.get("ollama_python", "Unknown")
+        current_card.append(
+            Markdown(
+                f"**System:** `{system_version}` | **Python Client:** `{python_version}`"
+            )
+        )
+
+        # Cache policy information
+        cache_info = data.get("cache", {})
+        cache_policy = cache_info.get("policy", "auto")
+        current_card.append(Markdown(f"**Cache Policy:** `{cache_policy}`"))
+
+        current_card.append(
+            Markdown(f"_Last updated: {datetime.now().strftime('%H:%M:%S')}_")
+        )
+
+        # Circuit Breaker Status
+        cb_data = data["circuit_breaker"]
+        cb_emoji = self.get_circuit_breaker_emoji(cb_data["state"])
+        cb_status = f"{cb_emoji} **{cb_data['state']}**"
+        if cb_data["failure_count"] > 0:
+            cb_status += f" (failures: {cb_data['failure_count']})"
+
+        # Server Status
+        server_data = data["server"]
+        uptime = self.get_uptime_string(server_data.get("uptime_start"))
+        server_status = f"**{server_data['status']}**"
+        if server_data["restart_attempts"] > 0:
+            server_status += f" (restarts: {server_data['restart_attempts']})"
+
+        # Status Overview Table
+        status_table = [
+            ["Circuit Breaker", Markdown(cb_status)],
+            ["Server Status", Markdown(server_status)],
+            ["Server Uptime", Markdown(uptime)],
+            [
+                "Last Health Check",
+                Markdown(server_data.get("health_status", "Unknown")),
+            ],
+        ]
+
+        current_card.append(Markdown("## System Status"))
+        current_card.append(Table(status_table, headers=["Component", "Status"]))
+
+        # Models Status
+        if data["models"]:
+            current_card.append(Markdown("## Models"))
+            model_table = []
+            cache_model_status = cache_info.get("model_status", {})
+
+            for model_name, model_info in data["models"].items():
+                status = model_info.get("status", "Unknown")
+                pull_time = model_info.get("pull_time", "N/A")
+                if isinstance(pull_time, (int, float)):
+                    pull_time = f"{pull_time:.1f}s"
+
+                # Add cache status indicator
+                cache_status = cache_model_status.get(model_name, "unknown")
+                cache_emoji = {
+                    "exists": "💾",
+                    "missing": "❌",
+                    "error": "⚠️",
+                    "unknown": "❓",
+                }.get(cache_status, "❓")
+
+                # Get model metadata
+                size_formatted = model_info.get("size_formatted", "Unknown")
+                blob_count = model_info.get("blob_count", "Unknown")
+                if blob_count == 0:
+                    blob_count = "Unknown"
+
+                model_table.append(
+                    [
+                        f"{model_name} {cache_emoji}",
+                        status,
+                        pull_time,
+                        size_formatted,
+                        str(blob_count),
+                    ]
+                )
+
+            current_card.append(
+                Table(
+                    model_table,
+                    headers=["Model (Cache)", "Status", "Pull Time", "Size", "Blobs"],
+                )
+            )
+
+        # Performance Metrics
+        perf_data = data["performance"]
+        if any(v is not None for v in perf_data.values()):
+            current_card.append(Markdown("## Performance"))
+
+            # Separate initialization and shutdown metrics
+            init_metrics = []
+            shutdown_metrics = []
+            other_metrics = []
+
+            for metric, value in perf_data.items():
+                if value is not None:
+                    display_value = value
+                    if isinstance(value, (int, float)):
+                        display_value = f"{value:.1f}s"
+
+                    metric_display = metric.replace("_", " ").title()
+
+                    if "shutdown" in metric.lower():
+                        shutdown_metrics.append([metric_display, display_value])
+                    elif metric in [
+                        "install_time",
+                        "server_startup_time",
+                        "total_initialization_time",
+                    ]:
+                        init_metrics.append([metric_display, display_value])
+                    else:
+                        other_metrics.append([metric_display, display_value])
+
+            # Display metrics in organized sections
+            if init_metrics:
+                current_card.append(Markdown("### Initialization"))
+                current_card.append(Table(init_metrics, headers=["Metric", "Duration"]))
+
+            if shutdown_metrics:
+                current_card.append(Markdown("### Shutdown"))
+                current_card.append(
+                    Table(shutdown_metrics, headers=["Metric", "Value"])
+                )
+
+            if other_metrics:
+                current_card.append(Markdown("### Other"))
+                current_card.append(Table(other_metrics, headers=["Metric", "Value"]))
+
+        # Recent Events
+        if data["events"]:
+            current_card.append(Markdown("## Recent Events"))
+            events_table = []
+            for event in data["events"][:5]:  # Show last 5 events
+                timestamp = event["timestamp"].strftime("%H:%M:%S")
+                event_type = event["type"]
+                message = event["message"]
+
+                # Add emoji based on event type
+                type_emoji = {
+                    "info": "ℹ️",
+                    "success": "✅",
+                    "warning": "⚠️",
+                    "error": "❌",
+                }.get(event_type, "ℹ️")
+
+                events_table.append([timestamp, f"{type_emoji} {message}"])
+
+            current_card.append(Table(events_table, headers=["Time", "Event"]))
+
+        current_card.refresh()
+
+    def on_error(self, current_card, error_message):
+        """Handle errors in card rendering"""
+        if not self._already_rendered:
+            current_card.clear()
+            current_card.append(Markdown("# 🦙 `@ollama` Status Dashboard"))
+            current_card.append(Markdown(f"## ❌ Error: {str(error_message)}"))
+            current_card.refresh()
+
+    def on_update(self, current_card, data_object):
+        """Update the card with new data"""
+        with self._lock:
+            current_data = self.status_data.copy()
+
+        if not self._already_rendered:
+            self.render_card_fresh(current_card, current_data)
+        else:
+            # For frequent updates, we could implement incremental updates here
+            # For now, just re-render the whole card
+            self.render_card_fresh(current_card, current_data)
+
+    def sqlite_fetch_func(self, conn):
+        """Required by CardRefresher (which needs a refactor), but we use in-memory data instead"""
+        with self._lock:
+            return {"status": self.status_data}

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_programmatic_access()
+    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/perimeters.py

@@ -3,12 +3,16 @@ import fcntl
 from os import path
 import json
 from metaflow.exception import MetaflowException
+from typing import Union
 
 CURRENT_PERIMETER_KEY = "OB_CURRENT_PERIMETER"
 CURRENT_PERIMETER_URL = "OB_CURRENT_PERIMETER_MF_CONFIG_URL"
+CURRENT_PERIMETER_URL_LEGACY_KEY = (
+    "OB_CURRENT_PERIMETER_URL"  # For backwards compatibility with workstations.
+)
 
 
-def set_current_perimeter_config_url_in_environment():
+def get_perimeter_config_url_if_set_in_ob_config() -> Union[str, None]:
     # If OBP_CONFIG_DIR is set, use that, otherwise use METAFLOW_HOME
     # If neither are set, use ~/.metaflowconfig
     obp_config_dir = path.expanduser(
@@ -36,18 +40,28 @@ def set_current_perimeter_config_url_in_environment():
         with open(file_path, "r") as f:
             ob_config = json.loads(f.read())
 
-        if CURRENT_PERIMETER_KEY in ob_config and CURRENT_PERIMETER_URL in ob_config:
+        if CURRENT_PERIMETER_KEY in ob_config and (
+            CURRENT_PERIMETER_URL in ob_config
+            or CURRENT_PERIMETER_URL_LEGACY_KEY in ob_config
+        ):
             os.environ[CURRENT_PERIMETER_KEY] = ob_config[CURRENT_PERIMETER_KEY]
-            os.environ[CURRENT_PERIMETER_URL] = ob_config[CURRENT_PERIMETER_URL]
+            if CURRENT_PERIMETER_URL in ob_config:
+                os.environ[CURRENT_PERIMETER_URL] = ob_config[CURRENT_PERIMETER_URL]
+            elif CURRENT_PERIMETER_URL_LEGACY_KEY in ob_config:
+                os.environ[CURRENT_PERIMETER_URL] = ob_config[
+                    CURRENT_PERIMETER_URL_LEGACY_KEY
+                ]
+            return os.environ[CURRENT_PERIMETER_URL]
         else:
             raise MetaflowException(
-                "%s does not contain the key %s".format(
+                "{} does not contain the key {}".format(
                     file_path, CURRENT_PERIMETER_KEY
                 )
             )
     elif "OBP_CONFIG_DIR" in os.environ:
         raise MetaflowException(
-            "Environment variable OBP_CONFIG_DIR is set to %s but this directory does not contain an ob_config.json file.".format(
+            "Environment variable OBP_CONFIG_DIR is set to {} but this directory does not contain an ob_config.json file.".format(
                 os.environ["OBP_CONFIG_DIR"]
             )
         )
+    return None

metaflow_extensions/outerbounds/plugins/profilers/deco_injector.py

@@ -0,0 +1,70 @@
+from metaflow.exception import MetaflowException
+from collections import defaultdict
+
+
+class CardDecoratorInjector:
+    """
+    Mixin Useful for injecting @card decorators from other first class Metaflow decorators.
+    """
+
+    _first_time_init = defaultdict(dict)
+
+    @classmethod
+    def _get_first_time_init_cached_value(cls, step_name, card_id):
+        return cls._first_time_init.get(step_name, {}).get(card_id, None)
+
+    @classmethod
+    def _set_first_time_init_cached_value(cls, step_name, card_id, value):
+        cls._first_time_init[step_name][card_id] = value
+
+    def _card_deco_already_attached(self, step, card_id):
+        for decorator in step.decorators:
+            if decorator.name == "card":
+                if decorator.attributes["id"] and card_id in decorator.attributes["id"]:
+                    return True
+        return False
+
+    def _get_step(self, flow, step_name):
+        for step in flow:
+            if step.name == step_name:
+                return step
+        return None
+
+    def _first_time_init_check(self, step_dag_node, card_id):
+        """ """
+        return not self._card_deco_already_attached(step_dag_node, card_id)
+
+    def attach_card_decorator(
+        self,
+        flow,
+        step_name,
+        card_id,
+        card_type,
+        refresh_interval=5,
+    ):
+        """
+        This method is called `step_init` in your StepDecorator code since
+        this class is used as a Mixin
+        """
+        from metaflow import decorators as _decorators
+
+        if not all([card_id, card_type]):
+            raise MetaflowException(
+                "`card_id` and `card_type` must be set in the `CardDecoratorInjector` Mixin"
+            )
+
+        step_dag_node = self._get_step(flow, step_name)
+        if (
+            self._get_first_time_init_cached_value(step_name, card_id) is None
+        ):  # First check class level setting.
+            if self._first_time_init_check(step_dag_node, card_id):
+                self._set_first_time_init_cached_value(step_name, card_id, True)
+                _decorators._attach_decorators_to_step(
+                    step_dag_node,
+                    [
+                        "card:type=%s,id=%s,refresh_interval=%s"
+                        % (card_type, card_id, str(refresh_interval))
+                    ],
+                )
+            else:
+                self._set_first_time_init_cached_value(step_name, card_id, False)

metaflow_extensions/outerbounds/plugins/profilers/gpu_profile_decorator.py

@@ -0,0 +1,88 @@
+from datetime import datetime
+from metaflow.decorators import StepDecorator
+from ...profilers.gpu import GPUProfiler  # Fix import
+from .deco_injector import CardDecoratorInjector
+import threading
+
+
+class GPUProfileDecorator(StepDecorator):
+    name = "gpu_profile"
+
+    defaults = {
+        "include_artifacts": True,
+        "artifact_prefix": "gpu_profile_",
+        "interval": 1,
+    }
+
+    def step_init(
+        self, flow, graph, step_name, decorators, environment, flow_datastore, logger
+    ):
+        self.deco_injector = CardDecoratorInjector()
+        self.deco_injector.attach_card_decorator(
+            flow,
+            step_name,
+            "gpu_profile",
+            "blank",
+            refresh_interval=self.attributes["interval"],
+        )
+
+    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,
+    ):
+        self._profiler = GPUProfiler(
+            interval=self.attributes["interval"],
+            artifact_name=self.attributes["artifact_prefix"] + "data",
+        )
+
+    def task_decorate(
+        self, step_func, flow, graph, retry_count, max_user_code_retries, ubf_context
+    ):
+        from metaflow import current
+        from metaflow.cards import Markdown
+
+        if self.attributes["include_artifacts"]:
+            setattr(
+                flow,
+                self.attributes["artifact_prefix"] + "num_gpus",
+                len(self._profiler.devices),
+            )
+
+        current.card["gpu_profile"].append(
+            Markdown("# GPU profile for `%s`" % current.pathspec)
+        )
+        current.card["gpu_profile"].append(
+            Markdown(
+                "_Started at: %s_"
+                % datetime.now().astimezone().strftime("%Y-%m-%dT%H:%M:%S %z")
+            )
+        )
+        self._profiler._setup_card()
+        current.card["gpu_profile"].refresh()
+        self._update_thread = threading.Thread(
+            target=self._profiler._update_card, daemon=True
+        )
+        self._update_thread.start()
+
+        def wrapped_step_func():
+            try:
+                step_func()
+            finally:
+                try:
+                    results = self._profiler.finish()
+                except:
+                    results = {"error": "couldn't read profiler results"}
+                if self.attributes["include_artifacts"]:
+                    setattr(flow, self.attributes["artifact_prefix"] + "data", results)
+
+        return wrapped_step_func

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"]