Addresses most PR comments

Author: sarthakag

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

@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 
-# Copyright 2025 "Google LLC"
+# 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.
@@ -28,40 +28,47 @@
 
 def is_node_being_repaired(node):
     """Check if a node is currently being repaired."""
-    operations = get_operations()
+    operations = _get_operations()
     return node in operations and operations[node]["status"] == "REPAIR_IN_PROGRESS"
 
-def get_operations():
+def _get_operations():
     """Get all repair operations from the file."""
     if not REPAIR_FILE.exists():
         return {}
     with open(REPAIR_FILE, 'r', encoding='utf-8') as f:
         try:
             return json.load(f)
         except json.JSONDecodeError:
+            log.error(f"Failed to decode JSON from {REPAIR_FILE}, returning empty operations list.")
             return {}
 
-def store_operations(operations):
+def _write_all_operations(operations):
     """Store the operations to the file."""
-    with open(REPAIR_FILE, 'w', encoding='utf-8') as f:
-        fcntl.lockf(f, fcntl.LOCK_EX)
-        try:
-            json.dump(operations, f, indent=4)
-        finally:
-            fcntl.lockf(f, fcntl.LOCK_UN)
+    try:
+        with open(REPAIR_FILE, 'w', encoding='utf-8') as f:
+            fcntl.lockf(f, fcntl.LOCK_EX)
+            try:
+                json.dump(operations, f, indent=4)
+                return True
+            finally:
+                fcntl.lockf(f, fcntl.LOCK_UN)
+    except (IOError, TypeError) as e:
+        log.error(f"Failed to store repair operations to {REPAIR_FILE}: {e}")
+        return False
 
 def store_operation(node, operation_id, reason):
     """Store a single repair operation."""
-    operations = get_operations()
+    operations = _get_operations()
     operations[node] = {
         "operation_id": operation_id,
         "reason": reason,
         "status": "REPAIR_IN_PROGRESS",
         "timestamp": datetime.now(timezone.utc).isoformat(),
     }
-    store_operations(operations)
+    if not _write_all_operations(operations):
+        log.error(f"Failed to persist repair operation for node {node}.")
 
-def call_rr_api(node, reason):
+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)
@@ -71,66 +78,57 @@ def call_rr_api(node, reason):
     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)
+        log.info(f"gcloud compute instances report-host-as-faulty stdout: {result.stdout.strip()}")
         op = json.loads(result.stdout)
         if isinstance(op, list):
             op = op[0]
         return op["name"]
-    except subprocess.CalledProcessError as e:
-        log.error(f"Failed to call R&R API for {node} due to command execution error: {e}")
-        return None
-    except json.JSONDecodeError as e:
-        log.error(f"Failed to parse R&R API response for {node} due to JSON decode error: {e}")
+    except (subprocess.CalledProcessError, json.JSONDecodeError) as e:
+        log.error(f"Failed to call or parse R&R API response for {node}: {e}")
         return None
     except Exception as e:
         log.error(f"An unexpected error occurred while calling R&R API for {node}: {e}")
         return None
 
-def get_operation_status(operation_id):
+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:
+        if operations_list:
             return operations_list[0]
 
         return None 
-    except subprocess.CalledProcessError as e:
-        log.error(f"Failed to get operation status for {operation_id} due to command execution error: {e}")
-        return None
-    except json.JSONDecodeError as e:
-        log.error(f"Failed to parse operation status for {operation_id} due to JSON decode error: {e}")
+    except (subprocess.CalledProcessError, json.JSONDecodeError) as e:
+        log.error(f"Failed to get or parse operation status for {operation_id}: {e}")
         return None
     except Exception as e:
         log.error(f"An unexpected error occurred while getting operation status for {operation_id}: {e}")
         return None
 
 def poll_operations():
     """Poll the status of ongoing repair operations."""
-    operations = get_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:
+            gcp_op_status = _get_operation_status(op_details["operation_id"])
+            if not gcp_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']}")
+            if gcp_op_status.get("status") == "DONE":
+                if gcp_op_status.get("error"):
+                    log.error(f"Repair operation for {node} failed: {gcp_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. Powering down the VM")
                     run(f"{lookup().scontrol} update nodename={node} state=power_down reason='Repair succeeded'")
                     op_details["status"] = "SUCCESS"
-        elif op_details["status"] == "SUCCESS":
-            inst = lookup().instance(node)
-            if inst and inst.status == "RUNNING":
-                log.info(f"Node {node} is back online.")
-                op_details["status"] = "RECOVERED"
 
-    store_operations(operations)
+    if not _write_all_operations(operations):
+        log.error("Failed to persist updated repair operations state after polling.")

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

@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 
-# Copyright 2025 "Google LLC"
+# 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.
@@ -30,7 +30,7 @@ def setUp(self):
         # Reset the REPAIR_FILE path for each test
         repair.REPAIR_FILE = Path("/slurm/repair_operations.json")
 
