Report job stage and progress on annual economy job status

[PavelMakarchuk]

Fixes #780

Same-repository draft PR check

• This PR is a draft.
• This PR is from a branch in PolicyEngine/policyengine-sim-api, not a personal fork.

Summary

GET /jobs/{job_id} reports only running, complete or failed, so clients estimate progress from elapsed time. This PR makes the worker record which stage an annual economy job is in and adds optional stage, progress and timing fields to the status response. Clients that read only status see no change.

Phase 2 of #780 (runtime statistics: p50/p90 by country, region type and execution path) is not in this PR.

Response shape

Running, monolithic (for example a state run):

{
  "status": "running",
  "result": null,
  "error": null,
  "resolved_app_name": "policyengine-simulation-py6-2-7",
  "policyengine_bundle": { "...": "..." },
  "submitted_at": "2026-10-13T12:00:00Z",
  "started_at": "2026-10-13T12:00:30Z",
  "stage": "reform",
  "stage_started_at": "2026-10-13T12:01:00Z",
  "progress": 67,
  "progress_max": 96,
  "stage_durations": { "loading_data": 12.5, "baseline": 17.5 }
}

Running, segmented US national: "stage": "simulating", "segments": {"completed": 10, "total": 20}, "progress": 5, "progress_max": 96.

Complete: "progress": 100, started_at, finished_at and the full stage_durations; no stage. Failed: stage is the stage the job stopped in; no progress.

Stages and where each starts

Stage	Starts at
queued	Gateway accepted the job; the worker has not written a record
loading_data	run_simulation entry (country model, region, dataset, simulation build)
baseline	First ensure() of the baseline simulation, inside economic_impact_analysis
reform	First ensure() of the reform simulation
simulating	Segmented national only: before the 20 segment children spawn; segments.completed rises as each one returns
outputs	After economic_impact_analysis returns (monolithic) or before the national reduce (segmented)

economic_impact_analysis calls baseline.ensure() and then reform.ensure() internally, so the runtime cannot observe the boundary from outside. StageReportingSimulation, a small Simulation subclass, calls back once when its first ensure() starts. ArtifactBaselineSimulation now extends it, and _build_simulation returns it on the plain path.

Progress values

progress is the percent at which the current stage starts in a typical successful production run. The values depend on the execution path, which the worker records as monolithic or segmented. They come from the production runtime analysis of 6–13 October 2026 (submission-to-completion time):

Stage	Monolithic (UK national, US state)	Segmented (US national)
queued	0	0
loading_data	25	3
baseline	55	—
reform	67	—
simulating	—	5
outputs	97	97

• Queue time is about 39% of a UK national run and 13% of a US state run, so the worker start is at 25 for monolithic runs.
• Outputs after economic_impact_analysis take almost no time (US state: 0.25 s), so outputs starts at 97.
• Main-path telemetry cannot split baseline from reform. The 55/67 split uses the Stage 12 state runs (baseline is about 29% of their calculation).
• Before the worker records its path (a few milliseconds after it starts), a stage uses the lowest floor of any path, so progress never moves backwards.
• A running job is capped at 99.
• Completed segments do not move progress. Anthony confirmed that the 20 regional children finish at about the same time and California finishes last. A segment count then jumps late and waits; a backtest on the October figures gave it 16–27 points of mean error. segments stays in the response for captions such as "19 of 20 regions done".

progress_max is the highest percent a running job reaches before its next stage (one below the next floor, or 99 for the last stage). For a smooth bar, a client moves its own time-based estimate between progress and progress_max and never moves the bar backwards. In the same backtest, this hybrid stayed at 3–15 points of mean error for normal runs, against 20–32 for the current fixed-duration bar on UK and US state runs. While a job is queued its path is unknown, so progress_max is the highest limit of any path (24).

Later versions may add stages. The schema and README tell clients to accept a stage value they do not know. Clients generated with strict enums from an older schema would reject a new value; API v1 reads plain dict keys and is not affected.

The values stay provisional until phase 2 of #780 measures them per country and region type.

Telemetry

run_simulation now closes its last stage inside its operation and sets job_stage_seconds_<stage> attributes, for example job_stage_seconds_baseline and job_stage_seconds_reform. These are the first main-path measurements that separate the baseline from the reform calculation. The attribute set is bounded by JobStage; docs/engineering/skills/observability.md documents it.

Compatibility and failure handling

• All new JobStatusResponse fields are optional and omitted when empty. New submissions add submitted_at; jobs submitted before deployment keep a byte-identical body.
• Versioned apps deployed before this change never write a record. Without a guard, the gateway would report their jobs as queued for the whole run. build_bundle_manifest_metadata now declares job_progress: 1; the gateway stores the routed app's capability in the job metadata and reports a stage only when it is present.
• Re-runs after preemption: Modal can run the same call again under the same call id. The new attempt continues the stored record: started_at and the ended-stage durations carry over, attempts counts up, and the published stage stays at the earlier attempt's stage until the new attempt reaches a stage at least as far along. progress therefore never moves backwards. run_simulation adds job_attempts to its telemetry.
• Worker reads and writes are non-fatal: a failure logs one warning (exception type only) and the calculation continues. A failed or malformed gateway read drops the stage fields and never changes the status, result or HTTP status code. The complete path stays inside the existing try, so result-validation failures still produce the redacted failed body.
• API v1 reads the job status with data.get(...), so the extra keys do not affect it.
• Deploy order does not matter. An old gateway ignores the new manifest key and the new Dict; a new gateway sees no capability on old manifests.

Changes

• Contract: job_progress.py (JobStage, JobProgressRecord, JobSegmentProgress, the provisional STAGE_PROGRESS_FLOOR, job_progress_fields) and job_progress_state.py (Modal Dict simulation-api-job-progress). JobStatusResponse gains the optional fields.
• Observability: registers modal_job_progress_read in the annual-impact plan; the gateway measures the read with it.
• Executor: StoredJobProgressReporter (created only in run_simulation, keyed by current_function_call_id()), the StageReportingSimulation hook, the execution path in dispatch_run_simulation, the outputs transition in SimulationOutputBuilder.analysis, segment counting in SegmentedNationalRunner, the stage-duration attributes, and the job_progress manifest key. Segment children, budget-window parents and Stage 12 use the no-op reporter. Budget-window children are ordinary run_simulation calls, so each writes a record under its own call id; /budget-window-jobs does not read those records yet.
• Gateway: stores submitted_at and the capability at submit, and adds the fields to the 200, 202 and 500 status responses. README.md documents the fields.

Testing

• Contract: 184 passed (39 new).
• Observability: 31 passed (1 new).
• Gateway: 227 passed (13 new), including the regenerated OpenAPI golden (new fields and descriptions only).
• Simulation entry: 158 passed; its contract test against the gateway golden passes unchanged.
• Executor: 767 passed, 23 skipped (34 new), including restart tests that check progress never decreases over all written records.
• ruff format --check passes for every src folder. ruff check on the changed source files reports no new findings compared with main (one existing import-order finding is fixed).
• Pyright: the new modules have no errors. The changed files show 5 more errors of two existing kinds: "possibly unbound" variables assigned inside with runtime.span(...) blocks, and modal.current_function_call_id, which app.py already calls.

Existing tests changed only at their seams: the plain-simulation test double now replaces StageReportingSimulation instead of policyengine.core.Simulation, two dispatch fakes accept progress, and the simulation doubles gained a no-op on_first_ensure.

Deployed verification: the beta deployment's integration tests (projects/policyengine-apis-integ/tests/simulation/test_calculate.py) submit real US national, US state and UK jobs. They now record every running and complete status body and check:

• while running: a known stage and progress ≤ progress_max ≤ 99;
• across polls: progress never moves backwards;
• on completion: progress 100, started_at, finished_at and stage_durations, and no stage.

A Stage 12-only deployment runs these tests against the existing gateway, which may predate stage reporting, so the check applies only when the gateway reports submitted_at (every gateway with this PR does). Six unit tests cover the check itself and run in PR CI.

Not run locally: no Modal deployment and no live economy job; the beta deployment runs the checks above.

Notes for review

• @anth-volk: the progress values above come from your runtime analysis. Please check that I read the stage shares correctly. You are checking whether the ~25 s of unattributed US state worker time is a step without telemetry coverage. The PR measures stages by wall-clock boundaries, so that time is counted inside whichever stage surrounds it, and the job_stage_seconds_* attributes will show which one.
• Phase 2 of Expose job stage and progress on annual economy job status #780: your analysis computed category percentiles from telemetry that already exists. A scheduled job that publishes p50/p90 per country, region type and execution path from that telemetry looks better than the Modal Dict ring buffer proposed in the issue. Do you agree?
• Stage 12 runs baseline and reform at the same time. If it becomes the authoritative path, its jobs fit simulating better than separate baseline and reform stages.
• libs/policyengine-simulation-contract/uv.lock on main still pins policyengine-observability 3.0.1, while feat: measure Stage 12 coordinator preparation #708 raised the floor to 3.0.2. uv sync rewrites it locally. This PR leaves the lockfile unchanged.

Downstream

• policyengine-api (API v1): Relay simulation run stage and progress on computing economy responses policyengine-api#3916 relays stage, progress, progress_max, the timestamps and segments on the economy computing response. It can merge before or after this PR.
• policyengine-app-v2: prefer the API progress in SocietyWideCalcStrategy and name the stage in the RunProgress caption. Keep the time-based estimate as the fallback.

🤖 Generated with Claude Code

[anth-volk]

Two comments from a review of the stage reporting: a possible progress regression when Modal runs the same call again, and a mismatch between the description and what budget-window children do.