diff --git a/runpod/apps/dev.py b/runpod/apps/dev.py index c4e95ad1c..1c955f604 100644 --- a/runpod/apps/dev.py +++ b/runpod/apps/dev.py @@ -6,6 +6,7 @@ api is the only source of truth; nothing is persisted locally. """ +import asyncio import logging import os as _os from typing import Dict, List, Optional, Union @@ -288,7 +289,7 @@ async def start(self) -> None: log.info("adopted dev endpoint %s (%s)", name, endpoint_id) else: self._emit("provisioning", spec.name, spec.kind.value, hardware) - result = await self.api.save_endpoint(payload) + result = await self._create_tracked(name, spec.name, payload) endpoint_id = result["id"] log.info("provisioned dev endpoint %s (%s)", name, endpoint_id) @@ -360,8 +361,9 @@ async def refresh(self, apps: List[App]) -> None: if existing_id: payload["id"] = existing_id - - result = await self.api.save_endpoint(payload) + result = await self.api.save_endpoint(payload) + else: + result = await self._create_tracked(name, handle.spec.name, payload) endpoint_id = result["id"] self._endpoints[name] = endpoint_id @@ -393,6 +395,28 @@ async def refresh(self, apps: List[App]) -> None: self.generation, ) + async def _create_tracked(self, name: str, resource: str, payload: Dict) -> Dict: + """create an endpoint, recording it even if the session is + cancelled mid-request, so teardown deletes it.""" + create = asyncio.ensure_future(self.api.save_endpoint(payload)) + try: + return await asyncio.shield(create) + except asyncio.CancelledError: + await self._record_late_create(create, name, resource) + raise + + async def _record_late_create( + self, create: "asyncio.Future[Dict]", name: str, resource: str + ) -> None: + """wait for a create the session was cancelled during, and record + the endpoint it made.""" + try: + result = await create + except Exception: # noqa: BLE001 - the create failed; nothing to track + return + self._endpoints[name] = result["id"] + self._resource_names[name] = resource + async def stop(self, events: Optional[object] = None) -> None: """delete every endpoint this session owns. diff --git a/runpod/rp_cli/main.py b/runpod/rp_cli/main.py index 5c4b4bfdc..5be6aaef0 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,18 @@ async def _wait_for_rerun(watcher: "FileWatcher") -> str: pass await asyncio.sleep(0.5) + # set by a second stop signal: teardown stops waiting for a running + # entrypoint + force_stop = asyncio.Event() + # once teardown starts, signals leave it to finish + tearing_down = False + 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 +430,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,24 +445,72 @@ 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 + await _settle_after_stop(done) raise + async def _settle_after_stop(done: "asyncio.Future[None]") -> None: + """after a stop, wait for the still-running entrypoint so its remote + calls cannot race teardown, unless a second signal forces the stop.""" + if done.done(): + _discard_result(done) + return + forced = asyncio.ensure_future(force_stop.wait()) + await asyncio.wait({done, forced}, return_when=asyncio.FIRST_COMPLETED) + forced.cancel() + done.add_done_callback(_discard_result) + + 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. signals during teardown + are ignored so every endpoint still gets deleted. + """ + loop = asyncio.get_running_loop() + task = asyncio.current_task() + stopping = False + + def _on_signal() -> None: + nonlocal stopping + if tearing_down: + # cancelling teardown would skip the remaining deletes + return + if stopping: + force_stop.set() + else: + stopping = True + task.cancel() + + 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 + nonlocal apps, entrypoint, tearing_down ui.set_name_width(_all_resource_names(apps)) session = DevSession(apps, events=ui.DevEvents()) 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)) @@ -472,6 +535,7 @@ async def _session() -> int: ui.set_name_width(_all_resource_names(apps)) await session.refresh(apps) finally: + tearing_down = True events = ui.CleanupEvents() try: await session.stop(events=events) @@ -486,6 +550,13 @@ async def _session() -> int: raise click.ClickException(str(exc)) from exc +def _discard_result(future: "asyncio.Future[None]") -> None: + """mark an abandoned entrypoint's error as retrieved, so asyncio does + not log it as never retrieved.""" + if not future.cancelled(): + future.exception() + + # ------------------------------------------------------------------ logs diff --git a/tests/test_cli/test_deploy_command.py b/tests/test_cli/test_deploy_command.py index 2fe5942e7..59ea111ee 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,192 @@ 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): + 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") + @pytest.mark.parametrize("stop_signal", ["SIGINT", "SIGTERM"]) + def test_stop_signal_preserves_cancellation_on_entrypoint_failure( + self, tmp_path, monkeypatch, stop_signal + ): + import asyncio + import signal + + handlers = {} + shield = asyncio.shield + + def install_handler(loop, sig, callback): + handlers[sig] = callback + + def stop_on_completion(future): + waiter = shield(future) + if not isinstance(future, asyncio.Task): + future.add_done_callback( + lambda _: handlers[getattr(signal, stop_signal)]() + ) + return waiter + + def entrypoint(events): + events.append("failed") + raise RuntimeError("entrypoint failed while stopping") + + with ( + patch.object( + asyncio.SelectorEventLoop, "add_signal_handler", install_handler + ), + patch("asyncio.shield", side_effect=stop_on_completion), + patch( + "runpod.apps.watch.FileWatcher.changed", + side_effect=AssertionError("a stopped session must not resume watching"), + ) as changed, + ): + result, events = self._signal_session(tmp_path, monkeypatch, entrypoint) + + assert result.exit_code == 0, result.output + assert events == ["failed", "deleted"] + changed.assert_not_called() + + @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" + + @pytest.mark.skipif(sys.platform == "win32", reason="posix signals") + def test_signal_during_teardown_still_deletes(self, tmp_path, monkeypatch): + import asyncio + import os + import signal + + 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): + os.kill(os.getpid(), signal.SIGTERM) + await asyncio.sleep(0.1) + events.append(("deleted", endpoint_id)) + return True + + api.delete_endpoint.side_effect = delete + with ( + patch("runpod.apps.dev.AppsApiClient", return_value=api), + patch("runpod.apps.entrypoint.run_entrypoint"), + ): + result = _runner().invoke(cli, ["flash", "dev", str(module), "--once"]) + assert result.exit_code == 0, result.output + assert events == [("deleted", "ep-active")] + + @pytest.mark.skipif(sys.platform == "win32", reason="posix signals") + def test_signal_during_create_deletes_the_new_endpoint(self, tmp_path, monkeypatch): + import asyncio + import os + import signal + + module = tmp_path / "main.py" + module.write_text(ENTRYPOINT_SOURCE) + monkeypatch.chdir(tmp_path) + deleted = [] + api = AsyncMock() + api.list_my_endpoints.return_value = [] + + async def create(payload): + # the platform creates the endpoint before the reply arrives + os.kill(os.getpid(), signal.SIGTERM) + await asyncio.sleep(0.1) + return {"id": "ep-new"} + + async def delete(endpoint_id): + deleted.append(endpoint_id) + return True + + api.save_endpoint.side_effect = create + api.delete_endpoint.side_effect = delete + with ( + patch("runpod.apps.dev.AppsApiClient", return_value=api), + patch("runpod.apps.entrypoint.run_entrypoint"), + ): + result = _runner().invoke(cli, ["flash", "dev", str(module)]) + assert result.exit_code == 0, result.output + assert deleted == ["ep-new"] + def test_module_without_entrypoint(self, tmp_path, monkeypatch): module = tmp_path / "main.py" module.write_text(APP_SOURCE)