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

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.
 

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
 
 

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)