Record report runs and simulation runs from the v2 Modal coordinator and workers

changelog.d/report-run-recording.added.md

@@ -0,0 +1 @@
+The v2 Modal report coordinator and its US and UK simulation workers can record a report's run in the Public API's run tables: when a call names a report run and its baseline and reform simulation runs, they write those runs' status, output file and economy results with the Modal worker role's credentials, from the Modal secret `policyengine-modal-worker-database`. The report run is marked succeeded in the same transaction as its result rows. On pre-emption the coordinator of such a call cancels its child calls and resumes when Modal restarts it. Comparison runs behave as before.

docs/engineering/skills/database-access.md

@@ -34,6 +34,11 @@ It builds on the Stage 12 rules:
   (`sim.stage12.no-migration-chain`). This repository's code reads and
   writes their rows, and deployment validation compares their live schema
   with its SQLAlchemy mappings; nothing here changes them.
+- policyengine-api also owns the run tables the v2 Modal coordinator and
+  workers write as the Modal worker role: `report_runs`, `simulation_runs`,
+  `report_run_simulation_runs`, `report_run_comparisons` and the result
+  tables. Their mappings in `policyengine_stage12_persistence.report_run_tables`
+  and `report_run_results` declare only the columns those writes touch.
 - **Once #847 merges**, this repository's own Alembic chain at `migrations/`
   owns the `simulation` schema, and the Stage 12 tables stay with
   policyengine-api.

libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_execution/__init__.py

@@ -38,6 +38,7 @@
     ComparisonSimulationRecord,
     ErrorCode,
     ErrorSummary,
+    ReportRunIds,
     ResultComparisonStatus,
     Stage12InvocationContext,
 )
@@ -95,6 +96,7 @@
     "ReportAggregate",
     "ReportExecutionInput",
     "ReportOutputRequirements",
+    "ReportRunIds",
     "ResultComparisonArtifactPayload",
     "ResultComparisonStatus",
     "ResultDifference",

libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_execution/comparison_runs.py

@@ -23,13 +23,49 @@
 )
 
 
+class ReportRunIds(StrictContractModel):
+    """The Public API's report run and simulation runs one report call records on.
+
+    The simulation service creates the three rows before it starts the
+    coordinator. The coordinator records the report run's status and results,
+    and each worker records its own simulation run's.
+    """
+
+    report_run_id: UUID
+    baseline_simulation_run_id: UUID
+    reform_simulation_run_id: UUID
+
+    @model_validator(mode="after")
+    def validate_distinct_runs(self) -> ReportRunIds:
+        """Require the report run and both simulation runs to be different rows."""
+        run_ids = {
+            self.report_run_id,
+            self.baseline_simulation_run_id,
+            self.reform_simulation_run_id,
+        }
+        if len(run_ids) != 3:  # noqa: PLR2004  # one report run and two simulation runs
+            msg = "report run IDs must name three different runs"
+            raise ValueError(msg)
+        return self
+
+    def simulation_run_id(self, role: SimulationRole) -> UUID:
+        """Return the simulation run of ``role``; a standalone role has none."""
+        if role is SimulationRole.BASELINE:
+            return self.baseline_simulation_run_id
+        if role is SimulationRole.REFORM:
+            return self.reform_simulation_run_id
+        msg = "a standalone simulation has no simulation run in a report run"
+        raise ValueError(msg)
+
+
 class Stage12InvocationContext(StrictContractModel):
     """Deployment and retention context for one comparison-run invocation.
 
     Records the originating request, the environments, the worker version and
     callable, the manifest digests and the artifact prefix. Automatic runs also
     name the production function call they are compared with. Artifacts are kept
-    for at most 30 days after ``created_at``.
+    for at most 30 days after ``created_at``. A call that also records its
+    outcome on the Public API's run tables names those runs in ``report_run``.
     """
 
     contract_version: Literal[1] = 1
@@ -51,13 +87,23 @@ class Stage12InvocationContext(StrictContractModel):
     production_function_call_id: ContractText | None = None
     created_at: datetime
     retention_expires_at: datetime
+    # Comparison runs leave it unset and record only the Stage 12 tables.
+    report_run: ReportRunIds | None = None
 
     @model_validator(mode="after")
     def validate_retention(self) -> Stage12InvocationContext:
         """Require aware timestamps and a retention window of at most 30 days."""
         _validate_retention(self.created_at, self.retention_expires_at)
         return self
 
+    @model_validator(mode="after")
+    def validate_report_run_is_direct(self) -> Stage12InvocationContext:
+        """Reject report runs on a shadow run compared with a production call."""
+        if self.report_run is not None and self.production_function_call_id:
+            msg = "a run compared with a production function call has no report run"
+            raise ValueError(msg)
+        return self
+
 
 ErrorCode = Annotated[str, Field(min_length=1, max_length=64)]
 ErrorSummary = Annotated[str, Field(min_length=1, max_length=512)]

libs/policyengine-simulation-contract/tests/test_stage12_report_run_ids.py

