ob-metaflow-extensions 1.1.166rc6__tar.gz → 1.1.168__tar.gz

{ob_metaflow_extensions-1.1.166rc6 → ob_metaflow_extensions-1.1.168}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ob_metaflow_extensions
-Version: 1.1.166rc6
+Version: 1.1.168
 Summary: Outerbounds Platform Extensions for Metaflow
 Author: Outerbounds, Inc.
 License: Commercial

{ob_metaflow_extensions-1.1.166rc6 → ob_metaflow_extensions-1.1.168}/metaflow_extensions/outerbounds/plugins/__init__.py

@@ -52,10 +52,6 @@ def get_boto3_session(role_arn=None, session_vars=None):
 
     token_info = get_token("/generate/aws")
 
-    # Check if the assume_role decorator has set a role ARN via environment variable
-    # This takes precedence over CSPR role
-    decorator_role_arn = os.environ.get(OBP_ASSUME_ROLE_ARN_ENV_VAR)
-
     # Write token to a file. The file name is derived from the user name
     # so it works with multiple users on the same machine.
     #
@@ -76,13 +72,18 @@ def get_boto3_session(role_arn=None, session_vars=None):
     if token_info.get("cspr_role_arn"):
         cspr_role = token_info["cspr_role_arn"]
 
-    # If assume_role decorator is used, prioritize it over CSPR role
-    effective_role = decorator_role_arn or cspr_role
-    if effective_role:
-        # If we have either a decorator role or CSPR role, set up AWS config
+    # Check if the assume_role decorator has set a CSPR ARN via environment variable
+    # This takes precedence over CSPR role that comes from the token_info response
+    decorator_role_arn = os.environ.get(OBP_ASSUME_ROLE_ARN_ENV_VAR)
+    if decorator_role_arn:
+        cspr_role = decorator_role_arn
+
+    if cspr_role:
+        # If CSPR role is set, we set it as the default role to assume
+        # for the AWS SDK. We do this by writing an AWS config file
         # with two profiles. One to get credentials for the task role
         # in exchange for the OIDC token, and second to assume the
-        # effective role using the task role credentials.
+        # CSPR role using the task role credentials.
         import configparser
         from io import StringIO
 
@@ -94,9 +95,9 @@ def get_boto3_session(role_arn=None, session_vars=None):
             "web_identity_token_file": token_file,
         }
 
-        # Effective role profile (decorator role or CSPR role)
-        aws_config["profile effective"] = {
-            "role_arn": effective_role,
+        # CSPR role profile (default)
+        aws_config["profile cspr"] = {
+            "role_arn": cspr_role,
             "source_profile": "task",
         }
 
@@ -112,7 +113,7 @@ def get_boto3_session(role_arn=None, session_vars=None):
             tmp_aws_config_file = f.name
         os.rename(tmp_aws_config_file, aws_config_file)
         os.environ["AWS_CONFIG_FILE"] = aws_config_file
-        os.environ["AWS_PROFILE"] = "effective"
+        os.environ["AWS_PROFILE"] = "cspr"
     else:
         os.environ["AWS_WEB_IDENTITY_TOKEN_FILE"] = token_file
         os.environ["AWS_ROLE_ARN"] = token_info["role_arn"]
@@ -124,29 +125,21 @@ def get_boto3_session(role_arn=None, session_vars=None):
     if token_info.get("region"):
         os.environ["AWS_DEFAULT_REGION"] = token_info["region"]
 
-    if effective_role:
-        if role_arn == decorator_role_arn or role_arn == USE_CSPR_ROLE_ARN_IF_SET:
-            # We have either a decorator role or CSPR role, use the effective profile
-            # The generated AWS config will be used here since we set the
-            # AWS_CONFIG_FILE environment variable above.
-            session = boto3.session.Session(profile_name="effective")
+    if cspr_role:
+        #  The generated AWS config will be used here since we set the
+        #  AWS_CONFIG_FILE environment variable above.
+        if role_arn == USE_CSPR_ROLE_ARN_IF_SET:
+            # Otherwise start from the default profile, assuming CSPR role
+            session = boto3.session.Session(profile_name="cspr")
         else:
             session = boto3.session.Session(profile_name="task")
     else:
-        # No decorator role or CSPR role, use default session
         # Not using AWS config, just AWS_WEB_IDENTITY_TOKEN_FILE + AWS_ROLE_ARN
         session = boto3.session.Session()
 
-    if (
-        role_arn
-        and role_arn != USE_CSPR_ROLE_ARN_IF_SET
-        and role_arn != decorator_role_arn
-    ):
-        # If the user provided a role_arn that's different from the decorator role,
-        # we assume that role using the current session credentials.
-        # This works for both cases:
-        # 1. No decorator role: Task role -> Secrets role
-        # 2. With decorator role: Decorator role -> Secrets role
+    if role_arn and role_arn != USE_CSPR_ROLE_ARN_IF_SET:
+        # If the user provided a role_arn, we assume that role
+        # using the task role credentials. CSPR role is not used.
         fetcher = botocore.credentials.AssumeRoleCredentialFetcher(
             client_creator=session._session.create_client,
             source_credentials=session._session.get_credentials(),
@@ -162,8 +155,8 @@ def get_boto3_session(role_arn=None, session_vars=None):
     else:
         # If the user didn't provide a role_arn, or if the role_arn
         # is set to USE_CSPR_ROLE_ARN_IF_SET, we return the default
-        # session which would use the effective role (decorator or CSPR) if set,
-        # or the task role otherwise.
+        # session which would use the CSPR role if it is set on the
+        # server, and the task role otherwise.
         return session
 
 
@@ -331,7 +324,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 +337,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"),
 ]
 

ob_metaflow_extensions-1.1.168/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

ob_metaflow_extensions-1.1.168/metaflow_extensions/outerbounds/plugins/vllm/constants.py

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

ob_metaflow_extensions-1.1.168/metaflow_extensions/outerbounds/plugins/vllm/exceptions.py

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

ob_metaflow_extensions-1.1.168/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}