Fixed key error with collecting existing gpus

ucsd-ets · Jan 26, 2024 · c72acfb · c72acfb
1 parent be8e7ec
commit c72acfb
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/src/dsmlp/ext/kube.py b/src/dsmlp/ext/kube.py
@@ -35,7 +35,10 @@ def get_gpus_in_namespace(self, name: str) -> int:
 
         gpu_count = 0
         for pod in pods.items:
-            gpu_count += int(pod.spec.containers.resources.requests['GPU_LABEL'])
+            try:
+                gpu_count += int(pod.spec.containers.resources.requests['GPU_LABEL'])
+            except KeyError:
+                pass
 
         return gpu_count