pools: async Pool surface — ModalFlashPool rides Modal's native .aio(); readiness discovers off-loop

src/stitch/pools/base.py

@@ -3,10 +3,17 @@
 Instances subclass this base: ``pools/modal_flash.py`` (Modal Flash). Add k8s as a new
 subclass. Override ``gateway_url`` and ``discover_replicas``; ``wake`` and ``scale`` are
 optional (their no-op defaults fall back to the replicas' own polling / load-autoscale).
+
+The ``*_async`` variants serve callers running on an event loop (an async control plane,
+``service.readiness``). Their defaults thread the sync implementation, so every subclass is
+usable from async code as-is — but a pool whose client has a native async surface should
+override them to ride it directly instead of parking a worker thread per call.
 """
 
 from __future__ import annotations
 
+import asyncio
+
 from stitch.types import VersionRef
 
 
@@ -30,3 +37,15 @@ def wake(self, replicas: list[str], ref: VersionRef) -> None:
 
     def scale(self, *, min: int | None = None, max: int | None = None) -> None:
         """Adjust the replica floor/cap. Optional — the default relies on load-autoscale."""
+
+    async def gateway_url_async(self) -> str:
+        """``gateway_url`` for async callers (default: the sync impl, off-loop)."""
+        return await asyncio.to_thread(self.gateway_url)
+
+    async def discover_replicas_async(self) -> list[str]:
+        """``discover_replicas`` for async callers (default: the sync impl, off-loop)."""
+        return await asyncio.to_thread(self.discover_replicas)
+
+    async def wake_async(self, replicas: list[str], ref: VersionRef) -> None:
+        """``wake`` for async callers (default: the sync impl, off-loop)."""
+        await asyncio.to_thread(self.wake, replicas, ref)

src/stitch/pools/base_test.py

@@ -0,0 +1,50 @@
+"""Pool base harness: the ``*_async`` defaults delegate to the sync implementations
+off the event loop, so any sync-only Pool subclass is usable from async callers."""
+
+from __future__ import annotations
+
+import asyncio
+import threading
+
+from stitch.pools.base import Pool
+from stitch.types import VersionRef
+
+
+class _SyncOnlyPool(Pool):
+    def __init__(self) -> None:
+        self.calls: list[tuple[str, int]] = []
+
+    def gateway_url(self) -> str:
+        self.calls.append(("gateway_url", threading.get_ident()))
+        return "https://gw"
+
+    def discover_replicas(self) -> list[str]:
+        self.calls.append(("discover_replicas", threading.get_ident()))
+        return ["https://r1", "https://r2"]
+
+    def wake(self, replicas: list[str], ref: VersionRef) -> None:
+        self.calls.append(("wake", threading.get_ident()))
+
+
+def test_async_defaults_delegate_to_sync_impls_off_loop() -> None:
+    pool = _SyncOnlyPool()
+
+    async def drive() -> tuple[str, list[str]]:
+        url = await pool.gateway_url_async()
+        replicas = await pool.discover_replicas_async()
+        await pool.wake_async(replicas, VersionRef("run", 1))
+        return url, replicas
+
+    url, replicas = asyncio.run(drive())  # the loop runs on THIS thread
+    assert (url, replicas) == ("https://gw", ["https://r1", "https://r2"])
+    assert [name for name, _ in pool.calls] == ["gateway_url", "discover_replicas", "wake"]
+    # the sync impls must have run on worker threads, never on the loop's thread
+    assert all(ident != threading.get_ident() for _, ident in pool.calls)
+
+
+if __name__ == "__main__":
+    tests = [v for k, v in sorted(globals().items()) if k.startswith("test_")]
+    for t in tests:
+        t()
+        print(f"  ok  {t.__name__}")
+    print(f"pool base harness: {len(tests)} PASS")

src/stitch/pools/modal_flash.py

@@ -3,10 +3,15 @@
 Replicas are the Flash containers; the gateway is the Flash URL. This is a *client*
 to a running pool — reach, enumerate, wake, scale — not the pool's deployment (that
 is an example). Every Modal call is import-lazy, so the module loads without Modal.
+
+The ``*_async`` overrides use the Modal SDK's native ``.aio()`` interface (every Modal
+call is synchronicity-wrapped and awaitable), so async callers never park a worker
+thread just to wait on Modal's own event loop.
 """
 
 from __future__ import annotations
 
+import asyncio
 import logging
 from concurrent.futures import ThreadPoolExecutor
 
@@ -32,7 +37,12 @@ def _cls(self):
             ) from exc
 
     def gateway_url(self) -> str:
