vec-inf 0.5.0__py3-none-any.whl → 0.6.1__py3-none-any.whl

vec_inf/cli/_utils.py

@@ -1,162 +1,38 @@
-"""Utility functions for the CLI."""
+"""Helper functions for the CLI.
 
-import json
-import os
-import subprocess
-from pathlib import Path
-from typing import Any, Optional, Union, cast
+This module provides utility functions for creating consistent table displays
+in the command-line interface.
+"""
 
-import requests
-import yaml
 from rich.table import Table
 
-from vec_inf.cli._config import ModelConfig
-
-
-MODEL_READY_SIGNATURE = "INFO:     Application startup complete."
-CACHED_CONFIG = Path("/", "model-weights", "vec-inf-shared", "models.yaml")
-
-
-def run_bash_command(command: str) -> tuple[str, str]:
-    """Run a bash command and return the output."""
-    process = subprocess.Popen(
-        command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True
-    )
-    return process.communicate()
-
-
-def read_slurm_log(
-    slurm_job_name: str,
-    slurm_job_id: int,
-    slurm_log_type: str,
-    log_dir: Optional[Union[str, Path]],
-) -> Union[list[str], str, dict[str, str]]:
-    """Read the slurm log file."""
-    if not log_dir:
-        # Default log directory
-        models_dir = Path.home() / ".vec-inf-logs"
-        if not models_dir.exists():
-            return "LOG DIR NOT FOUND"
-        # Iterate over all dirs in models_dir, sorted by dir name length in desc order
-        for directory in sorted(
-            [d for d in models_dir.iterdir() if d.is_dir()],
-            key=lambda d: len(d.name),
-            reverse=True,
-        ):
-            if directory.name in slurm_job_name:
-                log_dir = directory
-                break
-    else:
-        log_dir = Path(log_dir)
-
-    # If log_dir is still not set, then didn't find the log dir at default location
-    if not log_dir:
-        return "LOG DIR NOT FOUND"
-
-    try:
-        file_path = (
-            log_dir
-            / Path(f"{slurm_job_name}.{slurm_job_id}")
-            / f"{slurm_job_name}.{slurm_job_id}.{slurm_log_type}"
-        )
-        if slurm_log_type == "json":
-            with file_path.open("r") as file:
-                json_content: dict[str, str] = json.load(file)
-                return json_content
-        else:
-            with file_path.open("r") as file:
-                return file.readlines()
-    except FileNotFoundError:
-        return f"LOG FILE NOT FOUND: {file_path}"
-
-
-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, "json", log_dir)
-    if isinstance(log_content, str):
-        return log_content
-
-    server_addr = cast(dict[str, str], log_content).get("server_address")
-    return server_addr if server_addr else "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_config() -> list[ModelConfig]:
-    """Load the model configuration."""
-    default_path = (
-        CACHED_CONFIG
-        if CACHED_CONFIG.exists()
-        else Path(__file__).resolve().parent.parent / "config" / "models.yaml"
-    )
-
-    config: dict[str, Any] = {}
-    with open(default_path) as f:
-        config = yaml.safe_load(f) or {}
-
-    user_path = os.getenv("VEC_INF_CONFIG")
-    if user_path:
-        user_path_obj = Path(user_path)
-        if user_path_obj.exists():
-            with open(user_path_obj) as f:
-                user_config = yaml.safe_load(f) or {}
-                for name, data in user_config.get("models", {}).items():
-                    if name in config.get("models", {}):
-                        config["models"][name].update(data)
-                    else:
-                        config.setdefault("models", {})[name] = data
-        else:
-            print(
-                f"WARNING: Could not find user config: {user_path}, revert to default config located at {default_path}"
-            )
-
-    return [
-        ModelConfig(model_name=name, **model_data)
-        for name, model_data in config.get("models", {}).items()
-    ]

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

vec_inf/client/_client_vars.py

