vec-inf 0.4.1__py3-none-any.whl → 0.6.0__py3-none-any.whl

vec_inf/cli/_helper.py

@@ -0,0 +1,400 @@
+"""Helper classes for the CLI.
+
+This module provides formatting and display classes for the command-line interface,
+handling the presentation of model information, status updates, and metrics.
+"""
+
+from pathlib import Path
+from typing import Any, Union
+
+import click
+from rich.columns import Columns
+from rich.console import Console
+from rich.panel import Panel
+from rich.table import Table
+
+from vec_inf.cli._utils import create_table
+from vec_inf.cli._vars import MODEL_TYPE_COLORS, MODEL_TYPE_PRIORITY
+from vec_inf.client import ModelConfig, ModelInfo, StatusResponse
+
+
+class LaunchResponseFormatter:
+    """CLI Helper class for formatting LaunchResponse.
+
+    A formatter class that handles the presentation of model launch information
+    in both table and JSON formats.
+
+    Parameters
+    ----------
+    model_name : str
+        Name of the launched model
+    params : dict[str, Any]
+        Launch parameters and configuration
+    """
+
+    def __init__(self, model_name: str, params: dict[str, Any]):
+        self.model_name = model_name
+        self.params = params
+
+    def format_table_output(self) -> Table:
+        """Format output as rich Table.
+
+        Returns
+        -------
+        Table
+            Rich table containing formatted launch information including:
+            - Job configuration
+            - Model details
+            - Resource allocation
+            - vLLM configuration
+        """
+        table = create_table(key_title="Job Config", value_title="Value")
+
+        # Add key information with consistent styling
+        table.add_row("Slurm Job ID", self.params["slurm_job_id"], style="blue")
+        table.add_row("Job Name", self.model_name)
+
+        # Add model details
+        table.add_row("Model Type", self.params["model_type"])
+        table.add_row("Vocabulary Size", self.params["vocab_size"])
+
+        # Add resource allocation details
+        table.add_row("Partition", self.params["partition"])
+        table.add_row("QoS", self.params["qos"])
+        table.add_row("Time Limit", self.params["time"])
+        table.add_row("Num Nodes", self.params["num_nodes"])
+        table.add_row("GPUs/Node", self.params["gpus_per_node"])
+        table.add_row("CPUs/Task", self.params["cpus_per_task"])
+        table.add_row("Memory/Node", self.params["mem_per_node"])
+
+        # Add job config details
+        table.add_row(
+            "Model Weights Directory",
+            str(Path(self.params["model_weights_parent_dir"], self.model_name)),
+        )
+        table.add_row("Log Directory", self.params["log_dir"])
+
+        # Add vLLM configuration details
+        table.add_row("vLLM Arguments:", style="magenta")
+        for arg, value in self.params["vllm_args"].items():
+            table.add_row(f"  {arg}:", str(value))
+
+        return table
+
+
+class StatusResponseFormatter:
+    """CLI Helper class for formatting StatusResponse.
+
+    A formatter class that handles the presentation of model status information
+    in both table and JSON formats.
+
+    Parameters
+    ----------
+    status_info : StatusResponse
+        Status information to format
+    """
+
+    def __init__(self, status_info: StatusResponse):
+        self.status_info = status_info
+
+    def output_json(self) -> None:
+        """Format and output JSON data.
+
+        Outputs a JSON object containing:
+        - model_name
+        - model_status
+        - base_url
+        - pending_reason (if applicable)
+        - failed_reason (if applicable)
+        """
+        json_data = {
+            "model_name": self.status_info.model_name,
+            "model_status": self.status_info.server_status,
+            "base_url": self.status_info.base_url,
+        }
+        if self.status_info.pending_reason:
+            json_data["pending_reason"] = self.status_info.pending_reason
+        if self.status_info.failed_reason:
+            json_data["failed_reason"] = self.status_info.failed_reason
+        click.echo(json_data)
+
+    def output_table(self) -> Table:
+        """Create and display rich table.
+
+        Returns
+        -------
+        Table
+            Rich table containing formatted status information including:
+            - Model name
+            - Status
+            - Base URL
+            - Error information (if applicable)
+        """
+        table = create_table(key_title="Job Status", value_title="Value")
+        table.add_row("Model Name", self.status_info.model_name)
+        table.add_row("Model Status", self.status_info.server_status, style="blue")
+
+        if self.status_info.pending_reason:
+            table.add_row("Pending Reason", self.status_info.pending_reason)
+        if self.status_info.failed_reason:
+            table.add_row("Failed Reason", self.status_info.failed_reason)
+
+        table.add_row("Base URL", self.status_info.base_url)
+        return table
+
+
+class MetricsResponseFormatter:
+    """CLI Helper class for formatting MetricsResponse.
+
+    A formatter class that handles the presentation of model metrics
+    in a table format.
+
+    Parameters
+    ----------
+    metrics : Union[dict[str, float], str]
+        Dictionary of metrics or error message
+    """
+
+    def __init__(self, metrics: Union[dict[str, float], str]):
+        self.metrics = self._set_metrics(metrics)
+        self.table = create_table("Metric", "Value")
+        self.enabled_prefix_caching = self._check_prefix_caching()
+
+    def _set_metrics(self, metrics: Union[dict[str, float], str]) -> dict[str, float]:
+        """Set the metrics attribute.
+
+        Parameters
+        ----------
+        metrics : Union[dict[str, float], str]
+            Raw metrics data
+
+        Returns
+        -------
+        dict[str, float]
+            Processed metrics dictionary
+        """
+        return metrics if isinstance(metrics, dict) else {}
+
+    def _check_prefix_caching(self) -> bool:
+        """Check if prefix caching is enabled.
+
+        Returns
+        -------
+        bool
+            True if prefix caching metrics are present
+        """
+        return self.metrics.get("gpu_prefix_cache_hit_rate") is not None
+
+    def format_failed_metrics(self, message: str) -> None:
+        """Format error message for failed metrics collection.
+
+        Parameters
+        ----------
+        message : str
+            Error message to display
+        """
+        self.table.add_row("ERROR", message)
+
+    def format_metrics(self) -> None:
+        """Format and display all available metrics.
+
+        Formats and adds to the table:
+        - Throughput metrics
+        - Request queue metrics
+        - Cache usage metrics
+        - Prefix cache metrics (if enabled)
+        - Latency metrics
+        - Token counts
+        """
+        # Throughput metrics
+        self.table.add_row(
+            "Prompt Throughput",
+            f"{self.metrics.get('prompt_tokens_per_sec', 0):.1f} tokens/s",
+        )
+        self.table.add_row(
+            "Generation Throughput",
+            f"{self.metrics.get('generation_tokens_per_sec', 0):.1f} tokens/s",
+        )
+
+        # Request queue metrics
+        self.table.add_row(
+            "Requests Running",
+            f"{self.metrics.get('requests_running', 0):.0f} reqs",
+        )
+        self.table.add_row(
+            "Requests Waiting",
+            f"{self.metrics.get('requests_waiting', 0):.0f} reqs",
+        )
+        self.table.add_row(
+            "Requests Swapped",
+            f"{self.metrics.get('requests_swapped', 0):.0f} reqs",
+        )
+
+        # Cache usage metrics
+        self.table.add_row(
+            "GPU Cache Usage",
+            f"{self.metrics.get('gpu_cache_usage', 0) * 100:.1f}%",
+        )
+        self.table.add_row(
+            "CPU Cache Usage",
+            f"{self.metrics.get('cpu_cache_usage', 0) * 100:.1f}%",
+        )
+
+        if self.enabled_prefix_caching:
+            self.table.add_row(
+                "GPU Prefix Cache Hit Rate",
+                f"{self.metrics.get('gpu_prefix_cache_hit_rate', 0) * 100:.1f}%",
+            )
+            self.table.add_row(
+                "CPU Prefix Cache Hit Rate",
+                f"{self.metrics.get('cpu_prefix_cache_hit_rate', 0) * 100:.1f}%",
+            )
+
+        # Show average latency if available
+        if "avg_request_latency" in self.metrics:
+            self.table.add_row(
+                "Avg Request Latency",
+                f"{self.metrics['avg_request_latency']:.1f} s",
+            )
+
+        # Token counts
+        self.table.add_row(
+            "Total Prompt Tokens",
+            f"{self.metrics.get('total_prompt_tokens', 0):.0f} tokens",
+        )
+        self.table.add_row(
+            "Total Generation Tokens",
+            f"{self.metrics.get('total_generation_tokens', 0):.0f} tokens",
+        )
+        self.table.add_row(
+            "Successful Requests",
+            f"{self.metrics.get('successful_requests_total', 0):.0f} reqs",
+        )
+
+
+class ListCmdDisplay:
+    """CLI Helper class for displaying model listing functionality.
+
+    A display class that handles the presentation of model listings
+    in both table and JSON formats.
+
+    Parameters
+    ----------
+    console : Console
+        Rich console instance for output
+    json_mode : bool, default=False
+        Whether to output in JSON format
+    """
+
+    def __init__(self, console: Console, json_mode: bool = False):
+        self.console = console
+        self.json_mode = json_mode
+        self.model_config = None
+        self.model_names: list[str] = []
+
+    def _format_single_model_output(
+        self, config: ModelConfig
+    ) -> Union[dict[str, Any], Table]:
+        """Format output table for a single model.
+
+        Parameters
+        ----------
+        config : ModelConfig
+            Model configuration to format
+
+        Returns
+        -------
+        Union[dict[str, Any], Table]
+            Either a dictionary for JSON output or a Rich table
+        """
+        if self.json_mode:
+            # Exclude non-essential fields from JSON output
+            excluded = {"venv", "log_dir"}
+            config_dict = config.model_dump(exclude=excluded)
+            # Convert Path objects to strings
+            config_dict["model_weights_parent_dir"] = str(
+                config_dict["model_weights_parent_dir"]
+            )
+            return config_dict
+
+        table = create_table(key_title="Model Config", value_title="Value")
+        for field, value in config.model_dump().items():
+            if field not in {"venv", "log_dir", "vllm_args"}:
+                table.add_row(field, str(value))
+            if field == "vllm_args":
+                table.add_row("vLLM Arguments:", style="magenta")
+                for vllm_arg, vllm_value in value.items():
+                    table.add_row(f"  {vllm_arg}:", str(vllm_value))
+        return table
+
+    def _format_all_models_output(
+        self, model_infos: list[ModelInfo]
+    ) -> Union[list[str], list[Panel]]:
+        """Format output table for all models.
+
+        Parameters
+        ----------
+        model_infos : list[ModelInfo]
+            List of model information to format
+
+        Returns
+        -------
+        Union[list[str], list[Panel]]
+            Either a list of model names or a list of formatted panels
+
+        Notes
+        -----
+        Models are sorted by type priority and color-coded based on their type.
+        """
+        # Sort by model type priority
+        sorted_model_infos = sorted(
+            model_infos,
+            key=lambda x: MODEL_TYPE_PRIORITY.get(x.model_type, 4),
+        )
+
+        # Create panels with color coding
+        panels = []
+        for model_info in sorted_model_infos:
+            color = MODEL_TYPE_COLORS.get(model_info.model_type, "white")
+            variant = model_info.variant or ""
+            display_text = f"[magenta]{model_info.family}[/magenta]"
+            if variant:
+                display_text += f"-{variant}"
+            panels.append(Panel(display_text, expand=True, border_style=color))
+
+        return panels
+
+    def display_single_model_output(self, config: ModelConfig) -> None:
+        """Display the output for a single model.
+
+        Parameters
+        ----------
+        config : ModelConfig
+            Model configuration to display
+        """
+        output = self._format_single_model_output(config)
+        if self.json_mode:
+            click.echo(output)
+        else:
+            self.console.print(output)
+
+    def display_all_models_output(self, model_infos: list[ModelInfo]) -> None:
+        """Display the output for all models.
+
+        Parameters
+        ----------
+        model_infos : list[ModelInfo]
+            List of model information to display
+
+        Notes
+        -----
+        Output format depends on json_mode:
+        - JSON: List of model names
+        - Table: Color-coded panels with model information
+        """
+        if self.json_mode:
+            model_names = [info.name for info in model_infos]
+            click.echo(model_names)
+        else:
+            panels = self._format_all_models_output(model_infos)
+            self.console.print(Columns(panels, equal=True))

