Fix wrong var names and data access for client, removed unnecessary try excepts in api.py, update client tests accordingly

Author: XkunW

tests/vec_inf/client/test_api.py

@@ -113,10 +113,10 @@ def test_wait_until_ready():
     with patch.object(VecInfClient, "get_status") as mock_status:
         # First call returns LAUNCHING, second call returns READY
         status1 = MagicMock()
-        status1.status = ModelStatus.LAUNCHING
+        status1.server_status = ModelStatus.LAUNCHING
 
         status2 = MagicMock()
-        status2.status = ModelStatus.READY
+        status2.server_status = ModelStatus.READY
         status2.base_url = "http://gpu123:8080/v1"
 
         mock_status.side_effect = [status1, status2]
@@ -125,6 +125,6 @@ def test_wait_until_ready():
             client = VecInfClient()
             result = client.wait_until_ready("12345678", timeout_seconds=5)
 
-            assert result.status == ModelStatus.READY
+            assert result.server_status == ModelStatus.READY
             assert result.base_url == "http://gpu123:8080/v1"
             assert mock_status.call_count == 2

tests/vec_inf/client/test_models.py

@@ -10,14 +10,14 @@ def test_model_info_creation():
         family="test-family",
         variant="test-variant",
         type=ModelType.LLM,
-        config={"num_gpus": 1},
+        config={"gpus_per_node": 1},
     )
 
     assert model.name == "test-model"
     assert model.family == "test-family"
     assert model.variant == "test-variant"
     assert model.type == ModelType.LLM
-    assert model.config["num_gpus"] == 1
+    assert model.config["gpus_per_node"] == 1
 
 
 def test_model_info_optional_fields():
@@ -40,7 +40,7 @@ def test_launch_options_default_values():
     """Test LaunchOptions with default values."""
     options = LaunchOptions()
 
-    assert options.num_gpus is None
+    assert options.gpus_per_node is None
     assert options.partition is None
     assert options.data_type is None
     assert options.num_nodes is None

vec_inf/client/_helper.py

@@ -74,9 +74,7 @@ def _get_model_configuration(self) -> ModelConfig:
         )
 
         if not model_weights_parent_dir:
-            raise ValueError(
-                f"Could not determine model_weights_parent_dir and '{self.model_name}' not found in configuration"
-            )
+            raise ModelNotFoundError("Could not determine model weights parent directory")
 
         model_weights_path = Path(model_weights_parent_dir, self.model_name)
 
@@ -266,53 +264,47 @@ def _get_base_status_data(self) -> StatusResponse:
     def _check_model_health(self) -> None:
         """Check model health and update status accordingly."""
         status, status_code = utils.model_health_check(
-            cast(str, self.status_info["model_name"]), self.slurm_job_id, self.log_dir
+            self.status_info.model_name, self.slurm_job_id, self.log_dir
         )
         if status == ModelStatus.READY:
