Consolidate bastion _kill_job logic (#1193)

Author: raytomato

axlearn/cloud/common/bastion.py

@@ -61,6 +61,7 @@
 import tempfile
 import time
 from concurrent.futures import ThreadPoolExecutor, wait
+from contextlib import suppress
 from datetime import datetime, timezone
 from subprocess import CalledProcessError
 from typing import IO, Any, NamedTuple, Optional, Union
@@ -884,8 +885,20 @@ def _kill_job(self, job: Job):
         """
         if job.command_proc is not None:
             self._wait_and_close_proc(job.command_proc, kill=True)
+            job.command_proc = None
         if job.cleanup_proc is not None:
             self._wait_and_close_proc(job.cleanup_proc, kill=True)
+            job.cleanup_proc = None
+
+    def _remove_local_job(self, job: Job):
+        """Removes a job that is being tracked by self._active_jobs locally."""
+        try:
+            self._kill_job(job)
+            del self._active_jobs[job.spec.name]
+            logging.info("Removed job %s.", job.spec.name)
+        except Exception as e:  # pylint: disable=broad-except
+            logging.warning("Fail to remove a job %s with error: %s", job.spec.name, e)
+            raise
 
     def _sync_jobs(self):
         """Makes the local bastion state consistent with the remote state.
@@ -940,9 +953,7 @@ def _sync_jobs(self):
                 job = self._active_jobs[job_name]
                 if job.state.status != JobStatus.COMPLETED:
                     logging.warning("Detected orphaned job %s! Killing it...", job.spec.name)
-                    self._kill_job(job)
-                logging.info("Removed job %s.", job_name)
-                del self._active_jobs[job_name]
+                self._remove_local_job(job)
             # Detected updated job: exists in both.
             else:
                 curr_job = self._active_jobs[job_name]
@@ -1286,12 +1297,11 @@ def execute(self):
             self._execute()
         except Exception:
             logging.error("Caught exception, will cleanup all child jobs.")
-            for job in self._active_jobs.values():
-                try:
-                    self._kill_job(job)
-                except Exception as e:  # pylint: disable=broad-except
-                    logging.warning("Fail to kill a job with error: %s", e)
-            self._active_jobs = {}
+            for job in [*self._active_jobs.values()]:
+                # Gracefully attempt to kill each job, ignoring any exception without
+                # affecting cleanup of the next job.
+                with suppress(Exception):  # pylint: disable=broad-except
+                    self._remove_local_job(job)
             self._uploader.cleanup()
             raise  # Re-raise.
 

axlearn/cloud/common/bastion_test.py

@@ -1573,15 +1573,17 @@ def test_update_scheduler(
     def test_exception(self):
         patch_signal = mock.patch(f"{bastion.__name__}.send_signal")
         with patch_signal, self._patch_bastion() as mock_bastion:
+            mock_command_proc = mock.Mock()
+            mock_cleanup_proc = mock.Mock()
             mock_job = Job(
                 spec=mock.Mock(),
                 state=mock.Mock(),
-                command_proc=mock.Mock(),
-                cleanup_proc=mock.Mock(),
+                command_proc=mock_command_proc,
+                cleanup_proc=mock_cleanup_proc,
             )
 
             def mock_execute():
-                mock_bastion._active_jobs["test"] = mock_job
+                mock_bastion._active_jobs[mock_job.spec.name] = mock_job
                 raise ValueError("Mock error")
 
             with mock.patch.multiple(
@@ -1591,14 +1593,19 @@ def mock_execute():
             ) as mock_methods:
                 with self.assertRaisesRegex(ValueError, "Mock error"):
                     mock_bastion.execute()
-                self.assertIn(
-                    mock_job.command_proc, mock_methods["_wait_and_close_proc"].call_args_list[0][0]
-                )
-                self.assertIn(
-                    mock_job.cleanup_proc, mock_methods["_wait_and_close_proc"].call_args_list[1][0]
-                )
+                mock_wait_and_close_proc = mock_methods["_wait_and_close_proc"]
+                self.assertIn(mock_command_proc, mock_wait_and_close_proc.call_args_list[0][0])
+                self.assertIn(mock_cleanup_proc, mock_wait_and_close_proc.call_args_list[1][0])
 
-    def test_execute_with_exception_and_job_failure(self):
+    @parameterized.product(
+        kill_job1_error=[None, Exception("Cannot kill job1")],
+        kill_job2_error=[None, Exception("Cannot kill job2")],
+    )
+    def test_execute_with_exception_and_job_failure(
+        self,
+        kill_job1_error: Optional[Exception],
+        kill_job2_error: Optional[Exception],
+    ):
         job_1 = Job(
             spec=mock.Mock(),
             state=mock.Mock(),
@@ -1612,22 +1619,31 @@ def test_execute_with_exception_and_job_failure(self):
             cleanup_proc=mock.Mock(),
         )
         active_jobs = {
-            "job1": job_1,
-            "job2": job_2,
+            job_1.spec.name: job_1,
+            job_2.spec.name: job_2,
         }
 
         with self._patch_bastion() as mock_bastion:
             mock_bastion._execute = mock.Mock(side_effect=Exception("Execution failed"))
-            mock_bastion._kill_job = mock.Mock(side_effect=[Exception("Cannot kill job"), None])
-
+            mock_bastion._kill_job = mock.Mock(side_effect=[kill_job1_error, kill_job2_error])
+            mock_bastion._remove_local_job = mock.Mock(wraps=mock_bastion._remove_local_job)
             mock_bastion._active_jobs = active_jobs
 
             with self.assertRaises(Exception):
                 mock_bastion.execute()
 
-            self.assertEqual(mock_bastion._kill_job.call_count, 2)
+            self.assertEqual(mock_bastion._remove_local_job.call_count, 2)
             expected_calls = [mock.call(job_1), mock.call(job_2)]
-            self.assertEqual(mock_bastion._kill_job.call_args_list, expected_calls)
+            self.assertEqual(mock_bastion._remove_local_job.call_args_list, expected_calls)
+            # A job remains if and only if there is exception during the clean up process.
+            self.assertEqual(
+                job_1 in mock_bastion._active_jobs.values(),
+                kill_job1_error is not None,
+            )
+            self.assertEqual(
+                job_2 in mock_bastion._active_jobs.values(),
+                kill_job2_error is not None,
+            )
 
     def test_sync_jobs_for_valid_pending_to_sudden_invalid_jobs(self):
         """Test behavior of state transition for pending invalid jobs."""