mypy fixes

Author: XkunW

examples/api/basic_usage.py

@@ -33,7 +33,7 @@
 # Get metrics
 print("\nRetrieving metrics...")
 metrics = client.get_metrics(job_id)
-if metrics.metrics:
+if isinstance(metrics.metrics, dict):
     for key, value in metrics.metrics.items():
         print(f"- {key}: {value}")
 

vec_inf/cli/_cli.py

@@ -7,14 +7,13 @@
 from rich.console import Console
 from rich.live import Live
 
-import vec_inf.client._utils as utils
 from vec_inf.cli._helper import (
     LaunchResponseFormatter,
     ListCmdDisplay,
     MetricsResponseFormatter,
     StatusResponseFormatter,
 )
-from vec_inf.client._models import LaunchOptions
+from vec_inf.client._models import LaunchOptions, LaunchOptionsDict
 from vec_inf.client.api import VecInfClient
 
 
@@ -129,14 +128,15 @@ def cli() -> None:
 )
 def launch(
     model_name: str,
-    **cli_kwargs: Optional[Union[str, int, bool]],
+    **cli_kwargs: Optional[Union[str, int, float, bool]],
 ) -> None:
     """Launch a model on the cluster."""
     try:
         # Convert cli_kwargs to LaunchOptions
-        launch_options = LaunchOptions(
-            **{k: v for k, v in cli_kwargs.items() if k != "json_mode"}
-        )
+        kwargs = {k: v for k, v in cli_kwargs.items() if k != "json_mode"}
+        # Cast the dictionary to LaunchOptionsDict
+        options_dict: LaunchOptionsDict = kwargs  # type: ignore
+        launch_options = LaunchOptions(**options_dict)
 
         # Start the client and launch model inference server
         client = VecInfClient()
@@ -194,8 +194,12 @@ def status(
 @click.argument("slurm_job_id", type=int, nargs=1)
 def shutdown(slurm_job_id: int) -> None:
     """Shutdown a running model on the cluster."""
-    utils.shutdown_model(slurm_job_id)
-    click.echo(f"Shutting down model with Slurm Job ID: {slurm_job_id}")
+    try:
+        client = VecInfClient()
+        client.shutdown_model(slurm_job_id)
+        click.echo(f"Shutting down model with Slurm Job ID: {slurm_job_id}")
+    except Exception as e:
+        raise click.ClickException(f"Shutdown failed: {str(e)}") from e
 
 
 @cli.command("list")

vec_inf/cli/_helper.py

@@ -105,19 +105,18 @@ def output_table(self) -> Table:
 class MetricsResponseFormatter:
     """CLI Helper class for formatting MetricsResponse."""
 
-    def __init__(self, metrics: dict[str, float]):
-        self.metrics = metrics
+    def __init__(self, metrics: Union[dict[str, float], str]):
+        self.metrics = self._set_metrics(metrics)
         self.table = utils.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."""
+        return metrics if isinstance(metrics, dict) else {}
+
     def _check_prefix_caching(self) -> bool:
         """Check if prefix caching is enabled by looking for prefix cache metrics."""
-        if isinstance(self.metrics, str):
-            # If metrics is a string, it's an error message
-            return False
-
-        cache_rate = self.metrics.get("gpu_prefix_cache_hit_rate")
-        return cache_rate is not None
+        return self.metrics.get("gpu_prefix_cache_hit_rate") is not None
 
     def format_failed_metrics(self, message: str) -> None:
         self.table.add_row("ERROR", message)

vec_inf/client/_helper.py

@@ -294,7 +294,7 @@ def _process_running_state(self) -> None:
         if server_status == "RUNNING":
             self._check_model_health()
         else:
-            self.status_info.server_status = server_status
+            self.status_info.server_status = cast(ModelStatus, server_status)
 
     def _process_pending_state(self) -> None:
         """Process PENDING job state."""

vec_inf/client/api.py

@@ -5,7 +5,7 @@
 """
 
 import time
-from typing import Any, Optional
+from typing import Any, Optional, Union
 
 from vec_inf.client._config import ModelConfig
 from vec_inf.client._exceptions import (
@@ -147,6 +147,7 @@ def get_metrics(
             slurm_job_id, log_dir
         )
 
+        metrics: Union[dict[str, float], str]
         if not performance_metrics_collector.metrics_url.startswith("http"):
             metrics = performance_metrics_collector.metrics_url
         else: