Merge branch 'main' into dependabot/pip/documentation-dependencies-47…

…543b79e7
pitt-crc · Jun 13, 2024 · 5a3714b · 5a3714b
2 parents 1555f1a + fc78384
commit 5a3714b
Show file tree

Hide file tree

Showing 4 changed files with 27 additions and 7 deletions.
diff --git a/apps/crc_idle.py b/apps/crc_idle.py
@@ -7,6 +7,7 @@
 
 from argparse import Namespace
 from typing import Dict, Tuple
+import re
 
 from .utils.cli import BaseParser
 from .utils.system_info import Shell, Slurm
@@ -109,7 +110,7 @@ def _idle_gpu_resources(cluster: str, partition: str) -> Dict[int, int]:
             _, total, allocated, state = node_info.split('_')
 
             # If the node is in a downed state, report 0 resource availability.
-            if state in ['drain']:
+            if re.search("drain", state):
                 idle = 0
 
             else:

diff --git a/apps/crc_sus.py b/apps/crc_sus.py
@@ -69,7 +69,10 @@ def app_logic(self, args: Namespace) -> None:
             print("Showing SUs for most recently expired Resource Allocation Request:\033[0m")
             alloc_requests = get_most_recent_expired_request(KEYSTONE_URL, keystone_group_id, auth_header)
 
-        per_cluster_totals = get_per_cluster_totals(alloc_requests, auth_header)
+        per_cluster_totals = get_per_cluster_totals(alloc_requests,
+                                                    get_enabled_cluster_ids(KEYSTONE_URL, auth_header),
+                                                    auth_header)
+
         earliest_date = get_earliest_startdate(alloc_requests)
 
         for cluster in per_cluster_totals:

diff --git a/apps/crc_usage.py b/apps/crc_usage.py
@@ -109,10 +109,12 @@ def app_logic(self, args: Namespace) -> None:
             print("Showing usage information for most recently expired Resource Allocation Request: \033[0m")
             alloc_requests = get_most_recent_expired_request(KEYSTONE_URL, keystone_group_id, auth_header)
 
+        clusters = get_enabled_cluster_ids(KEYSTONE_URL, auth_header)
+
         self.print_summary_table(alloc_requests,
                                  args.account,
-                                 get_per_cluster_totals(alloc_requests, auth_header, per_request=True))
+                                 get_per_cluster_totals(alloc_requests, clusters, auth_header, per_request=True))
 
         self.print_usage_table(args.account,
-                               get_per_cluster_totals(alloc_requests, auth_header),
+                               get_per_cluster_totals(alloc_requests, clusters, auth_header),
                                get_earliest_startdate(alloc_requests))
diff --git a/apps/utils/keystone.py b/apps/utils/keystone.py
@@ -5,7 +5,6 @@
 import requests
 
 KEYSTONE_URL = "https://keystone.crc.pitt.edu"
-CLUSTERS = {1: 'MPI', 2: 'SMP', 3: 'HTC', 4: 'GPU'}
 RAWUSAGE_RESET_DATE = date.fromisoformat('2024-05-07')
 
 
@@ -78,15 +77,30 @@ def get_most_recent_expired_request(keystone_url: str, group_pk: int, auth_heade
     return [response.json()[0]]
 
 
-def get_per_cluster_totals(alloc_requests: [dict], auth_header: dict, per_request: bool = False) -> dict:
+def get_enabled_cluster_ids(keystone_url: str, auth_header: dict) -> dict():
+    """Get the list of enabled clusters defined in Keystone along with their IDs"""
+
+    response = requests.get(f"{keystone_url}/allocations/clusters/?enabled=True", headers=auth_header)
+    response.raise_for_status()
+    clusters = {}
+    for cluster in response.json():
+        clusters[cluster['id']] = cluster['name']
+
+    return clusters
+
+
+def get_per_cluster_totals(alloc_requests: [dict],
+                           clusters: dict,
+                           auth_header: dict,
+                           per_request: bool = False) -> dict:
     """Gather the awarded totals across the given requests on each cluster into a dictionary"""
 
     per_cluster_totals = {}
     for request in alloc_requests:
         if per_request:
             per_cluster_totals[request['id']] = {}
         for allocation in get_request_allocations(KEYSTONE_URL, request['id'], auth_header):
-            cluster = CLUSTERS[allocation['cluster']]
+            cluster = clusters[allocation['cluster']]
             if per_request:
                 per_cluster_totals[request['id']].setdefault(cluster, 0)
                 per_cluster_totals[request['id']][cluster] += allocation['awarded']