Addresses the remaining comments

Author: sarthakag

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

@@ -43,12 +43,20 @@ def _get_operations():
             return {}
 
 def _write_all_operations(operations):
-    """Store the operations to the file."""
+    """Store the operations to the file safely."""
     try:
-        with open(REPAIR_FILE, 'w', encoding='utf-8') as f:
-            fcntl.lockf(f, fcntl.LOCK_EX)
+        with open(REPAIR_FILE, 'a', encoding='utf-8') as f:
             try:
+                fcntl.lockf(f, fcntl.LOCK_EX | fcntl.LOCK_NB)
+            except (IOError, BlockingIOError):
+                log.warning(f"Could not acquire lock on {REPAIR_FILE}. Another process may be running.")
+                return False
+
+            try:
+                f.seek(0)
+                f.truncate()
                 json.dump(operations, f, indent=4)
+                f.flush()
                 return True
             finally:
                 fcntl.lockf(f, fcntl.LOCK_UN)

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

@@ -251,33 +251,31 @@ 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."""
+    """Get the reason for a node's state using JSON output."""
     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
+        # Use --json to get structured data
+        result = run(f"{lookup().scontrol} show node {nodename} --json")
+        data = json.loads(result.stdout)
+
+        # Access the reason field directly from the JSON structure
+        nodes = data.get('nodes', [])
+        if nodes:
+            reason = nodes[0].get('reason')
+            # Handle the specific formatting logic for brackets if needed
+            if reason and "[" in reason:
+                return reason.split("[")[0].strip()
+            return reason
+    except (subprocess.CalledProcessError, json.JSONDecodeError) as e:
+        log.error(f"Failed to execute scontrol or parse its JSON output for node {nodename}: {e}")
     except Exception as e:
-        log.error(f"Failed to get reason for node {nodename}: {e}")
+        log.error(f"An unexpected error occurred while getting 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()
 
@@ -299,6 +297,14 @@ def get_node_action(nodename: str) -> NodeAction:
         ("POWER_DOWN", "POWERING_UP", "POWERING_DOWN", "POWERED_DOWN")
     ) & (state.flags if state is not None else set())
 
+    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()
+            if inst:
+                return NodeActionRepair(reason=reason)
+
     if (state is None) and (inst is None):
         # Should never happen
         return NodeActionUnknown(None, None)