Split the simulation runtime and its output builder tests under the file-length limit

changelog_entry.yaml

@@ -16,6 +16,7 @@
       - Ran Stage 12 US national simulations as 20 region-group calls each, and built an independent, immutable Stage 12 cache of three annual US datasets and 60 current-law baselines before each deployment, publishing it only after a recorded determinism check and recording its release in the v2 version manifest under a new key, while the existing key keeps a copy that earlier readers accept
       - Made the budget refuse a household count outside its country's plausible range (above 200 million in the US or 60 million in the UK, negative or not finite), raising ImplausibleEntityCountError instead of publishing a count like #767's
     changed:
+      - Moved the worker's credential setup, request year and policy parsing, and region resolution out of simulation_runtime.py (1,111 lines, now 651) into gcp_credentials.py, request_policy.py and region_resolution.py; simulation_runtime.py keeps dataset selection and loading, simulation building, local-area breakdowns and the run itself, which tests replace by name there, and re-exports the moved names, so every caller still imports them from simulation_runtime; the shared helpers _parse_year, _normalise_policy, _compile_interval_policy, _runtime_policy (now wrapper_policy), _resolve_region, _resolve_dataset_reference and _require_local_area_breakdown are now public, and the credential messages still log under the simulation_runtime logger; the 3,048-line output builder test file is split into seven test files and a module of shared fakes, all named test_simulation_output_*.py so the weighted-frame rule still checks them
       - Split the artifact precompute (precompute.py, 931 lines) and the Stage 12 cache precompute (stage12_cache/precompute.py, 844 lines) each into planning, dataset-building and baseline modules, with the manifest and run_precompute left in the original module, which re-exports every public name, so the Modal precompute apps and their tests import from the same paths; the cache precompute's _scoped_planned_simulation and _baseline_identity are now public, and the worker validation imports scoped_planned_simulation from its new module; the 1,835-line precompute test file is split into four, and the cache precompute tests patch each name in the module whose code reads it
       - Split the 4,709-line Stage 12 runtime test file into thirteen test files by behaviour (dataset loading, provenance, the single worker and its failures, aggregates, the coordinator and its failures, runtime store retries, failure records, cancellation, child state errors and the wait window) and three support modules for the shared builders, fake stores and fake child calls, whose helpers are now public; the other Stage 12 tests import them from those modules, every test keeps its name and body, and two parametrized cases are named after the renamed bundle helper
       - Split Stage 12 report coordination (stage12_runtime/coordination.py, 1,091 lines) into modules for the report coordinator class, its child simulations' records, failures and cancellation, and the parent report's checks and claim, with coordinate_report and every public name still in coordination.py and child failures still logged under its logger; and turned the contract's stage12_execution.py (929 lines) into a stage12_execution package of modules for the shared field types and simulation input, report inputs and output plans, result artifacts, and comparison-run records, which re-exports every name from the same import path, with unchanged JSON schemas; models subclass the base models through the package path ruff.toml lists, so no ruff setting changes

docs/generated/rules.md

@@ -66,7 +66,7 @@ Every rule that governs `policyengine-sim-api` is listed here with how it is enf
 
 - `sim.ci.multiline-run-is-script`: 13 remaining in `rules/baselines/ci-multiline-run.txt`
 - `sim.microdata.weight-stripping-confined`: 0 remaining in `rules/baselines/sim-microdata-weight-stripping.txt`
-- `sim.quality.python-file-length`: 4 remaining in `rules/baselines/python-file-length.txt`
+- `sim.quality.python-file-length`: 2 remaining in `rules/baselines/python-file-length.txt`
 
 ### Exceptions and review dates
 

projects/policyengine-simulation-executor/src/modal/segmented_national.py

@@ -581,7 +581,7 @@ def _reduce(
         self, child_results: list[dict], *, country_module: ModuleType
     ) -> dict[str, Any]:
         from policyengine_simulation_executor.simulation_runtime import (  # noqa: PLC0415  # Modal worker code: bound per call, like this module's other simulation_runtime imports
-            _parse_year,
+            parse_year,
         )
 
         # The canonical internal-key stripper, plus the opt-out knob.
@@ -597,7 +597,7 @@ def _reduce(
                     country=self.country,
                     simulation_params=simulation_params,
                     country_module=country_module,
-                    year=_parse_year(simulation_params),
+                    year=parse_year(simulation_params),
                     resolved_data_version=None,
                 ),
                 runtime=self.runtime,
@@ -606,7 +606,7 @@ def _reduce(
             combine_spm_results(
                 child_results,
                 simulation_params.get("spm"),
-                expected_year=_parse_year(simulation_params),
+                expected_year=parse_year(simulation_params),
             )
         )
         for key in ("model_version", "data_version"):

projects/policyengine-simulation-executor/src/policyengine_simulation_executor/gcp_credentials.py

@@ -0,0 +1,119 @@
+"""Write the worker's Google Cloud credentials from its environment."""
+
+from __future__ import annotations
+
+import contextlib
+import json
+import logging
+import os
+import tempfile
+from typing import TYPE_CHECKING
+
+from policyengine_simulation_observability import scopes
+from policyengine_simulation_observability.stages import ANNUAL_IMPACT_STAGES, Stage
+
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+
+    from policyengine_observability import ObservabilityRuntime
+
+# Moved out of simulation_runtime.py; it still logs under that module's name.
+logger = logging.getLogger("policyengine_simulation_executor.simulation_runtime")
+
+# Credential-adjacent environment values are cut to this many characters.
+_ENV_VALUE_PREVIEW_CHARS = 50
+
+
+def _normalize_credentials_blob(creds_json: str) -> str:
+    """Return the raw JSON blob, decoding the outer escape if present.
+
+    The upstream Modal secret sometimes stores the credentials payload
+    double-encoded (the entire JSON object is wrapped in quotes with
+    backslash-escaped interior quotes). Historically we always attempted
+    the unescape as a fallback which could accidentally parse an already
+    clean blob. Only unwrap when the payload looks wrapped.
+    """
+    try:
+        json.loads(creds_json)
+    except json.JSONDecodeError:
+        looks_escaped = creds_json.lstrip().startswith('"') or '\\"' in creds_json
+        if looks_escaped:
+            return json.loads(f'"{creds_json}"')
+        raise
+    return creds_json
+
+
+@contextlib.contextmanager
+def setup_gcp_credentials(
+    runtime: ObservabilityRuntime | None = None,
+) -> Iterator[None]:
+    """Set up GCP credentials from environment variable.
+
+    Modal secrets are injected as environment variables. The GCP library
+    expects GOOGLE_APPLICATION_CREDENTIALS to point to a file path. If
+    credentials JSON is provided, write it to a temp file that's deleted
+    on exit. This runs as a context manager to guarantee cleanup even if
+    the caller raises mid-simulation; the previous fire-and-forget
+    ``tempfile.mkstemp`` path leaked credential material on disk every
+    time a container served a request.
+    """
+    previous = os.environ.get("GOOGLE_APPLICATION_CREDENTIALS")
+    creds_file = None
+    with contextlib.ExitStack() as credentials_file_stack:
+        try:
+            credential_span = (
+                scopes.span(runtime, ANNUAL_IMPACT_STAGES.name(Stage.CREDENTIAL_SETUP))
+                if runtime is not None
+                else contextlib.nullcontext()
+            )
+            with credential_span:
+                # Log available GCP-related env vars for debugging.
+                gcp_vars = {
+                    k: v[:_ENV_VALUE_PREVIEW_CHARS] + "..."
+                    if len(v) > _ENV_VALUE_PREVIEW_CHARS
+                    else v
+                    for k, v in os.environ.items()
+                    if "GOOGLE" in k or "GCP" in k or "CREDENTIAL" in k
+                }
+                logger.info("GCP-related env vars: %s", list(gcp_vars.keys()))
+
+                if previous:
+                    logger.info("GOOGLE_APPLICATION_CREDENTIALS already set")
+                else:
+                    creds_json = (
+                        os.environ.get("GOOGLE_APPLICATION_CREDENTIALS_JSON")
+                        or os.environ.get("GCP_CREDENTIALS_JSON")
+                        or os.environ.get("GOOGLE_CREDENTIALS")
+                        or os.environ.get("SERVICE_ACCOUNT_JSON")
+                    )
+
+                    if not creds_json:
+                        logger.warning(
+                            "No GCP credentials found in environment variables"
+                        )
+                    else:
+                        normalized = _normalize_credentials_blob(creds_json)
+                        # The exit stack keeps the file open across the yield,
+                        # and closing it (``delete=True``) removes the file.
+                        # The finally block first restores any prior value of
+                        # ``GOOGLE_APPLICATION_CREDENTIALS`` so a retry in the same
+                        # container doesn't silently pick up a stale path.
+                        creds_file = credentials_file_stack.enter_context(
+                            tempfile.NamedTemporaryFile(
+                                mode="w",
+                                suffix=".json",
+                                delete=True,
+                            )
+                        )
+                        creds_file.write(normalized)
+                        creds_file.flush()
+                        os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = creds_file.name
+                        logger.info("GCP credentials written to %s", creds_file.name)
+
+            yield
+        finally:
+            if creds_file is not None:
+                if previous is None:
+                    os.environ.pop("GOOGLE_APPLICATION_CREDENTIALS", None)
+                else:
+                    os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = previous

projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py

@@ -21,7 +21,7 @@
 means deleting its object first, then re-running.
 
 Cohort baselines are built through the executor's own request path
-(``_resolve_region`` + ``_build_simulation`` on a synthetic child request),
+(``resolve_region`` + ``_build_simulation`` on a synthetic child request),
 so scoping, extra variables, and the deterministic simulation id match the
 production child's by construction, not by convention — and a regression
 test asserts the writer and reader derive identical identities. All

projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute_planning.py

@@ -53,12 +53,12 @@ def cohort_resolution(
     """
     from policyengine_simulation_executor.simulation_runtime import (  # noqa: PLC0415  # resolved per call so tests can patch simulation_runtime helpers
         _country_module,
-        _resolve_region,
+        resolve_region,
     )
 
     params = cohort_params(year, group)
     country_module = _country_module(PRECOMPUTE_COUNTRY)
-    resolution = _resolve_region(
+    resolution = resolve_region(
         country_module=country_module,
         country=PRECOMPUTE_COUNTRY,
         params=params,
@@ -123,14 +123,14 @@ class StateDatasetTarget:
 def state_dataset_targets() -> list[StateDatasetTarget]:
     """Return every US state region that loads a non-default dataset, by state code.
 
-    Each state is resolved exactly as a request for it is (``_resolve_region``
+    Each state is resolved exactly as a request for it is (``resolve_region``
     then ``_resolve_dataset_selection``), so a bundle that sends states back to
     the default dataset gets no state files.
     """
     from policyengine_simulation_executor.simulation_runtime import (  # noqa: PLC0415  # resolved per call so tests can patch simulation_runtime helpers
         _country_module,
         _resolve_dataset_selection,
-        _resolve_region,
+        resolve_region,
     )
 
     country_module = _country_module(PRECOMPUTE_COUNTRY)
@@ -141,7 +141,7 @@ def state_dataset_targets() -> list[StateDatasetTarget]:
             "region": region.code,
             "time_period": PRECOMPUTE_YEARS[0],
         }
-        resolution = _resolve_region(
+        resolution = resolve_region(
             country_module=country_module, country=PRECOMPUTE_COUNTRY, params=params
         )
         selection = _resolve_dataset_selection(params, region_resolution=resolution)