[bugfix]: cover signals during worker startup

Author: Will Lin

fastvideo/tests/worker/test_signal_terminated_worker_is_reported.py

@@ -12,9 +12,9 @@
 memory drops below six percent, so this is the common failure there and it is
 completely silent.
 
-The test replaces worker construction with a lightweight signal injection, so
-it exercises the real handler and outer exception path without initializing a
-distributed process group or model.
+The tests replace startup and worker construction with lightweight signal
+injection, so they exercise the real handler and outer exception path without
+initializing a distributed process group or model.
 """
 from __future__ import annotations
 
@@ -30,6 +30,10 @@ class _ReadyPipe:
 
     def __init__(self) -> None:
         self.closed = False
+        self.messages = []
+
+    def send(self, message) -> None:
+        self.messages.append(message)
 
     def close(self) -> None:
         self.closed = True
@@ -48,22 +52,40 @@ def parent(self):
         (signal.SIGINT, "user interrupted", "out-of-memory daemon"),
     ],
 )
+@pytest.mark.parametrize("injection_point", ["startup", "worker_construction", "worker_loop"])
 def test_worker_main_reports_received_signal(monkeypatch: pytest.MonkeyPatch, signum: int, expected_reason: str,
-                                             unexpected_reason: str) -> None:
+                                             unexpected_reason: str, injection_point: str) -> None:
     installed_handlers = {}
     logged_messages = []
+    workers = []
     ready_pipe = _ReadyPipe()
 
     def install_handler(installed_signum, handler):
         installed_handlers[installed_signum] = handler
 
+    def inject_signal() -> None:
+        installed_handlers[signum](signum, None)
+
     class SignalledWorker:
 
+        READY_STR = "READY"
+
         def __init__(self, *args, **kwargs) -> None:
-            installed_handlers[signum](signum, None)
+            self.shutdown_called = False
+            workers.append(self)
+            if injection_point == "worker_construction":
+                inject_signal()
+
+        def worker_busy_loop(self) -> None:
+            assert injection_point == "worker_loop"
+            inject_signal()
+
+        def shutdown(self) -> None:
+            self.shutdown_called = True
 
     monkeypatch.setattr(multiproc_executor.signal, "signal", install_handler)
-    monkeypatch.setattr(multiproc_executor, "kill_itself_when_parent_died", lambda: None)
+    monkeypatch.setattr(multiproc_executor, "kill_itself_when_parent_died",
+                        inject_signal if injection_point == "startup" else lambda: None)
     monkeypatch.setattr(multiproc_executor.faulthandler, "enable", lambda: None)
     monkeypatch.setattr(multiproc_executor.psutil, "Process", _ParentlessProcess)
     monkeypatch.setattr(multiproc_executor, "WorkerMultiprocProc", SignalledWorker)
@@ -80,3 +102,31 @@ def __init__(self, *args, **kwargs) -> None:
     assert f"Worker 7 received {signal.Signals(signum).name} ({signum})" in logged_messages[0]
     assert expected_reason in logged_messages[0]
     assert unexpected_reason not in logged_messages[0]
+    if injection_point == "worker_loop":
+        assert ready_pipe.messages == [{"status": "READY"}]
+        assert workers[0].shutdown_called
+
+
+def test_worker_main_preserves_unrelated_system_exit(monkeypatch: pytest.MonkeyPatch) -> None:
+    logged_messages = []
+    ready_pipe = _ReadyPipe()
+
+    class ExitingWorker:
+
+        def __init__(self, *args, **kwargs) -> None:
+            raise SystemExit(23)
+
+    monkeypatch.setattr(multiproc_executor.signal, "signal", lambda *args: None)
+    monkeypatch.setattr(multiproc_executor, "kill_itself_when_parent_died", lambda: None)
+    monkeypatch.setattr(multiproc_executor.faulthandler, "enable", lambda: None)
+    monkeypatch.setattr(multiproc_executor.psutil, "Process", _ParentlessProcess)
+    monkeypatch.setattr(multiproc_executor, "WorkerMultiprocProc", ExitingWorker)
+    monkeypatch.setattr(multiproc_executor.logger, "exception",
+                        lambda message, *args: logged_messages.append(message % args))
+
+    with pytest.raises(SystemExit) as exc_info:
+        WorkerMultiprocProc.worker_main(ready_pipe=ready_pipe, rank=7)
+
+    assert exc_info.value.code == 23
+    assert ready_pipe.closed
+    assert logged_messages == []

fastvideo/worker/multiproc_executor.py

@@ -551,18 +551,22 @@ def signal_handler(signum, frame):
                 shutdown_requested = True
                 raise _WorkerSignalExit(signum)
 
-        # Either SIGTERM or SIGINT will terminate the worker
-        signal.signal(signal.SIGTERM, signal_handler)
-        signal.signal(signal.SIGINT, signal_handler)
-        kill_itself_when_parent_died()
-        faulthandler.enable()
-        parent_process = psutil.Process().parent()
-
         worker = None
         ready_pipe = kwargs.pop("ready_pipe")
         rank = kwargs.get("rank")
+        parent_process = None
 
         try:
+            # Keep all setup after handler installation inside this guarded
+            # region. The parent may terminate peer workers as soon as one
+            # worker fails, including while another peer is still starting.
+            # Either SIGTERM or SIGINT will terminate the worker.
+            signal.signal(signal.SIGTERM, signal_handler)
+            signal.signal(signal.SIGINT, signal_handler)
+            kill_itself_when_parent_died()
+            faulthandler.enable()
+            parent_process = psutil.Process().parent()
+
             worker = WorkerMultiprocProc(*args, **kwargs)
 
             # Send READY once we know everything is loaded