diff --git a/flashdreams/flashdreams/serving/webrtc/bootstrap.py b/flashdreams/flashdreams/serving/webrtc/bootstrap.py index 8326219e..8ced6150 100644 --- a/flashdreams/flashdreams/serving/webrtc/bootstrap.py +++ b/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(): diff --git a/flashdreams/tests/test_webrtc_bootstrap.py b/flashdreams/tests/test_webrtc_bootstrap.py new file mode 100644 index 00000000..dfda5908 --- /dev/null +++ b/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()