vec-inf 0.4.0.post1__py3-none-any.whl → 0.5.0__py3-none-any.whl

vec_inf/__init__.py

@@ -0,0 +1 @@
+"""vec_inf package."""

vec_inf/cli/__init__.py

@@ -0,0 +1 @@
+"""vec_inf cli package."""

vec_inf/cli/_cli.py

@@ -1,23 +1,22 @@
-import os
+"""Command line interface for Vector Inference."""
+
 import time
-from typing import Optional, cast
+from typing import Optional, Union
 
 import click
-
-import polars as pl
-from rich.columns import Columns
 from rich.console import Console
 from rich.live import Live
-from rich.panel import Panel
 
 import vec_inf.cli._utils as utils
+from vec_inf.cli._helper import LaunchHelper, ListHelper, MetricsHelper, StatusHelper
+
 
 CONSOLE = Console()
 
 
 @click.group()
-def cli():
-    """Vector Inference CLI"""
+def cli() -> None:
+    """Vector Inference CLI."""
     pass
 
 
@@ -35,11 +34,30 @@ def cli():
     type=int,
     help="Maximum number of sequences to process in a single request",
 )
+@click.option(
+    "--gpu-memory-utilization",
+    type=float,
+    help="GPU memory utilization, default to 0.9",
+)
+@click.option(
+    "--enable-prefix-caching",
+    is_flag=True,
+    help="Enables automatic prefix caching",
+)
+@click.option(
+    "--enable-chunked-prefill",
+    is_flag=True,
+    help="Enable chunked prefill, enabled by default if max number of sequences > 32k",
+)
+@click.option(
+    "--max-num-batched-tokens",
+    type=int,
+    help="Maximum number of batched tokens per iteration, defaults to 2048 if --enable-chunked-prefill is set, else None",
+)
 @click.option(
     "--partition",
     type=str,
-    default="a40",
-    help="Type of compute partition, default to a40",
+    help="Type of compute partition",
 )
 @click.option(
     "--num-nodes",
@@ -47,7 +65,7 @@ def cli():
     help="Number of nodes to use, default to suggested resource allocation for model",
 )
 @click.option(
-    "--num-gpus",
+    "--gpus-per-node",
     type=int,
     help="Number of GPUs/node to use, default to suggested resource allocation for model",
 )
@@ -66,36 +84,36 @@ def cli():
     type=int,
     help="Vocabulary size, this option is intended for custom models",
 )
-@click.option(
-    "--data-type", type=str, default="auto", help="Model data type, default to auto"
-)
+@click.option("--data-type", type=str, help="Model data type")
 @click.option(
     "--venv",
     type=str,
-    default="singularity",
-    help="Path to virtual environment, default to preconfigured singularity container",
+    help="Path to virtual environment",
 )
 @click.option(
     "--log-dir",
     type=str,
-    default="default",
-    help="Path to slurm log directory, default to .vec-inf-logs in user home directory",
+    help="Path to slurm log directory",
 )
 @click.option(
     "--model-weights-parent-dir",
     type=str,
-    default="/model-weights",
-    help="Path to parent directory containing model weights, default to '/model-weights' for supported models",
+    help="Path to parent directory containing model weights",
 )
 @click.option(
     "--pipeline-parallelism",
-    type=str,
-    help="Enable pipeline parallelism, accepts 'True' or 'False', default to 'True' for supported models",
+    is_flag=True,
+    help="Enable pipeline parallelism, enabled by default for supported models",
+)
+@click.option(
+    "--compilation-config",
+    type=click.Choice(["0", "3"]),
+    help="torch.compile optimization level, accepts '0' or '3', default to '0', which means no optimization is applied",
 )
 @click.option(
     "--enforce-eager",
-    type=str,
-    help="Always use eager-mode PyTorch, accepts 'True' or 'False', default to 'False' for custom models if not set",
+    is_flag=True,
+    help="Always use eager-mode PyTorch",
 )
 @click.option(
     "--json-mode",
@@ -104,74 +122,23 @@ def cli():
 )
 def launch(
     model_name: str,
-    model_family: Optional[str] = None,
-    model_variant: Optional[str] = None,
-    max_model_len: Optional[int] = None,
-    max_num_seqs: Optional[int] = None,
-    partition: Optional[str] = None,
-    num_nodes: Optional[int] = None,
-    num_gpus: Optional[int] = None,
-    qos: Optional[str] = None,
-    time: Optional[str] = None,
-    vocab_size: Optional[int] = None,
-    data_type: Optional[str] = None,
-    venv: Optional[str] = None,
-    log_dir: Optional[str] = None,
-    model_weights_parent_dir: Optional[str] = None,
-    pipeline_parallelism: Optional[str] = None,
-    enforce_eager: Optional[str] = None,
-    json_mode: bool = False,
+    **cli_kwargs: Optional[Union[str, int, bool]],
 ) -> None:
