Fix KeyError in SocketInterface.shutdown() when worker exits during wait

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

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(