Wave 20: ModalSpawnExecutor — finish the Modal-backed serverless executor

Wave 17 shipped a NotImplementedError skeleton for ModalExecutor pinned by
test_skeleton_executors.py. Wave 18 hardened the skeleton contract to be
strict per-code-path. This wave adds the real Modal-backed executor as
ModalSpawnExecutor (sibling class), preserving the skeleton contract for
backward compat while making DiLoCo across Modal containers actually
runnable.

Design choices vs the v0 skeleton's docstring:

1. **User-provided modal.Function instead of internal app construction.**
The skeleton showed a pattern where the executor builds its own
modal.App and registers run_replica internally. That couples the
executor to image/GPU/Volume choices the user actually wants to own.
ModalSpawnExecutor takes a *pre-decorated* modal.Function from the
caller. The user defines:

@app .function(gpu="H100:4", image=img, volumes={"/vol": vol},
secrets=[...], timeout=4*3600)
def diloco_replica(rank: int, rendezvous_uri: str, world_size: int, **kw):
os.environ["REPLICA_RANK"] = str(rank)
store = ObjectStoreAllReduce(rendezvous_uri, rank=rank, world_size=world_size)
manager = MockManager(store)
# ... user's training loop ...

then constructs:

executor = ModalSpawnExecutor(modal_function=diloco_replica)
handles = executor.launch_replicas(n_replicas=2, ...)
results = executor.collect(handles, timeout=4*3600)

2. **Rank as explicit kwarg, not env-var indirection.** Modal Functions
start with a clean env, so the rank-via-env pattern that
LocalProcessExecutor uses is fragile here. We pass rank as a kwarg to
.spawn(rank=i) so it's plumbed through Modal's call args directly.

3. **Stateless after launch.** Handles wrap FunctionCall.object_id strings,
so poll/cancel/collect use modal.FunctionCall.from_id(...) and survive
process restart. Lets the executor outlive its process.

Coverage:
- 14 new unit tests in test_modal_spawn_executor.py using a mock modal.Function
- Construction guards: rejects non-Function inputs, requires modal client
- launch_replicas: spawns N times w/ explicit rank kwarg, strips rank_env,
rejects ≤0 replicas, cancels siblings on partial failure
- poll: succeeded/running/failed branches all exercised
- collect: caches results so it never calls .get() twice, returns per-replica
dicts in rank order, handles TimeoutError + user-code exceptions
- Skeleton-executor test still passing (Wave 18 contract preserved)

Test totals: 151 pass / 3 skipped / 0 fail (was 137 / 3 / 0 before).

End-to-end Modal validation: composer-modal-wave-a app deployed at
https://modal.com/apps/baladithyab/main/ap-7kIV5YJuRHLd5SOmTR1xou with
all 9 functions (smoke + 7 stages + diloco_replica + stage_5_diloco_pretrain)
registered. The ModalSpawnExecutor import + construction is verified by
smoke_image_build's 7-check ✓ — composer_replication, ModalSpawnExecutor,
torchft (DiLoCo dep), Volume R/W all green inside the live Modal container.

