Refactoring client for CLI use

Author: XkunW

vec_inf/client/_helper.py

@@ -244,42 +244,26 @@ def _get_raw_status_output(self) -> str:
             raise SlurmJobError(f"Error: {stderr}")
         return output
 
-    def _get_base_status_data(self) -> dict[str, Union[str, None]]:
+    def _get_base_status_data(self) -> StatusResponse:
         """Extract basic job status information from scontrol output."""
         try:
             job_name = self.output.split(" ")[1].split("=")[1]
             job_state = self.output.split(" ")[9].split("=")[1]
         except IndexError:
-            job_name = ModelStatus.UNAVAILABLE
+            job_name = "UNAVAILABLE"
             job_state = ModelStatus.UNAVAILABLE
 
-        return {
-            "model_name": job_name,
-            "status": ModelStatus.UNAVAILABLE,
-            "base_url": ModelStatus.UNAVAILABLE,
-            "state": job_state,
-            "pending_reason": None,
-            "failed_reason": None,
-        }
-
-    def process_model_status(self) -> StatusResponse:
-        """Process different job states and update status information."""
-        if self.status_info["state"] == ModelStatus.PENDING:
-            self.process_pending_state()
-        elif self.status_info["state"] == "RUNNING":
-            self.process_running_state()
-
         return StatusResponse(
-            slurm_job_id=self.slurm_job_id,
-            model_name=cast(str, self.status_info["model_name"]),
-            status=cast(ModelStatus, self.status_info["status"]),
+            model_name=job_name,
+            server_status=ModelStatus.UNAVAILABLE,
+            job_state=job_state,
             raw_output=self.output,
-            base_url=self.status_info["base_url"],
-            pending_reason=self.status_info["pending_reason"],
-            failed_reason=self.status_info["failed_reason"],
+            base_url="UNAVAILABLE",
+            pending_reason=None,
+            failed_reason=None,
         )
 
-    def check_model_health(self) -> None:
+    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
@@ -290,40 +274,49 @@ def check_model_health(self) -> None:
                 self.slurm_job_id,
                 self.log_dir,
             )
-            self.status_info["status"] = status
+            self.status_info["server_status"] = status
         else:
-            self.status_info["status"], self.status_info["failed_reason"] = (
+            self.status_info["server_status"], self.status_info["failed_reason"] = (
                 status,
                 cast(str, status_code),
             )
 
-    def process_running_state(self) -> None:
+    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
         )
 
         if isinstance(server_status, tuple):
-            self.status_info["status"], self.status_info["failed_reason"] = (
+            self.status_info["server_status"], self.status_info["failed_reason"] = (
                 server_status
             )
             return
 
         if server_status == "RUNNING":
-            self.check_model_health()
+            self._check_model_health()
         else:
-            self.status_info["status"] = server_status
+            self.status_info["server_status"] = server_status
 
-    def process_pending_state(self) -> None:
+    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["status"] = ModelStatus.PENDING
+            self.status_info["server_status"] = ModelStatus.PENDING
         except IndexError:
             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:
+            self._process_pending_state()
+        elif self.status_info["job_state"] == "RUNNING":
+            self._process_running_state()
+
+        return self.status_info
+
 
 class PerformanceMetricsCollector:
     """Class for handling metrics collection and processing."""
@@ -340,18 +333,18 @@ def __init__(self, slurm_job_id: int, log_dir: Optional[str] = None):
         self._last_updated: Optional[float] = None
         self._last_throughputs = {"prompt": 0.0, "generation": 0.0}
 
-    def _get_status_info(self) -> dict[str, Union[str, None]]:
+    def _get_status_info(self) -> StatusResponse:
         """Retrieve status info using existing StatusHelper."""
         status_helper = ModelStatusMonitor(self.slurm_job_id, self.log_dir)
-        return status_helper.status_info
+        return status_helper.process_model_status()
 
     def _build_metrics_url(self) -> str:
         """Construct metrics endpoint URL from base URL with version stripping."""
-        if self.status_info.get("state") == "PENDING":
+        if self.status_info.job_state == ModelStatus.PENDING:
             return "Pending resources for server initialization"
 
         base_url = utils.get_base_url(
-            cast(str, self.status_info["model_name"]),
+            self.status_info.model_name,
             self.slurm_job_id,
             self.log_dir,
         )

vec_inf/client/_models.py

@@ -7,7 +7,7 @@
 
 from dataclasses import dataclass, field
 from enum import Enum
-from typing import Any, Optional, TypedDict
+from typing import Any, Optional, TypedDict, Union
 
 from typing_extensions import NotRequired
 
@@ -46,9 +46,9 @@ class LaunchResponse:
 class StatusResponse:
     """Response from checking a model's status."""
 
-    slurm_job_id: int
     model_name: str
-    status: ModelStatus
+    server_status: ModelStatus
+    job_state: Union[str, ModelStatus]
     raw_output: str = field(repr=False)
     base_url: Optional[str] = None
     pending_reason: Optional[str] = None
@@ -59,9 +59,8 @@ class StatusResponse:
 class MetricsResponse:
     """Response from retrieving model metrics."""
 
-    slurm_job_id: int
     model_name: str
-    metrics: dict[str, float]
+    metrics: Union[dict[str, float], str]
     timestamp: float
 
 

vec_inf/client/api.py

@@ -5,9 +5,7 @@
 """
 
 import time
-from typing import Any, Optional, cast
-
-import requests
+from typing import Any, Optional
 
 from vec_inf.client._config import ModelConfig
 from vec_inf.client._exceptions import (
@@ -138,7 +136,6 @@ def launch_model(
 
             # Create and use the API Launch Helper
             model_launcher = ModelLauncher(model_name, options_dict)
-
             return model_launcher.launch()
 
         except ValueError as e:
@@ -211,20 +208,12 @@ def get_metrics(
             )
 
             if not performance_metrics_collector.metrics_url.startswith("http"):
-                raise ServerError(
-                    f"Metrics endpoint unavailable or server not ready - {performance_metrics_collector.metrics_url}"
-                )
-
-            metrics = performance_metrics_collector.fetch_metrics()
-
-            if isinstance(metrics, str):
-                raise requests.RequestException(metrics)
+                metrics = performance_metrics_collector.metrics_url
+            else:
+                metrics = performance_metrics_collector.fetch_metrics()
 
             return MetricsResponse(
-                slurm_job_id=slurm_job_id,
-                model_name=cast(
-                    str, performance_metrics_collector.status_info["model_name"]
-                ),
+                model_name=performance_metrics_collector.status_info.model_name,
                 metrics=metrics,
                 timestamp=time.time(),
             )