Let a failed US state economy run be retried after two minutes

changelog.d/us-state-failure-retry.changed.md

@@ -0,0 +1 @@
+A failed US state economy job (`region=state/...`) now keeps its error for two minutes instead of 30 days, and the error message says when the state can be requested again; the next request after that submits a new job. Failed US national, congressional district and place jobs, UK national jobs and typed SPM failures are still kept for 30 days.

docs/canonical-spm.md

@@ -345,8 +345,14 @@ completes without `kind: "local_areas"`, also returns HTTP 502. Its result is
 not stored. Its error is kept for two minutes on that job's own record, so
 every request in that time gets the same error, which says when a new
 breakdown can be requested; then the record expires and the next request
-submits a new job. Other jobs' records are not changed. A failed national job
-is still stored and returned again for the cache lifetime of 30 days.
+submits a new job. Other jobs' records are not changed. A failed US state job
+(`country_id` `us` with region `state/...`, or a bare state code) keeps its
+error for two minutes in the same way, with a message that says when the
+state can be requested again, so a client can retry one state of a
+nationwide report's congressional districts. Any other failed job, such as
+a national, congressional district or place job, is still stored and
+returned again for the cache lifetime of 30 days, and so is a typed SPM
+failure.
 
 Before cache access or submission,
 the API validates the selected worker application's `canonical-spm-v1` capability

policyengine_api/services/economy_impact_execution.py

