Transition DWS Flex-Start to Regional MIGs

Author: abbas1902

community/modules/scheduler/schedmd-slurm-gcp-v6-controller/modules/cleanup_compute/scripts/cleanup_compute.sh

@@ -21,6 +21,7 @@ nodeset_name="$3"
 universe_domain="$4"
 compute_endpoint_version="$5"
 gcloud_dir="$6"
+MAX_ATTEMPTS=3
 
 if [[ $# -ne 5 ]] && [[ $# -ne 6 ]]; then
 	echo "Usage: $0 <project> <cluster_name> <nodeset_name> <universe_domain> <compute_endpoint_version> [<gcloud_dir>]"
@@ -44,14 +45,21 @@ tmpfile=$(mktemp) # have to use a temp file, since `< <(gcloud ...)` doesn't wor
 trap 'rm -f "$tmpfile"' EXIT
 
 echo "Deleting managed instance groups"
-mig_filter="name:${cluster_name}-${nodeset_name}-* AND status!=STOPPING"
+mig_filter="name:${cluster_name}-${nodeset_name}-*"
 gcloud compute instance-groups managed list --format="value(self_link)" --filter="${mig_filter}" >"$tmpfile"
-while batch="$(head -n 2)" && [[ ${#batch} -gt 0 ]]; do
+while batch="$(head -n 5)" && [[ ${#batch} -gt 0 ]]; do
 	groups=$(echo "$batch" | paste -sd " " -) # concat into a single space-separated line
 	# The lack of quotes around ${groups} is intentional and causes each new space-separated "word" to
 	# be treated as independent arguments. See PR#2523
 	# shellcheck disable=SC2086
-	gcloud compute instance-groups managed delete --quiet ${groups} || echo "Failed to delete some instance groups"
+	for _ in $( #occasionally MIGs will fail to delete due to some active transformation happening, so let's retry
+		seq 1 $MAX_ATTEMPTS
+	); do
+		if gcloud compute instance-groups managed delete --quiet ${groups}; then
+			break
+		fi
+		echo "MIG deletion failed, retrying"
+	done
 done <"$tmpfile"
 true >"$tmpfile" # Wipe contents of tmp file
 

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

@@ -45,9 +45,8 @@ def resume_flex_chunk(nodes: List[str], job_id: Optional[int], lkp: util.Lookup)
   assert nodes
   model = nodes[0]
   nodeset = lkp.node_nodeset(model)
-  zones = nodeset.zone_policy_allow
-  assert len(zones) == 1
-  zone = zones[0]
+  assert len(nodeset.zone_policy_allow) > 0
+  region = lkp.node_region(model)
 
   assert nodeset.dws_flex.enabled
 
@@ -58,24 +57,20 @@ def resume_flex_chunk(nodes: List[str], job_id: Optional[int], lkp: util.Lookup)
     mig_name = f"{lkp.cfg.slurm_cluster_name}-{nodeset.nodeset_name}-{uid}"
 
   # Create MIG
-  req = lkp.compute.instanceGroupManagers().insert(
+  req = lkp.compute.regionInstanceGroupManagers().insert(
     project=lkp.project,
-    zone=zone,
+    region=region,
     body=dict(
       name=mig_name,
-      versions=[dict(
-        instanceTemplate=nodeset.instance_template)],
+      versions=[dict(instanceTemplate=nodeset.instance_template)],
       targetSize=0,
-      # TODO(FLEX): uncomment once moved to RMIG
-      # distributionPolicy=dict(
-      #   zones=[
-      #      dict(zone=f"zones/{z}") for z in nodeset.zone_policy_allow
-      #   ],
-      #   targetShape="ANY_SINGLE_ZONE" ),
-      #updatePolicy = dict(
-      #  instanceRedistributionType = "NONE" ),
-      instanceLifecyclePolicy=dict(
-          defaultActionOnFailure= "DO_NOTHING" ), # TODO(FLEX): Not supported yet, migrate once supported
+      distributionPolicy=dict(
+        zones=[
+           dict(zone=f"zones/{z}") for z in nodeset.zone_policy_allow
+        ],
+        targetShape="ANY_SINGLE_ZONE" ),
+      updatePolicy = dict(instanceRedistributionType = "NONE" ),
+      instanceLifecyclePolicy=dict(defaultActionOnFailure= "DO_NOTHING" ), # TODO(FLEX): Not supported yet, migrate once supported
     )
   )
   util.log_api_request(req)
@@ -84,9 +79,9 @@ def resume_flex_chunk(nodes: List[str], job_id: Optional[int], lkp: util.Lookup)
   assert "error" not in res, f"{res}"
 
   # Create resize request
-  req = lkp.compute.instanceGroupManagerResizeRequests().insert(
+  req = lkp.compute.regionInstanceGroupManagerResizeRequests().insert(
     project=lkp.project,
-    zone=zone,
+    region=region,
     instanceGroupManager=mig_name,
     body=dict(
       name="initial-resize",
@@ -105,9 +100,8 @@ def _suspend_flex_mig(mig_self_link: str, nodes: List[str], lkp: util.Lookup) ->
   assert nodes
   model = nodes[0]
   nodeset = lkp.node_nodeset(model)
-  zones = nodeset.zone_policy_allow
-  assert len(zones) == 1
-  zone = zones[0]
+  assert len(nodeset.zone_policy_allow) > 0
+  region = lkp.node_region(model)
   project=lkp.project
   instanceGroupManager=util.trim_self_link(mig_self_link)
 
@@ -118,7 +112,7 @@ def _suspend_flex_mig(mig_self_link: str, nodes: List[str], lkp: util.Lookup) ->
     ] if inst
   ]
 
-  target_mig=lkp.get_mig(lkp.project, zone, instanceGroupManager)
+  target_mig=lkp.get_mig(lkp.project, region, instanceGroupManager)
   assert target_mig
 
   # TODO(FLEX): This will not work if MIG didn't obtain capacity yet.
@@ -130,15 +124,15 @@ def _suspend_flex_mig(mig_self_link: str, nodes: List[str], lkp: util.Lookup) ->
   # - Need to `down_nodes_notify_jobs` for all nodes in MIG, make sure that it doesn't interfere with Slurm suspend-flow. 
 
   if target_mig["targetSize"] == len(nodes): #We can just delete the whole MIG in this case
-    req = lkp.compute.instanceGroupManagers().delete(
+    req = lkp.compute.regionInstanceGroupManagers().delete(
     project=project,
-    zone=zone,
+    region=region,
     instanceGroupManager=instanceGroupManager,
     )
   else:
-    req = lkp.compute.instanceGroupManagers().deleteInstances(
+    req = lkp.compute.regionInstanceGroupManagers().deleteInstances(
       project=project,
-      zone=zone,
+      region=region,
       instanceGroupManager=instanceGroupManager,
       body=dict(
         instances=links,
@@ -156,11 +150,10 @@ def _suspend_provisioning_inst(nodes:List[str], node_template:str, lkp: util.Loo
   assert nodes
   model = nodes[0]
   nodeset = lkp.node_nodeset(model)
-  zones = nodeset.zone_policy_allow
-  assert len(zones) == 1
-  zone = zones[0]
+  assert len(nodeset.zone_policy_allow) > 0
+  region = lkp.node_region(model)
 
-  mig_list=lkp.get_mig_list(lkp.project, zone) #Validated via terraform that this is one
+  mig_list=lkp.get_mig_list(lkp.project, region)
 
   # FLEX (#TODO): If we enter this conditional it's likely this was called so early that MIG creation hasn't started
   # Consider potentially retrying? No natural mechanism for retry currently but we could
@@ -169,18 +162,18 @@ def _suspend_provisioning_inst(nodes:List[str], node_template:str, lkp: util.Loo
   # so until we do this is slurmsync this is a temporary workaround.
 
   if not mig_list or not mig_list.get("items"):
-    log.info("No matching MIG found to delete!")
+    log.info("No matching MIG found to delete! Retrying...")
     sleep(5)
-    mig_list=lkp.get_mig_list(lkp.project, zone)
+    mig_list=lkp.get_mig_list(lkp.project, region)
     if not mig_list or not mig_list.get("items"):
       return
 
   for mig in mig_list["items"]:
     if mig["instanceTemplate"] == node_template:
       if mig["currentActions"]["creating"] > 0 and mig["targetSize"] == mig["currentActions"]["creating"]:
-        req = lkp.compute.instanceGroupManagers().delete(
+        req = lkp.compute.regionInstanceGroupManagers().delete(
           project=lkp.project,
-          zone=zone,
+          region=region,
           instanceGroupManager=util.trim_self_link(mig["selfLink"]),
         )
 

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

@@ -81,6 +81,15 @@ def apply(self, nodes:List[str]) -> None:
         log.info(f"{len(nodes)} instances to power down ({hostlist})")
         run(f"{lookup().scontrol} update nodename={hostlist} state=power_down")
 
+
+@dataclass(frozen=True)
+class NodeActionPowerDownForce():
+    def apply(self, nodes:List[str]) -> None:
+        hostlist = util.to_hostlist(nodes)
+        log.info(f"{len(nodes)} instances to power down ({hostlist})")
+        run(f"{lookup().scontrol} update nodename={hostlist} state=power_down_force")
+
+
 @dataclass(frozen=True)
 class NodeActionDelete():
     def apply(self, nodes:List[str]) -> None:
@@ -292,7 +301,11 @@ def get_node_action(nodename: str) -> NodeAction:
     elif state is None:
         # if state is None here, the instance exists but it's not in Slurm
         return NodeActionUnknown(slurm_state=state, instance_state=inst.status)
-
+    elif lkp.is_flex_node(nodename) and "POWERING_UP" in state.flags:
+        threshold = timedelta(seconds=int(lkp.cfg.compute_startup_scripts_timeout) * 2) #extra buffer for unexpectedly long startup scripts
+        if util.now() - inst.creation_timestamp > threshold:
+            log.info(f"{nodename} was unable to join the cluster after {threshold.seconds}s, potential failure on VM startup. Powering down...")
+            return NodeActionPowerDownForce()
     return NodeActionUnchanged()
 
 

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

@@ -1784,18 +1784,18 @@ def _get_reservation(self, project: str, zone: str, name: str) -> Any:
             project=project, zone=zone, reservation=name).execute()
 
     @lru_cache()
-    def get_mig(self, project: str, zone: str, self_link:str) -> Any:
-        """https://cloud.google.com/compute/docs/reference/rest/v1/instanceGroupManagers"""
-        return self.compute.instanceGroupManagers().get(project=project, zone=zone, instanceGroupManager=self_link).execute()
+    def get_mig(self, project: str, region: str, self_link:str) -> Any:
+        """https://cloud.google.com/compute/docs/reference/rest/v1/regionInstanceGroupManagers"""
+        return self.compute.regionInstanceGroupManagers().get(project=project, region=region, instanceGroupManager=self_link).execute()
 
     @lru_cache
-    def get_mig_instances(self, project: str, zone: str, self_link:str) -> Any:
-        return self.compute.instanceGroupManagers().listManagedInstances(project=project, zone=zone, instanceGroupManager=self_link).execute() 
+    def get_mig_instances(self, project: str, region: str, self_link:str) -> Any:
+        return self.compute.regionInstanceGroupManagers().listManagedInstances(project=project, region=region, instanceGroupManager=self_link).execute() 
 
     @lru_cache()
-    def get_mig_list(self, project: str, zone: str) -> Any:
-        """https://cloud.google.com/compute/docs/reference/rest/v1/instanceGroupManagers"""
-        return self.compute.instanceGroupManagers().list(project=project, zone=zone).execute()
+    def get_mig_list(self, project: str, region: str) -> Any:
+        """https://cloud.google.com/compute/docs/reference/rest/v1/regionInstanceGroupManagers"""
+        return self.compute.regionInstanceGroupManagers().list(project=project, region=region).execute()
 
     @lru_cache()
     def _get_future_reservation(self, project:str, zone:str, name: str) -> Any:
@@ -2087,11 +2087,11 @@ def is_provisioning_flex_node(self, node:str) -> bool:
 
         nodeset = self.node_nodeset(node)
         zones = nodeset.zone_policy_allow
-        assert len(zones) == 1
-        zone = zones[0]
+        assert len(zones) > 0
+        region = self.node_region(node)
 
         potential_migs=[]
-        mig_list=self.get_mig_list(self.project, zone)
+        mig_list=self.get_mig_list(self.project, region)
 
         if not mig_list or not mig_list.get("items"):
             return False
@@ -2100,7 +2100,7 @@ def is_provisioning_flex_node(self, node:str) -> bool:
             if not mig.get("instanceTemplate"): #possibly an old MIG
                 return False
             if mig["instanceTemplate"] == self.node_template(node) and mig["currentActions"]["creating"] > 0:
-                potential_migs.append(self.get_mig_instances(self.project, zone, trim_self_link(mig["selfLink"])))
+                potential_migs.append(self.get_mig_instances(self.project, region, trim_self_link(mig["selfLink"])))
 
         if not potential_migs:
             return False