Merge pull request GoogleCloudPlatform#4276 from alyssa-sm/large-scale-testing

SlurmGCP Resume Improvements

Author: alyssa-sm

community/modules/scheduler/schedmd-slurm-gcp-v6-controller/modules/slurm_files/etc/slurm.conf.tpl

@@ -28,7 +28,7 @@ DebugFlags=Power
 #
 #
 # TIMERS
-MessageTimeout=60
+MessageTimeout=120
 
 ################################################################################
 #              vvvvv  WARNING: DO NOT MODIFY SECTION BELOW  vvvvv              #

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

@@ -118,7 +118,7 @@ def get(key, default):
             "salloc_wait_nodes",
             "ignore_prefer_validation",
         ]),
-        "ResumeProgram": f"{scripts_dir}/resume.py",
+        "ResumeProgram": f"{scripts_dir}/resume_wrapper.sh",
         "ResumeFailProgram": f"{scripts_dir}/suspend.py",
         "ResumeRate": get("resume_rate", 0),
         "ResumeTimeout": get("resume_timeout", 300),

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

@@ -67,6 +67,7 @@ class ResumeJobData:
 class ResumeData:
     jobs: List[ResumeJobData]
 
+
 def get_resume_file_data() -> Optional[ResumeData]:
     if not (path := os.getenv("SLURM_RESUME_FILE")):
         log.error("SLURM_RESUME_FILE was not in environment. Cannot get detailed job, node, partition allocation data.")
@@ -77,7 +78,6 @@ def get_resume_file_data() -> Optional[ResumeData]:
 
     jobs = []
     for jo in data.get("jobs", []):
-
         job = ResumeJobData(
             job_id = jo.get("job_id"),
             partition = jo.get("partition"),
@@ -181,7 +181,7 @@ def create_instances_request(nodes: List[str], placement_group: Optional[str], e
         ),
     )
 
-    if placement_group:
+    if placement_group and excl_job_id is not None:
         pass # do not set minCount to force "all or nothing" behavior
     else:
         body["minCount"] = 1
@@ -653,11 +653,10 @@ def main(nodelist: str) -> None:
     if not nodes:
         log.info("No nodes to resume")
         return
-
     resume_data = get_resume_file_data()
     log.info(f"resume {util.to_hostlist(nodes)}")
     resume_nodes(nodes, resume_data)
-    
+
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     parser.add_argument("nodelist", help="list of nodes to resume")

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

@@ -0,0 +1,41 @@
+#!/bin/bash
+#
+# Copyright 2024 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.
+
+set -e -o pipefail
+
+SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)
+PYTHON_SCRIPT="${SCRIPT_DIR}/resume.py"
+
+# Capture all arguments passed by Slurm (the nodelist).
+ALL_ARGS=("$@")
+
+# This array will hold extra argument for resume.py, like the resume data file.
+UNIQUE_RESUME_FILE=""
+
+# Handle SLURM_RESUME_FILE if provided
+if [ -n "${SLURM_RESUME_FILE-}" ] && [ -f "$SLURM_RESUME_FILE" ]; then
+	SAFE_DIR="/tmp/slurm_resume_data"
+	mkdir -p "$SAFE_DIR"
+
+	UNIQUE_RESUME_FILE="${SAFE_DIR}/resumedata.$$.json"
+	cp "$SLURM_RESUME_FILE" "$UNIQUE_RESUME_FILE"
+fi
+
+SLURM_RESUME_FILE="${UNIQUE_RESUME_FILE}"
+"${PYTHON_SCRIPT}" "${ALL_ARGS[@]}" &
+disown
+
+exit 0

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

@@ -93,7 +93,7 @@ def test_dict_to_conf(value: dict, want: str):
          """LaunchParameters=enable_nss_slurm,use_interactive_step
 SlurmctldParameters=cloud_dns,enable_configless,idle_on_node_suspend
 SchedulerParameters=bf_continue,salloc_wait_nodes,ignore_prefer_validation
-ResumeProgram=ukulele/resume.py
+ResumeProgram=ukulele/resume_wrapper.sh
 ResumeFailProgram=ukulele/suspend.py
 ResumeRate=0
 ResumeTimeout=300
@@ -119,7 +119,7 @@ def test_dict_to_conf(value: dict, want: str):
             },
         ),
          """SchedulerParameters=bf_continue,salloc_wait_nodes,ignore_prefer_validation
-ResumeProgram=ukulele/resume.py
+ResumeProgram=ukulele/resume_wrapper.sh
 ResumeFailProgram=ukulele/suspend.py
 ResumeRate=0
 ResumeTimeout=300
@@ -154,7 +154,7 @@ def test_dict_to_conf(value: dict, want: str):
         ),
          """PrivateData=events,jobs
 SchedulerParameters=bf_busy_nodes,bf_continue,ignore_prefer_validation,nohold_on_prolog_fail
-ResumeProgram=ukulele/resume.py
+ResumeProgram=ukulele/resume_wrapper.sh
 ResumeFailProgram=ukulele/suspend.py
 ResumeRate=1
 ResumeTimeout=2
@@ -174,7 +174,7 @@ def test_dict_to_conf(value: dict, want: str):
 TaskProlog=/slurm/custom_scripts/task_prolog.d/task-prolog
 TaskEpilog=/slurm/custom_scripts/task_epilog.d/task-epilog
 SchedulerParameters=bf_continue,salloc_wait_nodes,ignore_prefer_validation
-ResumeProgram=ukulele/resume.py
+ResumeProgram=ukulele/resume_wrapper.sh
 ResumeFailProgram=ukulele/suspend.py
 ResumeRate=0
 ResumeTimeout=300