@@ -17,8 +17,9 @@
 from policyengine_api.runtime_cache.reform_impacts import CachedReformImpact
 from policyengine_api.services.economic_impact_request import LOCAL_AREA_BREAKDOWN
 from policyengine_api.services.economy_impact_types import (
-    LOCAL_AREA_BREAKDOWN_FAILURE_SECONDS,
     LOCAL_AREA_BREAKDOWN_RETRY_NOTE,
+    RETRYABLE_FAILURE_SECONDS,
+    US_STATE_RETRY_NOTE,
     EconomicImpactResult,
     EconomicImpactSetupOptions,
     ImpactPolicies,
@@ -51,8 +52,11 @@ def _handle_execution_state(
             raw_result = self._simulation_gateway.get_execution_result(execution)
             kind_error = self._breakdown_result_kind_error(setup_options, raw_result)
             if kind_error is not None:
-                return self._record_breakdown_failure(
-                    setup_options, reform_impact, kind_error
+                return self._record_retryable_failure(
+                    setup_options,
+                    reform_impact,
+                    kind_error,
+                    retry_note=LOCAL_AREA_BREAKDOWN_RETRY_NOTE,
                 )
             result = self._with_policyengine_bundle(
                 result=raw_result,
@@ -86,9 +90,13 @@ def _handle_execution_state(
                 error_message = (
                     f"Simulation entrypoint execution failed: {execution.error}"
                 )
-            if setup_options.options.get("breakdown") is not None:
-                return self._record_breakdown_failure(
-                    setup_options, reform_impact, error_message
+            retry_note = self._failure_retry_note(setup_options)
+            if retry_note is not None:
+                return self._record_retryable_failure(
+                    setup_options,
+                    reform_impact,
+                    error_message,
+                    retry_note=retry_note,
                 )
 
             self._set_reform_impact_error(
@@ -137,31 +145,51 @@ def _breakdown_result_kind_error(
             f"{LOCAL_AREA_BREAKDOWN!r}; the result was not stored"
         )
 
-    def _record_breakdown_failure(
+    @staticmethod
+    def _failure_retry_note(setup_options: EconomicImpactSetupOptions) -> str | None:
+        """Return the retry note for a failed job whose error is kept briefly.
+
+        A local-area breakdown job and a US state job (region ``state/...``,
+        already normalized from a bare state code) get a note; every other
+        failed job, including US national, congressional district and place
+        jobs, gets None, and its error is kept for the cache lifetime.
+        """
+        if setup_options.options.get("breakdown") is not None:
+            return LOCAL_AREA_BREAKDOWN_RETRY_NOTE
+        if setup_options.country_id == "us" and setup_options.region.startswith(
+            "state/"
+        ):
+            return US_STATE_RETRY_NOTE
+        return None
+
+    def _record_retryable_failure(
         self,
         setup_options: EconomicImpactSetupOptions,
         reform_impact: CachedReformImpact,
         reason: str,
+        *,
+        retry_note: str,
     ) -> EconomicImpactResult:
-        """Keep a failed breakdown job's error for two minutes, then let it go.
-
-        A breakdown runs on the large local-area dataset, where a timeout or
-        memory failure may not repeat, so its failure is not kept for the usual
-        30 days. This job's own record, found by its execution ID, becomes an
-        error that expires after ``LOCAL_AREA_BREAKDOWN_FAILURE_SECONDS``.
-        Until then every request gets the same error, which says when to try
-        again; after it, the next request finds no record and submits a new
-        job. No other job's record is read or changed, and no submission claim
-        is held at this point: it was released once the job was handed off.
-        Failures of requests without a breakdown are stored as before.
+        """Keep a failed job's error for two minutes, then let it go.
+
+        Used for local-area breakdown and US state jobs, which run on large
+        local-area datasets, where a timeout or memory failure may not repeat,
+        so their failures are not kept for the usual 30 days. This job's own
+        record, found by its execution ID, becomes an error that expires after
+        ``RETRYABLE_FAILURE_SECONDS``. Until then every request gets the same
+        error, which ends with ``retry_note`` to say when to try again; after
+        it, the next request finds no record and submits a new job. No other
+        job's record is read or changed, and no submission claim is held at
+        this point: it was released once the job was handed off. Other failed
+        jobs, and typed SPM failures, are stored as before.
         """
-        message = f"{reason}. {LOCAL_AREA_BREAKDOWN_RETRY_NOTE}"
+        message = f"{reason}. {retry_note}"
         if reform_impact.execution_id:
             self._set_reform_impact_error(
                 setup_options=setup_options,
                 message=message,
                 execution_id=reform_impact.execution_id,
-                keep_seconds=LOCAL_AREA_BREAKDOWN_FAILURE_SECONDS,
+                keep_seconds=RETRYABLE_FAILURE_SECONDS,
             )
         return EconomicImpactResult.error(message=message)
 

policyengine_api/services/economy_impact_types.py

@@ -40,14 +40,22 @@ class ImpactStatus(Enum):
 # 400 from a gateway whose routed .py release does not declare breakdowns, and
 # 422 from an entry deployment that does not know the field.
 LOCAL_AREA_BREAKDOWN_REFUSAL_STATUS_CODES = frozenset({400, 422})
-# How long a failed breakdown's error is returned before the next request may
-# submit a new job. Short, so the app's Retry works soon, but long enough that
-# a client retrying on every 502 does not start a run on the large local-area
-# dataset each time.
-LOCAL_AREA_BREAKDOWN_FAILURE_SECONDS = 120
+# How long the error of a failed local-area breakdown or US state job is
+# returned before the next request may submit a new job. Both run on large
+# local-area datasets, where a timeout or memory failure may not repeat. The
+# app offers Retry for a UK breakdown, and will offer it for each state it runs
+# for a US nationwide report's congressional districts. Short, so Retry works
+# soon, but long enough that a client retrying on every 502 does not start a
+# run on that dataset each time. Every other failed job's error is kept for the
+# cache lifetime of 30 days.
+RETRYABLE_FAILURE_SECONDS = 120
 LOCAL_AREA_BREAKDOWN_RETRY_NOTE = (
     "A new local-area breakdown can be requested "
-    f"{LOCAL_AREA_BREAKDOWN_FAILURE_SECONDS // 60} minutes after this failure."
+    f"{RETRYABLE_FAILURE_SECONDS // 60} minutes after this failure."
+)
+US_STATE_RETRY_NOTE = (
+    "A new run for this state can be requested "
+    f"{RETRYABLE_FAILURE_SECONDS // 60} minutes after this failure."
 )
 # The telemetry simulation_kind of a breakdown run. The simulation API's
 # telemetry envelope forbids extra fields, so the run is marked here rather

policyengine_api/services/economy_service.py

@@ -49,11 +49,12 @@
     BUDGET_WINDOW_SUBMISSION_VALIDATION_ERROR_STATUS_CODES,
     COMPLETE_STATUSES,
     COMPUTING_STATUS,
-    LOCAL_AREA_BREAKDOWN_FAILURE_SECONDS,
     LOCAL_AREA_BREAKDOWN_REFUSAL_STATUS_CODES,
     LOCAL_AREA_BREAKDOWN_RETRY_NOTE,
     LOCAL_AREA_BREAKDOWN_SIMULATION_KIND,
     REFORM_IMPACT_START_CLAIM_ATTEMPTS,
+    RETRYABLE_FAILURE_SECONDS,
+    US_STATE_RETRY_NOTE,
     BudgetWindowEconomicImpactResult,
     BudgetWindowYearProgress,
     EconomicImpactResult,
@@ -85,11 +86,12 @@
     "BUDGET_WINDOW_SUBMISSION_VALIDATION_ERROR_STATUS_CODES",
     "COMPLETE_STATUSES",
     "COMPUTING_STATUS",
-    "LOCAL_AREA_BREAKDOWN_FAILURE_SECONDS",
     "LOCAL_AREA_BREAKDOWN_REFUSAL_STATUS_CODES",
     "LOCAL_AREA_BREAKDOWN_RETRY_NOTE",
     "LOCAL_AREA_BREAKDOWN_SIMULATION_KIND",
     "REFORM_IMPACT_START_CLAIM_ATTEMPTS",
+    "RETRYABLE_FAILURE_SECONDS",
+    "US_STATE_RETRY_NOTE",
     "BudgetWindowEconomicImpactResult",
     "BudgetWindowYearProgress",
     "EconomicImpactResult",

tests/unit/services/test_economy_service_impact_execution.py

@@ -10,6 +10,7 @@
     SimulationRunProgress,
 )
 from policyengine_api.services.economy_service import (
+    EconomicImpactResult,
     EconomicImpactSetupOptions,
     EconomyService,
     ImpactStatus,
@@ -19,6 +20,7 @@
     MOCK_BASELINE_POLICY_ID,
     MOCK_COUNTRY_ID,
     MOCK_DATA_VERSION,
+    MOCK_EXECUTION_ID,
     MOCK_MODEL_VERSION,
     MOCK_OBSERVABILITY_ID,
     MOCK_OPTIONS,
@@ -238,6 +240,79 @@ def test__given_modal_failure_with_error_message__includes_error_in_message(
             call_args = mock_reform_impacts_service.set_error_reform_impact.call_args
             assert "Simulation timed out" in call_args[1]["message"]
 
+        def test__given_failed_us_state_job__keeps_its_error_for_two_minutes(
+            self,
+            economy_service,
+            setup_options,
+            mock_reform_impacts_service,
+            mock_logger,
+        ):
+            # Given
+            state_setup_options = setup_options.model_copy(
+                update={"region": "state/ca"}
+            )
+            reform_impact = create_mock_reform_impact(status="computing")
+            execution = ModalSimulationExecution(
+                job_id=MOCK_EXECUTION_ID, status="failed", error="out of memory"
+            )
+
+            # When
+            result = economy_service._handle_execution_state(
+                state_setup_options, "failed", reform_impact, execution
+            )
+
+            # Then
+            expected_message = (
+                "Simulation entrypoint execution failed: out of memory. "
+                "A new run for this state can be requested 2 minutes after this "
+                "failure."
+            )
+            # The stored error expires, so the log keeps the failure's reason.
+            # Checked first: building the expected result below logs too.
+            mock_logger.log_struct.assert_any_call(
+                {"message": expected_message}, severity="ERROR"
+            )
+            assert result == EconomicImpactResult.error(expected_message)
+            mock_reform_impacts_service.set_error_reform_impact.assert_called_once_with(
+                execution_id=MOCK_EXECUTION_ID,
+                message=expected_message,
+                keep_seconds=120,
+            )
+
+        @pytest.mark.usefixtures("mock_logger")
+        @pytest.mark.parametrize(
+            "region",
+            [MOCK_REGION, "congressional_district/CA-37", "place/NJ-57000"],
+            ids=["national", "district", "place"],
+        )
+        def test__given_failed_us_job_outside_a_state__keeps_its_error_for_30_days(
+            self,
+            economy_service,
+            setup_options,
+            mock_reform_impacts_service,
+            region,
+        ):
+            # Given
+            region_setup_options = setup_options.model_copy(update={"region": region})
+            reform_impact = create_mock_reform_impact(status="computing")
+            execution = ModalSimulationExecution(
+                job_id=MOCK_EXECUTION_ID, status="failed", error="out of memory"
+            )
+
+            # When
+            result = economy_service._handle_execution_state(
+                region_setup_options, "failed", reform_impact, execution
+            )
+
+            # Then
+            expected_message = "Simulation entrypoint execution failed: out of memory"
+            assert result == EconomicImpactResult.error(expected_message)
+            # Without keep_seconds the error has the cache's 30-day lifetime.
+            mock_reform_impacts_service.set_error_reform_impact.assert_called_once_with(
+                execution_id=MOCK_EXECUTION_ID,
+                message=expected_message,
+            )
+
         @pytest.mark.usefixtures("mock_logger")
         def test__given_modal_running_state__then_returns_computing_result(
             self, economy_service, setup_options