diff --git a/.github/scripts/cloud_run_env.sh b/.github/scripts/cloud_run_env.sh index 8aad5620f..893934f83 100755 --- a/.github/scripts/cloud_run_env.sh +++ b/.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}" diff --git a/changelog.d/startup-warmup-readiness.fixed.md b/changelog.d/startup-warmup-readiness.fixed.md new file mode 100644 index 000000000..b18236383 --- /dev/null +++ b/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. diff --git a/docs/migration/cloud-run-operations.md b/docs/migration/cloud-run-operations.md index 666491a62..f3f477e6f 100644 --- a/docs/migration/cloud-run-operations.md +++ b/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 diff --git a/policyengine_api/api.py b/policyengine_api/api.py index 6e1e9a98d..e839d5744 100644 --- a/policyengine_api/api.py +++ b/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"}) diff --git a/policyengine_api/asgi.py b/policyengine_api/asgi.py index f54392915..901f60b59 100644 --- a/policyengine_api/asgi.py +++ b/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() diff --git a/policyengine_api/openapi_spec.yaml b/policyengine_api/openapi_spec.yaml index 805426bfd..3c68f55be 100644 --- a/policyengine_api/openapi_spec.yaml +++ b/policyengine_api/openapi_spec.yaml @@ -1069,7 +1069,10 @@ 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. @@ -1077,6 +1080,12 @@ paths: 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. diff --git a/policyengine_api/readiness.py b/policyengine_api/readiness.py new file mode 100644 index 000000000..fd56751a8 --- /dev/null +++ b/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 diff --git a/policyengine_api/warmup.py b/policyengine_api/warmup.py new file mode 100644 index 000000000..92200b95d --- /dev/null +++ b/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, + ) diff --git a/tests/integration/test_warmup_household.py b/tests/integration/test_warmup_household.py new file mode 100644 index 000000000..770c09992 --- /dev/null +++ b/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" + ) diff --git a/tests/unit/test_cloud_run_deploy_scripts.py b/tests/unit/test_cloud_run_deploy_scripts.py index 34d8b3e76..4a66ee3c9 100644 --- a/tests/unit/test_cloud_run_deploy_scripts.py +++ b/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 diff --git a/tests/unit/test_readiness.py b/tests/unit/test_readiness.py new file mode 100644 index 000000000..1f5f706c6 --- /dev/null +++ b/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 diff --git a/tests/unit/test_warmup.py b/tests/unit/test_warmup.py new file mode 100644 index 000000000..32aea3579 --- /dev/null +++ b/tests/unit/test_warmup.py @@ -0,0 +1,106 @@ +import sys +import types +from pathlib import Path + +from policyengine_api import warmup + +REPO = Path(__file__).resolve().parents[2] + + +def test_requested_countries_default(monkeypatch): + monkeypatch.delenv("POLICYENGINE_API_WARMUP_COUNTRIES", raising=False) + assert warmup._requested_countries() == ["us"] + + +def test_requested_countries_parses_env(monkeypatch): + monkeypatch.setenv("POLICYENGINE_API_WARMUP_COUNTRIES", "US, uk ,") + assert warmup._requested_countries() == ["us", "uk"] + + +def test_us_warmup_household_requests_a_broad_output(): + household = warmup.WARMUP_HOUSEHOLDS["us"] + assert household["people"] + net_income = household["households"]["household"]["household_net_income"] + # A null value marks it as a requested computation for calculate(). + assert None in net_income.values() + + +def test_all_warmup_household_members_reference_defined_people(): + # Cheap structural guard (no tax-benefit-system build): a typo'd member would + # make the real warmup calculate throw, which is swallowed and would silently + # leave the service unwarmed. The full validity check is the slow integration + # test in tests/integration/test_warmup_household.py. + for country_id, household in warmup.WARMUP_HOUSEHOLDS.items(): + people = set(household.get("people", {})) + assert people, f"{country_id} warmup household defines no people" + for entity_plural, entities in household.items(): + if entity_plural == "people" or not isinstance(entities, dict): + continue + for entity in entities.values(): + for member in entity.get("members", []): + assert member in people, ( + f"{country_id}: {entity_plural} member {member!r} is not a " + "defined person" + ) + + +def _inject_fake_countries(monkeypatch, countries): + module = types.ModuleType("policyengine_api.country") + module.COUNTRIES = countries + monkeypatch.setitem(sys.modules, "policyengine_api.country", module) + + +def test_run_startup_warmup_calculates_each_country(monkeypatch): + calls = [] + + class FakeCountry: + def calculate(self, household, reform): + calls.append((household, reform)) + + _inject_fake_countries(monkeypatch, {"us": FakeCountry()}) + monkeypatch.setenv("POLICYENGINE_API_WARMUP_COUNTRIES", "us") + + warmup.run_startup_warmup() + + assert len(calls) == 1 + household, reform = calls[0] + assert household["people"] + assert reform == {} + + +def test_run_startup_warmup_swallows_calculate_errors(monkeypatch): + class BoomCountry: + def calculate(self, household, reform): + raise RuntimeError("boom") + + _inject_fake_countries(monkeypatch, {"us": BoomCountry()}) + monkeypatch.setenv("POLICYENGINE_API_WARMUP_COUNTRIES", "us") + + # Must not raise: a warmup failure cannot stop the worker from serving. + warmup.run_startup_warmup() + + +def test_run_startup_warmup_skips_unknown_country(monkeypatch): + calls = [] + + class FakeCountry: + def calculate(self, household, reform): + calls.append(1) + + _inject_fake_countries(monkeypatch, {"us": FakeCountry()}) + monkeypatch.setenv("POLICYENGINE_API_WARMUP_COUNTRIES", "zz") + + warmup.run_startup_warmup() + assert calls == [] + + +def test_asgi_runs_warmup_and_marks_ready(): + src = (REPO / "policyengine_api/asgi.py").read_text(encoding="utf-8") + assert "run_startup_warmup" in src + assert "mark_ready" in src + + +def test_readiness_check_gates_on_readiness(): + src = (REPO / "policyengine_api/api.py").read_text(encoding="utf-8") + assert "is_ready" in src + assert "status=503" in src