Skip to content

Commit 59f4df2

Browse files
sarthakagsarthakag
authored
R&R Slurm integration (GoogleCloudPlatform#5003)
1 parent 53926f5 commit 59f4df2

4 files changed

Lines changed: 370 additions & 1 deletion

File tree

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -192,6 +192,7 @@ locals {
192192
"mig_flex.py",
193193
"resume_wrapper.sh",
194194
"resume.py",
195+
"repair.py",
195196
"setup_network_storage.py",
196197
"setup.py",
197198
"slurmsync.py",
@@ -218,6 +219,7 @@ locals {
218219
"setup.py",
219220
"slurmsync.py",
220221
"sort_nodes.py",
222+
"repair.py",
221223
"suspend.py",
222224
"tpu.py",
223225
"util.py",

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

Lines changed: 142 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,142 @@
1+
#!/usr/bin/env python3
2+
3+
# Copyright 2026 "Google LLC"
4+
#
5+
# Licensed under the Apache License, Version 2.0 (the "License");
6+
# you may not use this file except in compliance with the License.
7+
# You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing, software
12+
# distributed under the License is distributed on an "AS IS" BASIS,
13+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
# See the License for the specific language governing permissions and
15+
# limitations under the License.
16+
17+
import fcntl
18+
import json
19+
import logging
20+
from datetime import datetime, timezone
21+
from pathlib import Path
22+
import subprocess
23+
import util
24+
from util import run, to_hostlist, lookup
25+
log = logging.getLogger()
26+
REPAIR_FILE = Path("/slurm/repair_operations.json")
27+
REPAIR_REASONS = frozenset(["PERFORMANCE", "SDC", "XID", "unspecified"])
28+
29+
def is_node_being_repaired(node):
30+
"""Check if a node is currently being repaired."""
31+
operations = _get_operations()
32+
return node in operations and operations[node]["status"] == "REPAIR_IN_PROGRESS"
33+
34+
def _get_operations():
35+
"""Get all repair operations from the file."""
36+
if not REPAIR_FILE.exists():
37+
return {}
38+
with open(REPAIR_FILE, 'r', encoding='utf-8') as f:
39+
try:
40+
return json.load(f)
41+
except json.JSONDecodeError:
42+
log.error(f"Failed to decode JSON from {REPAIR_FILE}, returning empty operations list.")
43+
return {}
44+
45+
def _write_all_operations(operations):
46+
"""Store the operations to the file safely."""
47+
try:
48+
with open(REPAIR_FILE, 'a', encoding='utf-8') as f:
49+
try:
50+
fcntl.lockf(f, fcntl.LOCK_EX | fcntl.LOCK_NB)
51+
except (IOError, BlockingIOError):
52+
log.warning(f"Could not acquire lock on {REPAIR_FILE}. Another process may be running.")
53+
return False
54+
55+
try:
56+
f.seek(0)
57+
f.truncate()
58+
json.dump(operations, f, indent=4)
59+
f.flush()
60+
return True
61+
finally:
62+
fcntl.lockf(f, fcntl.LOCK_UN)
63+
except (IOError, TypeError) as e:
64+
log.error(f"Failed to store repair operations to {REPAIR_FILE}: {e}")
65+
return False
66+
67+
def store_operation(node, operation_id, reason):
68+
"""Store a single repair operation."""
69+
operations = _get_operations()
70+
operations[node] = {
71+
"operation_id": operation_id,
72+
"reason": reason,
73+
"status": "REPAIR_IN_PROGRESS",
74+
"timestamp": datetime.now(timezone.utc).isoformat(),
75+
}
76+
if not _write_all_operations(operations):
77+
log.error(f"Failed to persist repair operation for node {node}.")
78+
79+
def call_rr_api(node, reason):
80+
"""Call the R&R API for a given node."""
81+
log.info(f"Calling R&R API for node {node} with reason {reason}")
82+
inst = lookup().instance(node)
83+
if not inst:
84+
log.error(f"Instance {node} not found, cannot report fault.")
85+
return None
86+
cmd = f"gcloud compute instances report-host-as-faulty {node} --async --disruption-schedule=IMMEDIATE --fault-reasons=behavior={reason},description='VM is managed by Slurm' --zone={inst.zone} --format=json"
87+
try:
88+
result = run(cmd)
89+
log.info(f"gcloud compute instances report-host-as-faulty stdout: {result.stdout.strip()}")
90+
op = json.loads(result.stdout)
91+
if isinstance(op, list):
92+
op = op[0]
93+
return op["name"]
94+
except (subprocess.CalledProcessError, json.JSONDecodeError) as e:
95+
log.error(f"Failed to call or parse R&R API response for {node}: {e}")
96+
return None
97+
except Exception as e:
98+
log.error(f"An unexpected error occurred while calling R&R API for {node}: {e}")
99+
return None
100+
101+
def _get_operation_status(operation_id):
102+
"""Get the status of a GCP operation."""
103+
cmd = f'gcloud compute operations list --filter="name={operation_id}" --format=json'
104+
try:
105+
result = run(cmd)
106+
operations_list = json.loads(result.stdout)
107+
if operations_list:
108+
return operations_list[0]
109+
110+
return None
111+
except (subprocess.CalledProcessError, json.JSONDecodeError) as e:
112+
log.error(f"Failed to get or parse operation status for {operation_id}: {e}")
113+
return None
114+
except Exception as e:
115+
log.error(f"An unexpected error occurred while getting operation status for {operation_id}: {e}")
116+
return None
117+
118+
def poll_operations():
119+
"""Poll the status of ongoing repair operations."""
120+
operations = _get_operations()
121+
if not operations:
122+
return
123+
124+
log.info("Polling repair operations")
125+
for node, op_details in operations.items():
126+
if op_details["status"] == "REPAIR_IN_PROGRESS":
127+
gcp_op_status = _get_operation_status(op_details["operation_id"])
128+
if not gcp_op_status:
129+
continue
130+
131+
if gcp_op_status.get("status") == "DONE":
132+
if gcp_op_status.get("error"):
133+
log.error(f"Repair operation for {node} failed: {gcp_op_status['error']}")
134+
op_details["status"] = "FAILURE"
135+
run(f"{lookup().scontrol} update nodename={node} state=down reason='Repair failed'")
136+
else:
137+
log.info(f"Repair operation for {node} succeeded. Powering down the VM")
138+
run(f"{lookup().scontrol} update nodename={node} state=power_down reason='Repair succeeded'")
139+
op_details["status"] = "SUCCESS"
140+
141+
if not _write_all_operations(operations):
142+
log.error("Failed to persist updated repair operations state after polling.")

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

100755100644
Lines changed: 48 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
import re
2222
import sys
2323
import shlex
24+
import subprocess
2425
from datetime import datetime, timedelta
2526
from itertools import chain
2627
from pathlib import Path
@@ -48,6 +49,7 @@
4849
import conf
4950
import conf_v2411
5051
import watch_delete_vm_op
52+
import repair
5153

5254
log = logging.getLogger()
5355

@@ -120,6 +122,17 @@ def apply(self, nodes: List[str]) -> None:
120122
log.info(f"{len(nodes)} nodes set down ({hostlist}) with reason={self.reason}")
121123
run(f"{lookup().scontrol} update nodename={hostlist} state=down reason={shlex.quote(self.reason)}")
122124

125+
@dataclass(frozen=True)
126+
class NodeActionRepair():
127+
reason: str
128+
def apply(self, nodes: List[str]) -> None:
129+
hostlist = util.to_hostlist(nodes)
130+
log.info(f"{len(nodes)} nodes to repair ({hostlist}) with reason={self.reason}")
131+
for node in nodes:
132+
op_id = repair.call_rr_api(node, self.reason)
133+
if op_id:
134+
repair.store_operation(node, op_id, self.reason)
135+
123136
@dataclass(frozen=True)
124137
class NodeActionUnknown():
125138
slurm_state: Optional[NodeState]
@@ -238,11 +251,32 @@ def _find_tpu_node_action(nodename, state) -> NodeAction:
238251

239252
return NodeActionUnchanged()
240253

254+
def get_node_reason(nodename: str) -> Optional[str]:
255+
"""Get the reason for a node's state using JSON output."""
256+
try:
257+
# Use --json to get structured data
258+
result = run(f"{lookup().scontrol} show node {nodename} --json")
259+
data = json.loads(result.stdout)
260+
261+
# Access the reason field directly from the JSON structure
262+
nodes = data.get('nodes', [])
263+
if nodes:
264+
reason = nodes[0].get('reason')
265+
# Handle the specific formatting logic for brackets if needed
266+
if reason and "[" in reason:
267+
return reason.split("[")[0].strip()
268+
return reason
269+
except (subprocess.CalledProcessError, json.JSONDecodeError) as e:
270+
log.error(f"Failed to execute scontrol or parse its JSON output for node {nodename}: {e}")
271+
except Exception as e:
272+
log.error(f"An unexpected error occurred while getting reason for node {nodename}: {e}")
273+
return None
274+
275+
241276
def get_node_action(nodename: str) -> NodeAction:
242277
"""Determine node/instance status that requires action"""
243278
lkp = lookup()
244279
state = lkp.node_state(nodename)
245-
246280
if lkp.node_is_gke(nodename):
247281
return NodeActionUnchanged()
248282

@@ -264,6 +298,14 @@ def get_node_action(nodename: str) -> NodeAction:
264298
("POWER_DOWN", "POWERING_UP", "POWERING_DOWN", "POWERED_DOWN")
265299
) & (state.flags if state is not None else set())
266300

301+
if state is not None and "DRAIN" in state.flags:
302+
reason = get_node_reason(nodename)
303+
if reason in repair.REPAIR_REASONS:
304+
if repair.is_node_being_repaired(nodename):
305+
return NodeActionUnchanged()
306+
if inst:
307+
return NodeActionRepair(reason=reason)
308+
267309
if (state is None) and (inst is None):
268310
# Should never happen
269311
return NodeActionUnknown(None, None)
@@ -643,6 +685,11 @@ def main():
643685
except Exception:
644686
log.exception("failed to sync instances")
645687

688+
try:
689+
repair.poll_operations()
690+
except Exception:
691+
log.exception("failed to poll repair operations")
692+
646693
try:
647694
sync_flex_migs(lkp)
648695
except Exception:

0 commit comments

Comments
 (0)