fix(webrtc): reap checkpoint download workers

flashdreams/flashdreams/serving/webrtc/bootstrap.py

@@ -7,6 +7,8 @@
 
 import gc
 import logging
+import multiprocessing
+import time
 from typing import Protocol
 
 import torch
@@ -16,6 +18,8 @@
 
 from flashdreams.core.distributed import configure_loguru_for_distributed
 
+_CHILD_PROCESS_TERMINATION_TIMEOUT_S = 5.0
+
 
 class WebRTCServerLifecycle(Protocol):
     """Rank-coordination surface the serve loop needs from a session manager."""
@@ -30,6 +34,36 @@ def configure_logging(*, world_rank: int | None = None) -> None:
         logging.getLogger(logger_name).setLevel(logging.WARNING)
 
 
+def _terminate_child_processes() -> None:
+    """Terminate and reap subprocesses still owned by this server rank."""
+    children = multiprocessing.active_children()
+    if not children:
+        return
+
+    child_summary = ", ".join(f"{child.name} (pid={child.pid})" for child in children)
+    logger.warning("Terminating child processes during shutdown: {}", child_summary)
+    for child in children:
+        try:
+            child.terminate()
+        except ProcessLookupError:
+            pass
+
+    deadline = time.monotonic() + _CHILD_PROCESS_TERMINATION_TIMEOUT_S
+    for child in children:
+        child.join(timeout=max(0.0, deadline - time.monotonic()))
+
+    survivors = [child for child in children if child.is_alive()]
+    for child in survivors:
+        logger.warning(
+            "Child process {} (pid={}) did not terminate; killing it.",
+            child.name,
+            child.pid,
+        )
+        child.kill()
+    for child in survivors:
+        child.join()
+
+
 def run_webrtc_server(
     *,
     world_rank: int,
@@ -39,18 +73,24 @@ def run_webrtc_server(
     port: int,
 ) -> None:
     """Serve on rank 0, idle on worker ranks, then tear the runtime down."""
-    if world_rank == 0:
-        if app is None:
-            raise ValueError("Rank 0 requires an aiohttp app to serve.")
-        try:
+    if world_rank == 0 and app is None:
+        raise ValueError("Rank 0 requires an aiohttp app to serve.")
+
+    try:
+        if world_rank == 0:
+            assert app is not None
             web.run_app(app, host=host, port=port)
-        finally:
-            session_manager.send_exit_signal()
-    else:
+        else:
+            try:
+                session_manager.wait_for_termination()
+            except KeyboardInterrupt:
+                logger.warning("Worker rank interrupted, shutting down.")
+    finally:
         try:
-            session_manager.wait_for_termination()
-        except KeyboardInterrupt:
-            logger.warning("Worker rank interrupted, shutting down.")
+            _terminate_child_processes()
+        finally:
+            if world_rank == 0:
+                session_manager.send_exit_signal()
 
     gc.collect()
     if torch.cuda.is_available():

flashdreams/tests/test_webrtc_bootstrap.py

@@ -0,0 +1,73 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import multiprocessing
+import time
+
+import pytest
+
+from flashdreams.serving.webrtc import bootstrap
+
+pytestmark = pytest.mark.ci_cpu
+
+
+class _FakeSessionManager:
+    def __init__(self) -> None:
+        self.exit_called = False
+
+    def send_exit_signal(self) -> None:
+        self.exit_called = True
+
+    def wait_for_termination(self) -> None:
+        raise AssertionError("Rank 0 must not wait for a termination signal.")
+
+
+def test_interrupted_server_terminates_child_processes(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    child = multiprocessing.Process(
+        target=time.sleep,
+        args=(60.0,),
+        name="checkpoint-download-worker",
+    )
+    manager = _FakeSessionManager()
+
+    def interrupted_run_app(*args: object, **kwargs: object) -> None:
+        del args, kwargs
+        child.start()
+        raise KeyboardInterrupt
+
+    monkeypatch.setattr(bootstrap.web, "run_app", interrupted_run_app)
+
+    try:
+        with pytest.raises(KeyboardInterrupt):
+            bootstrap.run_webrtc_server(
+                world_rank=0,
+                session_manager=manager,
+                app=bootstrap.web.Application(),
+                host="127.0.0.1",
+                port=8080,
+            )
+
+        child.join(timeout=5.0)
+        assert not child.is_alive()
+        assert child.exitcode is not None
+        assert manager.exit_called
+    finally:
+        if child.is_alive():
+            child.kill()
+        child.join()