ENG-525: provider-aware warm pool

docs/architecture.md

@@ -92,7 +92,8 @@ keys in `hosts.schemas.RESERVED_HOST_ENV_KEYS` are rejected.
 
 Two maintenance commands run as cron jobs from the same image:
 `hosts.janitor` reaps expired and orphaned hosts, `hosts.pool` keeps a
-warm pool of pre-provisioned hosts (`POOL_SIZE`) to hide provider cold
+warm pool of pre-provisioned hosts per provider (`POOL_SIZES`, with
+`POOL_SIZE` as the default provider's target) to hide provider cold
 starts.
 
 ## Diagnostics

docs/deploy.md

@@ -25,7 +25,8 @@ docker run --rm --env-file drukbox.env "$IMAGE" .venv/bin/python -m hosts.pool
 ```
 
 The janitor reaps expired and orphaned hosts. The pool maintainer
-pre-provisions warm hosts and only does anything when `POOL_SIZE > 0`.
+pre-provisions warm hosts per provider and only does anything when at
+least one provider has a warm target (`POOL_SIZES` / `POOL_SIZE`).
 Schedule both under your cron infrastructure (k8s `CronJob`, systemd
 timer) from the same image and env file.
 
@@ -139,9 +140,10 @@ Core, optional:
 | `UVICORN_HOST` | `0.0.0.0` | API bind address. Set `127.0.0.1` to restrict to loopback. |
 | `PROVISIONING_GRACE_SECONDS` | `600` | Safety TTL on in-flight hosts so the janitor reaps row + VM if the client disconnects mid-provision. Must exceed the worst-case provision duration. |
 | `IDEMPOTENCY_KEY_TTL_HOURS` | `24` | Retention period for successful `Idempotency-Key` mappings. |
-| `POOL_SIZE` | `0` | Warm hosts to keep ready. `0` disables pooling. |
+| `POOL_SIZES` | `{}` | Warm hosts to keep ready per provider, as JSON (e.g. `{"exe": 2, "hetzner": 1}`). Overrides `POOL_SIZE` for the providers it names. |
+| `POOL_SIZE` | `0` | Warm hosts to keep ready for the default provider. `0` disables its pool. |
 | `POOL_HOST_MAX_AGE_HOURS` | `4` | Max age before the janitor reaps an unclaimed pool host. |
-| `POOL_MAX_CREATES_PER_TICK` | `2` | Upper bound on pool provisions per tick; caps over-provision blast radius when ticks overlap. |
+| `POOL_MAX_CREATES_PER_TICK` | `2` | Upper bound on pool provisions per tick, across all providers; caps over-provision blast radius when ticks overlap. |
 
 Tailscale (required when `TAILSCALE_ENABLED=true`):
 

src/core/settings.py

@@ -70,11 +70,22 @@ class Settings(BaseSettings):
         validation_alias="PROVISIONING_GRACE_SECONDS",
         description="Safety TTL on the host row while provisioning is in flight.",
     )
+    pool_sizes: dict[str, Annotated[int, Field(ge=0)]] = Field(
+        default_factory=dict,
+        validation_alias="POOL_SIZES",
+        description=(
+            'Pre-warmed host targets per provider, as JSON (e.g. {"exe": 2, "hetzner": 1}). '
+            "Overrides POOL_SIZE for the providers it names."
+        ),
+    )
     pool_size: int = Field(
         default=0,
         ge=0,
         validation_alias="POOL_SIZE",
-        description="Number of pre-warmed hosts to keep ready. 0 disables pooling.",
+        description=(
+            "Number of pre-warmed hosts to keep ready for the default provider. "
+            "0 disables its pool. A POOL_SIZES entry for that provider wins."
+        ),
     )
     pool_host_max_age_hours: int = Field(
         default=4,
@@ -86,9 +97,15 @@ class Settings(BaseSettings):
         default=2,
         ge=0,
         validation_alias="POOL_MAX_CREATES_PER_TICK",
-        description="Upper bound on pool-maintainer provisions per tick.",
+        description="Upper bound on pool-maintainer provisions per tick, across all providers.",
     )
 
+    def get_pool_targets(self) -> dict[str, int]:
+        # POOL_SIZE seeds the default provider's target and POOL_SIZES
+        # overrides per provider; providers at zero drop out entirely.
+        targets = {self.default_host_provider: self.pool_size, **self.pool_sizes}
+        return {provider: target for provider, target in targets.items() if target > 0}
+
 
 @lru_cache
 def get_settings() -> Settings:

src/core/tests/test_settings.py

@@ -100,6 +100,49 @@ def test_numeric_settings_reject_negative_values(monkeypatch: pytest.MonkeyPatch
         _settings_with(monkeypatch, env)
 
 
+def test_pool_size_seeds_the_default_providers_target(monkeypatch: pytest.MonkeyPatch) -> None:
+    env: dict[str, str | None] = {
+        **_base_env(),
+        "DEFAULT_HOST_PROVIDER": None,
+        "POOL_SIZE": "2",
+        "POOL_SIZES": None,
+    }
+    settings = _settings_with(monkeypatch, env)
+    assert settings.get_pool_targets() == {"exe": 2}
+
+
+def test_pool_sizes_overrides_the_alias_for_the_same_provider(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    env: dict[str, str | None] = {
+        **_base_env(),
+        "DEFAULT_HOST_PROVIDER": None,
+        "POOL_SIZE": "5",
+        "POOL_SIZES": '{"exe": 2, "hetzner": 1}',
+    }
+    settings = _settings_with(monkeypatch, env)
+    assert settings.get_pool_targets() == {"exe": 2, "hetzner": 1}
+
+
+def test_pool_targets_omit_zeroed_providers(monkeypatch: pytest.MonkeyPatch) -> None:
+    # An explicit zero in POOL_SIZES disables that provider's pool even when
+    # the POOL_SIZE alias would seed it.
+    env: dict[str, str | None] = {
+        **_base_env(),
+        "DEFAULT_HOST_PROVIDER": None,
+        "POOL_SIZE": "5",
+        "POOL_SIZES": '{"exe": 0}',
+    }
+    settings = _settings_with(monkeypatch, env)
+    assert settings.get_pool_targets() == {}
+
+
+def test_pool_sizes_rejects_negative_targets(monkeypatch: pytest.MonkeyPatch) -> None:
+    env: dict[str, str | None] = {**_base_env(), "POOL_SIZES": '{"exe": -1}'}
+    with pytest.raises(ValueError, match="POOL_SIZES"):
+        _settings_with(monkeypatch, env)
+
+
 def test_load_test_env_overrides_ambient_values(monkeypatch: pytest.MonkeyPatch) -> None:
     monkeypatch.setenv("TAILSCALE_ENABLED", "false")
     conftest.load_test_env()

src/hosts/pool.py

@@ -1,5 +1,7 @@
 import logging
-from dataclasses import dataclass
+import random
+import uuid
+from dataclasses import dataclass, field
 from datetime import timedelta
 
 from sqlalchemy import func, or_, select
@@ -13,61 +15,71 @@
 log = logging.getLogger(__name__)
 
 # Bail out after this many consecutive create failures (broker outage scenario)
-# rather than logging the same failure pool_size times per tick.
+# rather than logging the same failure once per deficit slot per tick.
 _MAX_CONSECUTIVE_CREATE_FAILURES = 3
 
 
 @dataclass(frozen=True)
 class PoolMaintenanceSummary:
     created: int = 0
     removed_excess: int = 0
-    pool_size: int = 0
+    targets: dict[str, int] = field(default_factory=dict)
 
 
 async def maintain_pool() -> PoolMaintenanceSummary:
-    """Top up / shrink the warm-host pool and reap failed pool members.
+    """Top up / shrink each provider's warm-host pool and reap failed pool members.
 
-    Each tick creates at most ``POOL_MAX_CREATES_PER_TICK`` hosts. Overlapping
-    ticks or accidental scheduler replicas can therefore over-provision by at
-    most that batch size, which the next tick sheds via the excess path. This
-    intentionally avoids any cross-tick locking — see CONTRIBUTING for context.
+    Each tick creates at most ``POOL_MAX_CREATES_PER_TICK`` hosts across all
+    providers. Overlapping ticks or accidental scheduler replicas can therefore
+    over-provision by at most that batch size, which the next tick sheds via the
+    excess path. This intentionally avoids any cross-tick locking — see
+    CONTRIBUTING for context.
     """
     settings = get_settings()
-    if settings.pool_size <= 0:
-        return PoolMaintenanceSummary(pool_size=settings.pool_size)
+    targets = settings.get_pool_targets()
+    if not targets:
+        return PoolMaintenanceSummary()
     tailscale: Tailscale | None = Tailscale.from_settings() if settings.tailscale_enabled else None
     try:
-        return await _maintain(settings, tailscale)
+        return await _maintain(settings, targets, tailscale)
     finally:
         if tailscale is not None:
             await tailscale.aclose()
 
 
-async def _maintain(settings: Settings, tailscale: Tailscale | None) -> PoolMaintenanceSummary:
+async def _maintain(
+    settings: Settings, targets: dict[str, int], tailscale: Tailscale | None
+) -> PoolMaintenanceSummary:
     now = utc_now()
 
     async with async_session_factory() as session:
-        # Count warm pool members that are still fresh and unclaimed. Only
-        # pool_member hosts count — demand-provisioned hosts also keep
-        # claimed_at NULL but must never be counted or shed. Exclude hosts whose
-        # expires_at has already passed — those are queued for janitor reaping.
-        current = await session.scalar(
-            select(func.count())
-            .select_from(Host)
+        # Count warm pool members per provider that are still fresh and
+        # unclaimed. Only pool_member hosts count — demand-provisioned hosts
+        # also keep claimed_at NULL but must never be counted or shed. Exclude
+        # hosts whose expires_at has already passed — those are queued for
+        # janitor reaping.
+        counted = await session.execute(
+            select(Host.provider, func.count())
             .where(Host.pool_member.is_(True))
             .where(Host.claimed_at.is_(None))
             .where(Host.status != HostStatus.ERROR.value)
             .where(or_(Host.expires_at.is_(None), Host.expires_at > now))
+            .group_by(Host.provider)
         )
-        current = current or 0
-
-        excess_ids: list = []
-        if current > settings.pool_size:
-            shed = current - settings.pool_size
-            excess_ids = list(
+        current: dict[str, int] = {provider: count for provider, count in counted.all()}
+
+        excess_ids: list[uuid.UUID] = []
+        # Sweep the union of configured and present providers: one dropped from
+        # the targets sheds to zero instead of idling until its max-age TTL.
+        for provider in sorted(current.keys() | targets.keys()):
+            shed = current.get(provider, 0) - targets.get(provider, 0)
+            if shed <= 0:
+                continue
+            excess_ids.extend(
                 (
                     await session.execute(
                         select(Host.id)
+                        .where(Host.provider == provider)
                         .where(Host.pool_member.is_(True))
                         .where(Host.claimed_at.is_(None))
                         .where(Host.status != HostStatus.ERROR.value)
@@ -78,12 +90,29 @@ async def _maintain(settings: Settings, tailscale: Tailscale | None) -> PoolMain
                 ).scalars()
             )
 
-    deficit = max(settings.pool_size - current, 0)
-    batch = min(deficit, settings.pool_max_creates_per_tick)
+    underfilled = [
+        provider for provider, target in targets.items() if target > current.get(provider, 0)
+    ]
+    # Spread the per-tick budget round-robin across providers so one large
+    # deficit can't starve the others. The order is shuffled each tick: with
+    # no cross-tick state, a fixed order would starve the last provider
+    # whenever the cap is smaller than the number of underfilled providers —
+    # permanently so when an earlier provider's creates keep failing.
+    random.shuffle(underfilled)
+    deficits = {provider: targets[provider] - current.get(provider, 0) for provider in underfilled}
+    batch: list[str] = []
+    while deficits and len(batch) < settings.pool_max_creates_per_tick:
+        for provider in list(deficits):
+            if len(batch) == settings.pool_max_creates_per_tick:
+                break
+            batch.append(provider)
+            deficits[provider] -= 1
+            if not deficits[provider]:
+                del deficits[provider]
 
     created = 0
     consecutive_failures = 0
-    for _ in range(batch):
+    for provider in batch:
         try:
             async with async_session_factory() as session:
                 service = HostService(session, settings=settings, tailscale=tailscale)
@@ -92,13 +121,14 @@ async def _maintain(settings: Settings, tailscale: Tailscale | None) -> PoolMain
                     env={},
                     image=None,
                     expires_at=expires_at,
+                    provider=provider,
                     pool_member=True,
                 )
                 created += 1
                 consecutive_failures = 0
         except Exception:
             consecutive_failures += 1
-            log.exception("pool: failed to create pool host")
+            log.exception("pool: failed to create pool host provider=%s", provider)
             if consecutive_failures >= _MAX_CONSECUTIVE_CREATE_FAILURES:
                 log.error(
                     "pool: aborting top-up after %d consecutive failures",
@@ -118,16 +148,16 @@ async def _maintain(settings: Settings, tailscale: Tailscale | None) -> PoolMain
 
     if created or removed_excess:
         log.info(
-            "pool: maintained (created=%d, removed_excess=%d, target=%d)",
+            "pool: maintained (created=%d, removed_excess=%d, targets=%s)",
             created,
             removed_excess,
-            settings.pool_size,
+            targets,
         )
 
     return PoolMaintenanceSummary(
         created=created,
         removed_excess=removed_excess,
-        pool_size=settings.pool_size,
+        targets=targets,
     )
 
 

src/hosts/service.py

@@ -96,10 +96,14 @@ async def get_or_create_host(
                 return existing
 
         host: Host | None = None
-        # Pool hosts are warmed with the default provider, so a pinned provider
-        # always provisions fresh.
-        if self.settings.pool_size > 0 and not env and image is None and not provider:
-            host = await self._try_claim_pool_host(expires_at=expires_at)
+        # Warm hosts are provider-specific, so the claim is scoped to the
+        # requested provider's pool. A request is pool-eligible only when it
+        # doesn't customize the host: default image and no env.
+        requested_provider = provider or self.settings.default_host_provider
+        if not env and image is None and self.settings.get_pool_targets().get(requested_provider):
+            host = await self._try_claim_pool_host(
+                provider=requested_provider, expires_at=expires_at
+            )
         if host is None:
             host = await self.create_host(
                 env=env, image=image, expires_at=expires_at, provider=provider
@@ -119,7 +123,9 @@ async def get_or_create_host(
             return winner
         return host
 
-    async def _try_claim_pool_host(self, *, expires_at: datetime | None) -> Host | None:
+    async def _try_claim_pool_host(
+        self, *, provider: str, expires_at: datetime | None
+    ) -> Host | None:
         # Pick a candidate, then atomically claim it with UPDATE ... WHERE
         # claimed_at IS NULL ... RETURNING. The WHERE predicate is the actual
         # race guard — concurrent claimants resolve to a single winner per
@@ -130,6 +136,7 @@ async def _try_claim_pool_host(self, *, expires_at: datetime | None) -> Host | N
         candidate_id = (
             await self.session.execute(
                 select(Host.id)
+                .where(Host.provider == provider)
                 .where(Host.pool_member.is_(True))
                 .where(Host.claimed_at.is_(None))
                 .where(Host.status == HostStatus.ACTIVE.value)

src/hosts/tests/test_api.py

@@ -385,7 +385,7 @@ async def test_idempotency_loser_that_claimed_pool_host_is_returned_to_pool(monk
 
     async with async_session_factory() as session:
         service = HostService(session)
-        claimed = await service._try_claim_pool_host(expires_at=None)
+        claimed = await service._try_claim_pool_host(provider="exe", expires_at=None)
         assert claimed is not None
         assert claimed.claimed_at is not None
         await service._release_idempotency_loser(claimed)