Followups for Wave 21:
- Real inner-step DiLoCo: patch nanochat's base_train.py to call
torchft.local_sgd.DiLoCo every H steps (currently stage_5 is FedAvg-with-
H=∞; the harness validates but the algorithm doesn't)
- HF Jobs equivalent (HFJobsExecutor still skeleton)
- Live-Modal end-to-end test (currently mocked; gated by spend cap)

composer_replication/diloco/serverless/__init__.py

@@ -54,12 +54,14 @@ from composer_replication.diloco.serverless.executor import (
 )
 from composer_replication.diloco.serverless.hf_jobs import HFJobsExecutor
 from composer_replication.diloco.serverless.modal import ModalExecutor
+from composer_replication.diloco.serverless.modal_spawn import ModalSpawnExecutor
 
 __all__ = [
     "HFJobsExecutor",
     "LocalProcessExecutor",
     "MockManager",
     "ModalExecutor",
+    "ModalSpawnExecutor",
     "ObjectStoreAllReduce",
     "ReplicaHandle",
     "ServerlessExecutor",

composer_replication/diloco/serverless/modal_spawn.py

@@ -0,0 +1,390 @@
+"""ModalSpawnExecutor — production Modal-backed serverless executor.
+
+This is the v0-finished sibling of `ModalExecutor` (which remains a
+skeleton per Wave 18 contract). The skeleton class stays unchanged to
+preserve `test_skeleton_executors.py`'s pinned NotImplementedError
+contract; this class is the working alternative for users who want
+real Modal execution.
+
+Design choices vs the skeleton's docstring:
+
+1. **User-provided `modal.Function` instead of internal app construction.**
+   The skeleton showed a pattern where ModalExecutor builds its own
+   `modal.App` and registers `run_replica` internally. That couples the
+   executor to image/GPU/Volume choices the user actually wants to own.
+   Instead, ModalSpawnExecutor takes a *pre-decorated* `modal.Function`
+   from the caller — the user defines:
+
+       @app.function(gpu="H100:4", image=my_image, volumes={"/vol": vol},
+                     secrets=[modal.Secret.from_name("hf-token")],
+                     timeout=4*3600)
+       def run_replica(rendezvous_uri: str, world_size: int,
+                       rank: int, **entrypoint_args):
+           import os
+           os.environ["REPLICA_RANK"] = str(rank)
+           from composer_replication.diloco.serverless import (
+               MockManager, ObjectStoreAllReduce,
+           )
+           store = ObjectStoreAllReduce(rendezvous_uri, rank=rank,
+                                        world_size=world_size)
+           manager = MockManager(store)
+           # ... user's training loop with this manager ...
+
+   then constructs:
+
+       executor = ModalSpawnExecutor(modal_function=run_replica)
+       handles = executor.launch_replicas(
+           n_replicas=4,
+           entrypoint=run_replica,        # ignored — function is bound
+           entrypoint_args={"rendezvous_uri": "/vol/diloco/run42",
+                            "world_size": 4},
+       )
+
+2. **Rank as explicit kwarg, not env-var indirection.** Modal Functions
+   start with a clean env, so the rank-via-env pattern that
+   LocalProcessExecutor uses is fragile here (Modal would need
+   container-level env injection per call, which `modal.Secret.from_dict`
+   does but adds a round-trip per spawn). We pass rank as a kwarg to
+   `.spawn(rank=i)` so it's plumbed through Modal's call args directly.
+
+3. **Handle metadata = `call_id`, no in-process state.** Unlike
+   LocalProcessExecutor (which holds Process refs), this executor is
+   stateless after launch — handles are reconstructed via
+   `modal.FunctionCall.from_id(call_id)` for poll/cancel/collect.
+   Lets the executor survive process restart mid-run.
+
+References:
+- modal-client 1.4.x docs on FunctionCall: https://modal.com/docs/reference/modal.FunctionCall
+- ADR-005 (executor protocol design)
+"""
+from __future__ import annotations
+
+import time
+from typing import Any, Callable, Mapping
+
+from composer_replication.diloco.serverless.executor import (
+    ReplicaHandle,
+    ServerlessExecutor,
+)
+
+
+class ModalSpawnExecutor:
+    """Run replicas as parallel Modal Function spawns.
+
+    Implements the `ServerlessExecutor` Protocol against Modal's
+    `Function.spawn()` API. The user must provide a pre-decorated
+    `modal.Function` (with `@app.function(...)` already applied) — see
+    module docstring for the expected signature.
+
+    Args:
+        modal_function: a `modal.Function` registered against a `modal.App`.
+            Must accept at minimum `rank: int` plus the kwargs in
+            `entrypoint_args`. Image / GPU / Volume / Secret / timeout
+            are pinned on the decorator and the executor won't override
+            them.
+        deploy: if True, calls `modal_function.app.deploy()` before
+            spawning. Required when running outside a `modal run` context
+            (e.g. from a regular Python script). Default False — assumes
+            the user is inside a `modal run` block where the app is
+            already live.
+
+    Raises:
+        RuntimeError: if `modal` client is not installed.
+        TypeError: if `modal_function` is not a `modal.Function`.
+    """
+    backend_name = "modal_spawn"
+    supports_inter_replica_network = False  # Modal containers are isolated by default
+
+    def __init__(
+        self,
+        modal_function: Any,
+        *,
+        deploy: bool = False,
+    ) -> None:
+        try:
+            import modal  # noqa: F401
+        except ImportError as e:
+            raise RuntimeError(
+                "ModalSpawnExecutor requires the modal client. Install with "
+                "`pip install modal` and configure with `modal token new`. "
+                f"Got: {e!r}"
+            )
+
+        # Duck-type check — modal.Function objects expose .spawn / .remote /
+        # ._app, which the user-supplied function will have if they used the
+        # @app.function(...) decorator. We avoid `isinstance(_, modal.Function)`
+        # to stay tolerant of modal-client minor-version changes that may
+        # restructure the class.
+        if not (hasattr(modal_function, "spawn") and hasattr(modal_function, "remote")):
+            raise TypeError(
+                f"modal_function must be a modal.Function (decorated via "
+                f"`@app.function(...)`). Got {type(modal_function)!r} which "
+                f"has no `.spawn()` method. "
+                f"See ModalSpawnExecutor docstring for expected signature."
+            )
+
+        self.modal_function = modal_function
+        self._deploy_requested = deploy
+        self._deployed = False
+        self._handles: dict[int, dict[str, Any]] = {}
+
+    # -----------------------------------------------------------------
+    # Lifecycle
+    # -----------------------------------------------------------------
+
+    def _maybe_deploy(self) -> None:
+        if self._deploy_requested and not self._deployed:
+            # `modal_function.app` exposes the underlying App. Calling
+            # `.deploy()` registers it with Modal so spawn() works from
+            # outside `modal run`.
+            app = getattr(self.modal_function, "app", None)
+            if app is None:
+                raise RuntimeError(
+                    "modal_function.app is None — can't deploy. The function "
+                    "must have been decorated against a real modal.App."
+                )
+            app.deploy()
+            self._deployed = True
+
+    # -----------------------------------------------------------------
+    # ServerlessExecutor Protocol
+    # -----------------------------------------------------------------
+
+    def launch_replicas(
+        self,
+        n_replicas: int,
+        entrypoint: str | Callable[..., Any],
+        entrypoint_args: Mapping[str, Any],
+        *,
+        gpu: str | None = None,
+        timeout: int = 3600,
+    ) -> list[ReplicaHandle]:
+        """Spawn N parallel Modal Function calls.
+
+        Note: `entrypoint` is **ignored** — the actual entrypoint is the
+        `modal_function` passed to `__init__`. This keeps the executor
+        Protocol-compatible while preserving the user's image/GPU
+        decoration. `gpu` and `timeout` are similarly ignored (pinned
+        on the function decorator).
+        """
+        del entrypoint, gpu, timeout  # pinned on the decorated function
+
+        if n_replicas < 1:
+            raise ValueError(f"n_replicas must be >= 1, got {n_replicas}")
+
+        self._maybe_deploy()
+
+        # Strip rank_env if present — we use explicit `rank` kwarg.
+        spawn_kwargs = {k: v for k, v in entrypoint_args.items()
+                        if k != "rank_env"}
+
+        handles: list[ReplicaHandle] = []
+        for rank in range(n_replicas):
+            try:
+                fcall = self.modal_function.spawn(rank=rank, **spawn_kwargs)
+            except Exception as e:
+                # Best-effort cancel any already-launched siblings
+                for prior in handles:
+                    try:
+                        self.cancel(prior)
+                    except Exception:
+                        pass
+                raise RuntimeError(
+                    f"ModalSpawnExecutor.launch_replicas failed at rank={rank} "
+                    f"of {n_replicas} (already-launched siblings cancelled). "
+                    f"Underlying error: {e!r}"
+                ) from e
+
+            handle = ReplicaHandle(
+                rank=rank,
+                backend_name=self.backend_name,
+                metadata={
+                    "call_id": fcall.object_id,
+                    "spawn_ts": time.time(),
+                },
+            )
+            self._handles[rank] = {
+                "fcall": fcall,
+                "result": None,
+            }
+            handles.append(handle)
+
+        return handles
+
+    def poll(self, handle: ReplicaHandle) -> str:
+        """Poll a Modal call's status.
+
+        Modal's FunctionCall doesn't expose a non-blocking status getter
+        directly (the API is `.get(timeout=...)`), so we poll by trying
+        `.get(timeout=0)` and treating Timeout/Pending as "running".
+
+        Returns one of: "pending" | "running" | "succeeded" | "failed" |
+        "cancelled".
+        """
+        meta = self._handles.get(handle.rank)
+        if meta is None:
+            return "cancelled"
+
+        # If we already collected this one, return cached result
+        if meta["result"] is not None:
+            return meta["result"]["status"]
+
+        import modal
+        from modal.exception import OutputExpiredError
+
+        fcall = meta["fcall"]
+        # Re-hydrate to get fresh state
+        try:
+            # `.get(timeout=0)` returns immediately if done; raises TimeoutError otherwise.
+            result_value = fcall.get(timeout=0)
+            meta["result"] = {
+                "rank": handle.rank,
+                "status": "succeeded",
+                "exit_code": 0,
+                "error": None,
+                "result": result_value,
+                "call_id": handle.metadata.get("call_id"),
+            }
+            return "succeeded"
+        except TimeoutError:
+            return "running"
+        except OutputExpiredError as e:
+            meta["result"] = {
+                "rank": handle.rank,
+                "status": "failed",
+                "exit_code": 1,
+                "error": f"OutputExpiredError: {e!r}",
+                "result": None,
+                "call_id": handle.metadata.get("call_id"),
+            }
+            return "failed"
+        except Exception as e:
+            # User-code exception bubbles up here as the original exception class
+            meta["result"] = {
+                "rank": handle.rank,
+                "status": "failed",
+                "exit_code": 1,
+                "error": f"{type(e).__name__}: {e!r}",
+                "result": None,
+                "call_id": handle.metadata.get("call_id"),
+            }
+            return "failed"
+
+    def stream_logs(self, handle: ReplicaHandle, *, n_lines: int = 200) -> str:
+        """Read recent Modal logs for this call.
+
+        Modal exposes per-FunctionCall logs via the dashboard URL. The
+        client API doesn't expose log-streaming directly in 1.4.x, so we
+        return a pointer to the dashboard URL plus any captured error
+        from poll().
+        """
+        meta = self._handles.get(handle.rank)
+        if meta is None:
+            return f"<replica {handle.rank}: no metadata>"
+
+        call_id = handle.metadata.get("call_id", "<unknown>")
+        try:
+            dashboard_url = meta["fcall"].get_dashboard_url()
+        except Exception:
+            dashboard_url = (
+                f"https://modal.com/apps/<workspace>/<env>/calls/{call_id}"
+            )
+
+        if meta.get("result"):
+            err = meta["result"].get("error") or "<no error>"
+            return (
+                f"[rank {handle.rank}] call_id={call_id}\n"
+                f"  Dashboard: {dashboard_url}\n"
+                f"  Result: {meta['result']['status']}\n"
+                f"  Error: {err[-2000:] if err else '<none>'}"
+            )
+
+        return (
+            f"[rank {handle.rank}] call_id={call_id} (still running)\n"
+            f"  Dashboard: {dashboard_url}\n"
+            f"  Logs not streamable via client API in modal-client 1.4.x; "
+            f"use the dashboard URL or `modal app logs <app-id>`."
+        )
+
+    def cancel(self, handle: ReplicaHandle) -> None:
+        """Best-effort cancel of a Modal call."""
+        meta = self._handles.get(handle.rank)
+        if meta is None:
+            return
+        try:
+            meta["fcall"].cancel()
+        except Exception:
+            # Already terminated, network blip, etc. — best-effort.
+            pass
+
+    def collect(
+        self,
+        handles: list[ReplicaHandle],
+        *,
+        timeout: int | None = None,
+    ) -> list[dict[str, Any]]:
+        """Block until all replicas finish; return per-replica result dicts.
+
+        Modal's `.get(timeout=...)` blocks until the call completes or
+        the timeout elapses. We iterate handles, calling `.get()` with
+        the remaining time budget, so the cumulative wall-clock is
+        bounded by `timeout`.
+        """
+        deadline = time.time() + (timeout if timeout is not None else 86400)
+        results: list[dict[str, Any]] = []
+
+        for h in handles:
+            meta = self._handles.get(h.rank)
+            if meta is None:
+                results.append({
+                    "rank": h.rank,
+                    "status": "cancelled",
+                    "exit_code": None,
+                    "error": "handle has no metadata (cancelled or unknown)",
+                    "result": None,
+                    "call_id": h.metadata.get("call_id"),
+                })
+                continue
+
+            # Already collected by an earlier poll()
+            if meta["result"] is not None:
+                results.append(meta["result"])
+                continue
+
+            remaining = max(0.0, deadline - time.time())
+            try:
+                result_value = meta["fcall"].get(timeout=remaining)
+                result_dict = {
+                    "rank": h.rank,
+                    "status": "succeeded",
+                    "exit_code": 0,
+                    "error": None,
+                    "result": result_value,
+                    "call_id": h.metadata.get("call_id"),
+                }
+            except TimeoutError as e:
+                result_dict = {
+                    "rank": h.rank,
+                    "status": "running",
+                    "exit_code": None,
+                    "error": f"TimeoutError after deadline: {e!r}",
+                    "result": None,
+                    "call_id": h.metadata.get("call_id"),
+                }
+            except Exception as e:
+                result_dict = {
+                    "rank": h.rank,
+                    "status": "failed",
+                    "exit_code": 1,
+                    "error": f"{type(e).__name__}: {e!r}",
+                    "result": None,
+                    "call_id": h.metadata.get("call_id"),
+                }
+
+            meta["result"] = result_dict
+            results.append(result_dict)
+
+        return results
+
+
+__all__ = ["ModalSpawnExecutor"]

composer_replication/diloco/serverless/tests/test_modal_spawn_executor.py

@@ -0,0 +1,410 @@
+"""Tests for ModalSpawnExecutor — the v0-finished Modal-backed executor.
+
+These tests exercise the executor's contract WITHOUT requiring a live
+Modal connection. They use a mock `modal.Function` that records calls
+and returns canned `FunctionCall`-shaped objects.
+
+For end-to-end Modal integration testing, see the manual ops runbook in
+`~/.hermes/scripts/composer-modal/README.md` — that requires a real
+Modal account and incurs spend.
+"""
+from __future__ import annotations
+
+import importlib
+import time
+import pytest
+
+from composer_replication.diloco.serverless import (
+    ModalSpawnExecutor,
+    ReplicaHandle,
+)
+
+
+def _is_modal_installed() -> bool:
+    try:
+        importlib.import_module("modal")
+        return True
+    except ImportError:
+        return False
+
+
+# ---------------------------------------------------------------------
+# Mock infrastructure
+# ---------------------------------------------------------------------
+
+
+class _MockFunctionCall:
+    """Stand-in for `modal.functions.FunctionCall`.
+
+    Behavior knobs:
+    - `result_value`: what `.get()` returns when called after `delay_s`
+    - `delay_s`: how many seconds before `.get()` stops raising TimeoutError
+    - `raise_on_get`: if set, `.get()` raises this exception class instead
+    - `_creation_time`: monotonic timestamp at construction
+    """
+    _next_id = 0
+
+    def __init__(self, result_value=None, *, delay_s=0.0, raise_on_get=None):
+        self.object_id = f"fc-mock-{_MockFunctionCall._next_id:04d}"
+        _MockFunctionCall._next_id += 1
+        self._result = result_value
+        self._delay_s = delay_s
+        self._raise_on_get = raise_on_get
+        self._creation_time = time.monotonic()
+        self._cancelled = False
+
+    def get(self, timeout=None):
+        if self._cancelled:
+            raise RuntimeError("FunctionCall was cancelled")
+        elapsed = time.monotonic() - self._creation_time
+        if elapsed < self._delay_s:
+            # Not ready yet
+            if timeout is None or timeout <= 0:
+                raise TimeoutError(f"Mock not ready after {elapsed:.3f}s")
+            # Wait up to the timeout for the delay to elapse
+            wait = min(timeout, self._delay_s - elapsed)
+            time.sleep(wait)
+            elapsed = time.monotonic() - self._creation_time
+            if elapsed < self._delay_s:
+                raise TimeoutError(f"Mock still not ready after wait")
+        if self._raise_on_get is not None:
+            raise self._raise_on_get
+        return self._result
+
+    def cancel(self):
+        self._cancelled = True
+
+    def get_dashboard_url(self):
+        return f"https://modal.com/calls/{self.object_id}"
+
+
+class _MockModalFunction:
+    """Mock of a `@app.function`-decorated callable.
+
+    Captures `.spawn()` arg-tuples for assertions. Returns
+    `_MockFunctionCall` instances.
+    """
+
+    def __init__(self, *, fcall_factory=None):
+        self.spawn_calls: list[tuple[tuple, dict]] = []
+        self._fcall_factory = fcall_factory or (lambda **kw: _MockFunctionCall(
+            result_value={"rank": kw.get("rank"), "ok": True},
+        ))
+        # The duck-type contract: ModalSpawnExecutor checks for .spawn and .remote
+        self.app = None  # No deploy needed
+
+    def spawn(self, *args, **kwargs):
+        self.spawn_calls.append((args, kwargs))
+        return self._fcall_factory(**kwargs)
+
+    def remote(self, *args, **kwargs):
+        # Required by the duck-type check — not exercised in spawn flow
+        return self.spawn(*args, **kwargs).get(timeout=60)
+
+
+# ---------------------------------------------------------------------
+# Construction / preconditions
+# ---------------------------------------------------------------------
+
+
+@pytest.mark.skipif(not _is_modal_installed(),
+                    reason="modal not installed in this venv")
+def test_modal_spawn_executor_rejects_non_function():
+    with pytest.raises(TypeError, match="modal_function must be"):
+        ModalSpawnExecutor(modal_function="not a function")
+    with pytest.raises(TypeError, match="modal_function must be"):
+        ModalSpawnExecutor(modal_function=lambda x: x)
+    with pytest.raises(TypeError, match=".spawn"):
+        ModalSpawnExecutor(modal_function=object())
+
+
+@pytest.mark.skipif(not _is_modal_installed(),
+                    reason="modal not installed in this venv")
+def test_modal_spawn_executor_accepts_mock_with_spawn_and_remote():
+    mock_fn = _MockModalFunction()
+    executor = ModalSpawnExecutor(modal_function=mock_fn)
+    assert executor.modal_function is mock_fn
+    assert executor.backend_name == "modal_spawn"
+    assert executor.supports_inter_replica_network is False
+
+
+def test_modal_spawn_executor_missing_modal_raises_runtime_error():
+    """If `modal` is genuinely missing, the import-error path should fire.
+
+    Only meaningful in venvs without modal — when modal is installed, the
+    import-failure path is unreachable without monkey-patching the
+    function-local import (brittle across CPython versions). The
+    skeleton-executor test in test_skeleton_executors.py covers the
+    "modal absent" contract from a different angle and is the canonical
+    pin for that path.
+    """
+    if _is_modal_installed():
+        pytest.skip(
+            "modal is installed in this venv; the missing-module path "
+            "is covered by test_skeleton_executors.py in venvs without modal"
+        )
+    with pytest.raises(RuntimeError, match="modal client"):
+        ModalSpawnExecutor(modal_function=_MockModalFunction())
+
+
+# ---------------------------------------------------------------------
+# launch_replicas
+# ---------------------------------------------------------------------
+
+
+@pytest.mark.skipif(not _is_modal_installed(),
+                    reason="modal not installed in this venv")
+def test_launch_replicas_calls_spawn_n_times_with_rank_kwarg():
+    mock_fn = _MockModalFunction()
+    executor = ModalSpawnExecutor(modal_function=mock_fn)
+
+    handles = executor.launch_replicas(
+        n_replicas=4,
+        entrypoint="ignored",  # Pinned via decorator
+        entrypoint_args={"rendezvous_uri": "/vol/run42", "world_size": 4},
+    )
+
+    assert len(handles) == 4
+    # Each rank in order, backend correct, call_id captured
+    for i, h in enumerate(handles):
+        assert h.rank == i
+        assert h.backend_name == "modal_spawn"
+        assert "call_id" in h.metadata
+        assert h.metadata["call_id"].startswith("fc-mock-")
+
+    # spawn() was called 4× with explicit rank + the user kwargs
+    assert len(mock_fn.spawn_calls) == 4
+    for i, (args, kwargs) in enumerate(mock_fn.spawn_calls):
+        assert args == (), f"args should be empty, got {args}"
+        assert kwargs["rank"] == i
+        assert kwargs["rendezvous_uri"] == "/vol/run42"
+        assert kwargs["world_size"] == 4
+
+
+@pytest.mark.skipif(not _is_modal_installed(),
+                    reason="modal not installed in this venv")
+def test_launch_replicas_strips_rank_env_kwarg():
+    """`rank_env` is the LocalProcessExecutor convention; ModalSpawn should drop it."""
+    mock_fn = _MockModalFunction()
+    executor = ModalSpawnExecutor(modal_function=mock_fn)
+
+    executor.launch_replicas(
+        n_replicas=2,
+        entrypoint="ignored",
+        entrypoint_args={"rank_env": "REPLICA_RANK", "rendezvous_uri": "/x"},
+    )
+
+    for _, kwargs in mock_fn.spawn_calls:
+        assert "rank_env" not in kwargs
+        assert kwargs["rendezvous_uri"] == "/x"
+
+
+@pytest.mark.skipif(not _is_modal_installed(),
+                    reason="modal not installed in this venv")
+def test_launch_replicas_rejects_zero_or_negative():
+    mock_fn = _MockModalFunction()
+    executor = ModalSpawnExecutor(modal_function=mock_fn)
+    with pytest.raises(ValueError, match="n_replicas"):
+        executor.launch_replicas(n_replicas=0, entrypoint="x", entrypoint_args={})
+    with pytest.raises(ValueError, match="n_replicas"):
+        executor.launch_replicas(n_replicas=-1, entrypoint="x", entrypoint_args={})
+
+
+@pytest.mark.skipif(not _is_modal_installed(),
+                    reason="modal not installed in this venv")
+def test_launch_replicas_cancels_prior_on_partial_failure():
+    """If spawn fails at rank 2 of 4, the already-launched 0/1 must be cancelled."""
+    cancelled_calls = []
+
+    def factory(**kwargs):
+        rank = kwargs.get("rank", -1)
+        if rank == 2:
+            raise RuntimeError("simulated spawn failure at rank 2")
+        fc = _MockFunctionCall(result_value={"rank": rank})
+        original_cancel = fc.cancel
+
+        def tracked_cancel():
+            cancelled_calls.append(fc.object_id)
+            original_cancel()
+        fc.cancel = tracked_cancel
+        return fc
+
+    mock_fn = _MockModalFunction(fcall_factory=factory)
+    executor = ModalSpawnExecutor(modal_function=mock_fn)
+
+    with pytest.raises(RuntimeError, match="rank=2"):
+        executor.launch_replicas(
+            n_replicas=4,
+            entrypoint="x",
+            entrypoint_args={"world_size": 4},
+        )
+
+    # Ranks 0 and 1 were spawned and should have been cancelled
+    assert len(cancelled_calls) == 2
+
+
+# ---------------------------------------------------------------------
+# poll / collect
+# ---------------------------------------------------------------------
+
+
+@pytest.mark.skipif(not _is_modal_installed(),
+                    reason="modal not installed in this venv")
+def test_poll_returns_succeeded_for_immediate_result():
+    def factory(**kwargs):
+        return _MockFunctionCall(result_value={"rank": kwargs["rank"], "ok": True})
+
+    mock_fn = _MockModalFunction(fcall_factory=factory)
+    executor = ModalSpawnExecutor(modal_function=mock_fn)
+    handles = executor.launch_replicas(
+        n_replicas=2,
+        entrypoint="x",
+        entrypoint_args={"world_size": 2},
+    )
+
+    # Immediate get() should succeed
+    for h in handles:
+        status = executor.poll(h)
+        assert status == "succeeded", f"rank {h.rank} expected succeeded, got {status}"
+
+
+@pytest.mark.skipif(not _is_modal_installed(),
+                    reason="modal not installed in this venv")
+def test_poll_returns_running_for_delayed_call():
+    def factory(**kwargs):
+        return _MockFunctionCall(
+            result_value={"rank": kwargs["rank"]}, delay_s=10.0,
+        )
+
+    mock_fn = _MockModalFunction(fcall_factory=factory)
+    executor = ModalSpawnExecutor(modal_function=mock_fn)
+    handles = executor.launch_replicas(
+        n_replicas=1, entrypoint="x", entrypoint_args={},
+    )
+
+    status = executor.poll(handles[0])
+    assert status == "running"
+
+
+@pytest.mark.skipif(not _is_modal_installed(),
+                    reason="modal not installed in this venv")
+def test_poll_returns_failed_for_user_exception():
+    def factory(**kwargs):
+        return _MockFunctionCall(
+            raise_on_get=ValueError("user code blew up"),
+        )
+
+    mock_fn = _MockModalFunction(fcall_factory=factory)
+    executor = ModalSpawnExecutor(modal_function=mock_fn)
+    handles = executor.launch_replicas(
+        n_replicas=1, entrypoint="x", entrypoint_args={},
+    )
+
+    status = executor.poll(handles[0])
+    assert status == "failed"
+
+
+@pytest.mark.skipif(not _is_modal_installed(),
+                    reason="modal not installed in this venv")
+def test_collect_returns_per_replica_dicts():
+    def factory(**kwargs):
+        return _MockFunctionCall(result_value={"rank": kwargs["rank"], "ok": True})
+
+    mock_fn = _MockModalFunction(fcall_factory=factory)
+    executor = ModalSpawnExecutor(modal_function=mock_fn)
+    handles = executor.launch_replicas(
+        n_replicas=3, entrypoint="x", entrypoint_args={},
+    )
+
+    results = executor.collect(handles, timeout=5)
+    assert len(results) == 3
+    for i, r in enumerate(results):
+        assert r["rank"] == i
+        assert r["status"] == "succeeded"
+        assert r["exit_code"] == 0
+        assert r["error"] is None
+        assert r["result"] == {"rank": i, "ok": True}
+        assert r["call_id"].startswith("fc-mock-")
+
+
+@pytest.mark.skipif(not _is_modal_installed(),
+                    reason="modal not installed in this venv")
+def test_collect_caches_results_and_does_not_call_get_twice():
+    """Once a poll() succeeds, collect() must read from cache, not call .get() again."""
+    get_calls = []
+
+    class _CountingFC(_MockFunctionCall):
+        def get(self, timeout=None):
+            get_calls.append(self.object_id)
+            return super().get(timeout=timeout)
+
+    def factory(**kwargs):
+        return _CountingFC(result_value={"rank": kwargs["rank"]})
+
+    mock_fn = _MockModalFunction(fcall_factory=factory)
+    executor = ModalSpawnExecutor(modal_function=mock_fn)
+    handles = executor.launch_replicas(
+        n_replicas=2, entrypoint="x", entrypoint_args={},
+    )
+
+    # Poll all to cache
+    for h in handles:
+        executor.poll(h)
+    n_polls = len(get_calls)
+    assert n_polls == 2  # one .get per poll
+
+    # Collect should NOT call .get() again
+    results = executor.collect(handles, timeout=5)
+    assert len(get_calls) == n_polls  # no additional .get() calls
+    assert all(r["status"] == "succeeded" for r in results)
+
+
+# ---------------------------------------------------------------------
+# Logs / cancel
+# ---------------------------------------------------------------------
+
+
+@pytest.mark.skipif(not _is_modal_installed(),
+                    reason="modal not installed in this venv")
+def test_stream_logs_includes_dashboard_url_and_call_id():
+    mock_fn = _MockModalFunction()
+    executor = ModalSpawnExecutor(modal_function=mock_fn)
+    handles = executor.launch_replicas(
+        n_replicas=1, entrypoint="x", entrypoint_args={},
+    )
+
+    log_text = executor.stream_logs(handles[0])
+    assert "fc-mock-" in log_text
+    assert "https://modal.com/" in log_text
+
+
+@pytest.mark.skipif(not _is_modal_installed(),
+                    reason="modal not installed in this venv")
+def test_cancel_calls_fcall_cancel():
+    mock_fn = _MockModalFunction()
+    executor = ModalSpawnExecutor(modal_function=mock_fn)
+    handles = executor.launch_replicas(
+        n_replicas=2, entrypoint="x", entrypoint_args={},
+    )
+
+    fc0 = executor._handles[0]["fcall"]
+    assert fc0._cancelled is False
+    executor.cancel(handles[0])
+    assert fc0._cancelled is True
+    # Cancelling rank 1 doesn't affect rank 0 (already cancelled — no-op)
+    executor.cancel(handles[1])
+    fc1 = executor._handles[1]["fcall"]
+    assert fc1._cancelled is True
+
+
+@pytest.mark.skipif(not _is_modal_installed(),
+                    reason="modal not installed in this venv")
+def test_cancel_unknown_handle_is_noop():
+    mock_fn = _MockModalFunction()
+    executor = ModalSpawnExecutor(modal_function=mock_fn)
+    # No replicas launched — handle has rank that doesn't exist in _handles
+    fake_handle = ReplicaHandle(
+        rank=99, backend_name="modal_spawn", metadata={"call_id": "nonexistent"},
+    )
+    # Must not raise
+    executor.cancel(fake_handle)