Enhance Slurm-GCP resilience to capacity constraints and quota errors

Author: vikramvs-gg

community/modules/scheduler/schedmd-slurm-gcp-v6-controller/modules/slurm_files/README.md

@@ -16,8 +16,13 @@ limitations under the License.
 ## Requirements
 
 | Name | Version |
+<<<<<<< HEAD
 | ---- | ------- |
 | <a name="requirement_terraform"></a> [terraform](#requirement\_terraform) | >= 1.12.2 |
+=======
+|------|---------|
+| <a name="requirement_terraform"></a> [terraform](#requirement\_terraform) | = 1.12.2 |
+>>>>>>> be1982f3a (Enhance Slurm-GCP resilience to capacity constraints and quota errors)
 | <a name="requirement_archive"></a> [archive](#requirement\_archive) | ~> 2.0 |
 | <a name="requirement_google"></a> [google](#requirement\_google) | >= 6.41 |
 | <a name="requirement_local"></a> [local](#requirement\_local) | ~> 2.0 |
@@ -26,7 +31,7 @@ limitations under the License.
 ## Providers
 
 | Name | Version |
-| ---- | ------- |
+|------|---------|
 | <a name="provider_archive"></a> [archive](#provider\_archive) | ~> 2.0 |
 | <a name="provider_google"></a> [google](#provider\_google) | >= 6.41 |
 | <a name="provider_local"></a> [local](#provider\_local) | ~> 2.0 |
@@ -39,7 +44,7 @@ No modules.
 ## Resources
 
 | Name | Type |
-| ---- | ---- |
+|------|------|
 | [google_storage_bucket_object.config](https://registry.terraform.io/providers/hashicorp/google/latest/docs/resources/storage_bucket_object) | resource |
 | [google_storage_bucket_object.controller_startup_scripts](https://registry.terraform.io/providers/hashicorp/google/latest/docs/resources/storage_bucket_object) | resource |
 | [google_storage_bucket_object.devel](https://registry.terraform.io/providers/hashicorp/google/latest/docs/resources/storage_bucket_object) | resource |
@@ -64,9 +69,13 @@ No modules.
 ## Inputs
 
 | Name | Description | Type | Default | Required |
+<<<<<<< HEAD
 | ---- | ----------- | ---- | ------- | :------: |
 | <a name="input_accounting_storage_backup_host"></a> [accounting\_storage\_backup\_host](#input\_accounting\_storage\_backup\_host) | The backup accounting storage host. | `string` | `null` | no |
 | <a name="input_backup_controller_key_timeout"></a> [backup\_controller\_key\_timeout](#input\_backup\_controller\_key\_timeout) | The time in seconds for the backup controller to wait for the shared key to become available. | `number` | `300` | no |
+=======
+|------|-------------|------|---------|:--------:|
+>>>>>>> be1982f3a (Enhance Slurm-GCP resilience to capacity constraints and quota errors)
 | <a name="input_bucket_dir"></a> [bucket\_dir](#input\_bucket\_dir) | Bucket directory for cluster files to be put into. | `string` | `null` | no |
 | <a name="input_bucket_name"></a> [bucket\_name](#input\_bucket\_name) | Name of GCS bucket to use. | `string` | n/a | yes |
 | <a name="input_cgroup_conf_tpl"></a> [cgroup\_conf\_tpl](#input\_cgroup\_conf\_tpl) | Slurm cgroup.conf template file path. | `string` | `null` | no |
@@ -117,7 +126,7 @@ No modules.
 ## Outputs
 
 | Name | Description |
-| ---- | ----------- |
+|------|-------------|
 | <a name="output_bucket_dir"></a> [bucket\_dir](#output\_bucket\_dir) | Path directory within `bucket_name` for Slurm cluster file storage. |
 | <a name="output_bucket_name"></a> [bucket\_name](#output\_bucket\_name) | GCS Bucket name of Slurm cluster file storage. |
 | <a name="output_config"></a> [config](#output\_config) | Cluster configuration. |

community/modules/scheduler/schedmd-slurm-gcp-v6-controller/modules/slurm_files/main.tf

@@ -197,6 +197,7 @@ locals {
     "resume_wrapper.sh",
     "resume.py",
     "repair.py",
+    "error_handler.py",
     "setup_network_storage.py",
     "setup.py",
     "slurmsync.py",
@@ -223,6 +224,7 @@ locals {
     "setup.py",
     "slurmsync.py",
     "sort_nodes.py",
+    "error_handler.py",
     "repair.py",
     "suspend.py",
     "tpu.py",

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

@@ -0,0 +1,38 @@
+# Copyright 2026 "Google LLC"
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import enum
+
+class Action(enum.Enum):
+    REQUEUE = "REQUEUE"
+    IGNORE = "IGNORE"
+
+def classify_gcp_error(error_code: str, error_message: str) -> tuple[Action, str]:
+    """
+    Classifies a GCP API error into an actionable Slurm response.
+    Returns: (Action Enum, Normalized Reason String)
+    """
+    msg_lower = error_message.lower()
+    
+    # 1. Quota Errors -> Transient (active usage limits), require REQUEUE
+    if "quotaexceeded" in error_code.lower() or "quotaexceeded" in msg_lower or ("quota" in msg_lower and "exceeded" in msg_lower) or error_code == "QUOTA_EXCEEDED":
+        return Action.REQUEUE, f"GCP Quota Exceeded: {error_message}"
+        
+    # 2. Capacity Errors -> Transient, require REQUEUE
+    capacity_codes = ["ZONE_RESOURCE_POOL_EXHAUSTED", "VM_MIN_COUNT_NOT_REACHED", "INSUFFICIENT_RESOURCE_CAPACITY"]
+    if any(code in error_code for code in capacity_codes) or "sufficient capacity" in msg_lower:
+        return Action.REQUEUE, f"GCP Capacity Exhausted: {error_message}"
+        
+    # 3. Default fallback
+    return Action.IGNORE, f"GCP Error [{error_code}]: {error_message}"

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

@@ -17,10 +17,15 @@
 
 from typing import List, Optional, Dict, Any
 import argparse
+import random
+import time
+import subprocess
+import datetime
 from datetime import timedelta
 import shlex
 import json
 import logging
+import error_handler
 import os
 import yaml
 import collections
@@ -159,7 +164,7 @@ def dws_flex_duration(dws_flex: NSDict, job_id: Optional[int]) -> int:
             log.info("Job TimeLimit cannot be less than 30 seconds or exceed one week")
     return max_duration
 
-def create_instances_request(nodes: List[str], placement_group: Optional[str], excl_job_id: Optional[int]):
+def create_instances_request(nodes: List[str], placement_group: Optional[str], excl_job_id: Optional[int], is_job_request: bool):
     """Call regionInstances.bulkInsert to create instances"""
     assert 0 < len(nodes) <= BULK_INSERT_LIMIT
 
@@ -181,9 +186,14 @@ def create_instances_request(nodes: List[str], placement_group: Optional[str], e
         ),
     )
 
-    if placement_group and excl_job_id is not None:
-        pass # do not set minCount to force "all or nothing" behavior
+    if is_job_request:
+        # By omitting minCount, GCP will default minCount to equal count.
+        # This guarantees ATOMIC PROVISIONING. If GCP cannot provide all N nodes, it fails instantly
+        # rather than partially provisioning M nodes which waste budget while waiting for the rest.
+        pass
     else:
+        # Static pools and non-job allocations don't require all nodes to boot simultaneously.
+        # minCount=1 allows partial fulfillment so available static nodes can still service smaller jobs.
         body["minCount"] = 1
 
     zone_allow = nodeset.zone_policy_allow or []
@@ -221,6 +231,7 @@ class BulkChunk:
     chunk_idx: int
     excl_job_id: Optional[int]
     placement_group: Optional[str] = None
+    is_job_request: bool = False
 
     @property
     def name(self):
@@ -240,21 +251,25 @@ def group_nodes_bulk(nodes: List[str], resume_data: Optional[ResumeData], lkp: u
     non_excl = nodes_set.copy()
     groups : Dict[Optional[int], List[PlacementAndNodes]] = {} # excl_job_id|none -> PlacementAndNodes
 
-    # expand all exclusive job nodelists
+    # expand all job nodelists (both exclusive and non-exclusive)
     for job in resume_data.jobs:
-        if not lkp.cfg.partitions[job.partition].enable_job_exclusive: 
+        # Intersect the job's allocated nodes with the nodes currently being resumed.
+        # This fixes a bug where partially provisioned nodes in previous resume calls were improperly included
+        # in the new bulkInsert, causing duplicate provisioning and blocking minCount atomic requests.
+        job_nodes_in_resume = set(job.nodes_alloc) & non_excl
+        if not job_nodes_in_resume:
             continue
 
         groups[job.job_id] = []
-        # placement group assignment is based on all allocated nodes, ...
+        # placement group assignment requires full job.nodes_alloc to be accurate
         for pn in create_placements(job.nodes_alloc, job.job_id, lkp):
             groups[job.job_id].append(
                 PlacementAndNodes(
                     placement=pn.placement,
-                    #... but we only want to handle nodes in nodes_resume in this run.
-                    nodes = sorted(set(pn.nodes) & nodes_set)
+                    # safely intersect with job_nodes_in_resume to prevent duplicate overlapping node requests
+                    nodes = sorted(set(pn.nodes) & job_nodes_in_resume)
                 ))
-        non_excl.difference_update(job.nodes_alloc)
+        non_excl.difference_update(job_nodes_in_resume)
 
     groups[None] = create_placements(sorted(non_excl), excl_job_id=None, lkp=lkp)
 
@@ -280,7 +295,8 @@ def chunk_nodes(nodes: List[str]):
             prefix=lkp.node_prefix(nodes_chunk[0]), # <cluster_name>-<nodeset_name>
             excl_job_id = job_id,
             placement_group=pn.placement,
-            chunk_idx=i)
+            chunk_idx=i,
+            is_job_request=(job_id is not None))
 
         for job_id, placements in groups.items()
         for pn in placements if pn.nodes
@@ -331,7 +347,7 @@ def resume_nodes(nodes: List[str], resume_data: Optional[ResumeData]):
             flex_chunks.append(chunk)
         else:
             bi_inserts[group] = create_instances_request(
-                chunk.nodes, chunk.placement_group, chunk.excl_job_id
+                chunk.nodes, chunk.placement_group, chunk.excl_job_id, chunk.is_job_request
             )
 
     for chunk in flex_chunks:
@@ -353,7 +369,9 @@ def resume_nodes(nodes: List[str], resume_data: Optional[ResumeData]):
         failed_reqs = [str(e) for e in failed.items()]
         log.error("bulkInsert API failures: {}".format("; ".join(failed_reqs)))
         for ident, exc in failed.items():
-            down_nodes_notify_jobs(grouped_nodes[ident].nodes, f"GCP Error: {exc._get_reason()}", resume_data) # type: ignore
+            reason = exc._get_reason() if hasattr(exc, "_get_reason") else str(exc) # type: ignore
+            action, admin_comment = error_handler.classify_gcp_error(reason, reason)
+            handle_resume_failure(grouped_nodes[ident].nodes, f"GCP Error: {reason}", resume_data, action, admin_comment)
 
     if log.isEnabledFor(logging.DEBUG):
         for group, op in started.items():
@@ -449,7 +467,8 @@ def _handle_bulk_insert_op(op: Dict, nodes: List[str], resume_data: Optional[Res
             for err in failed_ops[0]["error"]["errors"]
         )
         if code != "RESOURCE_ALREADY_EXISTS":
-            down_nodes_notify_jobs(failed_nodes, f"GCP Error: {msg}", resume_data)
+            action, admin_comment = error_handler.classify_gcp_error(code, msg)
+            handle_resume_failure(failed_nodes, f"GCP Error: {msg}", resume_data, action, admin_comment)
         log.error(
             f"errors from insert for node '{failed_nodes[0]}' ({failed_ops[0]['name']}): {msg}"
         )
@@ -470,6 +489,52 @@ def down_nodes_notify_jobs(nodes: List[str], reason: str, resume_data: Optional[
     nodelist = util.to_hostlist(nodes)
     log.error(f"Marking nodes {nodelist} as DOWN, reason: {reason}")
     run(f"{lookup().scontrol} update nodename={nodelist} state=down reason={reason_quoted}", check=False)
+
+
+def handle_resume_failure(nodes: List[str], reason: str, resume_data: Optional[ResumeData], action: error_handler.Action, admin_comment: str) -> None:
+    """Handle resume failures via node and job management based on error action."""
+    nodes_set = set(nodes)
+    jobs = resume_data.jobs if resume_data else []
+    
+    admin_reason_quoted = shlex.quote(admin_comment)
+    fallback_reason_quoted = shlex.quote(reason)
+
+    nodelist = util.to_hostlist(nodes)
+    if action == error_handler.Action.REQUEUE:
+        log.error(f"Resetting nodes {nodelist} to power_down. Reason: {reason}")
+        # 1. Force DOWN to instantly strip the POWERING_UP (#) flag and bypass ResumeTimeout
+        run(f"{lookup().scontrol} update nodename={nodelist} state=down reason='Force clear booting state'", check=False)
+        # 2. Return to power_down so Slurm transitions it safely to idle~
+        run(f"{lookup().scontrol} update nodename={nodelist} state=power_down", check=False)
+    else:
+        log.error(f"Marking nodes {nodelist} as DOWN. Reason: {reason}")
+        run(f"{lookup().scontrol} update nodename={nodelist} state=down reason={fallback_reason_quoted}", check=False)
+
+    for job in jobs:
+        if not (set(job.nodes_alloc) & nodes_set):
+            continue
+            
+        run(f"{lookup().scontrol} update jobid={job.job_id} admincomment={admin_reason_quoted}", check=False)
+
+        if action == error_handler.Action.REQUEUE:
+            # To prevent the requeued job from immediately retrying and hammering
+            # the GCP API during a stockout, we calculate a randomized backoff (e.g. 3-5 mins) and
+            # explicitly delay the job's next evaluation via the StartTime parameter.
+            delay_seconds = random.randint(180, 300)
+            backoff_time = (datetime.datetime.now() + datetime.timedelta(seconds=delay_seconds)).strftime("%Y-%m-%dT%H:%M:%S")
+            log.info(f"Setting job {job.job_id} backoff: {delay_seconds}s (until {backoff_time})")
+            
+            # Wait up to 10 seconds for Slurm to natively requeue the job to PENDING
+            for _ in range(10):
+                try:
+                    res = subprocess.run([lookup().scontrol, "show", "job", str(job.job_id)], capture_output=True, text=True)
+                    if "JobState=PENDING" in res.stdout:
+                        break
+                except Exception:
+                    pass
+                time.sleep(1)
+            
+            run(f"{lookup().scontrol} update JobId={job.job_id} StartTime={backoff_time}", check=False)
 
 
 

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

@@ -124,13 +124,13 @@ def mock_tpu_se(ns: str, lkp) -> TstTPU:
     "c-n:jobNone:g0:0": BulkChunk(
       nodes=["c-n-1", "c-n-2"], prefix="c-n", chunk_idx=0, excl_job_id=None, placement_group="g0"),
     "c-n:job1:g10:0": BulkChunk(
-      nodes=["c-n-0"], prefix="c-n", chunk_idx=0, excl_job_id=1, placement_group="g10"),
+      nodes=["c-n-0"], prefix="c-n", chunk_idx=0, excl_job_id=1, placement_group="g10", is_job_request=True),
     "c-t:0": BulkChunk(
       nodes=["c-t-8", "c-t-9"], prefix="c-t", chunk_idx=0, excl_job_id=None, placement_group=None),
     "c-t:job2:0": BulkChunk(
-      nodes=["c-t-0", "c-t-1"], prefix="c-t", chunk_idx=0, excl_job_id=2, placement_group=None),
+      nodes=["c-t-0", "c-t-1"], prefix="c-t", chunk_idx=0, excl_job_id=2, placement_group=None, is_job_request=True),
     "c-t:job2:1": BulkChunk(
-      nodes=["c-t-2", "c-t-3"], prefix="c-t", chunk_idx=1, excl_job_id=2, placement_group=None),
+      nodes=["c-t-2", "c-t-3"], prefix="c-t", chunk_idx=1, excl_job_id=2, placement_group=None, is_job_request=True),
   }