ob-metaflow-extensions 1.1.166rc5__py2.py3-none-any.whl → 1.1.167__py2.py3-none-any.whl

metaflow_extensions/outerbounds/plugins/__init__.py

@@ -331,7 +331,6 @@ CLIS_DESC = [
     ("nvct", ".nvct.nvct_cli.cli"),
     ("fast-bakery", ".fast_bakery.fast_bakery_cli.cli"),
     ("snowpark", ".snowpark.snowpark_cli.cli"),
-    ("app", ".apps.app_cli.cli"),
 ]
 STEP_DECORATORS_DESC = [
     ("nvidia", ".nvcf.nvcf_decorator.NvcfDecorator"),
@@ -345,6 +344,7 @@ STEP_DECORATORS_DESC = [
     ("gpu_profile", ".profilers.gpu_profile_decorator.GPUProfileDecorator"),
     ("nim", ".nim.nim_decorator.NimDecorator"),
     ("ollama", ".ollama.OllamaDecorator"),
+    ("vllm", ".vllm.VLLMDecorator"),
     ("app_deploy", ".apps.deploy_decorator.WorkstationAppDeployDecorator"),
 ]
 

metaflow_extensions/outerbounds/plugins/vllm/__init__.py

@@ -0,0 +1,177 @@
+from metaflow.decorators import StepDecorator
+from metaflow import current
+import functools
+import os
+import threading
+from metaflow.unbounded_foreach import UBF_CONTROL, UBF_TASK
+from metaflow.metaflow_config import from_conf
+
+from .vllm_manager import VLLMManager
+from .status_card import VLLMStatusCard, CardDecoratorInjector
+
+__mf_promote_submodules__ = ["plugins.vllm"]
+
+
+class VLLMDecorator(StepDecorator, CardDecoratorInjector):
+    """
+    This decorator is used to run vllm APIs as Metaflow task sidecars.
+
+    User code call
+    --------------
+    @vllm(
+        model="...",
+        ...
+    )
+
+    Valid backend options
+    ---------------------
+    - 'local': Run as a separate process on the local task machine.
+
+    Valid model options
+    -------------------
+    Any HuggingFace model identifier, e.g. 'meta-llama/Llama-3.2-1B'
+
+    NOTE: vLLM's OpenAI-compatible server serves ONE model per server instance.
+    If you need multiple models, you must create multiple @vllm decorators.
+
+    Parameters
+    ----------
+    model: str
+        HuggingFace model identifier to be served by vLLM.
+    backend: str
+        Determines where and how to run the vLLM process.
+    debug: bool
+        Whether to turn on verbose debugging logs.
+    kwargs : Any
+        Any other keyword arguments are passed directly to the vLLM engine.
+        This allows for flexible configuration of vLLM server settings.
+        For example, `tensor_parallel_size=2`.
+    """
+
+    name = "vllm"
+    defaults = {
+        "model": None,
+        "backend": "local",
+        "debug": False,
+        "stream_logs_to_card": False,
+        "card_refresh_interval": 10,
+        "engine_args": {},
+    }
+
+    def step_init(
+        self, flow, graph, step_name, decorators, environment, flow_datastore, logger
+    ):
+        super().step_init(
+            flow, graph, step_name, decorators, environment, flow_datastore, logger
+        )
+
+        # Validate that a model is specified
+        if not self.attributes["model"]:
+            raise ValueError(
+                f"@vllm decorator on step '{step_name}' requires a 'model' parameter. "
+                f"Example: @vllm(model='meta-llama/Llama-3.2-1B')"
+            )
+
+        # Attach the vllm status card
+        self.attach_card_decorator(
+            flow,
+            step_name,
+            "vllm_status",
+            "blank",
+            refresh_interval=self.attributes["card_refresh_interval"],
+        )
+
+    def task_decorate(
+        self, step_func, flow, graph, retry_count, max_user_code_retries, ubf_context
+    ):
+        @functools.wraps(step_func)
+        def vllm_wrapper():
+            self.vllm_manager = None
+            self.status_card = None
+            self.card_monitor_thread = None
+
+            try:
+                self.status_card = VLLMStatusCard(
+                    refresh_interval=self.attributes["card_refresh_interval"]
+                )
+
+                def monitor_card():
+                    try:
+                        self.status_card.on_startup(current.card["vllm_status"])
+
+                        while not getattr(
+                            self.card_monitor_thread, "_stop_event", False
+                        ):
+                            try:
+                                self.status_card.on_update(
+                                    current.card["vllm_status"], None
+                                )
+                                import time
+
+                                time.sleep(self.attributes["card_refresh_interval"])
+                            except Exception as e:
+                                if self.attributes["debug"]:
+                                    print(f"[@vllm] Card monitoring error: {e}")
+                                break
+                    except Exception as e:
+                        if self.attributes["debug"]:
+                            print(f"[@vllm] Card monitor thread error: {e}")
+                        self.status_card.on_error(current.card["vllm_status"], str(e))
+
+                self.card_monitor_thread = threading.Thread(
+                    target=monitor_card, daemon=True
+                )
+                self.card_monitor_thread._stop_event = False
+                self.card_monitor_thread.start()
+                self.vllm_manager = VLLMManager(
+                    model=self.attributes["model"],
+                    backend=self.attributes["backend"],
+                    debug=self.attributes["debug"],
+                    status_card=self.status_card,
+                    stream_logs_to_card=self.attributes["stream_logs_to_card"],
+                    **self.attributes["engine_args"],
+                )
+                if self.attributes["debug"]:
+                    print("[@vllm] VLLMManager initialized.")
+
+            except Exception as e:
+                if self.status_card:
+                    self.status_card.add_event(
+                        "error", f"Initialization failed: {str(e)}"
+                    )
+                    try:
+                        self.status_card.on_error(current.card["vllm_status"], str(e))
+                    except:
+                        pass
+                print(f"[@vllm] Error initializing VLLMManager: {e}")
+                raise
+
+            try:
+                if self.status_card:
+                    self.status_card.add_event("info", "Starting user step function")
+                step_func()
+                if self.status_card:
+                    self.status_card.add_event(
+                        "success", "User step function completed successfully"
+                    )
+            finally:
+                if self.vllm_manager:
+                    self.vllm_manager.terminate_models()
+
+                if self.card_monitor_thread and self.status_card:
+                    import time
+
+                    try:
+                        self.status_card.on_update(current.card["vllm_status"], None)
+                    except Exception as e:
+                        if self.attributes["debug"]:
+                            print(f"[@vllm] Final card update error: {e}")
+                    time.sleep(2)
+
+                if self.card_monitor_thread:
+                    self.card_monitor_thread._stop_event = True
+                    self.card_monitor_thread.join(timeout=5)
+                    if self.attributes["debug"]:
+                        print("[@vllm] Card monitoring thread stopped.")
+
+        return vllm_wrapper

metaflow_extensions/outerbounds/plugins/vllm/constants.py

@@ -0,0 +1 @@
+VLLM_SUFFIX = "mf.vllm"

metaflow_extensions/outerbounds/plugins/vllm/exceptions.py

@@ -0,0 +1 @@
+from metaflow.exception import MetaflowException

metaflow_extensions/outerbounds/plugins/vllm/status_card.py

@@ -0,0 +1,352 @@
+from metaflow.cards import Markdown, Table, VegaChart
+from metaflow.metaflow_current import current
+from datetime import datetime
+import threading
+import time
+
+
+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 == 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(
+                "`INJECTED_CARD_ID` and `INJECTED_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)
+
+
+class CardRefresher:
+
+    CARD_ID = None
+
+    def on_startup(self, current_card):
+        raise NotImplementedError("make_card method must be implemented")
+
+    def on_error(self, current_card, error_message):
+        raise NotImplementedError("error_card method must be implemented")
+
+    def on_update(self, current_card, data_object):
+        raise NotImplementedError("update_card method must be implemented")
+
+    def sqlite_fetch_func(self, conn):
+        raise NotImplementedError("sqlite_fetch_func must be implemented")
+
+
+class VLLMStatusCard(CardRefresher):
+    """
+    Real-time status card for vLLM system monitoring.
+    Shows server health, model status, and recent events.
+
+    Intended to be inherited from in a step decorator like this:
+        class VLLMDecorator(StepDecorator, VLLMStatusCard):
+    """
+
+    CARD_ID = "vllm_status"
+
+    def __init__(self, refresh_interval=10):
+        self.refresh_interval = refresh_interval
+        self.status_data = {
+            "server": {
+                "status": "Starting",
+                "uptime_start": None,
+                "last_health_check": None,
+                "health_status": "Unknown",
+                "models": [],
+            },
+            "models": {},  # model_name -> {status, load_time, etc}
+            "performance": {
+                "install_time": None,
+                "server_startup_time": None,
+                "total_initialization_time": None,
+            },
+            "versions": {
+                "vllm": "Detecting...",
+            },
+            "events": [],  # Recent events log
+            "logs": [],
+        }
+        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_log_line(self, log_line):
+        """Add a log line to the logs."""
+        with self._lock:
+            self.status_data["logs"].append(log_line)
+            # Keep only last 20 lines
+            self.status_data["logs"] = self.status_data["logs"][-20:]
+
+    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("# 🚀 `@vllm` Status Dashboard"))
+        current_card.append(Markdown("_Initializing vLLM 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()
+
+        current_card.append(Markdown("# 🚀 `@vllm` Status Dashboard"))
+
+        versions = data.get("versions", {})
+        vllm_version = versions.get("vllm", "Unknown")
+        current_card.append(Markdown(f"**vLLM Version:** `{vllm_version}`"))
+
+        current_card.append(
+            Markdown(f"_Last updated: {datetime.now().strftime('%H:%M:%S')}_")
+        )
+
+        server_data = data["server"]
+        uptime = self.get_uptime_string(server_data.get("uptime_start"))
+        server_status = server_data.get("status", "Unknown")
+        model = server_data.get("model", "Unknown")
+
+        # Determine status emoji
+        if server_status == "Running":
+            status_emoji = "🟢"
+            model_emoji = "✅"
+        elif server_status == "Failed":
+            status_emoji = "🔴"
+            model_emoji = "❌"
+        elif server_status == "Starting":
+            status_emoji = "🟡"
+            model_emoji = "⏳"
+        else:  # Stopped, etc.
+            status_emoji = "⚫"
+            model_emoji = "⏹️"
+
+        # Main status section
+        current_card.append(
+            Markdown(f"## {status_emoji} Server Status: {server_status}")
+        )
+
+        if server_status == "Running" and uptime:
+            current_card.append(Markdown(f"**Uptime:** {uptime}"))
+
+        # Model information - only show detailed status if server is running
+        if server_status == "Running":
+            current_card.append(Markdown(f"## {model_emoji} Model: `{model}`"))
+
+            # Show model-specific status if available
+            models_data = data.get("models", {})
+            if models_data and model in models_data:
+                model_info = models_data[model]
+                model_status = model_info.get("status", "Unknown")
+                load_time = model_info.get("load_time")
+                location = model_info.get("location")
+
+                current_card.append(Markdown(f"**Status:** {model_status}"))
+                if location:
+                    current_card.append(Markdown(f"**Location:** `{location}`"))
+                if load_time and isinstance(load_time, (int, float)):
+                    current_card.append(Markdown(f"**Load Time:** {load_time:.1f}s"))
+        elif model != "Unknown":
+            current_card.append(
+                Markdown(f"## {model_emoji} Model: `{model}` (Server Stopped)")
+            )
+
+        # Simplified monitoring note
+        current_card.append(
+            Markdown(
+                "## 🔧 Monitoring\n**Advanced Features:** Disabled (Circuit Breaker, Request Interception)"
+            )
+        )
+
+        # Performance metrics
+        perf_data = data["performance"]
+        if any(v is not None for v in perf_data.values()):
+            current_card.append(Markdown("## ⚡ Performance"))
+
+            init_metrics = []
+            shutdown_metrics = []
+
+            for metric, value in perf_data.items():
+                if value is not None:
+                    display_value = (
+                        f"{value:.1f}s" if isinstance(value, (int, float)) else value
+                    )
+                    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])
+
+            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"])
+                )
+
+        # Recent events
+        events = data.get("events", [])
+        if events:
+            current_card.append(Markdown("## 📝 Recent Events"))
+            for event in events[:5]:  # Show last 5 events
+                event_type = event.get("type", "info")
+                message = event.get("message", "")
+                timestamp = event.get("timestamp", datetime.now())
+
+                emoji_map = {
+                    "info": "ℹ️",
+                    "success": "✅",
+                    "warning": "⚠️",
+                    "error": "❌",
+                }
+                emoji = emoji_map.get(event_type, "ℹ️")
+
+                time_str = (
+                    timestamp.strftime("%H:%M:%S")
+                    if isinstance(timestamp, datetime)
+                    else str(timestamp)
+                )
+                current_card.append(Markdown(f"- {emoji} `{time_str}` {message}"))
+
+        # Server Logs
+        logs = data.get("logs", [])
+        if logs:
+            current_card.append(Markdown("## 📜 Server Logs"))
+            # The logs are appended, so they are in chronological order.
+            log_content = "\n".join(logs)
+            current_card.append(Markdown(f"```\n{log_content}\n```"))
+
+        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("# 🚀 `@vllm` 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/vllm/vllm_manager.py

@@ -0,0 +1,471 @@
+import subprocess
+from concurrent.futures import ThreadPoolExecutor, as_completed
+import time
+import socket
+import sys
+import os
+import functools
+import json
+import requests
+from enum import Enum
+import threading
+from datetime import datetime
+
+from .constants import VLLM_SUFFIX
+
+
+class ProcessStatus:
+    RUNNING = "RUNNING"
+    FAILED = "FAILED"
+    SUCCESSFUL = "SUCCESSFUL"
+
+
+class VLLMManager:
+    """
+    A process manager for vLLM runtimes.
+    Implements interface @vllm(model=..., ...) to provide a local backend.
+    It wraps the vLLM OpenAI-compatible API server to make it easier to profile vLLM use on Outerbounds.
+
+    NOTE: vLLM's OpenAI-compatible server serves ONE model per server instance.
+    If you need multiple models, you must create multiple server instances.
+
+    Example usage:
+        from vllm import LLM
+        llm = LLM(model="meta-llama/Llama-3.2-1B")
+        llm.generate("Hello, world!")
+
+    Or via OpenAI-compatible API:
+        import openai
+        client = openai.OpenAI(
+            base_url="http://localhost:8000/v1",
+            api_key="token-abc123"
+        )
+        response = client.chat.completions.create(
+            model="meta-llama/Llama-3.2-1B",
+            messages=[{"role": "user", "content": "Hello"}]
+        )
+    """
+
+    def __init__(
+        self,
+        model,
+        backend="local",
+        debug=False,
+        status_card=None,
+        port=8000,
+        host="127.0.0.1",
+        stream_logs_to_card=False,
+        **vllm_args,
+    ):
+        # Validate that only a single model is provided
+        if isinstance(model, list):
+            if len(model) != 1:
+                raise ValueError(
+                    f"vLLM server can only serve one model per instance. "
+                    f"Got {len(model)} models: {model}. "
+                    f"Please specify a single model or create multiple @vllm decorators."
+                )
+            self.model = model[0]
+        else:
+            self.model = model
+
+        self.processes = {}
+        self.debug = debug
+        self.stream_logs_to_card = stream_logs_to_card
+        self.stats = {}
+        self.port = port
+        self.host = host
+        self.vllm_url = f"http://{host}:{port}"
+        self.status_card = status_card
+        self.initialization_start = time.time()
+        self.server_process = None
+        self.vllm_args = vllm_args
+
+        if backend != "local":
+            raise ValueError(
+                "VLLMManager only supports the 'local' backend at this time."
+            )
+
+        self._log_event("info", "Starting vLLM initialization")
+        self._update_server_status("Initializing")
+
+        self._timeit(self._install_vllm, "install_vllm")
+        self._timeit(self._launch_vllm_server, "launch_server")
+        self._collect_version_info()
+
+        total_init_time = time.time() - self.initialization_start
+        self._update_performance("total_initialization_time", total_init_time)
+        self._log_event(
+            "success", f"vLLM initialization completed in {total_init_time:.1f}s"
+        )
+
+    def _log_event(self, event_type, message):
+        if self.status_card:
+            self.status_card.add_event(event_type, message)
+        if self.debug:
+            print(f"[@vllm] {event_type.upper()}: {message}")
+
+    def _update_server_status(self, status, **kwargs):
+        if self.status_card:
+            update_data = {"status": status}
+            update_data.update(kwargs)
+            self.status_card.update_status("server", update_data)
+
+    def _update_model_status(self, model_name, **kwargs):
+        if self.status_card:
+            current_models = self.status_card.status_data.get("models", {})
+            if model_name not in current_models:
+                current_models[model_name] = {}
+            current_models[model_name].update(kwargs)
+            self.status_card.update_status("models", current_models)
+
+    def _update_performance(self, metric, value):
+        if self.status_card:
+            self.status_card.update_status("performance", {metric: value})
+
+    def _timeit(self, f, name):
+        t0 = time.time()
+        f()
+        tf = time.time()
+        duration = tf - t0
+        self.stats[name] = {"process_runtime": duration}
+
+        if name == "install_vllm":
+            self._update_performance("install_time", duration)
+        elif name == "launch_server":
+            self._update_performance("server_startup_time", duration)
+
+    def _stream_output(self, stream, prefix):
+        """Reads and logs output from a stream."""
+        for line in iter(stream.readline, ""):
+            if line:
+                line = line.strip()
+                if self.stream_logs_to_card and self.status_card:
+                    self.status_card.add_log_line(f"[{prefix}] {line}")
+                elif self.debug:
+                    print(f"[{prefix}] {line}")
+        stream.close()
+
+    def _is_port_open(self, host, port, timeout=1):
+        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
+            sock.settimeout(timeout)
+            try:
+                sock.connect((host, port))
+                return True
+            except socket.error:
+                return False
+
+    def _install_vllm(self):
+        self._log_event("info", "Checking for existing vLLM installation")
+        try:
+            import vllm
+
+            self._log_event("success", f"vLLM {vllm.__version__} is already installed")
+            if self.debug:
+                print(f"[@vllm] vLLM {vllm.__version__} is already installed.")
+            return
+        except ImportError as e:
+            self._log_event(
+                "Error", "vLLM not installed. Please add it to your environment."
+            )
+            if self.debug:
+                print(
+                    "[@vllm] vLLM not found. The user is responsible for installation."
+                )
+            raise e
+            # We are not installing it automatically to respect user's environment management.
+
+    def _launch_vllm_server(self):
+        self._update_server_status("Starting")
+        self._log_event("info", f"Starting vLLM server with model: {self.model}")
+
+        # Check if the model is cached
+        hf_home = os.environ.get("HF_HOME")
+        if hf_home:
+            # Construct the expected cache path for the model
+            model_path_id = f"models--{self.model.replace('/', '--')}"
+            model_cache_path = os.path.join(hf_home, model_path_id)
+            if os.path.exists(model_cache_path):
+                self._log_event("info", f"Found cached model at: {model_cache_path}")
+                self._update_model_status(
+                    self.model, status="Found in cache", location=model_cache_path
+                )
+            else:
+                self._log_event(
+                    "warning",
+                    f"Cached model not found at {model_cache_path}. vLLM will attempt to download it.",
+                )
+                self._update_model_status(self.model, status="Downloading")
+        else:
+            self._log_event(
+                "warning",
+                "HF_HOME environment variable not set. vLLM will use default cache location and may re-download.",
+            )
+
+        if not self.model:
+            raise ValueError("At least one model must be specified for @vllm.")
+
+        try:
+            if self.debug:
+                print(
+                    f"[@vllm] Starting vLLM OpenAI-compatible server for model: {self.model}"
+                )
+
+            cmd = [
+                sys.executable,
+                "-m",
+                "vllm.entrypoints.openai.api_server",
+                "--model",
+                self.model,
+                "--host",
+                self.host,
+                "--port",
+                str(self.port),
+            ]
+
+            vllm_args_copy = self.vllm_args.copy()
+            if self.debug or self.stream_logs_to_card:
+                # Note: This is an undocumented argument for the vLLM OpenAI server entrypoint.
+                vllm_args_copy.setdefault("uvicorn_log_level", "debug")
+
+            for key, value in vllm_args_copy.items():
+                arg_name = f"--{key.replace('_', '-')}"
+                if isinstance(value, bool):
+                    if value:
+                        cmd.append(arg_name)
+                elif value is not None:
+                    cmd.append(arg_name)
+                    cmd.append(str(value))
+
+            # For debugging, log the exact command being run to the status card
+            command_str = " ".join(cmd)
+            self._log_event("info", f"Launch Command: `{command_str}`")
+            if self.debug:
+                print(f"[@vllm] Launching vLLM with command: {command_str}")
+
+            process = subprocess.Popen(
+                cmd,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+                text=True,
+                bufsize=1,  # Line-buffered
+            )
+
+            # Threads to stream subprocess output
+            if self.debug or self.stream_logs_to_card:
+                stdout_thread = threading.Thread(
+                    target=self._stream_output,
+                    args=(process.stdout, "@vllm-server-out"),
+                )
+                stderr_thread = threading.Thread(
+                    target=self._stream_output,
+                    args=(process.stderr, "@vllm-server-err"),
+                )
+                stdout_thread.daemon = True
+                stderr_thread.daemon = True
+                stdout_thread.start()
+                stderr_thread.start()
+
+            self.server_process = process
+            self.processes[process.pid] = {
+                "p": process,
+                "properties": {
+                    "type": "vllm-server",
+                    "model": self.model,
+                    "error_details": None,
+                },
+                "status": ProcessStatus.RUNNING,
+            }
+
+            if self.debug:
+                print(f"[@vllm] Started vLLM server process with PID {process.pid}")
+
+            retries = 0
+            max_retries = 240
+            while (
+                not self._is_port_open(self.host, self.port, timeout=2)
+                and retries < max_retries
+            ):
+                if retries == 0:
+                    print("[@vllm] Waiting for server to be ready...")
+                elif retries % 10 == 0:
+                    print(
+                        f"[@vllm] Still waiting for server... ({retries}/{max_retries})"
+                    )
+
+                returncode = process.poll()
+                if returncode is not None:
+                    if self.debug or self.stream_logs_to_card:
+                        # Threads are handling output, can't use communicate.
+                        # The error has already been printed to the log by the thread.
+                        if self.stream_logs_to_card:
+                            details_msg = "See card for logs."
+                        else:
+                            details_msg = "See logs from @vllm-server-err for details."
+                        error_details = f"Return code: {returncode}. {details_msg}"
+                    else:
+                        # No threads, so we can and should use communicate to get stderr.
+                        stdout, stderr = process.communicate()
+                        error_details = f"Return code: {returncode}, stderr: {stderr}"
+
+                    self.processes[process.pid]["properties"][
+                        "error_details"
+                    ] = error_details
+                    self.processes[process.pid]["status"] = ProcessStatus.FAILED
+                    self._update_server_status("Failed", error_details=error_details)
+                    self._log_event(
+                        "error", f"vLLM server failed to start: {error_details}"
+                    )
+                    raise RuntimeError(f"vLLM server failed to start: {error_details}")
+
+                time.sleep(2)
+                retries += 1
+
+            if not self._is_port_open(self.host, self.port, timeout=2):
+                error_details = f"vLLM server did not start listening on {self.host}:{self.port} after {max_retries*2}s"
+                self.processes[process.pid]["properties"][
+                    "error_details"
+                ] = error_details
+                self.processes[process.pid]["status"] = ProcessStatus.FAILED
+                self._update_server_status("Failed", error_details=error_details)
+                self._log_event("error", f"Server startup timeout: {error_details}")
+                raise RuntimeError(f"vLLM server failed to start: {error_details}")
+
+            if not self._verify_server_health():
+                error_details = "vLLM server started but failed health check"
+                self.processes[process.pid]["status"] = ProcessStatus.FAILED
+                self._update_server_status("Failed", error_details=error_details)
+                self._log_event("error", error_details)
+                raise RuntimeError(error_details)
+
+            self._update_server_status(
+                "Running", uptime_start=datetime.now(), model=self.model
+            )
+            self._log_event("success", "vLLM server is ready and listening")
+
+            self._update_model_status(self.model, status="Ready")
+
+            if self.debug:
+                print("[@vllm] Server is ready.")
+
+        except Exception as e:
+            if process and process.pid in self.processes:
+                self.processes[process.pid]["status"] = ProcessStatus.FAILED
+                self.processes[process.pid]["properties"]["error_details"] = str(e)
+            self._update_server_status("Failed", error_details=str(e))
+            self._log_event("error", f"Error starting vLLM server: {str(e)}")
+            raise RuntimeError(f"Error starting vLLM server: {e}") from e
+
+    def _verify_server_health(self):
+        try:
+            response = requests.get(f"{self.vllm_url}/v1/models", timeout=10)
+            if response.status_code == 200:
+                if self.debug:
+                    models_data = response.json()
+                    available_models = [
+                        m.get("id", "unknown") for m in models_data.get("data", [])
+                    ]
+                    print(
+                        f"[@vllm] Health check OK. Available models: {available_models}"
+                    )
+                return True
+            else:
+                if self.debug:
+                    print(
+                        f"[@vllm] Health check failed with status {response.status_code}"
+                    )
+                return False
+        except Exception as e:
+            if self.debug:
+                print(f"[@vllm] Health check exception: {e}")
+            return False
+
+    def _collect_version_info(self):
+        version_info = {}
+        try:
+            import vllm
+
+            version_info["vllm"] = getattr(vllm, "__version__", "Unknown")
+        except ImportError:
+            version_info["vllm"] = "Not installed"
+        except Exception as e:
+            version_info["vllm"] = "Error detecting"
+            if self.debug:
+                print(f"[@vllm] Error getting vLLM version: {e}")
+
+        if self.status_card:
+            self.status_card.update_status("versions", version_info)
+            self._log_event(
+                "info", f"vLLM version: {version_info.get('vllm', 'Unknown')}"
+            )
+
+    def terminate_models(self):
+        shutdown_start_time = time.time()
+        self._log_event("info", "Starting vLLM shutdown sequence")
+        if self.debug:
+            print("[@vllm] Shutting down vLLM server...")
+
+        server_shutdown_cause = "graceful"
+
+        if self.server_process:
+            try:
+                self._update_server_status("Stopping")
+                self._log_event("info", "Stopping vLLM server")
+
+                # Clear model status since server is shutting down
+                self._update_model_status(self.model, status="Stopping")
+
+                self.server_process.terminate()
+                try:
+                    self.server_process.wait(timeout=10)
+                    if self.debug:
+                        print("[@vllm] Server terminated gracefully")
+                except subprocess.TimeoutExpired:
+                    server_shutdown_cause = "force_kill"
+                    self._log_event(
+                        "warning",
+                        "vLLM server did not terminate gracefully, killing...",
+                    )
+                    if self.debug:
+                        print("[@vllm] Server did not terminate, killing...")
+                    self.server_process.kill()
+                    self.server_process.wait()
+
+                if self.server_process.pid in self.processes:
+                    self.processes[self.server_process.pid][
+                        "status"
+                    ] = ProcessStatus.SUCCESSFUL
+
+                self._update_server_status("Stopped")
+                if self.status_card:
+                    self.status_card.update_status("models", {})
+
+                self._log_event(
+                    "success", f"vLLM server stopped ({server_shutdown_cause})"
+                )
+
+            except Exception as e:
+                server_shutdown_cause = "failed"
+                if self.server_process.pid in self.processes:
+                    self.processes[self.server_process.pid][
+                        "status"
+                    ] = ProcessStatus.FAILED
+                    self.processes[self.server_process.pid]["properties"][
+                        "error_details"
+                    ] = str(e)
+                self._update_server_status("Failed to stop")
+                if self.status_card:
+                    self.status_card.update_status("models", {})
+                self._log_event("error", f"vLLM server shutdown error: {str(e)}")
+                if self.debug:
+                    print(f"[@vllm] Warning: Error terminating vLLM server: {e}")
+
+        total_shutdown_time = time.time() - shutdown_start_time
+        self._update_performance("total_shutdown_time", total_shutdown_time)
+        self._update_performance("shutdown_cause", server_shutdown_cause)
+
+        self._log_event(
+            "success", f"vLLM shutdown completed in {total_shutdown_time:.1f}s"
+        )
+        if self.debug:
+            print("[@vllm] vLLM server shutdown complete.")

metaflow_extensions/outerbounds/toplevel/global_aliases_for_metaflow_package.py

@@ -75,4 +75,3 @@ from .. import profilers
 from ..plugins.snowflake import Snowflake
 from ..plugins.checkpoint_datastores import nebius_checkpoints, coreweave_checkpoints
 from ..plugins.aws import assume_role
-from . import ob_internal

metaflow_extensions/outerbounds/toplevel/plugins/vllm/__init__.py

@@ -0,0 +1 @@
+__mf_promote_submodules__ = ["plugins.vllm"]

{ob_metaflow_extensions-1.1.166rc5.dist-info → ob_metaflow_extensions-1.1.167.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ob-metaflow-extensions
-Version: 1.1.166rc5
+Version: 1.1.167
 Summary: Outerbounds Platform Extensions for Metaflow
 Author: Outerbounds, Inc.
 License: Commercial

{ob_metaflow_extensions-1.1.166rc5.dist-info → ob_metaflow_extensions-1.1.167.dist-info}/RECORD

@@ -1,11 +1,10 @@
 metaflow_extensions/outerbounds/__init__.py,sha256=Gb8u06s9ClQsA_vzxmkCzuMnigPy7kKcDnLfb7eB-64,514
 metaflow_extensions/outerbounds/remote_config.py,sha256=pEFJuKDYs98eoB_-ryPjVi9b_c4gpHMdBHE14ltoxIU,4672
 metaflow_extensions/outerbounds/config/__init__.py,sha256=JsQGRuGFz28fQWjUvxUgR8EKBLGRdLUIk_buPLJplJY,1225
-metaflow_extensions/outerbounds/plugins/__init__.py,sha256=HB298DqOlmM96-j_FjloDSgt1TZO9I-K2eJ7kevpHVQ,14092
+metaflow_extensions/outerbounds/plugins/__init__.py,sha256=gQIy3VaDg5e8tbx5q5TflddxhS6N3B7TaCWZEtnqJbo,14095
 metaflow_extensions/outerbounds/plugins/auth_server.py,sha256=_Q9_2EL0Xy77bCRphkwT1aSu8gQXRDOH-Z-RxTUO8N4,2202
 metaflow_extensions/outerbounds/plugins/perimeters.py,sha256=QXh3SFP7GQbS-RAIxUOPbhPzQ7KDFVxZkTdKqFKgXjI,2697
 metaflow_extensions/outerbounds/plugins/apps/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-metaflow_extensions/outerbounds/plugins/apps/app_cli.py,sha256=erfKCC6zKuwax0ye9j-tqIkAVdCg7MJfbRGhMhViSzU,575
 metaflow_extensions/outerbounds/plugins/apps/app_utils.py,sha256=sw9whU17lAzlD2K2kEDNjlk1Ib-2xE2UNhJkmzD8Qv8,8543
 metaflow_extensions/outerbounds/plugins/apps/consts.py,sha256=iHsyqbUg9k-rgswCs1Jxf5QZIxR1V-peCDRjgr9kdBM,177
 metaflow_extensions/outerbounds/plugins/apps/deploy_decorator.py,sha256=VkmiMdNYHhNdt-Qm9AVv7aE2LWFsIFEc16YcOYjwF6Q,8568
@@ -21,7 +20,6 @@ metaflow_extensions/outerbounds/plugins/checkpoint_datastores/__init__.py,sha256
 metaflow_extensions/outerbounds/plugins/checkpoint_datastores/coreweave.py,sha256=_WzoOROFjoFa8TzsMNFp-r_1Zz7NUp-5ljn_kKlczXA,4534
 metaflow_extensions/outerbounds/plugins/checkpoint_datastores/nebius.py,sha256=zgqDLFewCeF5jqh-hUNKmC_OAjld09ln0bb8Lkeqapc,4659
 metaflow_extensions/outerbounds/plugins/fast_bakery/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-metaflow_extensions/outerbounds/plugins/fast_bakery/baker.py,sha256=ShE5omFBr83wkvEhL_ptRFvDNMs6wefg4BjaafQjTcM,3602
 metaflow_extensions/outerbounds/plugins/fast_bakery/docker_environment.py,sha256=Tl520HdBteg-aDOM7mnnJJpdDCZc49BmFFmLUc_vTi8,15018
 metaflow_extensions/outerbounds/plugins/fast_bakery/fast_bakery.py,sha256=PE81ZB54OAMXkMGSB7JqgvgMg7N9kvoVclrWL-6jc2U,5626
 metaflow_extensions/outerbounds/plugins/fast_bakery/fast_bakery_cli.py,sha256=kqFyu2bJSnc9_9aYfBpz5xK6L6luWFZK_NMuh8f1eVk,1494
@@ -69,18 +67,23 @@ metaflow_extensions/outerbounds/plugins/snowpark/snowpark_job.py,sha256=aQphxX6j
 metaflow_extensions/outerbounds/plugins/snowpark/snowpark_service_spec.py,sha256=AI_kcm1hZV3JRxJkookcH6twiGnAYjk9Dx-MeoYz60Y,8511
 metaflow_extensions/outerbounds/plugins/tensorboard/__init__.py,sha256=9lUM4Cqi5RjrHBRfG6AQMRz8-R96eZC8Ih0KD2lv22Y,1858
 metaflow_extensions/outerbounds/plugins/torchtune/__init__.py,sha256=TOXNeyhcgd8VxplXO_oEuryFEsbk0tikn5GL0-44SU8,5853
+metaflow_extensions/outerbounds/plugins/vllm/__init__.py,sha256=O04DPVoEdCZhPbvdldaE4ztoAxJNXU-ExosBCqe43v8,6463
+metaflow_extensions/outerbounds/plugins/vllm/constants.py,sha256=ODX_uM5iYrzpVltsAdSf9Jo0DAOMiZ3647DcKdCnlS0,24
+metaflow_extensions/outerbounds/plugins/vllm/exceptions.py,sha256=8m65k2L17zXgSkgU299DWqxr1wGUMsZgSJw0hBRizJ0,49
+metaflow_extensions/outerbounds/plugins/vllm/status_card.py,sha256=ofTuBkhZgJq8cs-KwUOc85TB4hP9FT5qxagVEomY8jA,12931
+metaflow_extensions/outerbounds/plugins/vllm/vllm_manager.py,sha256=E6SCBoanNNsFjm-IGGGf4VMVe1mueDiRWw0ZvCXizDQ,18496
 metaflow_extensions/outerbounds/profilers/__init__.py,sha256=wa_jhnCBr82TBxoS0e8b6_6sLyZX0fdHicuGJZNTqKw,29
 metaflow_extensions/outerbounds/profilers/gpu.py,sha256=3Er8uKQzfm_082uadg4yn_D4Y-iSCgzUfFmguYxZsz4,27485
 metaflow_extensions/outerbounds/toplevel/__init__.py,sha256=qWUJSv_r5hXJ7jV_On4nEasKIfUCm6_UjkjXWA_A1Ts,90
-metaflow_extensions/outerbounds/toplevel/global_aliases_for_metaflow_package.py,sha256=_fKWv_-O1k5Nk5A1q05Ioh-PSsFXGL-jiAt7zfl8pIE,2999
-metaflow_extensions/outerbounds/toplevel/ob_internal.py,sha256=53xM6d_UYT3uGFFA59UzxN23H5QMO5_F39pALpmGy04,51
+metaflow_extensions/outerbounds/toplevel/global_aliases_for_metaflow_package.py,sha256=FS0ZKQJsJKlw9PgtLqVV4kH9o7_qwYcIVtyu2Kqwa_U,2973
 metaflow_extensions/outerbounds/toplevel/plugins/azure/__init__.py,sha256=WUuhz2YQfI4fz7nIcipwwWq781eaoHEk7n4GAn1npDg,63
 metaflow_extensions/outerbounds/toplevel/plugins/gcp/__init__.py,sha256=BbZiaH3uILlEZ6ntBLKeNyqn3If8nIXZFq_Apd7Dhco,70
 metaflow_extensions/outerbounds/toplevel/plugins/kubernetes/__init__.py,sha256=5zG8gShSj8m7rgF4xgWBZFuY3GDP5n1T0ktjRpGJLHA,69
 metaflow_extensions/outerbounds/toplevel/plugins/ollama/__init__.py,sha256=GRSz2zwqkvlmFS6bcfYD_CX6CMko9DHQokMaH1iBshA,47
 metaflow_extensions/outerbounds/toplevel/plugins/snowflake/__init__.py,sha256=LptpH-ziXHrednMYUjIaosS1SXD3sOtF_9_eRqd8SJw,50
 metaflow_extensions/outerbounds/toplevel/plugins/torchtune/__init__.py,sha256=uTVkdSk3xZ7hEKYfdlyVteWj5KeDwaM1hU9WT-_YKfI,50
-ob_metaflow_extensions-1.1.166rc5.dist-info/METADATA,sha256=EuJI-hYrzPPZYZcyQhz8GP0PHTOpJ10pM1SQdH2vFPo,524
-ob_metaflow_extensions-1.1.166rc5.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
-ob_metaflow_extensions-1.1.166rc5.dist-info/top_level.txt,sha256=NwG0ukwjygtanDETyp_BUdtYtqIA_lOjzFFh1TsnxvI,20
-ob_metaflow_extensions-1.1.166rc5.dist-info/RECORD,,
+metaflow_extensions/outerbounds/toplevel/plugins/vllm/__init__.py,sha256=ekcgD3KVydf-a0xMI60P4uy6ePkSEoFHiGnDq1JM940,45
+ob_metaflow_extensions-1.1.167.dist-info/METADATA,sha256=0WnYu65sqUQZp6jEd6aJrLMBEx6Nm1VIB7lFsfmQXaU,521
+ob_metaflow_extensions-1.1.167.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
+ob_metaflow_extensions-1.1.167.dist-info/top_level.txt,sha256=NwG0ukwjygtanDETyp_BUdtYtqIA_lOjzFFh1TsnxvI,20
+ob_metaflow_extensions-1.1.167.dist-info/RECORD,,

metaflow_extensions/outerbounds/plugins/apps/app_cli.py

@@ -1,26 +0,0 @@
-from metaflow._vendor import click
-
-OUTERBOUNDS_APP_CLI_AVAILABLE = True
-try:
-    import outerbounds.apps.app_cli as ob_apps_cli
-except ImportError:
-    OUTERBOUNDS_APP_CLI_AVAILABLE = False
-
-
-if not OUTERBOUNDS_APP_CLI_AVAILABLE:
-
-    @click.group()
-    def _cli():
-        pass
-
-    @_cli.group(help="Dummy Group to append to CLI for Safety")
-    def app():
-        pass
-
-    @app.command(help="Dummy Command to append to CLI for Safety")
-    def cannot_deploy():
-        raise Exception("Outerbounds App CLI not available")
-
-    cli = _cli
-else:
-    cli = ob_apps_cli.cli

metaflow_extensions/outerbounds/plugins/fast_bakery/baker.py

@@ -1,110 +0,0 @@
-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/toplevel/ob_internal.py

@@ -1 +0,0 @@
-from ..plugins.fast_bakery.baker import bake_image