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

metaflow_extensions/outerbounds/plugins/vllm/__init__.py

@@ -1,17 +1,30 @@
 from metaflow.decorators import StepDecorator
 from metaflow import current
 import functools
-import os
+from enum import Enum
 import threading
 from metaflow.unbounded_foreach import UBF_CONTROL, UBF_TASK
 from metaflow.metaflow_config import from_conf
 
-from .vllm_manager import VLLMManager
+from .vllm_manager import VLLMOpenAIManager, VLLMPyManager
 from .status_card import VLLMStatusCard, CardDecoratorInjector
 
 __mf_promote_submodules__ = ["plugins.vllm"]
 
 
+### The following classes are used to store the vLLM information in the current environment.
+# Then, Metaflow users can access the vLLM information through the current environment.
+class OpenAIAPIInfo:
+    def __init__(self, local_endpoint, local_api_key):
+        self.local_endpoint = local_endpoint
+        self.local_api_key = local_api_key
+
+
+class VLLM:
+    def __init__(self, llm):
+        self.llm = llm
+
+
 class VLLMDecorator(StepDecorator, CardDecoratorInjector):
     """
     This decorator is used to run vllm APIs as Metaflow task sidecars.
@@ -40,11 +53,23 @@ class VLLMDecorator(StepDecorator, CardDecoratorInjector):
         HuggingFace model identifier to be served by vLLM.
     backend: str
         Determines where and how to run the vLLM process.
+    openai_api_server: bool
+        Whether to use OpenAI-compatible API server mode (subprocess) instead of native engine.
+        Default is False (uses native engine).
+        Set to True for backward compatibility with existing code.
     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.
+    card_refresh_interval: int
+        Interval in seconds for refreshing the vLLM status card.
+        Only used when openai_api_server=True.
+    max_retries: int
+        Maximum number of retries checking for vLLM server startup.
+        Only used when openai_api_server=True.
+    retry_alert_frequency: int
+        Frequency of alert logs for vLLM server startup retries.
+        Only used when openai_api_server=True.
+    engine_args : dict
+        Additional keyword arguments to pass to the vLLM engine.
         For example, `tensor_parallel_size=2`.
     """
 
@@ -52,9 +77,12 @@ class VLLMDecorator(StepDecorator, CardDecoratorInjector):
     defaults = {
         "model": None,
         "backend": "local",
+        "openai_api_server": False,  # Default to native engine
         "debug": False,
         "stream_logs_to_card": False,
         "card_refresh_interval": 10,
+        "max_retries": 60,
+        "retry_alert_frequency": 5,
         "engine_args": {},
     }
 
@@ -72,106 +100,156 @@ class VLLMDecorator(StepDecorator, CardDecoratorInjector):
                 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"],
-        )
+        # Attach the vllm status card only for API server mode
+        if self.attributes["openai_api_server"]:
+            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
+            # FIXME: Kind of ugly branch. Causing branching elsewhere.
+            # Other possibile code paths:
+            # - OpenAI batch API
+            # - Embedding
+            # - Special types of models
+            if self.attributes["openai_api_server"]:
+                # API Server mode (existing functionality)
+                self._run_api_server_mode(step_func)
+            else:
+                # Native engine mode (new functionality)
+                self._run_native_engine_mode(step_func)
 
-            try:
-                self.status_card = VLLMStatusCard(
-                    refresh_interval=self.attributes["card_refresh_interval"]
-                )
+        return vllm_wrapper
 
-                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
+    def _run_api_server_mode(self, step_func):
+        """Run vLLM in API server mode (subprocess, existing functionality)"""
+        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 = VLLMOpenAIManager(
+                model=self.attributes["model"],
+                backend=self.attributes["backend"],
+                debug=self.attributes["debug"],
+                status_card=self.status_card,
+                max_retries=self.attributes["max_retries"],
+                retry_alert_frequency=self.attributes["retry_alert_frequency"],
+                stream_logs_to_card=self.attributes["stream_logs_to_card"],
+                **self.attributes["engine_args"],
+            )
+            current._update_env(
+                dict(
+                    vllm=OpenAIAPIInfo(
+                        local_endpoint=f"http://127.0.0.1:{self.vllm_manager.port}/v1",
+                        local_api_key="token123",
+                    )
                 )
-                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] API server mode 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 API server mode: {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"
                 )
-                if self.attributes["debug"]:
-                    print("[@vllm] VLLMManager initialized.")
+        finally:
+            if self.vllm_manager:
+                self.vllm_manager.terminate_models()
 
-            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.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("[@vllm] Card monitoring thread stopped.")
+                        print(f"[@vllm] Final card update error: {e}")
+                time.sleep(2)
 
-        return vllm_wrapper
+            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.")
+
+    def _run_native_engine_mode(self, step_func):
+        """Run vLLM in native engine mode (direct LLM API access)"""
+        self.vllm = None
+
+        try:
+            if self.attributes["debug"]:
+                print("[@vllm] Initializing native engine mode")
+
+            self.vllm = VLLMPyManager(
+                model=self.attributes["model"],
+                debug=self.attributes["debug"],
+                **self.attributes["engine_args"],
+            )
+            current._update_env(dict(vllm=VLLM(llm=self.vllm.engine)))
+
+            if self.attributes["debug"]:
+                print("[@vllm] Native engine mode initialized.")
+
+        except Exception as e:
+            print(f"[@vllm] Error initializing native engine mode: {e}")
+            raise
+
+        try:
+            step_func()
+        finally:
+            if self.vllm:
+                self.vllm.terminate_engine()

