Adding chunking of placements based on topology size

Author: cdunbar13

community/modules/compute/schedmd-slurm-gcp-v6-nodeset/outputs.tf

@@ -37,6 +37,11 @@ output "nodeset" {
     error_message = "accelerator_topology requires enable_placement to be set to true."
   }
 
+  precondition {
+    condition     = (var.accelerator_topology == null) || try(tonumber(split("x", var.accelerator_topology)[1]) % local.guest_accelerator[0].count == 0, false)
+    error_message = "accelerator_topology must be divisible by number of gpus in machine."
+  }
+
   precondition {
     condition     = var.placement_max_distance == null || var.enable_placement
     error_message = "placement_max_distance requires enable_placement to be set to true."

community/modules/scheduler/schedmd-slurm-gcp-v6-controller/modules/slurm_files/scripts/resume.py

@@ -475,14 +475,14 @@ def down_nodes_notify_jobs(nodes: List[str], reason: str, resume_data: Optional[
 
 
 
-def create_placement_request(pg_name: str, region: str, max_distance: Optional[int], gpu_topology: Optional[str]):
+def create_placement_request(pg_name: str, region: str, max_distance: Optional[int], accelerator_topology: Optional[str]):
     config = {
         "name": pg_name,
         "region": region,
         "groupPlacementPolicy": {
             "collocation": "COLLOCATED",
             "maxDistance": max_distance,
-            "gpuTopology": gpu_topology
+            "gpuTopology": accelerator_topology,
         },
     }
 
@@ -554,14 +554,46 @@ def _allocate_nodes_to_placements(nodes: List[str], excl_job_id:Optional[int], l
 
     return placements
 
+def calculate_hosts_per_topo(accelerator_topology: str, machine_type: NSDict) -> int:
+    # Calculate total number of hosts per topology (Assumes format: '1x72')
+    try:
+        top_split = [int(x) for x in accelerator_topology.split("x")]
+    except Exception as e:
+        log.error(f"Accelerator topology {accelerator_topology} is formatted incorrectly.")
+        raise e
+
+    if len(machine_type.accelerators) == 0:
+        gpus_per_machine = 0
+    else: 
+        gpus_per_machine = machine_type.accelerators[0].count
+
+    if len(top_split) != 2:
+        log.error(f"Accelerator topology {accelerator_topology} is formatted incorrectly.")
+    elif top_split[0] <= 0 or top_split[1] <= 0:
+        log.error(f"Accelerator topology {accelerator_topology} is formatted incorrectly.")
+    elif gpus_per_machine <= 0:
+        log.error(f"The machine type has no accelerators. Cannot use accelerator topology {accelerator_topology}.")
+    elif top_split[1] % gpus_per_machine:
+        log.error(f"The GPU count {gpus_per_machine} per node is not a factor of the accelerator topology {accelerator_topology}")
+    
+    return (top_split[0] * top_split[1]) // gpus_per_machine
+ 
 def calculate_chunk_size(nodeset: NSDict, lkp: util.Lookup) -> int:
-    # Calculates the chunk size based on max distance value received
-    machine_type = lkp.template_info(nodeset.instance_template).machine_type.family
+    # Calculates the chunk size based on max distance value received or accelerator topology
+    # Assuming nodeset is not tpu
+    machine_type = lkp.template_info(nodeset.instance_template).machine_type
     max_distance = nodeset.placement_max_distance
+    accelerator_topology = nodeset.accelerator_topology
+
+    # Look for accelerator topology first
+    if accelerator_topology:
+        hosts_per_topo = calculate_hosts_per_topo(accelerator_topology, machine_type)
+        return hosts_per_topo
+
     if max_distance == 1:
         return 22
     elif max_distance == 2:
-        if machine_type.startswith("a3"):
+        if machine_type.family.startswith("a3"):
             return 256
         else:
             return 150
@@ -574,7 +606,7 @@ def create_nodeset_placements(nodes: List[str], excl_job_id:Optional[int], lkp:
     placements = _allocate_nodes_to_placements(nodes, excl_job_id, lkp)
     region = lkp.node_region(nodes[0])
     max_distance = lkp.node_nodeset(nodes[0]).get('placement_max_distance')
-    gpu_topology = lkp.node_nodeset(nodes[0]).get('accelerator_topology') if not lkp.node_is_tpu(nodes[0]) else None
+    accelerator_topology = lkp.nodeset_accelerator_topology(lkp.node_nodeset_name(nodes[0]))
 
     if log.isEnabledFor(logging.DEBUG):
         debug_p = {p.placement: to_hostlist(p.nodes) for p in placements}
@@ -583,7 +615,7 @@ def create_nodeset_placements(nodes: List[str], excl_job_id:Optional[int], lkp:
         )
 
     requests = {
-        p.placement: create_placement_request(p.placement, region, max_distance, gpu_topology) for p in placements if p.placement
+        p.placement: create_placement_request(p.placement, region, max_distance, accelerator_topology) for p in placements if p.placement
     }
     if not requests:
         return placements

community/modules/scheduler/schedmd-slurm-gcp-v6-controller/modules/slurm_files/scripts/util.py

@@ -1615,6 +1615,9 @@ def node_is_tpu(self, node_name=None):
         nodeset_name = self.node_nodeset_name(node_name)
         return self.cfg.nodeset_tpu.get(nodeset_name) is not None
 
+    def nodeset_is_tpu(self, nodeset_name=None) -> bool:
+        return self.cfg.nodeset_tpu.get(nodeset_name) is not None
+
     def node_is_fr(self, node_name:str) -> bool:
         return bool(self.node_nodeset(node_name).future_reservation)
 
@@ -1639,6 +1642,11 @@ def node_region(self, node_name=None):
         nodeset = self.node_nodeset(node_name)
         return parse_self_link(nodeset.subnetwork).region
 
+    def nodeset_accelerator_topology(self, nodeset_name: str) -> Optional[str]:
+        if not self.nodeset_is_tpu(nodeset_name):
+            return self.cfg.nodeset[nodeset_name].get('accelerator_topology')
+        return None
+
     def nodeset_prefix(self, nodeset_name):
         return f"{self.cfg.slurm_cluster_name}-{nodeset_name}"