fix: only retrieve GPU models from _eligible_ nodes

ai4os · Apr 22, 2024 · 3733159 · 3733159
1 parent 6f95d64
commit 3733159
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/ai4papi/nomad/common.py b/ai4papi/nomad/common.py
@@ -364,8 +364,14 @@ def get_gpu_models():
     gpu_models = set()
     nodes = Nomad.nodes.get_nodes(resources=True)
     for node in nodes:
+        # Discard GPU models of nodes that are not eligible
+        if node['SchedulingEligibility'] != 'eligible':
+            continue
+
+        # Retrieve GPU models of the node
         devices = node['NodeResources']['Devices']
         gpus = [d for d in devices if d['Type'] == 'gpu'] if devices else []
         for gpu in gpus:
             gpu_models.add(gpu['Name'])
+
     return list(gpu_models)