-    @patch('repair.get_operations')
+    @patch('repair._get_operations')
     def test_is_node_being_repaired(self, mock_get_operations):
         mock_get_operations.return_value = {
             "node-1": {"status": "REPAIR_IN_PROGRESS"},
@@ -43,26 +43,26 @@ def test_is_node_being_repaired(self, mock_get_operations):
     @patch('builtins.open', new_callable=mock_open, read_data='{"node-1": {"status": "REPAIR_IN_PROGRESS"}}')
     @patch('pathlib.Path.exists', return_value=True)
     def test_get_operations_success(self, mock_exists, mock_open_file):
-        ops = repair.get_operations()
+        ops = repair._get_operations()
         self.assertEqual(ops, {"node-1": {"status": "REPAIR_IN_PROGRESS"}})
         mock_open_file.assert_called_with(repair.REPAIR_FILE, 'r', encoding='utf-8')
 
     @patch('builtins.open', new_callable=mock_open, read_data='invalid json')
     @patch('pathlib.Path.exists', return_value=True)
     def test_get_operations_json_decode_error(self, mock_exists, mock_open_file):
-        ops = repair.get_operations()
+        ops = repair._get_operations()
         self.assertEqual(ops, {})
 
     @patch('pathlib.Path.exists', return_value=False)
     def test_get_operations_file_not_found(self, mock_exists):
-        ops = repair.get_operations()
+        ops = repair._get_operations()
         self.assertEqual(ops, {})
 
     @patch('builtins.open', new_callable=mock_open)
     @patch('fcntl.lockf')
-    def test_store_operations(self, mock_lockf, mock_open_file):
+    def test_write_all_operations(self, mock_lockf, mock_open_file):
         operations = {"node-1": {"status": "SUCCESS"}}
-        repair.store_operations(operations)
+        repair._write_all_operations(operations)
         mock_open_file.assert_called_with(repair.REPAIR_FILE, 'w', encoding='utf-8')
         handle = mock_open_file()
         expected_json_string = json.dumps(operations, indent=4)
@@ -72,10 +72,10 @@ def test_store_operations(self, mock_lockf, mock_open_file):
         mock_lockf.assert_any_call(handle, fcntl.LOCK_EX)
         mock_lockf.assert_any_call(handle, fcntl.LOCK_UN)
 
-    @patch('repair.get_operations', return_value={})
-    @patch('repair.store_operations')
+    @patch('repair._get_operations', return_value={})
+    @patch('repair._write_all_operations')
     @patch('repair.datetime')
-    def test_store_operation(self, mock_datetime, mock_store_operations, mock_get_operations):
+    def test_store_operation(self, mock_datetime, mock_write_all_operations, mock_get_operations):
         mock_now = datetime(2025, 1, 1, tzinfo=timezone.utc)
         mock_datetime.now.return_value = mock_now
 
@@ -89,7 +89,7 @@ def test_store_operation(self, mock_datetime, mock_store_operations, mock_get_op
                 "timestamp": mock_now.isoformat(),
             }
         }
-        mock_store_operations.assert_called_with(expected_operations)
+        mock_write_all_operations.assert_called_with(expected_operations)
 
     @patch('repair.lookup')
     @patch('repair.run')
@@ -104,14 +104,14 @@ def test_call_rr_api_success(self, mock_run, mock_lookup):
             returncode=0
         )
 
-        op_id = repair.call_rr_api("node-1", "XID")
+        op_id = repair._call_rr_api("node-1", "XID")
         self.assertEqual(op_id, "op-123")
         mock_run.assert_called_once()
 
     @patch('repair.lookup')
     def test_call_rr_api_instance_not_found(self, mock_lookup):
         mock_lookup.return_value.instance.return_value = None
-        op_id = repair.call_rr_api("node-1", "XID")
+        op_id = repair._call_rr_api("node-1", "XID")
         self.assertIsNone(op_id)
 
     @patch('repair.lookup')
@@ -120,24 +120,24 @@ def test_call_rr_api_run_error(self, mock_run, mock_lookup):
         mock_instance = MagicMock()
         mock_instance.zone = "us-central1-a"
         mock_lookup.return_value.instance.return_value = mock_instance
-        op_id = repair.call_rr_api("node-1", "XID")
+        op_id = repair._call_rr_api("node-1", "XID")
         self.assertIsNone(op_id)
 
     @patch('repair.run')
     def test_get_operation_status_success(self, mock_run):
         mock_run.return_value.stdout = '[{"status": "DONE"}]'
-        status = repair.get_operation_status("op-123")
+        status = repair._get_operation_status("op-123")
         self.assertEqual(status, {"status": "DONE"})
 
     @patch('repair.run')
     def test_get_operation_status_empty_list(self, mock_run):
         mock_run.return_value.stdout = '[]'
-        status = repair.get_operation_status("op-123")
+        status = repair._get_operation_status("op-123")
         self.assertIsNone(status)
 
-    @patch('repair.get_operations')
-    @patch('repair.store_operations')
-    @patch('repair.get_operation_status')
+    @patch('repair._get_operations')
+    @patch('repair._write_all_operations')
+    @patch('repair._get_operation_status')
     @patch('repair.lookup')
     @patch('repair.run')
     def test_poll_operations(self, mock_run, mock_lookup, mock_get_op_status, mock_store_ops, mock_get_ops):