-            self.status_info["base_url"] = utils.get_base_url(
-                cast(str, self.status_info["model_name"]),
+            self.status_info.base_url = utils.get_base_url(
+                self.status_info.model_name,
                 self.slurm_job_id,
                 self.log_dir,
             )
-            self.status_info["server_status"] = status
+            self.status_info.server_status = status
         else:
-            self.status_info["server_status"], self.status_info["failed_reason"] = (
-                status,
-                cast(str, status_code),
-            )
+            self.status_info.server_status = status
+            self.status_info.failed_reason = cast(str, status_code)
 
     def _process_running_state(self) -> None:
         """Process RUNNING job state and check server status."""
         server_status = utils.is_server_running(
-            cast(str, self.status_info["model_name"]), self.slurm_job_id, self.log_dir
+            self.status_info.model_name, self.slurm_job_id, self.log_dir
         )
 
         if isinstance(server_status, tuple):
-            self.status_info["server_status"], self.status_info["failed_reason"] = (
-                server_status
-            )
+            self.status_info.server_status, self.status_info.failed_reason = server_status
             return
 
         if server_status == "RUNNING":
             self._check_model_health()
         else:
-            self.status_info["server_status"] = server_status
+            self.status_info.server_status = server_status
 
     def _process_pending_state(self) -> None:
         """Process PENDING job state."""
         try:
-            self.status_info["pending_reason"] = self.output.split(" ")[10].split("=")[
-                1
-            ]
-            self.status_info["server_status"] = ModelStatus.PENDING
+            self.status_info.pending_reason = self.output.split(" ")[10].split("=")[1]
+            self.status_info.server_status = ModelStatus.PENDING
         except IndexError:
-            self.status_info["pending_reason"] = "Unknown pending reason"
+            self.status_info.pending_reason = "Unknown pending reason"
 
     def process_model_status(self) -> StatusResponse:
         """Process different job states and update status information."""
-        if self.status_info["job_state"] == ModelStatus.PENDING:
+        if self.status_info.job_state == ModelStatus.PENDING:
             self._process_pending_state()
-        elif self.status_info["job_state"] == "RUNNING":
+        elif self.status_info.job_state == "RUNNING":
             self._process_running_state()
 
         return self.status_info
@@ -360,7 +352,7 @@ def _build_metrics_url(self) -> str:
     def _check_prefix_caching(self) -> bool:
         """Check if prefix caching is enabled."""
         job_json = utils.read_slurm_log(
-            cast(str, self.status_info["model_name"]),
+            self.status_info.model_name,
             self.slurm_job_id,
             "json",
             self.log_dir,
@@ -369,6 +361,43 @@ def _check_prefix_caching(self) -> bool:
             return False
         return bool(cast(dict[str, str], job_json).get("enable_prefix_caching", False))
 
+    def _parse_metrics(self, metrics_text: str) -> dict[str, float]:
+        """Parse metrics with latency count and sum."""
+        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",
+        }
+
+        if self.enabled_prefix_caching:
+            key_metrics["vllm:gpu_prefix_cache_hit_rate"] = "gpu_prefix_cache_hit_rate"
+            key_metrics["vllm:cpu_prefix_cache_hit_rate"] = "cpu_prefix_cache_hit_rate"
+
+        parsed: dict[str, float] = {}
+        for line in metrics_text.split("\n"):
+            if line.startswith("#") or not line.strip():
+                continue
+
+            parts = line.split()
+            if len(parts) < 2:
+                continue
+
+            metric_name = parts[0].split("{")[0]
+            if metric_name in key_metrics:
+                try:
+                    parsed[key_metrics[metric_name]] = float(parts[1])
+                except (ValueError, IndexError):
+                    continue
+        return parsed
+
     def fetch_metrics(self) -> Union[dict[str, float], str]:
         """Fetch metrics from the endpoint."""
         try:
@@ -443,43 +472,6 @@ def fetch_metrics(self) -> Union[dict[str, float], str]:
         except requests.RequestException as e:
             return f"Metrics request failed, `metrics` endpoint might not be ready yet: {str(e)}"
 
-    def _parse_metrics(self, metrics_text: str) -> dict[str, float]:
-        """Parse metrics with latency count and sum."""
-        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",
-        }
-
-        if self.enabled_prefix_caching:
-            key_metrics["vllm:gpu_prefix_cache_hit_rate"] = "gpu_prefix_cache_hit_rate"
-            key_metrics["vllm:cpu_prefix_cache_hit_rate"] = "cpu_prefix_cache_hit_rate"
-
-        parsed: dict[str, float] = {}
-        for line in metrics_text.split("\n"):
-            if line.startswith("#") or not line.strip():
-                continue
-
-            parts = line.split()
-            if len(parts) < 2:
-                continue
-
-            metric_name = parts[0].split("{")[0]
-            if metric_name in key_metrics:
-                try:
-                    parsed[key_metrics[metric_name]] = float(parts[1])
-                except (ValueError, IndexError):
-                    continue
-        return parsed
-
 
 class ModelRegistry:
     """Class for handling model listing and configuration management."""

vec_inf/client/_models.py

@@ -78,7 +78,7 @@ class LaunchOptions:
     max_num_batched_tokens: Optional[int] = None
     partition: Optional[str] = None
     num_nodes: Optional[int] = None
-    num_gpus: Optional[int] = None
+    gpus_per_node: Optional[int] = None
     qos: Optional[str] = None
     time: Optional[str] = None
     vocab_size: Optional[int] = None
@@ -104,7 +104,7 @@ class LaunchOptionsDict(TypedDict):
     max_num_batched_tokens: NotRequired[Optional[int]]
     partition: NotRequired[Optional[str]]
     num_nodes: NotRequired[Optional[int]]
-    num_gpus: NotRequired[Optional[int]]
+    gpus_per_node: NotRequired[Optional[int]]
     qos: NotRequired[Optional[str]]
     time: NotRequired[Optional[str]]
     vocab_size: NotRequired[Optional[int]]

vec_inf/client/_utils.py

@@ -3,6 +3,7 @@
 import json
 import os
 import subprocess
+import warnings
 from pathlib import Path
 from typing import Any, Optional, Union, cast
 
@@ -151,8 +152,9 @@ def load_config() -> list[ModelConfig]:
                     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}"
+            warnings.warn(
+                f"WARNING: Could not find user config: {user_path}, revert to default config located at {default_path}", UserWarning,
+                stacklevel=2
             )
 
     return [
@@ -161,12 +163,6 @@ def load_config() -> list[ModelConfig]:
     ]
 
 
-def shutdown_model(slurm_job_id: int) -> None:
-    """Shutdown a running model on the cluster."""
-    shutdown_cmd = f"scancel {slurm_job_id}"
-    run_bash_command(shutdown_cmd)
-
-
 def parse_launch_output(output: str) -> tuple[str, dict[str, str]]:
     """Parse output from model launch command.