diff --git a/src/executorlib/standalone/interactive/communication.py b/src/executorlib/standalone/interactive/communication.py index a5623dd2..a9f4ea59 100644 --- a/src/executorlib/standalone/interactive/communication.py +++ b/src/executorlib/standalone/interactive/communication.py @@ -150,9 +150,11 @@ def shutdown(self, wait: bool = True): """ result = None if self._spawner.poll(): - result = self.send_and_receive_dict( + output = self.send_and_receive_dict( input_dict={"shutdown": True, "wait": wait} - )["result"] + ) + if "result" in output: + result = output["result"] self._spawner.shutdown(wait=wait) self._reset_socket() return result diff --git a/tests/unit/standalone/interactive/test_communication.py b/tests/unit/standalone/interactive/test_communication.py index 6c0832d4..9d5be030 100644 --- a/tests/unit/standalone/interactive/test_communication.py +++ b/tests/unit/standalone/interactive/test_communication.py @@ -31,6 +31,20 @@ class BrokenSpawner(MpiExecSpawner): def bootup(self, command_lst: list[str], stop_function: Optional[Callable] = None,): return False + +class DelayedExitSpawner(MpiExecSpawner): + """Spawner that reports the process as alive for the first poll() call only, + emulating a worker which exits while shutdown() is waiting for its reply.""" + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self._poll_call_count = 0 + + def poll(self) -> bool: + self._poll_call_count += 1 + return self._poll_call_count == 1 + + class TestInterface(unittest.TestCase): @unittest.skipIf( skip_mpi4py_test, "mpi4py is not installed, so the mpi4py tests are skipped." @@ -146,6 +160,24 @@ def test_interface_serial_with_error(self): self.assertFalse(interface._spawner.poll()) interface.shutdown(wait=True) + def test_interface_shutdown_with_process_exiting_during_wait(self): + cloudpickle_register(ind=1) + interface = SocketInterface( + spawner=DelayedExitSpawner(cwd=None, cores=1, openmpi_oversubscribe=False), + log_obj_size=False, + time_out_ms=100, + ) + port = interface.bind_to_random_port() + # Connect a peer so the PAIR socket is not in the ZMQ "mute state" and + # send_dict() does not block forever. The peer never replies, emulating + # a worker process that exits while shutdown() is waiting for its reply. + context, socket = interface_connect(host="localhost", port=str(port)) + try: + self.assertIsNone(interface.shutdown(wait=True)) + finally: + socket.close() + context.term() + def test_interface_serial_wrong_input(self): cloudpickle_register(ind=1) interface = SocketInterface(