diff --git a/runpod/rp_cli/main.py b/runpod/rp_cli/main.py index 5c4b4bfdc..01aa7fe6f 100644 --- a/runpod/rp_cli/main.py +++ b/runpod/rp_cli/main.py @@ -8,6 +8,7 @@ import asyncio import os +import signal import sys from pathlib import Path @@ -405,11 +406,16 @@ async def _wait_for_rerun(watcher: "FileWatcher") -> str: pass await asyncio.sleep(0.5) + # set once a second stop signal arrives: teardown stops waiting for + # a running entrypoint + force_stop: "list[asyncio.Event]" = [] + async def _run_entrypoint_cancellable(fn) -> None: """run user code without blocking the event loop. cancellation waits for the entrypoint to finish so its remote - calls cannot race with session teardown. + calls cannot race with session teardown, unless a second stop + signal forces teardown first. """ loop = asyncio.get_running_loop() done: asyncio.Future[None] = loop.create_future() @@ -422,6 +428,13 @@ def _finish(exc: "BaseException | None") -> None: else: done.set_exception(exc) + def _report(exc: "BaseException | None") -> None: + try: + loop.call_soon_threadsafe(_finish, exc) + except RuntimeError: + # a forced stop closed the loop while this entrypoint ran + pass + def _runner() -> None: try: run_entrypoint(fn) @@ -430,17 +443,56 @@ def _runner() -> None: except (Exception, KeyboardInterrupt, SystemExit) as exc: # bind the exception now: `as exc` is unbound once the # except block exits, before the loop callback runs - loop.call_soon_threadsafe(_finish, exc) + _report(exc) else: - loop.call_soon_threadsafe(_finish, None) + _report(None) - threading.Thread(target=_runner).start() + # daemon: a forced stop must not leave process exit waiting on + # a long remote call + threading.Thread(target=_runner, daemon=True).start() try: return await asyncio.shield(done) except asyncio.CancelledError: - await done + if force_stop and not done.done(): + forced = asyncio.ensure_future(force_stop[0].wait()) + await asyncio.wait({done, forced}, return_when=asyncio.FIRST_COMPLETED) + forced.cancel() + # an abandoned entrypoint's error is not retrieved + done.add_done_callback(lambda f: f.cancelled() or f.exception()) + else: + await done raise + def _install_stop_signals() -> None: + """stop cleanly on sigint and sigterm. + + python only turns sigint into KeyboardInterrupt when it was not + ignored at startup, and shells start background jobs with sigint + ignored; sigterm has no handler at all. without these a stopped + session would leave its endpoints running. the first signal + cancels the session (teardown deletes the endpoints), the second + stops waiting for a running entrypoint. + """ + loop = asyncio.get_running_loop() + task = asyncio.current_task() + force_stop.append(asyncio.Event()) + seen = [] + + def _on_signal() -> None: + seen.append(True) + if len(seen) == 1: + task.cancel() + else: + force_stop[0].set() + + for sig in (signal.SIGINT, signal.SIGTERM): + try: + loop.add_signal_handler(sig, _on_signal) + except (NotImplementedError, RuntimeError, ValueError): + # windows event loops have no signal handlers; ctrl-c + # still raises KeyboardInterrupt there + pass + async def _session() -> int: nonlocal apps, entrypoint ui.set_name_width(_all_resource_names(apps)) @@ -448,6 +500,7 @@ async def _session() -> int: watcher = FileWatcher([module.parent]) ui.dev_banner([a.name for a in apps], str(module.name)) + _install_stop_signals() try: await session.start() ui.resources_table(_table_rows(session)) diff --git a/tests/test_cli/test_deploy_command.py b/tests/test_cli/test_deploy_command.py index 2fe5942e7..29b4ca6ff 100644 --- a/tests/test_cli/test_deploy_command.py +++ b/tests/test_cli/test_deploy_command.py @@ -1,5 +1,6 @@ """tests for rp flash deploy and rp flash dev command wiring.""" +import sys from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -290,6 +291,90 @@ async def delete(endpoint_id): assert result.exit_code == 0, result.output assert events == ["cancelled", "finished", "deleted"] + def _signal_session(self, tmp_path, monkeypatch, entrypoint): + import asyncio + + module = tmp_path / "main.py" + module.write_text(ENTRYPOINT_SOURCE) + monkeypatch.chdir(tmp_path) + events = [] + api = AsyncMock() + api.list_my_endpoints.return_value = [] + api.save_endpoint.return_value = {"id": "ep-active"} + + async def delete(endpoint_id): + events.append("deleted") + return True + + api.delete_endpoint.side_effect = delete + with ( + patch("runpod.apps.dev.AppsApiClient", return_value=api), + patch( + "runpod.apps.entrypoint.run_entrypoint", + side_effect=lambda fn: entrypoint(events), + ), + ): + result = _runner().invoke(cli, ["flash", "dev", str(module)]) + return result, events + + @pytest.mark.skipif(sys.platform == "win32", reason="posix signals") + def test_sigterm_deletes_endpoints(self, tmp_path, monkeypatch): + import os + import signal + import time + + def entrypoint(events): + os.kill(os.getpid(), signal.SIGTERM) + time.sleep(0.1) + events.append("finished") + + result, events = self._signal_session(tmp_path, monkeypatch, entrypoint) + assert result.exit_code == 0, result.output + assert events == ["finished", "deleted"] + + @pytest.mark.skipif(sys.platform == "win32", reason="posix signals") + def test_sigint_cleans_up_even_when_ignored_at_start(self, tmp_path, monkeypatch): + import os + import signal + import time + + # shells start background jobs with sigint ignored + previous = signal.signal(signal.SIGINT, signal.SIG_IGN) + + def entrypoint(events): + os.kill(os.getpid(), signal.SIGINT) + time.sleep(0.1) + events.append("finished") + + try: + result, events = self._signal_session(tmp_path, monkeypatch, entrypoint) + finally: + signal.signal(signal.SIGINT, previous) + assert result.exit_code == 0, result.output + assert events == ["finished", "deleted"] + + @pytest.mark.skipif(sys.platform == "win32", reason="posix signals") + def test_second_signal_forces_teardown(self, tmp_path, monkeypatch): + import os + import signal + import threading + + release = threading.Event() + + def entrypoint(events): + os.kill(os.getpid(), signal.SIGTERM) + threading.Timer(0.2, os.kill, (os.getpid(), signal.SIGTERM)).start() + # a long remote call that outlives the stop request + release.wait(5) + events.append("finished") + + try: + result, events = self._signal_session(tmp_path, monkeypatch, entrypoint) + finally: + release.set() + assert result.exit_code == 0, result.output + assert events[0] == "deleted" + def test_module_without_entrypoint(self, tmp_path, monkeypatch): module = tmp_path / "main.py" module.write_text(APP_SOURCE)