metaflow_extensions/outerbounds/plugins/vllm/status_card.py

@@ -160,10 +160,10 @@ class VLLMStatusCard(CardRefresher):
             # 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_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"""
@@ -252,11 +252,11 @@ class VLLMStatusCard(CardRefresher):
             )
 
         # Simplified monitoring note
-        current_card.append(
-            Markdown(
-                "## 🔧 Monitoring\n**Advanced Features:** Disabled (Circuit Breaker, Request Interception)"
-            )
-        )
+        # current_card.append(
+        #     Markdown(
+        #         "## 🔧 Monitoring\n**Advanced Features:** Disabled (Circuit Breaker, Request Interception)"
+        #     )
+        # )
 
         # Performance metrics
         perf_data = data["performance"]

metaflow_extensions/outerbounds/plugins/vllm/vllm_manager.py

@@ -4,10 +4,7 @@ import time
 import socket
 import sys
 import os
-import functools
-import json
 import requests
-from enum import Enum
 import threading
 from datetime import datetime
 
@@ -20,7 +17,147 @@ class ProcessStatus:
     SUCCESSFUL = "SUCCESSFUL"
 
 
-class VLLMManager:
+class VLLMPyManager:
+    """
+    A native vLLM engine manager that provides direct access to the vLLM LLM class.
+    This replaces the subprocess-based API server approach with direct Python API access.
+
+    Example usage:
+        from vllm.sampling_params import SamplingParams, GuidedDecodingParams
+
+        engine = current.vllm.engine
+        sampling_params = SamplingParams(temperature=0.7, max_tokens=150)
+        outputs = engine.generate(["Hello, world!"], sampling_params)
+
+        # Structured outputs
+        guided_params = GuidedDecodingParams(json=my_schema)
+        sampling_params = SamplingParams(guided_decoding=guided_params)
+        outputs = engine.generate(prompts, sampling_params)
+    """
+
+    def __init__(
+        self,
+        model,
+        debug=False,
+        **engine_args,
+    ):
+        if isinstance(model, list):
+            if len(model) != 1:
+                raise ValueError(
+                    f"vLLM native engine 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.debug = debug
+        self.engine_args = engine_args
+        self.engine = None
+        self.initialization_start = time.time()
+
+        if self.debug:
+            print(
+                f"[@vllm-native] Initializing native vLLM engine for model: {self.model}"
+            )
+
+        self._validate_vllm_installation()
+        self._initialize_engine()
+
+        total_init_time = time.time() - self.initialization_start
+        if self.debug:
+            print(
+                f"[@vllm-native] Native engine initialization completed in {total_init_time:.1f}s"
+            )
+
+    def _validate_vllm_installation(self):
+        """Validate that vLLM is properly installed"""
+        try:
+            import vllm
+
+            if self.debug:
+                print(f"[@vllm-native] vLLM {vllm.__version__} is available")
+        except ImportError as e:
+            raise ImportError(
+                "vLLM not installed. Please add vLLM to your environment."
+            ) from e
+
+    def _map_engine_args(self, engine_args):
+        """
+        Map CLI-style engine_args to LLM constructor parameters.
+        Most parameters map directly from the API server CLI args to LLM constructor.
+        """
+        llm_params = {}
+
+        # Direct mappings (parameter names are the same)
+        direct_mapping = [
+            "tensor_parallel_size",
+            "max_model_len",
+            "gpu_memory_utilization",
+            "swap_space",
+            "dtype",
+            "quantization",
+            "seed",
+            "trust_remote_code",
+            "revision",
+            "tokenizer_revision",
+            "enforce_eager",
+            "max_seq_len_to_capture",
+            "disable_custom_all_reduce",
+        ]
+
+        for param in direct_mapping:
+            if param in engine_args:
+                llm_params[param] = engine_args[param]
+
+        # Handle special mappings if needed
+        # (Most/all vLLM CLI args map directly to LLM constructor args)
+
+        return llm_params
+
+    def _initialize_engine(self):
+        """Initialize the native vLLM LLM engine"""
+        try:
+            from vllm import LLM
+
+            # Map engine args to LLM constructor parameters
+            llm_params = self._map_engine_args(self.engine_args)
+
+            if self.debug:
+                print(f"[@vllm] Initializing LLM with params: {llm_params}")
+
+            # Initialize the native vLLM engine
+            self.engine = LLM(model=self.model, **llm_params)
+
+            if self.debug:
+                print(f"[@vllm] LLM engine initialized successfully")
+
+        except Exception as e:
+            error_msg = f"Failed to initialize vLLM engine: {str(e)}"
+            if self.debug:
+                print(f"[@vllm-native] ERROR: {error_msg}")
+            raise RuntimeError(error_msg) from e
+
+    def terminate_engine(self):
+        """
+        Clean up the native engine.
+        The LLM class handles cleanup automatically when the object is destroyed.
+        """
+        if self.debug:
+            print("[@vllm-] Cleaning up vLLM engine")
+
+        # The vLLM LLM class handles cleanup automatically
+        # We just need to clear our reference
+        if self.engine:
+            del self.engine
+            self.engine = None
+
+        if self.debug:
+            print("[@vllm] Engine cleanup completed")
+
+
+class VLLMOpenAIManager:
     """
     A process manager for vLLM runtimes.
     Implements interface @vllm(model=..., ...) to provide a local backend.
@@ -55,6 +192,8 @@ class VLLMManager:
         port=8000,
         host="127.0.0.1",
         stream_logs_to_card=False,
+        max_retries=60,
+        retry_alert_frequency=5,
         **vllm_args,
     ):
         # Validate that only a single model is provided
@@ -79,6 +218,8 @@ class VLLMManager:
         self.status_card = status_card
         self.initialization_start = time.time()
         self.server_process = None
+        self.max_retries = max_retries
+        self.retry_alert_frequency = retry_alert_frequency
         self.vllm_args = vllm_args
 
         if backend != "local":
@@ -211,6 +352,13 @@ class VLLMManager:
                     f"[@vllm] Starting vLLM OpenAI-compatible server for model: {self.model}"
                 )
 
+            ### NOTE: This is not the only way to start the vLLM server.
+            #  https://docs.vllm.ai/en/v0.9.0/api/vllm/entrypoints/openai/api_server.html
+
+            # There are other APIs we should consider using in a future extension:
+            # https://docs.vllm.ai/en/stable/api/vllm/entrypoints/openai/run_batch.html#vllm.entrypoints.openai.run_batch
+            # https://docs.vllm.ai/en/v0.9.0/api/vllm/entrypoints/openai/serving_embedding.html
+            # MANY MORE!!! Wait for some feedback and we can add more.
             cmd = [
                 sys.executable,
                 "-m",
@@ -226,6 +374,8 @@ class VLLMManager:
             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.
+                # It was useful for debugging the vLLM server startup,
+                # likely more confusion potential than its worth for end user.
                 vllm_args_copy.setdefault("uvicorn_log_level", "debug")
 
             for key, value in vllm_args_copy.items():
@@ -281,16 +431,15 @@ class VLLMManager:
                 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
+                and retries < self.max_retries
             ):
                 if retries == 0:
                     print("[@vllm] Waiting for server to be ready...")
-                elif retries % 10 == 0:
+                elif retries % self.retry_alert_frequency == 0:
                     print(
-                        f"[@vllm] Still waiting for server... ({retries}/{max_retries})"
+                        f"[@vllm] Still waiting for server... ({retries}/{self.max_retries})"
                     )
 
                 returncode = process.poll()
@@ -322,7 +471,7 @@ class VLLMManager:
                 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"
+                error_details = f"vLLM server did not start listening on {self.host}:{self.port} after {self.max_retries*2}s"
                 self.processes[process.pid]["properties"][
                     "error_details"
                 ] = error_details
@@ -342,6 +491,7 @@ class VLLMManager:
                 "Running", uptime_start=datetime.now(), model=self.model
             )
             self._log_event("success", "vLLM server is ready and listening")
+            print(f"[@vllm] Server ready!")
 
             self._update_model_status(self.model, status="Ready")
 

metaflow_extensions/outerbounds/remote_config.py

@@ -11,6 +11,11 @@ from metaflow_extensions.outerbounds.plugins.perimeters import (
     get_perimeter_config_url_if_set_in_ob_config,
 )
 
+
+class OuterboundsConfigException(MetaflowException):
+    _OB_CONFIG_EXCEPTION = True
+
+
 OBP_REMOTE_CONFIG_KEY = "OBP_METAFLOW_CONFIG_URL"
 HOSTNAME_KEY = "OBP_API_SERVER"
 AUTH_KEY = "METAFLOW_SERVICE_AUTH_KEY"
@@ -31,7 +36,7 @@ def read_config_from_local() -> Optional[Path]:
 
     # we should error because the user wants a specific config
     if profile:
-        raise MetaflowException(
+        raise OuterboundsConfigException(
             f"Unable to locate METAFLOW_PROFILE {profile} in {config_path}"
         )
 
@@ -55,7 +60,7 @@ def resolve_config_from_remote(remote_url: str, auth_token: str) -> Dict[str, st
         data = response.json()
         return data["config"]
     except HTTPError:
-        raise MetaflowException(
+        raise OuterboundsConfigException(
             "Error fetching resolving configuration. Make sure you have run \
                                 `outerbounds configure` with the correct value"
         )
@@ -81,7 +86,7 @@ def init_config() -> Dict[str, str]:
     try:
         remote_config = json.loads(config_path.read_text())
     except ValueError:
-        raise MetaflowException(
+        raise OuterboundsConfigException(
             "Error decoding your metaflow config. Please run the `outerbounds configure` \
                                  command with the string provided in the Outerbounds dashboard"
         )