vec_inf/cli/_utils.py

@@ -1,147 +1,38 @@
-"""Utility functions for the CLI."""
+"""Helper functions for the CLI.
 
-import os
-import subprocess
-from typing import Dict, List, Optional, Tuple, Union, cast
+This module provides utility functions for creating consistent table displays
+in the command-line interface.
+"""
 
-import polars as pl
-import requests
 from rich.table import Table
 
 
-MODEL_READY_SIGNATURE = "INFO:     Application startup complete."
-SERVER_ADDRESS_SIGNATURE = "Server address: "
-
-
-def run_bash_command(command: str) -> str:
-    """Run a bash command and return the output."""
-    process = subprocess.Popen(
-        command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True
-    )
-    stdout, _ = process.communicate()
-    return stdout
-
-
-def read_slurm_log(
-    slurm_job_name: str, slurm_job_id: int, slurm_log_type: str, log_dir: Optional[str]
-) -> Union[list[str], str]:
-    """Read the slurm log file."""
-    if not log_dir:
-        models_dir = os.path.join(os.path.expanduser("~"), ".vec-inf-logs")
-
-        for directory in sorted(os.listdir(models_dir), key=len, reverse=True):
-            if directory in slurm_job_name:
-                log_dir = os.path.join(models_dir, directory)
-                break
-
-    log_dir = cast(str, log_dir)
-
-    try:
-        file_path = os.path.join(
-            log_dir,
-            f"{slurm_job_name}.{slurm_job_id}.{slurm_log_type}",
-        )
-        with open(file_path, "r") as file:
-            lines = file.readlines()
-    except FileNotFoundError:
-        print(f"Could not find file: {file_path}")
-        return "LOG_FILE_NOT_FOUND"
-    return lines
-
-
-def is_server_running(
-    slurm_job_name: str, slurm_job_id: int, log_dir: Optional[str]
-) -> Union[str, tuple[str, str]]:
-    """Check if a model is ready to serve requests."""
-    log_content = read_slurm_log(slurm_job_name, slurm_job_id, "err", log_dir)
-    if isinstance(log_content, str):
-        return log_content
-
-    status: Union[str, tuple[str, str]] = "LAUNCHING"
-
-    for line in log_content:
-        if "error" in line.lower():
-            status = ("FAILED", line.strip("\n"))
-        if MODEL_READY_SIGNATURE in line:
-            status = "RUNNING"
-
-    return status
-
-
-def get_base_url(slurm_job_name: str, slurm_job_id: int, log_dir: Optional[str]) -> str:
-    """Get the base URL of a model."""
-    log_content = read_slurm_log(slurm_job_name, slurm_job_id, "out", log_dir)
-    if isinstance(log_content, str):
-        return log_content
-
-    for line in log_content:
-        if SERVER_ADDRESS_SIGNATURE in line:
-            return line.split(SERVER_ADDRESS_SIGNATURE)[1].strip("\n")
-    return "URL_NOT_FOUND"
-
-
-def model_health_check(
-    slurm_job_name: str, slurm_job_id: int, log_dir: Optional[str]
-) -> Tuple[str, Union[str, int]]:
-    """Check the health of a running model on the cluster."""
-    base_url = get_base_url(slurm_job_name, slurm_job_id, log_dir)
-    if not base_url.startswith("http"):
-        return ("FAILED", base_url)
-    health_check_url = base_url.replace("v1", "health")
-
-    try:
-        response = requests.get(health_check_url)
-        # Check if the request was successful
-        if response.status_code == 200:
-            return ("READY", response.status_code)
-        return ("FAILED", response.status_code)
-    except requests.exceptions.RequestException as e:
-        return ("FAILED", str(e))
-
-
 def create_table(
     key_title: str = "", value_title: str = "", show_header: bool = True
 ) -> Table:
