Ensure proper clean up when ctrl+c'ing cluster jobs (#1081)

* also clean up main path and log in case clean up fails silently

* also remove pickle input files when ctrl+c'ing slurm jobs; log when deletion does not work

* fix test

Author: philippotto

cluster_tools/cluster_tools/schedulers/cluster_executor.py

@@ -176,6 +176,7 @@ def handle_kill(
 
         self.inner_handle_kill(signum, frame)
         self.wait_thread.stop()
+        self.clean_up()
 
         if (
             existing_sigint_handler != signal.default_int_handler
@@ -248,15 +249,15 @@ def inner_submit(
     ) -> Tuple[List["Future[str]"], List[Tuple[int, int]]]:
         pass
 
-    def _cleanup(self, jobid: str) -> None:
+    def _maybe_mark_logs_for_cleanup(self, jobid: str) -> None:
         """Given a job ID as returned by _start, perform any necessary
         cleanup after the job has finished.
         """
         if self.keep_logs:
             return
 
-        outf = self.format_log_file_path(self.cfut_dir, jobid)
-        self.files_to_clean_up.append(outf)
+        log_path = self.format_log_file_path(self.cfut_dir, jobid)
+        self.files_to_clean_up.append(log_path)
 
     @staticmethod
     @abstractmethod
@@ -367,7 +368,7 @@ def _completion(self, jobid: str, failed_early: bool) -> None:
         if not should_keep_output:
             self.files_to_clean_up.append(outfile_name)
 
-        self._cleanup(jobid)
+        self._maybe_mark_logs_for_cleanup(jobid)
 
     def ensure_not_shutdown(self) -> None:
         if self.was_requested_to_shutdown:
@@ -462,10 +463,10 @@ def get_main_meta_path(cfut_dir: str, workerid: str) -> str:
         return os.path.join(cfut_dir, f"cfut.main_path.{workerid}.txt")
 
     def store_main_path_to_meta_file(self, workerid: str) -> None:
-        with open(
-            self.get_main_meta_path(self.cfut_dir, workerid), "w", encoding="utf-8"
-        ) as file:
+        main_meta_path = self.get_main_meta_path(self.cfut_dir, workerid)
+        with open(main_meta_path, "w", encoding="utf-8") as file:
             file.write(file_path_to_absolute_module(sys.argv[0]))
+        self.files_to_clean_up.append(main_meta_path)
 
     def map_to_futures(
         self,
@@ -615,11 +616,16 @@ def shutdown(self, wait: bool = True, cancel_futures: bool = True) -> None:
         self.wait_thread.stop()
         self.wait_thread.join()
 
+        self.clean_up()
+
+    def clean_up(self) -> None:
         for file_to_clean_up in self.files_to_clean_up:
             try:
                 os.unlink(file_to_clean_up)
-            except OSError:  # noqa: PERF203 `try`-`except` within a loop incurs performance overhead
-                pass
+            except OSError as exc:  # noqa: PERF203 `try`-`except` within a loop incurs performance overhead
+                logging.warning(
+                    f"Could not delete file during clean up. Path: {file_to_clean_up} Exception: {exc}. Continuing..."
+                )
         self.files_to_clean_up = []
 
     def map(  # type: ignore[override]

cluster_tools/tests/test_slurm.py

@@ -79,7 +79,7 @@ def test_slurm_cfut_dir() -> None:
         assert future.result() == 4
 
     assert os.path.exists(cfut_dir)
-    assert len(os.listdir(cfut_dir)) == 2
+    assert len(os.listdir(cfut_dir)) == 1  # only the log file should still exist
 
 
 def test_slurm_max_submit_user() -> None: