Hang killer (#8)

Implement a hang killer at BSS level

Author: xiki-tempula

python/BioSimSpace/Sandpit/Exscientia/Process/_process.py

@@ -32,6 +32,7 @@
 import traceback
 
 import pandas as pd
+from loguru import logger
 
 from .._Utils import _try_import
 
@@ -53,6 +54,7 @@
 from ..Protocol._protocol import Protocol as _Protocol
 from .._SireWrappers import System as _System
 from ..Types._type import Type as _Type
+from ..Types import Time as _Time
 from .. import Units as _Units
 from .. import _Utils
 from ..FreeEnergy._restraint import Restraint as _Restraint
@@ -898,7 +900,7 @@ def setSeed(self, seed):
         else:
             self._seed = seed
 
-    def wait(self, max_time=None):
+    def wait(self, max_time=None, inactivity_timeout: None | _Time = None):
         """
         Wait for the process to finish.
 
@@ -939,11 +941,52 @@ def wait(self, max_time=None):
             self._process.wait(max_time)
 
         else:
-            # Wait for the process to finish.
-            self._process.wait()
+            if inactivity_timeout is None:
+                # Wait for the process to finish.
+                self._process.wait()
 
-            # Store the final run time.
-            self.runTime()
+                # Store the final run time.
+                self.runTime()
+            else:
+                inactivity_timeout = int(inactivity_timeout.milliseconds().value())
+                last_time = self._getLastTime()
+                if last_time is None:
+                    # Wait for the process to finish.
+                    self._process.wait()
+
+                    # Store the final run time.
+                    self.runTime()
+                else:
+                    while self.isRunning():
+                        self._process.wait(inactivity_timeout)
+                        if self.isRunning():
+                            current_time = self._getLastTime()
+                            if current_time > last_time:
+                                logger.info(
+                                    f"Current simulation time ({current_time})."
+                                )
+                                last_time = current_time
+                            else:
+                                logger.warning(
+                                    f"Current simulation time ({current_time}) has not advanced compared "
+                                    f"to the last time ({last_time}). The process "
+                                    f"might have hung and will be killed."
+                                )
+                                with open(
+                                    f"{self.workDir()}/{self._name}.out", "a+"
+                                ) as f:
+                                    f.write("Process Hung. Killed.")
+                                self.kill()
+
+    def _getLastTime(self) -> float | None:
+        """This is the base method in the Process base class.
+        Each subclass, such as AMBER or GROMACS, is expected to override this method
+        to provide their own implementation for returning the current time.
+
+        If this method is not overridden, it will return None,
+        and the `inactivity_timeout` feature will be skipped.
+        """
+        return None
 
     def isQueued(self):
         """

python/BioSimSpace/tests/Sandpit/Exscientia/Process/test_process_wait.py

@@ -0,0 +1,55 @@
+from unittest.mock import MagicMock
+
+import BioSimSpace.Sandpit.Exscientia as BSS
+from BioSimSpace.Sandpit.Exscientia.Process._process import Process
+
+
+def test_max_time():
+    process = MagicMock()
+    Process.wait(process, max_time=1)
+    process._process.wait.assert_called_once_with(60000)
+
+
+def test_None_inactivity_timeout():
+    process = MagicMock()
+    Process.wait(process, max_time=None, inactivity_timeout=None)
+    process._process.wait.assert_called_once()
+
+
+def test_inactivity_timeout_no_getLastTime():
+    process = MagicMock()
+    process._getLastTime.return_value = None
+    Process.wait(process, max_time=None, inactivity_timeout=BSS.Units.Time.nanosecond)
+    process._process.wait.assert_called_once()
+
+
+def test_hang(tmp_path):
+    process = MagicMock()
+    process.workDir.return_value = str(tmp_path)
+    process._name = "test"
+    # Using TEST_HANG_COUNTER to mimic simulation progress
+    global TEST_HANG_COUNTER
+    TEST_HANG_COUNTER = 0
+    process.isRunning.return_value = True
+
+    def _getLastTime():
+        global TEST_HANG_COUNTER
+        TEST_HANG_COUNTER += 1
+        # Mimic simulation hang after 10 calls
+        return min(TEST_HANG_COUNTER, 10)
+
+    process._getLastTime = _getLastTime
+
+    def mock_kill():
+        # Mock kill to stop the simulation
+        process.isRunning.return_value = False
+
+    process.kill.side_effect = mock_kill
+
+    Process.wait(process, max_time=None, inactivity_timeout=BSS.Units.Time.nanosecond)
+
+    assert process._process.wait.call_count == 10
+    process.kill.assert_called_once()
+
+    with open(f"{tmp_path}/test.out", "r") as f:
+        assert f.read() == "Process Hung. Killed."