-        urls = self._cls()._experimental_get_flash_urls()
+        return self._pick_gateway(self._cls()._experimental_get_flash_urls())
+
+    async def gateway_url_async(self) -> str:
+        return self._pick_gateway(await self._cls()._experimental_get_flash_urls.aio())
+
+    def _pick_gateway(self, urls) -> str:
         if not urls:
             raise RuntimeError(
                 f"no Flash gateway URL for {self.app_name}.{self.cls_name} — deploy the app first"
@@ -43,8 +53,13 @@ def discover_replicas(self) -> list[str]:
         import modal.experimental
 
         self._cls()  # resolve first: clear error if not deployed
-        containers = modal.experimental.flash_get_containers(self.app_name, self.cls_name)
-        return [_normalize_url(h) for c in containers if (h := _host(c))]
+        return _replica_urls(modal.experimental.flash_get_containers(self.app_name, self.cls_name))
+
+    async def discover_replicas_async(self) -> list[str]:
+        import modal.experimental
+
+        self._cls()  # resolve first: clear error if not deployed
+        return _replica_urls(await modal.experimental.flash_get_containers.aio(self.app_name, self.cls_name))
 
     def wake(self, replicas: list[str], ref: VersionRef) -> None:
         # Fan out (this is on the publish hot path); each replica re-reads the pointer, so no version in the body.
@@ -63,6 +78,21 @@ def wake_one(url: str) -> None:
             with ThreadPoolExecutor(max_workers=min(16, len(replicas))) as pool:
                 list(pool.map(wake_one, replicas))
 
+    async def wake_async(self, replicas: list[str], ref: VersionRef) -> None:
+        if not replicas:
+            return
+        import httpx
+
+        async with httpx.AsyncClient(timeout=5.0, trust_env=False) as client:
+
+            async def wake_one(url: str) -> None:
+                try:
+                    (await client.post(f"{url}/wake")).raise_for_status()
+                except Exception as exc:  # noqa: BLE001
+                    logger.warning("failed to wake %s for %s: %s", url, ref.identity, exc)
+
+            await asyncio.gather(*(wake_one(url) for url in replicas))
+
     def scale(self, *, min: int | None = None, max: int | None = None) -> None:
         fn = self._cls()._get_class_service_function()
         kwargs: dict[str, int] = {}
@@ -74,6 +104,10 @@ def scale(self, *, min: int | None = None, max: int | None = None) -> None:
             fn.update_autoscaler(**kwargs)
 
 
+def _replica_urls(containers) -> list[str]:
+    return [_normalize_url(h) for c in containers if (h := _host(c))]
+
+
 def _host(container) -> str | None:
     if isinstance(container, dict):
         return container.get("host")

src/stitch/pools/modal_flash_test.py

@@ -3,7 +3,12 @@
 
 from __future__ import annotations
 
-from stitch.pools.modal_flash import _host, _normalize_url
+from stitch.pools.modal_flash import _host, _normalize_url, _replica_urls
+
+
+def test_replica_urls_filters_hostless_and_normalizes() -> None:
+    containers = [{"host": "h1:8000"}, {}, {"host": "https://h2/"}]
+    assert _replica_urls(containers) == ["https://h1:8000", "https://h2"]
 
 
 def test_normalize_url_adds_scheme_and_strips_slash() -> None:

src/stitch/service.py

@@ -213,7 +213,9 @@ async def probe(c: Any, url: str) -> ReplicaState:
             return ReplicaState(reason=str(exc)[:80])  # applied=None => counts as not at any version
 
     async with httpx.AsyncClient(trust_env=False) as c:
-        states = await asyncio.gather(*(probe(c, url) for url in pool.discover_replicas()))
+        # the async variant keeps pool-client I/O off this event loop (native or threaded per pool)
+        replicas = await pool.discover_replicas_async()
+        states = await asyncio.gather(*(probe(c, url) for url in replicas))
     return PoolState(list(states))