-    """Create a table for displaying model status."""
+    """Create a table for displaying model status.
+
+    Creates a two-column Rich table with consistent styling for displaying
+    key-value pairs in the CLI.
+
+    Parameters
+    ----------
+    key_title : str, default=""
+        Title for the key column
+    value_title : str, default=""
+        Title for the value column
+    show_header : bool, default=True
+        Whether to display column headers
+
+    Returns
+    -------
+    Table
+        Rich Table instance with configured styling:
+        - Headers in bold magenta
+        - Key column in dim style
+        - Value column in default style
+    """
     table = Table(show_header=show_header, header_style="bold magenta")
     table.add_column(key_title, style="dim")
     table.add_column(value_title)
     return table
-
-
-def load_models_df() -> pl.DataFrame:
-    """Load the models dataframe."""
-    return pl.read_csv(
-        os.path.join(
-            os.path.dirname(os.path.dirname(os.path.realpath(__file__))),
-            "models/models.csv",
-        )
-    )
-
-
-def load_default_args(models_df: pl.DataFrame, model_name: str) -> Dict[str, str]:
-    """Load the default arguments for a model."""
-    row_data = models_df.filter(models_df["model_name"] == model_name)
-    default_args = row_data.to_dicts()[0]
-    default_args.pop("model_name", None)
-    return default_args
-
-
-def get_latest_metric(log_lines: List[str]) -> Union[str, Dict[str, str]]:
-    """Read the latest metric entry from the log file."""
-    latest_metric = {}
-
-    try:
-        for line in reversed(log_lines):
-            if "Avg prompt throughput" in line:
-                # Parse the metric values from the line
-                metrics_str = line.split("] ")[1].strip().strip(".")
-                metrics_list = metrics_str.split(", ")
-                for metric in metrics_list:
-                    key, value = metric.split(": ")
-                    latest_metric[key] = value
-                break
-    except Exception as e:
-        return f"[red]Error reading log file: {e}[/red]"
-
-    return latest_metric

