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,7 +7,6 @@
 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,
@@ -129,7 +128,7 @@ 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:
@@ -194,8 +193,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,7 +105,7 @@ def output_table(self) -> Table:
 class MetricsResponseFormatter:
     """CLI Helper class for formatting MetricsResponse."""
 
-    def __init__(self, metrics: dict[str, float]):
+    def __init__(self, metrics: Union[dict[str, float], str]):
         self.metrics = metrics
         self.table = utils.create_table("Metric", "Value")
         self.enabled_prefix_caching = self._check_prefix_caching()

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[str, dict[str, float]] = None
         if not performance_metrics_collector.metrics_url.startswith("http"):
             metrics = performance_metrics_collector.metrics_url
         else: