diff --git a/docs/engineering/skills/observability.md b/docs/engineering/skills/observability.md index 488f16822..4f9e6b1c5 100644 --- a/docs/engineering/skills/observability.md +++ b/docs/engineering/skills/observability.md @@ -102,6 +102,20 @@ 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. + +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/gateway_models.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/gateway_models.py index f4580ed20..d7718b068 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 @@ -7,6 +7,7 @@ from policyengine_simulation_observability.telemetry import TelemetryEnvelope from pydantic import ( + AwareDatetime, BaseModel, ConfigDict, Field, @@ -15,6 +16,10 @@ model_validator, ) +from policyengine_simulation_contract.job_progress import ( + JobSegmentProgress, + JobStage, +) from policyengine_simulation_contract.json_types import JsonObject from policyengine_simulation_contract.macro_output import SingleYearMacroOutput from policyengine_simulation_contract.spm import ( @@ -204,6 +209,56 @@ class JobStatusResponse(BaseModel): error: str | None = None resolved_app_name: str | None = None policyengine_bundle: PolicyEngineBundle | None = 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. Later versions may add " + "stages; clients must accept a value they do not know." + ), + ) + progress: int | None = Field( + default=None, + ge=0, + le=100, + 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: 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." + ) + 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..cf362ecb9 --- /dev/null +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/job_progress.py @@ -0,0 +1,247 @@ +"""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``. + + Later versions may add stages. Clients must accept a value they do not + know, for example by showing a generic "running" state. + """ + + QUEUED = "queued" + LOADING_DATA = "loading_data" + BASELINE = "baseline" + REFORM = "reform" + SIMULATING = "simulating" + 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.""" + + 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 + # 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) + + +# 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; 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( + { + 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. +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 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 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( + *, + 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(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) + if status != "complete": + fields["stage"] = record.stage.value + fields["stage_started_at"] = _json_datetime(record.stage_started_at) + if status == "running": + 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: + 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 36ff49a60..7ccaff8cb 100644 --- a/libs/policyengine-simulation-contract/tests/conftest.py +++ b/libs/policyengine-simulation-contract/tests/conftest.py @@ -6,7 +6,7 @@ """ import pytest -from policyengine_simulation_contract import budget_window_state +from policyengine_simulation_contract import budget_window_state, job_progress_state class MockDict: @@ -28,7 +28,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: @@ -44,4 +44,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..d45335c58 --- /dev/null +++ b/libs/policyengine-simulation-contract/tests/test_job_progress.py @@ -0,0 +1,358 @@ +"""Tests for client-visible job stage fields.""" + +from datetime import UTC, datetime + +import pytest +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, + JobExecution, + JobProgressRecord, + JobSegmentProgress, + JobStage, + job_progress_fields, + stage_progress_floor, + stage_progress_max, + supported_job_progress_capability, +) +from policyengine_simulation_contract.job_progress_state import ( + JOB_PROGRESS_DICT_NAME, + get_job_progress, + put_job_progress, +) +from pydantic import ValidationError + +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, + "execution": JobExecution.MONOLITHIC, + "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, + # The path is unknown while queued, so the widest limit applies. + "progress_max": 24, + } + + 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": 67, + "progress_max": 96, + "stage_durations": {"loading_data": 12.5, "baseline": 17.5}, + } + + @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, + 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"] == 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): + 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 + ) + assert fields["progress_max"] == 54 + + def test__running_progress_never_reaches_100(self): + assert all( + floor <= RUNNING_PROGRESS_CAP + for floors in STAGE_PROGRESS_FLOOR.values() + for floor in floors.values() + ) + record = _record(stage=JobStage.OUTPUTS) + + fields = job_progress_fields( + status="running", + submitted_at=SUBMITTED, + record=record, + reports_stages=True, + ) + + 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) + + fields = job_progress_fields( + status="complete", + submitted_at=SUBMITTED, + record=record, + reports_stages=True, + ) + + 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 + + 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 + assert "progress_max" 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 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 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 ( + 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_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)) + + +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 6b8d6afbc..b023ff0bb 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 e3413410b..fbba9f721 100644 --- a/libs/policyengine-simulation-observability/tests/test_stages.py +++ b/libs/policyengine-simulation-observability/tests/test_stages.py @@ -48,3 +48,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-apis-integ/tests/simulation/test_calculate.py b/projects/policyengine-apis-integ/tests/simulation/test_calculate.py index 942539177..01ed2a1da 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 @@ -36,6 +37,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): @@ -70,6 +72,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. @@ -79,6 +82,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 @@ -94,6 +98,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", ( @@ -122,6 +132,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, @@ -337,6 +383,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, @@ -366,7 +521,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" @@ -414,7 +575,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" @@ -464,7 +631,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" @@ -510,7 +683,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" 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 4e5786b01..8de1784e6 100644 --- a/projects/policyengine-simulation-executor/src/modal/app.py +++ b/projects/policyengine-simulation-executor/src/modal/app.py @@ -258,6 +258,59 @@ 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. + + 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 ( + 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, 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. 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 + + @app.function( image=simulation_image, cpu=8.0, @@ -277,7 +330,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 @@ -301,16 +357,25 @@ 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, - ) + 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 afe42d494..b769a6181 100644 --- a/projects/policyengine-simulation-executor/src/modal/segmented_national.py +++ b/projects/policyengine-simulation-executor/src/modal/segmented_national.py @@ -30,11 +30,16 @@ from typing import Any from policyengine_observability import ObservabilityRuntime +from policyengine_simulation_contract.job_progress import JobExecution, JobStage from policyengine_simulation_contract.spm import ( SPMInputError, combine_spm_results, spm_error_detail, ) +from policyengine_simulation_executor.job_progress import ( + NULL_JOB_PROGRESS, + JobProgressReporter, +) from policyengine_simulation_executor.national_partition import ( national_region_groups, ) @@ -185,10 +190,12 @@ def __init__( PARENT_FUNCTION_TIMEOUT_SECONDS - PARENT_COMPLETION_MARGIN_SECONDS ), clock: Callable[[], float] = time.monotonic, + 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 @@ -223,8 +230,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)) children = self._spawn_all_children() child_results = self._collect(children) + self.progress.enter(JobStage.OUTPUTS) return self._reduce(child_results, country_module=country_module) def _country_module(self): @@ -357,6 +366,7 @@ def _collect(self, children: list[ChildExecution]) -> 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) @@ -537,12 +547,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() @@ -551,6 +566,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.""" @@ -560,10 +576,12 @@ 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, runtime=runtime, + progress=progress, ) from policyengine_simulation_executor.simulation_runtime import ( @@ -574,4 +592,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") - return run_simulation_impl(params, runtime=runtime) + progress.set_execution(JobExecution.MONOLITHIC) + 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 4d091f342..6443e6f53 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 @@ -52,6 +52,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): @@ -119,6 +122,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() @@ -127,6 +133,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 e66a59092..3bb7f5039 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 @@ -23,10 +23,12 @@ import logging from typing import Any -from policyengine.core import Simulation from pydantic import PrivateAttr from policyengine_simulation_executor import artifact_keys +from policyengine_simulation_executor.stage_reporting_simulation import ( + StageReportingSimulation, +) logger = logging.getLogger(__name__) @@ -139,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..d3c2e4684 --- /dev/null +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/job_progress.py @@ -0,0 +1,260 @@ +"""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 read and 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 ( + JobExecution, + JobProgressRecord, + JobSegmentProgress, + JobStage, + stage_progress_floor, +) + +logger = logging.getLogger(__name__) + +ProgressWriter = Callable[[str, JobProgressRecord], None] +ProgressReader = Callable[[str], JobProgressRecord | None] + + +def _utc_now() -> datetime: + return datetime.now(UTC) + + +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: ... + + def segment_completed(self) -> None: ... + + def finish(self) -> None: ... + + def stage_durations(self) -> dict[JobStage, float]: ... + + def attempts(self) -> int: ... + + +class NullJobProgressReporter: + """Reports nothing.""" + + def start(self) -> None: + pass + + def set_execution(self, execution: JobExecution) -> 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 + + 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. + + 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 + def record(self) -> JobProgressRecord | None: + return self._record + + def start(self) -> None: + if self._record is not None: + return + now = self._clock() + 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: + 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() + 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) + # 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 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) + 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._published(self._record)) + 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 cd0636b7f..db9106831 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_observability.stages import ( ANNUAL_IMPACT_STAGES, @@ -22,6 +23,10 @@ simulation_output_labor, simulation_output_poverty, ) +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, @@ -53,6 +58,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: @@ -70,6 +76,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 6004f9283..b1058ad2d 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,12 +21,17 @@ 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_observability.telemetry import split_internal_payload +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, @@ -151,6 +156,7 @@ def run_simulation_impl( params: dict, *, runtime: ObservabilityRuntime, + progress: JobProgressReporter = NULL_JOB_PROGRESS, ) -> dict: """ Execute economic simulation. @@ -162,7 +168,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. @@ -684,13 +690,14 @@ 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.spm import normalize_runtime_spm + from policyengine_simulation_executor.stage_reporting_simulation import ( + StageReportingSimulation, + ) selection = normalize_runtime_spm(params) params = {**params, **({"spm": selection} if selection is not None else {})} @@ -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 8c2f06865..b0a3eb398 100644 --- a/projects/policyengine-simulation-executor/tests/test_baseline_artifacts.py +++ b/projects/policyengine-simulation-executor/tests/test_baseline_artifacts.py @@ -349,8 +349,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, @@ -368,15 +370,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", @@ -456,6 +457,7 @@ def test_bounded_reform_build_keeps_current_law_artifact_reuse(monkeypatch): from policyengine import core as policyengine_core 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( @@ -467,8 +469,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 c06fda7b5..b6c63cf6a 100644 --- a/projects/policyengine-simulation-executor/tests/test_canonical_spm.py +++ b/projects/policyengine-simulation-executor/tests/test_canonical_spm.py @@ -600,9 +600,11 @@ def test_fanout_and_budget_window_preserve_settings_and_receipts(): def test_worker_build_passes_selection_to_baseline_and_reform(monkeypatch): - import policyengine.core import policyengine_simulation_executor.spm as spm - from policyengine_simulation_executor import baseline_artifacts + from policyengine_simulation_executor import ( + baseline_artifacts, + stage_reporting_simulation, + ) calls = [] monkeypatch.setattr(spm, "normalize_runtime_spm", lambda params: SELECTION) @@ -617,7 +619,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..ea1e1bde9 --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_job_progress.py @@ -0,0 +1,820 @@ +"""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.macro_output import CURRENT_SINGLE_YEAR_MACRO_RESULT +from fixtures.spm_doubles import installed_spm_selection, spm_capable_simulation +from policyengine_simulation_contract.job_progress import ( + JobExecution, + JobProgressRecord, + JobSegmentProgress, + JobStage, + job_progress_fields, +) +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 set_execution(self, execution: JobExecution) -> None: + self.calls.append(("set_execution", execution)) + + 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 stage_durations(self) -> dict[JobStage, float]: + return {} + + +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__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() + + 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.set_execution(JobExecution.SEGMENTED) + 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 + + +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 + 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 + 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 +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 _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, earlier=None): + 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( + job_progress_state, "get_job_progress", lambda job_id: earlier + ) + monkeypatch.setattr( + app, + "_configure_modal_observability", + lambda *, service_role: _RecordingRuntime(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 + [(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) + raise RuntimeError("model failed") + + 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 + + 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 6e1a4aadf..09ea172f7 100644 --- a/projects/policyengine-simulation-executor/tests/test_segmented_national.py +++ b/projects/policyengine-simulation-executor/tests/test_segmented_national.py @@ -733,7 +733,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} @@ -907,7 +907,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 2fbb378f7..14c336a17 100644 --- a/projects/policyengine-simulation-executor/tests/test_simulation_output_builder.py +++ b/projects/policyengine-simulation-executor/tests/test_simulation_output_builder.py @@ -32,6 +32,7 @@ UKLocalAuthorityBoundaryVersion, UKLocalAuthorityMetadata, ) +from policyengine_simulation_executor.job_progress import NULL_JOB_PROGRESS from policyengine_simulation_executor.release_bundle import ( BUNDLE_RECEIPT_FILENAME, get_country_release_bundle, @@ -1126,6 +1127,7 @@ def serialize(self): "resolved_region_code": "us", "runtime": runtime, "uk_local_authority_metadata": None, + "progress": NULL_JOB_PROGRESS, } ] @@ -1241,7 +1243,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 74d260851..6aa96bfe6 100644 --- a/projects/policyengine-simulation-executor/tests/test_update_version_registry.py +++ b/projects/policyengine-simulation-executor/tests/test_update_version_registry.py @@ -163,6 +163,26 @@ 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..a6196235f 100644 --- a/projects/policyengine-simulation-gateway/README.md +++ b/projects/policyengine-simulation-gateway/README.md @@ -37,6 +37,57 @@ 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. 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, for captions. | +| `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`. 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 +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`) depends on the execution +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 +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 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 f5ac439a0..2617a8e84 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 mock_func = MockFunction() @@ -354,6 +358,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 d23832d98..1b1e68746 100644 --- a/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/endpoints.py +++ b/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/endpoints.py @@ -5,6 +5,7 @@ import logging import re from dataclasses import dataclass +from datetime import UTC, datetime from typing import TypedDict import modal @@ -36,6 +37,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, @@ -69,6 +77,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): @@ -662,14 +674,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, @@ -707,10 +731,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, *, @@ -853,6 +912,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( @@ -862,7 +922,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 @@ -1018,16 +1084,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) @@ -1037,7 +1111,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( @@ -1046,10 +1120,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 bffd7ec05..3b4cfde8f 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``.\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": { "status": { @@ -1312,6 +1346,126 @@ "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. Later versions may add stages; clients must accept a value they do not know." + }, + "progress": { + "anyOf": [ + { + "type": "integer", + "maximum": 100.0, + "minimum": 0.0 + }, + { + "type": "null" + } + ], + "title": "Progress", + "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": [ + { + "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..8cc4586a0 --- /dev/null +++ b/projects/policyengine-simulation-gateway/tests/test_job_progress.py @@ -0,0 +1,300 @@ +"""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, + JobExecution, + 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, + "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}, + } + 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["progress_max"] == 24 + 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"] == 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_segments_without_moving_progress( + 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, + execution=JobExecution.SEGMENTED, + 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"] == 5 + assert body["progress_max"] == 96 + + 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 "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 + 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