@@ -0,0 +1,231 @@
+"""Global variables for Vector Inference.
+
+This module contains configuration constants and templates used throughout the
+Vector Inference package, including SLURM script templates, model configurations,
+and metric definitions.
+
+Constants
+---------
+MODEL_READY_SIGNATURE : str
+    Signature string indicating successful model server startup
+SRC_DIR : str
+    Absolute path to the package source directory
+REQUIRED_FIELDS : set
+    Set of required fields for model configuration
+KEY_METRICS : dict
+    Mapping of vLLM metrics to their human-readable names
+SLURM_JOB_CONFIG_ARGS : dict
+    Mapping of SLURM configuration arguments to their parameter names
+"""
+
+from pathlib import Path
+from typing import TypedDict
+
+from vec_inf.client.slurm_vars import (
+    LD_LIBRARY_PATH,
+    SINGULARITY_IMAGE,
+    SINGULARITY_LOAD_CMD,
+    VLLM_NCCL_SO_PATH,
+)
+
+
+MODEL_READY_SIGNATURE = "INFO:     Application startup complete."
+SRC_DIR = str(Path(__file__).parent.parent)
+
+
+# Required fields for model configuration
+REQUIRED_FIELDS = {
+    "model_family",
+    "model_type",
+    "gpus_per_node",
+    "num_nodes",
+    "vocab_size",
+}
+
+# Key production metrics for inference servers
+KEY_METRICS = {
+    "vllm:prompt_tokens_total": "total_prompt_tokens",
+    "vllm:generation_tokens_total": "total_generation_tokens",
+    "vllm:e2e_request_latency_seconds_sum": "request_latency_sum",
+    "vllm:e2e_request_latency_seconds_count": "request_latency_count",
+    "vllm:request_queue_time_seconds_sum": "queue_time_sum",
+    "vllm:request_success_total": "successful_requests_total",
+    "vllm:num_requests_running": "requests_running",
+    "vllm:num_requests_waiting": "requests_waiting",
+    "vllm:num_requests_swapped": "requests_swapped",
+    "vllm:gpu_cache_usage_perc": "gpu_cache_usage",
+    "vllm:cpu_cache_usage_perc": "cpu_cache_usage",
+}
+
+# Slurm job configuration arguments
+SLURM_JOB_CONFIG_ARGS = {
+    "job-name": "model_name",
+    "partition": "partition",
+    "account": "account",
+    "qos": "qos",
+    "time": "time",
+    "nodes": "num_nodes",
+    "exclude": "exclude",
+    "nodelist": "node_list",
+    "gpus-per-node": "gpus_per_node",
+    "cpus-per-task": "cpus_per_task",
+    "mem": "mem_per_node",
+    "output": "out_file",
+    "error": "err_file",
+}
+
+# vLLM engine args mapping between short and long names
+VLLM_SHORT_TO_LONG_MAP = {
+    "-tp": "--tensor-parallel-size",
+    "-pp": "--pipeline-parallel-size",
+    "-dp": "--data-parallel-size",
+    "-dpl": "--data-parallel-size-local",
+    "-dpa": "--data-parallel-address",
+    "-dpp": "--data-parallel-rpc-port",
+    "-O": "--compilation-config",
+    "-q": "--quantization",
+}
+
+
+# Slurm script templates
+class ShebangConfig(TypedDict):
+    """TypedDict for SLURM script shebang configuration.
+
+    Parameters
+    ----------
+    base : str
+        Base shebang line for all SLURM scripts
+    multinode : list[str]
+        Additional SLURM directives for multi-node configurations
+    """
+
+    base: str
+    multinode: list[str]
+
+
+class ServerSetupConfig(TypedDict):
+    """TypedDict for server setup configuration.
+
+    Parameters
+    ----------
+    single_node : list[str]
+        Setup commands for single-node deployments
+    multinode : list[str]
+        Setup commands for multi-node deployments, including Ray initialization
+    """
+
+    single_node: list[str]
+    multinode: list[str]
+
+
+class SlurmScriptTemplate(TypedDict):
+    """TypedDict for complete SLURM script template configuration.
+
+    Parameters
+    ----------
+    shebang : ShebangConfig
+        Shebang and SLURM directive configuration
+    singularity_setup : list[str]
+        Commands for Singularity container setup
+    imports : str
+        Import statements and source commands
+    env_vars : list[str]
+        Environment variables to set
+    singularity_command : str
+        Template for Singularity execution command
+    activate_venv : str
+        Template for virtual environment activation
+    server_setup : ServerSetupConfig
+        Server initialization commands for different deployment modes
+    find_vllm_port : list[str]
+        Commands to find available ports for vLLM server
+    write_to_json : list[str]
+        Commands to write server configuration to JSON
+    launch_cmd : list[str]
+        vLLM server launch commands
+    """
+
+    shebang: ShebangConfig
+    singularity_setup: list[str]
+    imports: str
+    env_vars: list[str]
+    singularity_command: str
+    activate_venv: str
+    server_setup: ServerSetupConfig
+    find_vllm_port: list[str]
+    write_to_json: list[str]
+    launch_cmd: list[str]
+
+
+SLURM_SCRIPT_TEMPLATE: SlurmScriptTemplate = {
+    "shebang": {
+        "base": "#!/bin/bash",
+        "multinode": [
+            "#SBATCH --exclusive",
+            "#SBATCH --tasks-per-node=1",
+        ],
+    },
+    "singularity_setup": [
+        SINGULARITY_LOAD_CMD,
+        f"singularity exec {SINGULARITY_IMAGE} ray stop",
+    ],
+    "imports": "source {src_dir}/find_port.sh",
+    "env_vars": [
+        f"export LD_LIBRARY_PATH={LD_LIBRARY_PATH}",
+        f"export VLLM_NCCL_SO_PATH={VLLM_NCCL_SO_PATH}",
+    ],
+    "singularity_command": f"singularity exec --nv --bind {{model_weights_path}}{{additional_binds}} --containall {SINGULARITY_IMAGE}",
+    "activate_venv": "source {venv}/bin/activate",
+    "server_setup": {
+        "single_node": [
+            "\n# Find available port",
+            "head_node_ip=${SLURMD_NODENAME}",
+        ],
+        "multinode": [
+            "\n# Get list of nodes",
+            'nodes=$(scontrol show hostnames "$SLURM_JOB_NODELIST")',
+            "nodes_array=($nodes)",
+            "head_node=${nodes_array[0]}",
+            'head_node_ip=$(srun --nodes=1 --ntasks=1 -w "$head_node" hostname --ip-address)',
+            "\n# Start Ray head node",
+            "head_node_port=$(find_available_port $head_node_ip 8080 65535)",
+            "ray_head=$head_node_ip:$head_node_port",
+            'echo "Ray Head IP: $ray_head"',
+            'echo "Starting HEAD at $head_node"',
+            'srun --nodes=1 --ntasks=1 -w "$head_node" \\',
+            "    SINGULARITY_PLACEHOLDER \\",
+            '    ray start --head --node-ip-address="$head_node_ip" --port=$head_node_port \\',
+            '    --num-cpus "$SLURM_CPUS_PER_TASK" --num-gpus "$SLURM_GPUS_PER_NODE" --block &',
+            "sleep 10",
+            "\n# Start Ray worker nodes",
+            "worker_num=$((SLURM_JOB_NUM_NODES - 1))",
+            "for ((i = 1; i <= worker_num; i++)); do",
+            "    node_i=${nodes_array[$i]}",
+            '    echo "Starting WORKER $i at $node_i"',
+            '    srun --nodes=1 --ntasks=1 -w "$node_i" \\',
+            "        SINGULARITY_PLACEHOLDER \\",
+            '        ray start --address "$ray_head" \\',
+            '        --num-cpus "$SLURM_CPUS_PER_TASK" --num-gpus "$SLURM_GPUS_PER_NODE" --block &',
+            "    sleep 5",
+            "done",
+        ],
+    },
+    "find_vllm_port": [
+        "\nvllm_port_number=$(find_available_port $head_node_ip 8080 65535)",
+        'server_address="http://${head_node_ip}:${vllm_port_number}/v1"',
+    ],
+    "write_to_json": [
+        '\njson_path="{log_dir}/{model_name}.$SLURM_JOB_ID/{model_name}.$SLURM_JOB_ID.json"',
+        'jq --arg server_addr "$server_address" \\',
+        "    '. + {{\"server_address\": $server_addr}}' \\",
+        '    "$json_path" > temp.json \\',
+        '    && mv temp.json "$json_path"',
+    ],
+    "launch_cmd": [
+        "vllm serve {model_weights_path} \\",
+        "    --served-model-name {model_name} \\",
+        '    --host "0.0.0.0" \\',
+        "    --port $vllm_port_number \\",
+        "    --trust-remote-code \\",
+    ],
+}

vec_inf/client/_exceptions.py

@@ -0,0 +1,37 @@
+"""Exceptions for the vector inference package."""
+
+
+class ModelConfigurationError(Exception):
+    """Raised when the model config or weights are missing or invalid."""
+
+    pass
+
+
+class MissingRequiredFieldsError(ValueError):
+    """Raised when required fields are missing from the provided parameters."""
+
+    pass
+
+
+class ModelNotFoundError(KeyError):
+    """Raised when the specified model name is not found in the configuration."""
+
+    pass
+
+
+class SlurmJobError(RuntimeError):
+    """Raised when there's an error with a Slurm job."""
+
+    pass
+
+
+class APIError(Exception):
+    """Base exception for API errors."""
+
+    pass
+
+
+class ServerError(Exception):
+    """Exception raised when there's an error with the inference server."""
+
+    pass