From c14cfd0e48f95c890e7c2277007314fa0a49fa3b Mon Sep 17 00:00:00 2001 From: PavelMakarchuk Date: Tue, 13 Oct 2026 10:41:37 -0400 Subject: [PATCH 1/7] Report job stage and progress on annual economy job status The worker records which stage an annual job is in (loading_data, baseline, reform or simulating, outputs) in a new Modal Dict keyed by job_id. GET /jobs/{job_id} adds optional stage, progress, timestamps, segment counts and per-stage durations. Clients that read only status see no change. The gateway reports a stage only for jobs routed to a worker whose bundle manifest declares job_progress, so jobs on older versioned apps are never shown as queued for their whole run. Store failures and malformed records drop the stage fields and never change the status, result or HTTP status code. Fixes #780 Co-Authored-By: Claude Opus 5.5 --- .../gateway_models.py | 44 ++ .../job_progress.py | 175 ++++++ .../job_progress_state.py | 24 + .../tests/conftest.py | 5 +- .../tests/test_job_progress.py | 247 ++++++++ .../stages.py | 2 + .../tests/test_stages.py | 6 + .../fixtures/spm_doubles.py | 1 + .../src/modal/app.py | 29 + .../src/modal/segmented_national.py | 21 +- .../modal/utils/update_version_registry.py | 7 + .../baseline_artifacts.py | 7 +- .../job_progress.py | 161 +++++ .../simulation_output_builder.py | 8 + .../simulation_runtime.py | 26 +- .../stage_reporting_simulation.py | 35 ++ .../tests/test_baseline_artifacts.py | 17 +- .../tests/test_canonical_spm.py | 8 +- .../tests/test_job_progress.py | 563 ++++++++++++++++++ .../tests/test_segmented_national.py | 4 +- .../tests/test_simulation_output_builder.py | 4 +- .../tests/test_update_version_registry.py | 21 + .../policyengine-simulation-gateway/README.md | 32 + .../fixtures/gateway_endpoints.py | 7 +- .../endpoints.py | 87 ++- .../tests/golden/openapi.json | 140 +++++ .../tests/test_job_progress.py | 293 +++++++++ 27 files changed, 1943 insertions(+), 31 deletions(-) create mode 100644 libs/policyengine-simulation-contract/src/policyengine_simulation_contract/job_progress.py create mode 100644 libs/policyengine-simulation-contract/src/policyengine_simulation_contract/job_progress_state.py create mode 100644 libs/policyengine-simulation-contract/tests/test_job_progress.py create mode 100644 projects/policyengine-simulation-executor/src/policyengine_simulation_executor/job_progress.py create mode 100644 projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage_reporting_simulation.py create mode 100644 projects/policyengine-simulation-executor/tests/test_job_progress.py create mode 100644 projects/policyengine-simulation-gateway/tests/test_job_progress.py diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/gateway_models.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/gateway_models.py index 97c16a11b..342dfbd62 100644 --- a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/gateway_models.py +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/gateway_models.py @@ -6,6 +6,7 @@ from typing import ClassVar, Literal, Optional from pydantic import ( + AwareDatetime, BaseModel, ConfigDict, Field, @@ -14,6 +15,10 @@ model_validator, ) +from policyengine_simulation_contract.job_progress import ( + JobSegmentProgress, + JobStage, +) from policyengine_simulation_contract.spm import ( SPMSelection, SPMCapability, @@ -194,6 +199,45 @@ class JobStatusResponse(BaseModel): error: Optional[str] = None resolved_app_name: Optional[str] = None policyengine_bundle: Optional[PolicyEngineBundle] = None + # Stage reporting. Present only for jobs whose worker reports stages; + # see policyengine_simulation_contract.job_progress. + stage: JobStage | None = Field( + default=None, + description=( + "Current stage of a running job, or the stage in which a failed " + "job stopped. Absent for complete jobs." + ), + ) + progress: int | None = Field( + default=None, + ge=0, + le=100, + description=( + "Coarse completion percent. It advances at stage boundaries and " + "as national segments complete, stays below 100 while running, " + "and is 100 when complete." + ), + ) + submitted_at: AwareDatetime | None = Field( + default=None, description="When the gateway accepted the job." + ) + started_at: AwareDatetime | None = Field( + default=None, description="When a worker started the job." + ) + stage_started_at: AwareDatetime | None = Field( + default=None, description="When the current stage started." + ) + finished_at: AwareDatetime | None = Field( + default=None, description="When the worker finished or stopped the job." + ) + segments: JobSegmentProgress | None = Field( + default=None, + description="Region segments of a segmented US national run.", + ) + stage_durations: dict[JobStage, float] | None = Field( + default=None, + description="Seconds spent in each stage that has ended.", + ) class BudgetWindowBatchRequest(GatewayRequestBase): diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/job_progress.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/job_progress.py new file mode 100644 index 000000000..a076b256e --- /dev/null +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/job_progress.py @@ -0,0 +1,175 @@ +"""Client-visible stage and progress of an annual economy job. + +The worker that runs a job records which stage it is in. The gateway reads +that record on every status poll and adds the optional stage fields to the +``JobStatusResponse``. Clients that read only ``status`` see no change. + +Records are keyed by the job's Modal function-call id, which is the public +``job_id``. Older worker apps do not write records, so the gateway reports a +stage only for jobs routed to a worker whose bundle manifest declares +``JOB_PROGRESS_CAPABILITY_VERSION``. Without that declaration a missing +record would be indistinguishable from a job that has not started. +""" + +from __future__ import annotations + +from collections.abc import Mapping +from datetime import datetime +from enum import StrEnum +from types import MappingProxyType +from typing import Any, Literal + +from pydantic import ( + AwareDatetime, + BaseModel, + Field, + NonNegativeFloat, + TypeAdapter, + ValidationError, + model_validator, +) + +JOB_PROGRESS_CAPABILITY_KEY = "job_progress" +JOB_PROGRESS_CAPABILITY_VERSION = 1 + + +class JobStage(StrEnum): + """Coarse stages of an annual economy job, in the order they run. + + A monolithic run moves through ``loading_data``, ``baseline``, ``reform`` + and ``outputs``. A segmented US national run computes baseline and reform + together in parallel region segments, so it moves through + ``loading_data``, ``simulating`` and ``outputs``. + """ + + QUEUED = "queued" + LOADING_DATA = "loading_data" + BASELINE = "baseline" + REFORM = "reform" + SIMULATING = "simulating" + OUTPUTS = "outputs" + + +class JobSegmentProgress(BaseModel): + """Completed region segments of a segmented national run.""" + + completed: int = Field(ge=0) + total: int = Field(ge=1) + + @model_validator(mode="after") + def _completed_within_total(self) -> JobSegmentProgress: + if self.completed > self.total: + raise ValueError("completed segments cannot exceed total segments") + return self + + +class JobProgressRecord(BaseModel): + """Worker-owned stage state of one job. Internal; not a public schema.""" + + schema_version: Literal[1] = 1 + stage: JobStage + started_at: AwareDatetime + stage_started_at: AwareDatetime + finished_at: AwareDatetime | None = None + segments: JobSegmentProgress | None = None + stage_durations: dict[JobStage, NonNegativeFloat] = Field(default_factory=dict) + + +# Percent of a typical run that is complete when each stage starts. +# PROVISIONAL: chosen by hand, not measured. Replace these with measured +# stage shares once per-stage runtime statistics exist (#780, phase 2). +# ``simulating`` advances from its floor to the ``outputs`` floor as +# segments complete. +STAGE_PROGRESS_FLOOR: Mapping[JobStage, int] = MappingProxyType( + { + JobStage.QUEUED: 0, + JobStage.LOADING_DATA: 5, + JobStage.BASELINE: 15, + JobStage.SIMULATING: 15, + JobStage.REFORM: 45, + JobStage.OUTPUTS: 85, + } +) +# A running job never reports 100: only a complete result does. +RUNNING_PROGRESS_CAP = 99 + +JobStatus = Literal["running", "complete", "failed"] + +_DATETIME = TypeAdapter(AwareDatetime) + + +def supported_job_progress_capability(value: object) -> int | None: + """Return the declared capability version when this gateway reads it.""" + + if isinstance(value, bool) or not isinstance(value, int): + return None + if value != JOB_PROGRESS_CAPABILITY_VERSION: + return None + return value + + +def _json_datetime(value: datetime | None) -> str | None: + return None if value is None else _DATETIME.dump_python(value, mode="json") + + +def _parse_datetime(value: object) -> datetime | None: + if value is None: + return None + try: + return _DATETIME.validate_python(value) + except ValidationError: + return None + + +def _running_progress(record: JobProgressRecord) -> int: + floor = STAGE_PROGRESS_FLOOR[record.stage] + if record.stage is JobStage.SIMULATING and record.segments is not None: + ceiling = STAGE_PROGRESS_FLOOR[JobStage.OUTPUTS] + share = record.segments.completed / record.segments.total + floor += int((ceiling - floor) * share) + return min(floor, RUNNING_PROGRESS_CAP) + + +def job_progress_fields( + *, + status: JobStatus, + submitted_at: object, + record: JobProgressRecord | None, + reports_stages: bool, +) -> dict[str, Any]: + """Build the optional stage fields of a job status response. + + Returns JSON-ready values and omits every field that has no value, so a + job without stage reporting keeps its existing response body. + + ``reports_stages`` is true only when the job's worker declares the stage + capability and the gateway read its record (or found none) without an + error. Then a running job without a record is ``queued``. + """ + + fields: dict[str, Any] = { + "submitted_at": _json_datetime(_parse_datetime(submitted_at)) + } + if reports_stages: + if status == "complete": + fields["progress"] = 100 + if record is None: + if status == "running": + fields["stage"] = JobStage.QUEUED.value + fields["progress"] = STAGE_PROGRESS_FLOOR[JobStage.QUEUED] + else: + fields["started_at"] = _json_datetime(record.started_at) + fields["finished_at"] = _json_datetime(record.finished_at) + if status != "complete": + fields["stage"] = record.stage.value + fields["stage_started_at"] = _json_datetime(record.stage_started_at) + if status == "running": + fields["progress"] = _running_progress(record) + if record.segments is not None: + fields["segments"] = record.segments.model_dump(mode="json") + if record.stage_durations: + fields["stage_durations"] = { + stage.value: seconds + for stage, seconds in record.stage_durations.items() + } + return {key: value for key, value in fields.items() if value is not None} diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/job_progress_state.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/job_progress_state.py new file mode 100644 index 000000000..d538cdc7e --- /dev/null +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/job_progress_state.py @@ -0,0 +1,24 @@ +"""Modal Dict storage for annual job stage records.""" + +from __future__ import annotations + +import modal + +from policyengine_simulation_contract.job_progress import JobProgressRecord + +JOB_PROGRESS_DICT_NAME = "simulation-api-job-progress" + + +def _job_progress_store(): + return modal.Dict.from_name(JOB_PROGRESS_DICT_NAME, create_if_missing=True) + + +def get_job_progress(job_id: str) -> JobProgressRecord | None: + payload = _job_progress_store().get(job_id) + if payload is None: + return None + return JobProgressRecord.model_validate(payload) + + +def put_job_progress(job_id: str, record: JobProgressRecord) -> None: + _job_progress_store()[job_id] = record.model_dump(mode="json") diff --git a/libs/policyengine-simulation-contract/tests/conftest.py b/libs/policyengine-simulation-contract/tests/conftest.py index 1d30dfbde..0f71b7c95 100644 --- a/libs/policyengine-simulation-contract/tests/conftest.py +++ b/libs/policyengine-simulation-contract/tests/conftest.py @@ -7,7 +7,7 @@ import pytest -from policyengine_simulation_contract import budget_window_state +from policyengine_simulation_contract import budget_window_state, job_progress_state class MockDict: @@ -29,7 +29,7 @@ def get(self, key, default=None): @pytest.fixture def mock_modal(monkeypatch): - """Patch modal.Dict in budget_window_state with an in-memory store.""" + """Patch modal.Dict in the Dict-backed state modules with an in-memory store.""" mock_dicts: dict[str, dict] = {} class MockModalDict: @@ -45,4 +45,5 @@ class MockModal: Dict = MockModalDict monkeypatch.setattr(budget_window_state, "modal", MockModal) + monkeypatch.setattr(job_progress_state, "modal", MockModal) return {"dicts": mock_dicts} diff --git a/libs/policyengine-simulation-contract/tests/test_job_progress.py b/libs/policyengine-simulation-contract/tests/test_job_progress.py new file mode 100644 index 000000000..14a981831 --- /dev/null +++ b/libs/policyengine-simulation-contract/tests/test_job_progress.py @@ -0,0 +1,247 @@ +"""Tests for client-visible job stage fields.""" + +from datetime import UTC, datetime + +import pytest +from pydantic import ValidationError + +from policyengine_simulation_contract.gateway_models import JobStatusResponse +from policyengine_simulation_contract.job_progress import ( + JOB_PROGRESS_CAPABILITY_VERSION, + RUNNING_PROGRESS_CAP, + STAGE_PROGRESS_FLOOR, + JobProgressRecord, + JobSegmentProgress, + JobStage, + job_progress_fields, + supported_job_progress_capability, +) +from policyengine_simulation_contract.job_progress_state import ( + JOB_PROGRESS_DICT_NAME, + get_job_progress, + put_job_progress, +) + +SUBMITTED = "2026-10-13T12:00:00+00:00" +STARTED = datetime(2026, 10, 13, 12, 0, 30, tzinfo=UTC) +STAGE_STARTED = datetime(2026, 10, 13, 12, 1, 0, tzinfo=UTC) +FINISHED = datetime(2026, 10, 13, 12, 4, 0, tzinfo=UTC) + + +def _record(**overrides) -> JobProgressRecord: + values = { + "stage": JobStage.REFORM, + "started_at": STARTED, + "stage_started_at": STAGE_STARTED, + "stage_durations": {JobStage.LOADING_DATA: 12.5, JobStage.BASELINE: 17.5}, + } + return JobProgressRecord(**{**values, **overrides}) + + +class TestJobProgressFieldsWithoutStageReporting: + def test__legacy_job_adds_no_fields(self): + assert ( + job_progress_fields( + status="running", + submitted_at=None, + record=None, + reports_stages=False, + ) + == {} + ) + + def test__only_submission_time_is_reported(self): + fields = job_progress_fields( + status="running", + submitted_at=SUBMITTED, + record=_record(), + reports_stages=False, + ) + + assert fields == {"submitted_at": "2026-10-13T12:00:00Z"} + + def test__malformed_submission_time_is_dropped(self): + assert ( + job_progress_fields( + status="running", + submitted_at="not a time", + record=None, + reports_stages=False, + ) + == {} + ) + + +class TestJobProgressFieldsWithStageReporting: + def test__running_job_without_record_is_queued(self): + fields = job_progress_fields( + status="running", + submitted_at=SUBMITTED, + record=None, + reports_stages=True, + ) + + assert fields == { + "submitted_at": "2026-10-13T12:00:00Z", + "stage": "queued", + "progress": 0, + } + + def test__running_job_reports_stage_floor_and_timing(self): + fields = job_progress_fields( + status="running", + submitted_at=SUBMITTED, + record=_record(), + reports_stages=True, + ) + + assert fields == { + "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": STAGE_PROGRESS_FLOOR[JobStage.REFORM], + "stage_durations": {"loading_data": 12.5, "baseline": 17.5}, + } + + @pytest.mark.parametrize( + ("completed", "expected"), + [(0, 15), (10, 50), (19, 81), (20, 85)], + ) + def test__segmented_run_advances_with_completed_segments(self, completed, expected): + record = _record( + stage=JobStage.SIMULATING, + segments=JobSegmentProgress(completed=completed, total=20), + ) + + fields = job_progress_fields( + status="running", + submitted_at=SUBMITTED, + record=record, + reports_stages=True, + ) + + assert fields["stage"] == "simulating" + assert fields["progress"] == expected + assert fields["segments"] == {"completed": completed, "total": 20} + + def test__running_progress_never_reaches_100(self): + assert all( + floor <= RUNNING_PROGRESS_CAP for floor in STAGE_PROGRESS_FLOOR.values() + ) + record = _record(stage=JobStage.OUTPUTS) + + fields = job_progress_fields( + status="running", + submitted_at=SUBMITTED, + record=record, + reports_stages=True, + ) + + assert fields["progress"] < 100 + + def test__complete_job_reports_100_and_no_current_stage(self): + record = _record(stage=JobStage.OUTPUTS, finished_at=FINISHED) + + fields = job_progress_fields( + status="complete", + submitted_at=SUBMITTED, + record=record, + reports_stages=True, + ) + + assert fields["progress"] == 100 + assert fields["finished_at"] == "2026-10-13T12:04:00Z" + assert "stage" not in fields + assert "stage_started_at" not in fields + + def test__complete_job_without_record_still_reports_100(self): + fields = job_progress_fields( + status="complete", + submitted_at=SUBMITTED, + record=None, + reports_stages=True, + ) + + assert fields == {"submitted_at": "2026-10-13T12:00:00Z", "progress": 100} + + def test__failed_job_reports_the_stage_it_stopped_in(self): + record = _record(finished_at=FINISHED) + + fields = job_progress_fields( + status="failed", + submitted_at=SUBMITTED, + record=record, + reports_stages=True, + ) + + assert fields["stage"] == "reform" + assert fields["finished_at"] == "2026-10-13T12:04:00Z" + assert "progress" not in fields + + def test__failed_job_without_record_reports_no_stage(self): + fields = job_progress_fields( + status="failed", + submitted_at=SUBMITTED, + record=None, + reports_stages=True, + ) + + assert fields == {"submitted_at": "2026-10-13T12:00:00Z"} + + @pytest.mark.parametrize("status", ["running", "complete", "failed"]) + def test__fields_are_valid_public_status_fields(self, status): + record = _record( + stage=JobStage.SIMULATING, + segments=JobSegmentProgress(completed=3, total=20), + finished_at=FINISHED if status != "running" else None, + ) + fields = job_progress_fields( + status=status, + submitted_at=SUBMITTED, + record=record, + reports_stages=True, + ) + + response = JobStatusResponse(status=status, **fields) + + assert response.model_dump(mode="json", exclude_none=True) == { + "status": status, + **fields, + } + + +class TestSupportedJobProgressCapability: + def test__current_version_is_supported(self): + assert ( + supported_job_progress_capability(JOB_PROGRESS_CAPABILITY_VERSION) + == JOB_PROGRESS_CAPABILITY_VERSION + ) + + @pytest.mark.parametrize("value", [None, True, "1", 0, 2, 1.0]) + def test__other_values_are_unsupported(self, value): + assert supported_job_progress_capability(value) is None + + +class TestJobProgressModels: + def test__completed_segments_cannot_exceed_total(self): + with pytest.raises(ValidationError): + JobSegmentProgress(completed=21, total=20) + + def test__record_requires_timezone_aware_times(self): + with pytest.raises(ValidationError): + _record(started_at=datetime(2026, 10, 13, 12, 0, 30)) + + +class TestJobProgressState: + def test__record_round_trips_through_the_store(self, mock_modal): + record = _record(segments=JobSegmentProgress(completed=2, total=20)) + + put_job_progress("fc-job-123", record) + + assert get_job_progress("fc-job-123") == record + stored = mock_modal["dicts"][JOB_PROGRESS_DICT_NAME]["fc-job-123"] + assert stored["stage"] == "reform" + + def test__missing_record_is_none(self, mock_modal): + assert get_job_progress("fc-unknown") is None diff --git a/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/stages.py b/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/stages.py index d923c3379..df1c9e0dd 100644 --- a/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/stages.py +++ b/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/stages.py @@ -57,6 +57,7 @@ class Stage(StrEnum): SEGMENTED_NATIONAL_CHILD_SPAWN = "segmented_national_child_spawn" SEGMENTED_NATIONAL_REDUCE = "segmented_national_reduce" MODAL_JOB_METADATA_READ = "modal_job_metadata_read" + MODAL_JOB_PROGRESS_READ = "modal_job_progress_read" SIMULATION_OUTPUT_BUILD = "simulation_output_build" SIMULATION_OUTPUT_MODEL_DUMP = "simulation_output_model_dump" @@ -153,6 +154,7 @@ def name(self, stage: Stage) -> str: Stage.MODAL_JOB_METADATA_WRITE, Stage.MODAL_JOB_METADATA_READ, Stage.MODAL_JOB_STATUS_POLL, + Stage.MODAL_JOB_PROGRESS_READ, Stage.MODAL_DICT_READ, *_COMMON_SIMULATION, ) diff --git a/libs/policyengine-simulation-observability/tests/test_stages.py b/libs/policyengine-simulation-observability/tests/test_stages.py index 9481c909b..1978f025c 100644 --- a/libs/policyengine-simulation-observability/tests/test_stages.py +++ b/libs/policyengine-simulation-observability/tests/test_stages.py @@ -49,3 +49,9 @@ def test_stage12_report_plans_include_output_planning_breakdown( assert Stage.STAGE12_OUTPUT_CONFIGURATION in stages assert Stage.STAGE12_OUTPUT_VARIABLE_RESOLUTION in stages assert Stage.STAGE12_CHILD_INPUT_PLANNING in stages + + +def test_annual_status_polls_measure_the_job_progress_read() -> None: + annual = RUN_STAGE_REGISTRY[RunConfiguration.ANNUAL_IMPACT] + + assert annual.name(Stage.MODAL_JOB_PROGRESS_READ) == "modal_job_progress_read" diff --git a/projects/policyengine-simulation-executor/fixtures/spm_doubles.py b/projects/policyengine-simulation-executor/fixtures/spm_doubles.py index 790e2ac9b..e8d3ff8e1 100644 --- a/projects/policyengine-simulation-executor/fixtures/spm_doubles.py +++ b/projects/policyengine-simulation-executor/fixtures/spm_doubles.py @@ -33,6 +33,7 @@ def spm_receipt(selection, year="2026"): def spm_capable_simulation(selection, year="2026", **attributes): """Return a stand-in for a simulation that measured ``selection``.""" + attributes.setdefault("on_first_ensure", lambda callback: None) if selection is None: return SimpleNamespace(**attributes) return SimpleNamespace( diff --git a/projects/policyengine-simulation-executor/src/modal/app.py b/projects/policyengine-simulation-executor/src/modal/app.py index 4e96b2276..cda8a108c 100644 --- a/projects/policyengine-simulation-executor/src/modal/app.py +++ b/projects/policyengine-simulation-executor/src/modal/app.py @@ -258,6 +258,30 @@ def _set_modal_call_attributes(runtime) -> None: pass +def _job_progress_reporter(): + """A stage reporter for this call, keyed by its public ``job_id``. + + The gateway's ``job_id`` is the spawned call's ``object_id``, which is + this call's ``current_function_call_id()``. Without one there is nothing + to key the record by, so nothing is reported. + """ + from policyengine_simulation_contract.job_progress_state import ( + put_job_progress, + ) + from policyengine_simulation_executor.job_progress import ( + NULL_JOB_PROGRESS, + StoredJobProgressReporter, + ) + + try: + job_id = modal.current_function_call_id() + except Exception: # noqa: BLE001 - stage reporting is non-fatal + job_id = None + if not job_id: + return NULL_JOB_PROGRESS + return StoredJobProgressReporter(job_id, write=put_job_progress) + + @app.function( image=simulation_image, cpu=8.0, @@ -277,7 +301,10 @@ def run_simulation( Imports the snapshotted implementation at runtime. Emits redacted operation data through the shared observability runtime. + Records the job's stage for status polls (see ``job_progress``). """ + progress = _job_progress_reporter() + progress.start() runtime = _configure_modal_observability(service_role="simulation_worker") # We deliberately avoid sending full ``params`` or ``result`` blobs to @@ -309,8 +336,10 @@ def run_simulation( params, app_name=APP_NAME, runtime=runtime, + progress=progress, ) finally: + progress.finish() runtime.shutdown() diff --git a/projects/policyengine-simulation-executor/src/modal/segmented_national.py b/projects/policyengine-simulation-executor/src/modal/segmented_national.py index 9e5ed5cc2..863bdecd0 100644 --- a/projects/policyengine-simulation-executor/src/modal/segmented_national.py +++ b/projects/policyengine-simulation-executor/src/modal/segmented_national.py @@ -25,6 +25,7 @@ from typing import Any from policyengine_observability import ObservabilityRuntime +from policyengine_simulation_contract.job_progress import JobStage from policyengine_simulation_contract.spm import ( SPMInputError, combine_spm_results, @@ -38,6 +39,10 @@ from policyengine_simulation_observability.telemetry import split_internal_payload import modal +from policyengine_simulation_executor.job_progress import ( + NULL_JOB_PROGRESS, + JobProgressReporter, +) from policyengine_simulation_executor.national_partition import ( national_region_groups, ) @@ -145,10 +150,12 @@ def __init__( modal_module=None, poll_interval_seconds: float = POLL_INTERVAL_INITIAL_SECONDS, poll_interval_max_seconds: float = POLL_INTERVAL_MAX_SECONDS, + progress: JobProgressReporter = NULL_JOB_PROGRESS, ): self.params = params self.app_name = app_name self.runtime = runtime + self.progress = progress self.modal = modal if modal_module is None else modal_module self.poll_interval_initial_seconds = poll_interval_seconds self.poll_interval_max_seconds = poll_interval_max_seconds @@ -171,8 +178,10 @@ def __init__( def run(self) -> dict[str, Any]: country_module = self._country_module() self._extend_last_group_with_uncovered_states(country_module) + self.progress.start_segments(len(self.groups)) handles = self._spawn_all_children() child_results = self._collect(handles) + self.progress.enter(JobStage.OUTPUTS) return self._reduce(child_results, country_module=country_module) def _country_module(self): @@ -283,6 +292,7 @@ def _collect(self, handles: list[tuple[list[str], Any]]) -> list[dict]: pending.discard(index) poll_errors.pop(index, None) progress_made = True + self.progress.segment_completed() # Bounded aggregate instead of one segment per probe (see the # Stage registry note about long poll loops). self.runtime.set_context(segmented_poll_count=poll_count) @@ -339,12 +349,17 @@ def _reduce(self, child_results: list[dict], *, country_module) -> dict[str, Any def run_segmented_national_impl( - params: dict[str, Any], *, app_name: str, runtime: ObservabilityRuntime + params: dict[str, Any], + *, + app_name: str, + runtime: ObservabilityRuntime, + progress: JobProgressReporter = NULL_JOB_PROGRESS, ) -> dict[str, Any]: return SegmentedNationalRunner( params, app_name=app_name, runtime=runtime, + progress=progress, ).run() @@ -353,6 +368,7 @@ def dispatch_run_simulation( *, app_name: str, runtime: ObservabilityRuntime, + progress: JobProgressReporter = NULL_JOB_PROGRESS, ) -> dict[str, Any]: """The run_simulation entrypoint's routing: segmented national fan-out for eligible requests, the monolithic path for everything else.""" @@ -366,6 +382,7 @@ def dispatch_run_simulation( params, app_name=app_name, runtime=runtime, + progress=progress, ) from policyengine_simulation_executor.simulation_runtime import ( @@ -376,4 +393,4 @@ def dispatch_run_simulation( # and UK requests must not pollute the segmented-vs-monolithic metric. if is_plain_national_macro(params): runtime.set_context(national_execution="monolithic") - return run_simulation_impl(params, runtime=runtime) + return run_simulation_impl(params, runtime=runtime, progress=progress) diff --git a/projects/policyengine-simulation-executor/src/modal/utils/update_version_registry.py b/projects/policyengine-simulation-executor/src/modal/utils/update_version_registry.py index 649a652a6..bcec9b95d 100644 --- a/projects/policyengine-simulation-executor/src/modal/utils/update_version_registry.py +++ b/projects/policyengine-simulation-executor/src/modal/utils/update_version_registry.py @@ -39,6 +39,9 @@ class BundleManifestMetadata(TypedDict): us: CountryBundleMetadata uk: CountryBundleMetadata spm: NotRequired[dict] + # Apps deployed before stage reporting have no key; the gateway then + # reports no stage for their jobs. + job_progress: NotRequired[int] class LatestVersions(TypedDict): @@ -105,6 +108,9 @@ def build_bundle_manifest_metadata( app_name: str, policyengine_version: str, ) -> BundleManifestMetadata: + from policyengine_simulation_contract.job_progress import ( + JOB_PROGRESS_CAPABILITY_VERSION, + ) from policyengine_simulation_executor.spm import runtime_spm_capability capability = runtime_spm_capability() @@ -113,6 +119,7 @@ def build_bundle_manifest_metadata( "policyengine_version": policyengine_version, "us": _country_bundle_metadata("us"), "uk": _country_bundle_metadata("uk"), + "job_progress": JOB_PROGRESS_CAPABILITY_VERSION, } if capability is not None: metadata["spm"] = capability.model_dump(mode="json") diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/baseline_artifacts.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/baseline_artifacts.py index a2af47817..bf53844fa 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/baseline_artifacts.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/baseline_artifacts.py @@ -25,9 +25,10 @@ from pydantic import PrivateAttr -from policyengine.core import Simulation - from policyengine_simulation_executor import artifact_keys +from policyengine_simulation_executor.stage_reporting_simulation import ( + StageReportingSimulation, +) logger = logging.getLogger(__name__) @@ -140,7 +141,7 @@ def deterministic_baseline_id( return None if identity is None else identity.simulation_id -class ArtifactBaselineSimulation(Simulation): +class ArtifactBaselineSimulation(StageReportingSimulation): """A Simulation whose ``ensure()`` validates what it loads.""" _computed_this_process: bool = PrivateAttr(default=False) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/job_progress.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/job_progress.py new file mode 100644 index 000000000..7f00654c7 --- /dev/null +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/job_progress.py @@ -0,0 +1,161 @@ +"""Record the stage of an annual job for clients that poll its status. + +The runtime calls a ``JobProgressReporter`` at stage boundaries. The +gateway-facing Modal function creates a ``StoredJobProgressReporter`` for its +own job; every other caller (segment children, Stage 12, tests) gets the +no-op ``NULL_JOB_PROGRESS``. + +Stage reporting is diagnostic for clients and must never change a +calculation: every write failure is caught and logged, and the job goes on. +""" + +from __future__ import annotations + +import logging +from collections.abc import Callable +from datetime import UTC, datetime +from typing import Protocol + +from policyengine_simulation_contract.job_progress import ( + JobProgressRecord, + JobSegmentProgress, + JobStage, +) + +logger = logging.getLogger(__name__) + +ProgressWriter = Callable[[str, JobProgressRecord], None] + + +def _utc_now() -> datetime: + return datetime.now(UTC) + + +class JobProgressReporter(Protocol): + def start(self) -> None: ... + + def enter(self, stage: JobStage) -> None: ... + + def start_segments(self, total: int) -> None: ... + + def segment_completed(self) -> None: ... + + def finish(self) -> None: ... + + +class NullJobProgressReporter: + """Reports nothing.""" + + def start(self) -> None: + pass + + def enter(self, stage: JobStage) -> None: + pass + + def start_segments(self, total: int) -> None: + pass + + def segment_completed(self) -> None: + pass + + def finish(self) -> None: + pass + + +NULL_JOB_PROGRESS = NullJobProgressReporter() + + +class StoredJobProgressReporter: + """Writes the job's ``JobProgressRecord`` after every change.""" + + def __init__( + self, + job_id: str, + *, + write: ProgressWriter, + clock: Callable[[], datetime] = _utc_now, + ): + self.job_id = job_id + self._write = write + self._clock = clock + self._record: JobProgressRecord | None = None + self._write_failed = False + + @property + def record(self) -> JobProgressRecord | None: + return self._record + + def start(self) -> None: + if self._record is not None: + return + now = self._clock() + self._record = JobProgressRecord( + stage=JobStage.LOADING_DATA, + started_at=now, + stage_started_at=now, + ) + self._publish() + + def enter(self, stage: JobStage) -> None: + if self._record is None: + self.start() + record = self._record + if record is None or record.finished_at is not None: + return + if record.stage is stage: + return + now = self._close_stage(record) + record.stage = stage + record.stage_started_at = now + self._publish() + + def start_segments(self, total: int) -> None: + self.enter(JobStage.SIMULATING) + record = self._record + if record is None or record.finished_at is not None or total < 1: + return + record.segments = JobSegmentProgress(completed=0, total=total) + self._publish() + + def segment_completed(self) -> None: + record = self._record + if record is None or record.finished_at is not None: + return + segments = record.segments + if segments is None or segments.completed >= segments.total: + return + record.segments = JobSegmentProgress( + completed=segments.completed + 1, + total=segments.total, + ) + self._publish() + + def finish(self) -> None: + """Close the current stage. Runs after success and after failure.""" + + record = self._record + if record is None or record.finished_at is not None: + return + record.finished_at = self._close_stage(record) + self._publish() + + def _close_stage(self, record: JobProgressRecord) -> datetime: + now = self._clock() + elapsed = max((now - record.stage_started_at).total_seconds(), 0.0) + previous = record.stage_durations.get(record.stage, 0.0) + record.stage_durations[record.stage] = round(previous + elapsed, 2) + return now + + def _publish(self) -> None: + if self._record is None: + return + try: + self._write(self.job_id, self._record.model_copy(deep=True)) + except Exception as exc: # noqa: BLE001 - stage reporting is non-fatal + if not self._write_failed: + logger.warning( + "Could not record job progress for %s: %s", + self.job_id, + type(exc).__name__, + ) + self._write_failed = True diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_builder.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_builder.py index 5805c8bee..3d6ffc68b 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_builder.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_builder.py @@ -7,6 +7,7 @@ from typing import Any from policyengine_observability import ObservabilityRuntime +from policyengine_simulation_contract.job_progress import JobStage from policyengine_simulation_contract.uk_geography import UKLocalAuthorityMetadata from policyengine_simulation_executor import simulation_output_budget @@ -20,6 +21,10 @@ ANNUAL_IMPACT_STAGES, Stage, ) +from policyengine_simulation_executor.job_progress import ( + NULL_JOB_PROGRESS, + JobProgressReporter, +) from policyengine_simulation_executor.release_bundle import get_country_release_bundle from policyengine_simulation_executor.simulation_macro_output import ( BudgetaryImpact, @@ -51,6 +56,7 @@ class SimulationOutputBuilder: resolved_region_code: str | None = None runtime: ObservabilityRuntime | None = None uk_local_authority_metadata: UKLocalAuthorityMetadata | None = None + progress: JobProgressReporter = NULL_JOB_PROGRESS _analysis: Any = field(default=None, init=False) def __post_init__(self) -> None: @@ -68,6 +74,8 @@ def analysis(self) -> Any: self.reform, include_cliff_impacts=self._include_cliff_impacts(), ) + # Both simulations have run; the rest is output aggregation. + self.progress.enter(JobStage.OUTPUTS) return self._analysis def _include_cliff_impacts(self) -> bool: diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py index 5018b0f79..c1963edd0 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py @@ -21,10 +21,15 @@ from policyengine_observability import ObservabilityRuntime from policyengine_simulation_contract.dataset_uri import runtime_dataset_uri +from policyengine_simulation_contract.job_progress import JobStage from policyengine_simulation_observability.stages import ( ANNUAL_IMPACT_STAGES, Stage, ) +from policyengine_simulation_executor.job_progress import ( + NULL_JOB_PROGRESS, + JobProgressReporter, +) from policyengine_simulation_executor.release_bundle import ( get_country_release_bundle, resolve_bundle_dataset_name, @@ -150,6 +155,7 @@ def run_simulation_impl( params: dict, *, runtime: ObservabilityRuntime, + progress: JobProgressReporter = NULL_JOB_PROGRESS, ) -> dict: """ Execute economic simulation. @@ -161,7 +167,7 @@ def run_simulation_impl( # cleaned up on exit so we never leave signed JSON material on disk. with setup_gcp_credentials(runtime): try: - return _run_simulation_impl_core(params, runtime=runtime) + return _run_simulation_impl_core(params, runtime=runtime, progress=progress) except ValueError as exc: # Gateway images intentionally do not install country packages. # Transport the public error through the shared contract instead. @@ -683,12 +689,13 @@ def _build_simulation( scoping_strategy=None, region_code: str | None = None, ): - from policyengine.core import Simulation - from policyengine_simulation_executor.baseline_artifacts import ( ArtifactBaselineSimulation, deterministic_baseline_id, ) + from policyengine_simulation_executor.stage_reporting_simulation import ( + StageReportingSimulation, + ) from policyengine_simulation_executor.spm import normalize_runtime_spm @@ -721,7 +728,7 @@ def _build_simulation( policy=policy, scoping_strategy=scoping_strategy, ) - return Simulation( + return StageReportingSimulation( **spm_kwargs, dataset=dataset, tax_benefit_model_version=country_module.model, @@ -730,10 +737,17 @@ def _build_simulation( ) +def _report_calculation_start( + simulation, progress: JobProgressReporter, stage: JobStage +) -> None: + simulation.on_first_ensure(lambda: progress.enter(stage)) + + def _run_simulation_impl_core( params: dict, *, runtime: ObservabilityRuntime, + progress: JobProgressReporter = NULL_JOB_PROGRESS, ) -> dict: with runtime.span(ANNUAL_IMPACT_STAGES.name(Stage.REQUEST_PARSE)): simulation_params, telemetry, metadata = split_internal_payload(params) @@ -823,6 +837,9 @@ def _run_simulation_impl_core( region_code=region_resolution.code, ) + _report_calculation_start(baseline, progress, JobStage.BASELINE) + _report_calculation_start(reform, progress, JobStage.REFORM) + logger.info("Calculating economic impact") builder = SimulationOutputBuilder( country=country, @@ -835,6 +852,7 @@ def _run_simulation_impl_core( resolved_region_code=region_resolution.code, runtime=runtime, uk_local_authority_metadata=uk_local_authority_metadata, + progress=progress, ) output = builder.serialize() output.update( diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage_reporting_simulation.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage_reporting_simulation.py new file mode 100644 index 000000000..47c75b08c --- /dev/null +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage_reporting_simulation.py @@ -0,0 +1,35 @@ +"""A Simulation that reports when its calculation starts. + +``economic_impact_analysis`` configures both simulations and then calls +``baseline.ensure()`` before ``reform.ensure()``. The runtime cannot observe +that boundary from outside the call, so each simulation reports the start of +its own first ``ensure()``. Later ``ensure()`` calls (the decile impacts call +it again) report nothing. +""" + +from __future__ import annotations + +from collections.abc import Callable + +from policyengine.core import Simulation +from pydantic import PrivateAttr + + +class StageReportingSimulation(Simulation): + _on_first_ensure: Callable[[], None] | None = PrivateAttr(default=None) + + def on_first_ensure(self, callback: Callable[[], None]) -> None: + """Call ``callback`` once, when the first ``ensure()`` starts. + + ``callback`` must not raise: it runs inside the calculation. + """ + + self._on_first_ensure = callback + + def ensure(self) -> None: + callback = self._on_first_ensure + # Clear before the call so cached copies never carry the callback. + self._on_first_ensure = None + if callback is not None: + callback() + super().ensure() diff --git a/projects/policyengine-simulation-executor/tests/test_baseline_artifacts.py b/projects/policyengine-simulation-executor/tests/test_baseline_artifacts.py index 6f9fdfaf2..b4bedf9d9 100644 --- a/projects/policyengine-simulation-executor/tests/test_baseline_artifacts.py +++ b/projects/policyengine-simulation-executor/tests/test_baseline_artifacts.py @@ -350,9 +350,10 @@ def test_second_ensure_keeps_an_incomplete_outcome(self, fresh_cache): class TestBuildSimulationSelection: @pytest.fixture def wired(self, monkeypatch): - from policyengine import core as policyengine_core - from policyengine_simulation_executor import simulation_runtime as sr + from policyengine_simulation_executor import ( + stage_reporting_simulation as srs, + ) recorded = SimpleNamespace( artifact_kwargs=None, @@ -370,15 +371,14 @@ class FakePlainSimulation: def __init__(self, **kwargs): recorded.plain_kwargs = kwargs - # Replacing ``Simulation`` also replaces a class inspected by the SPM - # capability resolver. Preserve the capability derived from the real - # installed class before installing the test double. + # Preserve the SPM capability derived from the real installed + # Simulation class before installing the test doubles. from policyengine_simulation_executor import spm as executor_spm installed_capability = executor_spm.runtime_spm_capability() monkeypatch.setattr(ba, "ArtifactBaselineSimulation", FakeArtifactSimulation) - monkeypatch.setattr(policyengine_core, "Simulation", FakePlainSimulation) + monkeypatch.setattr(srs, "StageReportingSimulation", FakePlainSimulation) monkeypatch.setattr( executor_spm, "runtime_spm_capability", @@ -459,6 +459,7 @@ def test_bounded_reform_build_keeps_current_law_artifact_reuse(monkeypatch): from policyengine_simulation_executor import simulation_runtime as sr from policyengine_simulation_executor import spm as executor_spm + from policyengine_simulation_executor import stage_reporting_simulation as srs installed_capability = executor_spm.runtime_spm_capability() monkeypatch.setattr( @@ -470,8 +471,8 @@ def test_bounded_reform_build_keeps_current_law_artifact_reuse(monkeypatch): lambda **kwargs: SimpleNamespace(kind="artifact", **kwargs), ) monkeypatch.setattr( - policyengine_core, - "Simulation", + srs, + "StageReportingSimulation", lambda **kwargs: SimpleNamespace(kind="plain", **kwargs), ) identities = [] diff --git a/projects/policyengine-simulation-executor/tests/test_canonical_spm.py b/projects/policyengine-simulation-executor/tests/test_canonical_spm.py index b091a9a23..9bc6e3f73 100644 --- a/projects/policyengine-simulation-executor/tests/test_canonical_spm.py +++ b/projects/policyengine-simulation-executor/tests/test_canonical_spm.py @@ -601,8 +601,8 @@ def test_fanout_and_budget_window_preserve_settings_and_receipts(): def test_worker_build_passes_selection_to_baseline_and_reform(monkeypatch): - import policyengine.core from policyengine_simulation_executor import baseline_artifacts + from policyengine_simulation_executor import stage_reporting_simulation import policyengine_simulation_executor.spm as spm calls = [] @@ -618,7 +618,11 @@ def test_worker_build_passes_selection_to_baseline_and_reform(monkeypatch): monkeypatch.setattr( baseline_artifacts, "ArtifactBaselineSimulation", lambda **k: calls.append(k) ) - monkeypatch.setattr(policyengine.core, "Simulation", lambda **k: calls.append(k)) + monkeypatch.setattr( + stage_reporting_simulation, + "StageReportingSimulation", + lambda **k: calls.append(k), + ) for policy in (None, {"test": 123}): simulation_runtime._build_simulation( {"country": "us", "spm": SELECTION}, diff --git a/projects/policyengine-simulation-executor/tests/test_job_progress.py b/projects/policyengine-simulation-executor/tests/test_job_progress.py new file mode 100644 index 000000000..cb83a38a0 --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_job_progress.py @@ -0,0 +1,563 @@ +"""Tests for job stage reporting in the simulation worker.""" + +from __future__ import annotations + +import importlib +import logging +import sys +from datetime import UTC, datetime, timedelta +from types import SimpleNamespace + +import pytest +from conftest import NoOpObservabilityRuntime +from fixtures.fake_modal import install_fake_modal +from fixtures.spm_doubles import installed_spm_selection, spm_capable_simulation +from fixtures.test_simulation_api_contracts import CURRENT_SINGLE_YEAR_MACRO_RESULT +from policyengine_simulation_contract.job_progress import ( + JobProgressRecord, + JobSegmentProgress, + JobStage, +) + +from policyengine_simulation_executor import simulation_runtime as sr +from policyengine_simulation_executor.job_progress import ( + NULL_JOB_PROGRESS, + StoredJobProgressReporter, +) +from policyengine_simulation_executor.release_bundle import ( + get_bundled_package_version, +) +from policyengine_simulation_executor.simulation_output_builder import ( + SimulationOutputBuilder, +) +from policyengine_simulation_executor.simulation_runtime import ( + RegionResolution, + _run_simulation_impl_core, +) +from src.modal import segmented_national as sn + +T0 = datetime(2026, 10, 13, 12, 0, 0, tzinfo=UTC) + + +class _Clock: + def __init__(self): + self.now = T0 + + def __call__(self) -> datetime: + return self.now + + def advance(self, seconds: float) -> None: + self.now += timedelta(seconds=seconds) + + +class _RecordingWriter: + def __init__(self): + self.writes: list[tuple[str, JobProgressRecord]] = [] + + def __call__(self, job_id: str, record: JobProgressRecord) -> None: + self.writes.append((job_id, record)) + + @property + def last(self) -> JobProgressRecord: + return self.writes[-1][1] + + +class _RecordingReporter: + """Records reporter calls in order.""" + + def __init__(self): + self.calls: list[tuple] = [] + + def start(self) -> None: + self.calls.append(("start",)) + + def enter(self, stage: JobStage) -> None: + self.calls.append(("enter", stage)) + + def start_segments(self, total: int) -> None: + self.calls.append(("start_segments", total)) + + def segment_completed(self) -> None: + self.calls.append(("segment_completed",)) + + def finish(self) -> None: + self.calls.append(("finish",)) + + +def _reporter(): + clock = _Clock() + writer = _RecordingWriter() + reporter = StoredJobProgressReporter("fc-job-1", write=writer, clock=clock) + return reporter, writer, clock + + +class TestStoredJobProgressReporter: + def test__monolithic_lifecycle_records_each_stage_duration(self): + reporter, writer, clock = _reporter() + + reporter.start() + clock.advance(10) + reporter.enter(JobStage.BASELINE) + clock.advance(20) + reporter.enter(JobStage.REFORM) + clock.advance(30) + reporter.enter(JobStage.OUTPUTS) + clock.advance(5) + reporter.finish() + + assert [record.stage for _, record in writer.writes] == [ + JobStage.LOADING_DATA, + JobStage.BASELINE, + JobStage.REFORM, + JobStage.OUTPUTS, + JobStage.OUTPUTS, + ] + assert {job_id for job_id, _ in writer.writes} == {"fc-job-1"} + final = writer.last + assert final.started_at == T0 + assert final.stage_started_at == T0 + timedelta(seconds=60) + assert final.finished_at == T0 + timedelta(seconds=65) + assert final.stage_durations == { + JobStage.LOADING_DATA: 10.0, + JobStage.BASELINE: 20.0, + JobStage.REFORM: 30.0, + JobStage.OUTPUTS: 5.0, + } + + def test__written_records_are_snapshots(self): + reporter, writer, clock = _reporter() + + reporter.start() + clock.advance(10) + reporter.enter(JobStage.BASELINE) + + assert writer.writes[0][1].stage is JobStage.LOADING_DATA + assert writer.writes[0][1].stage_durations == {} + + def test__entering_the_current_stage_writes_nothing(self): + reporter, writer, _ = _reporter() + + reporter.start() + reporter.enter(JobStage.LOADING_DATA) + reporter.start() + + assert len(writer.writes) == 1 + + def test__enter_before_start_starts_the_record(self): + reporter, writer, _ = _reporter() + + reporter.enter(JobStage.BASELINE) + + assert [record.stage for _, record in writer.writes] == [ + JobStage.LOADING_DATA, + JobStage.BASELINE, + ] + + def test__segments_count_up_to_their_total(self): + reporter, writer, _ = _reporter() + + reporter.start() + reporter.start_segments(2) + reporter.segment_completed() + reporter.segment_completed() + reporter.segment_completed() + + assert writer.last.stage is JobStage.SIMULATING + assert writer.last.segments == JobSegmentProgress(completed=2, total=2) + assert len(writer.writes) == 5 + + def test__segment_completion_without_segments_writes_nothing(self): + reporter, writer, _ = _reporter() + + reporter.start() + reporter.segment_completed() + + assert len(writer.writes) == 1 + + def test__nothing_changes_after_finish(self): + reporter, writer, clock = _reporter() + + reporter.start() + clock.advance(3) + reporter.finish() + reporter.finish() + reporter.enter(JobStage.REFORM) + reporter.segment_completed() + + assert len(writer.writes) == 2 + assert writer.last.stage is JobStage.LOADING_DATA + assert writer.last.stage_durations == {JobStage.LOADING_DATA: 3.0} + + def test__finish_without_start_writes_nothing(self): + reporter, writer, _ = _reporter() + + reporter.finish() + + assert writer.writes == [] + + def test__write_failures_are_logged_once_and_never_raised(self, caplog): + calls = [] + + def failing_write(job_id, record): + calls.append(record.stage) + raise ConnectionError("modal unavailable") + + reporter = StoredJobProgressReporter("fc-job-1", write=failing_write) + + with caplog.at_level(logging.WARNING): + reporter.start() + reporter.enter(JobStage.BASELINE) + reporter.finish() + + assert calls == [JobStage.LOADING_DATA, JobStage.BASELINE, JobStage.BASELINE] + warnings = [ + record for record in caplog.records if "job progress" in record.message + ] + assert len(warnings) == 1 + assert "ConnectionError" in warnings[0].message + assert "modal unavailable" not in warnings[0].message + + +class TestStageReportingSimulation: + def test__callback_runs_once_before_the_first_ensure(self, monkeypatch): + from policyengine.core import Simulation + + from policyengine_simulation_executor.stage_reporting_simulation import ( + StageReportingSimulation, + ) + + events = [] + monkeypatch.setattr(Simulation, "ensure", lambda self: events.append("ensure")) + simulation = StageReportingSimulation.model_construct() + simulation.on_first_ensure(lambda: events.append("callback")) + + simulation.ensure() + simulation.ensure() + + assert events == ["callback", "ensure", "ensure"] + + def test__callback_is_cleared_before_the_calculation_runs(self, monkeypatch): + from policyengine.core import Simulation + + from policyengine_simulation_executor.stage_reporting_simulation import ( + StageReportingSimulation, + ) + + seen = [] + monkeypatch.setattr( + Simulation, + "ensure", + lambda self: seen.append(self._on_first_ensure), + ) + simulation = StageReportingSimulation.model_construct() + simulation.on_first_ensure(lambda: None) + + simulation.ensure() + + assert seen == [None] + + def test__artifact_baseline_reports_its_first_ensure(self): + from policyengine_simulation_executor.baseline_artifacts import ( + ArtifactBaselineSimulation, + ) + from policyengine_simulation_executor.stage_reporting_simulation import ( + StageReportingSimulation, + ) + + assert issubclass(ArtifactBaselineSimulation, StageReportingSimulation) + + +class TestSimulationOutputBuilderStage: + def _builder(self, analysis, progress): + country_module = SimpleNamespace(economic_impact_analysis=analysis) + return SimulationOutputBuilder( + country="us", + simulation_params={}, + country_module=country_module, + dataset=None, + baseline="baseline", + reform="reform", + progress=progress, + ) + + def test__outputs_stage_starts_after_the_impact_analysis(self): + progress = _RecordingReporter() + calls = [] + + def analysis(baseline, reform, include_cliff_impacts): + calls.append(list(progress.calls)) + return "analysis" + + builder = self._builder(analysis, progress) + + assert builder.analysis == "analysis" + assert builder.analysis == "analysis" + assert calls == [[]] + assert progress.calls == [("enter", JobStage.OUTPUTS)] + + def test__failed_analysis_does_not_enter_outputs(self): + progress = _RecordingReporter() + + def analysis(baseline, reform, include_cliff_impacts): + raise RuntimeError("model failed") + + builder = self._builder(analysis, progress) + + with pytest.raises(RuntimeError): + builder.analysis + assert progress.calls == [] + + +def test_runtime_reports_baseline_then_reform_when_each_calculation_starts( + monkeypatch, +): + selection = installed_spm_selection() + callbacks = {} + + def simulation_double(role): + return spm_capable_simulation( + selection, + on_first_ensure=lambda callback: callbacks.__setitem__(role, callback), + ) + + simulations = iter([simulation_double("baseline"), simulation_double("reform")]) + progress = _RecordingReporter() + builders = [] + + class FakeSimulationOutputBuilder: + def __init__(self, **kwargs): + builders.append(kwargs) + + def serialize(self): + # economic_impact_analysis ensures the baseline, then the reform. + callbacks["baseline"]() + callbacks["reform"]() + return dict(CURRENT_SINGLE_YEAR_MACRO_RESULT) + + monkeypatch.setattr( + sr, + "_country_module", + lambda country: SimpleNamespace(model=SimpleNamespace(version="1.715.2")), + ) + monkeypatch.setattr( + sr, + "_resolve_region", + lambda **kwargs: RegionResolution(code="us", dataset_reference=None), + ) + monkeypatch.setattr( + sr, "_load_dataset", lambda params, country_module, selection: object() + ) + monkeypatch.setattr(sr, "_build_simulation", lambda *a, **k: next(simulations)) + monkeypatch.setattr(sr, "SimulationOutputBuilder", FakeSimulationOutputBuilder) + + _run_simulation_impl_core( + {"country": "us", "scope": "macro", "reform": {}}, + runtime=NoOpObservabilityRuntime(), + progress=progress, + ) + + assert progress.calls == [ + ("enter", JobStage.BASELINE), + ("enter", JobStage.REFORM), + ] + assert builders[0]["progress"] is progress + + +class _FakeCall: + def __init__(self, result, polls_until_ready=0, error=None): + self.result = result + self.polls_until_ready = polls_until_ready + self.error = error + self.cancelled = False + + def get(self, timeout): + if self.error is not None: + raise self.error + if self.polls_until_ready > 0: + self.polls_until_ready -= 1 + raise TimeoutError + return self.result + + def cancel(self): + self.cancelled = True + + +def _segmented_runner(calls, progress, monkeypatch): + monkeypatch.setattr( + sr, + "_country_module", + lambda country: SimpleNamespace( + model=SimpleNamespace(get_region=lambda code: None, region_registry=None) + ), + ) + remaining = iter(calls) + fake_modal = SimpleNamespace( + Function=SimpleNamespace( + from_name=lambda app_name, function_name: SimpleNamespace( + spawn=lambda payload, *, observability_context=None: next(remaining) + ) + ) + ) + return sn.SegmentedNationalRunner( + {"country": "us", "scope": "macro", "time_period": "2026"}, + app_name="policyengine-simulation-py-test", + runtime=NoOpObservabilityRuntime(), + modal_module=fake_modal, + poll_interval_seconds=0.001, + poll_interval_max_seconds=0.002, + progress=progress, + ) + + +class TestSegmentedNationalStages: + def test__reports_segments_then_outputs_before_the_reduce(self, monkeypatch): + progress = _RecordingReporter() + calls = [_FakeCall({"child": i}, polls_until_ready=i % 3) for i in range(20)] + runner = _segmented_runner(calls, progress, monkeypatch) + reduce_saw = [] + monkeypatch.setattr( + runner, + "_reduce", + lambda results, country_module: reduce_saw.append(list(progress.calls)), + ) + + runner.run() + + assert progress.calls[0] == ("start_segments", len(runner.groups)) + assert progress.calls.count(("segment_completed",)) == 20 + assert progress.calls[-1] == ("enter", JobStage.OUTPUTS) + assert reduce_saw == [progress.calls] + + def test__child_failure_never_enters_outputs(self, monkeypatch): + progress = _RecordingReporter() + calls = [_FakeCall({"child": i}, polls_until_ready=5) for i in range(20)] + calls[4] = _FakeCall(None, error=ValueError("boom")) + runner = _segmented_runner(calls, progress, monkeypatch) + + with pytest.raises(RuntimeError): + runner.run() + + assert ("enter", JobStage.OUTPUTS) not in progress.calls + + def test__dispatch_passes_the_reporter_to_the_monolithic_path(self, monkeypatch): + progress = _RecordingReporter() + seen = {} + + def fake_impl(params, *, runtime, progress): + seen["progress"] = progress + return {"ok": True} + + monkeypatch.setattr(sr, "run_simulation_impl", fake_impl) + + sn.dispatch_run_simulation( + {"country": "us", "scope": "macro", "region": "state/ut"}, + app_name="app-x", + runtime=NoOpObservabilityRuntime(), + progress=progress, + ) + + assert seen["progress"] is progress + + +@pytest.fixture +def modal_app(monkeypatch): + install_fake_modal(monkeypatch) + for env_var, package in ( + ("POLICYENGINE_VERSION", "policyengine"), + ("POLICYENGINE_CORE_VERSION", "policyengine-core"), + ("POLICYENGINE_US_VERSION", "policyengine-us"), + ("POLICYENGINE_UK_VERSION", "policyengine-uk"), + ): + monkeypatch.setenv(env_var, get_bundled_package_version(package)) + sys.modules.pop("src.modal.app", None) + app = importlib.import_module("src.modal.app") + yield app + sys.modules.pop("src.modal.app", None) + + +class _ShutdownRuntime(NoOpObservabilityRuntime): + def __init__(self, events): + self.events = events + + def shutdown(self): + self.events.append("shutdown") + + +class TestModalRunSimulation: + def _wire(self, app, monkeypatch, dispatch): + from policyengine_simulation_contract import job_progress_state + + events = [] + writes = [] + monkeypatch.setattr( + app.modal, "current_function_call_id", lambda: "fc-job-1", raising=False + ) + monkeypatch.setattr( + job_progress_state, + "put_job_progress", + lambda job_id, record: writes.append((job_id, record)), + ) + monkeypatch.setattr( + app, + "_configure_modal_observability", + lambda *, service_role: _ShutdownRuntime(events), + ) + monkeypatch.setattr(sn, "dispatch_run_simulation", dispatch) + return events, writes + + def test__records_stages_for_the_spawned_job(self, modal_app, monkeypatch): + def dispatch(params, *, app_name, runtime, progress): + progress.enter(JobStage.REFORM) + return {"ok": True} + + events, writes = self._wire(modal_app, monkeypatch, dispatch) + + assert modal_app.run_simulation({"country": "us"}) == {"ok": True} + + assert {job_id for job_id, _ in writes} == {"fc-job-1"} + assert [record.stage for _, record in writes] == [ + JobStage.LOADING_DATA, + JobStage.REFORM, + JobStage.REFORM, + ] + assert writes[-1][1].finished_at is not None + assert events == ["shutdown"] + + def test__failed_job_still_records_where_it_stopped(self, modal_app, monkeypatch): + def dispatch(params, *, app_name, runtime, progress): + progress.enter(JobStage.BASELINE) + raise RuntimeError("model failed") + + _, writes = self._wire(modal_app, monkeypatch, dispatch) + + with pytest.raises(RuntimeError, match="model failed"): + modal_app.run_simulation({"country": "us"}) + + assert writes[-1][1].stage is JobStage.BASELINE + assert writes[-1][1].finished_at is not None + + def test__store_failure_does_not_change_the_result(self, modal_app, monkeypatch): + from policyengine_simulation_contract import job_progress_state + + def dispatch(params, *, app_name, runtime, progress): + progress.enter(JobStage.OUTPUTS) + return {"ok": True} + + self._wire(modal_app, monkeypatch, dispatch) + + def failing_write(job_id, record): + raise ConnectionError("modal unavailable") + + monkeypatch.setattr(job_progress_state, "put_job_progress", failing_write) + + assert modal_app.run_simulation({"country": "us"}) == {"ok": True} + + def test__no_call_id_means_no_reporting(self, modal_app, monkeypatch): + def missing_call_id(): + raise RuntimeError("not inside a Modal function call") + + monkeypatch.setattr( + modal_app.modal, "current_function_call_id", missing_call_id, raising=False + ) + + assert modal_app._job_progress_reporter() is NULL_JOB_PROGRESS diff --git a/projects/policyengine-simulation-executor/tests/test_segmented_national.py b/projects/policyengine-simulation-executor/tests/test_segmented_national.py index 547c202e3..f5d8e412f 100644 --- a/projects/policyengine-simulation-executor/tests/test_segmented_national.py +++ b/projects/policyengine-simulation-executor/tests/test_segmented_national.py @@ -349,7 +349,7 @@ class TestDispatchRunSimulation: def test__eligible_national_takes_the_segmented_path(self, monkeypatch): calls = {} - def fake_segmented(params, *, app_name, runtime): + def fake_segmented(params, *, app_name, runtime, progress): calls["segmented"] = (params, app_name) return {"segmented": True} @@ -523,7 +523,7 @@ def test__selection_is_resolved_before_the_fan_out(self, monkeypatch): ) seen = {} - def fake_segmented(params, *, app_name, runtime): + def fake_segmented(params, *, app_name, runtime, progress): seen["params"] = params return {"segmented": True} diff --git a/projects/policyengine-simulation-executor/tests/test_simulation_output_builder.py b/projects/policyengine-simulation-executor/tests/test_simulation_output_builder.py index 2e200af21..f0aa0b204 100644 --- a/projects/policyengine-simulation-executor/tests/test_simulation_output_builder.py +++ b/projects/policyengine-simulation-executor/tests/test_simulation_output_builder.py @@ -29,6 +29,7 @@ FakeModelOutput, fake_analysis, ) +from policyengine_simulation_executor.job_progress import NULL_JOB_PROGRESS from policyengine_simulation_executor.release_bundle import BUNDLE_RECEIPT_FILENAME from policyengine_simulation_executor.release_bundle import get_country_release_bundle from policyengine_simulation_executor.release_bundle import ( @@ -1131,6 +1132,7 @@ def serialize(self): "resolved_region_code": "us", "runtime": runtime, "uk_local_authority_metadata": None, + "progress": NULL_JOB_PROGRESS, } ] @@ -1246,7 +1248,7 @@ def serialize(self): ) monkeypatch.setattr( "policyengine_simulation_executor.simulation_runtime._build_simulation", - lambda *args, **kwargs: object(), + lambda *args, **kwargs: SimpleNamespace(on_first_ensure=lambda callback: None), ) monkeypatch.setattr( "policyengine_simulation_executor.simulation_runtime.SimulationOutputBuilder", diff --git a/projects/policyengine-simulation-executor/tests/test_update_version_registry.py b/projects/policyengine-simulation-executor/tests/test_update_version_registry.py index 9bd1544c0..4aef4eaea 100644 --- a/projects/policyengine-simulation-executor/tests/test_update_version_registry.py +++ b/projects/policyengine-simulation-executor/tests/test_update_version_registry.py @@ -148,6 +148,27 @@ def test_bundle_manifest_records_policyengine_6_spm_capability( } +def test_bundle_manifest_declares_job_progress_capability( + fake_bundle_metadata, + monkeypatch, +): + from policyengine_simulation_contract.job_progress import ( + JOB_PROGRESS_CAPABILITY_VERSION, + ) + + from policyengine_simulation_executor import spm as executor_spm + + monkeypatch.setattr(executor_spm, "runtime_spm_capability", lambda: None) + + metadata = registry.build_bundle_manifest_metadata( + app_name="policyengine-simulation-py6-2-7", + policyengine_version="6.2.7", + ) + + # The gateway reports a queued stage only for workers that declare this. + assert metadata["job_progress"] == JOB_PROGRESS_CAPABILITY_VERSION + + def test_validate_routing_state_rejects_missing_latest_route(fake_bundle_metadata): state = registry.build_next_routing_state( current_state=None, diff --git a/projects/policyengine-simulation-gateway/README.md b/projects/policyengine-simulation-gateway/README.md index 07ecdc0df..3b9516b78 100644 --- a/projects/policyengine-simulation-gateway/README.md +++ b/projects/policyengine-simulation-gateway/README.md @@ -37,6 +37,38 @@ A manifest that states *no* model version for the country — entry absent, or the value absent, non-string, empty or whitespace — contradicts nothing and still routes. +## Job stage reporting + +`GET /jobs/{job_id}` adds optional fields to annual economy jobs, so a +client can show real progress instead of a fixed duration estimate. Clients +that read only `status` see no change. + +| Field | Meaning | +|---|---| +| `stage` | `queued`, `loading_data`, `baseline`, `reform`, `simulating` or `outputs`. Running jobs report their current stage; failed jobs report the stage they stopped in; complete jobs omit it. | +| `progress` | Coarse percent, 0–100. It advances at stage boundaries and as national segments complete, stays below 100 while running, and is 100 when complete. | +| `submitted_at`, `started_at`, `stage_started_at`, `finished_at` | UTC timestamps. `started_at - submitted_at` is the queue time. | +| `segments` | `{completed, total}` region segments of a segmented US national run. | +| `stage_durations` | Seconds spent in each stage that has ended. | + +A monolithic run moves through `loading_data`, `baseline`, `reform` and +`outputs`. A segmented US national run computes baseline and reform together +in 20 region segments, so it moves through `loading_data`, `simulating` and +`outputs`. + +The `run_simulation` worker writes a stage record to the Modal Dict +`simulation-api-job-progress`, keyed by `job_id`. The gateway reports +`queued` for a job without a record only when the worker app's bundle +manifest declares `job_progress`; apps deployed before stage reporting never +write a record, so their jobs get no stage fields. A failed read or a +malformed record drops the stage fields; it never changes the status, the +result or the HTTP status code. + +The `progress` value at each stage start (`STAGE_PROGRESS_FLOOR` in +`policyengine_simulation_contract.job_progress`) is provisional and chosen by +hand. Replace it with measured stage shares once per-stage runtime +statistics exist (#780, phase 2). + ## Image dependencies The Modal image installs with `uv_sync(frozen=True)` from this project's diff --git a/projects/policyengine-simulation-gateway/fixtures/gateway_endpoints.py b/projects/policyengine-simulation-gateway/fixtures/gateway_endpoints.py index 7caf5e078..bdbf89f1e 100644 --- a/projects/policyengine-simulation-gateway/fixtures/gateway_endpoints.py +++ b/projects/policyengine-simulation-gateway/fixtures/gateway_endpoints.py @@ -320,7 +320,11 @@ class OutputExpiredError(Exception): @pytest.fixture def mock_modal(monkeypatch): """Patch Modal calls in the gateway endpoints module.""" - from policyengine_simulation_contract import budget_window_state, dataset_uri + from policyengine_simulation_contract import ( + budget_window_state, + dataset_uri, + job_progress_state, + ) from policyengine_simulation_gateway import endpoints @@ -355,6 +359,7 @@ class MockModal: monkeypatch.setattr(endpoints, "modal", MockModal) monkeypatch.setattr(budget_window_state, "modal", MockModal) + monkeypatch.setattr(job_progress_state, "modal", MockModal) monkeypatch.setattr( dataset_uri, "with_hf_revision", diff --git a/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/endpoints.py b/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/endpoints.py index c8315c881..e9a9a9fd0 100644 --- a/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/endpoints.py +++ b/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/endpoints.py @@ -4,6 +4,7 @@ import logging from dataclasses import dataclass +from datetime import UTC, datetime from typing import TypedDict import modal @@ -35,6 +36,13 @@ from policyengine_simulation_contract.hf_dataset import ( HuggingFaceDatasetReferenceError, ) +from policyengine_simulation_contract.job_progress import ( + JOB_PROGRESS_CAPABILITY_KEY, + JobStatus, + job_progress_fields, + supported_job_progress_capability, +) +from policyengine_simulation_contract.job_progress_state import get_job_progress from policyengine_simulation_contract.spm import ( SPMCapability, SPMInputError, @@ -68,6 +76,10 @@ ROUTING_STATE_DICT_NAME = "simulation-api-routing-state" ROUTING_STATE_ACTIVE_KEY = "active" SUPPORTED_ROUTE_KINDS = ("policyengine", "us", "uk") +# Job metadata keys that the status response never echoes directly. +_INTERNAL_JOB_METADATA_KEYS = frozenset( + {"observability_id", "submitted_at", JOB_PROGRESS_CAPABILITY_KEY} +) class VersionRoutingState(TypedDict, total=False): @@ -585,14 +597,26 @@ def _serialize_job_metadata( resolved_app_name: str, bundle: PolicyEngineBundle, observability_id: str | None = None, + *, + submitted_at: datetime | None = None, + job_progress: int | None = None, ) -> dict: return { "resolved_app_name": resolved_app_name, "policyengine_bundle": _bundle_payload(bundle), "observability_id": observability_id, + "submitted_at": None if submitted_at is None else submitted_at.isoformat(), + JOB_PROGRESS_CAPABILITY_KEY: job_progress, } +def _route_job_progress_capability(route: RouteResolution) -> int | None: + manifest = route.bundle_manifest + if not isinstance(manifest, dict): + return None + return supported_job_progress_capability(manifest.get(JOB_PROGRESS_CAPABILITY_KEY)) + + def _bind_observability_id( http_request: Request, runtime: ObservabilityRuntime, @@ -630,10 +654,45 @@ def _public_job_metadata(job_metadata: dict | None) -> dict: if not isinstance(job_metadata, dict): return {} return { - key: value for key, value in job_metadata.items() if key != "observability_id" + key: value + for key, value in job_metadata.items() + if key not in _INTERNAL_JOB_METADATA_KEYS } +def _job_progress_payload( + runtime: ObservabilityRuntime, + job_id: str, + job_metadata: dict, + status: JobStatus, +) -> dict: + """Optional stage fields for a status response. + + Stage reporting never changes the status, the result or the HTTP code: a + failed or malformed read drops the stage fields instead of guessing. + """ + reports_stages = ( + supported_job_progress_capability(job_metadata.get(JOB_PROGRESS_CAPABILITY_KEY)) + is not None + ) + record = None + if reports_stages: + try: + with runtime.span(ANNUAL_IMPACT_STAGES.name(Stage.MODAL_JOB_PROGRESS_READ)): + record = get_job_progress(job_id) + except Exception as exc: # noqa: BLE001 - stage reporting is optional + logger.warning( + "Job progress unavailable for %s: %s", job_id, type(exc).__name__ + ) + reports_stages = False + return job_progress_fields( + status=status, + submitted_at=job_metadata.get("submitted_at"), + record=record, + reports_stages=reports_stages, + ) + + def _build_budget_window_parent_payload( request: BudgetWindowBatchRequest, *, @@ -776,6 +835,7 @@ async def submit_simulation( # Spawn the job (returns immediately). ``Function.from_name`` is a lazy # handle — the control-plane RPC (hydration) happens inside ``spawn`` — # so both live under the spawn segment to time the real network cost. + submitted_at = datetime.now(UTC) with runtime.span(ANNUAL_IMPACT_STAGES.name(Stage.MODAL_FUNCTION_SPAWN)): sim_func = modal.Function.from_name(route.app_name, "run_simulation") call = sim_func.spawn( @@ -785,7 +845,13 @@ async def submit_simulation( runtime.set_context(job_id=call.object_id) - job_metadata = _serialize_job_metadata(route.app_name, bundle, observability_id) + job_metadata = _serialize_job_metadata( + route.app_name, + bundle, + observability_id, + submitted_at=submitted_at, + job_progress=_route_job_progress_capability(route), + ) with runtime.span(ANNUAL_IMPACT_STAGES.name(Stage.MODAL_JOB_METADATA_WRITE)): _job_metadata_store()[call.object_id] = job_metadata @@ -941,16 +1007,24 @@ async def get_job_status(job_id: str, request: Request, response: Response): runtime.record_exception(exc, handled=False, status_code=500) raise + def with_progress(status: JobStatus) -> dict: + return { + **public_metadata, + **_job_progress_payload(runtime, job_id, job_metadata, status), + } + try: with runtime.span(ANNUAL_IMPACT_STAGES.name(Stage.MODAL_JOB_STATUS_POLL)): result = call.get(timeout=0) + complete_metadata = with_progress("complete") with runtime.span(ANNUAL_IMPACT_STAGES.name(Stage.RESPONSE_SERIALIZATION)): return JobStatusResponse( - status="complete", result=result, **public_metadata + status="complete", result=result, **complete_metadata ) except TimeoutError: + running_metadata = with_progress("running") with runtime.span(ANNUAL_IMPACT_STAGES.name(Stage.RESPONSE_SERIALIZATION)): - return running_job_response(public_metadata, headers=response_headers) + return running_job_response(running_metadata, headers=response_headers) except Exception as exc: if _is_modal_job_not_found(exc): runtime.record_exception(exc, handled=True, status_code=404) @@ -960,7 +1034,7 @@ async def get_job_status(job_id: str, request: Request, response: Response): return failed_job_response( error=detail.message, errors=[detail.model_dump()], - job_metadata=public_metadata, + job_metadata=with_progress("failed"), headers=response_headers, ) redacted = log_and_redact_exception( @@ -969,10 +1043,11 @@ async def get_job_status(job_id: str, request: Request, response: Response): scope="simulation_job_status", context={"job_id": job_id}, ) + failed_metadata = with_progress("failed") with runtime.span(ANNUAL_IMPACT_STAGES.name(Stage.RESPONSE_SERIALIZATION)): return failed_job_response( error=redacted, - job_metadata=public_metadata, + job_metadata=failed_metadata, headers=response_headers, ) diff --git a/projects/policyengine-simulation-gateway/tests/golden/openapi.json b/projects/policyengine-simulation-gateway/tests/golden/openapi.json index 346a301cb..0d8ab4ef9 100644 --- a/projects/policyengine-simulation-gateway/tests/golden/openapi.json +++ b/projects/policyengine-simulation-gateway/tests/golden/openapi.json @@ -1246,6 +1246,40 @@ ], "title": "IntraDecileOutput" }, + "JobSegmentProgress": { + "properties": { + "completed": { + "type": "integer", + "minimum": 0.0, + "title": "Completed" + }, + "total": { + "type": "integer", + "minimum": 1.0, + "title": "Total" + } + }, + "type": "object", + "required": [ + "completed", + "total" + ], + "title": "JobSegmentProgress", + "description": "Completed region segments of a segmented national run." + }, + "JobStage": { + "type": "string", + "enum": [ + "queued", + "loading_data", + "baseline", + "reform", + "simulating", + "outputs" + ], + "title": "JobStage", + "description": "Coarse stages of an annual economy job, in the order they run.\n\nA monolithic run moves through ``loading_data``, ``baseline``, ``reform``\nand ``outputs``. A segmented US national run computes baseline and reform\ntogether in parallel region segments, so it moves through\n``loading_data``, ``simulating`` and ``outputs``." + }, "JobStatusResponse": { "properties": { "status": { @@ -1312,6 +1346,112 @@ "type": "null" } ] + }, + "stage": { + "anyOf": [ + { + "$ref": "#/components/schemas/JobStage" + }, + { + "type": "null" + } + ], + "description": "Current stage of a running job, or the stage in which a failed job stopped. Absent for complete jobs." + }, + "progress": { + "anyOf": [ + { + "type": "integer", + "maximum": 100.0, + "minimum": 0.0 + }, + { + "type": "null" + } + ], + "title": "Progress", + "description": "Coarse completion percent. It advances at stage boundaries and as national segments complete, stays below 100 while running, and is 100 when complete." + }, + "submitted_at": { + "anyOf": [ + { + "type": "string", + "format": "date-time" + }, + { + "type": "null" + } + ], + "title": "Submitted At", + "description": "When the gateway accepted the job." + }, + "started_at": { + "anyOf": [ + { + "type": "string", + "format": "date-time" + }, + { + "type": "null" + } + ], + "title": "Started At", + "description": "When a worker started the job." + }, + "stage_started_at": { + "anyOf": [ + { + "type": "string", + "format": "date-time" + }, + { + "type": "null" + } + ], + "title": "Stage Started At", + "description": "When the current stage started." + }, + "finished_at": { + "anyOf": [ + { + "type": "string", + "format": "date-time" + }, + { + "type": "null" + } + ], + "title": "Finished At", + "description": "When the worker finished or stopped the job." + }, + "segments": { + "anyOf": [ + { + "$ref": "#/components/schemas/JobSegmentProgress" + }, + { + "type": "null" + } + ], + "description": "Region segments of a segmented US national run." + }, + "stage_durations": { + "anyOf": [ + { + "additionalProperties": { + "type": "number" + }, + "propertyNames": { + "$ref": "#/components/schemas/JobStage" + }, + "type": "object" + }, + { + "type": "null" + } + ], + "title": "Stage Durations", + "description": "Seconds spent in each stage that has ended." } }, "type": "object", diff --git a/projects/policyengine-simulation-gateway/tests/test_job_progress.py b/projects/policyengine-simulation-gateway/tests/test_job_progress.py new file mode 100644 index 000000000..e8689f64e --- /dev/null +++ b/projects/policyengine-simulation-gateway/tests/test_job_progress.py @@ -0,0 +1,293 @@ +"""Job status stage fields. + +The worker writes a stage record; the gateway adds optional stage fields to +``GET /jobs/{job_id}``. Jobs routed to workers that do not declare the +capability keep their existing response body. +""" + +from datetime import datetime + +import pytest +from fastapi.testclient import TestClient +from policyengine_simulation_contract.job_progress import ( + JOB_PROGRESS_CAPABILITY_VERSION, + JobProgressRecord, + JobSegmentProgress, + JobStage, +) +from policyengine_simulation_contract.job_progress_state import ( + JOB_PROGRESS_DICT_NAME, +) +from policyengine_simulation_observability.stages import Stage + +JOB_METADATA_DICT_NAME = "simulation-api-job-metadata" +LEGACY_RUNNING_KEYS = { + "status", + "result", + "error", + "resolved_app_name", + "policyengine_bundle", +} + + +@pytest.fixture +def stage_reporting_worker(mock_modal): + """Declare the capability on the routing state's latest bundle.""" + bundles = mock_modal["dicts"]["simulation-api-routing-state"]["active"]["bundles"] + bundles["4.10.0"] = { + **bundles["4.10.0"], + "job_progress": JOB_PROGRESS_CAPABILITY_VERSION, + } + return mock_modal + + +def _submit(client: TestClient) -> str: + response = client.post( + "/simulate/economy/comparison", + json={"country": "us", "scope": "macro", "reform": {}}, + ) + assert response.status_code == 200 + return response.json()["job_id"] + + +def _store_record(mock_modal, job_id: str, record: JobProgressRecord) -> None: + mock_modal["dicts"].setdefault(JOB_PROGRESS_DICT_NAME, {})[job_id] = ( + record.model_dump(mode="json") + ) + + +def _reform_record(**overrides) -> JobProgressRecord: + values = { + "stage": JobStage.REFORM, + "started_at": "2026-10-13T12:00:30Z", + "stage_started_at": "2026-10-13T12:01:00Z", + "stage_durations": {JobStage.LOADING_DATA: 12.5, JobStage.BASELINE: 17.5}, + } + return JobProgressRecord.model_validate({**values, **overrides}) + + +class TestSubmission: + def test__records_submission_time_and_worker_capability( + self, stage_reporting_worker, client: TestClient + ): + job_id = _submit(client) + + metadata = stage_reporting_worker["dicts"][JOB_METADATA_DICT_NAME][job_id] + assert datetime.fromisoformat(metadata["submitted_at"]).tzinfo is not None + assert metadata["job_progress"] == JOB_PROGRESS_CAPABILITY_VERSION + + def test__worker_without_capability_is_recorded_as_none( + self, mock_modal, client: TestClient + ): + job_id = _submit(client) + + metadata = mock_modal["dicts"][JOB_METADATA_DICT_NAME][job_id] + assert metadata["job_progress"] is None + + +class TestStageReportingWorker: + def test__running_job_without_record_is_queued( + self, stage_reporting_worker, client: TestClient + ): + job_id = _submit(client) + stage_reporting_worker["func"].last_call.running = True + + response = client.get(f"/jobs/{job_id}") + + assert response.status_code == 202 + body = response.json() + assert body["status"] == "running" + assert body["stage"] == "queued" + assert body["progress"] == 0 + assert body["submitted_at"].endswith("Z") + assert "job_progress" not in body + + def test__running_job_reports_the_worker_stage( + self, stage_reporting_worker, client: TestClient + ): + job_id = _submit(client) + stage_reporting_worker["func"].last_call.running = True + _store_record(stage_reporting_worker, job_id, _reform_record()) + + response = client.get(f"/jobs/{job_id}") + + assert response.status_code == 202 + body = response.json() + assert body["stage"] == "reform" + assert body["progress"] == 45 + assert body["started_at"] == "2026-10-13T12:00:30Z" + assert body["stage_started_at"] == "2026-10-13T12:01:00Z" + assert body["stage_durations"] == {"loading_data": 12.5, "baseline": 17.5} + assert body["result"] is None + assert body["error"] is None + + def test__segmented_run_reports_completed_segments( + self, stage_reporting_worker, client: TestClient + ): + job_id = _submit(client) + stage_reporting_worker["func"].last_call.running = True + _store_record( + stage_reporting_worker, + job_id, + _reform_record( + stage=JobStage.SIMULATING, + segments=JobSegmentProgress(completed=10, total=20), + ), + ) + + body = client.get(f"/jobs/{job_id}").json() + + assert body["stage"] == "simulating" + assert body["segments"] == {"completed": 10, "total": 20} + assert body["progress"] == 50 + + def test__complete_job_reports_durations_and_no_current_stage( + self, stage_reporting_worker, client: TestClient + ): + job_id = _submit(client) + _store_record( + stage_reporting_worker, + job_id, + _reform_record( + stage=JobStage.OUTPUTS, + finished_at="2026-10-13T12:04:00Z", + stage_durations={ + JobStage.LOADING_DATA: 12.5, + JobStage.BASELINE: 17.5, + JobStage.REFORM: 150.0, + JobStage.OUTPUTS: 30.0, + }, + ), + ) + + response = client.get(f"/jobs/{job_id}") + + assert response.status_code == 200 + body = response.json() + assert body["status"] == "complete" + assert body["result"] is not None + assert body["progress"] == 100 + assert body["finished_at"] == "2026-10-13T12:04:00Z" + assert body["stage_durations"]["reform"] == 150.0 + assert "stage" not in body + assert "stage_started_at" not in body + + def test__failed_job_reports_the_stage_it_stopped_in( + self, stage_reporting_worker, client: TestClient + ): + job_id = _submit(client) + stage_reporting_worker["func"].last_call.error = RuntimeError("worker crashed") + _store_record( + stage_reporting_worker, + job_id, + _reform_record(finished_at="2026-10-13T12:02:00Z"), + ) + + response = client.get(f"/jobs/{job_id}") + + assert response.status_code == 500 + body = response.json() + assert body["status"] == "failed" + assert body["stage"] == "reform" + assert "progress" not in body + assert "worker crashed" not in body["error"] + + def test__progress_read_uses_the_registered_stage( + self, stage_reporting_worker, client: TestClient, monkeypatch + ): + job_id = _submit(client) + stage_reporting_worker["func"].last_call.running = True + runtime = client.app.state.policyengine_observability + original_span = runtime.span + names = [] + + def recording_span(name, *args, **kwargs): + names.append(name) + return original_span(name, *args, **kwargs) + + monkeypatch.setattr(runtime, "span", recording_span) + + client.get(f"/jobs/{job_id}") + + assert Stage.MODAL_JOB_PROGRESS_READ.value in names + + +class TestStageReportingFailuresAreHarmless: + def test__unreadable_store_drops_stage_fields( + self, stage_reporting_worker, client: TestClient, monkeypatch + ): + from policyengine_simulation_gateway import endpoints + + job_id = _submit(client) + stage_reporting_worker["func"].last_call.running = True + + def failing_read(job_id): + raise ConnectionError("modal unavailable") + + monkeypatch.setattr(endpoints, "get_job_progress", failing_read) + + response = client.get(f"/jobs/{job_id}") + + assert response.status_code == 202 + body = response.json() + assert body["status"] == "running" + assert "stage" not in body + assert "progress" not in body + assert "submitted_at" in body + + def test__malformed_record_drops_stage_fields( + self, stage_reporting_worker, client: TestClient + ): + job_id = _submit(client) + stage_reporting_worker["func"].last_call.running = True + stage_reporting_worker["dicts"].setdefault(JOB_PROGRESS_DICT_NAME, {})[ + job_id + ] = {"stage": "teleporting"} + + response = client.get(f"/jobs/{job_id}") + + assert response.status_code == 202 + assert "stage" not in response.json() + + def test__malformed_record_keeps_a_complete_result( + self, stage_reporting_worker, client: TestClient + ): + job_id = _submit(client) + stage_reporting_worker["dicts"].setdefault(JOB_PROGRESS_DICT_NAME, {})[ + job_id + ] = {"schema_version": 99} + + response = client.get(f"/jobs/{job_id}") + + assert response.status_code == 200 + assert response.json()["status"] == "complete" + + +class TestWorkerWithoutCapability: + def test__running_body_adds_only_the_submission_time( + self, mock_modal, client: TestClient + ): + job_id = _submit(client) + mock_modal["func"].last_call.running = True + _store_record(mock_modal, job_id, _reform_record()) + + response = client.get(f"/jobs/{job_id}") + + assert response.status_code == 202 + assert set(response.json()) == LEGACY_RUNNING_KEYS | {"submitted_at"} + + def test__legacy_job_metadata_adds_no_fields(self, mock_modal, client: TestClient): + job_id = "legacy-job" + mock_modal["dicts"][JOB_METADATA_DICT_NAME] = { + job_id: { + "resolved_app_name": "legacy-worker", + "policyengine_bundle": {"model_version": "1.500.0"}, + "observability_id": None, + } + } + mock_modal["func"].call_for(job_id).running = True + + response = client.get(f"/jobs/{job_id}") + + assert response.status_code == 202 + assert set(response.json()) == LEGACY_RUNNING_KEYS From b78a2261bd591dfb63330450a5a01b7b78ed91f9 Mon Sep 17 00:00:00 2001 From: PavelMakarchuk Date: Tue, 13 Oct 2026 12:02:01 -0400 Subject: [PATCH 2/7] Set progress floors per execution path from production telemetry One week of production telemetry (6-13 October 2026) shows that the hand-chosen floors were far from the measured stage shares. Queue time is about 39% of a UK national run, outputs take almost no time, and the segmented calculation is about 96% of a US national run. A single table cannot fit both paths. - The worker records its execution path (monolithic or segmented), and the gateway picks the matching floor table. - Before the path is known, a stage uses the lowest floor that any path gives it, so progress never moves backwards. - run_simulation closes the last stage inside its operation and adds job_stage_seconds_ attributes, so telemetry can split the baseline from the reform calculation. Refs #780 Co-Authored-By: Claude Opus 5.5 --- docs/engineering/skills/observability.md | 10 ++ .../job_progress.py | 71 ++++++++++--- .../tests/test_job_progress.py | 63 +++++++++++- .../src/modal/app.py | 47 +++++++-- .../src/modal/segmented_national.py | 4 +- .../job_progress.py | 29 ++++++ .../tests/test_job_progress.py | 99 ++++++++++++++++++- .../policyengine-simulation-gateway/README.md | 13 ++- .../tests/test_job_progress.py | 7 +- 9 files changed, 305 insertions(+), 38 deletions(-) diff --git a/docs/engineering/skills/observability.md b/docs/engineering/skills/observability.md index 488f16822..36efbd399 100644 --- a/docs/engineering/skills/observability.md +++ b/docs/engineering/skills/observability.md @@ -102,6 +102,16 @@ operations: `simulation_role` attribute. The other planning operations include the `country` attribute. +### Annual job stage durations + +`run_simulation` records client-visible job stages (see +`policyengine_simulation_contract.job_progress`). Before its operation ends, +it sets one `job_stage_seconds_` attribute for each ended stage, for +example `job_stage_seconds_baseline` and `job_stage_seconds_reform`. These +are the only main-path measurements that separate the baseline from the +reform calculation, because both run inside `economic_impact_analysis`. The +attribute set is bounded by the `JobStage` enum. + ## Metric resource identity Every process that emits OpenTelemetry metrics must have a unique diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/job_progress.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/job_progress.py index a076b256e..ceb296fea 100644 --- a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/job_progress.py +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/job_progress.py @@ -50,6 +50,13 @@ class JobStage(StrEnum): OUTPUTS = "outputs" +class JobExecution(StrEnum): + """How the worker runs the job; the two paths spend time differently.""" + + MONOLITHIC = "monolithic" + SEGMENTED = "segmented" + + class JobSegmentProgress(BaseModel): """Completed region segments of a segmented national run.""" @@ -71,23 +78,41 @@ class JobProgressRecord(BaseModel): started_at: AwareDatetime stage_started_at: AwareDatetime finished_at: AwareDatetime | None = None + # None only between the worker start and its choice of execution path. + execution: JobExecution | None = None segments: JobSegmentProgress | None = None stage_durations: dict[JobStage, NonNegativeFloat] = Field(default_factory=dict) -# Percent of a typical run that is complete when each stage starts. -# PROVISIONAL: chosen by hand, not measured. Replace these with measured -# stage shares once per-stage runtime statistics exist (#780, phase 2). -# ``simulating`` advances from its floor to the ``outputs`` floor as -# segments complete. -STAGE_PROGRESS_FLOOR: Mapping[JobStage, int] = MappingProxyType( +# Percent of submission-to-completion time that has passed, in a typical +# successful production run, when each stage starts. Set from one week of +# production telemetry (6-13 October 2026); queue time counts towards the +# ``loading_data`` floor. PROVISIONAL: replace with stage shares measured per +# country and region type once runtime statistics exist (#780, phase 2). +# +# Monolithic: UK national and US state runs. The baseline/reform split comes +# from Stage 12 state runs (baseline is about 29% of their calculation). +# Segmented: US national runs; ``simulating`` advances from its floor to the +# ``outputs`` floor as segments complete, and the national reduce is short. +STAGE_PROGRESS_FLOOR: Mapping[JobExecution, Mapping[JobStage, int]] = MappingProxyType( { - JobStage.QUEUED: 0, - JobStage.LOADING_DATA: 5, - JobStage.BASELINE: 15, - JobStage.SIMULATING: 15, - JobStage.REFORM: 45, - JobStage.OUTPUTS: 85, + JobExecution.MONOLITHIC: MappingProxyType( + { + JobStage.QUEUED: 0, + JobStage.LOADING_DATA: 25, + JobStage.BASELINE: 55, + JobStage.REFORM: 67, + JobStage.OUTPUTS: 97, + } + ), + JobExecution.SEGMENTED: MappingProxyType( + { + JobStage.QUEUED: 0, + JobStage.LOADING_DATA: 3, + JobStage.SIMULATING: 5, + JobStage.OUTPUTS: 97, + } + ), } ) # A running job never reports 100: only a complete result does. @@ -121,10 +146,26 @@ def _parse_datetime(value: object) -> datetime | None: return None +def stage_progress_floor(execution: JobExecution | None, stage: JobStage) -> int: + """The progress percent at which ``stage`` starts on ``execution``. + + Before the worker chooses its path, use the lowest floor that any path + gives the stage, so progress never moves backwards once the path is known. + """ + + floors = STAGE_PROGRESS_FLOOR.get(execution) if execution is not None else None + if floors is not None and stage in floors: + return floors[stage] + return min( + (floors[stage] for floors in STAGE_PROGRESS_FLOOR.values() if stage in floors), + default=0, + ) + + def _running_progress(record: JobProgressRecord) -> int: - floor = STAGE_PROGRESS_FLOOR[record.stage] + floor = stage_progress_floor(record.execution, record.stage) if record.stage is JobStage.SIMULATING and record.segments is not None: - ceiling = STAGE_PROGRESS_FLOOR[JobStage.OUTPUTS] + ceiling = stage_progress_floor(record.execution, JobStage.OUTPUTS) share = record.segments.completed / record.segments.total floor += int((ceiling - floor) * share) return min(floor, RUNNING_PROGRESS_CAP) @@ -156,7 +197,7 @@ def job_progress_fields( if record is None: if status == "running": fields["stage"] = JobStage.QUEUED.value - fields["progress"] = STAGE_PROGRESS_FLOOR[JobStage.QUEUED] + fields["progress"] = stage_progress_floor(None, JobStage.QUEUED) else: fields["started_at"] = _json_datetime(record.started_at) fields["finished_at"] = _json_datetime(record.finished_at) diff --git a/libs/policyengine-simulation-contract/tests/test_job_progress.py b/libs/policyengine-simulation-contract/tests/test_job_progress.py index 14a981831..ce38622dd 100644 --- a/libs/policyengine-simulation-contract/tests/test_job_progress.py +++ b/libs/policyengine-simulation-contract/tests/test_job_progress.py @@ -10,10 +10,12 @@ JOB_PROGRESS_CAPABILITY_VERSION, RUNNING_PROGRESS_CAP, STAGE_PROGRESS_FLOOR, + JobExecution, JobProgressRecord, JobSegmentProgress, JobStage, job_progress_fields, + stage_progress_floor, supported_job_progress_capability, ) from policyengine_simulation_contract.job_progress_state import ( @@ -31,6 +33,7 @@ def _record(**overrides) -> JobProgressRecord: values = { "stage": JobStage.REFORM, + "execution": JobExecution.MONOLITHIC, "started_at": STARTED, "stage_started_at": STAGE_STARTED, "stage_durations": {JobStage.LOADING_DATA: 12.5, JobStage.BASELINE: 17.5}, @@ -100,17 +103,18 @@ def test__running_job_reports_stage_floor_and_timing(self): "started_at": "2026-10-13T12:00:30Z", "stage": "reform", "stage_started_at": "2026-10-13T12:01:00Z", - "progress": STAGE_PROGRESS_FLOOR[JobStage.REFORM], + "progress": 67, "stage_durations": {"loading_data": 12.5, "baseline": 17.5}, } @pytest.mark.parametrize( ("completed", "expected"), - [(0, 15), (10, 50), (19, 81), (20, 85)], + [(0, 5), (10, 51), (19, 92), (20, 97)], ) def test__segmented_run_advances_with_completed_segments(self, completed, expected): record = _record( stage=JobStage.SIMULATING, + execution=JobExecution.SEGMENTED, segments=JobSegmentProgress(completed=completed, total=20), ) @@ -125,9 +129,26 @@ def test__segmented_run_advances_with_completed_segments(self, completed, expect assert fields["progress"] == expected assert fields["segments"] == {"completed": completed, "total": 20} + def test__running_job_before_its_execution_path_uses_the_lowest_floor(self): + record = _record(stage=JobStage.LOADING_DATA, execution=None) + + fields = job_progress_fields( + status="running", + submitted_at=SUBMITTED, + record=record, + reports_stages=True, + ) + + assert fields["progress"] == 3 + assert fields["progress"] <= stage_progress_floor( + JobExecution.MONOLITHIC, JobStage.LOADING_DATA + ) + def test__running_progress_never_reaches_100(self): assert all( - floor <= RUNNING_PROGRESS_CAP for floor in STAGE_PROGRESS_FLOOR.values() + floor <= RUNNING_PROGRESS_CAP + for floors in STAGE_PROGRESS_FLOOR.values() + for floor in floors.values() ) record = _record(stage=JobStage.OUTPUTS) @@ -211,6 +232,42 @@ def test__fields_are_valid_public_status_fields(self, status): } +class TestStageProgressFloor: + @pytest.mark.parametrize( + ("execution", "stages"), + [ + ( + JobExecution.MONOLITHIC, + [ + JobStage.QUEUED, + JobStage.LOADING_DATA, + JobStage.BASELINE, + JobStage.REFORM, + JobStage.OUTPUTS, + ], + ), + ( + JobExecution.SEGMENTED, + [ + JobStage.QUEUED, + JobStage.LOADING_DATA, + JobStage.SIMULATING, + JobStage.OUTPUTS, + ], + ), + ], + ) + def test__each_path_covers_its_stages_in_increasing_order(self, execution, stages): + floors = STAGE_PROGRESS_FLOOR[execution] + + assert list(floors) == stages + assert [floors[stage] for stage in stages] == sorted(floors.values()) + + def test__a_stage_outside_the_path_uses_the_lowest_floor_of_any_path(self): + assert stage_progress_floor(JobExecution.MONOLITHIC, JobStage.SIMULATING) == 5 + assert stage_progress_floor(None, JobStage.QUEUED) == 0 + + class TestSupportedJobProgressCapability: def test__current_version_is_supported(self): assert ( diff --git a/projects/policyengine-simulation-executor/src/modal/app.py b/projects/policyengine-simulation-executor/src/modal/app.py index cda8a108c..9594554f2 100644 --- a/projects/policyengine-simulation-executor/src/modal/app.py +++ b/projects/policyengine-simulation-executor/src/modal/app.py @@ -282,6 +282,26 @@ def _job_progress_reporter(): return StoredJobProgressReporter(job_id, write=put_job_progress) +def _record_stage_durations(runtime, progress) -> None: + """Attach the job's stage durations to its operation telemetry. + + Main-path traces cannot split the baseline from the reform calculation: + both run inside ``economic_impact_analysis``. These attributes can. + """ + durations = progress.stage_durations() + if not durations: + return + try: + runtime.set_context( + **{ + f"job_stage_seconds_{stage.value}": seconds + for stage, seconds in durations.items() + } + ) + except Exception: # noqa: BLE001 - telemetry is non-fatal + pass + + @app.function( image=simulation_image, cpu=8.0, @@ -328,17 +348,24 @@ def run_simulation( _set_modal_call_attributes(runtime) from src.modal.segmented_national import dispatch_run_simulation - # Plain national macro requests fan out across region groups - # by default (segmented: false opts out). APP_NAME is this - # deployed app; children spawn into its dedicated - # run_simulation_segment pool. - return dispatch_run_simulation( - params, - app_name=APP_NAME, - runtime=runtime, - progress=progress, - ) + try: + # Plain national macro requests fan out across region groups + # by default (segmented: false opts out). APP_NAME is this + # deployed app; children spawn into its dedicated + # run_simulation_segment pool. + return dispatch_run_simulation( + params, + app_name=APP_NAME, + runtime=runtime, + progress=progress, + ) + finally: + # Close the last stage inside the operation, so its + # completion record carries every stage duration. + progress.finish() + _record_stage_durations(runtime, progress) finally: + # A no-op unless the operation never started. progress.finish() runtime.shutdown() diff --git a/projects/policyengine-simulation-executor/src/modal/segmented_national.py b/projects/policyengine-simulation-executor/src/modal/segmented_national.py index 863bdecd0..f8f33aa38 100644 --- a/projects/policyengine-simulation-executor/src/modal/segmented_national.py +++ b/projects/policyengine-simulation-executor/src/modal/segmented_national.py @@ -25,7 +25,7 @@ from typing import Any from policyengine_observability import ObservabilityRuntime -from policyengine_simulation_contract.job_progress import JobStage +from policyengine_simulation_contract.job_progress import JobExecution, JobStage from policyengine_simulation_contract.spm import ( SPMInputError, combine_spm_results, @@ -378,6 +378,7 @@ def dispatch_run_simulation( if selection is not None: params = {**params, "spm": selection} if should_run_segmented_national(params): + progress.set_execution(JobExecution.SEGMENTED) return run_segmented_national_impl( params, app_name=app_name, @@ -393,4 +394,5 @@ def dispatch_run_simulation( # and UK requests must not pollute the segmented-vs-monolithic metric. if is_plain_national_macro(params): runtime.set_context(national_execution="monolithic") + progress.set_execution(JobExecution.MONOLITHIC) return run_simulation_impl(params, runtime=runtime, progress=progress) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/job_progress.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/job_progress.py index 7f00654c7..efed6757d 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/job_progress.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/job_progress.py @@ -17,6 +17,7 @@ from typing import Protocol from policyengine_simulation_contract.job_progress import ( + JobExecution, JobProgressRecord, JobSegmentProgress, JobStage, @@ -34,6 +35,8 @@ def _utc_now() -> datetime: class JobProgressReporter(Protocol): def start(self) -> None: ... + def set_execution(self, execution: JobExecution) -> None: ... + def enter(self, stage: JobStage) -> None: ... def start_segments(self, total: int) -> None: ... @@ -42,6 +45,8 @@ def segment_completed(self) -> None: ... def finish(self) -> None: ... + def stage_durations(self) -> dict[JobStage, float]: ... + class NullJobProgressReporter: """Reports nothing.""" @@ -49,6 +54,9 @@ class NullJobProgressReporter: def start(self) -> None: pass + def set_execution(self, execution: JobExecution) -> None: + pass + def enter(self, stage: JobStage) -> None: pass @@ -61,6 +69,9 @@ def segment_completed(self) -> None: def finish(self) -> None: pass + def stage_durations(self) -> dict[JobStage, float]: + return {} + NULL_JOB_PROGRESS = NullJobProgressReporter() @@ -96,6 +107,17 @@ def start(self) -> None: ) self._publish() + def set_execution(self, execution: JobExecution) -> None: + if self._record is None: + self.start() + record = self._record + if record is None or record.finished_at is not None: + return + if record.execution is execution: + return + record.execution = execution + self._publish() + def enter(self, stage: JobStage) -> None: if self._record is None: self.start() @@ -139,6 +161,13 @@ def finish(self) -> None: record.finished_at = self._close_stage(record) self._publish() + def stage_durations(self) -> dict[JobStage, float]: + """Seconds per ended stage; complete only after ``finish()``.""" + + if self._record is None: + return {} + return dict(self._record.stage_durations) + def _close_stage(self, record: JobProgressRecord) -> datetime: now = self._clock() elapsed = max((now - record.stage_started_at).total_seconds(), 0.0) diff --git a/projects/policyengine-simulation-executor/tests/test_job_progress.py b/projects/policyengine-simulation-executor/tests/test_job_progress.py index cb83a38a0..5b3eaa8b5 100644 --- a/projects/policyengine-simulation-executor/tests/test_job_progress.py +++ b/projects/policyengine-simulation-executor/tests/test_job_progress.py @@ -14,6 +14,7 @@ from fixtures.spm_doubles import installed_spm_selection, spm_capable_simulation from fixtures.test_simulation_api_contracts import CURRENT_SINGLE_YEAR_MACRO_RESULT from policyengine_simulation_contract.job_progress import ( + JobExecution, JobProgressRecord, JobSegmentProgress, JobStage, @@ -71,6 +72,9 @@ def __init__(self): def start(self) -> None: self.calls.append(("start",)) + def set_execution(self, execution: JobExecution) -> None: + self.calls.append(("set_execution", execution)) + def enter(self, stage: JobStage) -> None: self.calls.append(("enter", stage)) @@ -83,6 +87,9 @@ def segment_completed(self) -> None: def finish(self) -> None: self.calls.append(("finish",)) + def stage_durations(self) -> dict[JobStage, float]: + return {} + def _reporter(): clock = _Clock() @@ -166,6 +173,31 @@ def test__segments_count_up_to_their_total(self): assert writer.last.segments == JobSegmentProgress(completed=2, total=2) assert len(writer.writes) == 5 + def test__execution_path_is_recorded_once(self): + reporter, writer, _ = _reporter() + + reporter.start() + reporter.set_execution(JobExecution.MONOLITHIC) + reporter.set_execution(JobExecution.MONOLITHIC) + + assert [record.execution for _, record in writer.writes] == [ + None, + JobExecution.MONOLITHIC, + ] + assert writer.last.stage is JobStage.LOADING_DATA + + def test__stage_durations_are_a_copy_of_ended_stages(self): + reporter, _, clock = _reporter() + + assert reporter.stage_durations() == {} + reporter.start() + clock.advance(4) + reporter.enter(JobStage.REFORM) + durations = reporter.stage_durations() + durations[JobStage.REFORM] = 1.0 + + assert reporter.stage_durations() == {JobStage.LOADING_DATA: 4.0} + def test__segment_completion_without_segments_writes_nothing(self): reporter, writer, _ = _reporter() @@ -182,6 +214,7 @@ def test__nothing_changes_after_finish(self): reporter.finish() reporter.finish() reporter.enter(JobStage.REFORM) + reporter.set_execution(JobExecution.SEGMENTED) reporter.segment_completed() assert len(writer.writes) == 2 @@ -457,6 +490,27 @@ def fake_impl(params, *, runtime, progress): ) assert seen["progress"] is progress + assert progress.calls == [("set_execution", JobExecution.MONOLITHIC)] + + def test__dispatch_marks_national_runs_as_segmented(self, monkeypatch): + progress = _RecordingReporter() + seen = {} + + def fake_segmented(params, *, app_name, runtime, progress): + seen["progress"] = progress + return {"ok": True} + + monkeypatch.setattr(sn, "run_segmented_national_impl", fake_segmented) + + sn.dispatch_run_simulation( + {"country": "us", "scope": "macro", "time_period": "2026"}, + app_name="app-x", + runtime=NoOpObservabilityRuntime(), + progress=progress, + ) + + assert seen["progress"] is progress + assert progress.calls == [("set_execution", JobExecution.SEGMENTED)] @pytest.fixture @@ -475,14 +529,41 @@ def modal_app(monkeypatch): sys.modules.pop("src.modal.app", None) -class _ShutdownRuntime(NoOpObservabilityRuntime): +class _OperationScope: def __init__(self, events): self.events = events + def __enter__(self): + self.events.append("operation_start") + + def __exit__(self, *exc_info): + self.events.append("operation_end") + return False + + +class _RecordingRuntime(NoOpObservabilityRuntime): + def __init__(self, events): + self.events = events + + def operation(self, *args, **kwargs): + return _OperationScope(self.events) + + def set_context(self, **attributes): + self.events.append(("context", attributes)) + def shutdown(self): self.events.append("shutdown") +def _stage_attributes(events) -> list[tuple[int, dict]]: + return [ + (index, event[1]) + for index, event in enumerate(events) + if isinstance(event, tuple) + and any(key.startswith("job_stage_seconds_") for key in event[1]) + ] + + class TestModalRunSimulation: def _wire(self, app, monkeypatch, dispatch): from policyengine_simulation_contract import job_progress_state @@ -500,7 +581,7 @@ def _wire(self, app, monkeypatch, dispatch): monkeypatch.setattr( app, "_configure_modal_observability", - lambda *, service_role: _ShutdownRuntime(events), + lambda *, service_role: _RecordingRuntime(events), ) monkeypatch.setattr(sn, "dispatch_run_simulation", dispatch) return events, writes @@ -521,20 +602,30 @@ def dispatch(params, *, app_name, runtime, progress): JobStage.REFORM, ] assert writes[-1][1].finished_at is not None - assert events == ["shutdown"] + [(index, attributes)] = _stage_attributes(events) + assert set(attributes) == { + "job_stage_seconds_loading_data", + "job_stage_seconds_reform", + } + # Durations reach the operation before it ends, then the runtime stops. + assert index < events.index("operation_end") + assert events[-1] == "shutdown" def test__failed_job_still_records_where_it_stopped(self, modal_app, monkeypatch): def dispatch(params, *, app_name, runtime, progress): progress.enter(JobStage.BASELINE) raise RuntimeError("model failed") - _, writes = self._wire(modal_app, monkeypatch, dispatch) + events, writes = self._wire(modal_app, monkeypatch, dispatch) with pytest.raises(RuntimeError, match="model failed"): modal_app.run_simulation({"country": "us"}) assert writes[-1][1].stage is JobStage.BASELINE assert writes[-1][1].finished_at is not None + [(index, attributes)] = _stage_attributes(events) + assert "job_stage_seconds_baseline" in attributes + assert index < events.index("operation_end") def test__store_failure_does_not_change_the_result(self, modal_app, monkeypatch): from policyengine_simulation_contract import job_progress_state diff --git a/projects/policyengine-simulation-gateway/README.md b/projects/policyengine-simulation-gateway/README.md index 3b9516b78..0a726d39c 100644 --- a/projects/policyengine-simulation-gateway/README.md +++ b/projects/policyengine-simulation-gateway/README.md @@ -65,9 +65,16 @@ malformed record drops the stage fields; it never changes the status, the result or the HTTP status code. The `progress` value at each stage start (`STAGE_PROGRESS_FLOOR` in -`policyengine_simulation_contract.job_progress`) is provisional and chosen by -hand. Replace it with measured stage shares once per-stage runtime -statistics exist (#780, phase 2). +`policyengine_simulation_contract.job_progress`) depends on the execution +path, which the worker records as `monolithic` or `segmented`. The values +come from one week of production telemetry (6–13 October 2026): queue time +is about 39% of a UK national run, while the segmented calculation is about +96% of a US national run. They stay provisional until per-stage runtime +statistics exist for each country and region type (#780, phase 2). + +The worker also adds `job_stage_seconds_` attributes to its +`run_simulation` operation, so telemetry can split the baseline from the +reform calculation. ## Image dependencies diff --git a/projects/policyengine-simulation-gateway/tests/test_job_progress.py b/projects/policyengine-simulation-gateway/tests/test_job_progress.py index e8689f64e..01faaa942 100644 --- a/projects/policyengine-simulation-gateway/tests/test_job_progress.py +++ b/projects/policyengine-simulation-gateway/tests/test_job_progress.py @@ -11,6 +11,7 @@ from fastapi.testclient import TestClient from policyengine_simulation_contract.job_progress import ( JOB_PROGRESS_CAPABILITY_VERSION, + JobExecution, JobProgressRecord, JobSegmentProgress, JobStage, @@ -59,6 +60,7 @@ def _store_record(mock_modal, job_id: str, record: JobProgressRecord) -> None: def _reform_record(**overrides) -> JobProgressRecord: values = { "stage": JobStage.REFORM, + "execution": JobExecution.MONOLITHIC, "started_at": "2026-10-13T12:00:30Z", "stage_started_at": "2026-10-13T12:01:00Z", "stage_durations": {JobStage.LOADING_DATA: 12.5, JobStage.BASELINE: 17.5}, @@ -114,7 +116,7 @@ def test__running_job_reports_the_worker_stage( assert response.status_code == 202 body = response.json() assert body["stage"] == "reform" - assert body["progress"] == 45 + assert body["progress"] == 67 assert body["started_at"] == "2026-10-13T12:00:30Z" assert body["stage_started_at"] == "2026-10-13T12:01:00Z" assert body["stage_durations"] == {"loading_data": 12.5, "baseline": 17.5} @@ -131,6 +133,7 @@ def test__segmented_run_reports_completed_segments( job_id, _reform_record( stage=JobStage.SIMULATING, + execution=JobExecution.SEGMENTED, segments=JobSegmentProgress(completed=10, total=20), ), ) @@ -139,7 +142,7 @@ def test__segmented_run_reports_completed_segments( assert body["stage"] == "simulating" assert body["segments"] == {"completed": 10, "total": 20} - assert body["progress"] == 50 + assert body["progress"] == 51 def test__complete_job_reports_durations_and_no_current_stage( self, stage_reporting_worker, client: TestClient From e89285e47f95378480256d5374edfe2b538a0036 Mon Sep 17 00:00:00 2001 From: PavelMakarchuk Date: Tue, 13 Oct 2026 13:44:28 -0400 Subject: [PATCH 3/7] Report progress_max and stop moving progress by segment count The regional children of a segmented national run finish at about the same time, and California finishes last. A completed-segment count therefore jumps late and then waits, and a backtest on the 6-13 October runtime figures put it at 16-27 points of mean error. A time-based estimate held inside the stage range stays at about 7 points. - progress is the floor of the current stage on every path; segments stay in the response for captions. - progress_max is the highest percent a running job reaches before its next stage. While the path is unknown it is the highest limit of any path, so a client estimate is not held back. - The schema and README tell clients to accept stage values they do not know, because later versions may add stages. Refs #780 Co-Authored-By: Claude Opus 5.5 --- .../gateway_models.py | 17 +++++- .../job_progress.py | 48 ++++++++++++--- .../tests/test_job_progress.py | 61 +++++++++++++++++-- .../policyengine-simulation-gateway/README.md | 15 +++-- .../tests/golden/openapi.json | 20 +++++- .../tests/test_job_progress.py | 8 ++- 6 files changed, 141 insertions(+), 28 deletions(-) diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/gateway_models.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/gateway_models.py index 342dfbd62..57b26aad2 100644 --- a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/gateway_models.py +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/gateway_models.py @@ -205,7 +205,8 @@ class JobStatusResponse(BaseModel): default=None, description=( "Current stage of a running job, or the stage in which a failed " - "job stopped. Absent for complete jobs." + "job stopped. Absent for complete jobs. Later versions may add " + "stages; clients must accept a value they do not know." ), ) progress: int | None = Field( @@ -213,11 +214,21 @@ class JobStatusResponse(BaseModel): ge=0, le=100, description=( - "Coarse completion percent. It advances at stage boundaries and " - "as national segments complete, stays below 100 while running, " + "Coarse completion percent at the start of the current stage. It " + "changes only at stage boundaries, stays below 100 while running, " "and is 100 when complete." ), ) + progress_max: int | None = Field( + default=None, + ge=0, + le=100, + description=( + "Highest percent a running job reaches before its next stage " + "starts. A client can move a time-based estimate between " + "progress and progress_max." + ), + ) submitted_at: AwareDatetime | None = Field( default=None, description="When the gateway accepted the job." ) diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/job_progress.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/job_progress.py index ceb296fea..53884bb31 100644 --- a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/job_progress.py +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/job_progress.py @@ -40,6 +40,9 @@ class JobStage(StrEnum): and ``outputs``. A segmented US national run computes baseline and reform together in parallel region segments, so it moves through ``loading_data``, ``simulating`` and ``outputs``. + + Later versions may add stages. Clients must accept a value they do not + know, for example by showing a generic "running" state. """ QUEUED = "queued" @@ -92,8 +95,10 @@ class JobProgressRecord(BaseModel): # # Monolithic: UK national and US state runs. The baseline/reform split comes # from Stage 12 state runs (baseline is about 29% of their calculation). -# Segmented: US national runs; ``simulating`` advances from its floor to the -# ``outputs`` floor as segments complete, and the national reduce is short. +# Segmented: US national runs; the national reduce is short. Completed +# segments do not move progress: the regional children finish at about the +# same time and California finishes last, so a count jumps late and then +# waits. ``segments`` reports the count for captions instead. STAGE_PROGRESS_FLOOR: Mapping[JobExecution, Mapping[JobStage, int]] = MappingProxyType( { JobExecution.MONOLITHIC: MappingProxyType( @@ -162,13 +167,29 @@ def stage_progress_floor(execution: JobExecution | None, stage: JobStage) -> int ) -def _running_progress(record: JobProgressRecord) -> int: - floor = stage_progress_floor(record.execution, record.stage) - if record.stage is JobStage.SIMULATING and record.segments is not None: - ceiling = stage_progress_floor(record.execution, JobStage.OUTPUTS) - share = record.segments.completed / record.segments.total - floor += int((ceiling - floor) * share) - return min(floor, RUNNING_PROGRESS_CAP) +def stage_progress_max(execution: JobExecution | None, stage: JobStage) -> int: + """The highest progress percent a job reaches before its next stage. + + This is one below the next stage's floor, or ``RUNNING_PROGRESS_CAP`` for + the last stage. Before the worker chooses its path, use the highest value + that any path gives the stage, so a client estimate is not held back. + """ + + paths = [STAGE_PROGRESS_FLOOR[execution]] if execution is not None else [] + if not any(stage in floors for floors in paths): + paths = list(STAGE_PROGRESS_FLOOR.values()) + limits = [] + for floors in paths: + stages = list(floors) + if stage not in floors: + continue + index = stages.index(stage) + if index + 1 < len(stages): + limits.append(floors[stages[index + 1]] - 1) + else: + limits.append(RUNNING_PROGRESS_CAP) + floor = stage_progress_floor(execution, stage) + return min(max([floor, *limits]), RUNNING_PROGRESS_CAP) def job_progress_fields( @@ -198,6 +219,7 @@ def job_progress_fields( if status == "running": fields["stage"] = JobStage.QUEUED.value fields["progress"] = stage_progress_floor(None, JobStage.QUEUED) + fields["progress_max"] = stage_progress_max(None, JobStage.QUEUED) else: fields["started_at"] = _json_datetime(record.started_at) fields["finished_at"] = _json_datetime(record.finished_at) @@ -205,7 +227,13 @@ def job_progress_fields( fields["stage"] = record.stage.value fields["stage_started_at"] = _json_datetime(record.stage_started_at) if status == "running": - fields["progress"] = _running_progress(record) + fields["progress"] = min( + stage_progress_floor(record.execution, record.stage), + RUNNING_PROGRESS_CAP, + ) + fields["progress_max"] = stage_progress_max( + record.execution, record.stage + ) if record.segments is not None: fields["segments"] = record.segments.model_dump(mode="json") if record.stage_durations: diff --git a/libs/policyengine-simulation-contract/tests/test_job_progress.py b/libs/policyengine-simulation-contract/tests/test_job_progress.py index ce38622dd..1da0a40de 100644 --- a/libs/policyengine-simulation-contract/tests/test_job_progress.py +++ b/libs/policyengine-simulation-contract/tests/test_job_progress.py @@ -16,6 +16,7 @@ JobStage, job_progress_fields, stage_progress_floor, + stage_progress_max, supported_job_progress_capability, ) from policyengine_simulation_contract.job_progress_state import ( @@ -88,6 +89,8 @@ def test__running_job_without_record_is_queued(self): "submitted_at": "2026-10-13T12:00:00Z", "stage": "queued", "progress": 0, + # The path is unknown while queued, so the widest limit applies. + "progress_max": 24, } def test__running_job_reports_stage_floor_and_timing(self): @@ -104,14 +107,12 @@ def test__running_job_reports_stage_floor_and_timing(self): "stage": "reform", "stage_started_at": "2026-10-13T12:01:00Z", "progress": 67, + "progress_max": 96, "stage_durations": {"loading_data": 12.5, "baseline": 17.5}, } - @pytest.mark.parametrize( - ("completed", "expected"), - [(0, 5), (10, 51), (19, 92), (20, 97)], - ) - def test__segmented_run_advances_with_completed_segments(self, completed, expected): + @pytest.mark.parametrize("completed", [0, 10, 19, 20]) + def test__completed_segments_are_reported_but_do_not_move_progress(self, completed): record = _record( stage=JobStage.SIMULATING, execution=JobExecution.SEGMENTED, @@ -126,7 +127,8 @@ def test__segmented_run_advances_with_completed_segments(self, completed, expect ) assert fields["stage"] == "simulating" - assert fields["progress"] == expected + assert fields["progress"] == 5 + assert fields["progress_max"] == 96 assert fields["segments"] == {"completed": completed, "total": 20} def test__running_job_before_its_execution_path_uses_the_lowest_floor(self): @@ -143,6 +145,7 @@ def test__running_job_before_its_execution_path_uses_the_lowest_floor(self): assert fields["progress"] <= stage_progress_floor( JobExecution.MONOLITHIC, JobStage.LOADING_DATA ) + assert fields["progress_max"] == 54 def test__running_progress_never_reaches_100(self): assert all( @@ -160,6 +163,7 @@ def test__running_progress_never_reaches_100(self): ) assert fields["progress"] < 100 + assert fields["progress_max"] == RUNNING_PROGRESS_CAP def test__complete_job_reports_100_and_no_current_stage(self): record = _record(stage=JobStage.OUTPUTS, finished_at=FINISHED) @@ -172,6 +176,7 @@ def test__complete_job_reports_100_and_no_current_stage(self): ) assert fields["progress"] == 100 + assert "progress_max" not in fields assert fields["finished_at"] == "2026-10-13T12:04:00Z" assert "stage" not in fields assert "stage_started_at" not in fields @@ -199,6 +204,7 @@ def test__failed_job_reports_the_stage_it_stopped_in(self): assert fields["stage"] == "reform" assert fields["finished_at"] == "2026-10-13T12:04:00Z" assert "progress" not in fields + assert "progress_max" not in fields def test__failed_job_without_record_reports_no_stage(self): fields = job_progress_fields( @@ -268,6 +274,49 @@ def test__a_stage_outside_the_path_uses_the_lowest_floor_of_any_path(self): assert stage_progress_floor(None, JobStage.QUEUED) == 0 +class TestStageProgressMax: + @pytest.mark.parametrize( + ("execution", "expected"), + [ + ( + JobExecution.MONOLITHIC, + { + JobStage.QUEUED: 24, + JobStage.LOADING_DATA: 54, + JobStage.BASELINE: 66, + JobStage.REFORM: 96, + JobStage.OUTPUTS: 99, + }, + ), + ( + JobExecution.SEGMENTED, + { + JobStage.QUEUED: 2, + JobStage.LOADING_DATA: 4, + JobStage.SIMULATING: 96, + JobStage.OUTPUTS: 99, + }, + ), + ], + ) + def test__limit_is_one_below_the_next_stage(self, execution, expected): + assert { + stage: stage_progress_max(execution, stage) + for stage in STAGE_PROGRESS_FLOOR[execution] + } == expected + + @pytest.mark.parametrize("execution", [None, *JobExecution]) + def test__limit_is_never_below_the_floor_or_above_the_cap(self, execution): + for stage in JobStage: + floor = stage_progress_floor(execution, stage) + limit = stage_progress_max(execution, stage) + assert floor <= limit <= RUNNING_PROGRESS_CAP + + def test__unknown_path_uses_the_highest_limit_of_any_path(self): + assert stage_progress_max(None, JobStage.QUEUED) == 24 + assert stage_progress_max(None, JobStage.LOADING_DATA) == 54 + + class TestSupportedJobProgressCapability: def test__current_version_is_supported(self): assert ( diff --git a/projects/policyengine-simulation-gateway/README.md b/projects/policyengine-simulation-gateway/README.md index 0a726d39c..22730e7aa 100644 --- a/projects/policyengine-simulation-gateway/README.md +++ b/projects/policyengine-simulation-gateway/README.md @@ -45,10 +45,11 @@ that read only `status` see no change. | Field | Meaning | |---|---| -| `stage` | `queued`, `loading_data`, `baseline`, `reform`, `simulating` or `outputs`. Running jobs report their current stage; failed jobs report the stage they stopped in; complete jobs omit it. | -| `progress` | Coarse percent, 0–100. It advances at stage boundaries and as national segments complete, stays below 100 while running, and is 100 when complete. | +| `stage` | `queued`, `loading_data`, `baseline`, `reform`, `simulating` or `outputs`. Running jobs report their current stage; failed jobs report the stage they stopped in; complete jobs omit it. Later versions may add stages, so clients must accept a value they do not know. | +| `progress` | Coarse percent, 0–100, at the start of the current stage. It changes only at stage boundaries, stays below 100 while running, and is 100 when complete. | +| `progress_max` | The highest percent a running job reaches before its next stage starts. | | `submitted_at`, `started_at`, `stage_started_at`, `finished_at` | UTC timestamps. `started_at - submitted_at` is the queue time. | -| `segments` | `{completed, total}` region segments of a segmented US national run. | +| `segments` | `{completed, total}` region segments of a segmented US national run, for captions. | | `stage_durations` | Seconds spent in each stage that has ended. | A monolithic run moves through `loading_data`, `baseline`, `reform` and @@ -66,7 +67,13 @@ result or the HTTP status code. The `progress` value at each stage start (`STAGE_PROGRESS_FLOOR` in `policyengine_simulation_contract.job_progress`) depends on the execution -path, which the worker records as `monolithic` or `segmented`. The values +path, which the worker records as `monolithic` or `segmented`. +Completed segments do not move `progress`: the regional children finish at +about the same time and California finishes last, so a segment count jumps +late and then waits. For a smooth bar, a client moves its own time-based +estimate between `progress` and `progress_max` and never moves the bar +backwards. While a job is `queued`, its path is unknown, so `progress_max` +is the highest limit of any path. The values come from one week of production telemetry (6–13 October 2026): queue time is about 39% of a UK national run, while the segmented calculation is about 96% of a US national run. They stay provisional until per-stage runtime diff --git a/projects/policyengine-simulation-gateway/tests/golden/openapi.json b/projects/policyengine-simulation-gateway/tests/golden/openapi.json index 0d8ab4ef9..b448125ed 100644 --- a/projects/policyengine-simulation-gateway/tests/golden/openapi.json +++ b/projects/policyengine-simulation-gateway/tests/golden/openapi.json @@ -1278,7 +1278,7 @@ "outputs" ], "title": "JobStage", - "description": "Coarse stages of an annual economy job, in the order they run.\n\nA monolithic run moves through ``loading_data``, ``baseline``, ``reform``\nand ``outputs``. A segmented US national run computes baseline and reform\ntogether in parallel region segments, so it moves through\n``loading_data``, ``simulating`` and ``outputs``." + "description": "Coarse stages of an annual economy job, in the order they run.\n\nA monolithic run moves through ``loading_data``, ``baseline``, ``reform``\nand ``outputs``. A segmented US national run computes baseline and reform\ntogether in parallel region segments, so it moves through\n``loading_data``, ``simulating`` and ``outputs``.\n\nLater versions may add stages. Clients must accept a value they do not\nknow, for example by showing a generic \"running\" state." }, "JobStatusResponse": { "properties": { @@ -1356,7 +1356,7 @@ "type": "null" } ], - "description": "Current stage of a running job, or the stage in which a failed job stopped. Absent for complete jobs." + "description": "Current stage of a running job, or the stage in which a failed job stopped. Absent for complete jobs. Later versions may add stages; clients must accept a value they do not know." }, "progress": { "anyOf": [ @@ -1370,7 +1370,21 @@ } ], "title": "Progress", - "description": "Coarse completion percent. It advances at stage boundaries and as national segments complete, stays below 100 while running, and is 100 when complete." + "description": "Coarse completion percent at the start of the current stage. It changes only at stage boundaries, stays below 100 while running, and is 100 when complete." + }, + "progress_max": { + "anyOf": [ + { + "type": "integer", + "maximum": 100.0, + "minimum": 0.0 + }, + { + "type": "null" + } + ], + "title": "Progress Max", + "description": "Highest percent a running job reaches before its next stage starts. A client can move a time-based estimate between progress and progress_max." }, "submitted_at": { "anyOf": [ diff --git a/projects/policyengine-simulation-gateway/tests/test_job_progress.py b/projects/policyengine-simulation-gateway/tests/test_job_progress.py index 01faaa942..8cc4586a0 100644 --- a/projects/policyengine-simulation-gateway/tests/test_job_progress.py +++ b/projects/policyengine-simulation-gateway/tests/test_job_progress.py @@ -101,6 +101,7 @@ def test__running_job_without_record_is_queued( assert body["status"] == "running" assert body["stage"] == "queued" assert body["progress"] == 0 + assert body["progress_max"] == 24 assert body["submitted_at"].endswith("Z") assert "job_progress" not in body @@ -117,13 +118,14 @@ def test__running_job_reports_the_worker_stage( body = response.json() assert body["stage"] == "reform" assert body["progress"] == 67 + assert body["progress_max"] == 96 assert body["started_at"] == "2026-10-13T12:00:30Z" assert body["stage_started_at"] == "2026-10-13T12:01:00Z" assert body["stage_durations"] == {"loading_data": 12.5, "baseline": 17.5} assert body["result"] is None assert body["error"] is None - def test__segmented_run_reports_completed_segments( + def test__segmented_run_reports_segments_without_moving_progress( self, stage_reporting_worker, client: TestClient ): job_id = _submit(client) @@ -142,7 +144,8 @@ def test__segmented_run_reports_completed_segments( assert body["stage"] == "simulating" assert body["segments"] == {"completed": 10, "total": 20} - assert body["progress"] == 51 + assert body["progress"] == 5 + assert body["progress_max"] == 96 def test__complete_job_reports_durations_and_no_current_stage( self, stage_reporting_worker, client: TestClient @@ -170,6 +173,7 @@ def test__complete_job_reports_durations_and_no_current_stage( assert body["status"] == "complete" assert body["result"] is not None assert body["progress"] == 100 + assert "progress_max" not in body assert body["finished_at"] == "2026-10-13T12:04:00Z" assert body["stage_durations"]["reform"] == 150.0 assert "stage" not in body From 31c8463a3e78ea3c8f2421779947ca381d7b5aa8 Mon Sep 17 00:00:00 2001 From: PavelMakarchuk Date: Wed, 14 Oct 2026 17:46:22 -0400 Subject: [PATCH 4/7] Continue a preempted job's stage record instead of restarting it Modal can run the same call again after a preemption, under the same call id. run_simulation built a new reporter each time, so start() wrote a fresh loading_data record and progress dropped from reform (67) or outputs (97) back to 3 and 25. started_at and the stage durations also covered only the last attempt. start() now reads the stored record and continues it: started_at and the ended-stage durations carry over, attempts counts up, and the published stage stays at the earlier attempt's stage until this attempt reaches a stage that starts at least as far along. A finished earlier record is continued too, because a preempted attempt can run its finally block. run_simulation adds job_attempts to its telemetry. The docs now say that budget-window children, which are ordinary run_simulation calls, write their own records that nothing reads yet. Refs #780 Co-Authored-By: Claude Opus 5.5 --- docs/engineering/skills/observability.md | 4 + .../job_progress.py | 3 + .../tests/test_job_progress.py | 6 + .../src/modal/app.py | 15 +- .../job_progress.py | 88 ++++++++- .../tests/test_job_progress.py | 171 +++++++++++++++++- .../policyengine-simulation-gateway/README.md | 7 +- 7 files changed, 280 insertions(+), 14 deletions(-) diff --git a/docs/engineering/skills/observability.md b/docs/engineering/skills/observability.md index 36efbd399..4f9e6b1c5 100644 --- a/docs/engineering/skills/observability.md +++ b/docs/engineering/skills/observability.md @@ -112,6 +112,10 @@ are the only main-path measurements that separate the baseline from the reform calculation, because both run inside `economic_impact_analysis`. The attribute set is bounded by the `JobStage` enum. +Modal can run the same call again after a preemption, under the same call +id. A later attempt continues the stored stage record, so the durations add +up over attempts, and `job_attempts` counts them. + ## Metric resource identity Every process that emits OpenTelemetry metrics must have a unique diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/job_progress.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/job_progress.py index 53884bb31..cf362ecb9 100644 --- a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/job_progress.py +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/job_progress.py @@ -83,6 +83,9 @@ class JobProgressRecord(BaseModel): finished_at: AwareDatetime | None = None # None only between the worker start and its choice of execution path. execution: JobExecution | None = None + # Modal can run the same call again after a preemption; each run of the + # call is one attempt, and later attempts continue the same record. + attempts: int = Field(default=1, ge=1) segments: JobSegmentProgress | None = None stage_durations: dict[JobStage, NonNegativeFloat] = Field(default_factory=dict) diff --git a/libs/policyengine-simulation-contract/tests/test_job_progress.py b/libs/policyengine-simulation-contract/tests/test_job_progress.py index 1da0a40de..196d50b83 100644 --- a/libs/policyengine-simulation-contract/tests/test_job_progress.py +++ b/libs/policyengine-simulation-contract/tests/test_job_progress.py @@ -334,6 +334,12 @@ def test__completed_segments_cannot_exceed_total(self): with pytest.raises(ValidationError): JobSegmentProgress(completed=21, total=20) + def test__record_without_attempts_counts_one_attempt(self): + payload = _record().model_dump(mode="json") + payload.pop("attempts") + + assert JobProgressRecord.model_validate(payload).attempts == 1 + def test__record_requires_timezone_aware_times(self): with pytest.raises(ValidationError): _record(started_at=datetime(2026, 10, 13, 12, 0, 30)) diff --git a/projects/policyengine-simulation-executor/src/modal/app.py b/projects/policyengine-simulation-executor/src/modal/app.py index 9594554f2..64d6a51f2 100644 --- a/projects/policyengine-simulation-executor/src/modal/app.py +++ b/projects/policyengine-simulation-executor/src/modal/app.py @@ -264,8 +264,13 @@ def _job_progress_reporter(): The gateway's ``job_id`` is the spawned call's ``object_id``, which is this call's ``current_function_call_id()``. Without one there is nothing to key the record by, so nothing is reported. + + 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; the writes are a few small Dict puts per year. """ from policyengine_simulation_contract.job_progress_state import ( + get_job_progress, put_job_progress, ) from policyengine_simulation_executor.job_progress import ( @@ -279,24 +284,28 @@ def _job_progress_reporter(): job_id = None if not job_id: return NULL_JOB_PROGRESS - return StoredJobProgressReporter(job_id, write=put_job_progress) + return StoredJobProgressReporter( + job_id, write=put_job_progress, read=get_job_progress + ) def _record_stage_durations(runtime, progress) -> None: """Attach the job's stage durations to its operation telemetry. Main-path traces cannot split the baseline from the reform calculation: - both run inside ``economic_impact_analysis``. These attributes can. + both run inside ``economic_impact_analysis``. These attributes can. The + durations add up over attempts when Modal runs the call again. """ durations = progress.stage_durations() if not durations: return try: runtime.set_context( + job_attempts=progress.attempts(), **{ f"job_stage_seconds_{stage.value}": seconds for stage, seconds in durations.items() - } + }, ) except Exception: # noqa: BLE001 - telemetry is non-fatal pass diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/job_progress.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/job_progress.py index efed6757d..d3c2e4684 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/job_progress.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/job_progress.py @@ -6,7 +6,8 @@ no-op ``NULL_JOB_PROGRESS``. Stage reporting is diagnostic for clients and must never change a -calculation: every write failure is caught and logged, and the job goes on. +calculation: every read and write failure is caught and logged, and the job +goes on. """ from __future__ import annotations @@ -21,11 +22,13 @@ JobProgressRecord, JobSegmentProgress, JobStage, + stage_progress_floor, ) logger = logging.getLogger(__name__) ProgressWriter = Callable[[str, JobProgressRecord], None] +ProgressReader = Callable[[str], JobProgressRecord | None] def _utc_now() -> datetime: @@ -47,6 +50,8 @@ def finish(self) -> None: ... def stage_durations(self) -> dict[JobStage, float]: ... + def attempts(self) -> int: ... + class NullJobProgressReporter: """Reports nothing.""" @@ -72,24 +77,41 @@ def finish(self) -> None: def stage_durations(self) -> dict[JobStage, float]: return {} + def attempts(self) -> int: + return 1 + NULL_JOB_PROGRESS = NullJobProgressReporter() class StoredJobProgressReporter: - """Writes the job's ``JobProgressRecord`` after every change.""" + """Writes the job's ``JobProgressRecord`` after every change. + + Modal can run the same call again after a preemption, under the same + call id. ``start()`` therefore continues any record that an earlier + attempt wrote: ``started_at`` and the ended-stage durations carry over, + ``attempts`` counts up, and the published stage stays at the earlier + attempt's stage until this attempt reaches a stage that starts at least + as far along. ``progress`` therefore never moves backwards. + """ def __init__( self, job_id: str, *, write: ProgressWriter, + read: ProgressReader | None = None, clock: Callable[[], datetime] = _utc_now, ): self.job_id = job_id self._write = write + self._read = read self._clock = clock + # This attempt's own state; stage durations come from it. self._record: JobProgressRecord | None = None + # The stage and its start that an earlier attempt had reached. + self._held: tuple[JobStage, datetime] | None = None + self._read_failed = False self._write_failed = False @property @@ -100,11 +122,23 @@ def start(self) -> None: if self._record is not None: return now = self._clock() - self._record = JobProgressRecord( - stage=JobStage.LOADING_DATA, - started_at=now, - stage_started_at=now, - ) + earlier = self._earlier_attempt() + if earlier is None: + self._record = JobProgressRecord( + stage=JobStage.LOADING_DATA, + started_at=now, + stage_started_at=now, + ) + else: + self._record = JobProgressRecord( + stage=JobStage.LOADING_DATA, + started_at=earlier.started_at, + stage_started_at=now, + execution=earlier.execution, + stage_durations=dict(earlier.stage_durations), + attempts=earlier.attempts + 1, + ) + self._held = (earlier.stage, earlier.stage_started_at) self._publish() def set_execution(self, execution: JobExecution) -> None: @@ -159,15 +193,51 @@ def finish(self) -> None: if record is None or record.finished_at is not None: return record.finished_at = self._close_stage(record) + # A failed job reports the stage this attempt stopped in. + self._held = None self._publish() def stage_durations(self) -> dict[JobStage, float]: - """Seconds per ended stage; complete only after ``finish()``.""" + """Seconds per ended stage over all attempts; complete after ``finish()``.""" if self._record is None: return {} return dict(self._record.stage_durations) + def attempts(self) -> int: + return 1 if self._record is None else self._record.attempts + + def _earlier_attempt(self) -> JobProgressRecord | None: + if self._read is None: + return None + try: + return self._read(self.job_id) + except Exception as exc: # noqa: BLE001 - stage reporting is non-fatal + if not self._read_failed: + logger.warning( + "Could not read job progress for %s: %s", + self.job_id, + type(exc).__name__, + ) + self._read_failed = True + return None + + def _published(self, record: JobProgressRecord) -> JobProgressRecord: + """The record clients see: never behind an earlier attempt's stage.""" + + snapshot = record.model_copy(deep=True) + if self._held is None: + return snapshot + held_stage, held_started_at = self._held + if stage_progress_floor(record.execution, record.stage) >= ( + stage_progress_floor(record.execution, held_stage) + ): + self._held = None + return snapshot + snapshot.stage = held_stage + snapshot.stage_started_at = held_started_at + return snapshot + def _close_stage(self, record: JobProgressRecord) -> datetime: now = self._clock() elapsed = max((now - record.stage_started_at).total_seconds(), 0.0) @@ -179,7 +249,7 @@ def _publish(self) -> None: if self._record is None: return try: - self._write(self.job_id, self._record.model_copy(deep=True)) + self._write(self.job_id, self._published(self._record)) except Exception as exc: # noqa: BLE001 - stage reporting is non-fatal if not self._write_failed: logger.warning( diff --git a/projects/policyengine-simulation-executor/tests/test_job_progress.py b/projects/policyengine-simulation-executor/tests/test_job_progress.py index 5b3eaa8b5..6e53a4477 100644 --- a/projects/policyengine-simulation-executor/tests/test_job_progress.py +++ b/projects/policyengine-simulation-executor/tests/test_job_progress.py @@ -18,6 +18,7 @@ JobProgressRecord, JobSegmentProgress, JobStage, + job_progress_fields, ) from policyengine_simulation_executor import simulation_runtime as sr @@ -251,6 +252,153 @@ def failing_write(job_id, record): assert "modal unavailable" not in warnings[0].message +def _earlier_attempt(**overrides) -> JobProgressRecord: + values = { + "stage": JobStage.REFORM, + "execution": JobExecution.MONOLITHIC, + "started_at": T0 - timedelta(seconds=100), + "stage_started_at": T0 - timedelta(seconds=30), + "stage_durations": {JobStage.LOADING_DATA: 40.0, JobStage.BASELINE: 30.0}, + } + return JobProgressRecord(**{**values, **overrides}) + + +def _restarted_reporter(earlier): + clock = _Clock() + writer = _RecordingWriter() + reporter = StoredJobProgressReporter( + "fc-job-1", write=writer, read=lambda job_id: earlier, clock=clock + ) + return reporter, writer, clock + + +def _running_progress(record: JobProgressRecord) -> int: + return job_progress_fields( + status="running", submitted_at=None, record=record, reports_stages=True + )["progress"] + + +class TestRestartedAttempt: + """Modal can run the same call again after a preemption (same call id).""" + + def test__restart_continues_the_earlier_record(self): + reporter, writer, _ = _restarted_reporter(_earlier_attempt()) + + reporter.start() + + first = writer.writes[0][1] + assert first.started_at == T0 - timedelta(seconds=100) + assert first.attempts == 2 + assert first.execution is JobExecution.MONOLITHIC + assert first.stage_durations == { + JobStage.LOADING_DATA: 40.0, + JobStage.BASELINE: 30.0, + } + # Clients still see the stage the earlier attempt reached. + assert first.stage is JobStage.REFORM + assert first.stage_started_at == T0 - timedelta(seconds=30) + + def test__restart_never_moves_progress_backwards(self): + reporter, writer, clock = _restarted_reporter(_earlier_attempt()) + + reporter.start() + clock.advance(10) + reporter.set_execution(JobExecution.MONOLITHIC) + reporter.enter(JobStage.BASELINE) + clock.advance(20) + reporter.enter(JobStage.REFORM) + clock.advance(30) + reporter.enter(JobStage.OUTPUTS) + + progress = [_running_progress(record) for _, record in writer.writes] + assert progress == sorted(progress) + assert progress[0] == 67 + assert progress[-1] == 97 + # Held at the earlier attempt's reform stage, then this attempt's own. + reform_starts = [ + record.stage_started_at + for _, record in writer.writes + if record.stage is JobStage.REFORM + ] + assert reform_starts[0] == T0 - timedelta(seconds=30) + assert reform_starts[-1] == T0 + timedelta(seconds=30) + + def test__durations_add_up_over_attempts(self): + reporter, writer, clock = _restarted_reporter(_earlier_attempt()) + + reporter.start() + clock.advance(5) + reporter.enter(JobStage.BASELINE) + clock.advance(10) + reporter.finish() + + assert writer.last.stage_durations == { + JobStage.LOADING_DATA: 45.0, + JobStage.BASELINE: 40.0, + } + assert reporter.attempts() == 2 + + def test__failure_reports_the_stage_this_attempt_stopped_in(self): + reporter, writer, clock = _restarted_reporter(_earlier_attempt()) + + reporter.start() + clock.advance(5) + reporter.finish() + + assert writer.last.stage is JobStage.LOADING_DATA + assert writer.last.finished_at == T0 + timedelta(seconds=5) + + def test__a_finished_earlier_attempt_is_continued_too(self): + # A preempted attempt can run its ``finally`` block before it stops. + earlier = _earlier_attempt(finished_at=T0 - timedelta(seconds=1)) + reporter, writer, _ = _restarted_reporter(earlier) + + reporter.start() + + assert writer.last.attempts == 2 + assert writer.last.finished_at is None + assert writer.last.stage is JobStage.REFORM + + def test__segmented_restart_is_held_until_simulating(self): + earlier = _earlier_attempt( + stage=JobStage.SIMULATING, + execution=JobExecution.SEGMENTED, + segments=JobSegmentProgress(completed=15, total=20), + ) + reporter, writer, _ = _restarted_reporter(earlier) + + reporter.start() + reporter.set_execution(JobExecution.SEGMENTED) + reporter.start_segments(20) + + # The earlier record already says segmented, so set_execution writes + # nothing; every write shows simulating. + assert [record.stage for _, record in writer.writes] == [ + JobStage.SIMULATING, + JobStage.SIMULATING, + JobStage.SIMULATING, + ] + assert writer.last.segments == JobSegmentProgress(completed=0, total=20) + progress = [_running_progress(record) for _, record in writer.writes] + assert progress == sorted(progress) + + def test__unreadable_earlier_record_starts_fresh(self, caplog): + def failing_read(job_id): + raise ConnectionError("modal unavailable") + + writer = _RecordingWriter() + reporter = StoredJobProgressReporter( + "fc-job-1", write=writer, read=failing_read, clock=_Clock() + ) + + with caplog.at_level(logging.WARNING): + reporter.start() + + assert writer.last.attempts == 1 + assert writer.last.started_at == T0 + assert "ConnectionError" in caplog.text + + class TestStageReportingSimulation: def test__callback_runs_once_before_the_first_ensure(self, monkeypatch): from policyengine.core import Simulation @@ -565,7 +713,7 @@ def _stage_attributes(events) -> list[tuple[int, dict]]: class TestModalRunSimulation: - def _wire(self, app, monkeypatch, dispatch): + def _wire(self, app, monkeypatch, dispatch, earlier=None): from policyengine_simulation_contract import job_progress_state events = [] @@ -578,6 +726,9 @@ def _wire(self, app, monkeypatch, dispatch): "put_job_progress", lambda job_id, record: writes.append((job_id, record)), ) + monkeypatch.setattr( + job_progress_state, "get_job_progress", lambda job_id: earlier + ) monkeypatch.setattr( app, "_configure_modal_observability", @@ -604,13 +755,31 @@ def dispatch(params, *, app_name, runtime, progress): assert writes[-1][1].finished_at is not None [(index, attributes)] = _stage_attributes(events) assert set(attributes) == { + "job_attempts", "job_stage_seconds_loading_data", "job_stage_seconds_reform", } + assert attributes["job_attempts"] == 1 # Durations reach the operation before it ends, then the runtime stops. assert index < events.index("operation_end") assert events[-1] == "shutdown" + def test__rerun_of_the_same_call_continues_its_record(self, modal_app, monkeypatch): + def dispatch(params, *, app_name, runtime, progress): + progress.set_execution(JobExecution.MONOLITHIC) + return {"ok": True} + + events, writes = self._wire( + modal_app, monkeypatch, dispatch, earlier=_earlier_attempt() + ) + + assert modal_app.run_simulation({"country": "us"}) == {"ok": True} + + assert writes[0][1].stage is JobStage.REFORM + assert writes[0][1].attempts == 2 + [(_, attributes)] = _stage_attributes(events) + assert attributes["job_attempts"] == 2 + def test__failed_job_still_records_where_it_stopped(self, modal_app, monkeypatch): def dispatch(params, *, app_name, runtime, progress): progress.enter(JobStage.BASELINE) diff --git a/projects/policyengine-simulation-gateway/README.md b/projects/policyengine-simulation-gateway/README.md index 22730e7aa..a6196235f 100644 --- a/projects/policyengine-simulation-gateway/README.md +++ b/projects/policyengine-simulation-gateway/README.md @@ -58,7 +58,12 @@ in 20 region segments, so it moves through `loading_data`, `simulating` and `outputs`. The `run_simulation` worker writes a stage record to the Modal Dict -`simulation-api-job-progress`, keyed by `job_id`. The gateway reports +`simulation-api-job-progress`, keyed by `job_id`. If Modal runs the same call +again after a preemption, the new attempt continues that record and keeps +the reported stage until it reaches a stage at least as far along, so +`progress` never moves backwards. Budget-window children are ordinary +`run_simulation` calls and write their own records, which +`/budget-window-jobs` does not read yet. The gateway reports `queued` for a job without a record only when the worker app's bundle manifest declares `job_progress`; apps deployed before stage reporting never write a record, so their jobs get no stage fields. A failed read or a From c7388ef7c3d9ae9c7f02846cb5e8e095442ffc71 Mon Sep 17 00:00:00 2001 From: PavelMakarchuk Date: Wed, 14 Oct 2026 17:51:28 -0400 Subject: [PATCH 5/7] Import the macro output fixture from its new module main renamed fixtures/test_simulation_api_contracts.py to fixtures/macro_output.py, so the job progress tests failed to collect in CI against the merged branch. Co-Authored-By: Claude Opus 5.5 --- .../policyengine-simulation-executor/tests/test_job_progress.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/projects/policyengine-simulation-executor/tests/test_job_progress.py b/projects/policyengine-simulation-executor/tests/test_job_progress.py index 6e53a4477..20321d8f1 100644 --- a/projects/policyengine-simulation-executor/tests/test_job_progress.py +++ b/projects/policyengine-simulation-executor/tests/test_job_progress.py @@ -11,8 +11,8 @@ import pytest from conftest import NoOpObservabilityRuntime from fixtures.fake_modal import install_fake_modal +from fixtures.macro_output import CURRENT_SINGLE_YEAR_MACRO_RESULT from fixtures.spm_doubles import installed_spm_selection, spm_capable_simulation -from fixtures.test_simulation_api_contracts import CURRENT_SINGLE_YEAR_MACRO_RESULT from policyengine_simulation_contract.job_progress import ( JobExecution, JobProgressRecord, From 6736d7756ff5bfa7f3a7fd909d2919bcaf8b58b6 Mon Sep 17 00:00:00 2001 From: PavelMakarchuk Date: Thu, 15 Oct 2026 08:54:20 -0400 Subject: [PATCH 6/7] Verify stage reporting on real jobs in the beta integration tests The beta deployment's integration tests submit real US national, US state and UK jobs and poll them to completion. They now record every running and complete status body and check the stage fields: a known stage and progress <= progress_max <= 99 while running, progress that never moves backwards, and progress 100, started_at, finished_at and stage_durations on completion. A Stage 12-only deployment runs these tests against the existing gateway, which may predate stage reporting. The check applies only when the gateway reports submitted_at, which every gateway with stage reporting does. Refs #780 Co-Authored-By: Claude Opus 5.5 --- .../tests/simulation/test_calculate.py | 187 +++++++++++++++++- 1 file changed, 183 insertions(+), 4 deletions(-) diff --git a/projects/policyengine-apis-integ/tests/simulation/test_calculate.py b/projects/policyengine-apis-integ/tests/simulation/test_calculate.py index b636f165b..d2feb0d5a 100644 --- a/projects/policyengine-apis-integ/tests/simulation/test_calculate.py +++ b/projects/policyengine-apis-integ/tests/simulation/test_calculate.py @@ -5,6 +5,7 @@ that economy-wide simulations complete successfully. """ +import json import time from collections.abc import Mapping from dataclasses import dataclass, field @@ -37,6 +38,7 @@ "population", } _AT_LARGE_DISTRICT_IDS = {"AK-01", "DC-01", "DE-01", "ND-01", "SD-01", "VT-01", "WY-01"} +_JOB_STAGES = {"queued", "loading_data", "baseline", "reform", "simulating", "outputs"} class _ObjectWithToDict(Protocol): @@ -71,6 +73,7 @@ def poll_for_completion( job_id: str, max_wait_seconds: float, poll_interval: float, + bodies: list[Mapping] | None = None, ) -> JobStatusResponse: """ Poll for job completion. @@ -80,6 +83,7 @@ def poll_for_completion( job_id: The job ID to poll max_wait_seconds: Maximum time to wait for completion poll_interval: Time between polls in seconds + bodies: When given, receives every running and complete status body Returns: The final JobStatusResponse @@ -95,6 +99,12 @@ def poll_for_completion( job_id=job_id, client=client ) + if bodies is not None and response.status_code in { + HTTPStatus.OK, + HTTPStatus.ACCEPTED, + }: + bodies.append(json.loads(response.content)) + if response.status_code == HTTPStatus.OK: assert isinstance(response.parsed, JobStatusResponse) assert response.parsed.status == "complete", ( @@ -123,6 +133,42 @@ def poll_for_completion( raise TimeoutError(f"Job {job_id} did not complete within {max_wait_seconds}s") +def assert_job_progress_reported(bodies: list[Mapping]) -> None: + """Verify the stage fields of ``GET /jobs/{job_id}`` (policyengine-sim-api#781). + + A gateway with stage reporting adds ``submitted_at`` to every new job. A + Stage 12-only deployment runs these tests against the existing gateway, + which may predate stage reporting; then there is nothing to verify. + """ + assert bodies, "No status body was recorded" + final = bodies[-1] + if "submitted_at" not in final: + return + + progress_seen: list[int] = [] + for body in bodies[:-1]: + assert body["status"] == "running", f"Unexpected status body: {body}" + assert body.get("stage") in _JOB_STAGES, f"Missing or unknown stage: {body}" + progress = body["progress"] + assert 0 <= progress <= body["progress_max"] <= 99, ( + f"Running progress out of range: {body}" + ) + progress_seen.append(progress) + assert progress_seen == sorted(progress_seen), ( + f"Progress moved backwards: {progress_seen}" + ) + + assert final["status"] == "complete" + assert final.get("progress") == 100, f"Complete job without progress 100: {final}" + assert "stage" not in final + assert "progress_max" not in final + for key in ("started_at", "finished_at"): + assert key in final, f"Complete job without {key}: {final}" + assert final.get("stage_durations"), ( + f"Complete job without stage durations: {final}" + ) + + def submit_simulation_request( client: Client | AuthenticatedClient, request: SimulationRequest, @@ -338,6 +384,115 @@ def test_poll_for_completion_honors_retry_after_header( assert sleep_seconds == [10.0] +def _running_body(stage: str, progress: int, progress_max: int) -> dict: + return { + "status": "running", + "result": None, + "error": None, + "submitted_at": "2026-10-15T12:00:00Z", + "stage": stage, + "progress": progress, + "progress_max": progress_max, + } + + +_COMPLETE_BODY = { + "status": "complete", + "result": {}, + "submitted_at": "2026-10-15T12:00:00Z", + "started_at": "2026-10-15T12:00:10Z", + "finished_at": "2026-10-15T12:02:00Z", + "progress": 100, + "stage_durations": {"loading_data": 12.5, "baseline": 20.0, "reform": 60.0}, +} + + +def test_poll_for_completion_records_running_and_complete_bodies( + monkeypatch: pytest.MonkeyPatch, +) -> None: + running = _running_body("reform", 67, 96) + bodies: list[Mapping] = [] + response_iterator = iter( + [ + _StubJobStatusHttpResponse( + status_code=HTTPStatus.ACCEPTED, + content=json.dumps(running).encode(), + ), + _StubJobStatusHttpResponse( + status_code=HTTPStatus.OK, + parsed=JobStatusResponse.from_dict( + {"status": "complete", "result": {}} + ), + content=json.dumps(_COMPLETE_BODY).encode(), + ), + ] + ) + monkeypatch.setattr( + get_job_status_jobs_job_id_get, + "sync_detailed", + lambda *, job_id, client: next(response_iterator), + ) + monkeypatch.setattr(time, "sleep", lambda seconds: None) + + poll_for_completion( + cast(Client, object()), + "fc-test", + max_wait_seconds=1.0, + poll_interval=0.25, + bodies=bodies, + ) + + assert bodies == [running, _COMPLETE_BODY] + + +def test_job_progress_check_accepts_a_reporting_job() -> None: + assert_job_progress_reported( + [ + _running_body("queued", 0, 24), + _running_body("loading_data", 25, 54), + _running_body("reform", 67, 96), + _COMPLETE_BODY, + ] + ) + + +def test_job_progress_check_rejects_progress_that_moves_backwards() -> None: + with pytest.raises(AssertionError, match="backwards"): + assert_job_progress_reported( + [ + _running_body("reform", 67, 96), + _running_body("loading_data", 25, 54), + _COMPLETE_BODY, + ] + ) + + +def test_job_progress_check_rejects_a_complete_job_without_progress() -> None: + complete = { + key: value for key, value in _COMPLETE_BODY.items() if key != "progress" + } + + with pytest.raises(AssertionError, match="progress 100"): + assert_job_progress_reported([complete]) + + +def test_job_progress_check_rejects_a_running_job_without_a_stage() -> None: + running = _running_body("reform", 67, 96) + del running["stage"] + + with pytest.raises(AssertionError, match="stage"): + assert_job_progress_reported([running, _COMPLETE_BODY]) + + +def test_job_progress_check_skips_a_gateway_without_stage_reporting() -> None: + assert_job_progress_reported( + [ + {"status": "running", "result": None, "error": None}, + {"status": "complete", "result": {}}, + ] + ) + + @pytest.mark.beta_only def test_calculate_default_model( client: Client | AuthenticatedClient, @@ -367,7 +522,13 @@ def test_calculate_default_model( job_id = submit_response.job_id # When - poll for completion - result = poll_for_completion(client, job_id, max_wait_seconds, poll_interval) + bodies: list[Mapping] = [] + result = poll_for_completion( + client, job_id, max_wait_seconds, poll_interval, bodies=bodies + ) + + # Then - the job reported its stages while running + assert_job_progress_reported(bodies) # Then - verify result structure assert result.status == "complete" @@ -415,7 +576,13 @@ def test_calculate_us_state_region_model( job_id = submit_response.job_id # When - poll for completion - result = poll_for_completion(client, job_id, max_wait_seconds, poll_interval) + bodies: list[Mapping] = [] + result = poll_for_completion( + client, job_id, max_wait_seconds, poll_interval, bodies=bodies + ) + + # Then - the job reported its stages while running + assert_job_progress_reported(bodies) # Then - verify result structure assert result.status == "complete" @@ -465,7 +632,13 @@ def test_calculate_specific_model( job_id = submit_response.job_id # When - poll for completion - result = poll_for_completion(client, job_id, max_wait_seconds, poll_interval) + bodies: list[Mapping] = [] + result = poll_for_completion( + client, job_id, max_wait_seconds, poll_interval, bodies=bodies + ) + + # Then - the job reported its stages while running + assert_job_progress_reported(bodies) # Then - verify result structure assert result.status == "complete" @@ -511,7 +684,13 @@ def test_calculate_uk_model( job_id = submit_response.job_id # When - poll for completion - result = poll_for_completion(client, job_id, max_wait_seconds, poll_interval) + bodies: list[Mapping] = [] + result = poll_for_completion( + client, job_id, max_wait_seconds, poll_interval, bodies=bodies + ) + + # Then - the job reported its stages while running + assert_job_progress_reported(bodies) # Then - verify result structure assert result.status == "complete" From 7cb1dc5a3fd1e4d380a4c1bc59aad959502c5c18 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 15 Oct 2026 22:04:25 +0400 Subject: [PATCH 7/7] Keep a preempted job's stage record running for its next attempt Modal preempts a call with SIGINT, which run_simulation receives as KeyboardInterrupt, and then runs the call again under the same call id. The finally block called finish(), which marked the record finished, dropped the stage an earlier attempt had reached and published this attempt's own stage. A second preemption therefore moved progress backwards: an attempt preempted while loading data after an earlier one reached the reform published 25 instead of 67, and the next attempt held loading_data. run_simulation now calls the reporter's new interrupt() on KeyboardInterrupt, both during the calculation and before it starts. It counts the current stage's time but leaves the record running with any held stage, and finish() does nothing afterwards. A failed job still reports the stage it stopped in. The beta integration check now takes the stage list from the generated client's JobStage instead of a hand-written copy. Co-Authored-By: Claude Opus 5.5 --- .../tests/simulation/test_calculate.py | 4 +- .../src/modal/app.py | 9 ++ .../job_progress.py | 31 ++++- .../tests/test_job_progress.py | 114 +++++++++++++++++- .../policyengine-simulation-gateway/README.md | 5 +- 5 files changed, 157 insertions(+), 6 deletions(-) diff --git a/projects/policyengine-apis-integ/tests/simulation/test_calculate.py b/projects/policyengine-apis-integ/tests/simulation/test_calculate.py index 01ed2a1da..d17222c1a 100644 --- a/projects/policyengine-apis-integ/tests/simulation/test_calculate.py +++ b/projects/policyengine-apis-integ/tests/simulation/test_calculate.py @@ -19,6 +19,7 @@ submit_simulation_simulate_economy_comparison_post, ) from policyengine_api_simulation_client.models import ( + JobStage, JobStatusResponse, JobSubmitResponse, SimulationRequest, @@ -37,7 +38,8 @@ "population", } _AT_LARGE_DISTRICT_IDS = {"AK-01", "DC-01", "DE-01", "ND-01", "SD-01", "VT-01", "WY-01"} -_JOB_STAGES = {"queued", "loading_data", "baseline", "reform", "simulating", "outputs"} +# From the generated client, so a stage added to the schema is accepted here. +_JOB_STAGES = frozenset(stage.value for stage in JobStage) class _ObjectWithToDict(Protocol): diff --git a/projects/policyengine-simulation-executor/src/modal/app.py b/projects/policyengine-simulation-executor/src/modal/app.py index 8de1784e6..3cfcd0e27 100644 --- a/projects/policyengine-simulation-executor/src/modal/app.py +++ b/projects/policyengine-simulation-executor/src/modal/app.py @@ -368,11 +368,20 @@ def run_simulation( runtime=runtime, progress=progress, ) + except KeyboardInterrupt: + # Modal preempts with SIGINT and runs the call again; the + # record stays running for that attempt. + progress.interrupt() + raise finally: # Close the last stage inside the operation, so its # completion record carries every stage duration. progress.finish() _record_stage_durations(runtime, progress) + except KeyboardInterrupt: + # Preempted before the calculation started. + progress.interrupt() + raise finally: # A no-op unless the operation never started. progress.finish() diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/job_progress.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/job_progress.py index d3c2e4684..18bd36815 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/job_progress.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/job_progress.py @@ -48,6 +48,8 @@ def segment_completed(self) -> None: ... def finish(self) -> None: ... + def interrupt(self) -> None: ... + def stage_durations(self) -> dict[JobStage, float]: ... def attempts(self) -> int: ... @@ -74,6 +76,9 @@ def segment_completed(self) -> None: def finish(self) -> None: pass + def interrupt(self) -> None: + pass + def stage_durations(self) -> dict[JobStage, float]: return {} @@ -92,7 +97,10 @@ class StoredJobProgressReporter: attempt wrote: ``started_at`` and the ended-stage durations carry over, ``attempts`` counts up, and the published stage stays at the earlier attempt's stage until this attempt reaches a stage that starts at least - as far along. ``progress`` therefore never moves backwards. + as far along. A preempted attempt calls ``interrupt()`` rather than + ``finish()``, so its record stays running and keeps any held stage for + the next attempt. ``progress`` therefore never moves backwards, however + many times the call is preempted. """ def __init__( @@ -113,6 +121,7 @@ def __init__( self._held: tuple[JobStage, datetime] | None = None self._read_failed = False self._write_failed = False + self._interrupted = False @property def record(self) -> JobProgressRecord | None: @@ -190,13 +199,31 @@ def finish(self) -> None: """Close the current stage. Runs after success and after failure.""" record = self._record - if record is None or record.finished_at is not None: + if record is None or record.finished_at is not None or self._interrupted: return record.finished_at = self._close_stage(record) # A failed job reports the stage this attempt stopped in. self._held = None self._publish() + def interrupt(self) -> None: + """Leave the record running for the call's next attempt. + + Modal preempts a call with SIGINT, which a sync function receives as + ``KeyboardInterrupt``, and then runs the call again under the same + call id. The current stage's time still counts in the durations, but + the record does not finish and any held stage stays held, so the next + attempt continues from the furthest stage reached. ``finish()`` does + nothing afterwards. + """ + + record = self._record + if record is None or record.finished_at is not None or self._interrupted: + return + self._interrupted = True + self._close_stage(record) + self._publish() + def stage_durations(self) -> dict[JobStage, float]: """Seconds per ended stage over all attempts; complete after ``finish()``.""" diff --git a/projects/policyengine-simulation-executor/tests/test_job_progress.py b/projects/policyengine-simulation-executor/tests/test_job_progress.py index ea1e1bde9..47140a503 100644 --- a/projects/policyengine-simulation-executor/tests/test_job_progress.py +++ b/projects/policyengine-simulation-executor/tests/test_job_progress.py @@ -87,6 +87,9 @@ def segment_completed(self) -> None: def finish(self) -> None: self.calls.append(("finish",)) + def interrupt(self) -> None: + self.calls.append(("interrupt",)) + def stage_durations(self) -> dict[JobStage, float]: return {} @@ -348,7 +351,8 @@ def test__failure_reports_the_stage_this_attempt_stopped_in(self): assert writer.last.finished_at == T0 + timedelta(seconds=5) def test__a_finished_earlier_attempt_is_continued_too(self): - # A preempted attempt can run its ``finally`` block before it stops. + # An attempt stopped by something other than the SIGINT the worker + # handles (or an older worker) can leave a finished record behind. earlier = _earlier_attempt(finished_at=T0 - timedelta(seconds=1)) reporter, writer, _ = _restarted_reporter(earlier) @@ -358,6 +362,72 @@ def test__a_finished_earlier_attempt_is_continued_too(self): assert writer.last.finished_at is None assert writer.last.stage is JobStage.REFORM + def test__preemption_leaves_the_record_running_for_the_next_attempt(self): + reporter, writer, clock = _restarted_reporter(_earlier_attempt()) + + reporter.start() + clock.advance(5) + reporter.interrupt() + reporter.finish() + + # start and interrupt write; finish writes nothing after interrupt. + assert len(writer.writes) == 2 + assert writer.last.finished_at is None + # The earlier attempt's stage stays held for the next attempt. + assert writer.last.stage is JobStage.REFORM + assert writer.last.stage_started_at == T0 - timedelta(seconds=30) + # This attempt's time still counts. + assert writer.last.stage_durations[JobStage.LOADING_DATA] == 45.0 + assert writer.last.attempts == 2 + + def test__repeated_preemptions_never_move_progress_backwards(self): + store: dict[str, JobProgressRecord] = {} + published: list[JobProgressRecord] = [] + + def write(job_id: str, record: JobProgressRecord) -> None: + store[job_id] = record + published.append(record) + + def attempt() -> tuple[StoredJobProgressReporter, _Clock]: + clock = _Clock() + reporter = StoredJobProgressReporter( + "fc-job-1", write=write, read=store.get, clock=clock + ) + reporter.start() + return reporter, clock + + # Attempt 1 reaches the reform, then is preempted. + first, clock = attempt() + first.set_execution(JobExecution.MONOLITHIC) + first.enter(JobStage.BASELINE) + first.enter(JobStage.REFORM) + first.interrupt() + first.finish() + # Attempt 2 is preempted while still loading data. + second, clock = attempt() + clock.advance(5) + second.interrupt() + second.finish() + # Attempt 3 still starts from the reform. + attempt() + + progress = [_running_progress(record) for record in published] + assert progress == sorted(progress) + assert store["fc-job-1"].stage is JobStage.REFORM + assert store["fc-job-1"].attempts == 3 + assert all(record.finished_at is None for record in published) + + def test__interrupt_after_finish_writes_nothing(self): + reporter, writer, _ = _reporter() + reporter.start() + reporter.finish() + writes = len(writer.writes) + + reporter.interrupt() + + assert len(writer.writes) == writes + assert writer.last.finished_at is not None + def test__segmented_restart_is_held_until_simulating(self): earlier = _earlier_attempt( stage=JobStage.SIMULATING, @@ -793,6 +863,48 @@ def dispatch(params, *, app_name, runtime, progress): assert "job_stage_seconds_baseline" in attributes assert index < events.index("operation_end") + def test__preempted_job_stays_running_for_the_next_attempt( + self, modal_app, monkeypatch + ): + def dispatch(params, *, app_name, runtime, progress): + progress.enter(JobStage.BASELINE) + raise KeyboardInterrupt + + events, writes = self._wire(modal_app, monkeypatch, dispatch) + + with pytest.raises(KeyboardInterrupt): + modal_app.run_simulation({"country": "us"}) + + assert writes[-1][1].stage is JobStage.BASELINE + assert writes[-1][1].finished_at is None + # The preempted attempt's durations still reach its operation. + [(index, attributes)] = _stage_attributes(events) + assert "job_stage_seconds_baseline" in attributes + assert index < events.index("operation_end") + assert events[-1] == "shutdown" + + def test__preemption_before_the_calculation_stays_running( + self, modal_app, monkeypatch + ): + def dispatch(params, *, app_name, runtime, progress): + raise AssertionError("dispatch should not run") + + _, writes = self._wire( + modal_app, monkeypatch, dispatch, earlier=_earlier_attempt() + ) + + def preempted(runtime): + raise KeyboardInterrupt + + monkeypatch.setattr(modal_app, "_set_modal_call_attributes", preempted) + + with pytest.raises(KeyboardInterrupt): + modal_app.run_simulation({"country": "us"}) + + assert writes[-1][1].finished_at is None + assert writes[-1][1].stage is JobStage.REFORM + assert writes[-1][1].attempts == 2 + def test__store_failure_does_not_change_the_result(self, modal_app, monkeypatch): from policyengine_simulation_contract import job_progress_state diff --git a/projects/policyengine-simulation-gateway/README.md b/projects/policyengine-simulation-gateway/README.md index a6196235f..adea361f5 100644 --- a/projects/policyengine-simulation-gateway/README.md +++ b/projects/policyengine-simulation-gateway/README.md @@ -59,9 +59,10 @@ in 20 region segments, so it moves through `loading_data`, `simulating` and The `run_simulation` worker writes a stage record to the Modal Dict `simulation-api-job-progress`, keyed by `job_id`. If Modal runs the same call -again after a preemption, the new attempt continues that record and keeps +again after a preemption, the preempted attempt leaves its record running +rather than finishing it, and the new attempt continues that record and keeps the reported stage until it reaches a stage at least as far along, so -`progress` never moves backwards. Budget-window children are ordinary +`progress` never moves backwards, however many times the call is preempted. Budget-window children are ordinary `run_simulation` calls and write their own records, which `/budget-window-jobs` does not read yet. The gateway reports `queued` for a job without a record only when the worker app's bundle