Merge pull request #150 from VectorInstitute/bugfix/broken-throughput

Fix broken throughput computation

Author: XkunW

.github/workflows/code_checks.yml

@@ -49,3 +49,5 @@ jobs:
         uses: pypa/[email protected]
         with:
           virtual-environment: .venv/
+          # Temporary: ignore pip advisory until fixed in pip>=25.3
+          ignore-vulns: GHSA-4xh5-x5gv-qwph

vec_inf/cli/_cli.py

@@ -447,7 +447,7 @@ def metrics(slurm_job_id: str) -> None:
                     metrics_formatter.format_metrics()
 
                 live.update(metrics_formatter.table)
-                time.sleep(2)
+                time.sleep(1)
     except click.ClickException as e:
         raise e
     except Exception as e:

vec_inf/client/api.py

@@ -81,7 +81,7 @@ class VecInfClient:
 
     def __init__(self) -> None:
         """Initialize the Vector Inference client."""
-        pass
+        self._metrics_collectors: dict[str, PerformanceMetricsCollector] = {}
 
     def list_models(self) -> list[ModelInfo]:
         """List all available models.
@@ -218,7 +218,13 @@ def get_metrics(self, slurm_job_id: str) -> MetricsResponse:
             - Performance metrics or error message
             - Timestamp of collection
         """
-        performance_metrics_collector = PerformanceMetricsCollector(slurm_job_id)
+        # Use cached collector to preserve state between calls to compute throughput
+        if slurm_job_id not in self._metrics_collectors:
+            self._metrics_collectors[slurm_job_id] = PerformanceMetricsCollector(
+                slurm_job_id
+            )
+
+        performance_metrics_collector = self._metrics_collectors[slurm_job_id]
 
         metrics: Union[dict[str, float], str]
         if not performance_metrics_collector.metrics_url.startswith("http"):