Codeseys/composer-replication-framework

1"""Modal executor — skeleton for v0.2 3This file is a STUB. The full Modal integration requires the `modal`4client library installed (`pip install modal`) and a configured Modal5account (`~/.modal.toml`). The user's environment has both, but the6test suite must run without them, so we keep this file import-safe.7 8Real implementation lives in v0 polish; the docstring below is the9contract.10"""11from __future__ import annotations12 13from typing import Any, Callable, Mapping14 15from composer_replication.diloco.serverless.executor import (16    ReplicaHandle,17    ServerlessExecutor,18)19 20 21class ModalExecutor(ServerlessExecutor):22    """Run replicas as Modal Functions in parallel.23 24    Reference implementation pattern (per ADR-005):25 26        @app.function(gpu="A100-40GB", timeout=3600)27        def run_replica(rank: int, rendezvous_uri: str, **kwargs):28            os.environ["REPLICA_RANK"] = str(rank)29            from composer_replication.diloco.serverless import (30                MockManager, ObjectStoreAllReduce,31            )32            store = ObjectStoreAllReduce(rendezvous_uri,33                                         rank=rank, world_size=N)34            manager = MockManager(store)35            # ... run the trainer with this manager ...36 37    Then `launch_replicas` does:38        calls = [run_replica.spawn(rank=i, ...) for i in range(N)]39        return [ReplicaHandle(rank=i, backend_name="modal",40                              metadata={"call_id": calls[i].object_id})41                for i in range(N)]42 43    Pricing reference (2026-05-26): A100-40GB ≈ $1.95/hr, H100 ≈ $5.50/hr.44    Cold start ≈ 30s. Inter-job networking via cluster mode (opt-in,45    not used by default).46 47    Status: SKELETON. Real implementation pending v0 polish wave.48    """49    backend_name = "modal"50    supports_inter_replica_network = False  # default; cluster mode = True51 52    def __init__(self, *, app_name: str = "composer-replication-diloco") -> None:53        try:54            import modal  # noqa: F40155        except ImportError as e:56            raise RuntimeError(57                "ModalExecutor requires the modal client. Install with "58                "`pip install modal` and configure with `modal token new`. "59                "Got: " + repr(e)60            )61        self.app_name = app_name62        # Real implementation: build a `modal.App` and register `run_replica`63        # here so that subsequent `launch_replicas` can `.spawn()` it.64        raise NotImplementedError(65            "ModalExecutor is a v0 skeleton; full implementation pending. "66            "Use LocalProcessExecutor for testing."67        )68 69    # All Protocol methods raise NotImplementedError via __init__ — the70    # class never instantiates successfully in the skeleton. Sketch71    # signatures here for documentation:72 73    def launch_replicas(74        self,75        n_replicas: int,76        entrypoint: str | Callable[..., Any],77        entrypoint_args: Mapping[str, Any],78        *,79        gpu: str | None = "A100-40GB",80        timeout: int = 3600,81    ) -> list[ReplicaHandle]:82        raise NotImplementedError83 84    def poll(self, handle: ReplicaHandle) -> str:85        raise NotImplementedError86 87    def stream_logs(self, handle: ReplicaHandle, *, n_lines: int = 200) -> str:88        raise NotImplementedError89 90    def cancel(self, handle: ReplicaHandle) -> None:91        raise NotImplementedError92 93    def collect(94        self,95        handles: list[ReplicaHandle],96        *,97        timeout: int | None = None,98    ) -> list[dict[str, Any]]:99        raise NotImplementedError100 101 102__all__ = ["ModalExecutor"]103