Report job stage and progress on annual economy job status

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_<stage>` 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

libs/policyengine-simulation-contract/src/policyengine_simulation_contract/gateway_models.py

@@ -6,6 +6,7 @@
 from typing import ClassVar, Literal, Optional
 
 from pydantic import (
+    AwareDatetime,
     BaseModel,
     ConfigDict,
     Field,
@@ -14,6 +15,10 @@
     model_validator,
 )
 
+from policyengine_simulation_contract.job_progress import (
+    JobSegmentProgress,
+    JobStage,
+)
 from policyengine_simulation_contract.spm import (
     SPMSelection,
     SPMCapability,
@@ -204,6 +209,56 @@ class JobStatusResponse(BaseModel):
     error: Optional[str] = None
     resolved_app_name: Optional[str] = None
     policyengine_bundle: Optional[PolicyEngineBundle] = None
+    # Stage reporting. Present only for jobs whose worker reports stages;
+    # see policyengine_simulation_contract.job_progress.
+    stage: JobStage | None = Field(
+        default=None,
+        description=(
+            "Current stage of a running job, or the stage in which a failed "
+            "job stopped. Absent for complete jobs. 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):

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}

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")

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

@@ -7,7 +7,7 @@
 
 import pytest
 
-from policyengine_simulation_contract import budget_window_state
+from policyengine_simulation_contract import budget_window_state, job_progress_state
 
 
 class MockDict:
@@ -29,7 +29,7 @@ def get(self, key, default=None):
 
 @pytest.fixture
 def mock_modal(monkeypatch):
-    """Patch modal.Dict in budget_window_state with an in-memory store."""
+    """Patch modal.Dict in the Dict-backed state modules with an in-memory store."""
     mock_dicts: dict[str, dict] = {}
 
     class MockModalDict:
@@ -45,4 +45,5 @@ class MockModal:
         Dict = MockModalDict
 
     monkeypatch.setattr(budget_window_state, "modal", MockModal)
+    monkeypatch.setattr(job_progress_state, "modal", MockModal)
     return {"dicts": mock_dicts}