Skip to content

Commit

Permalink
fix #252: handle keyboard interrupts (#270)
Browse files Browse the repository at this point in the history
* fix #252: Handle KeyboardInterrupts when using MDRunner.run
* adding tests
* update CHANGES
  • Loading branch information
jandom authored Nov 15, 2023
1 parent 10516d0 commit 9ecbdea
Show file tree
Hide file tree
Showing 3 changed files with 73 additions and 5 deletions.
1 change: 1 addition & 0 deletions CHANGES
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@ orbeckst, jandom, njzjz
file (#261)
* fixed Python 3.12: No module named ('pkg_resources' #263)
* fixed AttributeDict does not support hasattr (#214)
* fixed handle KeyboardInterrupts when using MDRunner.run() (#252)


2023-09-16 0.8.5
Expand Down
30 changes: 25 additions & 5 deletions gromacs/run.py
Original file line number Diff line number Diff line change
Expand Up @@ -102,6 +102,7 @@ class MDrunnerMPI(gromacs.run.MDrunner):
import subprocess
import os.path
import errno
import signal

# logging
import logging
Expand Down Expand Up @@ -230,6 +231,20 @@ def __init__(self, dirname=os.path.curdir, **kwargs):
self.logname = os.path.realpath(
os.path.join(self.dirname, self.filename(logname, ext="log"))
)
self.process = None
self.signal_handled = False
signal.signal(signal.SIGINT, self.signal_handler)

def signal_handler(self, signum, frame):
"""Custom signal handler for SIGINT."""
if self.process is not None:
try:
self.process.terminate() # Attempt to terminate the subprocess
self.process.wait() # Wait for the subprocess to terminate
self.signal_handled = True
except Exception as e:
logger.error(f"Error terminating subprocess: {e}")
raise KeyboardInterrupt # Re-raise the KeyboardInterrupt to exit the main script

def commandline(self, **mpiargs):
"""Returns simple command line to invoke mdrun.
Expand Down Expand Up @@ -308,17 +323,22 @@ def run(self, pre=None, post=None, mdrunargs=None, **mpiargs):
try:
self.prehook(**pre)
logger.info(" ".join(cmd))
rc = subprocess.call(cmd)
self.process = subprocess.Popen(cmd) # Use Popen instead of call
returncode = self.process.wait() # Wait for the process to complete
except KeyboardInterrupt:
# Handle the keyboard interrupt gracefully
logger.info("Keyboard Interrupt received, terminating the subprocess.")
raise
except:
logger.exception("Failed MD run for unknown reasons.")
raise
finally:
self.posthook(**post)
if rc == 0:
logger.info("MDrun completed ok, returncode = {0:d}".format(rc))
if returncode == 0:
logger.info("MDrun completed ok, returncode = {0:d}".format(returncode))
else:
logger.critical("Failure in MDrun, returncode = {0:d}".format(rc))
return rc
logger.critical("Failure in MDrun, returncode = {0:d}".format(returncode))
return returncode

def run_check(self, **kwargs):
"""Run :program:`mdrun` and check if run completed when it finishes.
Expand Down
47 changes: 47 additions & 0 deletions tests/test_run.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,11 @@

from __future__ import division, absolute_import, print_function

import time
import pytest
import os
import signal
import threading

from .datafiles import datafile

Expand Down Expand Up @@ -54,6 +58,49 @@ def test_MDRunner():
assert rc == 0, "mdrun failed to run through MDrunner"


def test_MDRunner_keyboard_interrupt(monkeypatch):
"""Test that a keyboard interrupt is handled correctly."""

# Create a mock for subprocess.Popen
class MockPopen:
def __init__(self, *args, **kwargs):
pass

def wait(self):
time.sleep(2) # Simulate a long-running process

def terminate(self):
pass

# Use monkeypatch to replace subprocess.Popen with MockPopen
monkeypatch.setattr("gromacs.run.subprocess.Popen", MockPopen)

# Initialize MDrunner
mdrunner = gromacs.run.MDrunner()

# Function to send SIGINT after a delay
def send_interrupt():
time.sleep(1) # Short delay before sending SIGINT
os.kill(os.getpid(), signal.SIGINT)

# Start a thread to send the SIGINT
interrupt_thread = threading.Thread(target=send_interrupt)
interrupt_thread.start()

# Run the MDrunner in the main thread to handle the signal
try:
mdrunner.run()
except KeyboardInterrupt:
# Handle the KeyboardInterrupt within the test
pass

# Ensure the signal handler was called
assert mdrunner.signal_handled, "The signal handler was not called as expected"

# Ensure the interrupt thread has finished
interrupt_thread.join()


class Test_find_gromacs_command(object):
# Gromacs 4 or Gromacs 5 (in this order)
commands = ["grompp", "gmx grompp"]
Expand Down

0 comments on commit 9ecbdea

Please sign in to comment.