From 8a21f713c7f2b89384acf090c9946948fbf33433 Mon Sep 17 00:00:00 2001 From: protereus <63172654+protereus@users.noreply.github.com> Date: Wed, 22 Jul 2026 14:48:35 +0000 Subject: [PATCH] feat(runner,guard): make pod orchestration primitives fan-out safe ship_and_run/ship_and_run_detached already guaranteed pod cleanup via finally, but could still raise past it (only CancelledError was caught) -- fine for one pod, but the moment several run concurrently via asyncio.gather, one raised exception surfaces to the caller immediately while sibling tasks keep running unmonitored in the background. Add ShipAndRunResult.error (+ raise_if_error() to opt back into today's behavior) and capture generic exceptions into it instead. Add ship_and_run_many() as a thin concurrent fan-out on top, using gather(..., return_exceptions=True) so a failing job never starves or loses a sibling's result. local_roots is required (not defaulting to None-per-job) since ship_and_run's own local_root has no such default and unconditionally uploads from it -- a None default here would have made every job silently fail with an AttributeError instead of failing loudly at the call boundary. Add PodGuard.terminate(verify=True) for an optional post-terminate status re-check, opt-in and fully backward compatible with existing zero-arg call sites. Co-Authored-By: Claude Sonnet 5 --- src/runpod_lifecycle/guard.py | 37 +++++--- src/runpod_lifecycle/runner.py | 82 +++++++++++++++++- tests/test_guard.py | 60 +++++++++++++ tests/test_runner.py | 151 +++++++++++++++++++++++++++++++++ 4 files changed, 319 insertions(+), 11 deletions(-) diff --git a/src/runpod_lifecycle/guard.py b/src/runpod_lifecycle/guard.py index 65240ee..470015f 100644 --- a/src/runpod_lifecycle/guard.py +++ b/src/runpod_lifecycle/guard.py @@ -46,19 +46,36 @@ def attach(self, pod: Pod) -> None: self.pod = pod self._start_watchdog() - async def terminate(self) -> None: - """Cancel the watchdog and terminate the bound pod (idempotent).""" + async def terminate(self, *, verify: bool = False, verify_delay: float = 3.0) -> bool: + """Cancel the watchdog and terminate the bound pod (idempotent). + + When *verify* is True, waits `verify_delay` seconds then re-queries + the pod's status and returns whether it's confirmed no longer + RUNNING/PROVISIONING, instead of just trusting the terminate call + succeeded. Returns True when verified terminated (or, when + verify=False, whenever the terminate call itself didn't raise an + unhandled exception). + """ if self._watchdog is not None: self._watchdog.cancel() self._watchdog = None - if self.pod is not None: - try: - await self.pod.terminate() - except Exception as exc: - if "not found" in str(exc).lower(): - print(f"pod_already_terminated={self.pod.id}", flush=True) - else: - raise + if self.pod is None: + return True + try: + await self.pod.terminate() + except Exception as exc: + if "not found" in str(exc).lower(): + print(f"pod_already_terminated={self.pod.id}", flush=True) + else: + raise + if not verify: + return True + await asyncio.sleep(verify_delay) + status = await self.pod.status() + return (status is None) or status.get("desired_status") not in { + "RUNNING", + "PROVISIONING", + } # ------------------------------------------------------------------ # Internal watchdog diff --git a/src/runpod_lifecycle/runner.py b/src/runpod_lifecycle/runner.py index 71be8bb..a6383b5 100644 --- a/src/runpod_lifecycle/runner.py +++ b/src/runpod_lifecycle/runner.py @@ -7,7 +7,7 @@ import time from dataclasses import dataclass, field from pathlib import Path -from typing import Any, Callable, Literal +from typing import Any, Callable, Literal, Sequence from .config import RunPodConfig from .guard import PodGuard @@ -78,6 +78,12 @@ class ShipAndRunResult: breach_log: list[dict] = field(default_factory=list) terminated: bool = False upload_info: dict[str, Any] = field(default_factory=dict) + error: BaseException | None = None + + def raise_if_error(self) -> None: + """Re-raise the captured exception, if any (opt back into raising).""" + if self.error is not None: + raise self.error # --------------------------------------------------------------------------- @@ -207,6 +213,13 @@ async def ship_and_run( logger.info("ship_and_run cancelled, returning 130") return result + except Exception as exc: + logger.warning( + "ship_and_run failed pod=%s error=%s", getattr(pod, "id", None), exc + ) + result.error = exc + return result + finally: if terminate_after_exec: await guard.terminate() @@ -398,6 +411,15 @@ async def ship_and_run_detached( result.returncode = 130 return result + except Exception as exc: + logger.warning( + "ship_and_run_detached failed pod=%s error=%s", + getattr(active_pod, "id", None), + exc, + ) + result.error = exc + return result + finally: if terminate_after_exec and active_pod is not None: if guard is not None: @@ -409,3 +431,61 @@ async def ship_and_run_detached( result.pod = active_pod if guard is not None: result.breach_log = list(guard.breach_log) + + +async def ship_and_run_many( + config: RunPodConfig, + remote_scripts: Sequence[str], + *, + local_roots: Sequence[Path], + remote_root: str = "/workspace", + exclude: set[str] | None = None, + name_prefix: str = "pod", + **shared_kwargs: Any, +) -> list[ShipAndRunResult]: + """Launch ``len(remote_scripts)`` pods concurrently, one script per pod. + + Each pod's full lifecycle is independent via `ship_and_run`'s own + guarantees (see its docstring) — a failure or exception in one pod's + run never prevents another pod from completing normally or being torn + down, and never prevents this function from returning a result for + every pod. Results are returned in the same order as `remote_scripts`; + check `.error` on each (or call `.raise_if_error()`) to see whether + that pod's run failed. + + *local_roots* is required and must be the same length as + *remote_scripts* — one local directory to upload per pod (repeat the + same path across entries if every pod should get the same payload). + It's required rather than defaulting to "no upload" because + `ship_and_run` itself has no such default: its own *local_root* is a + required, non-optional `Path` that it unconditionally uploads from, so + a `None` here would only surface as an `AttributeError` captured deep + in every job's `.error` instead of failing loudly at the call site. + + Does not support coordination *between* jobs (e.g. one pod needing to + wait on and consume another's output) — that's caller-specific; write + a custom per-pod coroutine closing over a shared `asyncio.Future` (or + similar) for that instead of using this helper. + """ + if len(local_roots) != len(remote_scripts): + raise ValueError("local_roots must be the same length as remote_scripts") + + results = await asyncio.gather( + *( + ship_and_run( + config, + script, + local_root=root, + remote_root=remote_root, + exclude=exclude or set(), + name_prefix=f"{name_prefix}-{i}", + **shared_kwargs, + ) + for i, (script, root) in enumerate(zip(remote_scripts, local_roots)) + ), + return_exceptions=True, + ) + return [ + r if isinstance(r, ShipAndRunResult) else ShipAndRunResult(returncode=-1, error=r) + for r in results + ] diff --git a/tests/test_guard.py b/tests/test_guard.py index 59348ea..09de8b4 100644 --- a/tests/test_guard.py +++ b/tests/test_guard.py @@ -150,6 +150,66 @@ async def test_podguard_terminate_re_raises_other_errors() -> None: await guard.terminate() +# --------------------------------------------------------------------------- +# PodGuard.terminate — verify option +# --------------------------------------------------------------------------- + +@pytest.mark.asyncio +async def test_podguard_terminate_verify_true_confirms_terminated() -> None: + """verify=True with a non-active post-terminate status returns True.""" + mock_pod = MagicMock() + mock_pod.id = "pod-test-verify-1" + mock_pod.terminate = AsyncMock() + mock_pod.status = AsyncMock(return_value={"desired_status": "TERMINATED"}) + + guard = PodGuard(name_prefix="test", auto_terminate=True) + guard.pod = mock_pod + guard._watchdog = None + + with patch.object(asyncio, "sleep", new_callable=AsyncMock): + result = await guard.terminate(verify=True) + + assert result is True + mock_pod.status.assert_called_once() + + +@pytest.mark.asyncio +async def test_podguard_terminate_verify_true_still_running_returns_false() -> None: + """verify=True with status still RUNNING/PROVISIONING returns False.""" + mock_pod = MagicMock() + mock_pod.id = "pod-test-verify-2" + mock_pod.terminate = AsyncMock() + mock_pod.status = AsyncMock(return_value={"desired_status": "RUNNING"}) + + guard = PodGuard(name_prefix="test", auto_terminate=True) + guard.pod = mock_pod + guard._watchdog = None + + with patch.object(asyncio, "sleep", new_callable=AsyncMock): + result = await guard.terminate(verify=True) + + assert result is False + mock_pod.status.assert_called_once() + + +@pytest.mark.asyncio +async def test_podguard_terminate_verify_false_skips_status_check() -> None: + """verify=False (default) returns True without ever calling pod.status().""" + mock_pod = MagicMock() + mock_pod.id = "pod-test-verify-3" + mock_pod.terminate = AsyncMock() + mock_pod.status = AsyncMock(return_value={"desired_status": "RUNNING"}) + + guard = PodGuard(name_prefix="test", auto_terminate=True) + guard.pod = mock_pod + guard._watchdog = None + + result = await guard.terminate() + + assert result is True + mock_pod.status.assert_not_called() + + # --------------------------------------------------------------------------- # guard_factory injection # --------------------------------------------------------------------------- diff --git a/tests/test_runner.py b/tests/test_runner.py index 0cec6d3..c22c444 100644 --- a/tests/test_runner.py +++ b/tests/test_runner.py @@ -23,6 +23,7 @@ ShipAndRunResult, ship_and_run, ship_and_run_detached, + ship_and_run_many, _parse_detached_exit, ) @@ -229,6 +230,156 @@ async def slow_exec(cmd: str, timeout: int = 600) -> tuple[int, str, str]: assert result.returncode == 130 +# --------------------------------------------------------------------------- +# ship_and_run — generic exception is captured, not raised +# --------------------------------------------------------------------------- + +@pytest.mark.asyncio +async def test_ship_and_run_captures_exception_instead_of_raising(tmp_path: Path) -> None: + """A generic Exception during exec_ssh is captured on result.error, not raised.""" + config = RunPodConfig(api_key="***") + mock_pod = _make_mock_pod("pod-sar-err") + local_root = tmp_path / "local_err" + local_root.mkdir() + + boom = RuntimeError("ssh exec blew up") + mock_pod.exec_ssh = AsyncMock(side_effect=boom) + + with patch("runpod_lifecycle.runner._launch_pod", new_callable=AsyncMock) as mock_launch: + mock_launch.return_value = mock_pod + + result = await ship_and_run( + config, + "echo ok", + local_root=local_root, + remote_root="/tmp/remote", + exclude=set(), + timeout=30, + ) + + assert result.error is boom + assert result.terminated is True + mock_pod.terminate.assert_called() + with pytest.raises(RuntimeError, match="ssh exec blew up"): + result.raise_if_error() + + +@pytest.mark.asyncio +async def test_ship_and_run_detached_captures_exception_instead_of_raising(tmp_path: Path) -> None: + """A generic Exception in ship_and_run_detached is captured, not raised.""" + config = RunPodConfig(api_key="***") + mock_pod = _make_mock_pod("pod-det-err") + local_root = tmp_path / "local_det_err" + local_root.mkdir() + + boom = RuntimeError("detached exec blew up") + mock_pod.exec_ssh = AsyncMock(side_effect=boom) + + with patch("runpod_lifecycle.runner._launch_pod", new_callable=AsyncMock) as mock_launch: + mock_launch.return_value = mock_pod + + result = await ship_and_run_detached( + config, + "echo ok", + local_root=local_root, + remote_root="/tmp/remote", + exclude=set(), + timeout=30, + terminate_after_exec=True, + poll_interval=1, + ) + + assert result.error is boom + assert result.terminated is True + mock_pod.terminate.assert_called() + with pytest.raises(RuntimeError, match="detached exec blew up"): + result.raise_if_error() + + +# --------------------------------------------------------------------------- +# ship_and_run_many +# --------------------------------------------------------------------------- + +@pytest.mark.asyncio +async def test_ship_and_run_many_all_succeed(tmp_path: Path) -> None: + """All jobs succeed: N results returned in order, none with .error set.""" + config = RunPodConfig(api_key="***") + + async def fake_ship_and_run(config, script, **kwargs): + return ShipAndRunResult(returncode=0, stdout=script) + + with patch("runpod_lifecycle.runner.ship_and_run", side_effect=fake_ship_and_run): + results = await ship_and_run_many( + config, + ["script-a", "script-b", "script-c"], + local_roots=[tmp_path, tmp_path, tmp_path], + ) + + assert len(results) == 3 + assert [r.stdout for r in results] == ["script-a", "script-b", "script-c"] + assert all(r.error is None for r in results) + + +@pytest.mark.asyncio +async def test_ship_and_run_many_one_job_fails(tmp_path: Path) -> None: + """One job's ship_and_run raises; others still succeed and are captured in order.""" + config = RunPodConfig(api_key="***") + boom = RuntimeError("job b blew up") + + async def fake_ship_and_run(config, script, **kwargs): + if script == "script-b": + raise boom + return ShipAndRunResult(returncode=0, stdout=script) + + with patch("runpod_lifecycle.runner.ship_and_run", side_effect=fake_ship_and_run): + results = await ship_and_run_many( + config, + ["script-a", "script-b", "script-c"], + local_roots=[tmp_path, tmp_path, tmp_path], + ) + + assert len(results) == 3 + assert results[0].stdout == "script-a" + assert results[0].error is None + assert results[1].error is boom + assert results[2].stdout == "script-c" + assert results[2].error is None + + +@pytest.mark.asyncio +async def test_ship_and_run_many_local_roots_length_mismatch() -> None: + """A local_roots length mismatch raises ValueError.""" + config = RunPodConfig(api_key="***") + + with pytest.raises(ValueError, match="local_roots must be the same length"): + await ship_and_run_many( + config, + ["script-a", "script-b"], + local_roots=[None], + ) + + +@pytest.mark.asyncio +async def test_ship_and_run_many_requires_local_roots() -> None: + """Regression test: local_roots must be required, not default to None-per-job. + + ship_and_run's own local_root is a required, non-optional Path that it + unconditionally uploads from (no "skip upload" branch, unlike + ship_and_run_detached) — a None default here would silently make every + job fail with an AttributeError deep inside the upload step, captured + into .error rather than failing loudly at the call site. Omitting + local_roots must raise immediately instead. + """ + config = RunPodConfig(api_key="***") + + async def fake_ship_and_run(config, script, **kwargs): + return ShipAndRunResult(returncode=0, stdout=script) + + with patch("runpod_lifecycle.runner.ship_and_run", side_effect=fake_ship_and_run): + with pytest.raises(TypeError, match="local_roots"): + await ship_and_run_many(config, ["script-a", "script-b"]) # type: ignore[call-arg] + + # --------------------------------------------------------------------------- # ship_and_run_detached — provision path # ---------------------------------------------------------------------------