diff --git a/lean/components/docker/docker_manager.py b/lean/components/docker/docker_manager.py index c3bc80a0..69544e06 100644 --- a/lean/components/docker/docker_manager.py +++ b/lean/components/docker/docker_manager.py @@ -96,7 +96,9 @@ def run_image(self, image: DockerImage, **kwargs) -> bool: :param kwargs: the kwargs to forward to docker.containers.run :return: True if the command in the container exited successfully, False if not """ - from signal import signal, SIGINT, Signals + import signal as signals + from signal import signal, SIGINT, SIGTERM, Signals + SIGHUP = getattr(signals, "SIGHUP", None) from platform import node from sys import stdout, exit from threading import Thread @@ -224,7 +226,8 @@ def kill_container(force: bool) -> None: self._temp_manager.delete_temporary_directories() exit(1) - # Kill the container on Ctrl+C + # Kill the container on Ctrl+C, and when the CLI itself is terminated (SIGTERM, or SIGHUP when the + # terminal or SSH session closes), so the container never keeps running without the CLI def signal_handler(sig: Signals, frame: FrameType) -> None: nonlocal force_kill_next if not graceful_stop or force_kill_next or kwargs["stop_signal"] == "SIGKILL": @@ -239,7 +242,9 @@ def signal_handler(sig: Signals, frame: FrameType) -> None: kill_thread.daemon = True kill_thread.start() - signal(SIGINT, signal_handler) + # SIGHUP does not exist on Windows + for stop_signal in [SIGINT, SIGTERM] + ([SIGHUP] if SIGHUP is not None else []): + signal(stop_signal, signal_handler) # container.logs() is blocking, we run it on a separate thread so the SIGINT handler works properly # If we run this code on the current thread, SIGINT won't be triggered on Windows when Ctrl+C is triggered diff --git a/tests/components/docker/test_docker_manager.py b/tests/components/docker/test_docker_manager.py index c3a146d4..765676af 100644 --- a/tests/components/docker/test_docker_manager.py +++ b/tests/components/docker/test_docker_manager.py @@ -15,6 +15,7 @@ from unittest import mock from lean.components.docker.docker_manager import DockerManager +from lean.models.docker import DockerImage def _create_docker_manager() -> DockerManager: @@ -46,3 +47,45 @@ def test_write_to_file_does_not_let_the_host_shell_expand_the_payload() -> None: # $type must reach the container untouched (raw JSON: real double quotes, no backslash escaping) assert '"$type"' in echo_command assert "\\" not in echo_command + + +def _run_image_and_capture_signal_handlers(docker_manager: DockerManager, container: mock.Mock) -> dict: + handlers = {} + + docker_client = mock.Mock() + docker_client.containers.run.return_value = container + + with mock.patch.object(docker_manager, "image_installed", return_value=True), \ + mock.patch.object(docker_manager, "create_network"), \ + mock.patch.object(docker_manager, "_get_docker_client", return_value=docker_client), \ + mock.patch("signal.signal", side_effect=lambda sig, handler: handlers.__setitem__(sig, handler)), \ + mock.patch("sys.exit"): + docker_manager.run_image(DockerImage(name="quantconnect/lean", tag="latest")) + + return handlers + + +def test_run_image_stops_the_container_when_the_cli_is_terminated() -> None: + import signal + from time import sleep, time + + docker_manager = _create_docker_manager() + + container = mock.Mock() + container.status = "exited" + container.attrs = {"State": {"ExitCode": 0}} + + handlers = _run_image_and_capture_signal_handlers(docker_manager, container) + + stop_signals = [signal.SIGINT, signal.SIGTERM] + ([signal.SIGHUP] if hasattr(signal, "SIGHUP") else []) + for stop_signal in stop_signals: + assert stop_signal in handlers + + # A terminated CLI must stop the container exactly like Ctrl+C does + handlers[signal.SIGTERM](signal.SIGTERM, None) + + deadline = time() + 5 + while not container.kill.called and time() < deadline: + sleep(0.01) + + container.kill.assert_called_once()