add partition info to GRES selection

Author: sjpb

ansible/roles/openondemand/filter_plugins/filters.py

@@ -22,48 +22,51 @@ def to_gres_options(stdout):
     For example with a single GRES defined of `gpu:H200:8' the following
     entries are returned:
       - ['None', 'none']
-      - ['Any gpu (max count: 8)', 'gpu']
-      - ['H200 gpu' (max count: 8)', 'gpu:H200']
+      - ['Any gpu (max count=8, partitions=standard,long)', 'gpu']
+      - ['H200 gpu (max count=8, partitions=standard,long)', 'gpu:H200']
 
     [1] https://osc.github.io/ood-documentation/latest/how-tos/app-development/interactive/form-widgets.html#form-widgets
     [2] https://slurm.schedmd.com/srun.html#OPT_gres
     """  # noqa: E501 pylint: disable=line-too-long
 
     gres_data = {}
-    # key=gres_opt, what would be passed to --gres
-    # value=[label, max_count]
-    gres_data["none"] = ["None", 0]
+    # key=gres_opt - 'name' or 'name:type', i.e. what would be passed to --gres
+    # value={label:str, max_count: int, partitions=[]}
+    gres_data["none"] = {'label':'None', 'max_count':0, 'partitions':['all']}
 
     for line in stdout.splitlines():
-        partition, gres_definitions = (  # pylint: disable=unused-variable
-            line.split()
-        )  # e.g. 'part1 gpu:H200:8(S:0-1),test:foo:1', or 'part2 (null)'
+        # line examples:
+        # 'part1 gpu:H200:8(S:0-1),test:foo:1'
+        # 'part2 (null)'
+        # - First example shows multiple GRES per partition
+        # - Core suffix e.g. '(S:0-1)' only exists for auto-detected gres
+        # - stackhpc.openhpc role guarantees that name:type:count all exist
+        partition, gres_definitions = (line.split())
         for gres in gres_definitions.split(","):
             if "(null)" in gres:
                 continue
             gres_name, gres_type, gres_count_cores = gres.split(":", maxsplit=2)
-            gres_count = gres_count_cores.split("(")[
-                0
-            ]  # may or may not have the e.g. '(S:0-1)' core definition
+            gres_count = gres_count_cores.split("(")[0]
             for gres_opt in [gres_name, f"{gres_name}:{gres_type}"]:
                 if gres_opt not in gres_data:
                     label = (
                         f"{gres_type} {gres_name}"
                         if ":" in gres_opt
                         else f"Any {gres_opt}"
                     )
-                    gres_data[gres_opt] = [label, gres_count]
-                elif len(gres_data[gres_name]) == 1:
-                    raise ValueError(gres_data[gres_name])
-                elif gres_count > gres_data[gres_name][1]:
-                    gres_data[gres_opt][1] = gres_count
+                    gres_data[gres_opt] = {'label':label, 'max_count':gres_count, 'partitions':[partition]}
+                else:
+                    gres_data[gres_opt]['partitions'].append(partition)
+                    if gres_count > gres_data[gres_name]['max_count']:
+                        gres_data[gres_opt]['max_count'] = gres_count
 
     gres_options = []
     for gres_opt in gres_data:  # pylint: disable=consider-using-dict-items
-        max_count = gres_data[gres_opt][1]
-        label = gres_data[gres_opt][0]
+        max_count = gres_data[gres_opt]['max_count']
+        partitions = gres_data[gres_opt]['partitions']
+        label = gres_data[gres_opt]['label']
         if gres_opt != "none":
-            label += f" (max count: {max_count})"
+            label += f" (max count={max_count}, partitions={','.join(partitions)})"
         gres_options.append((label, gres_opt))
     return gres_options