-    """
-    Launch a model on the cluster
-    """
-
-    if isinstance(pipeline_parallelism, str):
-        pipeline_parallelism = (
-            "True" if pipeline_parallelism.lower() == "true" else "False"
-        )
-
-    launch_script_path = os.path.join(
-        os.path.dirname(os.path.dirname(os.path.realpath(__file__))), "launch_server.sh"
-    )
-    launch_cmd = f"bash {launch_script_path}"
-
-    models_df = utils.load_models_df()
-
-    if model_name in models_df["model_name"].to_list():
-        default_args = utils.load_default_args(models_df, model_name)
-        for arg in default_args:
-            if arg in locals() and locals()[arg] is not None:
-                default_args[arg] = locals()[arg]
-            renamed_arg = arg.replace("_", "-")
-            launch_cmd += f" --{renamed_arg} {default_args[arg]}"
-    else:
-        model_args = models_df.columns
-        model_args.remove("model_name")
-        model_args.remove("model_type")
-        for arg in model_args:
-            if locals()[arg] is not None:
-                renamed_arg = arg.replace("_", "-")
-                launch_cmd += f" --{renamed_arg} {locals()[arg]}"
-
-    output = utils.run_bash_command(launch_cmd)
-
-    slurm_job_id = output.split(" ")[-1].strip().strip("\n")
-    output_lines = output.split("\n")[:-2]
-
-    table = utils.create_table(key_title="Job Config", value_title="Value")
-    table.add_row("Slurm Job ID", slurm_job_id, style="blue")
-    output_dict = {"slurm_job_id": slurm_job_id}
+    """Launch a model on the cluster."""
+    try:
+        launch_helper = LaunchHelper(model_name, cli_kwargs)
 
-    for line in output_lines:
-        key, value = line.split(": ")
-        table.add_row(key, value)
-        output_dict[key.lower().replace(" ", "_")] = value
+        launch_helper.set_env_vars()
+        launch_command = launch_helper.build_launch_command()
+        command_output, stderr = utils.run_bash_command(launch_command)
+        if stderr:
+            raise click.ClickException(f"Error: {stderr}")
+        launch_helper.post_launch_processing(command_output, CONSOLE)
 
-    if json_mode:
-        click.echo(output_dict)
-    else:
-        CONSOLE.print(table)
+    except click.ClickException as e:
+        raise e
+    except Exception as e:
+        raise click.ClickException(f"Launch failed: {str(e)}") from e
 
 
 @cli.command("status")
@@ -189,79 +156,25 @@ def launch(
 def status(
     slurm_job_id: int, log_dir: Optional[str] = None, json_mode: bool = False
 ) -> None:
-    """
-    Get the status of a running model on the cluster
-    """
+    """Get the status of a running model on the cluster."""
     status_cmd = f"scontrol show job {slurm_job_id} --oneliner"
-    output = utils.run_bash_command(status_cmd)
+    output, stderr = utils.run_bash_command(status_cmd)
+    if stderr:
+        raise click.ClickException(f"Error: {stderr}")
 
-    slurm_job_name = "UNAVAILABLE"
-    status = "SHUTDOWN"
-    base_url = "UNAVAILABLE"
-
-    try:
-        slurm_job_name = output.split(" ")[1].split("=")[1]
-        slurm_job_state = output.split(" ")[9].split("=")[1]
-    except IndexError:
-        # Job ID not found
-        slurm_job_state = "UNAVAILABLE"
-
-    # If Slurm job is currently PENDING
-    if slurm_job_state == "PENDING":
-        slurm_job_pending_reason = output.split(" ")[10].split("=")[1]
-        status = "PENDING"
-    # If Slurm job is currently RUNNING
-    elif slurm_job_state == "RUNNING":
-        # Check whether the server is ready, if yes, run model health check to further determine status
-        server_status = utils.is_server_running(slurm_job_name, slurm_job_id, log_dir)
-        # If server status is a tuple, then server status is "FAILED"
-        if isinstance(server_status, tuple):
-            status = server_status[0]
-            slurm_job_failed_reason = server_status[1]
-        elif server_status == "RUNNING":
-            model_status = utils.model_health_check(
-                slurm_job_name, slurm_job_id, log_dir
-            )
-            if model_status == "READY":
-                # Only set base_url if model is ready to serve requests
-                base_url = utils.get_base_url(slurm_job_name, slurm_job_id, log_dir)
-                status = "READY"
-            else:
-                # If model is not ready, then status must be "FAILED"
-                status = model_status[0]
-                slurm_job_failed_reason = str(model_status[1])
-        else:
-            status = server_status
+    status_helper = StatusHelper(slurm_job_id, output, log_dir)
 
+    status_helper.process_job_state()
     if json_mode:
-        status_dict = {
-            "model_name": slurm_job_name,
-            "model_status": status,
-            "base_url": base_url,
-        }
-        if "slurm_job_pending_reason" in locals():
-            status_dict["pending_reason"] = slurm_job_pending_reason
-        if "slurm_job_failed_reason" in locals():
-            status_dict["failed_reason"] = slurm_job_failed_reason
-        click.echo(f"{status_dict}")
+        status_helper.output_json()
     else:
-        table = utils.create_table(key_title="Job Status", value_title="Value")
-        table.add_row("Model Name", slurm_job_name)
-        table.add_row("Model Status", status, style="blue")
-        if "slurm_job_pending_reason" in locals():
-            table.add_row("Reason", slurm_job_pending_reason)
-        if "slurm_job_failed_reason" in locals():
-            table.add_row("Reason", slurm_job_failed_reason)
-        table.add_row("Base URL", base_url)
-        CONSOLE.print(table)
+        status_helper.output_table(CONSOLE)
 
 
 @cli.command("shutdown")
 @click.argument("slurm_job_id", type=int, nargs=1)
 def shutdown(slurm_job_id: int) -> None:
-    """
-    Shutdown a running model on the cluster
-    """
+    """Shutdown a running model on the cluster."""
     shutdown_cmd = f"scancel {slurm_job_id}"
     utils.run_bash_command(shutdown_cmd)
     click.echo(f"Shutting down model with Slurm Job ID: {slurm_job_id}")
@@ -275,109 +188,41 @@ def shutdown(slurm_job_id: int) -> None:
     help="Output in JSON string",
 )
 def list_models(model_name: Optional[str] = None, json_mode: bool = False) -> None:
-    """
-    List all available models, or get default setup of a specific model
-    """
-
-    def list_model(model_name: str, models_df: pl.DataFrame, json_mode: bool):
-        if model_name not in models_df["model_name"].to_list():
-            raise ValueError(f"Model name {model_name} not found in available models")
-
-        excluded_keys = {"venv", "log_dir"}
-        model_row = models_df.filter(models_df["model_name"] == model_name)
-
-        if json_mode:
-            filtered_model_row = model_row.drop(excluded_keys, strict=False)
-            click.echo(filtered_model_row.to_dicts()[0])
-            return
-        table = utils.create_table(key_title="Model Config", value_title="Value")
-        for row in model_row.to_dicts():
-            for key, value in row.items():
-                if key not in excluded_keys:
-                    table.add_row(key, str(value))
-        CONSOLE.print(table)
-
-    def list_all(models_df: pl.DataFrame, json_mode: bool):
-        if json_mode:
-            click.echo(models_df["model_name"].to_list())
-            return
-        panels = []
-        model_type_colors = {
-            "LLM": "cyan",
-            "VLM": "bright_blue",
-            "Text Embedding": "purple",
-            "Reward Modeling": "bright_magenta",
-        }
-
-        models_df = models_df.with_columns(
-            pl.when(pl.col("model_type") == "LLM")
-            .then(0)
-            .when(pl.col("model_type") == "VLM")
-            .then(1)
-            .when(pl.col("model_type") == "Text Embedding")
-            .then(2)
-            .when(pl.col("model_type") == "Reward Modeling")
-            .then(3)
-            .otherwise(-1)
-            .alias("model_type_order")
-        )
-
-        models_df = models_df.sort("model_type_order")
-        models_df = models_df.drop("model_type_order")
-
-        for row in models_df.to_dicts():
-            panel_color = model_type_colors.get(row["model_type"], "white")
-            styled_text = (
-                f"[magenta]{row['model_family']}[/magenta]-{row['model_variant']}"
-            )
-            panels.append(Panel(styled_text, expand=True, border_style=panel_color))
-        CONSOLE.print(Columns(panels, equal=True))
-
-    models_df = utils.load_models_df()
-
-    if model_name:
-        list_model(model_name, models_df, json_mode)
-    else:
-        list_all(models_df, json_mode)
+    """List all available models, or get default setup of a specific model."""
+    list_helper = ListHelper(model_name, json_mode)
+    list_helper.process_list_command(CONSOLE)
 
 
 @cli.command("metrics")
 @click.argument("slurm_job_id", type=int, nargs=1)
 @click.option(
-    "--log-dir",
-    type=str,
-    help="Path to slurm log directory. This is required if --log-dir was set in model launch",
+    "--log-dir", type=str, help="Path to slurm log directory (if used during launch)"
 )
 def metrics(slurm_job_id: int, log_dir: Optional[str] = None) -> None:
