diff --git a/changelog.d/relay-simulation-run-progress.added.md b/changelog.d/relay-simulation-run-progress.added.md new file mode 100644 index 000000000..490594884 --- /dev/null +++ b/changelog.d/relay-simulation-run-progress.added.md @@ -0,0 +1 @@ +Annual economy responses relay the simulation entrypoint's optional stage fields (`stage`, `progress`, `progress_max`, timestamps and `segments`) while a run is computing. diff --git a/policyengine_api/libs/simulation_entrypoint.py b/policyengine_api/libs/simulation_entrypoint.py index 997cb0352..c76e4a58e 100644 --- a/policyengine_api/libs/simulation_entrypoint.py +++ b/policyengine_api/libs/simulation_entrypoint.py @@ -7,6 +7,7 @@ import httpx from policyengine_observability import instrument_httpx +from pydantic import AwareDatetime, BaseModel, ConfigDict, Field, ValidationError from policyengine_api.gcp_logging import logger from policyengine_api.libs.gateway_auth import ( @@ -69,6 +70,45 @@ def _attach_current_request_id(request: httpx.Request) -> None: request.headers[OBSERVABILITY_ID_HEADER] = observability_id +class SimulationRunSegments(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + + completed: int = Field(ge=0) + total: int = Field(ge=1) + + +class SimulationRunProgress(BaseModel): + """Stage fields of a running Simulation Entrypoint job. + + The entrypoint adds these optional fields to ``GET /jobs/{job_id}`` for + jobs on workers that report stages (policyengine-sim-api#781). The API + relays them while an economy run is computing. ``stage`` stays a string: + later entrypoint versions may add stages, and clients must accept them. + """ + + model_config = ConfigDict(extra="ignore", frozen=True) + + stage: str | None = None + progress: int | None = Field(default=None, ge=0, le=100) + progress_max: int | None = Field(default=None, ge=0, le=100) + submitted_at: AwareDatetime | None = None + started_at: AwareDatetime | None = None + stage_started_at: AwareDatetime | None = None + segments: SimulationRunSegments | None = None + + @classmethod + def from_status(cls, data: dict) -> "SimulationRunProgress | None": + """Read the stage fields; never fail a poll because of them.""" + try: + run_progress = cls.model_validate(data) + except ValidationError: + return None + return run_progress if run_progress.to_dict() else None + + def to_dict(self) -> dict: + return self.model_dump(mode="json", exclude_none=True) + + @dataclass class ModalSimulationExecution: """ @@ -81,6 +121,7 @@ class ModalSimulationExecution: error: str | None = None policyengine_bundle: dict | None = None resolved_app_name: str | None = None + run_progress: SimulationRunProgress | None = None @property def name(self) -> str: @@ -441,6 +482,7 @@ def get_execution_by_id(self, job_id: str) -> ModalSimulationExecution: error=data.get("error"), policyengine_bundle=data.get("policyengine_bundle"), resolved_app_name=data.get("resolved_app_name"), + run_progress=SimulationRunProgress.from_status(data), ) except httpx.HTTPStatusError as e: diff --git a/policyengine_api/openapi_spec.yaml b/policyengine_api/openapi_spec.yaml index 3155fe341..7925f6935 100644 --- a/policyengine_api/openapi_spec.yaml +++ b/policyengine_api/openapi_spec.yaml @@ -879,7 +879,7 @@ paths: # Query parameters are generated from policyengine_api.query_parameters. responses: 200: - description: Annual economic impact, progress, or persisted error state. Canonical completed results include actual baseline/reform receipts containing the selected SPM configuration. + description: Annual economic impact, progress, or persisted error state. Canonical completed results include actual baseline/reform receipts containing the selected SPM configuration. While computing, the optional stage fields come from the simulation entrypoint and appear only when its worker reports stages. content: application/json: schema: @@ -897,6 +897,41 @@ paths: properties: spm_provenance: $ref: '#/components/schemas/SPMComparisonProvenance' + stage: + type: string + description: Current stage of a computing run, such as queued, loading_data, baseline, reform, simulating or outputs. Later versions may add stages; clients must accept a value they do not know. + progress: + type: integer + minimum: 0 + maximum: 100 + description: Coarse completion percent at the start of the current stage. It changes only at stage boundaries. + progress_max: + type: integer + minimum: 0 + maximum: 100 + description: Highest percent the run reaches before its next stage starts. A client can move a time-based estimate between progress and progress_max. + submitted_at: + type: string + format: date-time + description: When the simulation entrypoint accepted the run. + started_at: + type: string + format: date-time + description: When a worker started the run. + stage_started_at: + type: string + format: date-time + description: When the current stage started. + segments: + type: object + description: Completed region segments of a segmented US national run. + properties: + completed: + type: integer + minimum: 0 + total: + type: integer + minimum: 1 400: $ref: '#/components/responses/SPMValidationError' 404: diff --git a/policyengine_api/routes/economy_routes.py b/policyengine_api/routes/economy_routes.py index 60837fa7d..000cb35e9 100644 --- a/policyengine_api/routes/economy_routes.py +++ b/policyengine_api/routes/economy_routes.py @@ -103,13 +103,19 @@ def get_economic_impact(country_id: str, policy_id: int, baseline_policy_id: int result=None, ) - return _json_response( - { - "status": result_dict["status"], - "message": None, - "result": result_dict["data"], - } - ) + body = { + "status": result_dict["status"], + "message": None, + "result": result_dict["data"], + } + # Stage fields from the Simulation Entrypoint, only while computing and + # only when the entrypoint reports them, so other bodies keep their shape. + if ( + result_dict["status"] == "computing" + and economic_impact_result.run_progress is not None + ): + body.update(economic_impact_result.run_progress.to_dict()) + return _json_response(body) @economy_bp.route( diff --git a/policyengine_api/services/economy_service.py b/policyengine_api/services/economy_service.py index c6169f4cf..173a3a4bc 100644 --- a/policyengine_api/services/economy_service.py +++ b/policyengine_api/services/economy_service.py @@ -25,7 +25,10 @@ from policyengine_api.data.places import validate_place_code from policyengine_api.data.v1_models import ReformImpact from policyengine_api.gcp_logging import logger -from policyengine_api.libs.simulation_entrypoint import simulation_entrypoint +from policyengine_api.libs.simulation_entrypoint import ( + SimulationRunProgress, + simulation_entrypoint, +) from policyengine_api.observability import ( runtime as observability_runtime, ) @@ -131,6 +134,9 @@ class EconomicImpactResult(BaseModel): status: ImpactStatus data: dict | None = None message: str | None = None + # Stage fields relayed from the Simulation Entrypoint while computing. + # Kept out of to_dict(); the route adds them only when present. + run_progress: SimulationRunProgress | None = None model_config = {"frozen": True} # Make model immutable @@ -152,11 +158,13 @@ def completed(cls, data: dict) -> "EconomicImpactResult": return cls(status=ImpactStatus.OK, data=data) @classmethod - def computing(cls) -> "EconomicImpactResult": + def computing( + cls, run_progress: SimulationRunProgress | None = None + ) -> "EconomicImpactResult": """ Create an EconomicImpactResult for a computing impact calculation. """ - return cls(status=ImpactStatus.COMPUTING, data=None) + return cls(status=ImpactStatus.COMPUTING, data=None, run_progress=run_progress) @classmethod def error(cls, message: str) -> "EconomicImpactResult": @@ -1276,7 +1284,9 @@ def _handle_execution_state( {"message": "Sim API execution is still running"}, severity="INFO", ) - return EconomicImpactResult.computing() + return EconomicImpactResult.computing( + run_progress=execution.run_progress if execution is not None else None + ) raise ValueError(f"Unexpected sim API execution state: {execution_state}") diff --git a/tests/fixtures/services/economy_service.py b/tests/fixtures/services/economy_service.py index 4f3738d96..0ec3481ee 100644 --- a/tests/fixtures/services/economy_service.py +++ b/tests/fixtures/services/economy_service.py @@ -264,6 +264,7 @@ def create_mock_modal_execution( result=None, error=None, policyengine_bundle=None, + run_progress=None, ): """ Helper function to create mock Modal execution objects. @@ -292,6 +293,7 @@ def create_mock_modal_execution( mock_execution.error = error mock_execution.policyengine_bundle = policyengine_bundle or MOCK_RUNTIME_BUNDLE mock_execution.resolved_app_name = MOCK_RESOLVED_APP_NAME + mock_execution.run_progress = run_progress return mock_execution diff --git a/tests/unit/libs/test_simulation_entrypoint.py b/tests/unit/libs/test_simulation_entrypoint.py index 5616b8dae..bb0c452b2 100644 --- a/tests/unit/libs/test_simulation_entrypoint.py +++ b/tests/unit/libs/test_simulation_entrypoint.py @@ -26,6 +26,7 @@ ModalSimulationExecution, SimulationAPIModal, SimulationEntrypointClient, + SimulationRunProgress, resolve_simulation_entrypoint_url, ) from policyengine_api.observability.identifiers import ( # noqa: E402 @@ -960,6 +961,131 @@ def test__given_failed_job__then_returns_error( assert execution.status == MODAL_EXECUTION_STATUS_FAILED assert execution.error == "Simulation timed out" + def test__given_running_job_with_stage_fields__then_reads_run_progress( + self, + mock_httpx_client, + mock_modal_logger, + ): + # Given + mock_httpx_client.get.return_value = create_mock_httpx_response( + status_code=202, + json_data={ + **MOCK_POLL_RESPONSE_RUNNING, + "stage": "simulating", + "progress": 5, + "progress_max": 96, + "submitted_at": "2026-10-13T12:00:00Z", + "started_at": "2026-10-13T12:00:10Z", + "stage_started_at": "2026-10-13T12:00:17Z", + "segments": {"completed": 19, "total": 20}, + "stage_durations": {"loading_data": 7.0}, + }, + ) + api = SimulationAPIModal() + + # When + execution = api.get_execution_by_id(MOCK_MODAL_JOB_ID) + + # Then + assert execution.status == MODAL_EXECUTION_STATUS_RUNNING + assert execution.run_progress.to_dict() == { + "stage": "simulating", + "progress": 5, + "progress_max": 96, + "submitted_at": "2026-10-13T12:00:00Z", + "started_at": "2026-10-13T12:00:10Z", + "stage_started_at": "2026-10-13T12:00:17Z", + "segments": {"completed": 19, "total": 20}, + } + + def test__given_running_job_without_stage_fields__then_has_no_run_progress( + self, + mock_httpx_client, + mock_modal_logger, + ): + # Given + mock_httpx_client.get.return_value = create_mock_httpx_response( + status_code=202, + json_data=MOCK_POLL_RESPONSE_RUNNING, + ) + api = SimulationAPIModal() + + # When + execution = api.get_execution_by_id(MOCK_MODAL_JOB_ID) + + # Then + assert execution.run_progress is None + + def test__given_only_a_submission_time__then_relays_it( + self, + mock_httpx_client, + mock_modal_logger, + ): + # Given - a job on a worker app without stages, after + # policyengine-sim-api#781 deploys its gateway + mock_httpx_client.get.return_value = create_mock_httpx_response( + status_code=202, + json_data={ + **MOCK_POLL_RESPONSE_RUNNING, + "submitted_at": "2026-10-13T12:00:00Z", + }, + ) + api = SimulationAPIModal() + + # When + execution = api.get_execution_by_id(MOCK_MODAL_JOB_ID) + + # Then + assert execution.run_progress.to_dict() == { + "submitted_at": "2026-10-13T12:00:00Z" + } + + def test__given_a_malformed_timestamp__then_drops_the_stage_fields( + self, + mock_httpx_client, + mock_modal_logger, + ): + # Given + mock_httpx_client.get.return_value = create_mock_httpx_response( + status_code=202, + json_data={ + **MOCK_POLL_RESPONSE_RUNNING, + "stage": "reform", + "stage_started_at": "yesterday", + }, + ) + api = SimulationAPIModal() + + # When + execution = api.get_execution_by_id(MOCK_MODAL_JOB_ID) + + # Then + assert execution.status == MODAL_EXECUTION_STATUS_RUNNING + assert execution.run_progress is None + + def test__given_malformed_stage_fields__then_still_returns_status( + self, + mock_httpx_client, + mock_modal_logger, + ): + # Given + mock_httpx_client.get.return_value = create_mock_httpx_response( + status_code=202, + json_data={ + **MOCK_POLL_RESPONSE_RUNNING, + "stage": "reform", + "progress": 150, + }, + ) + api = SimulationAPIModal() + + # When + execution = api.get_execution_by_id(MOCK_MODAL_JOB_ID) + + # Then + assert execution.status == MODAL_EXECUTION_STATUS_RUNNING + assert execution.run_progress is None + def test__given_job_id__then_polls_correct_endpoint( self, mock_httpx_client, @@ -1007,6 +1133,17 @@ def test__given_network_error__then_raises_exception( log_payload = mock_modal_logger.log_struct.call_args.args[0] assert MOCK_MODAL_JOB_ID in log_payload["message"] + class TestSimulationRunProgress: + def test__given_unknown_stage__then_keeps_it(self): + run_progress = SimulationRunProgress.from_status( + {"status": "running", "stage": "a_future_stage", "progress": 40} + ) + + assert run_progress.to_dict() == { + "stage": "a_future_stage", + "progress": 40, + } + class TestGetBudgetWindowBatchById: def test__given_running_batch__then_returns_running_status( self, diff --git a/tests/unit/routes/test_economy_dataset_validation.py b/tests/unit/routes/test_economy_dataset_validation.py index 4054448dc..237b79377 100644 --- a/tests/unit/routes/test_economy_dataset_validation.py +++ b/tests/unit/routes/test_economy_dataset_validation.py @@ -3,6 +3,7 @@ from flask import Flask +from policyengine_api.libs.simulation_entrypoint import SimulationRunProgress from policyengine_api.routes.economy_routes import economy_bp @@ -92,12 +93,91 @@ def test_economy_route_keeps_computing_response_as_http_200(mock_get_economic_im "data": None, "message": None, } + mock_result.run_progress = None mock_get_economic_impact.return_value = mock_result client = _client_with_economy_blueprint() response = client.get("/us/economy/123/over/456?region=us&time_period=2026") assert response.status_code == 200 + assert json.loads(response.data) == { + "status": "computing", + "message": None, + "result": None, + } + + +@patch("policyengine_api.routes.economy_routes.economy_service.get_economic_impact") +def test_economy_route_relays_stage_fields_while_computing(mock_get_economic_impact): + mock_result = Mock() + mock_result.to_dict.return_value = { + "status": "computing", + "data": None, + "message": None, + } + mock_result.run_progress = SimulationRunProgress( + stage="reform", + progress=67, + progress_max=96, + stage_started_at="2026-10-13T12:01:00Z", + ) + mock_get_economic_impact.return_value = mock_result + client = _client_with_economy_blueprint() + + response = client.get("/us/economy/123/over/456?region=us&time_period=2026") + + assert response.status_code == 200 + assert json.loads(response.data) == { + "status": "computing", + "message": None, + "result": None, + "stage": "reform", + "progress": 67, + "progress_max": 96, + "stage_started_at": "2026-10-13T12:01:00Z", + } + + +@patch("policyengine_api.routes.economy_routes.economy_service.get_economic_impact") +def test_economy_route_relays_submission_time_for_workers_without_stages( + mock_get_economic_impact, +): + # After policyengine-sim-api#781 deploys, the entrypoint reports + # submitted_at for every new job, also on worker apps without stages. + mock_result = Mock() + mock_result.to_dict.return_value = { + "status": "computing", + "data": None, + "message": None, + } + mock_result.run_progress = SimulationRunProgress( + submitted_at="2026-10-13T12:00:00Z" + ) + mock_get_economic_impact.return_value = mock_result + client = _client_with_economy_blueprint() + + response = client.get("/us/economy/123/over/456?region=us&time_period=2026") + + assert json.loads(response.data) == { + "status": "computing", + "message": None, + "result": None, + "submitted_at": "2026-10-13T12:00:00Z", + } + + +@patch("policyengine_api.routes.economy_routes.economy_service.get_economic_impact") +def test_economy_route_never_adds_stage_fields_to_completed_results( + mock_get_economic_impact, +): + mock_result = _mock_economic_result() + mock_result.run_progress = SimulationRunProgress(stage="outputs", progress=97) + mock_get_economic_impact.return_value = mock_result + client = _client_with_economy_blueprint() + + response = client.get("/us/economy/123/over/456?region=us&time_period=2026") + + assert set(json.loads(response.data)) == {"status", "message", "result"} @patch( diff --git a/tests/unit/routes/test_economy_spm_query.py b/tests/unit/routes/test_economy_spm_query.py index 90a14ced2..a65f39e06 100644 --- a/tests/unit/routes/test_economy_spm_query.py +++ b/tests/unit/routes/test_economy_spm_query.py @@ -46,6 +46,7 @@ def calculate(**kwargs): setups.append(setup) return SimpleNamespace( cache_status=None, + run_progress=None, to_dict=lambda: { "status": "computing", "message": None, diff --git a/tests/unit/services/test_economy_service.py b/tests/unit/services/test_economy_service.py index d50262bd2..3110b3512 100644 --- a/tests/unit/services/test_economy_service.py +++ b/tests/unit/services/test_economy_service.py @@ -5,6 +5,10 @@ import httpx import pytest +from policyengine_api.libs.simulation_entrypoint import ( + ModalSimulationExecution, + SimulationRunProgress, +) from policyengine_api.runtime_cache.core import CacheCoordinationError from policyengine_api.runtime_cache.reform_impacts import ReformImpactStartClaim from policyengine_api.services.budget_window_cache import BudgetWindowCacheState @@ -2270,6 +2274,49 @@ def test__given_modal_running_state__then_returns_computing_result( assert result.status == ImpactStatus.COMPUTING assert result.data is None + def test__given_modal_running_state_with_run_progress__then_relays_it( + self, economy_service, setup_options, mock_logger + ): + # Given + reform_impact = create_mock_reform_impact(status="computing") + execution = ModalSimulationExecution( + job_id="fc-job-1", + status="running", + run_progress=SimulationRunProgress( + stage="reform", progress=67, progress_max=96 + ), + ) + + # When + result = economy_service._handle_execution_state( + setup_options, "running", reform_impact, execution + ) + + # Then + assert result.status == ImpactStatus.COMPUTING + assert result.run_progress.to_dict() == { + "stage": "reform", + "progress": 67, + "progress_max": 96, + } + assert "run_progress" not in result.to_dict() + + def test__given_modal_running_state_without_run_progress__then_has_none( + self, economy_service, setup_options, mock_logger + ): + # Given + reform_impact = create_mock_reform_impact(status="computing") + execution = ModalSimulationExecution(job_id="fc-job-1", status="running") + + # When + result = economy_service._handle_execution_state( + setup_options, "running", reform_impact, execution + ) + + # Then + assert result.status == ImpactStatus.COMPUTING + assert result.run_progress is None + def test__given_modal_submitted_state__then_returns_computing_result( self, economy_service, setup_options, mock_logger ):