Fix gpuCount status update logic

stackhpc · Oct 2, 2023 · 308718e · 308718e
1 parent 932180b
commit 308718e
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/python/perftest/models/v1alpha1/pytorch.py b/python/perftest/models/v1alpha1/pytorch.py
@@ -221,9 +221,11 @@ async def pod_modified(
         fetch_pod_log: t.Callable[[], t.Awaitable[str]]
     ):  
         # Set default GPU count if none given in spec
-        gpu_count = pod.get("status", {}).get("gpuCount")
-        if gpu_count is None:
+        # (have to do this in status since spec is immutable)
+        if self.spec.gpu_count is None:
             self.status.gpu_count = (0 if self.spec.device == "cpu" else 1)
+        else:
+            self.status.gpu_count = self.spec.gpu_count
 
         pod_phase = pod.get("status", {}).get("phase", "Unknown")
         if pod_phase == "Running":