Fix Custodian only killing MPI parent process but leaving VASP ghost processes blocking nodes (#414)

* fix(vasp): kill entire process group when terminating VASP jobs

Use os.killpg() to send SIGTERM/SIGKILL to the entire process group
instead of just the main process. This ensures all MPI child processes
are properly terminated when custodian needs to kill a VASP job.

- Import signal module
- Use os.killpg() with SIGTERM for graceful termination
- Use os.killpg() with SIGKILL for force kill after timeout
- Handle ProcessLookupError when process group is already dead

* test(vasp): refactor VaspJob.terminate() tests for process group killing

Consolidate 13 tests across 3 classes into 9 focused tests in 1 class,
reducing ~180 lines to ~95 lines while maintaining full coverage.

Key changes:
- Use pytest-style class with fixture instead of unittest setUp/tearDown
- Use SimpleNamespace for cleaner mock access (mocks.process vs mocks["process"])
- Mock os.killpg and os.getpgid to test new process group termination logic

Tests now cover all code paths in terminate():
1. test_terminate_process_already_finished - early return when poll() != None
2. test_terminate_graceful_success - SIGTERM succeeds, verifies wait(timeout=10)
3. test_terminate_force_kill_after_timeout - SIGTERM times out, SIGKILL + kill()
4. test_terminate_oserror_during_sigterm - OSError propagates from killpg
5. test_terminate_oserror_during_wait - OSError propagates from final wait()
6. test_terminate_process_lookup_error_during_sigterm - graceful handling
7. test_terminate_process_lookup_error_during_sigkill - falls back to kill()
8. test_terminate_multiple_calls - second call returns early
9. test_terminate_integration_with_real_process - verifies real process group killed

Removed duplicate tests:
- test_terminate_wait_with_different_timeout_behavior (same as force_kill test)
- test_terminate_kills_process_group_with_sigterm (same as graceful_success)
- test_terminate_force_kills_process_group_with_sigkill (same as force_kill)
- Redundant integration test (kept the one verifying process group is dead)

* Add configurable terminate_timeout and remove unused function

- Add terminate_timeout parameter to VaspJob (default 10s)
- Large MPI jobs can use longer timeouts before SIGKILL escalation
- Log message now includes timeout value for debugging
- Remove unused copy_contcar_to_poscar_if_valid() from utils.py
- Add test_custom_timeout to verify configurable timeout works

* Skip terminate tests on Windows (os.killpg/getpgid are POSIX-only)

- Code now falls back to process.terminate() on Windows
- Tests use create=True to mock POSIX-only functions on Windows
- Only integration test with real subprocess skipped on Windows

* apply @Andrew-S-Rosen's suggested `terminate` changes and update test expectations

* try fix windows CI

* VaspJob.terminate() with signal escalation + added back wait confirmation

- Add wait after SIGTERM to confirm process death before returning
- Use try/except/else pattern instead of sigterm_sent flag for cleaner flow
- Skip wait and go straight to SIGKILL if SIGTERM raises OSError
- Use terminate_timeout consistently (remove hardcoded 5s for SIGKILL)
- Catch OSError specifically instead of broad Exception
- Add ProcessLookupError handling for race conditions
- Fall back to parent process kill on Windows or if process group methods fail
- Add type hint to directory parameter and document it as unused
- Distinguish "killed with SIGKILL" from fallback "killed" in logs
- Patch custodian.vasp.jobs.os.* in tests to ensure POSIX path is tested on all platforms
- Add test for SIGTERM OSError skipping wait and going to SIGKILL

Flow: SIGTERM → wait(timeout) → SIGKILL → wait(timeout) → fallback terminate/kill

* try fix flaky CI

* return on ProcessLookupError after SIGKILL

* revert flaky test fix causing slow CI

* try fix windows CI (again)

* Update logging for process termination signals

Clarify the logging when a `ProcessLookupError` is hit

* Update warning messages for process group handling

---------

Co-authored-by: Andrew S. Rosen <asrosen93@gmail.com>

Author: janosh

src/custodian/vasp/jobs.py

@@ -4,6 +4,7 @@
 import math
 import os
 import shutil
+import signal
 import subprocess
 from shutil import which
 
@@ -84,6 +85,7 @@ def __init__(
         copy_magmom=False,
         auto_continue=False,
         update_incar=False,
+        terminate_timeout: float = 10.0,
     ) -> None:
         """
         This constructor is necessarily complex due to the need for
@@ -142,6 +144,9 @@ def __init__(
                 already present in the INCAR will be updated, i.e., no new parameters will be
                 added even if they are in the final vasprun.xml. Note that settings_override take
                 precedence over updated params.
+            terminate_timeout (float): Timeout in seconds to wait for graceful
+                termination (SIGTERM) before escalating to SIGKILL. Large MPI
+                jobs may need longer timeouts. Defaults to 10.0 seconds.
         """
         self.vasp_cmd = tuple(vasp_cmd)
         self.output_file = output_file
@@ -156,6 +161,7 @@ def __init__(
         self.copy_magmom = copy_magmom
         self.auto_continue = auto_continue
         self.update_incar = update_incar
+        self.terminate_timeout = terminate_timeout
 
         if SENTRY_DSN:
             # if using Sentry logging, add specific VASP executable to scope
@@ -703,20 +709,78 @@ def constrained_opt_run(
             for key in sorted(energies):
                 file.write(f"{key} {energies[key]}\n")
 
-    def terminate(self, directory="./") -> None:
-        """Kill all VASP processes associated with the current job."""
+    def terminate(self, directory: str = "./") -> None:
+        """Kill all VASP processes associated with the current job.
+
+        Tries to kill the entire process group (safest for MPI jobs), then waits
+        to confirm termination. Escalates SIGTERM → SIGKILL → parent process fallback.
+
+        Note: The parent process fallback may leave behind ghost MPI child processes
+        (less likely with srun since SLURM purportedly cleans up process trees).
+
+        Args:
+            directory: Unused, kept for API compatibility with base class.
+        """
+        pid = self._vasp_process.pid
+
         if self._vasp_process.poll() is not None:
-            logger.warning("The process was already done!")
+            logger.warning(f"Process {pid} already terminated")
             return
 
+        if os.name != "nt":
+            # Look up process group ID
+            try:
+                pgid = os.getpgid(pid)
+            except ProcessLookupError:
+                logger.warning(f"Process group for {pid} not found")
+                return
+
+            # Send SIGTERM to the entire process group
+            logger.info(f"Sending SIGTERM to process group {pgid}")
+            try:
+                os.killpg(pgid, signal.SIGTERM)
+            except ProcessLookupError:
+                logger.warning(f"Process group {pgid} not found")
+                return
+            except OSError as exc:
+                logger.warning(f"SIGTERM to process group {pgid} failed: {exc}")
+            else:
+                # Wait for graceful termination (only if SIGTERM was sent)
+                try:
+                    self._vasp_process.wait(timeout=self.terminate_timeout)
+                    logger.info(f"Process {pid} terminated gracefully")
+                    return
+                except subprocess.TimeoutExpired:
+                    logger.warning(f"SIGTERM timeout ({self.terminate_timeout}s), sending SIGKILL")
+
+            # Escalate to SIGKILL
+            logger.info(f"Sending SIGKILL to process group {pgid}")
+            try:
+                os.killpg(pgid, signal.SIGKILL)
+            except ProcessLookupError:
+                logger.warning(f"Process group {pgid} not found")
+                return
+            except OSError as exc:
+                logger.warning(f"SIGKILL to process group {pgid} failed: {exc}")
+            else:
+                # Wait for process to die (only if SIGKILL was sent)
+                try:
+                    self._vasp_process.wait(timeout=self.terminate_timeout)
+                    logger.info(f"Process {pid} killed with SIGKILL")
+                    return
+                except subprocess.TimeoutExpired:
+                    pass  # Fall through to parent process fallback
+
+        # Fall back to killing the parent launcher process (Windows or if above failed)
+        logger.warning(f"Falling back to killing parent process {pid}")
         try:
-            logger.info(f"Killing PID {self._vasp_process.pid}")
             self._vasp_process.terminate()
-            self._vasp_process.wait(timeout=10)
+            self._vasp_process.wait(timeout=self.terminate_timeout)
+            logger.info(f"Process {pid} terminated")
         except subprocess.TimeoutExpired:
-            logger.warning(f"Graceful termination did not work. Force killing PID {self._vasp_process.pid}")
             self._vasp_process.kill()
             self._vasp_process.wait()
+            logger.info(f"Process {pid} killed")
 
 
 class VaspNEBJob(VaspJob):

tests/test_custodian.py

@@ -7,7 +7,15 @@
 import pytest
 from ruamel.yaml import YAML
 
-from custodian.custodian import (
+
+@pytest.fixture(autouse=True)
+def _seed_rng() -> None:
+    """Seed RNG for deterministic ExampleJob.run() behavior across all tests.
+    Needed for ExampleJob.run() with random.uniform(0, 1) to not be flaky."""
+    random.seed(42)
+
+
+from custodian.custodian import (  # noqa: E402
     Custodian,
     ErrorHandler,
     Job,
@@ -47,7 +55,7 @@ def setup(self, directory="./") -> None:
         self.params["total"] = 0
 
     def run(self, directory="./") -> None:
-        sequence = [random.uniform(0, 1) for i in range(100)]
+        sequence = [random.uniform(0, 1) for _ in range(100)]
         self.params["total"] = self.params["initial"] + sum(sequence)
 
     def postprocess(self, directory="./") -> None: