Make ship_and_run primitives safe for concurrent pod fan-out

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

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
+    ]

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

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