R&R Slurm integration

Author: sarthakag

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

@@ -191,6 +191,7 @@ locals {
     "mig_flex.py",
     "resume_wrapper.sh",
     "resume.py",
+    "repair.py",
     "setup_network_storage.py",
     "setup.py",
     "slurmsync.py",
@@ -215,6 +216,7 @@ locals {
     "setup.py",
     "slurmsync.py",
     "sort_nodes.py",
+    "repair.py",
     "suspend.py",
     "tpu.py",
     "util.py",

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

@@ -0,0 +1,123 @@
+#!/usr/bin/env python3
+
+# Copyright (C) SchedMD 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 fcntl
+import json
+import logging
+from datetime import datetime, timezone
+from pathlib import Path
+import util
+from util import run, to_hostlist, lookup
+log = logging.getLogger()
+REPAIR_FILE = Path("/slurm/repair_operations.json")
+REPAIR_REASONS = frozenset(["PERFORMANCE", "SDC", "XID", "unspecified"])
+
+def is_node_being_repaired(node):
+    """Check if a node is currently being repaired."""
+    operations = get_operations()
+    return node in operations and operations[node]["status"] == "REPAIR_IN_PROGRESS"
+
+def get_operations():
+    """Get all repair operations from the file."""
+    if not REPAIR_FILE.exists():
+        return {}
+    with open(REPAIR_FILE, 'r') as f:
+        try:
+            return json.load(f)
+        except json.JSONDecodeError:
+            return {}
+
+def store_operations(operations):
+    """Store the operations to the file."""
+    with open(REPAIR_FILE, 'w') as f:
+        fcntl.lockf(f, fcntl.LOCK_EX)
+        try:
+            json.dump(operations, f, indent=4)
+        finally:
+            fcntl.lockf(f, fcntl.LOCK_UN)
+
+def store_operation(node, operation_id, reason):
+    """Store a single repair operation."""
+    operations = get_operations()
+    operations[node] = {
+        "operation_id": operation_id,
+        "reason": reason,
+        "status": "REPAIR_IN_PROGRESS",
+        "timestamp": datetime.now(timezone.utc).isoformat(),
+    }
+    store_operations(operations)
+
+def call_rr_api(node, reason):
+    """Call the R&R API for a given node."""
+    log.info(f"Calling R&R API for node {node} with reason {reason}")
+    inst = lookup().instance(node)
+    if not inst:
+        log.error(f"Instance {node} not found, cannot report fault.")
+        return None
+    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"
+    try:
+        result = run(cmd)
+        op = json.loads(result.stdout)
+        if isinstance(op, list):
+            op = op[0]
+        return op["name"]
+    except Exception as e:
+        log.error(f"Failed to call R&R API for {node}: {e}")
+        return None
+
+def get_operation_status(operation_id):
+    """Get the status of a GCP operation."""
+    cmd = f'gcloud compute operations list --filter="name={operation_id}" --format=json'
+    try:
+        result = run(cmd)
+        operations_list = json.loads(result.stdout)
+        if operations_list and len(operations_list) > 0:
+            return operations_list[0]
+
+        return None 
+    except Exception as e:
+        log.error(f"Failed to get operation status for {operation_id}: {e}")
+        return None
+
+def poll_operations():
+    """Poll the status of ongoing repair operations."""
+    operations = get_operations()
+    if not operations:
+        return
+
+    log.info("Polling repair operations")
+    for node, op_details in operations.items():
+        if op_details["status"] == "REPAIR_IN_PROGRESS":
+            op_status = get_operation_status(op_details["operation_id"])
+            if not op_status:
+                continue
+
+            if op_status.get("status") == "DONE":
+                if op_status.get("error"):
+                    log.error(f"Repair operation for {node} failed: {op_status['error']}")
+                    op_details["status"] = "FAILURE"
+                    run(f"{lookup().scontrol} update nodename={node} state=down reason='Repair failed'")
+                else:
+                    log.info(f"Repair operation for {node} succeeded.")
+                    op_details["status"] = "SUCCESS"
+        if op_details["status"] == "SUCCESS":
+            inst = lookup().instance(node)
+            if inst and inst.status == "RUNNING":
+                log.info(f"Node {node} is back online, undraining.")
+                op_details["status"] = "RECOVERED"
+
+    store_operations(operations)
+

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

@@ -47,6 +47,7 @@
 import tpu
 import conf
 import watch_delete_vm_op
+import repair
 
 log = logging.getLogger()
 
@@ -119,6 +120,17 @@ def apply(self, nodes: List[str]) -> None:
         log.info(f"{len(nodes)} nodes set down ({hostlist}) with reason={self.reason}")
         run(f"{lookup().scontrol} update nodename={hostlist} state=down reason={shlex.quote(self.reason)}")
 
+@dataclass(frozen=True)
+class NodeActionRepair():
+    reason: str
+    def apply(self, nodes: List[str]) -> None:
+        hostlist = util.to_hostlist(nodes)
+        log.info(f"{len(nodes)} nodes to repair ({hostlist}) with reason={self.reason}")
+        for node in nodes:
+            op_id = repair.call_rr_api(node, self.reason)
+            if op_id:
+                repair.store_operation(node, op_id, self.reason)
+
 @dataclass(frozen=True)
 class NodeActionUnknown():
     slurm_state: Optional[NodeState]
@@ -237,10 +249,33 @@ def _find_tpu_node_action(nodename, state) -> NodeAction:
 
     return NodeActionUnchanged()
 
+def get_node_reason(nodename: str) -> Optional[str]:
+    """Get the reason for a node's state."""
+    try:
+        result = run(f"{lookup().scontrol} show node {nodename}")
+        for line in result.stdout.splitlines():
+            if "Reason=" in line:
+                reason = line.split("Reason=")[1].strip()
+                if "[" in reason:
+                    return reason.split("[")[0].strip()
+                return reason
+    except Exception as e:
+        log.error(f"Failed to get reason for node {nodename}: {e}")
+    return None
+
+
 def get_node_action(nodename: str) -> NodeAction:
     """Determine node/instance status that requires action"""
     lkp = lookup()
     state = lkp.node_state(nodename)
+    if state is not None and "DRAIN" in state.flags:
+        reason = get_node_reason(nodename)
+        if reason in repair.REPAIR_REASONS:
+            if repair.is_node_being_repaired(nodename):
+                return NodeActionUnchanged()
+            inst = lkp.instance(nodename.split(".")[0])
+            if inst:
+                return NodeActionRepair(reason=reason)
 
     if lkp.node_is_gke(nodename):
         return NodeActionUnchanged()
@@ -631,6 +666,11 @@ def main():
         except Exception:
             log.exception("failed to sync instances")
 
+        try:
+            repair.poll_operations()
+        except Exception:
+            log.exception("failed to poll repair operations")
+
         try:
             sync_flex_migs(lkp)
         except Exception: