fix(cli): delete dev endpoints when rp flash dev is stopped by a signal

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 @@
                     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 @@
             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 +443,64 @@
             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))
         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))

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)