Source purchasedNodes from Facility.computepurchases

modules/coact.py

@@ -964,7 +964,7 @@ def get(self, date: str) -> Iterator[OveragePoint]:
     def get_data(self) -> dict:
         """Fetch usage data from GraphQL."""
         per_window_template = Template(
-            """_$key: facilityRecentComputeUsage(pastMinutes:$minutes) { cluster: clustername, facility, percentUsed, purchasedNodes }"""
+            """_$key: facilityRecentComputeUsage(pastMinutes:$minutes) { cluster: clustername, facility, percentUsed }"""
         )
         logger.trace(f"Fetching windows {self.windows}")
         all_windows = []
@@ -974,6 +974,7 @@ def get_data(self) -> dict:
 
         query = "query usage {"
         query += "\n".join(all_windows) + ",\n"
+        query += "facilities(filter:{}) { name, computepurchases { clustername, purchased } },\n"
         query += "repos { facility, allocs: currentComputeAllocations { cluster: clustername, start, end } }"
         query += "\n}"
 
@@ -984,26 +985,29 @@ def get_data(self) -> dict:
 
     def format_data(self, result: dict) -> dict:
         """Format the raw data for processing."""
+        # Build purchased nodes lookup from Facility.computepurchases
+        fac_purchases = {}
+        for fac in result.pop("facilities", []):
+            for cp in fac.get("computepurchases") or []:
+                fac_purchases[(fac["name"].lower(), cp["clustername"].lower())] = cp["purchased"]
+
         current = {}
         for k in result["repos"]:
             f = k["facility"].lower()
             if f not in current:
                 current[f] = {}
             for item in k["allocs"]:
                 c = item["cluster"].lower()
-                current[f][c] = {"held": None, "percentUsed": [], "purchasedNodes": None}
+                current[f][c] = {"held": None, "percentUsed": [], "purchasedNodes": fac_purchases.get((f, c))}
         del result["repos"]
 
         for time, array in result.items():
             logger.trace(f"Looking at time {time} with {array}")
             for a in array:
                 f = a["facility"].lower()
                 c = a["cluster"].lower()
-                logger.trace(f"Setting {f} {c} to {a['percentUsed']} (nodes: {a.get('purchasedNodes')})")
+                logger.trace(f"Setting {f} {c} to {a['percentUsed']}")
                 current[f][c]["percentUsed"].append(int(a["percentUsed"]))
-                # Store purchased nodes (use the value from any time window since it's constant)
-                if a.get("purchasedNodes") is not None and current[f][c]["purchasedNodes"] is None:
-                    current[f][c]["purchasedNodes"] = a["purchasedNodes"]
 
         logger.trace(f"Overages: {current}")
 

tests/test_node_allocation.py

@@ -20,8 +20,11 @@ def create_graphql_response(usage_percent: float, nodes: int):
                 ]
             }
         ],
+        "facilities": [
+            {"name": "LCLS", "computepurchases": [{"clustername": "ada", "purchased": nodes}]}
+        ],
         "000060": [
-            {"facility": "LCLS", "cluster": "ada", "percentUsed": usage_percent, "purchasedNodes": nodes},
+            {"facility": "LCLS", "cluster": "ada", "percentUsed": usage_percent},
         ]
     }
 
@@ -51,7 +54,7 @@ def test_facility_lifecycle_goes_over_blocks_recovers_and_restores_nodes():
         dry_run=False
     )
 
-    # GraphQL response includes purchasedNodes from coact-api
+    # GraphQL response includes purchasedNodes from Facility.computepurchases
     graphql_response = {
         "repos": [
             {
@@ -61,8 +64,11 @@ def test_facility_lifecycle_goes_over_blocks_recovers_and_restores_nodes():
                 ]
             }
         ],
+        "facilities": [
+            {"name": "LCLS", "computepurchases": [{"clustername": "ada", "purchased": purchased_nodes}]}
+        ],
         "000060": [
-            {"facility": "LCLS", "cluster": "ada", "percentUsed": 85, "purchasedNodes": purchased_nodes},
+            {"facility": "LCLS", "cluster": "ada", "percentUsed": 85},
         ]
     }