From 98f8807fc3c2d79af9a00fe8e10e5397478f291c Mon Sep 17 00:00:00 2001 From: Jason Mancuso <7891333+jvmncs@users.noreply.github.com> Date: Wed, 22 Jul 2026 21:20:38 +0000 Subject: [PATCH] =?UTF-8?q?pools:=20async=20Pool=20surface=20=E2=80=94=20M?= =?UTF-8?q?odalFlashPool=20rides=20Modal's=20native=20.aio();=20readiness?= =?UTF-8?q?=20discovers=20off-loop?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Base Pool grows gateway_url_async/discover_replicas_async/wake_async with to_thread fallbacks over the sync impls, so any sync-only Pool stays usable from async callers; ModalFlashPool overrides them with the SDK's .aio() and async-httpx wake. service.readiness switches to the async discovery, keeping blocking pool-client I/O off the caller's event loop. Confirmed by probe (stitch-dev, read-only): sync and async surfaces return identical gateway/replica results against a live Flash pool; NotFoundError parity on an app without the cls; no event-loop stall (<25ms heartbeat gap) during the async calls. Modal's own AsyncUsageWarning fires on the sync surface in an async context and suggests exactly the .aio rewrite adopted here. uv run pytest: 52 passed. --- src/stitch/pools/base.py | 19 +++++++++++ src/stitch/pools/base_test.py | 50 ++++++++++++++++++++++++++++ src/stitch/pools/modal_flash.py | 40 ++++++++++++++++++++-- src/stitch/pools/modal_flash_test.py | 7 +++- src/stitch/service.py | 4 ++- 5 files changed, 115 insertions(+), 5 deletions(-) create mode 100644 src/stitch/pools/base_test.py diff --git a/src/stitch/pools/base.py b/src/stitch/pools/base.py index 085b10c..3f596af 100644 --- a/src/stitch/pools/base.py +++ b/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) diff --git a/src/stitch/pools/base_test.py b/src/stitch/pools/base_test.py new file mode 100644 index 0000000..c284709 --- /dev/null +++ b/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") diff --git a/src/stitch/pools/modal_flash.py b/src/stitch/pools/modal_flash.py index c1f1882..7b280cd 100644 --- a/src/stitch/pools/modal_flash.py +++ b/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") diff --git a/src/stitch/pools/modal_flash_test.py b/src/stitch/pools/modal_flash_test.py index 1b91d64..3c80a3a 100644 --- a/src/stitch/pools/modal_flash_test.py +++ b/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: diff --git a/src/stitch/service.py b/src/stitch/service.py index 320a942..7285576 100644 --- a/src/stitch/service.py +++ b/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))