vec_inf/cli/_vars.py

@@ -0,0 +1,32 @@
+"""Constants for CLI rendering.
+
+This module defines constant mappings for model type priorities and colors
+used in the CLI display formatting.
+
+Constants
+---------
+MODEL_TYPE_PRIORITY : dict
+    Mapping of model types to their display priority (lower numbers shown first)
+
+MODEL_TYPE_COLORS : dict
+    Mapping of model types to their display colors in Rich
+
+Notes
+-----
+These constants are used primarily by the ListCmdDisplay class to ensure
+consistent sorting and color coding of different model types in the CLI output.
+"""
+
+MODEL_TYPE_PRIORITY = {
+    "LLM": 0,
+    "VLM": 1,
+    "Text_Embedding": 2,
+    "Reward_Modeling": 3,
+}
+
+MODEL_TYPE_COLORS = {
+    "LLM": "cyan",
+    "VLM": "bright_blue",
+    "Text_Embedding": "purple",
+    "Reward_Modeling": "bright_magenta",
+}

vec_inf/client/__init__.py

@@ -0,0 +1,31 @@
+"""Programmatic API for Vector Inference.
+
+This module provides a Python API for launching and managing inference servers
+using `vec_inf`. It is an alternative to the command-line interface, and allows
+users direct control over the lifecycle of inference servers via python scripts.
+"""
+
+from vec_inf.client.api import VecInfClient
+from vec_inf.client.config import ModelConfig
+from vec_inf.client.models import (
+    LaunchOptions,
+    LaunchResponse,
+    MetricsResponse,
+    ModelInfo,
+    ModelStatus,
+    ModelType,
+    StatusResponse,
+)
+
+
+__all__ = [
+    "VecInfClient",
+    "LaunchResponse",
+    "StatusResponse",
+    "ModelInfo",
+    "MetricsResponse",
+    "ModelStatus",
+    "ModelType",
+    "LaunchOptions",
+    "ModelConfig",
+]