-    """
-    Stream performance metrics to the console
-    """
-    status_cmd = f"scontrol show job {slurm_job_id} --oneliner"
-    output = utils.run_bash_command(status_cmd)
-    slurm_job_name = output.split(" ")[1].split("=")[1]
+    """Stream real-time performance metrics from the model endpoint."""
+    helper = MetricsHelper(slurm_job_id, log_dir)
+
+    # Check if metrics URL is ready
+    if not helper.metrics_url.startswith("http"):
+        table = utils.create_table("Metric", "Value")
+        helper.display_failed_metrics(
+            table, f"Metrics endpoint unavailable - {helper.metrics_url}"
+        )
+        CONSOLE.print(table)
+        return
 
     with Live(refresh_per_second=1, console=CONSOLE) as live:
         while True:
-            out_logs = utils.read_slurm_log(
-                slurm_job_name, slurm_job_id, "out", log_dir
-            )
-            # if out_logs is a string, then it is an error message
-            if isinstance(out_logs, str):
-                live.update(out_logs)
-                break
-            out_logs = cast(list, out_logs)
-            latest_metrics = utils.get_latest_metric(out_logs)
-            # if latest_metrics is a string, then it is an error message
-            if isinstance(latest_metrics, str):
-                live.update(latest_metrics)
-                break
-            latest_metrics = cast(dict, latest_metrics)
-            table = utils.create_table(key_title="Metric", value_title="Value")
-            for key, value in latest_metrics.items():
-                table.add_row(key, value)
+            metrics = helper.fetch_metrics()
+            table = utils.create_table("Metric", "Value")
 
-            live.update(table)
+            if isinstance(metrics, str):
+                # Show status information if metrics aren't available
+                helper.display_failed_metrics(table, metrics)
+            else:
+                helper.display_metrics(table, metrics)
 
+            live.update(table)
             time.sleep(2)
 
 

vec_inf/cli/_config.py

@@ -0,0 +1,87 @@
+"""Model configuration."""
+
+from pathlib import Path
+from typing import Optional, Union
+
+from pydantic import BaseModel, ConfigDict, Field
+from typing_extensions import Literal
+
+
+QOS = Literal[
+    "normal",
+    "m",
+    "m2",
+    "m3",
+    "m4",
+    "m5",
+    "long",
+    "deadline",
+    "high",
+    "scavenger",
+    "llm",
+    "a100",
+]
+
+PARTITION = Literal["a40", "a100", "t4v1", "t4v2", "rtx6000"]
+
+DATA_TYPE = Literal["auto", "float16", "bfloat16", "float32"]
+
+
+class ModelConfig(BaseModel):
+    """Pydantic model for validating and managing model deployment configurations."""
+
+    model_name: str = Field(..., min_length=3, pattern=r"^[a-zA-Z0-9\-_\.]+$")
+    model_family: str = Field(..., min_length=2)
+    model_variant: Optional[str] = Field(
+        default=None, description="Specific variant/version of the model family"
+    )
+    model_type: Literal["LLM", "VLM", "Text_Embedding", "Reward_Modeling"] = Field(
+        ..., description="Type of model architecture"
+    )
+    gpus_per_node: int = Field(..., gt=0, le=8, description="GPUs per node")
+    num_nodes: int = Field(..., gt=0, le=16, description="Number of nodes")
+    vocab_size: int = Field(..., gt=0, le=1_000_000)
+    max_model_len: int = Field(
+        ..., gt=0, le=1_010_000, description="Maximum context length supported"
+    )
+    max_num_seqs: int = Field(
+        default=256, gt=0, le=1024, description="Maximum concurrent request sequences"
+    )
+    compilation_config: int = Field(
+        default=0,
+        gt=-1,
+        le=4,
+        description="torch.compile optimization level",
+    )
+    gpu_memory_utilization: float = Field(
+        default=0.9, gt=0.0, le=1.0, description="GPU memory utilization"
+    )
+    pipeline_parallelism: bool = Field(
+        default=True, description="Enable pipeline parallelism"
+    )
+    enforce_eager: bool = Field(default=False, description="Force eager mode execution")
+    qos: Union[QOS, str] = Field(default="m2", description="Quality of Service tier")
+    time: str = Field(
+        default="08:00:00",
+        pattern=r"^\d{2}:\d{2}:\d{2}$",
+        description="HH:MM:SS time limit",
+    )
+    partition: Union[PARTITION, str] = Field(
+        default="a40", description="GPU partition type"
+    )
+    data_type: Union[DATA_TYPE, str] = Field(
+        default="auto", description="Model precision format"
+    )
+    venv: str = Field(
+        default="singularity", description="Virtual environment/container system"
+    )
+    log_dir: Path = Field(
+        default=Path("~/.vec-inf-logs").expanduser(), description="Log directory path"
+    )
+    model_weights_parent_dir: Path = Field(
+        default=Path("/model-weights"), description="Base directory for model weights"
+    )
+
+    model_config = ConfigDict(
+        extra="forbid", str_strip_whitespace=True, validate_default=True, frozen=True
+    )