Warm simulation machinery at startup; make /readiness-check honest

.github/scripts/cloud_run_env.sh

@@ -29,10 +29,13 @@ cloud_run_set_defaults() {
   CLOUD_RUN_CONCURRENCY="${CLOUD_RUN_CONCURRENCY:-6}"
   CLOUD_RUN_WEB_CONCURRENCY="${CLOUD_RUN_WEB_CONCURRENCY:-2}"
   CLOUD_RUN_PORT="${CLOUD_RUN_PORT:-8080}"
-  # Readiness gate: a TCP probe passes at port-bind, minutes before the app can
-  # serve. Window is 180s + 24x10s = 420s; both halves are capped at 240s.
-  # Sizing rationale in docs/migration/cloud-run-operations.md.
-  CLOUD_RUN_STARTUP_PROBE="${CLOUD_RUN_STARTUP_PROBE:-httpGet.path=/readiness-check,httpGet.port=${CLOUD_RUN_PORT},initialDelaySeconds=180,periodSeconds=10,failureThreshold=24,timeoutSeconds=5}"
+  # Readiness gate: /readiness-check stays 503 until the worker has built the
+  # tax-benefit systems AND run the startup warmup calculate, so boot-to-ready is
+  # now the ~371s p90 import plus the warmup. Window sized to the platform maximum
+  # (240s + 24x10s = 480s); both halves are capped at 240s. initialDelay is high
+  # but adds no real scale-out delay because boot exceeds it either way. Sizing
+  # rationale in docs/migration/cloud-run-operations.md.
+  CLOUD_RUN_STARTUP_PROBE="${CLOUD_RUN_STARTUP_PROBE:-httpGet.path=/readiness-check,httpGet.port=${CLOUD_RUN_PORT},initialDelaySeconds=240,periodSeconds=10,failureThreshold=24,timeoutSeconds=5}"
   CLOUD_RUN_POLICYENGINE_DB_PASSWORD_SECRET="${CLOUD_RUN_POLICYENGINE_DB_PASSWORD_SECRET:-policyengine-api-prod-db-password:latest}"
   CLOUD_RUN_GITHUB_MICRODATA_TOKEN_SECRET="${CLOUD_RUN_GITHUB_MICRODATA_TOKEN_SECRET:-policyengine-api-prod-github-microdata-token:latest}"
   CLOUD_RUN_ANTHROPIC_API_KEY_SECRET="${CLOUD_RUN_ANTHROPIC_API_KEY_SECRET:-policyengine-api-prod-anthropic-api-key:latest}"

changelog.d/startup-warmup-readiness.fixed.md

@@ -0,0 +1 @@
+Run a warmup calculate per country at Cloud Run startup and gate /readiness-check on it, so the endpoint reports ready only once the simulation machinery is compiled and the first real request is fast (previously the first calculate on a fresh worker took ~2 minutes, which readiness did not reflect). The startup-probe window is widened to the 480s platform maximum to cover the added warmup time.

docs/migration/cloud-run-operations.md

@@ -95,35 +95,45 @@ Values measured and justified in
 
   Cloud Run caps `failureThreshold x periodSeconds` at 240s **and**
   `initialDelaySeconds` at 240s, and shuts the container down past their sum. We run
-  `initialDelaySeconds 180 + 24 x 10s = 420s`. The threshold/period half is already
-  at its ceiling, so `initialDelaySeconds` is the only way to widen the window; it
-  is additive (no probe runs during it) but also delays availability for instances
-  that boot faster than it. Trade-off against the distribution above:
-
-  | initialDelay | window | boots killed | boots delayed | median penalty |
-  |---|---|---|---|---|
-  | 120 | 360s | 12.5% | 8.3% | 31s |
-  | **180 (current)** | **420s** | **6.2%** | **22.9%** | **25s** |
-  | 240 (max) | 480s | 2.1% | 77.1% | 50s |
-
-  A delayed instance loses tens of seconds; a **killed** one loses its whole boot
-  plus a retry (400s+) and fails the deploy if it happens in CI — so the asymmetry
-  favours a wider window. 240 was rejected because it holds 77% of instances to a
-  full 240s, slowing every scale-out.
-
-  **Residual risk: ~6% of boots still exceed 420s and will be killed and retried**,
-  which can fail a deploy. That is accepted deliberately — the alternative (TCP)
-  served real users 5xx from instances that were never ready. Qualified 2026-07-22
-  on a `--no-traffic` revision: Cloud Run stored the config exactly as specified and
-  the revision reached `Ready` on a 181.6s boot.
-
-  The real fix is cutting boot time — **~82% of it is constructing the US
-  tax-benefit system** (`policyengine_api.country` builds all five countries at
-  import; US alone is ~90% of that, and `CountryTaxBenefitSystem()` is ~91% of the
-  per-country cost). At a 20s boot this would be `initialDelay 0` with a 240s window
-  covering every boot. Note that lazily deferring the build does **not** help: it
-  relocates the cost onto the first request, where readiness would lie and a user
-  would absorb it.
+  `initialDelaySeconds 240 + 24 x 10s = 480s` — the platform maximum.
+
+  Boot-to-ready is now the ~371s p90 import **plus** the startup warmup (below), so
+  every boot exceeds 240s. A high `initialDelaySeconds` therefore no longer delays
+  any real scale-out — the earlier reason for keeping it low (holding sub-240s boots
+  to a full 240s; see the import-only distribution above) no longer applies, because
+  no boot is that fast. Sizing to the 480s maximum minimises killed-and-retried
+  boots, which are the costly failure (a killed boot loses its whole boot plus a
+  retry and fails a CI deploy).
+
+  Be clear-eyed about what this buys: the warmup consumes most of the headroom the
+  wider window adds. p90 boot-to-ready rises from ~371s (import only) to ~420s
+  (import + warmup), so the slack under the ceiling stays ~60s — roughly what it was
+  at 420s before — and the **far tail can still exceed 480s and be killed-and-retried**.
+  We are now at the platform ceiling; the only remaining lever is cutting boot time
+  itself (prebuilding the tax-benefit system into the image — see below), not a wider
+  probe window.
+
+  **Startup warmup — why readiness gates on more than the import.** Building the
+  tax-benefit systems at import does not compile the per-simulation machinery
+  (parameter-tree materialisation, the formula graph). The **first** calculate on a
+  fresh worker paid that cost — measured at ~121s — and `/readiness-check` returned
+  200 before it, so the probe (and health checks, and smoke tests) saw a "ready"
+  instance whose first real request still took two minutes.
+  `policyengine_api.warmup.run_startup_warmup` now runs a throwaway calculate per
+  country (default US; override with `POLICYENGINE_API_WARMUP_COUNTRIES`, disable
+  with `POLICYENGINE_API_STARTUP_WARMUP=0`) from `asgi.py` before the worker serves,
+  and `/readiness-check` returns 503 until it completes
+  (`policyengine_api.readiness`). Verified locally: after the warmup a cold
+  test-payload calculate drops from ~121s to ~10s. This is the fix for the note that
+  used to live here — deferring the build lazily relocated the cost onto the first
+  request, where readiness lied; warming it at startup pays the cost off the request
+  path and makes readiness honest, at the cost of a longer (but bounded) boot.
+
+  The lasting fix is still cutting boot time: **~82% of the import is constructing
+  the US tax-benefit system** (`policyengine_api.country` builds all five countries
+  at import; US alone is ~90% of that, and `CountryTaxBenefitSystem()` is ~91% of the
+  per-country cost). Baking a prebuilt system into the image would shrink both the
+  import and the warmup and let the window come back down.
 - **Scaling pins live in `push.yml` per job**: production `max-instances 8`, staging
   `min 0 / max 1`. Production was `max-instances 4` through Stage 10; it was raised to
   8 for the 100% Cloud Run cutover because at the 50/50 split the service already sat

policyengine_api/api.py

@@ -82,6 +82,8 @@ def log_timing(message):
 
 log_timing("Legacy endpoints import completed")
 
+from policyengine_api.readiness import is_ready
+
 log_timing("Initialising API...")
 
 app = application = flask.Flask(__name__)
@@ -187,6 +189,12 @@ def liveness_check():
 
 @app.route("/readiness-check", methods=["GET"])
 def readiness_check():
+    # 503 until the startup warmup has compiled the simulation machinery
+    # (policyengine_api.readiness); /liveness-check stays unconditional.
+    if not is_ready():
+        return flask.Response(
+            "NOT READY", status=503, headers={"Content-Type": "text/plain"}
+        )
     return flask.Response("OK", status=200, headers={"Content-Type": "text/plain"})
 
 

policyengine_api/asgi.py

@@ -2,8 +2,20 @@
 
 from __future__ import annotations
 
+import os
+
 from policyengine_api.api import app as flask_app
 from policyengine_api.asgi_factory import create_asgi_app
-
+from policyengine_api.readiness import mark_not_ready, mark_ready
+from policyengine_api.warmup import run_startup_warmup
 
 app = application = create_asgi_app(flask_app)
+
+# Warm the simulation machinery before serving (see policyengine_api.warmup).
+# POLICYENGINE_API_STARTUP_WARMUP=0 skips it.
+if os.environ.get("POLICYENGINE_API_STARTUP_WARMUP", "1") != "0":
+    mark_not_ready()
+    try:
+        run_startup_warmup()
+    finally:
+        mark_ready()

policyengine_api/openapi_spec.yaml

@@ -1069,14 +1069,23 @@ paths:
     get:
       summary: Test for server readiness.
       operationId: readiness_check
-      description: Determine whether or not the PolicyEngine server is ready.
+      description: >-
+        Determine whether or not the PolicyEngine server is ready to serve
+        requests. Unlike /liveness-check, this returns 503 until the startup
+        warmup has compiled the simulation machinery.
       responses:
         200:
           description: Server is ready.
           content:
             text/plain:
               schema:
                 type: string
+        503:
+          description: Server is live but not yet ready (still warming up).
+          content:
+            text/plain:
+              schema:
+                type: string
   /specification:
     get:
       summary: Get OpenAPI specs.

policyengine_api/readiness.py

@@ -0,0 +1,33 @@
+"""Whether the API has finished warming up and can serve real requests.
+
+Defaults to ready, so contexts that do not run a startup warmup (App Engine,
+tests, tooling) report ready immediately. The Cloud Run startup path toggles it
+around the warmup (policyengine_api.warmup).
+"""
+
+from __future__ import annotations
+
+import threading
+
+_lock = threading.Lock()
+_ready = True
+
+
+def mark_not_ready() -> None:
+    """Report not-ready — call before running the startup warmup."""
+    global _ready
+    with _lock:
+        _ready = False
+
+
+def mark_ready() -> None:
+    """Report ready — call once the startup warmup has completed."""
+    global _ready
+    with _lock:
+        _ready = True
+
+
+def is_ready() -> bool:
+    """Whether the service is warmed up and can serve a real request quickly."""
+    with _lock:
+        return _ready

policyengine_api/warmup.py

@@ -0,0 +1,71 @@
+"""Warm the simulation machinery at startup so the first real request is fast.
+
+Building a tax-benefit system (at import) does not compile the per-simulation
+machinery, so the first calculate on a fresh worker pays a large one-time cost
+(~2 minutes for US on Cloud Run). Running a throwaway calculate here moves that
+off the first user request, so /readiness-check (via policyengine_api.readiness)
+only reports ready once a calculate is actually fast. Only
+POLICYENGINE_API_WARMUP_COUNTRIES (default "us") are warmed, to keep the added
+boot time inside the startup-probe window.
+"""
+
+from __future__ import annotations
+
+import copy
+import logging
+import os
+import time
+
+logger = logging.getLogger(__name__)
+
+# Minimal single-person household requesting household_net_income: computing that
+# broad output compiles most of the simulation graph. Period is arbitrary.
+_WARMUP_PERIOD = "2025"
+WARMUP_HOUSEHOLDS: dict[str, dict] = {
+    "us": {
+        "people": {"person": {"age": {_WARMUP_PERIOD: 40}}},
+        "households": {
+            "household": {
+                "members": ["person"],
+                "household_net_income": {_WARMUP_PERIOD: None},
+            }
+        },
+    },
+}
+
+
+def _requested_countries() -> list[str]:
+    raw = os.environ.get("POLICYENGINE_API_WARMUP_COUNTRIES", "us")
+    return [country.strip().lower() for country in raw.split(",") if country.strip()]
+
+
+def run_startup_warmup() -> None:
+    """Run a throwaway calculate per configured country (best-effort).
+
+    Failures are logged and swallowed: a warmup error must not stop the worker
+    from serving (the first real request would just be slow), and it must never
+    leave the service permanently unready.
+    """
+    from policyengine_api.country import COUNTRIES
+
+    for country_id in _requested_countries():
+        household = WARMUP_HOUSEHOLDS.get(country_id)
+        country = COUNTRIES.get(country_id)
+        if household is None:
+            logger.warning("No warmup household for %r; skipping", country_id)
+            continue
+        if country is None:
+            logger.warning("Country %r not loaded; skipping warmup", country_id)
+            continue
+
+        started = time.time()
+        try:
+            country.calculate(copy.deepcopy(household), {})
+            logger.info("Warmed up %s in %.1fs", country_id, time.time() - started)
+        except Exception:
+            logger.exception(
+                "Startup warmup calculate failed for %s after %.1fs; continuing "
+                "(first real request may be slow)",
+                country_id,
+                time.time() - started,
+            )

tests/integration/test_warmup_household.py

@@ -0,0 +1,69 @@
+"""Integration guard for the startup-warmup households.
+
+Runs the real ``policyengine_us`` simulation for each configured warmup household
+to prove it is a valid situation that computes a finite result — i.e. that
+``run_startup_warmup`` will actually warm the machinery rather than silently
+throw and leave the service unwarmed.
+
+Slow (~2 min: it builds the US tax-benefit system), so it is gated behind
+``RUN_WARMUP_INTEGRATION=1`` and skipped in the default suite. The staging smoke
+suite is the other, live guard.
+"""
+
+import importlib
+import math
+import os
+
+import pytest
+
+from policyengine_api.warmup import WARMUP_HOUSEHOLDS
+
+pytestmark = pytest.mark.skipif(
+    os.environ.get("RUN_WARMUP_INTEGRATION") != "1",
+    reason="slow (builds the tax-benefit system); set RUN_WARMUP_INTEGRATION=1",
+)
+
+_COUNTRY_PACKAGES = {
+    "us": "policyengine_us",
+    "uk": "policyengine_uk",
+    "ca": "policyengine_canada",
+    "ng": "policyengine_ng",
+    "il": "policyengine_il",
+}
+
+
+def _requested_computations(household):
+    # Mirror policyengine_api.country.get_requested_computations: null leaves at
+    # entity_plural/entity_id/variable/period.
+    requested = []
+    for entities in household.values():
+        if not isinstance(entities, dict):
+            continue
+        for variables in entities.values():
+            if not isinstance(variables, dict):
+                continue
+            for variable, periods in variables.items():
+                if not isinstance(periods, dict):
+                    continue
+                for period, value in periods.items():
+                    if value is None:
+                        requested.append((variable, period))
+    return requested
+
+
+@pytest.mark.parametrize("country_id", sorted(WARMUP_HOUSEHOLDS))
+def test_warmup_household_builds_and_computes_finite(country_id):
+    household = WARMUP_HOUSEHOLDS[country_id]
+    requested = _requested_computations(household)
+    assert requested, f"{country_id} warmup household requests no computation"
+
+    package = importlib.import_module(_COUNTRY_PACKAGES[country_id])
+    system = package.CountryTaxBenefitSystem()
+    simulation = package.Simulation(tax_benefit_system=system, situation=household)
+
+    for variable, period in requested:
+        result = simulation.calculate(variable, period)
+        assert len(result) >= 1
+        assert all(math.isfinite(float(value)) for value in result), (
+            f"{country_id}: {variable}@{period} produced a non-finite value"
+        )

tests/unit/test_cloud_run_deploy_scripts.py

@@ -445,13 +445,13 @@ def test_deploy_cloud_run_candidate_pins_http_startup_probe():
     assert threshold * period <= 240, "failureThreshold x periodSeconds > 240s cap"
     assert initial_delay <= 240, "initialDelaySeconds > 240s cap"
     # initialDelaySeconds is additive (no probe runs during it), so the real
-    # deadline is the sum. Keep it wide enough to cover the measured p90 boot
-    # (371s at time of writing) — see cloud-run-operations.md for the data.
-    assert initial_delay + threshold * period >= 340
-    # ...but initialDelaySeconds also delays availability, since no probe can
-    # succeed before it elapses. Keep it under the measured p50 boot (201s) so
-    # it does not needlessly slow down every scale-out.
-    assert initial_delay <= 200
+    # deadline is the sum. Readiness now gates on the startup warmup too, so
+    # boot-to-ready is the ~371s p90 import PLUS the warmup calculate; keep the
+    # window at (or near) the 480s platform maximum — see cloud-run-operations.md.
+    assert initial_delay + threshold * period >= 470
+    # initialDelaySeconds also delays availability, but boot (import + warmup)
+    # exceeds the p50 either way, so a high initialDelay adds no real scale-out
+    # delay.
     assert int(settings["timeoutSeconds"]) <= period
 
 

tests/unit/test_readiness.py

@@ -0,0 +1,21 @@
+import pytest
+
+from policyengine_api import readiness
+
+
+@pytest.fixture(autouse=True)
+def _restore_ready():
+    # readiness state is module-global; leave it ready for other tests.
+    yield
+    readiness.mark_ready()
+
+
+def test_defaults_to_ready():
+    assert readiness.is_ready() is True
+
+
+def test_mark_not_ready_then_ready():
+    readiness.mark_not_ready()
+    assert readiness.is_ready() is False
+    readiness.mark_ready()
+    assert readiness.is_ready() is True