@@ -0,0 +1,96 @@
+"""The report run and simulation runs a Stage 12 call records its outcome on."""
+
+from __future__ import annotations
+
+from datetime import UTC, datetime, timedelta
+from typing import Any
+from uuid import UUID
+
+import pytest
+from policyengine_simulation_contract.stage12_execution import (
+    ReportRunIds,
+    SimulationRole,
+    Stage12InvocationContext,
+)
+from pydantic import ValidationError
+
+NOW = datetime(2026, 11, 3, tzinfo=UTC)
+REPORT_RUN_ID = UUID("00000000-0000-0000-0000-0000000000a1")
+BASELINE_RUN_ID = UUID("00000000-0000-0000-0000-0000000000b1")
+REFORM_RUN_ID = UUID("00000000-0000-0000-0000-0000000000c1")
+
+
+def _context_fields(**overrides: Any) -> dict[str, Any]:
+    return {
+        "request_id": "request-1",
+        "environment": "staging",
+        "modal_environment": "staging",
+        "worker_version": "6.12.0",
+        "modal_application": "policyengine-simulation-v2-py6-12-0",
+        "simulation_callable": "run_single_simulation_us",
+        "version_manifest_sha256": "a" * 64,
+        "bundle_manifest_sha256": "b" * 64,
+        "artifact_prefix": "stage-12-runs/staging/2026/11/report",
+        "created_at": NOW,
+        "retention_expires_at": NOW + timedelta(days=30),
+        **overrides,
+    }
+
+
+def _run_ids() -> ReportRunIds:
+    return ReportRunIds(
+        report_run_id=REPORT_RUN_ID,
+        baseline_simulation_run_id=BASELINE_RUN_ID,
+        reform_simulation_run_id=REFORM_RUN_ID,
+    )
+
+
+def test_context_without_run_ids_keeps_its_comparison_fields() -> None:
+    context = Stage12InvocationContext(**_context_fields())
+
+    assert context.report_run is None
+    assert "report_run" not in context.model_dump(mode="json", exclude_none=True)
+
+
+def test_context_carries_the_report_run_and_both_simulation_runs() -> None:
+    context = Stage12InvocationContext(
+        **_context_fields(report_run=_run_ids().model_dump(mode="json"))
+    )
+
+    assert context.report_run == _run_ids()
+    restored = Stage12InvocationContext.model_validate(context.model_dump(mode="json"))
+    assert restored.report_run == _run_ids()
+
+
+def test_each_role_names_its_own_simulation_run() -> None:
+    run_ids = _run_ids()
+
+    assert run_ids.simulation_run_id(SimulationRole.BASELINE) == BASELINE_RUN_ID
+    assert run_ids.simulation_run_id(SimulationRole.REFORM) == REFORM_RUN_ID
+    with pytest.raises(ValueError, match="standalone"):
+        run_ids.simulation_run_id(SimulationRole.STANDALONE)
+
+
+@pytest.mark.parametrize(
+    "repeated_ids",
+    [
+        {"baseline_simulation_run_id": REPORT_RUN_ID},
+        {"reform_simulation_run_id": REPORT_RUN_ID},
+        {"reform_simulation_run_id": BASELINE_RUN_ID},
+    ],
+)
+def test_run_ids_must_name_three_different_runs(repeated_ids) -> None:
+    fields = {**_run_ids().model_dump(), **repeated_ids}
+
+    with pytest.raises(ValidationError, match="three different runs"):
+        ReportRunIds(**fields)
+
+
+def test_a_shadow_comparison_run_cannot_record_report_runs() -> None:
+    with pytest.raises(ValidationError, match="production function call"):
+        Stage12InvocationContext(
+            **_context_fields(
+                production_function_call_id="fc-production",
+                report_run=_run_ids().model_dump(mode="json"),
+            )
+        )

libs/policyengine-stage12-persistence/README.md

@@ -1,6 +1,26 @@
 # PolicyEngine Stage 12 persistence
 
-Temporary SQLAlchemy Core mappings and transactional persistence for the two
-Stage 12 comparison tables. `policyengine-api` remains the sole owner of the
-SQLModel definitions, Alembic revisions, and live schema. This package performs
-DML only and is removed with the temporary comparison path during Stage 14.
+SQLAlchemy Core mappings and transactional DML for two sets of tables, both
+owned by `policyengine-api`, which keeps the SQLModel definitions, Alembic
+revisions and live schema. This package never changes a schema.
+
+- **The Stage 12 comparison tables** (`Stage12PersistenceStore`): temporary
+  state of the comparison path, removed with it.
+- **The Public API's run tables** (`ReportRunStore`): `report_runs`,
+  `simulation_runs`, `report_run_simulation_runs`, the result tables, and
+  `report_run_comparisons`, which keeps a compared pair's results that no
+  result table holds. The v2 Modal coordinator and workers write them as the
+  Modal worker role (`policyengine_v2_modal_worker`), for a call that names a
+  report run. The mappings declare only the columns those writes touch. Every
+  status write is one `UPDATE` conditional on the run's status; a report
+  run's completion inserts its result rows in the same transaction.
+
+## Tests
+
+The report run tests need a disposable Postgres migrated by the Public API's
+v2 chain, with its declared privileges applied. They skip unless both URLs
+are set:
+
+- `REPORT_RUN_TEST_DATABASE_URL`: logs in as `policyengine_v2_modal_worker`;
+- `REPORT_RUN_TEST_SEED_DATABASE_URL`: logs in as a role that may insert the
+  rows the simulation service would create.

libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/__init__.py

@@ -1,5 +1,15 @@
-"""Direct DML for temporary Stage 12 comparison state."""
+"""Direct DML for temporary Stage 12 comparison state and for report runs."""
 
+from policyengine_stage12_persistence.report_run_results import (
+    DistrictCodes,
+    ReportRunResult,
+    report_run_result,
+)
+from policyengine_stage12_persistence.report_run_store import (
+    ReportRunStore,
+    RunStatus,
+    RunWrite,
+)
 from policyengine_stage12_persistence.store import (
     ConnectionOpener,
     Stage12PersistenceStore,
@@ -13,9 +23,15 @@
 
 __all__ = [
     "ConnectionOpener",
+    "DistrictCodes",
+    "ReportRunResult",
+    "ReportRunStore",
+    "RunStatus",
+    "RunWrite",
     "Stage12PersistenceStore",
     "Stage12RuntimeAccessError",
     "create_stage12_engine",
     "normalized_database_url",
+    "report_run_result",
     "validate_runtime_database",
 ]