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 # ---------------------------------------------------------------------------