From 89c5b193fd778498d709975c290a430faa6aed38 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 13 Oct 2026 16:56:47 +0400 Subject: [PATCH 01/18] Document the simulation contract library for ruff D Validators, state helpers and URI functions get docstrings. Every Pydantic model and enum whose class docstring would become a JSON-schema description carries a D101 noqa instead, so the published contract schemas and the OpenAPI documents built from them stay byte-identical. Co-Authored-By: Claude Opus 5.5 --- .../budget_window_state.py | 39 ++++++ .../dataset_uri.py | 13 +- .../gateway_models.py | 25 +++- .../hf_dataset.py | 8 +- .../macro_output.py | 48 ++++---- .../policyengine_simulation_contract/spm.py | 36 ++++-- .../stage12_bundle.py | 8 +- .../stage12_execution.py | 94 +++++++++++---- .../stage12_manifest.py | 64 ++++++++-- .../tests/test_budget_window_state.py | 10 +- .../tests/test_gateway_models.py | 113 +++++------------- .../tests/test_spm_selection.py | 7 +- 12 files changed, 301 insertions(+), 164 deletions(-) diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/budget_window_state.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/budget_window_state.py index 1e396ca0b..5d0e0b289 100644 --- a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/budget_window_state.py +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/budget_window_state.py @@ -58,6 +58,11 @@ def create_initial_batch_state( bundle: PolicyEngineBundle, observability_id: str | None = None, ) -> BudgetWindowBatchState: + """Build the ``submitted`` state for a new batch, with every year queued. + + The stored request payload omits ``telemetry``, and all timestamps are set + to the current UTC time. + """ years = _build_years(request.start_year, request.window_size) now = _utc_now_iso() @@ -90,6 +95,11 @@ def create_initial_batch_state( def get_batch_job_seed(batch_job_id: str) -> BudgetWindowBatchState | None: + """Load the submit-time seed state from its Modal Dict, or None if absent. + + The seed is written when the batch is submitted, before the batch runner + writes live progress, so pollers fall back to it while no live state exists. + """ payload = _budget_window_job_seed_store().get(batch_job_id) if payload is None: return None @@ -97,10 +107,12 @@ def get_batch_job_seed(batch_job_id: str) -> BudgetWindowBatchState | None: def put_batch_job_seed(state: BudgetWindowBatchState) -> None: + """Store ``state`` as the batch's seed in the seed Modal Dict.""" _budget_window_job_seed_store()[state.batch_job_id] = state.model_dump(mode="json") def get_batch_job_state(batch_job_id: str) -> BudgetWindowBatchState | None: + """Load the live state the batch runner last wrote, or None if none yet.""" payload = _budget_window_job_store().get(batch_job_id) if payload is None: return None @@ -108,11 +120,13 @@ def get_batch_job_state(batch_job_id: str) -> BudgetWindowBatchState | None: def put_batch_job_state(state: BudgetWindowBatchState) -> None: + """Store ``state`` as the batch's live state in the job Modal Dict.""" serialized = state.model_dump(mode="json") _budget_window_job_store()[state.batch_job_id] = serialized def mark_batch_running(state: BudgetWindowBatchState) -> BudgetWindowBatchState: + """Set the batch status to ``running`` in place and return the same state.""" state.status = "running" return _touch(state) @@ -123,6 +137,10 @@ def mark_child_started( year: str, child_job_id: str, ) -> BudgetWindowBatchState: + """Move ``year`` from queued to running and record its child job id. + + Mutates ``state`` in place, is safe to repeat, and returns the same object. + """ if year in state.queued_years: state.queued_years.remove(year) if year not in state.running_years: @@ -169,6 +187,11 @@ def mark_child_completed( year: str, annual_impact: BudgetWindowAnnualImpact, ) -> BudgetWindowBatchState: + """Move ``year`` from running to completed and keep its annual impact. + + Mutates ``state`` in place and returns it. A year with no tracked child job + gets a sentinel job id rather than raising. + """ if year in state.running_years: state.running_years.remove(year) if year not in state.completed_years: @@ -189,6 +212,11 @@ def mark_child_failed( year: str, error: str, ) -> BudgetWindowBatchState: + """Move ``year`` from running to failed and record ``error`` on its child. + + Mutates ``state`` in place and returns it; the batch itself keeps its + status. A year with no tracked child job gets a sentinel job id. + """ if year in state.running_years: state.running_years.remove(year) if year not in state.failed_years: @@ -208,6 +236,7 @@ def mark_batch_complete( *, result: BudgetWindowResult, ) -> BudgetWindowBatchState: + """Mark the batch complete with ``result``, clearing any error and pending years.""" state.status = "complete" state.result = result state.error = None @@ -221,6 +250,11 @@ def mark_batch_failed( *, error: str, ) -> BudgetWindowBatchState: + """Mark the batch failed and cancel every child that is still running. + + Cancelled children carry the batch ``error``; queued years are left as they + are. Mutates ``state`` in place and returns it. + """ for year in list(state.running_years): child = state.child_jobs.get(year) if child is None: @@ -240,6 +274,11 @@ def mark_batch_failed( def build_batch_status_response( state: BudgetWindowBatchState, ) -> BudgetWindowBatchStatusResponse: + """Project stored batch state onto the public polling response. + + ``progress`` is the whole-number percentage of years completed (0 for an + empty window); failed years do not count towards it. + """ total_years = len(state.years) progress = ( 0 if total_years == 0 else round(len(state.completed_years) / total_years * 100) diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/dataset_uri.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/dataset_uri.py index 40f1714af..c731db0c0 100644 --- a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/dataset_uri.py +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/dataset_uri.py @@ -11,6 +11,11 @@ def split_dataset_revision(dataset_uri: str) -> tuple[str, str | None]: + """Split a trailing ``@revision`` off a dataset URI. + + The revision is None when the URI has no ``@``. Raises ValueError when + either side of the last ``@`` is empty. + """ if "@" not in dataset_uri: return dataset_uri, None dataset_without_revision, revision = dataset_uri.rsplit("@", maxsplit=1) @@ -24,6 +29,11 @@ def select_dataset_revision( requested_revision: str | None, requested_data_version: str | None, ) -> str | None: + """Pick the dataset revision from the URI pin or ``data_version``. + + Either may be given alone; when both are given they must match, otherwise + ValueError is raised. Returns None when neither is set. + """ if ( requested_revision is not None and requested_data_version is not None @@ -72,7 +82,6 @@ def _policyengine_data_runtime_uri( artifact_revision: str | None, ) -> str | None: """Resolve a PolicyEngine data-repository artifact through its GCS mirror.""" - bucket = _policyengine_gcs_bucket_for_hf_repo(reference.repo_id) if bucket is None: return None @@ -103,7 +112,6 @@ def _non_policyengine_data_runtime_uri( validate_hf: bool, ) -> str: """Preserve the existing Hugging Face resolution behavior for other repos.""" - if override_revision is not None: if validate_hf: return with_hf_revision(dataset_uri, override_revision) @@ -133,7 +141,6 @@ def runtime_dataset_uri( can disable live HF validation to avoid requiring private Hub credentials in request-routing code. """ - if dataset_uri.startswith("gs://"): return _with_gs_revision(dataset_uri, override_revision or default_revision) 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 83e9a6ba8..0be6d2c90 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 @@ -1,6 +1,4 @@ -""" -Pydantic models for the Gateway API. -""" +"""Pydantic models for the Gateway API.""" import json from typing import ClassVar, Literal @@ -59,7 +57,6 @@ def _strip_internal_passthrough_fields(value: object) -> object: ``extra="forbid"`` keeps catching unknown caller fields without breaking internal round-trips. """ - if not isinstance(value, dict): return value if not INTERNAL_PASSTHROUGH_FIELDS.intersection(value): @@ -127,16 +124,25 @@ class GatewayRequestBase(BaseModel): @model_validator(mode="before") @classmethod def move_internal_telemetry_alias(cls, value: object) -> object: + """Accept the internal ``_telemetry`` key as ``telemetry`` before validation. + + ``telemetry`` wins when a payload carries both keys. + """ return _move_internal_telemetry_alias(value) @model_validator(mode="before") @classmethod def strip_internal_passthrough_fields(cls, value: object) -> object: + """Drop ``_metadata`` and ``_runtime_bundle`` before ``extra="forbid"`` runs.""" return _strip_internal_passthrough_fields(value) @model_validator(mode="before") @classmethod def enforce_max_payload_size(cls, value: object) -> object: + """Reject payloads whose JSON exceeds ``MAX_GATEWAY_REQUEST_BYTES``. + + Runs before field parsing so oversized bodies never reach Pydantic. + """ return _enforce_max_payload_size(value) @@ -147,7 +153,7 @@ class SimulationRequest(GatewayRequestBase): @model_validator(mode="after") def region_xor_region_group(self) -> "SimulationRequest": - """A request scopes to a single region OR a region group, never both.""" + """Reject a request that sets both ``region`` and ``region_group``.""" if self.region is not None and self.region_group is not None: raise ValueError("Provide `region` or `region_group`, not both.") return self @@ -222,6 +228,10 @@ class BudgetWindowBatchRequest(GatewayRequestBase): @field_validator("start_year") @classmethod def validate_start_year(cls, value: str) -> str: + """Normalise ``start_year`` to a canonical integer string. + + For example ``"02026"`` becomes ``"2026"``; non-integers are rejected. + """ try: return str(int(value)) except (TypeError, ValueError) as exc: @@ -229,6 +239,11 @@ def validate_start_year(cls, value: str) -> str: @model_validator(mode="after") def validate_end_year(self) -> "BudgetWindowBatchRequest": + """Reject windows past ``MAX_END_YEAR``, region groups and cliff requests. + + Budget windows run against a single ``region`` and do not compute cliff + impacts; those requests belong on the single-year endpoint. + """ if self.region_group is not None: raise ValueError("Provide `region` or `region_group`, not both.") end_year = int(self.start_year) + self.window_size - 1 diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/hf_dataset.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/hf_dataset.py index c4187d68e..bf4de1b84 100644 --- a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/hf_dataset.py +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/hf_dataset.py @@ -28,6 +28,11 @@ class HuggingFaceDatasetReferenceError(ValueError): @dataclass(frozen=True) class HFDatasetReference: + """An ``hf://`` dataset URI split into Hub repo, in-repo path and revision. + + ``revision`` is None when the URI carries no ``@revision`` pin. + """ + repo_id: str path: str revision: str | None @@ -44,7 +49,6 @@ def parse_hf_dataset_uri(dataset_uri: str) -> HFDatasetReference | None: Hub API needs ``org/repo`` and ``revision`` separately, while path validation needs the artifact path within the repo. """ - if not dataset_uri.startswith("hf://"): return None @@ -116,7 +120,6 @@ def _siblings_contain_path(payload: dict[str, Any], path: str) -> bool | None: def validate_hf_dataset_uri(dataset_uri: str) -> str: """Validate an explicit ``hf://`` dataset URI if it pins a revision.""" - parsed = parse_hf_dataset_uri(dataset_uri) if parsed is None or parsed.revision is None: return dataset_uri @@ -133,7 +136,6 @@ def validate_hf_dataset_uri(dataset_uri: str) -> str: def with_hf_revision(dataset_uri: str, revision: str) -> str: """Return ``dataset_uri`` pinned to ``revision`` and validate it on the Hub.""" - if not dataset_uri.startswith("hf://"): return dataset_uri without_revision = dataset_uri.rsplit("@", maxsplit=1)[0] diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/macro_output.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/macro_output.py index a46ca9203..32fd1f886 100644 --- a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/macro_output.py +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/macro_output.py @@ -21,7 +21,7 @@ class MacroRootModel(RootModel[T], Generic[T]): """Base model for named map and list output schemas.""" -class BudgetaryImpact(MacroOutputModel): +class BudgetaryImpact(MacroOutputModel): # noqa: D101 # docstring would change the published OpenAPI description tax_revenue_impact: float state_tax_revenue_impact: float benefit_spending_impact: float @@ -33,7 +33,7 @@ class BudgetaryImpact(MacroOutputModel): BudgetaryOutput = BudgetaryImpact -class DetailedBudgetProgramOutput(MacroOutputModel): +class DetailedBudgetProgramOutput(MacroOutputModel): # noqa: D101 # docstring would change the published OpenAPI description baseline: float reform: float difference: float @@ -43,82 +43,82 @@ class DetailedBudgetOutput(MacroRootModel[dict[str, DetailedBudgetProgramOutput] """Program name mapped to its baseline and reform totals.""" -class DecileOutput(MacroOutputModel): +class DecileOutput(MacroOutputModel): # noqa: D101 # docstring would change the published OpenAPI description average: dict[str, float] relative: dict[str, float] -class IntraDecileOutput(MacroOutputModel): +class IntraDecileOutput(MacroOutputModel): # noqa: D101 # docstring would change the published OpenAPI description deciles: dict[str, list[float]] all: dict[str, float] -class BaselineReformValue(MacroOutputModel): +class BaselineReformValue(MacroOutputModel): # noqa: D101 # docstring would change the published OpenAPI description baseline: float reform: float -class AgePovertyOutput(MacroOutputModel): +class AgePovertyOutput(MacroOutputModel): # noqa: D101 # docstring would change the published OpenAPI description child: BaselineReformValue adult: BaselineReformValue senior: BaselineReformValue all: BaselineReformValue -class GenderPovertyOutput(MacroOutputModel): +class GenderPovertyOutput(MacroOutputModel): # noqa: D101 # docstring would change the published OpenAPI description male: BaselineReformValue female: BaselineReformValue -class RacePovertyOutput(MacroOutputModel): +class RacePovertyOutput(MacroOutputModel): # noqa: D101 # docstring would change the published OpenAPI description white: BaselineReformValue black: BaselineReformValue hispanic: BaselineReformValue other: BaselineReformValue -class PovertyOutput(MacroOutputModel): +class PovertyOutput(MacroOutputModel): # noqa: D101 # docstring would change the published OpenAPI description poverty: AgePovertyOutput deep_poverty: AgePovertyOutput -class PovertyByGenderOutput(MacroOutputModel): +class PovertyByGenderOutput(MacroOutputModel): # noqa: D101 # docstring would change the published OpenAPI description poverty: GenderPovertyOutput deep_poverty: GenderPovertyOutput -class PovertyByRaceOutput(MacroOutputModel): +class PovertyByRaceOutput(MacroOutputModel): # noqa: D101 # docstring would change the published OpenAPI description poverty: RacePovertyOutput -class PovertyModuleOutputs(MacroOutputModel): +class PovertyModuleOutputs(MacroOutputModel): # noqa: D101 # docstring would change the published JSON schema description poverty: PovertyOutput poverty_by_gender: PovertyByGenderOutput poverty_by_race: PovertyByRaceOutput | None -class InequalityOutput(MacroOutputModel): +class InequalityOutput(MacroOutputModel): # noqa: D101 # docstring would change the published OpenAPI description gini: BaselineReformValue top_10_pct_share: BaselineReformValue top_1_pct_share: BaselineReformValue -class LaborSupplyRelativeResponse(MacroOutputModel): +class LaborSupplyRelativeResponse(MacroOutputModel): # noqa: D101 # docstring would change the published OpenAPI description income: float substitution: float -class LaborSupplyDecileMetric(MacroOutputModel): +class LaborSupplyDecileMetric(MacroOutputModel): # noqa: D101 # docstring would change the published OpenAPI description income: dict[int, float] substitution: dict[int, float] -class LaborSupplyDecileOutput(MacroOutputModel): +class LaborSupplyDecileOutput(MacroOutputModel): # noqa: D101 # docstring would change the published OpenAPI description average: LaborSupplyDecileMetric relative: LaborSupplyDecileMetric -class LaborSupplyHoursOutput(MacroOutputModel): +class LaborSupplyHoursOutput(MacroOutputModel): # noqa: D101 # docstring would change the published OpenAPI description baseline: float reform: float change: float @@ -126,7 +126,7 @@ class LaborSupplyHoursOutput(MacroOutputModel): substitution_effect: float -class LaborSupplyResponseOutput(MacroOutputModel): +class LaborSupplyResponseOutput(MacroOutputModel): # noqa: D101 # docstring would change the published OpenAPI description substitution_lsr: float income_lsr: float relative_lsr: LaborSupplyRelativeResponse @@ -136,17 +136,17 @@ class LaborSupplyResponseOutput(MacroOutputModel): hours: LaborSupplyHoursOutput -class CliffImpactInSimulation(MacroOutputModel): +class CliffImpactInSimulation(MacroOutputModel): # noqa: D101 # docstring would change the published OpenAPI description cliff_gap: float cliff_share: float -class CliffImpactOutput(MacroOutputModel): +class CliffImpactOutput(MacroOutputModel): # noqa: D101 # docstring would change the published OpenAPI description baseline: CliffImpactInSimulation reform: CliffImpactInSimulation -class ConstituencyImpactRecord(MacroOutputModel): +class ConstituencyImpactRecord(MacroOutputModel): # noqa: D101 # docstring would change the published OpenAPI description constituency_code: str constituency_name: str x: int | None @@ -156,7 +156,7 @@ class ConstituencyImpactRecord(MacroOutputModel): population: float -class LocalAuthorityImpactRecord(MacroOutputModel): +class LocalAuthorityImpactRecord(MacroOutputModel): # noqa: D101 # docstring would change the published OpenAPI description local_authority_code: str local_authority_name: str x: int | None @@ -172,7 +172,7 @@ class GeographicImpactOutput( """UK constituency or local-authority impact records.""" -class CongressionalDistrictImpactRecord(MacroOutputModel): +class CongressionalDistrictImpactRecord(MacroOutputModel): # noqa: D101 # docstring would change the published OpenAPI description district: str average_household_income_change: float relative_household_income_change: float @@ -182,7 +182,7 @@ class CongressionalDistrictImpactRecord(MacroOutputModel): population: float -class CongressionalDistrictImpactOutput(MacroOutputModel): +class CongressionalDistrictImpactOutput(MacroOutputModel): # noqa: D101 # docstring would change the published OpenAPI description districts: list[CongressionalDistrictImpactRecord] diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/spm.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/spm.py index 160175fa6..7cdd7b4b8 100644 --- a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/spm.py +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/spm.py @@ -78,6 +78,7 @@ def serialize_selection( # noqa: ANN201 # a return annotation would become the @field_validator("as_of") @classmethod def validate_as_of(cls, value: str | None) -> str | None: + """Require ``as_of`` to be a canonical ``YYYY-MM-DD`` date when given.""" if value is not None: if date.fromisoformat(value).isoformat() != value: raise ValueError("as_of must be an ISO calendar date (YYYY-MM-DD)") @@ -85,6 +86,11 @@ def validate_as_of(cls, value: str | None) -> str | None: @model_validator(mode="after") def validate_location(self) -> Self: + """Require ``geography_id`` for metro selections and forbid it otherwise. + + Only checked once both geography fields are explicitly set, because an + omitted option inherits the bundle default. + """ # Request options inherit bundle defaults. Validate coupled options # once both are explicit (including after defaults are resolved). if not {"geography_kind", "geography_id"} <= self.model_fields_set: @@ -114,7 +120,6 @@ def serialize_resolved_configuration( # noqa: ANN201 # a return annotation wou self, handler: SerializerFunctionWrapHandler, info: SerializationInfo ): """Keep required nullable fields in completed response bodies.""" - dumped = handler(self) if info.exclude_none: for name in ("geography_id", "as_of"): @@ -129,12 +134,14 @@ def serialize_resolved_configuration( # noqa: ANN201 # a return annotation wou @field_validator("as_of") @classmethod def validate_as_of(cls, value: str | None) -> str | None: + """Require ``as_of`` to be a canonical ``YYYY-MM-DD`` date when given.""" if value is not None and date.fromisoformat(value).isoformat() != value: raise ValueError("as_of must be an ISO calendar date (YYYY-MM-DD)") return value @model_validator(mode="after") def validate_location(self) -> "SPMResolvedConfiguration": + """Require ``geography_id`` for metro selections and forbid it otherwise.""" if self.geography_kind == "metro": if self.geography_id is None or not self.geography_id.strip(): raise ValueError("An SPM area selection requires geography_id") @@ -164,7 +171,6 @@ def serialize_provenance( # noqa: ANN201 # a return annotation would become th self, handler: SerializerFunctionWrapHandler, info: SerializationInfo ): """Keep required nullable configuration fields in response bodies.""" - dumped = handler(self) if info.exclude_none: for name in ("geography_id", "as_of"): @@ -179,6 +185,7 @@ def serialize_provenance( # noqa: ANN201 # a return annotation would become th @field_validator("years") @classmethod def validate_years(cls, value: tuple[str, ...]) -> tuple[str, ...]: + """Require four-digit years, sorted ascending with no duplicates.""" if any(re.fullmatch(r"[0-9]{4}", year) is None for year in value): raise ValueError("SPM provenance years must use four-digit years") if value != tuple(sorted(set(value))): @@ -187,6 +194,7 @@ def validate_years(cls, value: tuple[str, ...]) -> tuple[str, ...]: @model_validator(mode="after") def validate_provenance_location(self) -> "SPMProvenance": + """Require ``geography_id`` for metro receipts and forbid it otherwise.""" if self.geography_kind == "metro": if self.geography_id is None or not self.geography_id.strip(): raise ValueError("An SPM area receipt requires geography_id") @@ -228,7 +236,6 @@ def _resolved_spm_configuration( selection: SPMSelection, ) -> SPMResolvedConfiguration: """Convert a resolved request selection to the completed-result contract.""" - _require_complete_selection(selection) assert selection.forecast_content_sha256 is not None assert selection.scenario is not None @@ -246,7 +253,6 @@ def _resolved_spm_configuration_from_provenance( receipt: SPMProvenance, ) -> SPMResolvedConfiguration: """Derive internal configuration from a receipt without a sibling object.""" - return SPMResolvedConfiguration( forecast_content_sha256=receipt.forecast_sha256, scenario=receipt.scenario, @@ -266,7 +272,6 @@ def build_spm_provenance( runtime_versions: SPMRuntimeVersions, ) -> SPMProvenance: """Build one compact receipt from resolved calculation inputs.""" - _require_complete_selection(selection) assert selection.forecast_content_sha256 is not None assert selection.scenario is not None @@ -294,7 +299,6 @@ def validate_spm_calculation_provenance( receipt: SPMProvenance, ) -> SPMProvenance: """Validate and return the sole compact receipt for a calculation.""" - resolved = _resolved_spm_configuration(config) receipt_configuration = _resolved_spm_configuration_from_provenance(receipt) for field_name in SPMResolvedConfiguration.model_fields: @@ -319,7 +323,7 @@ def validate_spm_calculation_provenance( ) -class SPMErrorDetail(BaseModel): +class SPMErrorDetail(BaseModel): # noqa: D101 # docstring would change the published OpenAPI description code: str message: str @@ -328,14 +332,17 @@ class SPMInputError(ValueError): """Transportable error shared by the control plane and worker.""" def __init__(self, code: str, message: str) -> None: + """Keep the public error ``code`` and its human-readable ``message``.""" self.code = code self.message = message super().__init__(code, message) def __str__(self) -> str: + """Return the message alone, without the error code.""" return self.message def to_dict(self) -> dict[str, str]: + """Return the ``{"code", "message"}`` wire form of this error.""" return {"code": self.code, "message": self.message} @@ -352,13 +359,18 @@ def spm_error_detail(exc: BaseException) -> SPMErrorDetail | None: return None -class SPMCapability(BaseModel): +class SPMCapability(BaseModel): # noqa: D101 # docstring would change the published OpenAPI description model_config = ConfigDict(extra="forbid") contract_version: Literal["canonical-spm-v1"] = SPM_CONTRACT_VERSION defaults: SPMSelection @model_validator(mode="after") def pinned_defaults(self) -> Self: + """Require pinned defaults and mark every default field as explicitly set. + + Rebuilding ``defaults`` from all fields makes serialisation include each + default, since ``SPMSelection`` serialises only the fields that were set. + """ if not self.defaults.forecast_content_sha256 or not self.defaults.scenario: raise ValueError( "Certified SPM defaults must pin artifact hash and scenario" @@ -389,6 +401,7 @@ class SPMComparisonProvenance(BaseModel): @model_validator(mode="after") def validate_compatible_receipts(self) -> "SPMComparisonProvenance": + """Reject comparisons whose baseline and reform receipts differ.""" if self.baseline.receipt != self.reform.receipt: raise ValueError("baseline and reform SPM receipts differ") return self @@ -412,7 +425,6 @@ def build_spm_comparison_provenance( reform_receipts: Sequence[SPMProvenance], ) -> SPMComparisonProvenance: """Collapse identical child receipts into one receipt per policy side.""" - baseline = _collapse_spm_receipts("baseline", baseline_receipts) reform = _collapse_spm_receipts("reform", reform_receipts) if baseline.receipt != reform.receipt: @@ -552,6 +564,12 @@ def combine_spm_results( *, expected_year: int | str | None = None, ) -> dict: + """Merge per-segment SPM receipts into one ``spm_provenance`` result entry. + + Every segment is validated against ``selection``. Returns an empty dict + when no SPM selection was requested; raises SPMInputError when any + segment lacks a receipt or the receipts disagree. + """ receipts = [ validate_spm_result(result, selection, expected_year=expected_year) for result in results diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_bundle.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_bundle.py index d20fa1ccb..2b0b929ce 100644 --- a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_bundle.py +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_bundle.py @@ -12,11 +12,11 @@ SUPPORTED_COUNTRIES: tuple[CountryId, ...] = ("us", "uk") -class StrictBundleModel(BaseModel): +class StrictBundleModel(BaseModel): # noqa: D101 # docstring would change the published JSON schema description model_config = ConfigDict(extra="forbid", frozen=True) -class Stage12Dataset(StrictBundleModel): +class Stage12Dataset(StrictBundleModel): # noqa: D101 # docstring would change the published JSON schema description identity: NonEmptyText uri: NonEmptyText artifact_revision: NonEmptyText @@ -24,7 +24,7 @@ class Stage12Dataset(StrictBundleModel): repo_type: NonEmptyText -class Stage12CountryBundle(StrictBundleModel): +class Stage12CountryBundle(StrictBundleModel): # noqa: D101 # docstring would change the published JSON schema description country: CountryId country_package_name: NonEmptyText country_package_version: NonEmptyText @@ -51,6 +51,6 @@ class Stage12BundleManifest(StrictBundleModel): countries: tuple[Stage12CountryBundle, ...] -class ResolvedStage12Bundle(StrictBundleModel): +class ResolvedStage12Bundle(StrictBundleModel): # noqa: D101 # docstring would change the published JSON schema description bundle: Stage12BundleManifest bundle_manifest_sha256: Sha256Digest diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_execution.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_execution.py index 99c32bfc4..aff9f6b68 100644 --- a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_execution.py +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_execution.py @@ -53,17 +53,17 @@ ] -class StrictContractModel(BaseModel): +class StrictContractModel(BaseModel): # noqa: D101 # docstring would change the published JSON schema description model_config = ConfigDict(extra="forbid", allow_inf_nan=False, frozen=True) -class SimulationRole(StrEnum): +class SimulationRole(StrEnum): # noqa: D101 # docstring would change the published JSON schema description BASELINE = "baseline" REFORM = "reform" STANDALONE = "standalone" -class ArtifactMediaType(StrEnum): +class ArtifactMediaType(StrEnum): # noqa: D101 # docstring would change the published JSON schema description JSON = "application/json" PARQUET = "application/vnd.apache.parquet" @@ -104,13 +104,13 @@ class SimulationParquetPayloadContract(StrictContractModel): SIMULATION_PARQUET_PAYLOAD_CONTRACT = SimulationParquetPayloadContract() -class DatasetArtifactMediaType(StrEnum): +class DatasetArtifactMediaType(StrEnum): # noqa: D101 # docstring would change the published JSON schema description HDF5 = "application/x-hdf5" JSON = "application/json" PARQUET = "application/vnd.apache.parquet" -class DatasetProvenance(StrictContractModel): +class DatasetProvenance(StrictContractModel): # noqa: D101 # docstring would change the published JSON schema description identity: ContractText uri: Annotated[str, Field(min_length=1, max_length=2048)] artifact_revision: ContractText @@ -118,7 +118,7 @@ class DatasetProvenance(StrictContractModel): data_package_version: ContractText -class BundleProvenance(StrictContractModel): +class BundleProvenance(StrictContractModel): # noqa: D101 # docstring would change the published JSON schema description policyengine_version: ContractText country_package_name: ContractText country_package_version: ContractText @@ -126,26 +126,26 @@ class BundleProvenance(StrictContractModel): bundle_manifest_sha256: Sha256Digest -class ArtifactReference(StrictContractModel): +class ArtifactReference(StrictContractModel): # noqa: D101 # docstring would change the published JSON schema description uri: StorageUri media_type: ArtifactMediaType content_sha256: Sha256Digest size_bytes: Annotated[int, Field(ge=0)] -class DatasetArtifactReference(StrictContractModel): +class DatasetArtifactReference(StrictContractModel): # noqa: D101 # docstring would change the published JSON schema description uri: DatasetArtifactUri media_type: DatasetArtifactMediaType content_sha256: Sha256Digest size_bytes: Annotated[int, Field(ge=0)] | None = None -class DatasetPopulationInput(StrictContractModel): +class DatasetPopulationInput(StrictContractModel): # noqa: D101 # docstring would change the published JSON schema description kind: Literal["dataset"] = "dataset" artifact: DatasetArtifactReference -class HouseholdPopulationInput(StrictContractModel): +class HouseholdPopulationInput(StrictContractModel): # noqa: D101 # docstring would change the published JSON schema description kind: Literal["household"] = "household" document: dict[str, JsonValue] @@ -156,7 +156,7 @@ class HouseholdPopulationInput(StrictContractModel): ] -class GeographySelection(StrictContractModel): +class GeographySelection(StrictContractModel): # noqa: D101 # docstring would change the published JSON schema description country: CountryId region: ContractText filter_field: ContractText | None = None @@ -165,6 +165,7 @@ class GeographySelection(StrictContractModel): @model_validator(mode="after") def require_complete_filter(self) -> GeographySelection: + """Require the filter field and value together, and both for a strategy.""" if (self.filter_field is None) != (self.filter_value is None): raise ValueError("filter_field and filter_value must be supplied together") if self.filter_strategy is not None and self.filter_field is None: @@ -172,7 +173,7 @@ def require_complete_filter(self) -> GeographySelection: return self -class SimulationExecutionInput(StrictContractModel): +class SimulationExecutionInput(StrictContractModel): # noqa: D101 # docstring would change the published JSON schema description contract_version: Literal[1] = 1 evaluation_id: UUID simulation_execution_id: UUID @@ -187,6 +188,11 @@ class SimulationExecutionInput(StrictContractModel): @model_validator(mode="before") @classmethod def reject_combined_policy_input(cls, value: object) -> object: + """Reject report-shaped input that carries ``baseline`` or ``reform`` keys. + + One execution input describes one simulation; the pair belongs in + ``ReportExecutionInput``. + """ if isinstance(value, dict) and ({"baseline", "reform"} & value.keys()): raise ValueError( "single-simulation input cannot contain baseline or reform fields" @@ -194,7 +200,7 @@ def reject_combined_policy_input(cls, value: object) -> object: return value -class ReportAggregate(StrEnum): +class ReportAggregate(StrEnum): # noqa: D101 # docstring would change the published JSON schema description BUDGET = "budget" POVERTY = "poverty" INEQUALITY = "inequality" @@ -218,6 +224,7 @@ def require_complete_aggregate_profile( cls, value: tuple[ReportAggregate, ...], ) -> tuple[ReportAggregate, ...]: + """Require every ``ReportAggregate``, once each and in declaration order.""" if len(value) != len(set(value)): raise ValueError("report output aggregates must be unique") if value != tuple(ReportAggregate): @@ -245,6 +252,7 @@ class EntityOutputPlan(StrictContractModel): ) @classmethod def require_canonical_variables(cls, value: tuple[str, ...]) -> tuple[str, ...]: + """Require variable names to be unique and sorted, so plans hash stably.""" if len(value) != len(set(value)): raise ValueError("output-plan variables must be unique") if value != tuple(sorted(value)): @@ -253,6 +261,10 @@ def require_canonical_variables(cls, value: tuple[str, ...]) -> tuple[str, ...]: @model_validator(mode="after") def require_variable_subsets(self) -> EntityOutputPlan: + """Require the additional and dataset variables not to overlap. + + Both must also be subsets of ``materialized_variables``. + """ if not set(self.additional_variables).issubset(self.materialized_variables): raise ValueError( "additional output-plan variables must be a materialized subset" @@ -282,6 +294,7 @@ def require_canonical_entities( cls, value: tuple[EntityOutputPlan, ...], ) -> tuple[EntityOutputPlan, ...]: + """Require entity plans to be unique and sorted by entity name.""" names = tuple(entity.entity for entity in value) if len(names) != len(set(names)): raise ValueError("output-plan entities must be unique") @@ -292,7 +305,6 @@ def require_canonical_entities( def stage12_output_plan_sha256(plan: Stage12OutputPlan) -> str: """Return the stable digest recorded by each Stage 12 child artifact.""" - payload = json.dumps( plan.model_dump(mode="json"), sort_keys=True, @@ -310,12 +322,13 @@ class PlannedSimulationExecutionInput(SimulationExecutionInput): @model_validator(mode="after") def require_output_plan_country(self) -> PlannedSimulationExecutionInput: + """Reject an output plan built for a different country than the geography.""" if self.output_plan.country != self.geography.country: raise ValueError("output plan country must match simulation geography") return self -class RowIdentity(StrictContractModel): +class RowIdentity(StrictContractModel): # noqa: D101 # docstring would change the published JSON schema description schema_version: Literal[1] = 1 identifier_columns: Annotated[tuple[ContractText, ...], Field(min_length=1)] row_count: Annotated[int, Field(ge=0)] @@ -324,12 +337,13 @@ class RowIdentity(StrictContractModel): @field_validator("identifier_columns") @classmethod def require_unique_columns(cls, value: tuple[str, ...]) -> tuple[str, ...]: + """Reject repeated identifier column names.""" if len(value) != len(set(value)): raise ValueError("row identifier columns must be unique") return value -class SimulationArtifactDescriptor(StrictContractModel): +class SimulationArtifactDescriptor(StrictContractModel): # noqa: D101 # docstring would change the published JSON schema description contract_version: Literal[1] = 1 evaluation_id: UUID simulation_execution_id: UUID @@ -342,7 +356,7 @@ class SimulationArtifactDescriptor(StrictContractModel): calculation_provenance: SPMProvenance | None = None -class ReportExecutionInput(StrictContractModel): +class ReportExecutionInput(StrictContractModel): # noqa: D101 # docstring would change the published JSON schema description contract_version: Literal[1] = 1 evaluation_id: UUID baseline: SimulationExecutionInput @@ -351,6 +365,11 @@ class ReportExecutionInput(StrictContractModel): @model_validator(mode="after") def require_aligned_simulations(self) -> ReportExecutionInput: + """Check the two simulations form one comparable report. + + Roles must be baseline and reform, both must belong to this report with + distinct execution ids, and everything but the policy must match. + """ if self.baseline.role is not SimulationRole.BASELINE: raise ValueError("baseline input must use the baseline role") if self.reform.role is not SimulationRole.REFORM: @@ -377,7 +396,7 @@ def require_aligned_simulations(self) -> ReportExecutionInput: return self -class AggregateReportArtifactDescriptor(StrictContractModel): +class AggregateReportArtifactDescriptor(StrictContractModel): # noqa: D101 # docstring would change the published JSON schema description contract_version: Literal[1] = 1 evaluation_id: UUID artifact: ArtifactReference @@ -403,6 +422,7 @@ def require_unique_aggregates( cls, value: tuple[ReportAggregate, ...], ) -> tuple[ReportAggregate, ...]: + """Reject repeated aggregate names.""" if len(value) != len(set(value)): raise ValueError("requested aggregates must be unique") return value @@ -410,7 +430,6 @@ def require_unique_aggregates( @model_validator(mode="after") def validate_spm_provenance(self) -> AggregateReportArtifactPayload: """Reject duplicate or noncanonical SPM data in aggregate artifacts.""" - if "spm_config" in self.result: raise ValueError("aggregate result must not include legacy spm_config") # Reports without US SPM calculations serialize this optional field @@ -434,6 +453,11 @@ class ResultDifference(StrictContractModel): @model_validator(mode="after") def validate_presence(self) -> ResultDifference: + """Check presence flags against values, and that deltas come as a pair. + + An absent side must be null, and a relative delta needs an absolute one; + only an absolute delta may stand alone (a zero production value). + """ if not self.production_present and not self.stage12_present: raise ValueError( "a result difference must be present in at least one result" @@ -466,6 +490,7 @@ class ResultComparisonArtifactPayload(StrictContractModel): @model_validator(mode="after") def validate_differences(self) -> ResultComparisonArtifactPayload: + """Require a timezone-aware timestamp and a count and status that match.""" if self.compared_at.tzinfo is None: raise ValueError("comparison timestamp must include a timezone") if self.difference_count != len(self.differences): @@ -476,7 +501,7 @@ def validate_differences(self) -> ResultComparisonArtifactPayload: return self -class Stage12InvocationContext(StrictContractModel): +class Stage12InvocationContext(StrictContractModel): # noqa: D101 # docstring would change the published JSON schema description contract_version: Literal[1] = 1 request_id: ContractText # The logical environment labels durable records and artifact paths. @@ -499,6 +524,7 @@ class Stage12InvocationContext(StrictContractModel): @model_validator(mode="after") def validate_retention(self) -> Stage12InvocationContext: + """Require aware timestamps and a retention window of at most 30 days.""" _validate_retention(self.created_at, self.retention_expires_at) return self @@ -518,7 +544,7 @@ def _validate_retention(created_at: datetime, retention_expires_at: datetime) -> raise ValueError("comparison retention must not exceed 30 days") -class ComparisonRunLifecycleStatus(StrEnum): +class ComparisonRunLifecycleStatus(StrEnum): # noqa: D101 # docstring would change the published JSON schema description PENDING = "pending" RUNNING = "running" SUCCEEDED = "succeeded" @@ -527,14 +553,14 @@ class ComparisonRunLifecycleStatus(StrEnum): SKIPPED = "skipped" -class ComparisonRunAggregationStatus(StrEnum): +class ComparisonRunAggregationStatus(StrEnum): # noqa: D101 # docstring would change the published JSON schema description NOT_STARTED = "not_started" RUNNING = "running" SUCCEEDED = "succeeded" FAILED = "failed" -class ResultComparisonStatus(StrEnum): +class ResultComparisonStatus(StrEnum): # noqa: D101 # docstring would change the published JSON schema description NOT_REQUESTED = "not_requested" PENDING = "pending" RUNNING = "running" @@ -590,6 +616,12 @@ class ComparisonReportRecord(StrictContractModel): @model_validator(mode="after") def validate_state(self) -> ComparisonReportRecord: + """Check that the row's fields are consistent with its lifecycle states. + + Aggregate and comparison output fields come all-or-nothing, a success + needs its outputs and ``completed_at``, failures and skips need an + ``error_code``, and the retention window is at most 30 days. + """ _validate_retention(self.created_at, self.retention_expires_at) output_fields = ( self.aggregate_output_uri, @@ -692,6 +724,12 @@ class ComparisonSimulationRecord(StrictContractModel): @model_validator(mode="after") def validate_state(self) -> ComparisonSimulationRecord: + """Check that the row's fields are consistent with its lifecycle status. + + Output and row-identity fields come all-or-nothing, a success needs + them and ``completed_at``, failures and skips need an ``error_code``, + and the retention window is at most 30 days. + """ _validate_retention(self.created_at, self.retention_expires_at) output_fields = ( self.output_uri, @@ -729,11 +767,21 @@ def validate_state(self) -> ComparisonSimulationRecord: @dataclass(frozen=True) class ComparisonReportPersistenceResult: + """The stored parent row after an idempotent create. + + ``created`` is False when an identical row already existed and was reused. + """ + record: ComparisonReportRecord created: bool @dataclass(frozen=True) class ComparisonSimulationPersistenceResult: + """The stored child row after an idempotent create. + + ``created`` is False when an identical row already existed and was reused. + """ + record: ComparisonSimulationRecord created: bool diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_manifest.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_manifest.py index 98178064d..996e7db92 100644 --- a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_manifest.py +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_manifest.py @@ -30,16 +30,16 @@ CallableName = Annotated[str, Field(pattern=_CALLABLE_PATTERN.pattern)] -class StrictManifestModel(BaseModel): +class StrictManifestModel(BaseModel): # noqa: D101 # docstring would change the published JSON schema description model_config = ConfigDict(extra="forbid", frozen=True) -class V2CountryWorker(StrictManifestModel): +class V2CountryWorker(StrictManifestModel): # noqa: D101 # docstring would change the published JSON schema description country: CountryId single_simulation_callable: CallableName -class V2WorkerValidation(StrictManifestModel): +class V2WorkerValidation(StrictManifestModel): # noqa: D101 # docstring would change the published JSON schema description validated: Literal[True] application_name: ManifestText bundle_manifest_sha256: Sha256Digest @@ -48,7 +48,7 @@ class V2WorkerValidation(StrictManifestModel): country_validation_invocation_ids: dict[CountryId, ManifestText] -class V2WorkerVersion(StrictManifestModel): +class V2WorkerVersion(StrictManifestModel): # noqa: D101 # docstring would change the published JSON schema description application_name: ManifestText report_coordinator_callable: CallableName countries: tuple[V2CountryWorker, ...] @@ -58,6 +58,13 @@ class V2WorkerVersion(StrictManifestModel): @model_validator(mode="after") def validate_worker(self) -> V2WorkerVersion: + """Check the worker's name, validation receipt and countries match its bundle. + + The application name must be the v2 name derived from the bundle's + PolicyEngine.py version, the validation must name this application and + bundle digest, and the countries must equal the bundle's, in order, + each with a validation invocation. + """ if not self.application_name.startswith(V2_APPLICATION_PREFIX): raise ValueError("v2 application uses an invalid name prefix") if self.application_name.startswith(V1_APPLICATION_PREFIX): @@ -80,7 +87,7 @@ def validate_worker(self) -> V2WorkerVersion: return self -class V2VersionManifest(StrictManifestModel): +class V2VersionManifest(StrictManifestModel): # noqa: D101 # docstring would change the published JSON schema description schema_version: Literal[1] = 1 generation: ManifestText default_version: ManifestText @@ -88,6 +95,7 @@ class V2VersionManifest(StrictManifestModel): @model_validator(mode="after") def validate_versions(self) -> V2VersionManifest: + """Check the default has an entry and each key is its bundle's version.""" if not self.versions: raise ValueError("v2 manifest must contain a worker version") if self.default_version not in self.versions: @@ -99,12 +107,23 @@ def validate_versions(self) -> V2VersionManifest: class V2ManifestStore(Protocol): - def get(self, key: str, default: object = None) -> object: ... + """The part of a Modal Dict's mapping interface the manifest helpers use.""" + + def get(self, key: str, default: object = None) -> object: + """Return the value stored under ``key``, or ``default`` when absent.""" + ... - def __setitem__(self, key: str, value: object) -> None: ... + def __setitem__(self, key: str, value: object) -> None: + """Store ``value`` under ``key``, replacing any previous value.""" + ... def v2_application_name(policyengine_version: str) -> str: + """Derive the v2 Modal application name from a PolicyEngine.py version. + + Dots and plus signs become hyphens. Raises ValueError for a version with + characters a Modal name cannot hold, or a name over 64 characters. + """ if not policyengine_version or not re.fullmatch( r"[A-Za-z0-9][A-Za-z0-9._+-]*", policyengine_version ): @@ -118,6 +137,7 @@ def v2_application_name(policyengine_version: str) -> str: def manifest_digest(manifest: V2VersionManifest) -> str: + """Return the SHA-256 of the manifest's canonical JSON (sorted, compact).""" payload = json.dumps( manifest.model_dump(mode="json"), sort_keys=True, @@ -143,6 +163,12 @@ def build_next_v2_manifest( worker: V2WorkerVersion, force_default: bool = False, ) -> V2VersionManifest: + """Return ``current`` with ``worker`` added under its bundle version. + + ``current`` is the raw stored manifest, or None before the first publish. + The default moves to the new version only when it is newer by PEP 440 + ordering (an unparsable version never is) or ``force_default`` is set. + """ if current is None: versions: dict[str, V2WorkerVersion] = {} previous_default = None @@ -172,6 +198,11 @@ def publish_v2_manifest( worker: V2WorkerVersion, force_default: bool = False, ) -> V2VersionManifest: + """Add ``worker`` to the stored manifest, write it back and return it. + + The whole manifest is written as one store value, so readers never see a + partial update. + """ next_manifest = build_next_v2_manifest( current=store.get(ACTIVE_MANIFEST_KEY), worker=worker, @@ -187,10 +218,16 @@ class V2ManifestLoader: """Load v2 state while retaining only this loader's last valid v2 copy.""" def __init__(self, store: V2ManifestStore) -> None: + """Wrap ``store``; nothing is read until the first ``load``.""" self._store = store self._last_valid: V2VersionManifest | None = None def load(self) -> V2VersionManifest: + """Return the stored manifest, or this loader's last valid copy. + + The fallback covers a missing or invalid stored value; with no earlier + valid copy, ValueError is raised instead. + """ value = self._store.get(ACTIVE_MANIFEST_KEY) try: manifest = V2VersionManifest.model_validate(value) @@ -202,6 +239,10 @@ def load(self) -> V2VersionManifest: return manifest def resolve(self, version: str | None = None) -> tuple[str, V2WorkerVersion, str]: + """Return ``(version, worker, manifest digest)`` for ``version`` or the default. + + Raises ValueError when the manifest has no worker for that version. + """ manifest = self.load() selected = version or manifest.default_version try: @@ -214,6 +255,10 @@ def resolve(self, version: str | None = None) -> tuple[str, V2WorkerVersion, str def assert_separate_manifest_names(*, v1_name: str, v2_name: str) -> None: + """Raise ValueError unless both names are set and differ. + + This keeps a v2 manifest publish from overwriting the v1 routing state. + """ if not v1_name or not v2_name: raise ValueError("v1 and v2 manifest names must be non-empty") if v1_name == v2_name: @@ -221,6 +266,11 @@ def assert_separate_manifest_names(*, v1_name: str, v2_name: str) -> None: def validate_manifest_has_no_credentials(value: Mapping[str, object]) -> None: + """Raise ValueError if any key, at any depth, looks like a credential. + + Keys containing password, secret, token, credential or private_key match + case-insensitively; the error names the dotted path to the key. + """ prohibited = {"password", "secret", "token", "credential", "private_key"} def visit(item: object, path: tuple[str, ...]) -> None: diff --git a/libs/policyengine-simulation-contract/tests/test_budget_window_state.py b/libs/policyengine-simulation-contract/tests/test_budget_window_state.py index 966e70b1c..1b89f77b2 100644 --- a/libs/policyengine-simulation-contract/tests/test_budget_window_state.py +++ b/libs/policyengine-simulation-contract/tests/test_budget_window_state.py @@ -201,10 +201,12 @@ def test_state_transition_helpers_track_failed_child(): def test_mark_child_completed_handles_missing_child_jobs_entry(caplog): - """Regression test for #454: ``mark_child_completed`` crashed with a - ``KeyError`` when transitions were invoked without a preceding - ``mark_child_started`` call. The guard should synthesise a sentinel and - log a warning instead of aborting the batch.""" + """Complete a child that was never started without raising (regression #454). + + ``mark_child_completed`` crashed with a ``KeyError`` when transitions were + invoked without a preceding ``mark_child_started`` call. The guard should + synthesise a sentinel and log a warning instead of aborting the batch. + """ request = BudgetWindowBatchRequest( country="us", region="us", diff --git a/libs/policyengine-simulation-contract/tests/test_gateway_models.py b/libs/policyengine-simulation-contract/tests/test_gateway_models.py index a40576b7b..8fb126fe4 100644 --- a/libs/policyengine-simulation-contract/tests/test_gateway_models.py +++ b/libs/policyengine-simulation-contract/tests/test_gateway_models.py @@ -29,11 +29,7 @@ class TestPingRequest: """Tests for PingRequest model.""" def test_ping_request_accepts_integer_value(self): - """ - Given an integer value - When creating a PingRequest - Then the model is created with the value. - """ + """Given an integer value, a PingRequest is created with that value.""" # Given value = 42 @@ -44,11 +40,7 @@ def test_ping_request_accepts_integer_value(self): assert request.value == 42 def test_ping_request_accepts_negative_value(self): - """ - Given a negative integer value - When creating a PingRequest - Then the model is created with the negative value. - """ + """Given a negative integer, a PingRequest is created with that value.""" # Given value = -10 @@ -59,11 +51,7 @@ def test_ping_request_accepts_negative_value(self): assert request.value == -10 def test_ping_request_rejects_non_integer(self): - """ - Given a non-integer value - When creating a PingRequest - Then a ValidationError is raised. - """ + """Given a non-integer, creating a PingRequest raises ValidationError.""" # Given value = "not_an_integer" @@ -72,11 +60,7 @@ def test_ping_request_rejects_non_integer(self): PingRequest(value=value) def test_ping_request_rejects_missing_value(self): - """ - Given no value - When creating a PingRequest - Then a ValidationError is raised. - """ + """Given no value, creating a PingRequest raises ValidationError.""" # When/Then with pytest.raises(ValidationError): PingRequest() @@ -86,11 +70,7 @@ class TestPingResponse: """Tests for PingResponse model.""" def test_ping_response_accepts_integer_incremented(self): - """ - Given an integer incremented value - When creating a PingResponse - Then the model is created with the value. - """ + """Given an incremented integer, a PingResponse is created with that value.""" # Given incremented = 43 @@ -101,11 +81,7 @@ def test_ping_response_accepts_integer_incremented(self): assert response.incremented == 43 def test_ping_response_serializes_correctly(self): - """ - Given a PingResponse - When converting to dict - Then the correct structure is returned. - """ + """Given a PingResponse, model_dump returns the expected structure.""" # Given response = PingResponse(incremented=100) @@ -130,21 +106,13 @@ def _payload_with_encoded_size(target_size: int) -> dict: return payload def test_simulation_request_requires_country(self): - """ - Given no country - When creating a SimulationRequest - Then a ValidationError is raised. - """ + """Given no country, creating a SimulationRequest raises ValidationError.""" # When/Then with pytest.raises(ValidationError): SimulationRequest() def test_simulation_request_accepts_country_only(self): - """ - Given only a country - When creating a SimulationRequest - Then the model is created with version as None. - """ + """Given only a country, a SimulationRequest is created with version None.""" # Given country = "us" @@ -156,11 +124,7 @@ def test_simulation_request_accepts_country_only(self): assert request.version is None def test_simulation_request_accepts_country_and_version(self): - """ - Given a country and version - When creating a SimulationRequest - Then the model is created with both values. - """ + """Given a country and version, a SimulationRequest keeps both values.""" # Given country = "uk" version = "1.0.0" @@ -208,10 +172,10 @@ def test_simulation_request_maps_legacy_process_id(self): assert "process_id" not in request.telemetry.model_dump() def test_simulation_request_accepts_documented_simulation_fields(self): - """ - Given the documented simulation fields (reform, region, scope) - When creating a SimulationRequest - Then the model accepts and preserves them. + """Accept and preserve the documented simulation fields. + + Given reform, region and scope, when creating a SimulationRequest, then + the model accepts and preserves them. """ # Given data = { @@ -248,7 +212,7 @@ def test_simulation_request_accepts_neither_region_nor_region_group(self): assert request.region_group is None def test_simulation_request_rejects_region_and_region_group_together(self): - """region and region_group are mutually exclusive.""" + """Region and region_group are mutually exclusive.""" with pytest.raises(ValidationError, match="not both"): SimulationRequest( country="us", region="state/hi", region_group=["state/ia"] @@ -294,8 +258,11 @@ def test_simulation_request_rejects_dataset_overrides(self, field): ) def test_simulation_request_rejects_oversized_payload(self): - """Payloads that exceed the gateway max should 422 before Pydantic - allocates proxy objects for the nested reform dict.""" + """Reject payloads over the gateway maximum before parsing the reform. + + They should 422 before Pydantic allocates proxy objects for the nested + reform dict. + """ giant_reform = { f"mock.parameter[{i}]": {"2024-01-01": i} for i in range(10_000) } @@ -303,8 +270,7 @@ def test_simulation_request_rejects_oversized_payload(self): SimulationRequest(country="us", reform=giant_reform) def test_simulation_request_accepts_payload_just_below_256kb(self): - """256 KB boundary (#450): a payload just below the cap must be - accepted.""" + """Accept a payload just below the 256 KB cap (boundary from #450).""" payload = self._payload_with_encoded_size(MAX_GATEWAY_REQUEST_BYTES - 1) # Must not raise — this is the just-under-cap happy path. @@ -314,19 +280,17 @@ def test_simulation_request_accepts_payload_just_below_256kb(self): assert request.reform.root == payload["reform"] def test_simulation_request_rejects_payload_just_above_256kb(self): - """The cap is strict: a payload that crosses 262_144 bytes by even a - few bytes should be rejected with a ``too large`` ValidationError.""" + """Reject a payload a few bytes over the strict 262_144-byte cap. + + The ValidationError should report the payload as ``too large``. + """ payload = self._payload_with_encoded_size(MAX_GATEWAY_REQUEST_BYTES + 1) with pytest.raises(ValidationError, match="too large"): SimulationRequest(**payload) def test_simulation_request_accepts_typed_telemetry_envelope(self): - """ - Given a telemetry envelope - When creating a SimulationRequest - Then the envelope is validated and preserved. - """ + """Given a telemetry envelope, a SimulationRequest validates and keeps it.""" request = SimulationRequest( country="us", _telemetry={ @@ -343,11 +307,7 @@ class TestJobSubmitResponse: """Tests for JobSubmitResponse model.""" def test_job_submit_response_creates_with_all_fields(self): - """ - Given all required fields - When creating a JobSubmitResponse - Then the model is created correctly. - """ + """Given all required fields, a JobSubmitResponse is created correctly.""" # Given data = { "job_id": "fc-abc123", @@ -385,10 +345,11 @@ class TestJobStatusResponse: """Tests for JobStatusResponse model.""" def test_job_status_response_rejects_an_unstructured_result(self): - """ - Given a completed job with a result that does not match the macro schema - When creating a JobStatusResponse - Then validation fails instead of silently accepting an opaque mapping. + """Reject a completed result that does not match the macro output schema. + + Given a completed job with such a result, when creating a + JobStatusResponse, then validation fails instead of silently accepting + an opaque mapping. """ with pytest.raises(ValidationError): JobStatusResponse( @@ -397,11 +358,7 @@ def test_job_status_response_rejects_an_unstructured_result(self): ) def test_job_status_response_running_without_result(self): - """ - Given a running job - When creating a JobStatusResponse - Then result and error are None. - """ + """Given a running job, a JobStatusResponse has no result and no error.""" # When response = JobStatusResponse(status="running") @@ -411,11 +368,7 @@ def test_job_status_response_running_without_result(self): assert response.error is None def test_job_status_response_failed_with_error(self): - """ - Given a failed job with error message - When creating a JobStatusResponse - Then the error is captured. - """ + """Given a failed job with a message, a JobStatusResponse captures the error.""" # Given error_msg = "Simulation timed out" diff --git a/libs/policyengine-simulation-contract/tests/test_spm_selection.py b/libs/policyengine-simulation-contract/tests/test_spm_selection.py index 3ad50ac18..b73c38d0a 100644 --- a/libs/policyengine-simulation-contract/tests/test_spm_selection.py +++ b/libs/policyengine-simulation-contract/tests/test_spm_selection.py @@ -123,8 +123,11 @@ def test_pinned_pre_canonical_wrapper_is_historical_without_a_stated_model( ], ) def test_unpinned_or_contradicted_wrapper_is_not_historical(wrapper, model): - """A stated model version that contradicts the pin, and any wrapper - outside the pinned pre-canonical set, still fail closed.""" + """Fail closed for wrappers that are not provably pre-canonical. + + A stated model version that contradicts the pin, and any wrapper outside + the pinned pre-canonical set, still fail closed. + """ with pytest.raises(SPMInputError) as error: resolve_spm_selection( "us", From 3c86f37936d1279b0f34c860386bbee526e1aae9 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 13 Oct 2026 16:56:47 +0400 Subject: [PATCH 02/18] Document the FastAPI support library for ruff D The ping and health models and the /ping/alive route are mounted in the executor app, so they keep their existing (OpenAPI-visible) docstrings under noqa. The no-op observability methods keep their explicit return None under RET501 noqa, and two stray string statements in the GCP module become the docstrings they were meant to be. Co-Authored-By: Claude Opus 5.5 --- .../fixtures/ping/shared.py | 2 + .../src/policyengine_fastapi/__init__.py | 4 +- .../src/policyengine_fastapi/auth/__init__.py | 2 + .../policyengine_fastapi/auth/jwt_decoder.py | 20 ++++++- .../src/policyengine_fastapi/database.py | 14 +++-- .../src/policyengine_fastapi/exit.py | 17 +++++- .../policyengine_fastapi/health/__init__.py | 19 ++++-- .../observability/__init__.py | 3 +- .../observability/config.py | 10 +++- .../observability/contracts.py | 26 ++++++++ .../observability/correlation.py | 8 +++ .../observability/emitters.py | 60 ++++++++++++++----- .../observability/provider.py | 4 +- .../observability/stages.py | 10 ++++ .../opentelemetry/__init__.py | 2 + .../opentelemetry/console.py | 8 ++- .../policyengine_fastapi/opentelemetry/gcp.py | 38 ++++++++---- .../opentelemetry/instrumentor.py | 13 +++- .../opentelemetry/middleware.py | 12 ++++ .../src/policyengine_fastapi/ping/__init__.py | 21 ++++--- 20 files changed, 228 insertions(+), 65 deletions(-) diff --git a/libs/policyengine-fastapi/fixtures/ping/shared.py b/libs/policyengine-fastapi/fixtures/ping/shared.py index 525a46661..8f72ba1bd 100644 --- a/libs/policyengine-fastapi/fixtures/ping/shared.py +++ b/libs/policyengine-fastapi/fixtures/ping/shared.py @@ -9,11 +9,13 @@ @pytest.fixture def health_registry() -> HealthRegistry: + """Provide an empty registry; it reports healthy until a test registers a system.""" return HealthRegistry() @pytest.fixture def client(health_registry: HealthRegistry) -> TestClient: + """Return a client for a bare app serving ping routes from ``health_registry``.""" api = FastAPI() ping.include_all_routers(api, health_registry) return TestClient(api) diff --git a/libs/policyengine-fastapi/src/policyengine_fastapi/__init__.py b/libs/policyengine-fastapi/src/policyengine_fastapi/__init__.py index 04b09fae9..77965ebbf 100644 --- a/libs/policyengine-fastapi/src/policyengine_fastapi/__init__.py +++ b/libs/policyengine-fastapi/src/policyengine_fastapi/__init__.py @@ -1,6 +1,4 @@ -""" -Re-usable elements that support FastAPI and SQLModel indipendent of this specific app. -""" +"""Reusable FastAPI and SQLModel building blocks that are independent of any one app.""" from .observability import ( ObservabilityConfig as ObservabilityConfig, diff --git a/libs/policyengine-fastapi/src/policyengine_fastapi/auth/__init__.py b/libs/policyengine-fastapi/src/policyengine_fastapi/auth/__init__.py index 89c4db22d..b2c761138 100644 --- a/libs/policyengine-fastapi/src/policyengine_fastapi/auth/__init__.py +++ b/libs/policyengine-fastapi/src/policyengine_fastapi/auth/__init__.py @@ -1,3 +1,5 @@ +"""Bearer-token authentication dependencies for FastAPI routes.""" + from .jwt_decoder import JWTDecoder __all__ = ["JWTDecoder"] diff --git a/libs/policyengine-fastapi/src/policyengine_fastapi/auth/jwt_decoder.py b/libs/policyengine-fastapi/src/policyengine_fastapi/auth/jwt_decoder.py index a8ec9b929..a63578fb3 100644 --- a/libs/policyengine-fastapi/src/policyengine_fastapi/auth/jwt_decoder.py +++ b/libs/policyengine-fastapi/src/policyengine_fastapi/auth/jwt_decoder.py @@ -1,3 +1,5 @@ +"""FastAPI dependency that verifies RS256 bearer JWTs against an issuer's JWKS.""" + import logging import jwt @@ -11,12 +13,19 @@ class JWTDecoder: - """ - Parse and verify the JWT, verifying the correct issuer and audience and returning - the parsed token with all fields. + """Callable FastAPI dependency that verifies a bearer JWT and returns its claims. + + Tokens must be RS256-signed by a key from the issuer's JWKS and carry the + configured issuer and audience. A missing or invalid token raises HTTP 403, + or yields None when ``auto_error`` is False. """ def __init__(self, issuer: str, audience: str, auto_error: bool = True) -> None: + """Store the expected claims and create a JWKS client for the issuer. + + Keys are fetched from ``{issuer}.well-known/jwks.json``, so ``issuer`` + must end with a slash. + """ self.issuer = issuer self.audience = audience self.auto_error = auto_error @@ -29,6 +38,11 @@ def __call__( HTTPBearer(auto_error=False) ), ) -> dict[str, str] | None: + """Return the verified token's claims. + + A missing or invalid token raises HTTP 403, or returns None when + ``auto_error`` is False. + """ try: return self._validate(token) except HTTPException as err: diff --git a/libs/policyengine-fastapi/src/policyengine_fastapi/database.py b/libs/policyengine-fastapi/src/policyengine_fastapi/database.py index d1340a26e..6a8341fe7 100644 --- a/libs/policyengine-fastapi/src/policyengine_fastapi/database.py +++ b/libs/policyengine-fastapi/src/policyengine_fastapi/database.py @@ -1,3 +1,5 @@ +"""SQLModel engine and per-request session dependency helpers.""" + from collections.abc import Callable, Generator from typing import Any @@ -13,8 +15,10 @@ def create_sqlite_engine(filename: str | None = None) -> Engine: - """ - For desktop and testing, use sqlite to back sqlmodel + """Create a SQLite engine for desktop use and tests. + + Without ``filename`` the database is in memory; StaticPool keeps that single + in-memory connection shared across threads. """ sqlite_url = f"sqlite:///{filename}" if filename else "sqlite:///:memory:" connect_args = {"check_same_thread": False} @@ -25,9 +29,9 @@ def create_sqlite_engine(filename: str | None = None) -> Engine: def create_session_dep(engine: Engine) -> SessionGeneratorFactory: - """ - given an SQLModelEngine create a session dependency for use in - routers. + """Return a FastAPI dependency that yields one SQLModel session per request. + + The session is bound to ``engine`` and closed when the request finishes. """ def session_dep() -> Generator[Session]: diff --git a/libs/policyengine-fastapi/src/policyengine_fastapi/exit.py b/libs/policyengine-fastapi/src/policyengine_fastapi/exit.py index 5f7cde3c6..88e484d29 100644 --- a/libs/policyengine-fastapi/src/policyengine_fastapi/exit.py +++ b/libs/policyengine-fastapi/src/policyengine_fastapi/exit.py @@ -1,6 +1,8 @@ -""" -atexit does not work with uvicorn/fastapi for whatever reason and fastapi does not provide -a middleware-style way to register lifecycle events so I'm writing this +"""Process-wide registry of shutdown callbacks run at the end of the app lifespan. + +atexit handlers do not run reliably under uvicorn/FastAPI, and FastAPI has no +middleware-style way to register lifecycle events, so callbacks are registered +here and run from the app's lifespan instead. """ import logging @@ -14,10 +16,18 @@ class AppExit: + """Collect callbacks to run once, in registration order, when the app shuts down.""" + def __init__(self) -> None: + """Start with no registered callbacks.""" self.callbacks: list[Callable[[], None]] = [] def __call__(self, *args: object, **kwargs: object) -> Callable[[T], T]: + """Return a decorator that registers a function to run with these args at exit. + + The decorated function is returned unchanged. + """ + def accept(func: T) -> T: self.callbacks.append(lambda: func(*args, **kwargs)) return func @@ -32,6 +42,7 @@ def _exit(self) -> None: @asynccontextmanager async def lifespan(self) -> AsyncIterator[None]: + """Run every registered callback once the lifespan body exits without error.""" yield self._exit() diff --git a/libs/policyengine-fastapi/src/policyengine_fastapi/health/__init__.py b/libs/policyengine-fastapi/src/policyengine_fastapi/health/__init__.py index 66c9c64bd..4db20718f 100644 --- a/libs/policyengine-fastapi/src/policyengine_fastapi/health/__init__.py +++ b/libs/policyengine-fastapi/src/policyengine_fastapi/health/__init__.py @@ -1,21 +1,23 @@ +"""Health models and the registry behind the ``/ping/alive`` probe.""" + from collections.abc import Callable from pydantic import BaseModel -class ProbeStatus(BaseModel): +class ProbeStatus(BaseModel): # noqa: D101 # docstring would change the published OpenAPI description name: str healthy: bool message: str | None = None -class SystemStatus(BaseModel): +class SystemStatus(BaseModel): # noqa: D101 # docstring would change the published OpenAPI description name: str healthy: bool detail: list[ProbeStatus] -class HealthStatus(BaseModel): +class HealthStatus(BaseModel): # noqa: D101 # docstring would change the published OpenAPI description healthy: bool systems: list[SystemStatus] @@ -24,11 +26,15 @@ class HealthStatus(BaseModel): class HealthSystemReporter: + """Named group of health probes reported together as one system.""" + def __init__(self, name: str, probes: dict[str, HealthProbe]) -> None: + """Hold the system name and its probes, keyed by probe name.""" self.name = name self.probes = probes def report(self) -> SystemStatus: + """Run every probe; the system is healthy only if all of its probes are.""" detail = [self.probes[key]() for key in self.probes.keys()] first_unhealthy = next((d for d in detail if d.healthy is not True), None) return SystemStatus( @@ -37,8 +43,8 @@ def report(self) -> SystemStatus: class HealthRegistry: - """ - Mechanism for collecting and reporting sevice health. + """Registry of system reporters that together make up the service's health. + NOTE: any error will result in the service being restarted so only report real errors. """ @@ -46,12 +52,15 @@ class HealthRegistry: systems: dict[str, HealthSystemReporter] def __init__(self) -> None: + """Start with no systems, which reports as healthy.""" self.systems = {} def register(self, reporter: HealthSystemReporter) -> None: + """Add a system reporter, replacing any registered under the same name.""" self.systems[reporter.name] = reporter def report(self) -> HealthStatus: + """Report every system; the service is healthy only if all systems are.""" system_status = [self.systems[key].report() for key in self.systems.keys()] first_unhealthy = next( (s for s in system_status if s.healthy is not True), None diff --git a/libs/policyengine-fastapi/src/policyengine_fastapi/observability/__init__.py b/libs/policyengine-fastapi/src/policyengine_fastapi/observability/__init__.py index da530d37b..15f67934a 100644 --- a/libs/policyengine-fastapi/src/policyengine_fastapi/observability/__init__.py +++ b/libs/policyengine-fastapi/src/policyengine_fastapi/observability/__init__.py @@ -1,5 +1,4 @@ -""" -Shared observability contracts and utilities for the PolicyEngine simulation API services. +"""Shared observability contracts and utilities for the PolicyEngine simulation API services. This stays in `policyengine-fastapi` for now because only one active service needs it and the main value today is shared contract stability, not a separate diff --git a/libs/policyengine-fastapi/src/policyengine_fastapi/observability/config.py b/libs/policyengine-fastapi/src/policyengine_fastapi/observability/config.py index 465cf198a..9e1325481 100644 --- a/libs/policyengine-fastapi/src/policyengine_fastapi/observability/config.py +++ b/libs/policyengine-fastapi/src/policyengine_fastapi/observability/config.py @@ -1,3 +1,5 @@ +"""Settings for the shared observability emitters.""" + from __future__ import annotations from dataclasses import dataclass, field @@ -8,7 +10,6 @@ def parse_header_value_pairs(raw: str | None) -> dict[str, str]: """Parse OTLP headers from a comma or newline separated key=value string.""" - if raw is None: return {} @@ -33,6 +34,12 @@ def parse_header_value_pairs(raw: str | None) -> dict[str, str]: @dataclass(frozen=True) class ObservabilityConfig: + """Immutable observability settings. + + Covers whether emission is on, where OTLP telemetry and tracer artifacts go, + and when tracer output is captured. The defaults describe disabled emission. + """ + enabled: bool = False shadow_mode: bool = True service_name: str = "policyengine-observability" @@ -46,4 +53,5 @@ class ObservabilityConfig: @classmethod def disabled(cls, service_name: str = "policyengine-observability") -> Self: + """Return a config with emission off, setting only the service name.""" return cls(enabled=False, service_name=service_name) diff --git a/libs/policyengine-fastapi/src/policyengine_fastapi/observability/contracts.py b/libs/policyengine-fastapi/src/policyengine_fastapi/observability/contracts.py index 12c888709..7ac1ac103 100644 --- a/libs/policyengine-fastapi/src/policyengine_fastapi/observability/contracts.py +++ b/libs/policyengine-fastapi/src/policyengine_fastapi/observability/contracts.py @@ -1,3 +1,5 @@ +"""Pydantic payloads for run lifecycle events, tracer manifests and trace views.""" + from __future__ import annotations from datetime import datetime @@ -9,10 +11,18 @@ class ObservabilityModel(BaseModel): + """Observability payload base; rejects unknown fields and stores enum values.""" + model_config = ConfigDict(extra="forbid", use_enum_values=True) class CorrelatedRunFields(ObservabilityModel): + """Identifiers and version metadata that tie a record to one simulation run. + + Only ``observability_id`` is required; the rest are filled in as each + service learns them. + """ + observability_id: str submission_claim_id: str | None = None job_id: str | None = None @@ -31,6 +41,8 @@ class CorrelatedRunFields(ObservabilityModel): class SimulationLifecycleEvent(CorrelatedRunFields): + """One stage transition of a run, as reported by the service that saw it.""" + event_name: str stage: SimulationStage status: str @@ -41,6 +53,8 @@ class SimulationLifecycleEvent(CorrelatedRunFields): class TracerArtifactManifest(CorrelatedRunFields): + """Where a run's computation-tracer capture is stored, with its size and timings.""" + scenario: str capture_mode: TracerCaptureMode artifact_format: str @@ -55,6 +69,8 @@ class TracerArtifactManifest(CorrelatedRunFields): class SimulationRunSummary(CorrelatedRunFields): + """Overall status and end-to-end timing of one simulation run.""" + status: str requested_at: datetime | None = None returned_at: datetime | None = None @@ -62,6 +78,8 @@ class SimulationRunSummary(CorrelatedRunFields): class SimulationTimelineEntry(ObservabilityModel): + """Interval one service spent in one stage of a run.""" + stage: SimulationStage started_at: datetime ended_at: datetime @@ -70,6 +88,8 @@ class SimulationTimelineEntry(ObservabilityModel): class SimulationCompositeTraceResponse(ObservabilityModel): + """Combined view of a run: summary, stage timeline, spans, logs and tracer data.""" + run: SimulationRunSummary timeline: list[SimulationTimelineEntry] = Field(default_factory=list) spans: dict[str, Any] = Field(default_factory=dict) @@ -78,12 +98,16 @@ class SimulationCompositeTraceResponse(ObservabilityModel): class StageRuntimeSummary(ObservabilityModel): + """Mean, median and p95 stage duration in seconds; None where unobserved.""" + mean_seconds: float | None = None p50_seconds: float | None = None p95_seconds: float | None = None class VersionStageMetrics(ObservabilityModel): + """Per-stage runtime statistics for runs on one country package version.""" + country_package_version: str launch_source: str = "modal_version_registry" launched: bool = True @@ -92,6 +116,8 @@ class VersionStageMetrics(ObservabilityModel): class VersionStageMetricResponse(ObservabilityModel): + """Per-version stage runtime metrics for one country over a time window.""" + country: str window: dict[str, datetime] versions: list[VersionStageMetrics] = Field(default_factory=list) diff --git a/libs/policyengine-fastapi/src/policyengine_fastapi/observability/correlation.py b/libs/policyengine-fastapi/src/policyengine_fastapi/observability/correlation.py index 4282d0f3b..6bfdf0666 100644 --- a/libs/policyengine-fastapi/src/policyengine_fastapi/observability/correlation.py +++ b/libs/policyengine-fastapi/src/policyengine_fastapi/observability/correlation.py @@ -1,3 +1,5 @@ +"""Identifiers and hashes that correlate observability records across services.""" + from __future__ import annotations import json @@ -6,10 +8,16 @@ def generate_observability_id() -> str: + """Return a new random UUID4 string identifying one run across services.""" return str(uuid4()) def stable_config_hash(payload: object) -> str: + """Return a ``sha256:``-prefixed digest of the payload's canonical JSON. + + Keys are sorted and non-JSON values are stringified, so equal configurations + hash identically regardless of key order. + """ serialised = json.dumps( payload, sort_keys=True, diff --git a/libs/policyengine-fastapi/src/policyengine_fastapi/observability/emitters.py b/libs/policyengine-fastapi/src/policyengine_fastapi/observability/emitters.py index 05d4d7d3d..ffb05c79b 100644 --- a/libs/policyengine-fastapi/src/policyengine_fastapi/observability/emitters.py +++ b/libs/policyengine-fastapi/src/policyengine_fastapi/observability/emitters.py @@ -1,3 +1,5 @@ +"""Observability emitter interface and the no-op implementation used today.""" + from __future__ import annotations from contextlib import AbstractContextManager @@ -14,7 +16,10 @@ class NoOpSpan(AbstractContextManager["NoOpSpan"]): + """Span stand-in that accepts attributes and events and discards them.""" + def __enter__(self) -> Self: + """Return the span itself.""" return self def __exit__( @@ -23,49 +28,67 @@ def __exit__( exc_value: BaseException | None, traceback: TracebackType | None, ) -> None: - return None + """Do nothing, so any exception raised inside the span propagates.""" + return None # noqa: RET501 # explicit no-op body def set_attribute(self, key: str, value: object) -> None: - return None + """Discard the attribute.""" + return None # noqa: RET501 # explicit no-op body def add_event(self, name: str, attributes: Mapping[str, Any] | None = None) -> None: - return None + """Discard the event.""" + return None # noqa: RET501 # explicit no-op body class Observability(Protocol): + """Emitter interface for lifecycle events, metrics, tracer manifests and spans.""" + config: ObservabilityConfig - def emit_lifecycle_event(self, event: SimulationLifecycleEvent) -> None: ... + def emit_lifecycle_event(self, event: SimulationLifecycleEvent) -> None: + """Record one lifecycle stage event for a run.""" + ... def emit_counter( self, name: str, value: int = 1, attributes: Mapping[str, str] | None = None, - ) -> None: ... + ) -> None: + """Add ``value`` to the named counter.""" + ... def emit_histogram( self, name: str, value: float, attributes: Mapping[str, str] | None = None, - ) -> None: ... + ) -> None: + """Record one observation in the named histogram.""" + ... - def record_artifact_manifest(self, manifest: TracerArtifactManifest) -> None: ... + def record_artifact_manifest(self, manifest: TracerArtifactManifest) -> None: + """Record where a run's tracer capture was stored.""" + ... - def span( - self, name: str, attributes: Mapping[str, Any] | None = None - ) -> NoOpSpan: ... + def span(self, name: str, attributes: Mapping[str, Any] | None = None) -> NoOpSpan: + """Return a context manager covering the named span.""" + ... - def flush(self) -> None: ... + def flush(self) -> None: + """Send any buffered telemetry.""" + ... @dataclass class NoOpObservability: + """Observability that drops everything; returned even when emission is enabled.""" + config: ObservabilityConfig = field(default_factory=ObservabilityConfig.disabled) def emit_lifecycle_event(self, event: SimulationLifecycleEvent) -> None: - return None + """Discard the event.""" + return None # noqa: RET501 # explicit no-op body def emit_counter( self, @@ -73,7 +96,8 @@ def emit_counter( value: int = 1, attributes: Mapping[str, str] | None = None, ) -> None: - return None + """Discard the measurement.""" + return None # noqa: RET501 # explicit no-op body def emit_histogram( self, @@ -81,13 +105,17 @@ def emit_histogram( value: float, attributes: Mapping[str, str] | None = None, ) -> None: - return None + """Discard the measurement.""" + return None # noqa: RET501 # explicit no-op body def record_artifact_manifest(self, manifest: TracerArtifactManifest) -> None: - return None + """Discard the manifest.""" + return None # noqa: RET501 # explicit no-op body def span(self, name: str, attributes: Mapping[str, Any] | None = None) -> NoOpSpan: + """Return a span that records nothing.""" return NoOpSpan() def flush(self) -> None: - return None + """Do nothing; there is never anything buffered.""" + return None # noqa: RET501 # explicit no-op body diff --git a/libs/policyengine-fastapi/src/policyengine_fastapi/observability/provider.py b/libs/policyengine-fastapi/src/policyengine_fastapi/observability/provider.py index 69ac0fef5..b7993efd5 100644 --- a/libs/policyengine-fastapi/src/policyengine_fastapi/observability/provider.py +++ b/libs/policyengine-fastapi/src/policyengine_fastapi/observability/provider.py @@ -1,3 +1,5 @@ +"""Construction point for the observability emitter.""" + from __future__ import annotations from .config import ObservabilityConfig @@ -12,7 +14,6 @@ def build_observability( Commit 1 intentionally returns a no-op implementation even when enabled. Later commits can swap in a real backend here without changing callers. """ - if config is None: config = ObservabilityConfig.disabled() @@ -22,4 +23,5 @@ def build_observability( def get_observability( config: ObservabilityConfig | None = None, ) -> Observability: + """Return an emitter for ``config``, or a disabled one when it is omitted.""" return build_observability(config) diff --git a/libs/policyengine-fastapi/src/policyengine_fastapi/observability/stages.py b/libs/policyengine-fastapi/src/policyengine_fastapi/observability/stages.py index cb95cbb5a..7f51d9e99 100644 --- a/libs/policyengine-fastapi/src/policyengine_fastapi/observability/stages.py +++ b/libs/policyengine-fastapi/src/policyengine_fastapi/observability/stages.py @@ -1,7 +1,11 @@ +"""Simulation lifecycle stages and tracer capture modes.""" + from enum import Enum class SimulationStage(str, Enum): + """Lifecycle stage of a simulation run; values are dotted telemetry event names.""" + REQUEST_ACCEPTED = "request.accepted" REQUEST_VALIDATED = "request.validated" JOB_SETUP = "job.setup" @@ -23,6 +27,12 @@ class SimulationStage(str, Enum): class TracerCaptureMode(str, Enum): + """Which runs have their computation tracer captured. + + Ranges from none (DISABLED) through failed runs, slow runs and a sample, to + every run (ALWAYS). + """ + DISABLED = "disabled" FAILURES = "failures" THRESHOLD = "threshold" diff --git a/libs/policyengine-fastapi/src/policyengine_fastapi/opentelemetry/__init__.py b/libs/policyengine-fastapi/src/policyengine_fastapi/opentelemetry/__init__.py index 7a0d28bfb..198d35219 100644 --- a/libs/policyengine-fastapi/src/policyengine_fastapi/opentelemetry/__init__.py +++ b/libs/policyengine-fastapi/src/policyengine_fastapi/opentelemetry/__init__.py @@ -1,3 +1,5 @@ +"""OpenTelemetry exporters, GCP JSON logging and per-route metrics for FastAPI apps.""" + from .console import export_ot_to_console as export_ot_to_console from .gcp import ( GCPLoggingInstrumentor as GCPLoggingInstrumentor, diff --git a/libs/policyengine-fastapi/src/policyengine_fastapi/opentelemetry/console.py b/libs/policyengine-fastapi/src/policyengine_fastapi/opentelemetry/console.py index 31b21b294..86a26712e 100644 --- a/libs/policyengine-fastapi/src/policyengine_fastapi/opentelemetry/console.py +++ b/libs/policyengine-fastapi/src/policyengine_fastapi/opentelemetry/console.py @@ -1,3 +1,5 @@ +"""Console OpenTelemetry exporters for local debugging.""" + from opentelemetry import metrics, trace from opentelemetry.sdk.metrics import MeterProvider from opentelemetry.sdk.metrics.export import ( @@ -20,8 +22,10 @@ # 1. to include python logs and # 2. export to console (for demo purposes) def export_ot_to_console(resource: Resource) -> None: - """ - configure opentelemetry to dump messages to console for debugging on desktop + """Export OpenTelemetry traces and metrics to the console, for desktop debugging. + + Installs global tracer and meter providers for ``resource`` and registers an + ``exit`` callback that flushes both at shutdown. """ traceProvider = TracerProvider(resource=resource) processor = BatchSpanProcessor(ConsoleSpanExporter()) diff --git a/libs/policyengine-fastapi/src/policyengine_fastapi/opentelemetry/gcp.py b/libs/policyengine-fastapi/src/policyengine_fastapi/opentelemetry/gcp.py index 91f40ed8f..147433112 100644 --- a/libs/policyengine-fastapi/src/policyengine_fastapi/opentelemetry/gcp.py +++ b/libs/policyengine-fastapi/src/policyengine_fastapi/opentelemetry/gcp.py @@ -1,3 +1,5 @@ +"""Google Cloud exporters for OpenTelemetry, and JSON logs shaped for Cloud Logging.""" + import logging import os @@ -23,11 +25,15 @@ class AddGcpProjectToTraceFilter(logging.Filter): + """Logging filter that qualifies trace IDs as ``projects//traces/``.""" + def __init__(self, project: str) -> None: + """Remember the GCP project used to qualify trace IDs.""" super().__init__() self.project = project def filter(self, record: logging.LogRecord) -> bool | logging.LogRecord: + """Qualify the record's ``otelTraceID``, if set, then apply the base filter.""" if not hasattr(record, "otelTraceID"): return super().filter(record) record.otelTraceID = f"projects/{self.project}/traces/{record.otelTraceID}" # type: ignore @@ -58,15 +64,26 @@ def _get_project_id_from_metadata() -> str | None: class GCPLoggingInstrumentor: - def __init__(self, project_id: str | None = None) -> None: - self.project_id = project_id or _get_project_id() + """Logging setup that emits JSON logs in the shape Google Cloud Logging expects. + Configures the standard OpenTelemetry logging instrumentor and renames level, + timestamp and the trace and span IDs to the Cloud Logging fields, so log + lines link to their traces. """ - Configures the standard opentelemetry logging instrumentor to generate - json logs as per gcp expectations. - """ + + def __init__(self, project_id: str | None = None) -> None: + """Use ``project_id``, else look the project up. + + The lookup asks the GCE metadata server, then ``GOOGLE_CLOUD_PROJECT``. + """ + self.project_id = project_id or _get_project_id() def instrument(self) -> None: + """Send JSON logs to stderr and add OpenTelemetry trace IDs to log records. + + Uses ``logging.basicConfig``, so the handler is not installed if the root + logger already has handlers. + """ from opentelemetry.instrumentation.logging import LoggingInstrumentor from pythonjsonlogger.json import JsonFormatter @@ -94,16 +111,17 @@ def instrument(self) -> None: def export_ot_to_gcp(resource: Resource) -> None: + """Export traces and metrics straight to Cloud Trace and Cloud Monitoring. + + Meant for code running in Google Cloud. Installs global tracer and meter + providers for ``resource`` and registers an ``exit`` callback that flushes + both at shutdown. + """ from opentelemetry.exporter.cloud_monitoring import ( CloudMonitoringMetricsExporter, ) from opentelemetry.exporter.cloud_trace import CloudTraceSpanExporter - """ - configure opentelemetry to directly export to gcp cloudtrace/metrics - useful when running in the google cloud - """ - traceProvider = TracerProvider(resource=resource) processor = BatchSpanProcessor(CloudTraceSpanExporter()) traceProvider.add_span_processor(processor) diff --git a/libs/policyengine-fastapi/src/policyengine_fastapi/opentelemetry/instrumentor.py b/libs/policyengine-fastapi/src/policyengine_fastapi/opentelemetry/instrumentor.py index aab1ce0aa..7e8678516 100644 --- a/libs/policyengine-fastapi/src/policyengine_fastapi/opentelemetry/instrumentor.py +++ b/libs/policyengine-fastapi/src/policyengine_fastapi/opentelemetry/instrumentor.py @@ -1,3 +1,5 @@ +"""FastAPI instrumentation that adds per-route OpenTelemetry metrics.""" + from fastapi import FastAPI from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor @@ -5,12 +7,17 @@ class FastAPIEnhancedInstrumenter: - """ - Enhances the default FastAPIInstrumentor to generate per operation - metrics instead of global for the whole api. + """Enhanced FastAPIInstrumentor that also records per-operation metrics. + + The default instrumentor only produces metrics for the whole API. """ def instrument(self, app: FastAPI) -> None: + """Instrument ``app`` and add the per-route metrics middleware. + + Call this after every route is registered: routes added later get no + per-route metrics. + """ FastAPIInstrumentor.instrument_app(app) middleware = Middleware(app) app.middleware("http")(middleware) diff --git a/libs/policyengine-fastapi/src/policyengine_fastapi/opentelemetry/middleware.py b/libs/policyengine-fastapi/src/policyengine_fastapi/opentelemetry/middleware.py index f741b51f9..4d5ac0466 100644 --- a/libs/policyengine-fastapi/src/policyengine_fastapi/opentelemetry/middleware.py +++ b/libs/policyengine-fastapi/src/policyengine_fastapi/opentelemetry/middleware.py @@ -1,3 +1,5 @@ +"""HTTP middleware recording a request counter and latency histogram per route.""" + from dataclasses import dataclass from time import time @@ -10,13 +12,22 @@ @dataclass class Route: + """A FastAPI route with its request counter and latency histogram.""" + route: APIRoute execution: metrics.Counter latency: metrics.Histogram class Middleware: + """Count requests and record latency, in seconds, per matched API route. + + Metrics are named ``_count`` and ``_latency``. + Requests that match no route are passed through unmeasured. + """ + def __init__(self, app: FastAPI) -> None: + """Create metrics for every API route registered on ``app`` so far.""" self.app = app self.meter = metrics.get_meter("policyengine") @@ -35,6 +46,7 @@ def _create_route(self, route: APIRoute) -> Route: async def __call__( self, request: Request, call_next: RequestResponseEndpoint ) -> Response: + """Pass the request on; count and time it under its first matching route.""" route: Route | None = None for r in self.routes: if r.route.matches(request.scope)[0] != Match.NONE: diff --git a/libs/policyengine-fastapi/src/policyengine_fastapi/ping/__init__.py b/libs/policyengine-fastapi/src/policyengine_fastapi/ping/__init__.py index ce927959e..9cc0c622d 100644 --- a/libs/policyengine-fastapi/src/policyengine_fastapi/ping/__init__.py +++ b/libs/policyengine-fastapi/src/policyengine_fastapi/ping/__init__.py @@ -1,3 +1,5 @@ +"""Unauthenticated ``/ping`` liveness and health-probe routes.""" + from fastapi import APIRouter, FastAPI from fastapi.responses import JSONResponse from pydantic import BaseModel @@ -5,38 +7,34 @@ from policyengine_fastapi.health import HealthRegistry, HealthStatus -class PingRequest(BaseModel): +class PingRequest(BaseModel): # noqa: D101 # docstring would change the published OpenAPI description value: int -class PingResponse(BaseModel): +class PingResponse(BaseModel): # noqa: D101 # docstring would change the published OpenAPI description incremented: int def create_router(health_registry: HealthRegistry) -> APIRouter: + """Build the ping routes; ``/ping/alive`` is 503 while the registry is unhealthy.""" router = APIRouter() @router.post("/ping") async def ping(request: PingRequest) -> PingResponse: - """ - Verify the api is able to recieve and process unauthenticated requests. - """ + """Verify the api is able to recieve and process unauthenticated requests.""" return PingResponse(incremented=request.value + 1) @router.get("/ping/started", response_model=str) async def started() -> str: - """ - Verify the api is running. - """ + """Verify the api is running.""" return "alive" @router.get("/ping/alive", response_model=HealthStatus) async def alive() -> JSONResponse: - """ - Check if the service is healthy. This will always return + """Check if the service is healthy. This will always return data, but will return a 503 response code if the service is unhealthy. - """ + """ # noqa: D205 # docstring would change the published OpenAPI description content = health_registry.report() return JSONResponse( content=content.model_dump(exclude_none=True), @@ -47,4 +45,5 @@ async def alive() -> JSONResponse: def include_all_routers(api: FastAPI, health_registry: HealthRegistry) -> None: + """Mount the ping routes on ``api``, reporting health from ``health_registry``.""" api.include_router(create_router(health_registry)) From c776e967eecd9cc5f693564d4037e7c2f7a86e8d Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 13 Oct 2026 16:56:47 +0400 Subject: [PATCH 03/18] Document the simulation observability library for ruff D Co-Authored-By: Claude Opus 5.5 --- .../__init__.py | 3 +-- .../policyengine_simulation_observability/errors.py | 8 +++++--- .../identifiers.py | 3 --- .../observability.py | 12 ++++++++++-- .../policyengine_simulation_observability/stages.py | 7 ++++++- .../telemetry.py | 11 +++++++++-- 6 files changed, 31 insertions(+), 13 deletions(-) diff --git a/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/__init__.py b/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/__init__.py index f65a80295..debd9e137 100644 --- a/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/__init__.py +++ b/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/__init__.py @@ -1,2 +1 @@ -"""Shared observability plumbing for the simulation entrypoint, gateway, -and executor.""" +"""Shared observability plumbing for the simulation entrypoint, gateway and executor.""" diff --git a/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/errors.py b/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/errors.py index e4bf391a3..f90428554 100644 --- a/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/errors.py +++ b/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/errors.py @@ -36,6 +36,7 @@ def make_correlation_id() -> str: + """Return a random 32-character hex ID tying a redacted error to its log line.""" return uuid.uuid4().hex @@ -46,15 +47,16 @@ def log_and_redact_exception( scope: str, context: dict[str, Any] | None = None, ) -> str: - """Record ``exc`` with full context server-side and return a redacted - caller-safe message paired with a correlation id. + """Record ``exc`` server-side and return a redacted message with a correlation id. + + The full exception and ``context`` always reach the standard library logger + and, best effort, ``runtime``; a failure in ``runtime`` is swallowed. The public return string has the shape ``"Simulation failed (correlation_id=)"`` so clients can paste it into support tickets without needing to parse JSON. The same correlation id is attached to the server-side log line so operators can jump between the two. """ - correlation_id = make_correlation_id() payload = { "correlation_id": correlation_id, diff --git a/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/identifiers.py b/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/identifiers.py index e60ded480..896c1428a 100644 --- a/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/identifiers.py +++ b/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/identifiers.py @@ -9,13 +9,11 @@ def generate_observability_id() -> str: """Create a diagnostic identifier with no application semantics.""" - return str(uuid4()) def resolve_observability_id(value: object) -> str: """Use a valid caller value or create a new diagnostic identifier.""" - return normalize_observability_id(value) or generate_observability_id() @@ -25,7 +23,6 @@ def normalize_observability_id(value: object) -> str | None: Observability metadata is best effort. Callers must never reject or alter a simulation because this value is absent or malformed. """ - if not isinstance(value, str): return None try: diff --git a/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/observability.py b/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/observability.py index bdc96013e..37360d368 100644 --- a/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/observability.py +++ b/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/observability.py @@ -1,3 +1,5 @@ +"""Construction of the observability runtime for each simulation service.""" + from __future__ import annotations import os @@ -41,7 +43,6 @@ def modal_image_environment() -> dict[str, str]: """Return deployment-provided settings to bake into a Modal image.""" - configured = { name: value for name in MODAL_IMAGE_ENV_NAMES @@ -65,7 +66,6 @@ def build_runtime( service_version: str | None = None, ) -> ObservabilityRuntime: """Create one explicitly owned simulation observability runtime.""" - platform_instance_id = ( os.getenv("MODAL_TASK_ID") if platform == "modal" else os.getenv("K_REVISION") ) @@ -129,6 +129,10 @@ def init_simulation_observability( environment: str, service_version: str | None = None, ) -> ObservabilityRuntime: + """Build the service's observability runtime and instrument ``app`` with it. + + The caller owns the returned runtime and must shut it down on exit. + """ runtime = build_runtime( service_name=service_name, service_role=service_role, @@ -147,6 +151,10 @@ def init_process_observability( environment: str, service_version: str | None = None, ) -> ObservabilityRuntime: + """Build the observability runtime for a process with no FastAPI app. + + Used by Modal workers; the caller owns the runtime and must shut it down. + """ return build_runtime( service_name=service_name, service_role=service_role, 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 455854134..eac473f97 100644 --- a/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/stages.py +++ b/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/stages.py @@ -17,6 +17,8 @@ class RunConfiguration(StrEnum): + """Kind of simulation run; each has a fixed stage plan in ``RUN_STAGE_REGISTRY``.""" + ANNUAL_IMPACT = "annual_impact" BUDGET_WINDOW = "budget_window" SEGMENTED_NATIONAL = "segmented_national" @@ -26,6 +28,8 @@ class RunConfiguration(StrEnum): class Stage(StrEnum): + """Every traced stage of any run; each value is the span name.""" + ANNUAL_EXECUTION = "run_simulation" BUDGET_WINDOW_EXECUTION = "run_budget_window_batch" SEGMENTED_NATIONAL_EXECUTION = "run_simulation_segment" @@ -104,12 +108,13 @@ class Stage(StrEnum): @dataclass(frozen=True) class StagePlan: + """Stages one run configuration may emit, in their expected order.""" + configuration: RunConfiguration stages: tuple[Stage, ...] def name(self, stage: Stage) -> str: """Return a configured stage name, rejecting unregistered use.""" - if stage not in self.stages: raise ValueError( f"{stage.value!r} is not registered for {self.configuration.value!r}" diff --git a/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/telemetry.py b/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/telemetry.py index 26015dc0a..409f58ad7 100644 --- a/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/telemetry.py +++ b/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/telemetry.py @@ -24,6 +24,8 @@ class ObservabilityContext(BaseModel): + """Trace and correlation context transported separately from the payload.""" + traceparent: str | None = None tracestate: str | None = None captured_at: datetime @@ -37,6 +39,7 @@ class ObservabilityContext(BaseModel): @field_validator("observability_id") @classmethod def validate_observability_id(cls, value: str | None) -> str | None: + """Return the ID as a canonical UUID; raise ValueError if it is not one.""" if value is None: return None normalized = normalize_observability_id(value) @@ -66,7 +69,6 @@ def discard_noncanonical_context_fields(cls, value: object) -> object: This keeps the receiver compatible with the deployed API during the rollout. The HTTP header remains the only source of observability_id. """ - if not isinstance(value, dict): return value normalized = dict(value) @@ -86,7 +88,6 @@ def normalize_observability_context( context: object, ) -> dict[str, str] | None: """Validate context transported separately from an application payload.""" - if not isinstance(context, dict): return None try: @@ -102,6 +103,12 @@ def normalize_observability_context( def split_internal_payload( params: dict[str, Any], ) -> tuple[dict[str, Any], TelemetryEnvelope | None, dict[str, Any] | None]: + """Split worker params into simulation params, telemetry envelope and metadata. + + The reserved ``_telemetry`` and ``_metadata`` keys are removed from a copy of + ``params``. An invalid envelope becomes None rather than failing the run, and + non-dict metadata is dropped. + """ simulation_params = dict(params) raw_telemetry = simulation_params.pop("_telemetry", None) raw_metadata = simulation_params.pop("_metadata", None) From 8fba879a388ebf473f2dc87df678f64c39b584f5 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 13 Oct 2026 16:56:47 +0400 Subject: [PATCH 04/18] Document the Stage 12 persistence library for ruff D Co-Authored-By: Claude Opus 5.5 --- .../policyengine_stage12_persistence/store.py | 46 ++++++++++++++++++- .../validation.py | 1 - .../validators.py | 37 +++++++++++++++ 3 files changed, 81 insertions(+), 3 deletions(-) diff --git a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/store.py b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/store.py index 031a81330..9936cb36f 100644 --- a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/store.py +++ b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/store.py @@ -103,7 +103,6 @@ def normalized_database_url(database_url: str) -> URL: """Select psycopg 3 without changing the stored secret's URL format.""" - if not database_url or database_url != database_url.strip(): raise ValueError("Stage 12 database URL is missing or malformed") try: @@ -119,7 +118,6 @@ def normalized_database_url(database_url: str) -> URL: def create_stage12_engine(database_url: str) -> Engine: """Build the DML-only runtime engine; this never creates database objects.""" - return create_engine( normalized_database_url(database_url), # Coordinators and workers spend most of their lifetime calculating @@ -194,6 +192,10 @@ def __init__( engine: Engine | None = None, open_connection: ConnectionOpener = open_connection_once, ) -> None: + """Use ``engine`` if given, else build one from ``database_url``. + + Raises ValueError when neither is given or the URL is invalid. + """ if engine is None: if database_url is None: raise ValueError("Stage 12 database URL is required") @@ -222,6 +224,11 @@ def create_or_resolve_report( self, record: ComparisonReportRecord, ) -> ComparisonReportPersistenceResult: + """Insert the report, or return the stored one with the same identity. + + ``created`` on the result says which happened. A stored report must match + ``record``'s provenance, else ValueError. + """ values = _record_values(record) statement = ( insert(comparison_reports) @@ -258,6 +265,7 @@ def create_or_resolve_report( return ComparisonReportPersistenceResult(record=resolved, created=False) def get_report(self, evaluation_id: UUID) -> ComparisonReportRecord: + """Return the stored report; raise LookupError if there is none.""" with self._connection("get_report") as connection: row = ( connection.execute( @@ -276,6 +284,7 @@ def list_simulations( self, evaluation_id: UUID, ) -> tuple[ComparisonSimulationRecord, ...]: + """Return the report's simulations ordered by role, then execution ID.""" statement = ( select(comparison_simulations) .where(comparison_simulations.c.evaluation_id == evaluation_id) @@ -292,6 +301,13 @@ def replace_report( self, record: ComparisonReportRecord, ) -> ComparisonReportRecord: + """Write ``record``'s lifecycle and aggregation fields over the stored report. + + The row is locked for the update. Raises LookupError if the report is missing, + and ValueError for an identity or result-comparison change or a disallowed + status move. A succeeded report is returned as stored when ``record`` replays + it exactly. + """ values = _record_values(record) with self._transaction("replace_report") as connection: row = ( @@ -334,6 +350,13 @@ def replace_report_result_comparison( self, record: ComparisonReportRecord, ) -> ComparisonReportRecord: + """Write ``record``'s result-comparison fields over the stored report. + + The row is locked for the update and every other field must be unchanged. + Once the comparison is MATCHED or DIFFERENT the stored report is returned + as-is, provided ``record`` repeats its result. Raises LookupError if the + report is missing and ValueError for any other change or disallowed move. + """ values = _record_values(record) with self._transaction("replace_report_result_comparison") as connection: row = ( @@ -379,6 +402,11 @@ def create_or_resolve_simulation( self, record: ComparisonSimulationRecord, ) -> ComparisonSimulationPersistenceResult: + """Insert the simulation, or return the stored one with the same identity. + + The parent report must exist (else LookupError) and share ``record``'s + provenance; a stored simulation must match ``record`` (else ValueError). + """ values = _record_values(record) statement = ( insert(comparison_simulations) @@ -435,6 +463,7 @@ def get_simulation( self, simulation_execution_id: UUID, ) -> ComparisonSimulationRecord: + """Return the stored simulation; raise LookupError if there is none.""" with self._connection("get_simulation") as connection: row = ( connection.execute( @@ -456,6 +485,12 @@ def replace_simulation( self, record: ComparisonSimulationRecord, ) -> ComparisonSimulationRecord: + """Write ``record``'s mutable fields over the stored simulation. + + The row is locked for the update. Raises LookupError if the simulation is + missing, and ValueError for an identity change or a disallowed status move. A + succeeded simulation is returned as stored when ``record`` replays it exactly. + """ values = _record_values(record) with self._transaction("replace_simulation") as connection: row = ( @@ -506,6 +541,13 @@ def attach_simulation_invocation( modal_invocation_id: str, updated_at: datetime, ) -> ComparisonSimulationRecord: + """Replace the placeholder invocation ID with the real Modal invocation ID. + + The update applies only while ``expected_placeholder`` is still stored, so + repeating an attach that already succeeded returns the stored record. Raises + LookupError if the simulation is missing and ValueError if a different + invocation ID is already attached. + """ statement = ( update(comparison_simulations) .where( diff --git a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/validation.py b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/validation.py index 2850f1ac4..9f0e36ed4 100644 --- a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/validation.py +++ b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/validation.py @@ -365,7 +365,6 @@ def validate_runtime_database( environment: str, ) -> None: """Inspect the live schema and roll back a complete typed DML canary.""" - if environment not in {"staging", "production"}: raise Stage12RuntimeAccessError("environment must be staging or production") try: diff --git a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/validators.py b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/validators.py index 6f6179e93..f7bcdf264 100644 --- a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/validators.py +++ b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/validators.py @@ -118,6 +118,7 @@ def require_report_identity( existing: ComparisonReportRecord, candidate: ComparisonReportRecord, ) -> None: + """Raise ValueError if ``candidate`` changes a report's fixed fields.""" _require_equal_fields(existing, candidate, REPORT_IMMUTABLE_FIELDS) @@ -125,6 +126,10 @@ def require_report_conflict_matches( existing: ComparisonReportRecord, candidate: ComparisonReportRecord, ) -> None: + """Raise ValueError unless the report found on a conflict matches ``candidate``. + + A mismatch means two different runs claimed the same report identity. + """ _require_equal_fields(existing, candidate, REPORT_CONFLICT_FIELDS) @@ -132,6 +137,7 @@ def require_simulation_identity( existing: ComparisonSimulationRecord, candidate: ComparisonSimulationRecord, ) -> None: + """Raise ValueError if ``candidate`` changes a simulation's fixed fields.""" _require_equal_fields(existing, candidate, SIMULATION_IMMUTABLE_FIELDS) @@ -139,6 +145,7 @@ def require_simulation_conflict_matches( existing: ComparisonSimulationRecord, candidate: ComparisonSimulationRecord, ) -> None: + """Raise ValueError unless the row found on a conflict matches ``candidate``.""" _require_equal_fields(existing, candidate, SIMULATION_CONFLICT_FIELDS) @@ -146,6 +153,10 @@ def require_simulation_parent_matches( parent: ComparisonReportRecord, child: ComparisonSimulationRecord, ) -> None: + """Raise ValueError unless ``child`` belongs to ``parent`` and shares provenance. + + Versions, the Modal application, creation time and retention must all match. + """ if child.evaluation_id != parent.evaluation_id: raise ValueError("comparison child names a different parent evaluation_id") for field_name in ( @@ -164,6 +175,10 @@ def require_successful_report_replay( existing: ComparisonReportRecord, candidate: ComparisonReportRecord, ) -> None: + """Raise ValueError unless a write to a succeeded report repeats it exactly. + + Only ``updated_at`` and the result-comparison fields may differ. + """ _require_equal_fields(existing, candidate, REPORT_SUCCEEDED_REPLAY_FIELDS) @@ -171,6 +186,10 @@ def require_successful_simulation_replay( existing: ComparisonSimulationRecord, candidate: ComparisonSimulationRecord, ) -> None: + """Raise ValueError unless a write to a succeeded simulation repeats it exactly. + + Only ``updated_at`` may differ. + """ _require_equal_fields(existing, candidate, SIMULATION_SUCCEEDED_REPLAY_FIELDS) @@ -178,6 +197,10 @@ def require_comparison_update_only( existing: ComparisonReportRecord, candidate: ComparisonReportRecord, ) -> None: + """Raise ValueError if ``candidate`` changes more than the result-comparison fields. + + ``updated_at`` may also change. + """ fields = tuple( field_name for field_name in ComparisonReportRecord.model_fields @@ -191,6 +214,7 @@ def require_result_comparison_unchanged( existing: ComparisonReportRecord, candidate: ComparisonReportRecord, ) -> None: + """Raise ValueError if ``candidate`` changes any result-comparison field.""" _require_equal_fields(existing, candidate, REPORT_RESULT_COMPARISON_FIELDS) @@ -287,6 +311,10 @@ def require_lifecycle_transition( current: ComparisonRunLifecycleStatus, candidate: ComparisonRunLifecycleStatus, ) -> None: + """Raise ValueError unless the lifecycle status may make this move. + + Repeating the current status is allowed; SUCCEEDED and SKIPPED are terminal. + """ if candidate not in ALLOWED_LIFECYCLE_TRANSITIONS[current]: raise ValueError( "comparison-run lifecycle cannot move " @@ -298,6 +326,11 @@ def require_aggregation_transition( current: ComparisonRunAggregationStatus, candidate: ComparisonRunAggregationStatus, ) -> None: + """Raise ValueError unless the aggregation status may make this move. + + Repeating the current status is allowed; SUCCEEDED is terminal and FAILED may + return to RUNNING for a retry. + """ if candidate not in ALLOWED_AGGREGATION_TRANSITIONS[current]: raise ValueError( "comparison-run aggregation cannot move " @@ -309,6 +342,10 @@ def require_result_comparison_transition( current: ResultComparisonStatus, candidate: ResultComparisonStatus, ) -> None: + """Raise ValueError unless the result-comparison status may make this move. + + NOT_REQUESTED never changes, and MATCHED, DIFFERENT and FAILED are terminal. + """ if candidate not in ALLOWED_RESULT_COMPARISON_TRANSITIONS[current]: raise ValueError( f"result comparison cannot move from {current.value} to {candidate.value}" From 966e4d2b34792848e8aa29a5898f0009746b9799 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 13 Oct 2026 16:56:47 +0400 Subject: [PATCH 05/18] Document the simulation entrypoint for ruff D The entry routes are nested in create_app, so their docstrings (the OpenAPI descriptions) are untouched and the generated spec is unchanged. Co-Authored-By: Claude Opus 5.5 --- .../authenticated_tests/test_deployed_auth.py | 8 ++++ .../src/policyengine_simulation_entry/app.py | 24 ++++++++-- .../src/policyengine_simulation_entry/auth.py | 13 +++++- .../policyengine_simulation_entry/backend.py | 43 ++++++++++++++++-- .../policyengine_simulation_entry/config.py | 13 ++++++ .../generate_openapi.py | 5 +++ .../stage12_adapter.py | 10 ++++- .../stage12_backend.py | 45 ++++++++++++++++--- .../stage12_infrastructure.py | 5 +++ .../tests/test_deploy_workflow_scripts.py | 2 +- .../tests/test_workflow_scripts.py | 6 ++- .../tests/workflow_script_support.py | 4 +- 12 files changed, 157 insertions(+), 21 deletions(-) diff --git a/projects/policyengine-simulation-entry/authenticated_tests/test_deployed_auth.py b/projects/policyengine-simulation-entry/authenticated_tests/test_deployed_auth.py index 1d9609e0d..6c18f94f0 100644 --- a/projects/policyengine-simulation-entry/authenticated_tests/test_deployed_auth.py +++ b/projects/policyengine-simulation-entry/authenticated_tests/test_deployed_auth.py @@ -24,11 +24,16 @@ def _required(name: str) -> str: @pytest.fixture(scope="session") def base_url() -> str: + """Return the deployed candidate's base URL, without a trailing slash.""" return _required("SIMULATION_ENTRYPOINT_TEST_BASE_URL").rstrip("/") @pytest.fixture(scope="session") def access_token() -> str: + """Fetch a client-credentials token for the candidate's caller audience. + + Raise RuntimeError if a required variable is unset or Auth0 returns no token. + """ issuer = _required("SIMULATION_ENTRYPOINT_TEST_AUTH_ISSUER").rstrip("/") response = httpx.post( f"{issuer}/oauth/token", @@ -48,6 +53,7 @@ def access_token() -> str: def test_protected_route_rejects_missing_token(base_url: str) -> None: + """Check the deployed comparison route returns 403 without a bearer token.""" response = httpx.post( f"{base_url}{PROTECTED_PATH}", json=PROBE_PAYLOAD, @@ -58,6 +64,7 @@ def test_protected_route_rejects_missing_token(base_url: str) -> None: def test_protected_route_rejects_invalid_token(base_url: str) -> None: + """Check the deployed comparison route returns 403 for a malformed token.""" response = httpx.post( f"{base_url}{PROTECTED_PATH}", json=PROBE_PAYLOAD, @@ -72,6 +79,7 @@ def test_valid_token_reaches_authenticated_old_gateway( base_url: str, access_token: str, ) -> None: + """Check a valid token passes both auth hops to reach the old gateway.""" response = httpx.post( f"{base_url}{PROTECTED_PATH}", json=PROBE_PAYLOAD, diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/app.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/app.py index 506905399..dd95458c8 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/app.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/app.py @@ -84,6 +84,13 @@ class ComparisonBackend(Protocol): + """Stage 12 comparison dispatch and polling seam used by the entry routes. + + The real implementation is ``Stage12ComparisonBackend``; tests inject fakes. + The routes bound every call with a timeout and never let a failure here + alter the production v1 response. + """ + async def dispatch_after_production( self, *, @@ -91,7 +98,13 @@ async def dispatch_after_production( production_response: bytes, request_id: str, observability_id: str, - ) -> None: ... + ) -> None: + """Start a shadow comparison run for a production job the gateway accepted. + + ``production_response`` is the raw gateway response body; its ``job_id`` + identifies the incumbent run being compared. + """ + ... async def submit_temporary_report( self, @@ -99,12 +112,16 @@ async def submit_temporary_report( request_payload: dict[str, Any], request_id: str, observability_id: str, - ) -> ComparisonReportRecord: ... + ) -> ComparisonReportRecord: + """Submit an operator-requested report and return its acknowledged record.""" + ... async def get_temporary_report( self, evaluation_id: UUID, - ) -> tuple[ComparisonReportRecord, tuple[ComparisonSimulationRecord, ...]]: ... + ) -> tuple[ComparisonReportRecord, tuple[ComparisonSimulationRecord, ...]]: + """Return a report and its simulations; raise LookupError if it is unknown.""" + ... def _model_json(model: BaseModel) -> JsonObject: @@ -155,7 +172,6 @@ def create_app( comparison_backend: ComparisonBackend | None = None, ) -> FastAPI: """Build the app with injectable auth/backend seams for hermetic tests.""" - runtime_settings = settings or Settings.from_env() @asynccontextmanager diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/auth.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/auth.py index ca6f143f6..6344ad00d 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/auth.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/auth.py @@ -25,6 +25,7 @@ class JWTDecoder: """Validate Auth0 RS256 access tokens without database-layer dependencies.""" def __init__(self, issuer: str, audience: str) -> None: + """Bind the issuer and audience and point a JWKS client at the issuer.""" self.issuer = issuer self.audience = audience self.jwks_client = jwt.PyJWKClient(f"{issuer}.well-known/jwks.json") @@ -34,6 +35,11 @@ def __call__( token: HTTPAuthorizationCredentials | None, runtime: ObservabilityRuntime, ) -> CallerIdentity: + """Return the caller identity from a verified token. + + Raise HTTP 403 for a missing or invalid token, recording the rejection + reason as an observability event. + """ if token is None: runtime.event( "simulation_entry_auth_rejected", @@ -76,6 +82,7 @@ class CallerAuthenticator: """FastAPI dependency that preserves the gateway's JWT contract.""" def __init__(self, settings: Settings, runtime: ObservabilityRuntime) -> None: + """Keep the settings and runtime used to verify each request.""" self.settings = settings self.runtime = runtime @@ -83,6 +90,11 @@ def __call__( self, token: HTTPAuthorizationCredentials | None = Depends(_bearer), ) -> CallerIdentity | None: + """Return the verified caller, or None when auth is not required. + + Decoders are cached per issuer and audience, so the JWKS client is + reused across requests. + """ if not self.settings.auth_required: return None return _decoder( @@ -93,5 +105,4 @@ def __call__( def reset_decoder_cache() -> None: """Clear the decoder cache for tests and credential rotations.""" - _decoder.cache_clear() diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/backend.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/backend.py index a74f2ef70..f6a0edc91 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/backend.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/backend.py @@ -49,17 +49,30 @@ class BackendAuthenticationError(BackendUnavailable): @dataclass(frozen=True) class BackendResponse: + """A backend reply with only the safe-to-forward response headers kept.""" + status_code: int content: bytes headers: dict[str, str] class SimulationBackend(Protocol): - async def start(self) -> None: ... + """Backend the entry routes forward simulation requests to. + + The app starts it in the lifespan and closes it on shutdown. + """ - async def close(self) -> None: ... + async def start(self) -> None: + """Open connections; called once when the app starts.""" + ... + + async def close(self) -> None: + """Release connections; called when the app shuts down.""" + ... - async def ready(self) -> bool: ... + async def ready(self) -> bool: + """Return whether the backend can serve requests, for the /ready probe.""" + ... async def request( self, @@ -69,7 +82,12 @@ async def request( json_body: JsonObject | None = None, request_id: str | None = None, observability_id: str | None = None, - ) -> BackendResponse: ... + ) -> BackendResponse: + """Send one request and return the backend's response. + + Raise BackendUnavailable or BackendTimeout when no response arrives. + """ + ... class ClientCredentialsTokenProvider: @@ -78,6 +96,7 @@ class ClientCredentialsTokenProvider: REFRESH_MARGIN_SECONDS = 60 def __init__(self, settings: Settings, client: httpx.AsyncClient) -> None: + """Start with no cached token; the first get_token call fetches one.""" self.settings = settings self.client = client self._token: str | None = None @@ -85,6 +104,10 @@ def __init__(self, settings: Settings, client: httpx.AsyncClient) -> None: self._lock = asyncio.Lock() async def get_token(self) -> str: + """Return a bearer token, refetching it within a minute of expiry. + + Raise BackendAuthenticationError if Auth0 cannot issue one. + """ async with self._lock: now = time.time() if ( @@ -97,6 +120,7 @@ async def get_token(self) -> str: return self._token def invalidate(self) -> None: + """Drop the cached token so the next get_token call fetches a new one.""" self._token = None self._expires_at = 0.0 @@ -133,6 +157,7 @@ def __init__( transport: httpx.AsyncBaseTransport | None = None, runtime: ObservabilityRuntime | None = None, ) -> None: + """Configure the backend; no HTTP client exists until start is called.""" self.settings = settings self.transport = transport self.observability = runtime @@ -140,6 +165,7 @@ def __init__( self.token_provider: ClientCredentialsTokenProvider | None = None async def start(self) -> None: + """Create the instrumented HTTP client and token provider; idempotent.""" if self.client is not None: return timeout = httpx.Timeout( @@ -161,6 +187,7 @@ async def start(self) -> None: ) async def close(self) -> None: + """Close the HTTP client; start may be called again afterwards.""" if self.client is not None: await self.client.aclose() self.client = None @@ -172,6 +199,7 @@ def _runtime(self) -> tuple[httpx.AsyncClient, ClientCredentialsTokenProvider]: return self.client, self.token_provider async def ready(self) -> bool: + """Return whether the old gateway answers GET /health with 200.""" try: response = await self.request("GET", "/health") except (BackendUnavailable, BackendTimeout): @@ -187,6 +215,13 @@ async def request( request_id: str | None = None, observability_id: str | None = None, ) -> BackendResponse: + """Forward one request to the old gateway with a service bearer token. + + A connection failure is retried once for GET only. A 401 or 403 refreshes + the token and replays once for any method, then raises + BackendAuthenticationError. Raise BackendUnavailable if the client is + not started or the gateway is unreachable, and BackendTimeout on timeout. + """ client, token_provider = self._runtime() method = method.upper() connection_retries_remaining = 1 if method == "GET" else 0 diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/config.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/config.py index 003b5c785..14b20511b 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/config.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/config.py @@ -87,6 +87,11 @@ class Settings: @classmethod def from_env(cls) -> Settings: + """Read settings from environment variables, applying defaults. + + Only STAGE12_ENABLED is checked here (it must be 0 or 1); call validate + for everything else. Auth is required unless explicitly disabled. + """ return cls( environment=os.getenv("APP_ENVIRONMENT", "local").lower(), revision=os.getenv("K_REVISION", ""), @@ -130,10 +135,12 @@ def from_env(cls) -> Settings: @property def production(self) -> bool: + """Whether this is a production environment, where auth must stay on.""" return self.environment in PRODUCTION_ENVIRONMENTS @property def stage12_resources_configured(self) -> bool: + """Whether the Stage 12 manifest environment, database and bucket are set.""" return all( ( self.stage12_v2_manifest_environment, @@ -143,6 +150,12 @@ def stage12_resources_configured(self) -> bool: ) def validate(self) -> None: + """Raise ConfigurationError unless the service can start safely. + + Called from the app lifespan, so a misconfigured revision fails at + startup rather than on its first request. Stage 12 resources must be + configured all together or not at all. + """ if self.production and not self.auth_required: raise ConfigurationError( "Caller authentication cannot be disabled in production." diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/generate_openapi.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/generate_openapi.py index 732ed5167..70f1e9066 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/generate_openapi.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/generate_openapi.py @@ -9,6 +9,11 @@ def main() -> None: + """Write the entry app's OpenAPI spec to ``artifacts/openapi.json``. + + ``scripts/generate-clients.sh`` builds the published simulation client + from this file. + """ output = Path(__file__).resolve().parents[2] / "artifacts" / "openapi.json" output.parent.mkdir(parents=True, exist_ok=True) output.write_text( diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/stage12_adapter.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/stage12_adapter.py index 771a97148..ba598f9f2 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/stage12_adapter.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/stage12_adapter.py @@ -29,6 +29,12 @@ class ComparisonSkipReason(StrEnum): + """Why a request gets no Stage 12 comparison run. + + The value is reported as the skip event's ``reason`` attribute and in the + 422 body of the temporary direct-submission route. + """ + UNSUPPORTED_FLOW = "unsupported_flow" UNSUPPORTED_SCOPE = "unsupported_scope" UNSUPPORTED_BUDGET_WINDOW = "unsupported_budget_window" @@ -41,11 +47,14 @@ class ComparisonSkipReason(StrEnum): @dataclass(frozen=True) class AdaptedComparisonRun: + """Adapter outcome: exactly one of ``report`` and ``skip_reason`` is set.""" + report: ReportExecutionInput | None skip_reason: ComparisonSkipReason | None @property def eligible(self) -> bool: + """Whether the request was adapted into a report rather than skipped.""" return self.report is not None @@ -92,7 +101,6 @@ def adapt_annual_comparison( worker: V2WorkerVersion, ) -> AdaptedComparisonRun: """Convert only the reviewed production request shape without guessing.""" - if any( value is not None for key, value in payload.items() diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/stage12_backend.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/stage12_backend.py index a7afe6d03..63c474296 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/stage12_backend.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/stage12_backend.py @@ -50,15 +50,21 @@ class ComparisonStore(Protocol): """Read temporary comparison state from the API-owned schema.""" - def get_report(self, evaluation_id: UUID) -> ComparisonReportRecord: ... + def get_report(self, evaluation_id: UUID) -> ComparisonReportRecord: + """Return the parent report row; raise LookupError if there is none.""" + ... def list_simulations( self, evaluation_id: UUID, - ) -> tuple[ComparisonSimulationRecord, ...]: ... + ) -> tuple[ComparisonSimulationRecord, ...]: + """Return the simulation rows recorded so far for one report.""" + ... class ReportInvoker(Protocol): + """Seam that starts a worker's report coordinator, so tests can avoid Modal.""" + async def spawn( self, *, @@ -67,7 +73,9 @@ async def spawn( context_payload: dict[str, Any], parent_payload: dict[str, Any], observability_context: dict[str, Any] | None = None, - ) -> str: ... + ) -> str: + """Start the coordinator without waiting for it and return the call ID.""" + ... type PreparedReport = tuple[ @@ -84,6 +92,7 @@ class TemporaryStage12UnsupportedRequest(ValueError): """A temporary direct request cannot be normalized by the reviewed adapter.""" def __init__(self, reason: str) -> None: + """Keep the adapter's skip reason, which the route returns to the caller.""" super().__init__("Stage 12 does not support this request") self.reason = reason @@ -92,12 +101,16 @@ class TemporaryStage12DispatchFailed(RuntimeError): """Modal did not acknowledge a Stage 12 coordinator submission.""" def __init__(self, report: ComparisonReportRecord) -> None: + """Keep the report record, already marked failed, for the 502 response.""" super().__init__("Stage 12 comparison-run dispatch failed") self.report = report class ModalReportInvoker: + """Invoker that spawns report coordinators deployed as Modal functions.""" + def __init__(self, environment: str) -> None: + """Look up coordinator functions in this Modal environment.""" self._environment = environment async def spawn( @@ -109,6 +122,10 @@ async def spawn( parent_payload: dict[str, Any], observability_context: dict[str, Any] | None = None, ) -> str: + """Spawn the worker's report coordinator and return its Modal call ID. + + Raise RuntimeError if Modal returns no invocation identifier. + """ function = modal.Function.from_name( worker.application_name, worker.report_coordinator_callable, @@ -131,6 +148,12 @@ async def spawn( class Stage12ComparisonBackend: + """Backend that prepares Stage 12 comparison runs for Modal coordinators. + + It writes no database rows: the spawned coordinator owns all durable + state, and this class only reads it back for polling. + """ + def __init__( self, settings: Settings, @@ -140,6 +163,7 @@ def __init__( invoker: ReportInvoker, runtime: ObservabilityRuntime, ) -> None: + """Wire the manifest, persistence, invoker and telemetry dependencies.""" self._settings = settings self._manifest_loader = manifest_loader self._store = store @@ -153,6 +177,10 @@ def from_settings( *, runtime: ObservabilityRuntime, ) -> Stage12ComparisonBackend: + """Build the production backend from the Modal manifest and Stage 12 database. + + The manifest Dict must already exist; it is never created here. + """ manifest_store = modal.Dict.from_name( settings.stage12_v2_manifest_name, environment_name=settings.stage12_v2_manifest_environment, @@ -174,6 +202,14 @@ async def dispatch_after_production( request_id: str, observability_id: str, ) -> None: + """Spawn a shadow comparison of an accepted production job. + + The evaluation ID is derived from the environment, production job and + worker version, so a repeat dispatch maps to the same report. Requests + the adapter cannot represent are skipped with a telemetry event. + Raise ValueError if the production response has no ``job_id`` and + TemporaryStage12DispatchFailed if Modal does not accept the spawn. + """ prepared = await asyncio.to_thread( self._prepare_automatic_report, request_payload=request_payload, @@ -205,7 +241,6 @@ async def submit_temporary_report( observability_id: str, ) -> ComparisonReportRecord: """Submit one direct report and return after Modal acknowledges it.""" - prepared = await asyncio.to_thread( self._prepare_temporary_report, request_payload=request_payload, @@ -229,7 +264,6 @@ async def get_temporary_report( evaluation_id: UUID, ) -> tuple[ComparisonReportRecord, tuple[ComparisonSimulationRecord, ...]]: """Read coordinator-owned durable state without polling Modal.""" - return await asyncio.to_thread(self._get_temporary_report, evaluation_id) def _get_temporary_report( @@ -433,7 +467,6 @@ async def _spawn_report( observability_id: str, ) -> ComparisonReportRecord: """Return after Modal accepts the coordinator; perform no database writes.""" - artifact_prefix = ( f"stage-12-runs/{self._settings.environment}/" f"{parent.created_at:%Y}/{parent.created_at:%m}/" diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/stage12_infrastructure.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/stage12_infrastructure.py index ecfcaacc4..ccae557c9 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/stage12_infrastructure.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/stage12_infrastructure.py @@ -21,6 +21,11 @@ def _parser() -> argparse.ArgumentParser: def main() -> None: + """Check the runtime database credential from the command line. + + Exit with the error message if the role, environment or schema does not + match what the Stage 12 runtime expects. + """ args = _parser().parse_args() database_url = args.database_url_file.read_text(encoding="utf-8").strip() try: diff --git a/projects/policyengine-simulation-entry/tests/test_deploy_workflow_scripts.py b/projects/policyengine-simulation-entry/tests/test_deploy_workflow_scripts.py index d158ff522..c7c715b59 100644 --- a/projects/policyengine-simulation-entry/tests/test_deploy_workflow_scripts.py +++ b/projects/policyengine-simulation-entry/tests/test_deploy_workflow_scripts.py @@ -30,7 +30,7 @@ def fakes(tmp_path: Path) -> FakeCommands: def _step(workflow: str, name: str) -> str: - """The text of one workflow step, from its ``- name:`` to the next step.""" + """Return one workflow step's text, from its ``- name:`` to the next step.""" match = re.search( rf"^(?P +)- name: {re.escape(name)}\n(?P.*?)(?=^(?P=indent)- |^ \S|\Z)", workflow, diff --git a/projects/policyengine-simulation-entry/tests/test_workflow_scripts.py b/projects/policyengine-simulation-entry/tests/test_workflow_scripts.py index ab46e34cb..75887df7a 100644 --- a/projects/policyengine-simulation-entry/tests/test_workflow_scripts.py +++ b/projects/policyengine-simulation-entry/tests/test_workflow_scripts.py @@ -1,5 +1,7 @@ -"""Tests for the workflow scripts under .github/scripts/ that no deployment -test covers: each runs the script with fake commands first on PATH.""" +"""Tests for the .github/scripts/ workflow scripts no deployment test covers. + +Each test runs the script with fake commands first on PATH. +""" from __future__ import annotations diff --git a/projects/policyengine-simulation-entry/tests/workflow_script_support.py b/projects/policyengine-simulation-entry/tests/workflow_script_support.py index 8873cfe2a..f9111816e 100644 --- a/projects/policyengine-simulation-entry/tests/workflow_script_support.py +++ b/projects/policyengine-simulation-entry/tests/workflow_script_support.py @@ -57,7 +57,7 @@ def calls(self, name: str | None = None) -> list[list[str]]: return [call for call in recorded if name is None or call[0] == name] def env(self, **overrides: str) -> dict[str, str]: - """The environment to run a script in, with the fakes first on PATH.""" + """Return the environment to run a script in, with the fakes first on PATH.""" return { **os.environ, "PATH": f"{self.directory}{os.pathsep}{os.environ['PATH']}", @@ -84,7 +84,7 @@ def run_script( def github_output(path: Path) -> dict[str, str]: - """The ``name=value`` pairs a script appended to ``$GITHUB_OUTPUT``.""" + """Return the ``name=value`` pairs a script appended to ``$GITHUB_OUTPUT``.""" if not path.exists(): return {} pairs = {} From 1126ebb3fd69b86e2bc4add4ea113b58d613637a Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 13 Oct 2026 16:56:47 +0400 Subject: [PATCH 06/18] Document the simulation gateway for ruff D The /ping router endpoint keeps its OpenAPI description under a D205 noqa; the golden OpenAPI test still passes. Two placeholder pass statements in the test fixtures' mock exception classes go now that the classes have docstrings (PIE790). Co-Authored-By: Claude Opus 5.5 --- .../fixtures/__init__.py | 1 + .../fixtures/gateway_endpoints.py | 36 ++++++++- .../fixtures/package_imports.py | 1 + .../__init__.py | 4 +- .../policyengine_simulation_gateway/app.py | 14 ++-- .../policyengine_simulation_gateway/auth.py | 1 - .../endpoints.py | 36 ++++----- .../generate_openapi.py | 17 ++-- .../responses.py | 18 +++++ .../smoke_app.py | 5 ++ .../tests/test_auth.py | 45 ++++++----- .../tests/test_endpoints.py | 80 ++++++------------- .../tests/test_health.py | 12 +-- .../tests/test_ping.py | 30 ++----- .../tests/test_spm_routes.py | 60 +++++++++----- 15 files changed, 189 insertions(+), 171 deletions(-) diff --git a/projects/policyengine-simulation-gateway/fixtures/__init__.py b/projects/policyengine-simulation-gateway/fixtures/__init__.py index e69de29bb..bd6f7f7a0 100644 --- a/projects/policyengine-simulation-gateway/fixtures/__init__.py +++ b/projects/policyengine-simulation-gateway/fixtures/__init__.py @@ -0,0 +1 @@ +"""Pytest plugins shared by the gateway tests, loaded from tests/conftest.py.""" diff --git a/projects/policyengine-simulation-gateway/fixtures/gateway_endpoints.py b/projects/policyengine-simulation-gateway/fixtures/gateway_endpoints.py index 3866aff93..2530a4a9d 100644 --- a/projects/policyengine-simulation-gateway/fixtures/gateway_endpoints.py +++ b/projects/policyengine-simulation-gateway/fixtures/gateway_endpoints.py @@ -131,6 +131,12 @@ def resolve_test_dataset_uri( dataset: str | None, data_version: str | None = None, ) -> str | None: + """Return the dataset URI the gateway should report for this test request. + + Mirrors the gateway's resolution against ``TEST_APP_RELEASE_BUNDLE``: + aliases and bundle dataset names map to pinned URIs, and an explicit + revision or ``data_version`` overrides the bundle default. + """ country_bundle = TEST_APP_RELEASE_BUNDLE[country] if dataset is None: return _runtime_dataset_uri( @@ -162,23 +168,29 @@ class MockDict: """Mock for Modal.Dict to simulate version registry.""" def __init__(self, data: dict) -> None: + """Wrap ``data`` without copying, so writes reach the test's dict.""" self._data = data def __getitem__(self, key: str) -> object: + """Return the stored value; raise KeyError like a Modal Dict.""" if key not in self._data: raise KeyError(key) return self._data[key] def __setitem__(self, key: str, value: object) -> None: + """Store a value in the wrapped dict.""" self._data[key] = value def get(self, key: str, default: object = None) -> object: + """Return the stored value, or ``default`` when the key is absent.""" return self._data.get(key, default) def keys(self) -> KeysView[str]: + """Return a live view of the stored keys.""" return self._data.keys() def items(self) -> ItemsView[str, object]: + """Return a live view of the stored items.""" return self._data.items() @classmethod @@ -194,6 +206,10 @@ class MockFunctionCall: from_id_errors = {} def __init__(self, object_id: str = "mock-job-id-123") -> None: + """Register a complete call whose result is a valid zeroed macro output. + + Tests set ``running`` or ``error`` afterwards to change what get returns. + """ self.object_id = object_id value = BaselineReformValue(baseline=0.0, reform=0.0) age_poverty = AgePovertyOutput( @@ -244,6 +260,7 @@ def __init__(self, object_id: str = "mock-job-id-123") -> None: self.__class__.registry[object_id] = self def get(self, timeout: int = 0) -> object: + """Return the result; raise TimeoutError while running, or the set error.""" if self.running: raise TimeoutError if self.error is not None: @@ -252,6 +269,10 @@ def get(self, timeout: int = 0) -> object: @classmethod def from_id(cls, object_id: str) -> "MockFunctionCall": + """Return a registered call, raising any error staged for that ID. + + Raise KeyError for an unknown ID. + """ if object_id in cls.from_id_errors: raise cls.from_id_errors[object_id] if object_id not in cls.registry: @@ -263,6 +284,7 @@ class MockFunction: """Mock for Modal Function.""" def __init__(self) -> None: + """Start with no recorded spawns or lookups.""" self.last_payload = None self.last_observability_context = None self.last_from_name_call = None @@ -270,9 +292,11 @@ def __init__(self) -> None: self.calls = [] def bind(self, app_name: str, func_name: str) -> "BoundMockFunction": + """Return a handle for one app function that records into this mock.""" return BoundMockFunction(self, app_name, func_name) def call_for(self, object_id: str) -> MockFunctionCall: + """Register a completed call under ``object_id`` and make it the last call.""" call = MockFunctionCall(object_id=object_id) self.last_call = call return call @@ -282,6 +306,7 @@ class BoundMockFunction: """Function handle returned by Modal.Function.from_name.""" def __init__(self, recorder: MockFunction, app_name: str, func_name: str) -> None: + """Bind the app and function names that spawns will be recorded under.""" self.recorder = recorder self.app_name = app_name self.func_name = func_name @@ -292,6 +317,11 @@ def spawn( *, observability_context: dict | None = None, ) -> MockFunctionCall: + """Record the spawn and return a call with a fixed ID. + + A ``run_budget_window_batch`` spawn gets the batch ID and stays running; + any other spawn completes immediately. + """ self.recorder.last_payload = payload self.recorder.last_observability_context = observability_context is_batch = self.func_name == "run_budget_window_batch" @@ -312,11 +342,13 @@ def spawn( class MockModalException: + """Stand-in for ``modal.exception`` with the error types the gateway catches.""" + class NotFoundError(Exception): - pass + """Raised for a function call Modal has never heard of.""" class OutputExpiredError(Exception): - pass + """Raised for a function call whose result Modal no longer keeps.""" @pytest.fixture diff --git a/projects/policyengine-simulation-gateway/fixtures/package_imports.py b/projects/policyengine-simulation-gateway/fixtures/package_imports.py index d508167d1..a1e7232e2 100644 --- a/projects/policyengine-simulation-gateway/fixtures/package_imports.py +++ b/projects/policyengine-simulation-gateway/fixtures/package_imports.py @@ -36,6 +36,7 @@ class GatewayImportModuleNames: @pytest.fixture() def gateway_import_module_names() -> GatewayImportModuleNames: + """Return the module names the import-boundary tests assert on.""" return GatewayImportModuleNames() diff --git a/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/__init__.py b/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/__init__.py index d15c1d386..d6d47b71b 100644 --- a/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/__init__.py +++ b/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/__init__.py @@ -1,3 +1 @@ -""" -Gateway package for PolicyEngine Simulation API. -""" +"""Gateway package for PolicyEngine Simulation API.""" diff --git a/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/app.py b/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/app.py index ee05f39a1..23984108c 100644 --- a/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/app.py +++ b/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/app.py @@ -1,5 +1,4 @@ -""" -PolicyEngine Simulation Gateway - Modal App Definition +"""Modal app definition for the PolicyEngine Simulation Gateway. A lightweight, stable gateway that routes simulation requests to versioned simulation apps. This app rarely changes and provides a stable URL for consumers. @@ -44,9 +43,11 @@ def build_gateway_image() -> modal.Image: - """The gateway image, shared with the smoke app (smoke_app.py) so - both construct identical layer definitions and share Modal's - content-addressed cache.""" + """Build the gateway's Modal image. + + The smoke app (smoke_app.py) builds its image here too, so both construct + identical layer definitions and share Modal's content-addressed cache. + """ return ( modal.Image.debian_slim(python_version="3.13") .uv_sync( @@ -71,8 +72,7 @@ def build_gateway_image() -> modal.Image: @app.function(image=gateway_image, secrets=[gateway_auth_secret]) @modal.asgi_app() def web_app() -> FastAPI: - """ - FastAPI gateway for simulation job submission and polling. + """FastAPI gateway for simulation job submission and polling. Provides stable endpoints: POST /simulate/economy/comparison - Submit a simulation job diff --git a/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/auth.py b/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/auth.py index bd4265f28..ee0a27a71 100644 --- a/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/auth.py +++ b/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/auth.py @@ -239,7 +239,6 @@ def require_auth( silently make the gateway private. Partial auth configuration always returns 503 because it indicates an incomplete secret. """ - if _auth_disabled(): return None 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 6fd6796ac..5d6a62fa6 100644 --- a/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/endpoints.py +++ b/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/endpoints.py @@ -1,6 +1,4 @@ -""" -FastAPI endpoints for the Gateway API. -""" +"""FastAPI endpoints for the Gateway API.""" import logging import re @@ -83,6 +81,14 @@ class VersionRoutingState(TypedDict, total=False): @dataclass(frozen=True) class RouteResolution: + """The Modal app a request routes to, and the provenance the gateway reports. + + ``response_version`` is the version echoed to the caller, and + ``bundle_manifest`` is empty when the route came from the legacy version + dicts. ``route_provenance`` marks pre-canonical routes, which SPM + resolution treats differently. + """ + app_name: str response_version: str policyengine_version: str | None @@ -229,7 +235,7 @@ def _bundle_manifest(state: dict, policyengine_version: str | None) -> dict: def _policyengine_version_for_app( state: dict, app_name: str, *, country: str, model_version: str ) -> str | None: - """The wrapper version serving this country route, or None. + """Return the wrapper version serving this country route, or None. Behaviour change from the pre-canonical gateway, deliberate: a country route whose one candidate bundle states a *different* model version is @@ -616,7 +622,7 @@ def _bundle_identities(state: dict) -> dict[str, dict[str, PolicyEngineBundle]]: def _certified_model_version(country: str, route: RouteResolution) -> str | None: - """The country model version the registry states for this route. + """Return the country model version the registry states for this route. ``PolicyEngineBundle.model_version`` falls back to the routing response version so the response always carries one. On a policyengine-keyed @@ -789,8 +795,7 @@ async def submit_simulation( http_request: Request, response: Response, ) -> JobSubmitResponse | JSONResponse: - """ - Submit a simulation job. + """Submit a simulation job. Matches the existing Cloud Run API endpoint path. Routes to the appropriate app based on country and version params. @@ -899,9 +904,7 @@ async def submit_budget_window_batch( http_request: Request, response: Response, ) -> BudgetWindowBatchSubmitResponse | JSONResponse: - """ - Submit a budget-window batch job. - """ + """Submit a budget-window batch job.""" runtime = _runtime(http_request) observability_id = _start_workflow_observability(http_request, runtime) response.headers[OBSERVABILITY_ID_HEADER] = observability_id @@ -994,14 +997,14 @@ async def submit_budget_window_batch( async def get_job_status( job_id: str, request: Request, response: Response ) -> JobStatusResponse | JSONResponse: - """ - Poll for job status. + """Poll for job status. Returns: - 200 with status="complete" and result when done - 202 with status="running" while in progress - 500 with status="failed" and error on failure - 404 if job_id not found + """ runtime = _runtime(request) runtime.set_context(job_id=job_id) @@ -1077,9 +1080,7 @@ async def get_budget_window_job_status( request: Request, response: Response, ) -> JSONResponse: - """ - Poll for budget-window batch status. - """ + """Poll for budget-window batch status.""" runtime = _runtime(request) runtime.set_context(batch_job_id=batch_job_id) with runtime.span(BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_STATE_LOAD)): @@ -1264,8 +1265,7 @@ async def health() -> HealthResponse: @router.post("/ping", response_model=PingResponse) async def ping(request: PingRequest) -> PingResponse: - """ - Verify the API is able to receive and process requests. + """Verify the API is able to receive and process requests. Matches the policyengine_fastapi.ping endpoint for test compatibility. - """ + """ # noqa: D205 # docstring would change the published OpenAPI description return PingResponse(incremented=request.value + 1) diff --git a/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/generate_openapi.py b/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/generate_openapi.py index 207e5d76d..14770300d 100644 --- a/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/generate_openapi.py +++ b/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/generate_openapi.py @@ -1,5 +1,4 @@ -""" -Generate OpenAPI spec for the Modal Gateway API. +"""Generate OpenAPI spec for the Modal Gateway API. This creates a FastAPI app with the same route signatures as the gateway but without Modal dependencies, allowing OpenAPI generation without credentials. @@ -45,8 +44,7 @@ def create_openapi_app() -> FastAPI: }, ) async def submit_simulation(request: SimulationRequest) -> JobSubmitResponse: - """ - Submit a simulation job. + """Submit a simulation job. Routes to the appropriate simulation app based on country and version. Returns immediately with a job_id for polling. @@ -64,8 +62,7 @@ async def submit_simulation(request: SimulationRequest) -> JobSubmitResponse: async def submit_budget_window_batch( request: BudgetWindowBatchRequest, ) -> BudgetWindowBatchSubmitResponse: - """ - Submit a budget-window batch job. + """Submit a budget-window batch job. Returns immediately with a parent batch job ID for polling. """ @@ -82,14 +79,14 @@ async def submit_budget_window_batch( }, ) async def get_job_status(job_id: str) -> JobStatusResponse: - """ - Poll for job status. + """Poll for job status. Returns: - 200 with status="complete" and result when done - 202 with status="running" while in progress - 404 if job_id not found - 500 with status="failed" and error on failure + """ raise NotImplementedError("Stub for OpenAPI generation") @@ -109,9 +106,7 @@ async def get_job_status(job_id: str) -> JobStatusResponse: async def get_budget_window_job_status( batch_job_id: str, ) -> BudgetWindowBatchStatusResponse: - """ - Poll for budget-window batch status. - """ + """Poll for budget-window batch status.""" raise NotImplementedError("Stub for OpenAPI generation") @app.get("/versions", response_model=VersionsResponse) diff --git a/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/responses.py b/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/responses.py index 79428a456..967dc6b4d 100644 --- a/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/responses.py +++ b/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/responses.py @@ -16,10 +16,16 @@ class ServerErrorResponse(JSONResponse): """Shared 500 JSON response.""" def __init__(self, content: dict) -> None: + """Serialize ``content`` as the body of a 500 response.""" super().__init__(status_code=500, content=content) def batch_status_payload(response: BudgetWindowBatchStatusResponse) -> dict: + """Return the batch status as JSON, omitting null ``errors`` fields. + + ``errors`` is dropped at the top level and on each child job when it is + None, and ``policyengine_bundle`` omits its unset fields. + """ payload = response.model_dump(mode="json") if payload.get("errors") is None: payload.pop("errors", None) @@ -39,6 +45,11 @@ def batch_status_response( *, headers: dict[str, str] | None = None, ) -> JSONResponse: + """Return the batch status with the HTTP code the polling contract uses. + + Submitted or running batches get 202, complete ones 200, and failed ones + 400 when they carry structured caller errors, else 500. + """ payload = batch_status_payload(response) if response.status in {"submitted", "running"}: return JSONResponse(status_code=202, content=payload, headers=headers) @@ -56,6 +67,7 @@ def running_job_response( *, headers: dict[str, str] | None = None, ) -> JSONResponse: + """Return a 202 running-job body, merged with the job's public metadata.""" return JSONResponse( status_code=202, content={ @@ -75,6 +87,12 @@ def failed_job_response( errors: list[dict] | None = None, headers: dict[str, str] | None = None, ) -> JSONResponse: + """Return a failed-job body: 400 with structured ``errors``, otherwise 500. + + Pass structured errors only for caller-correctable failures such as an + invalid SPM selection. ``error`` reaches the caller verbatim, so redact + exception text first. + """ return JSONResponse( status_code=400 if errors else 500, content={ diff --git a/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/smoke_app.py b/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/smoke_app.py index b0651d821..6b53b8312 100644 --- a/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/smoke_app.py +++ b/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/smoke_app.py @@ -22,6 +22,10 @@ @app.function(image=build_gateway_image(), timeout=300) def smoke_import_gateway() -> dict: + """Import every shipped package module and build the ASGI app in-image. + + Return the number of modules imported and the app's ``METHOD path`` routes. + """ import importlib import os import pkgutil @@ -63,6 +67,7 @@ def smoke_import_gateway() -> dict: @app.local_entrypoint() def main() -> None: + """Run the in-image smoke remotely; fail unless modules and /health loaded.""" report = smoke_import_gateway.remote() print(report) assert report["modules_imported"] > 0 diff --git a/projects/policyengine-simulation-gateway/tests/test_auth.py b/projects/policyengine-simulation-gateway/tests/test_auth.py index 906bbcae5..610667ba3 100644 --- a/projects/policyengine-simulation-gateway/tests/test_auth.py +++ b/projects/policyengine-simulation-gateway/tests/test_auth.py @@ -41,9 +41,11 @@ def _request(): @pytest.fixture def unauthenticated_client(monkeypatch) -> TestClient: - """A client where the real auth dependency is active but no token is - attached. The underlying JWTDecoder is stubbed to preserve the 403 - contract without making a live JWKS fetch.""" + """Return a client with the real auth dependency active and no token. + + The underlying JWTDecoder is stubbed to preserve the 403 contract without + making a live JWKS fetch. + """ class FailingDecoder: def __call__(self, token): @@ -67,7 +69,6 @@ def test__given_no_bearer_token__then_gated_endpoint_returns_403( unauthenticated_client, method, path, body ): """Missing bearer tokens should be rejected on all private endpoints.""" - if method == "post": response = unauthenticated_client.post(path, json=body) else: @@ -135,7 +136,6 @@ def test__given_partial_auth_config__then_dependency_raises_503(monkeypatch): def test__given_health_endpoint__then_auth_not_required(monkeypatch): """Health/ping/versions endpoints remain public by design.""" - from policyengine_simulation_gateway.testing import create_gateway_app client = TestClient(create_gateway_app(authenticate=False)) @@ -144,10 +144,11 @@ def test__given_health_endpoint__then_auth_not_required(monkeypatch): def test__given_same_env__then_decoder_is_cached(monkeypatch): - """Calling ``_get_decoder`` repeatedly must return the same instance so - the wrapped ``PyJWKClient`` JWKS cache is reused across requests. - Rebuilding the decoder per request defeats the cache (#458 review).""" + """Repeated ``_get_decoder`` calls must return the same instance. + That keeps the wrapped ``PyJWKClient`` JWKS cache in use across requests; + rebuilding the decoder per request defeats the cache (#458 review). + """ monkeypatch.delenv(auth_module.GATEWAY_AUTH_DISABLED_ENV, raising=False) monkeypatch.setenv(auth_module.GATEWAY_AUTH_ISSUER_ENV, "https://issuer.example/") monkeypatch.setenv(auth_module.GATEWAY_AUTH_AUDIENCE_ENV, "aud-caching") @@ -160,9 +161,11 @@ def test__given_same_env__then_decoder_is_cached(monkeypatch): def test__given_rotated_audience__then_cache_returns_new_decoder(monkeypatch): - """Cache is keyed on issuer+audience so rotating the audience yields a - fresh decoder without polluting the previous one.""" + """Rotating the audience yields a fresh decoder. + The cache is keyed on issuer+audience, so the previous decoder is left + untouched. + """ monkeypatch.delenv(auth_module.GATEWAY_AUTH_DISABLED_ENV, raising=False) monkeypatch.setenv(auth_module.GATEWAY_AUTH_ISSUER_ENV, "https://issuer.example/") @@ -180,10 +183,11 @@ def test__given_rotated_audience__then_cache_returns_new_decoder(monkeypatch): def test__given_repeated_requests__then_decoder_not_reinstantiated(monkeypatch): - """Smoke test: hitting a gated endpoint many times must not rebuild the - decoder. We spy on ``JWTDecoder.__init__`` and assert it runs at most - once across many requests.""" + """Hitting a gated endpoint many times must not rebuild the decoder. + Smoke test: we spy on ``JWTDecoder.__init__`` and assert it runs at most + once across many requests. + """ from policyengine_fastapi.auth import jwt_decoder as jwt_decoder_module from policyengine_simulation_gateway.testing import create_gateway_app @@ -218,12 +222,14 @@ def _counting_init(self, *args, **kwargs): def test__given_dependency_override__then_gated_endpoint_returns_200( mock_modal, client: TestClient ): - """Positive-path auth test: when ``app.dependency_overrides`` bypasses + """Overriding ``require_auth`` lets a gated endpoint return 200. + + Positive-path auth test: when ``app.dependency_overrides`` bypasses ``require_auth`` (as real tests do), the gated endpoint returns 200 with the expected payload. This guards the override path that other tests rely on — if the override stops working, every gated-endpoint test would - start returning 403 instead of exercising real logic.""" - + start returning 403 instead of exercising real logic. + """ mock_modal["dicts"]["simulation-api-us-versions"] = { "latest": "1.500.0", "1.500.0": "policyengine-simulation-py4-10-0", @@ -242,8 +248,11 @@ def test__given_dependency_override__then_gated_endpoint_returns_200( class TestProductionAuthGuard: - """Tests for ``enforce_production_auth_guard`` — the startup-time check - that prevents ``GATEWAY_AUTH_DISABLED`` from slipping into production.""" + """Tests for ``enforce_production_auth_guard``. + + It is the startup-time check that prevents ``GATEWAY_AUTH_DISABLED`` from + slipping into production. + """ def test__given_auth_enabled__then_guard_noops(self, monkeypatch): monkeypatch.delenv(auth_module.GATEWAY_AUTH_DISABLED_ENV, raising=False) diff --git a/projects/policyengine-simulation-gateway/tests/test_endpoints.py b/projects/policyengine-simulation-gateway/tests/test_endpoints.py index 31099833f..60c12d5ab 100644 --- a/projects/policyengine-simulation-gateway/tests/test_endpoints.py +++ b/projects/policyengine-simulation-gateway/tests/test_endpoints.py @@ -1,5 +1,4 @@ -""" -Unit tests for gateway endpoints. +"""Unit tests for gateway endpoints. Tests verify the endpoint correctly resolves app names and routes simulation requests. @@ -71,11 +70,7 @@ def test__given_legacy_version_matching_policyengine_bundle__then_routes_by_bund assert app_name == "policyengine-simulation-py4-10-0" def test__given_us_country_no_version__then_returns_latest_app(self, mock_modal): - """ - Given country='us' and no version specified - When get_app_name is called - Then returns the app name for the 'latest' version. - """ + """Given 'us' and no version, when resolving, then use the 'latest' app.""" from policyengine_simulation_gateway.endpoints import get_app_name # When @@ -88,11 +83,7 @@ def test__given_us_country_no_version__then_returns_latest_app(self, mock_modal) def test__given_us_country_with_version__then_returns_specified_app( self, mock_modal ): - """ - Given country='us' and a specific version - When get_app_name is called - Then returns the app name for that version. - """ + """Given 'us' and a version, when resolving, then use that version's app.""" from policyengine_simulation_gateway.endpoints import get_app_name # When @@ -103,11 +94,7 @@ def test__given_us_country_with_version__then_returns_specified_app( assert app_name == "policyengine-simulation-py3-9-0" def test__given_uk_country__then_uses_uk_version_dict(self, mock_modal): - """ - Given country='uk' - When get_app_name is called - Then uses the UK version dictionary. - """ + """Given country='uk', when resolving, then use the UK version dictionary.""" from policyengine_simulation_gateway.endpoints import get_app_name # When @@ -118,11 +105,7 @@ def test__given_uk_country__then_uses_uk_version_dict(self, mock_modal): assert app_name == "policyengine-simulation-py4-10-0" def test__given_invalid_country__then_raises_value_error(self): - """ - Given an invalid country code - When get_app_name is called - Then raises ValueError. - """ + """Given an invalid country code, when resolving, then raise ValueError.""" from policyengine_simulation_gateway.endpoints import get_app_name # When / Then @@ -130,11 +113,7 @@ def test__given_invalid_country__then_raises_value_error(self): get_app_name("invalid", None) def test__given_unknown_version__then_raises_value_error(self, mock_modal): - """ - Given a version not in the registry - When get_app_name is called - Then raises ValueError. - """ + """Given an unregistered version, when resolving, then raise ValueError.""" from policyengine_simulation_gateway.endpoints import get_app_name # When / Then @@ -200,11 +179,7 @@ class TestSubmitSimulationEndpoint: def test__given_data_value__then_rejects_before_spawning( self, mock_modal, client: TestClient ): - """ - Given a request with a data value - When the simulation is submitted - Then request validation rejects it with HTTP 422. - """ + """Given a data value, when submitting, then validation rejects it with 422.""" # Given mock_modal["dicts"]["simulation-api-us-versions"] = { "latest": "1.500.0", @@ -229,11 +204,7 @@ def test__given_data_value__then_rejects_before_spawning( def test__given_no_data_value__then_routes_to_run_simulation( self, mock_modal, client: TestClient ): - """ - Given a request with no data field - When the simulation is submitted - Then routes to regular run_simulation function. - """ + """Given no data field, when submitting, then route to run_simulation.""" # Given mock_modal["dicts"]["simulation-api-us-versions"] = { "latest": "1.500.0", @@ -261,11 +232,7 @@ def test__given_no_data_value__then_routes_to_run_simulation( def test__given_submission__then_returns_job_id_and_poll_url( self, mock_modal, client: TestClient ): - """ - Given a simulation submission - When the request completes - Then returns job_id and poll_url for async polling. - """ + """Given a submission, when it completes, then return job_id and poll_url.""" # Given mock_modal["dicts"]["simulation-api-us-versions"] = { "latest": "1.500.0", @@ -381,7 +348,8 @@ def test__given_legacy_process_id__then_maps_submission_claim_id( def test__given_submission_header__then_propagates_observability_id( self, mock_modal, client: TestClient ): - """ + """Propagate the diagnostic identifier header to Modal and the response. + Given a simulation submission with the diagnostic identifier header When the request completes Then the Modal context and response header preserve the identifier. @@ -421,7 +389,8 @@ def test__given_submission_header__then_propagates_observability_id( def test__given_submission_without_data__then_returns_default_bundle_metadata( self, mock_modal, client: TestClient ): - """ + """Report the resolved app and default dataset provenance on submission. + Given a simulation submission without a dataset override When the request completes Then the response exposes the resolved app and submitted dataset provenance. @@ -864,11 +833,7 @@ def test__given_submission_with_unknown_dataset_name__then_rejected_before_spawn def test__given_submitted_job__then_job_status_includes_bundle_metadata( self, mock_modal, client: TestClient ): - """ - Given a submitted simulation job - When polling job status - Then the resolved bundle metadata is returned with the status response. - """ + """Given a submitted job, when polling, then return its bundle metadata.""" # Given mock_modal["dicts"]["simulation-api-us-versions"] = { "latest": "1.500.0", @@ -933,7 +898,8 @@ def test__given_submitted_job_with_header__then_polling_echoes_header( def test__given_unknown_job_id__then_polling_returns_404( self, mock_modal, client: TestClient, monkeypatch ): - """ + """Return 404 for a job ID the gateway never issued, without asking Modal. + Given a job id that the gateway never issued When polling job status Then the gateway returns 404 before asking Modal for a call result. @@ -962,7 +928,8 @@ def test__given_unknown_job_id__then_polling_returns_404( def test__given_lazy_modal_call_without_metadata__then_polling_returns_404( self, mock_modal, client: TestClient ): - """ + """Return 404 for a job ID with no gateway metadata, even if Modal has one. + Given Modal can construct a FunctionCall handle for an arbitrary id When the gateway has no metadata for that id Then the gateway still treats it as not found. @@ -1400,7 +1367,8 @@ def test__given_unknown_budget_window_job__then_records_handled_404( def test__given_parent_lookup_failure__then_degraded_poll_is_not_an_error( self, mock_modal, client: TestClient, monkeypatch ): - """ + """Serve the seed state when the parent lookup fails, without an error. + Given a submitted batch whose parent FunctionCall lookup fails When polling batch status Then the gateway serves the seed state (202) and records a degradation @@ -1729,12 +1697,14 @@ def test__given_max_parallel_above_active_limit__then_budget_window_submit_retur def test__given_parent_call_raises__then_failure_persists_across_polls( self, mock_modal, client: TestClient ): - """Regression test for #448: when the parent Modal FunctionCall raises + """Persist a parent-call failure so the failed state survives re-polling. + + Regression test for #448: when the parent Modal FunctionCall raises on .get() the first poll flipped seed state to failed in-memory only. Subsequent polls re-read the unmodified seed and flipped back to "submitted". Verify the mutation is persisted so the terminal state - survives a second poll.""" - + survives a second poll. + """ mock_modal["dicts"]["simulation-api-us-versions"] = { "latest": "1.500.0", "1.500.0": "policyengine-simulation-py4-10-0", diff --git a/projects/policyengine-simulation-gateway/tests/test_health.py b/projects/policyengine-simulation-gateway/tests/test_health.py index 67be999ab..826c76a04 100644 --- a/projects/policyengine-simulation-gateway/tests/test_health.py +++ b/projects/policyengine-simulation-gateway/tests/test_health.py @@ -8,11 +8,7 @@ class TestHealthEndpoint: """Tests for GET /health endpoint.""" def test_health_returns_healthy_status(self, client: TestClient): - """ - Given a running gateway service - When the health endpoint is called - Then the response indicates healthy status. - """ + """Given a running gateway, when calling /health, then report healthy.""" # When response = client.get("/health") @@ -34,11 +30,7 @@ def test_health_does_not_adopt_an_incoming_workflow_identifier( assert OBSERVABILITY_ID_HEADER not in response.headers def test_health_is_idempotent(self, client: TestClient): - """ - Given a running gateway service - When the health endpoint is called multiple times - Then each response is identical. - """ + """Given a running gateway, when calling /health repeatedly, then match.""" # When response1 = client.get("/health") response2 = client.get("/health") diff --git a/projects/policyengine-simulation-gateway/tests/test_ping.py b/projects/policyengine-simulation-gateway/tests/test_ping.py index 577d6a30a..d34c16a21 100644 --- a/projects/policyengine-simulation-gateway/tests/test_ping.py +++ b/projects/policyengine-simulation-gateway/tests/test_ping.py @@ -7,11 +7,7 @@ class TestPingEndpoint: """Tests for POST /ping endpoint.""" def test_ping_increments_value(self, client: TestClient): - """ - Given a ping request with a value - When the ping endpoint is called - Then the response contains the incremented value. - """ + """Given a value, when pinging, then the response has the value plus one.""" # Given request_body = {"value": 10} @@ -23,11 +19,7 @@ def test_ping_increments_value(self, client: TestClient): assert response.json() == {"incremented": 11} def test_ping_handles_zero(self, client: TestClient): - """ - Given a ping request with value zero - When the ping endpoint is called - Then the response contains 1. - """ + """Given value zero, when pinging, then the response contains 1.""" # Given request_body = {"value": 0} @@ -39,11 +31,7 @@ def test_ping_handles_zero(self, client: TestClient): assert response.json() == {"incremented": 1} def test_ping_handles_negative_value(self, client: TestClient): - """ - Given a ping request with a negative value - When the ping endpoint is called - Then the response contains the incremented (less negative) value. - """ + """Given a negative value, when pinging, then return it incremented.""" # Given request_body = {"value": -5} @@ -55,11 +43,7 @@ def test_ping_handles_negative_value(self, client: TestClient): assert response.json() == {"incremented": -4} def test_ping_rejects_missing_value(self, client: TestClient): - """ - Given a ping request without a value field - When the ping endpoint is called - Then a 422 validation error is returned. - """ + """Given no value field, when pinging, then return a 422 validation error.""" # Given request_body = {} @@ -70,11 +54,7 @@ def test_ping_rejects_missing_value(self, client: TestClient): assert response.status_code == 422 def test_ping_rejects_non_integer_value(self, client: TestClient): - """ - Given a ping request with a non-integer value - When the ping endpoint is called - Then a 422 validation error is returned. - """ + """Given a non-integer value, when pinging, then return a 422 error.""" # Given request_body = {"value": "not_an_integer"} diff --git a/projects/policyengine-simulation-gateway/tests/test_spm_routes.py b/projects/policyengine-simulation-gateway/tests/test_spm_routes.py index f1d1c367f..8ff06fd43 100644 --- a/projects/policyengine-simulation-gateway/tests/test_spm_routes.py +++ b/projects/policyengine-simulation-gateway/tests/test_spm_routes.py @@ -25,10 +25,13 @@ def legacy_route(mock_modal, source, model_version): - """Two shapes of pre-canonical route: the original per-country Modal - dicts, and a seeded routing-state country route the registry ties to no - wrapper bundle. For the latter, ``source`` is the state's ``generation`` - marker, which must not change how the route resolves.""" + """Stage one of two shapes of pre-canonical route for ``model_version``. + + The shapes are the original per-country Modal dicts, and a seeded + routing-state country route the registry ties to no wrapper bundle. For + the latter, ``source`` is the state's ``generation`` marker, which must + not change how the route resolves. + """ app_name = "policyengine-simulation-us1-715-2-uk2-88-20" if source == "legacy-country-dict": app_name = "legacy-app" @@ -380,9 +383,11 @@ def test_completed_budget_window_configuration_is_only_in_provenance( @pytest.mark.parametrize("source", LEGACY_SOURCES) def test_legacy_poll_bodies_carry_no_spm_key(mock_modal, client, source): - """The 202 and 500 bodies splat job metadata in raw, bypassing the - routes' ``response_model_exclude_none``. A legacy no-SPM job must stay - byte-identical to the pre-SPM gateway.""" + """Legacy no-SPM poll bodies must stay byte-identical to the pre-SPM gateway. + + The 202 and 500 bodies splat job metadata in raw, bypassing the routes' + ``response_model_exclude_none``. + """ legacy_route(mock_modal, source, "1.715.2") submitted = client.post("/simulate/economy/comparison", json={"country": "us"}) assert submitted.status_code == 200, submitted.text @@ -448,7 +453,7 @@ def test_publishing_does_not_revoke_a_seeded_country_route( def wrapper_route_without_manifest(mock_modal, wrapper): - """A seeded wrapper route the app-release snapshot had no bundle for.""" + """Seed a wrapper route the app-release snapshot had no bundle for.""" state = deepcopy(TEST_ROUTING_STATE) app_name = f"policyengine-simulation-py{wrapper.replace('.', '-')}" state["routes"]["policyengine"][wrapper] = app_name @@ -460,9 +465,12 @@ def wrapper_route_without_manifest(mock_modal, wrapper): def test_pinned_wrapper_route_without_manifest_submits_without_spm( mock_modal, client, endpoint, extra, wrapper ): - """With no manifest the response's ``model_version`` echoes the wrapper + """A pinned pre-canonical wrapper with no manifest still submits. + + With no manifest the response's ``model_version`` echoes the wrapper version. Reading that back as a country model version rejected the two - pinned pre-canonical bundles.""" + pinned pre-canonical bundles. + """ wrapper_route_without_manifest(mock_modal, wrapper) response = client.post( endpoint, json={"country": "us", "policyengine_version": wrapper, **extra} @@ -492,8 +500,11 @@ def test_pinned_wrapper_through_legacy_dicts_submits_without_spm( def test_unpinned_wrapper_route_without_manifest_still_fails_closed( mock_modal, client, endpoint, extra ): - """Only the two pinned pre-canonical wrappers are vouched for without a - manifest; an unrecognized bundle proves nothing.""" + """Fail closed for an unpinned wrapper route that has no manifest. + + Only the two pinned pre-canonical wrappers are vouched for without a + manifest; an unrecognized bundle proves nothing. + """ wrapper_route_without_manifest(mock_modal, "5.4.0") response = client.post( endpoint, json={"country": "us", "policyengine_version": "5.4.0", **extra} @@ -507,8 +518,10 @@ def test_unpinned_wrapper_route_without_manifest_still_fails_closed( def test_stated_model_version_still_contradicts_the_wrapper_pin( mock_modal, client, endpoint, extra ): - """A manifest that states a post-canonical model under a pinned wrapper - is a republish we cannot vouch for.""" + """Refuse a pinned wrapper whose manifest states a post-canonical model. + + That manifest is a republish we cannot vouch for. + """ state = deepcopy(TEST_ROUTING_STATE) original = state["bundles"]["4.10.0"] state["routes"]["policyengine"]["5.3.0"] = original["app_name"] @@ -527,7 +540,7 @@ def test_stated_model_version_still_contradicts_the_wrapper_pin( def republished_wrapper_state(mock_modal, *, model_version): - """The registry a re-publish leaves behind. + """Stage the registry a re-publish leaves behind. ``update_version_registry`` only ever adds country routes and it overwrites the bundle manifest for the wrapper version it deploys, so @@ -548,9 +561,12 @@ def republished_wrapper_state(mock_modal, *, model_version): def test_stale_country_route_is_refused_rather_than_served( mock_modal, client, endpoint, extra ): - """The behaviour change. Before, this route resolved and the 202 body - contradicted itself: ``version`` from the stale route, - ``policyengine_bundle.model_version`` from the live manifest.""" + """Refuse a stale country route instead of serving it (the behaviour change). + + Before, this route resolved and the 202 body contradicted itself: + ``version`` from the stale route, ``policyengine_bundle.model_version`` + from the live manifest. + """ republished_wrapper_state(mock_modal, model_version="1.470.0") response = client.post( endpoint, json={"country": "us", "version": "1.459.0", **extra} @@ -567,9 +583,11 @@ def test_stale_country_route_is_refused_rather_than_served( def test_a_country_route_its_bundle_still_states_is_served( mock_modal, client, endpoint, extra ): - """The other half: a route the manifest agrees with keeps working, so - the refusal above is about disagreement and not about single-candidate - country routes in general.""" + """Keep serving a country route the manifest agrees with. + + The other half: the refusal above is about disagreement and not about + single-candidate country routes in general. + """ republished_wrapper_state(mock_modal, model_version="1.459.0") response = client.post( endpoint, json={"country": "us", "version": "1.459.0", **extra} From 4464cf4077eb1dcab118be2c7e1ce949e49bb6d9 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 13 Oct 2026 16:56:47 +0400 Subject: [PATCH 07/18] Document the simulation executor for ruff D Covers the executor package, the Modal apps, fixtures, scripts and test docstrings. Docstrings are added to several Modal functions; images copy src/modal into a layer, so this changes those layers like any source edit. Protocol stubs keep their ... bodies. Co-Authored-By: Claude Opus 5.5 --- .../fixtures/fake_modal.py | 32 +++++++++ ...est_policyengine_package_update_scripts.py | 21 +++++- .../test_simulation_output_builder.py | 8 +++ .../fixtures/wrapper_spm.py | 1 + .../scripts/dump_package_version.py | 10 +++ .../src/modal/__init__.py | 3 +- .../src/modal/_image_setup.py | 3 +- .../src/modal/app.py | 12 ++-- .../src/modal/budget_window_batch.py | 5 ++ .../src/modal/budget_window_context.py | 10 +++ .../src/modal/budget_window_results.py | 16 +++-- .../src/modal/budget_window_scheduler.py | 26 ++++++++ .../src/modal/fanout.py | 7 +- .../src/modal/logging_redaction.py | 1 - .../src/modal/precompute_app.py | 10 +++ .../src/modal/segmented_national.py | 34 +++++++--- .../src/modal/smoke_app.py | 7 +- .../src/modal/static_runtime_files.py | 1 - .../modal/utils/extract_bundle_versions.py | 5 ++ .../src/modal/utils/extract_stage12_bundle.py | 5 ++ .../src/modal/utils/qualify_stage12_parity.py | 1 + .../src/modal/utils/record_deployment.py | 9 ++- .../modal/utils/update_v2_version_manifest.py | 9 +++ .../modal/utils/update_version_registry.py | 40 ++++++++++- .../src/modal/utils/validate_hf_access.py | 2 +- .../modal/utils/validate_v2_worker_release.py | 6 ++ .../src/modal/v2_app.py | 15 +++++ .../modal/v2_static_runtime_files_smoke.py | 7 +- .../__init__.py | 7 +- .../artifact_keys.py | 24 ++++++- .../artifact_store.py | 27 +++++++- .../baseline_artifacts.py | 14 +++- .../generate_openapi.py | 2 + .../hf_access_validation.py | 9 ++- .../policyengine_simulation_executor/main.py | 3 + .../national_partition.py | 2 +- .../precompute.py | 8 ++- .../precompute_models.py | 23 ++++++- .../release_bundle.py | 26 ++++++-- .../segmented_national_reduce.py | 2 +- .../settings.py | 8 +++ .../simulation.py | 6 ++ .../simulation_microdata.py | 19 +++--- .../simulation_output_budget.py | 11 ++++ .../simulation_output_builder.py | 14 ++++ .../simulation_output_cliff.py | 1 + .../simulation_output_distribution.py | 16 +++++ .../simulation_output_geographic.py | 27 +++++++- .../simulation_output_inequality.py | 4 ++ .../simulation_output_labor.py | 1 + .../simulation_output_poverty.py | 15 +++++ .../simulation_runtime.py | 29 ++++---- .../policyengine_simulation_executor/spm.py | 13 +++- .../stage12_artifacts.py | 43 +++++++++++- .../stage12_bundle/assertions.py | 5 +- .../stage12_bundle/manifest.py | 7 +- .../stage12_bundle/parsing.py | 8 ++- .../stage12_parity.py | 20 ++++++ .../stage12_qualification.py | 3 - .../stage12_result_comparison.py | 2 - .../stage12_runtime/aggregation.py | 16 ++++- .../stage12_runtime/comparison.py | 1 - .../stage12_runtime/coordination.py | 12 +++- .../stage12_runtime/dependencies.py | 65 ++++++++++++++---- .../stage12_runtime/output_planning.py | 10 +-- .../stage12_runtime/simulation.py | 15 ++++- .../stage12_worker_validation.py | 1 - .../uk_local_authority_metadata.py | 8 --- .../test_budget_window_ephemeral_modal.py | 1 - .../tests/test_app_redaction.py | 10 +-- .../tests/test_artifact_store.py | 7 +- .../tests/test_baseline_artifacts.py | 66 ++++++++++++------- .../tests/test_budget_window_batch.py | 17 ++--- .../tests/test_budget_window_context.py | 7 +- .../tests/test_budget_window_results.py | 10 +-- .../tests/test_budget_window_scheduler.py | 5 +- .../tests/test_image_setup_fetch.py | 37 ++++++----- .../tests/test_modal_bundle_image.py | 21 +++--- .../tests/test_modal_scripts.py | 24 +++---- .../tests/test_national_partition.py | 11 ++-- .../tests/test_precompute.py | 57 ++++++++++------ .../tests/test_precompute_app.py | 6 +- .../tests/test_release_bundle.py | 1 - .../tests/test_segmented_national.py | 10 +-- .../tests/test_segmented_national_reduce.py | 2 +- .../tests/test_simulation_output_builder.py | 7 +- .../tests/test_uk_weight_matrix_guard.py | 1 - 87 files changed, 878 insertions(+), 255 deletions(-) diff --git a/projects/policyengine-simulation-executor/fixtures/fake_modal.py b/projects/policyengine-simulation-executor/fixtures/fake_modal.py index 99bcc636b..5c48fd6bd 100644 --- a/projects/policyengine-simulation-executor/fixtures/fake_modal.py +++ b/projects/policyengine-simulation-executor/fixtures/fake_modal.py @@ -19,70 +19,95 @@ class FakeImage: + """Recording stand-in for ``modal.Image``: each builder call appends to ``calls``. + + Every method returns the same image, so a chained build leaves one ordered + list of ``(method_name, *arguments)`` tuples to assert layer order against. + """ + def __init__(self) -> None: + """Start with no recorded layers.""" self.calls: list[tuple[object, ...]] = [] @classmethod def debian_slim(cls, python_version: str) -> Self: + """Return a new image whose first recorded call is the base layer.""" image = cls() image.calls.append(("debian_slim", python_version)) return image def pip_install(self, *packages: str) -> Self: + """Record a ``pip_install`` layer with the package tuple.""" self.calls.append(("pip_install", packages)) return self def pip_install_from_requirements( self, requirements_txt: str, **kwargs: object ) -> Self: + """Record a requirements-file install layer.""" self.calls.append(("pip_install_from_requirements", requirements_txt, kwargs)) return self def uv_sync(self, uv_project_dir: str = "./", **kwargs: object) -> Self: + """Record a ``uv_sync`` layer with its project dir and options.""" self.calls.append(("uv_sync", uv_project_dir, kwargs)) return self def run_commands(self, *commands: str, **kwargs: object) -> Self: + """Record a ``run_commands`` layer with the command tuple and options.""" self.calls.append(("run_commands", commands, kwargs)) return self def env(self, env: Mapping[str, str]) -> Self: + """Record an ``env`` layer with the mapping as passed.""" self.calls.append(("env", env)) return self def add_local_python_source(self, *args: str, **kwargs: object) -> Self: + """Record the source packages added and their options (e.g. ``copy``).""" self.calls.append(("add_local_python_source", args, kwargs)) return self def add_local_file( self, local_path: str | Path, remote_path: str, **kwargs: object ) -> Self: + """Record a single file copied into the image.""" self.calls.append(("add_local_file", local_path, remote_path, kwargs)) return self def add_local_dir( self, local_path: str | Path, remote_path: str, **kwargs: object ) -> Self: + """Record a directory copied into the image.""" self.calls.append(("add_local_dir", local_path, remote_path, kwargs)) return self def run_function(self, function: Callable[..., object], **kwargs: object) -> Self: + """Record a build-time function layer by function name, without running it.""" self.calls.append(("run_function", function.__name__, kwargs)) return self class FakeSecret: + """Stand-in for ``modal.Secret`` that resolves nothing.""" + @staticmethod def from_name(*args: object, **kwargs: object) -> dict[str, object]: + """Return the lookup arguments so tests can assert which secret was named.""" return {"args": args, "kwargs": kwargs} class FakeApp: + """Stand-in for ``modal.App`` that records ``function`` decorator options.""" + def __init__(self, name: str) -> None: + """Name the app and start with no recorded function declarations.""" self.name = name self.function_calls: list[tuple[str, dict[str, object]]] = [] def function[F: Callable[..., object]](self, **kwargs: object) -> Callable[[F], F]: + """Return a decorator recording ``(name, kwargs)`` that returns the function.""" + def decorator(function: F) -> F: self.function_calls.append((function.__name__, kwargs)) return function @@ -92,6 +117,8 @@ def decorator(function: F) -> F: def local_entrypoint[F: Callable[..., object]]( self, **kwargs: object ) -> Callable[[F], F]: + """Return a pass-through decorator; entrypoints are not recorded.""" + def decorator(function: F) -> F: return function @@ -99,6 +126,11 @@ def decorator(function: F) -> F: def install_fake_modal(monkeypatch: pytest.MonkeyPatch) -> None: + """Install the recording fakes as ``sys.modules["modal"]`` for one test. + + ``is_local()`` reports True and ``asgi_app`` is a pass-through decorator. + Modules that already imported the real ``modal`` must be re-imported. + """ modal = ModuleType("modal") modal.Image = FakeImage modal.Secret = FakeSecret diff --git a/projects/policyengine-simulation-executor/fixtures/test_policyengine_package_update_scripts.py b/projects/policyengine-simulation-executor/fixtures/test_policyengine_package_update_scripts.py index 00b0b33be..84b849f4c 100644 --- a/projects/policyengine-simulation-executor/fixtures/test_policyengine_package_update_scripts.py +++ b/projects/policyengine-simulation-executor/fixtures/test_policyengine_package_update_scripts.py @@ -18,6 +18,7 @@ @pytest.fixture def fake_repo(tmp_path: Path) -> Path: + """Return a fake checkout whose ``simulation`` project pins policyengine 4.0.0.""" project = tmp_path / "simulation" project.mkdir(parents=True) @@ -59,12 +60,14 @@ def fake_repo(tmp_path: Path) -> Path: @pytest.fixture def fake_bin(tmp_path: Path) -> Path: + """Return an empty directory for fake executables, to put first on PATH.""" path = tmp_path / "bin" path.mkdir() return path def write_executable(path: Path, content: str) -> None: + """Write ``content`` to ``path`` and mark it executable.""" path.write_text(content, encoding="utf-8") path.chmod(0o755) @@ -78,8 +81,12 @@ def install_fake_git( diff_has_changes: bool = False, branch_pin: str = "4.1.0", ) -> None: - """A fake git; ``show`` prints the update branch's pyproject pinning ``branch_pin``.""" + """Install a fake ``git`` that appends each call to ``log``. + ``remote_branch_exists`` and ``diff_has_changes`` set the ``ls-remote`` + and ``diff`` exit codes, and ``show`` prints the update branch's pyproject + pinning ``branch_pin``; every other command succeeds silently. + """ write_executable( fake_bin / "git", f"""#!/usr/bin/env bash @@ -122,6 +129,11 @@ def install_fake_gh( open_pr: str = "", open_issue: str = "679", ) -> None: + """Install a fake ``gh`` that appends each call to ``log`` and answers canned JSON. + + ``open_issue`` and ``open_pr`` (empty for none) drive the existing issue + and PR lookups; the first line of a created PR's body is logged too. + """ write_executable( fake_bin / "gh", f"""#!/usr/bin/env bash @@ -190,6 +202,11 @@ def install_fake_uv( bundled_spm_version: str = "1.0.0", bundled_uk_dataset_sha256: str = "a" * 64, ) -> None: + """Install a fake ``uv`` that appends each call to ``log``. + + Only the updater's isolated ``bundle-versions`` query for policyengine + 4.1.0 prints anything: the ``bundled_*`` versions as ``key=value`` lines. + """ write_executable( fake_bin / "uv", f"""#!/usr/bin/env bash @@ -219,6 +236,7 @@ def install_fake_uv( def updater_env(fake_bin: Path, **extra: str) -> dict[str, str]: + """Return this environment with ``fake_bin`` first on PATH, plus ``extra``.""" env = os.environ.copy() env.update( { @@ -231,6 +249,7 @@ def updater_env(fake_bin: Path, **extra: str) -> dict[str, str]: def run_updater(*args: str, env: dict[str, str]) -> subprocess.CompletedProcess[str]: + """Run the updater from the repo root; a non-zero exit is returned, not raised.""" return subprocess.run( ["bash", str(SCRIPT), *args], cwd=REPO_ROOT, diff --git a/projects/policyengine-simulation-executor/fixtures/test_simulation_output_builder.py b/projects/policyengine-simulation-executor/fixtures/test_simulation_output_builder.py index ccc591b04..2d091ee87 100644 --- a/projects/policyengine-simulation-executor/fixtures/test_simulation_output_builder.py +++ b/projects/policyengine-simulation-executor/fixtures/test_simulation_output_builder.py @@ -8,20 +8,28 @@ class FakeCollection: + """Stand-in for a v4 output collection: only its ``dataframe`` is read.""" + def __init__(self, records: list[dict[str, object]]) -> None: + """Expose ``records`` as the collection's ``dataframe``.""" self.dataframe = pd.DataFrame(records) class FakeModelOutput: + """Stand-in for a v4 Pydantic output that the adapter dumps to JSON.""" + def __init__(self, payload: dict[str, object]) -> None: + """Hold the payload ``model_dump`` returns.""" self.payload = payload def model_dump(self, *, mode: str) -> dict[str, object]: + """Return the payload; fails unless the adapter asks for JSON mode.""" assert mode == "json" return self.payload def fake_analysis() -> SimpleNamespace: + """Return a small v4 ``economic_impact_analysis`` result with fixed known values.""" return SimpleNamespace( program_statistics=FakeCollection( [ diff --git a/projects/policyengine-simulation-executor/fixtures/wrapper_spm.py b/projects/policyengine-simulation-executor/fixtures/wrapper_spm.py index e8ddc4226..b5643c223 100644 --- a/projects/policyengine-simulation-executor/fixtures/wrapper_spm.py +++ b/projects/policyengine-simulation-executor/fixtures/wrapper_spm.py @@ -33,6 +33,7 @@ def storage_id(self) -> str: def wrapper_storage_id(simulation_id: str, spm_config: dict | None) -> str: + """Return the wrapper's ``storage_id`` for an already-resolved SPM config.""" if spm_config is None: return simulation_id encoded = json.dumps(spm_config, sort_keys=True, separators=(",", ":")).encode() diff --git a/projects/policyengine-simulation-executor/scripts/dump_package_version.py b/projects/policyengine-simulation-executor/scripts/dump_package_version.py index de7a603de..f704f5fd8 100644 --- a/projects/policyengine-simulation-executor/scripts/dump_package_version.py +++ b/projects/policyengine-simulation-executor/scripts/dump_package_version.py @@ -1,3 +1,8 @@ +"""Print a package's locked version: ``dump_package_version.py ``. + +Exits 1, printing nothing, when the lock file or package is missing. +""" + import sys import tomllib @@ -6,6 +11,11 @@ def get_package_version_from_lock( package_name: str, lock_file_path: str = "poetry.lock" ) -> str | None: + """Return ``package_name``'s version from a TOML lock file, or None if absent. + + A missing lock file also returns None; uv and poetry locks share the + ``[[package]]`` name/version layout read here. + """ try: with open(lock_file_path, "rb") as f: lock_data = tomllib.load(f) diff --git a/projects/policyengine-simulation-executor/src/modal/__init__.py b/projects/policyengine-simulation-executor/src/modal/__init__.py index 099fa6f22..04f84d996 100644 --- a/projects/policyengine-simulation-executor/src/modal/__init__.py +++ b/projects/policyengine-simulation-executor/src/modal/__init__.py @@ -1,5 +1,4 @@ -""" -Modal deployment package for PolicyEngine Simulation API. +"""Modal deployment package for PolicyEngine Simulation API. Deploy with: modal deploy src/modal/app.py diff --git a/projects/policyengine-simulation-executor/src/modal/_image_setup.py b/projects/policyengine-simulation-executor/src/modal/_image_setup.py index 54788bb6a..4451ee676 100644 --- a/projects/policyengine-simulation-executor/src/modal/_image_setup.py +++ b/projects/policyengine-simulation-executor/src/modal/_image_setup.py @@ -1,5 +1,4 @@ -""" -Standalone image setup functions. +"""Standalone image setup functions. These functions are executed during Modal image build and must not import any other modules from this package to avoid dependency issues. diff --git a/projects/policyengine-simulation-executor/src/modal/app.py b/projects/policyengine-simulation-executor/src/modal/app.py index b30670187..b7abd3540 100644 --- a/projects/policyengine-simulation-executor/src/modal/app.py +++ b/projects/policyengine-simulation-executor/src/modal/app.py @@ -1,5 +1,4 @@ -""" -PolicyEngine Simulation - Versioned Modal App +"""PolicyEngine Simulation versioned Modal app. This app contains the heavy simulation workload with snapshotted models. Each deployment creates a versioned app (e.g., policyengine-simulation-py4-10-0). @@ -95,8 +94,7 @@ def _version_from_env_or_local_bundle(env_var: str, country: str) -> str: def get_app_name(policyengine_version: str) -> str: - """ - Generate versioned app name from the policyengine.py package version. + """Generate versioned app name from the policyengine.py package version. Replaces dots with dashes for URL safety. Example: 4.10.0 -> policyengine-simulation-py4-10-0 @@ -171,8 +169,7 @@ def _deploy_time_artifact_inputs() -> tuple[str, dict | None]: def build_runtime_simulation_image() -> modal.Image: - """Image layers up to the version env — everything except the artifact - fetch and model snapshot. + """Build the image layers up to the version env (no artifact fetch or snapshot). Shared by the deployed app and the image smoke app (src/modal/smoke_app.py): all must construct these layers through @@ -276,8 +273,7 @@ def run_simulation( *, observability_context: dict | None = None, ) -> dict: - """ - Execute economic simulation. + """Execute economic simulation. Imports the snapshotted implementation at runtime. Emits redacted operation data through the shared observability runtime. diff --git a/projects/policyengine-simulation-executor/src/modal/budget_window_batch.py b/projects/policyengine-simulation-executor/src/modal/budget_window_batch.py index e83286c69..285b9ad1b 100644 --- a/projects/policyengine-simulation-executor/src/modal/budget_window_batch.py +++ b/projects/policyengine-simulation-executor/src/modal/budget_window_batch.py @@ -19,6 +19,11 @@ def run_budget_window_batch_impl( *, runtime: ObservabilityRuntime, ) -> dict[str, Any]: + """Run a budget-window batch to completion and return its final status payload. + + Must run inside a Modal call: the current function call id becomes the + batch job id that child jobs and stored batch state are keyed by. + """ from policyengine_simulation_executor.spm import normalize_runtime_spm selection = normalize_runtime_spm(params) diff --git a/projects/policyengine-simulation-executor/src/modal/budget_window_context.py b/projects/policyengine-simulation-executor/src/modal/budget_window_context.py index ce91f8073..6d8199725 100644 --- a/projects/policyengine-simulation-executor/src/modal/budget_window_context.py +++ b/projects/policyengine-simulation-executor/src/modal/budget_window_context.py @@ -56,6 +56,11 @@ def build_batch_context( *, batch_job_id: str, ) -> BudgetWindowBatchContext: + """Validate batch params and the gateway's routing ``_metadata`` into a context. + + Raises ValueError when ``_metadata`` lacks the resolved app name, the + resolved version, or the PolicyEngine bundle. + """ request = BudgetWindowBatchRequest.model_validate(params) metadata = params.get("_metadata") if not isinstance(metadata, dict): @@ -87,6 +92,11 @@ def build_child_simulation_request( *, simulation_year: str, ) -> ChildSimulationRequest: + """Build the single-year child request for ``simulation_year``. + + Batch-only fields are stripped and the child is pinned to a monolithic + (non-segmented) run. + """ payload = build_child_payload( context.raw_params, strip_fields=BATCH_ONLY_FIELDS, diff --git a/projects/policyengine-simulation-executor/src/modal/budget_window_results.py b/projects/policyengine-simulation-executor/src/modal/budget_window_results.py index 02464c0dc..719a5651b 100644 --- a/projects/policyengine-simulation-executor/src/modal/budget_window_results.py +++ b/projects/policyengine-simulation-executor/src/modal/budget_window_results.py @@ -25,11 +25,11 @@ def _as_decimal(value: float) -> Decimal: - """Convert an annual impact float to Decimal without reintroducing - binary-float quantisation noise. ``Decimal(str(...))`` is the canonical - idiom because it serialises the float to its shortest round-trippable - decimal form before parsing.""" + """Convert an annual impact float to Decimal without binary-float noise. + ``Decimal(str(...))`` is the canonical idiom because it serialises the + float to its shortest round-trippable decimal form before parsing. + """ return Decimal(str(value)) @@ -39,6 +39,12 @@ def extract_annual_impact( child_result: dict[str, Any], spm: object = None, ) -> BudgetWindowAnnualImpact: + """Extract one year's budget impact from a child simulation result. + + Raises when the result's SPM receipt does not match ``spm`` for + ``simulation_year``, or a required budget figure is missing or + non-numeric. A missing state tax revenue impact (UK results) counts as 0. + """ receipt = validate_spm_result(child_result, spm, expected_year=simulation_year) budget = child_result.get("budget", {}) if not isinstance(budget, dict): @@ -88,7 +94,6 @@ def sum_annual_impacts( bit-exact accounting should request the individual per-year impacts and sum them in their preferred numeric type. """ - totals: dict[str, Decimal] = { "taxRevenueImpact": Decimal(0), "federalTaxRevenueImpact": Decimal(0), @@ -119,6 +124,7 @@ def build_budget_window_result( window_size: int, annual_impacts: list[BudgetWindowAnnualImpact], ) -> BudgetWindowResult: + """Assemble the window result, with exact Decimal-summed totals.""" return BudgetWindowResult( startYear=start_year, endYear=str(int(start_year) + window_size - 1), diff --git a/projects/policyengine-simulation-executor/src/modal/budget_window_scheduler.py b/projects/policyengine-simulation-executor/src/modal/budget_window_scheduler.py index eb0596a9f..8d4e95c56 100644 --- a/projects/policyengine-simulation-executor/src/modal/budget_window_scheduler.py +++ b/projects/policyengine-simulation-executor/src/modal/budget_window_scheduler.py @@ -63,6 +63,7 @@ def serialize_batch_status( *, runtime: ObservabilityRuntime, ) -> dict[str, Any]: + """Render ``state`` as the JSON-ready batch status response.""" with runtime.span( BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_STATUS_SERIALIZATION) ): @@ -74,6 +75,7 @@ def load_or_create_batch_state( *, runtime: ObservabilityRuntime, ) -> BudgetWindowBatchState: + """Return the batch's submit-time seed state, creating and storing it if absent.""" with runtime.span(BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_STATE_LOAD)): state = get_batch_job_seed(context.batch_job_id) if state is None: @@ -102,6 +104,11 @@ def __init__( poll_interval_max_seconds: float = POLL_INTERVAL_MAX_SECONDS, poll_interval_backoff_factor: float = POLL_INTERVAL_BACKOFF_FACTOR, ) -> None: + """Load or seed the batch state and resolve a lazy child-function handle. + + Children run on the resolved app's ``run_simulation`` function. + ``modal_module`` replaces the ``modal`` module in tests. + """ self.context = context self.runtime = runtime self.modal = modal if modal_module is None else modal_module @@ -135,6 +142,12 @@ def __init__( self._backoff_sleep_count = 0 def run(self) -> dict[str, Any]: + """Spawn and poll children until all finish or one fails; return the status. + + At most ``max_parallel`` children run at once. Polling backs off + exponentially while no child makes progress, and the first child + failure fails the whole batch. + """ mark_batch_running(self.state) with self.runtime.span( BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_STATE_WRITE) @@ -164,9 +177,14 @@ def run(self) -> dict[str, Any]: return self.complete_batch() def has_pending_work(self) -> bool: + """Return True while any year is queued or running.""" return bool(self.state.queued_years or self.state.running_years) def spawn_until_capacity(self) -> None: + """Spawn queued years until ``max_parallel`` children are running. + + Each spawn is recorded in the stored batch state before the next one. + """ while ( len(self.state.running_years) < self.state.max_parallel and self.state.queued_years @@ -205,6 +223,11 @@ def spawn_until_capacity(self) -> None: put_batch_job_state(self.state) def poll_running_children_once(self) -> bool: + """Probe each running child once without blocking; return True if any finished. + + A failed child, or a result that cannot be parsed, fails the batch and + returns False immediately. + """ progress_made = False for simulation_year in list(self.state.running_years): @@ -293,6 +316,7 @@ def poll_running_children_once(self) -> bool: return progress_made def resolve_child_handle(self, simulation_year: str) -> ChildSimulationHandle: + """Return the year's call handle, rebuilt from the stored job id if needed.""" handle = self.child_handles.get(simulation_year) if handle is not None and handle.call is not None: return handle @@ -324,6 +348,7 @@ def fail_batch_for_child_error( simulation_year: str, error: str, ) -> None: + """Mark the year's child and the whole batch failed and store the state.""" mark_child_failed(self.state, year=simulation_year, error=error) mark_batch_failed(self.state, error=error) with self.runtime.span( @@ -332,6 +357,7 @@ def fail_batch_for_child_error( put_batch_job_state(self.state) def complete_batch(self) -> dict[str, Any]: + """Aggregate annual impacts, mark the batch complete, and return its status.""" with self.runtime.span( BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_AGGREGATION) ): diff --git a/projects/policyengine-simulation-executor/src/modal/fanout.py b/projects/policyengine-simulation-executor/src/modal/fanout.py index 13ce29bcb..42a9d4f12 100644 --- a/projects/policyengine-simulation-executor/src/modal/fanout.py +++ b/projects/policyengine-simulation-executor/src/modal/fanout.py @@ -22,6 +22,11 @@ def build_child_payload( strip_fields: Iterable[str], overrides: dict[str, Any], ) -> dict[str, Any]: + """Copy parent params minus ``strip_fields`` and apply the child ``overrides``. + + ``_telemetry`` is always re-attached from the parent (when it is a dict), + even if ``overrides`` sets it, so children join the parent's trace. + """ strip = set(strip_fields) | {"_telemetry"} payload = {key: value for key, value in raw_params.items() if key not in strip} payload.update(overrides) @@ -32,5 +37,5 @@ def build_child_payload( def next_backoff(current: float, *, factor: float, maximum: float) -> float: - """The next sleep in an exponential backoff walk.""" + """Return the next sleep in an exponential backoff walk, capped at ``maximum``.""" return min(current * factor, maximum) diff --git a/projects/policyengine-simulation-executor/src/modal/logging_redaction.py b/projects/policyengine-simulation-executor/src/modal/logging_redaction.py index 4f2542b51..152a952ac 100644 --- a/projects/policyengine-simulation-executor/src/modal/logging_redaction.py +++ b/projects/policyengine-simulation-executor/src/modal/logging_redaction.py @@ -31,7 +31,6 @@ def redact_params_for_logging(params: object) -> dict: dict so callers can splat the result into operation attributes without additional guards. """ - if not isinstance(params, dict): return {} return { diff --git a/projects/policyengine-simulation-executor/src/modal/precompute_app.py b/projects/policyengine-simulation-executor/src/modal/precompute_app.py index 1314b0f4c..a83b801c5 100644 --- a/projects/policyengine-simulation-executor/src/modal/precompute_app.py +++ b/projects/policyengine-simulation-executor/src/modal/precompute_app.py @@ -53,6 +53,7 @@ secrets=_worker_secrets, ) def plan_artifacts(bucket: str) -> dict: + """Return the dumped ``PrecomputePlan``: expected artifacts and store presence.""" from policyengine_simulation_executor.precompute import plan_artifacts_impl return plan_artifacts_impl(bucket).model_dump() @@ -66,6 +67,11 @@ def plan_artifacts(bucket: str) -> dict: secrets=_worker_secrets, ) def build_dataset(bucket: str, expected: dict) -> dict: + """Build and upload one single-year US dataset (wave 1). + + Raises if the dataset identity resolved in this container differs from the + planned one, so planner/worker drift cannot upload under the wrong key. + """ from policyengine_simulation_executor.precompute import build_dataset_impl from policyengine_simulation_executor.precompute_models import DatasetPlanEntry @@ -82,6 +88,7 @@ def build_dataset(bucket: str, expected: dict) -> dict: secrets=_worker_secrets, ) def compute_baseline(bucket: str, expected: dict) -> dict: + """Compute and upload one cohort's national baseline output (wave 2).""" from policyengine_simulation_executor.precompute import compute_baseline_impl from policyengine_simulation_executor.precompute_models import BaselinePlanEntry @@ -98,6 +105,7 @@ def compute_baseline(bucket: str, expected: dict) -> dict: secrets=_worker_secrets, ) def verify_determinism(bucket: str, expected: dict) -> dict: + """Recompute one stored cohort baseline afresh and report columns that differ.""" from policyengine_simulation_executor.precompute import verify_determinism_impl from policyengine_simulation_executor.precompute_models import BaselinePlanEntry @@ -114,6 +122,7 @@ def verify_determinism(bucket: str, expected: dict) -> dict: secrets=_worker_secrets, ) def publish_manifest(bucket: str, plan: dict) -> str: + """Write the deploy manifest for ``plan`` to the store and return its digest.""" from policyengine_simulation_executor.precompute import publish_manifest_impl from policyengine_simulation_executor.precompute_models import PrecomputePlan @@ -122,6 +131,7 @@ def publish_manifest(bucket: str, plan: dict) -> str: @app.local_entrypoint() def main(force: bool = False) -> None: + """Fill the store for the installed bundle; ``--force`` recomputes stored ones.""" from policyengine_simulation_executor.artifact_store import resolve_bucket_name from policyengine_simulation_executor.precompute import run_precompute diff --git a/projects/policyengine-simulation-executor/src/modal/segmented_national.py b/projects/policyengine-simulation-executor/src/modal/segmented_national.py index a2e25e823..01f414006 100644 --- a/projects/policyengine-simulation-executor/src/modal/segmented_national.py +++ b/projects/policyengine-simulation-executor/src/modal/segmented_national.py @@ -87,9 +87,9 @@ def _activates_labor_supply_response(params: dict[str, Any]) -> bool: def is_plain_national_macro(params: dict[str, Any]) -> bool: - """The request SHAPE of a plain US national macro run (ignores the - eligibility knobs — see ``should_run_segmented_national``). + """Return whether ``params`` has the request SHAPE of a plain US national macro run. + Ignores the eligibility knobs (see ``should_run_segmented_national``). National is spelled two ways in the wild: region omitted (the sim API's own convention) and ``region: "us"`` (what API v1 sends on every national run). Both must match, mirroring the worker's normalisation @@ -107,7 +107,7 @@ def is_plain_national_macro(params: dict[str, Any]) -> bool: def should_run_segmented_national(params: dict[str, Any]) -> bool: - """True when a request should fan out across the national partition. + """Return True when a request should fan out across the national partition. Eligible: a plain US national macro shape, not opted out via ``segmented: false``. Falls back to monolithic for ``include_cliffs`` @@ -150,6 +150,10 @@ def __init__( poll_interval_seconds: float = POLL_INTERVAL_INITIAL_SECONDS, poll_interval_max_seconds: float = POLL_INTERVAL_MAX_SECONDS, ) -> None: + """Resolve the country's region groups and a lazy segment-worker handle. + + Raises ValueError when the country has no national partition. + """ self.params = params self.app_name = app_name self.runtime = runtime @@ -173,6 +177,10 @@ def __init__( self.runtime.set_context(simulation_year=str(params["time_period"])) def run(self) -> dict[str, Any]: + """Spawn a child per region group, wait for all, and reduce to national output. + + Any child failure cancels the remaining children and raises. + """ country_module = self._country_module() self._extend_last_group_with_uncovered_states(country_module) handles = self._spawn_all_children() @@ -189,11 +197,13 @@ def _country_module(self) -> ModuleType: def _extend_last_group_with_uncovered_states( self, country_module: ModuleType ) -> None: - """Registry state regions missing from the static partition ride in - the last group, loudly, so a new state-level region added to the - model cannot silently drop its households from national results. - (Households whose state has NO registered region remain invisible - here; the partition test pins the registry at pin-bump time.) + """Append registry state regions missing from the partition to the last group. + + They ride along loudly (a warning and a run attribute) so a new + state-level region added to the model cannot silently drop its + households from national results. (Households whose state has NO + registered region remain invisible here; the partition test pins the + registry at pin-bump time.) """ # region_registry is populated eagerly in the model version's # __init__; None only occurs for stub models (unit tests). @@ -349,6 +359,7 @@ def _reduce( def run_segmented_national_impl( params: dict[str, Any], *, app_name: str, runtime: ObservabilityRuntime ) -> dict[str, Any]: + """Run ``params`` as a segmented national fan-out on ``app_name``'s segment pool.""" return SegmentedNationalRunner( params, app_name=app_name, @@ -362,8 +373,11 @@ def dispatch_run_simulation( app_name: str, runtime: ObservabilityRuntime, ) -> dict[str, Any]: - """The run_simulation entrypoint's routing: segmented national fan-out - for eligible requests, the monolithic path for everything else.""" + """Route a ``run_simulation`` request to segmented fan-out or the monolithic path. + + Segmented national fan-out serves eligible requests (see + ``should_run_segmented_national``); everything else runs monolithically. + """ from policyengine_simulation_executor.spm import normalize_runtime_spm selection = normalize_runtime_spm(params) diff --git a/projects/policyengine-simulation-executor/src/modal/smoke_app.py b/projects/policyengine-simulation-executor/src/modal/smoke_app.py index 78ecc52a1..f70815a48 100644 --- a/projects/policyengine-simulation-executor/src/modal/smoke_app.py +++ b/projects/policyengine-simulation-executor/src/modal/smoke_app.py @@ -97,7 +97,6 @@ def _summarize_bundle_status(status: dict) -> dict: def _validate_installed_uk_local_authority_dataset() -> str: """Require the installed UK dataset to match one packaged boundary version.""" - dataset_path = resolve_local_bundle_dataset_path("uk", None) if dataset_path is None: raise RuntimeError("installed certified UK dataset is unavailable") @@ -107,6 +106,11 @@ def _validate_installed_uk_local_authority_dataset() -> str: @app.function(image=smoke_image, timeout=600, memory=8192) def smoke_import_executor() -> dict: + """Check the installed bundle, then import the workers' lazy request-time modules. + + Raises on any bundle mismatch or import failure; otherwise returns the + bundle summary and how many shared-library modules imported. + """ import importlib import os import pkgutil @@ -162,6 +166,7 @@ def smoke_import_executor() -> dict: @app.local_entrypoint() def main() -> None: + """Run the image smoke remotely and print its report.""" report = smoke_import_executor.remote() print(report) print("executor image smoke OK") diff --git a/projects/policyengine-simulation-executor/src/modal/static_runtime_files.py b/projects/policyengine-simulation-executor/src/modal/static_runtime_files.py index 5ce0e2365..13478e201 100644 --- a/projects/policyengine-simulation-executor/src/modal/static_runtime_files.py +++ b/projects/policyengine-simulation-executor/src/modal/static_runtime_files.py @@ -16,7 +16,6 @@ def add_static_runtime_files( uv_project_dir: str, ) -> modal.Image: """Copy executor static files into a Modal image at the importable path.""" - local_directory = ( str( Path(uv_project_dir) diff --git a/projects/policyengine-simulation-executor/src/modal/utils/extract_bundle_versions.py b/projects/policyengine-simulation-executor/src/modal/utils/extract_bundle_versions.py index d66797fe3..4e41068fa 100644 --- a/projects/policyengine-simulation-executor/src/modal/utils/extract_bundle_versions.py +++ b/projects/policyengine-simulation-executor/src/modal/utils/extract_bundle_versions.py @@ -30,6 +30,11 @@ def _bundle_outputs() -> dict[str, str]: def main() -> None: + """Report the installed bundle's package and data versions. + + ``--shell`` prints only ``key=value`` lines. Otherwise the pairs are + appended to ``$GITHUB_OUTPUT`` when it is set and a summary line is printed. + """ outputs = _bundle_outputs() if "--shell" in sys.argv[1:]: diff --git a/projects/policyengine-simulation-executor/src/modal/utils/extract_stage12_bundle.py b/projects/policyengine-simulation-executor/src/modal/utils/extract_stage12_bundle.py index 0f5f4c3a9..881a54fd2 100644 --- a/projects/policyengine-simulation-executor/src/modal/utils/extract_stage12_bundle.py +++ b/projects/policyengine-simulation-executor/src/modal/utils/extract_stage12_bundle.py @@ -25,6 +25,11 @@ def _expectations(values: list[str]) -> dict[str, str]: def main() -> None: + """Check ``--expect`` assertions, then emit the resolved bundle as canonical JSON. + + The JSON goes to stdout, or to ``--output`` with only the bundle manifest + digest printed. + """ parser = argparse.ArgumentParser( description="Export the packaged PolicyEngine.py bundle for Stage 12", ) diff --git a/projects/policyengine-simulation-executor/src/modal/utils/qualify_stage12_parity.py b/projects/policyengine-simulation-executor/src/modal/utils/qualify_stage12_parity.py index 324c8fe5e..cc6085d4b 100644 --- a/projects/policyengine-simulation-executor/src/modal/utils/qualify_stage12_parity.py +++ b/projects/policyengine-simulation-executor/src/modal/utils/qualify_stage12_parity.py @@ -39,6 +39,7 @@ def _tolerances(path: Path | None) -> dict[str, NumericalTolerance]: def main() -> None: + """Qualify one report against the existing executor and print the receipt JSON.""" parser = argparse.ArgumentParser( description=( "Compare the existing combined executor with the Stage 12 independent " diff --git a/projects/policyengine-simulation-executor/src/modal/utils/record_deployment.py b/projects/policyengine-simulation-executor/src/modal/utils/record_deployment.py index a5d9fad46..3b0426453 100644 --- a/projects/policyengine-simulation-executor/src/modal/utils/record_deployment.py +++ b/projects/policyengine-simulation-executor/src/modal/utils/record_deployment.py @@ -25,9 +25,11 @@ def build_marker( manifest_digest: str, env: Mapping[str, str], ) -> DeployedMarker: - """Pure marker assembly: the deploy identity that referenced this - manifest. Missing env vars become empty fields, never errors — the - marker write itself is the part that must not fail silently.""" + """Assemble the marker for the deploy that referenced this manifest, without I/O. + + Missing env vars become empty fields, never errors — the marker write + itself is the part that must not fail silently. + """ server = env.get("GITHUB_SERVER_URL", "") repository = env.get("GITHUB_REPOSITORY", "") run_id = env.get("GITHUB_RUN_ID", "") @@ -52,6 +54,7 @@ def build_marker( def main() -> None: + """Write this deploy's ``deployed/.json`` marker to the store.""" parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--environment", required=True) parser.add_argument("--manifest-digest", required=True) diff --git a/projects/policyengine-simulation-executor/src/modal/utils/update_v2_version_manifest.py b/projects/policyengine-simulation-executor/src/modal/utils/update_v2_version_manifest.py index 4717bb628..c695073e7 100644 --- a/projects/policyengine-simulation-executor/src/modal/utils/update_v2_version_manifest.py +++ b/projects/policyengine-simulation-executor/src/modal/utils/update_v2_version_manifest.py @@ -22,6 +22,11 @@ def build_worker(validation_payload: object) -> V2WorkerVersion: + """Build the installed bundle's v2 worker entry from a validation record. + + Raises ValueError if the record is invalid or the entry carries a + credential-like key. + """ resolved = load_stage12_bundle() validation = V2WorkerValidation.model_validate(validation_payload) worker = V2WorkerVersion( @@ -43,6 +48,10 @@ def build_worker(validation_payload: object) -> V2WorkerVersion: def main() -> None: + """Publish the validated worker to the v2 manifest Dict and print the manifest. + + Refuses a ``--manifest-name`` that is empty or names the v1 routing state. + """ parser = argparse.ArgumentParser( description="Publish a validated Stage 12 v2 worker manifest", ) 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 b6cf9069b..ae24844be 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 @@ -32,6 +32,8 @@ class CountryBundleMetadata(TypedDict): + """One country's model package, data release and dataset URIs in a bundle.""" + country: str model_package_name: str model_version: str @@ -47,6 +49,11 @@ class CountryBundleMetadata(TypedDict): class BundleManifestMetadata(TypedDict): + """The release bundle served by one deployed app, keyed by its policyengine version. + + ``spm`` is present only when the runtime reports an SPM capability. + """ + app_name: str policyengine_version: str us: CountryBundleMetadata @@ -55,18 +62,24 @@ class BundleManifestMetadata(TypedDict): class LatestVersions(TypedDict): + """The versions that unpinned requests route to, per route kind.""" + policyengine: str us: str uk: str class RouteMaps(TypedDict): + """Version-to-app-name routes for each route kind.""" + policyengine: dict[str, str] us: dict[str, str] uk: dict[str, str] class RoutingState(TypedDict): + """The routing snapshot stored as one value under the routing Dict's active key.""" + schema_version: int generation: str latest: LatestVersions @@ -83,7 +96,6 @@ def _is_newer_version(candidate: str, current: str | None) -> bool: opted in via ``--force-latest``. That decision is made by the caller; this helper answers the strict-greater-than question for valid versions. """ - if current is None: return True try: @@ -119,6 +131,7 @@ def build_bundle_manifest_metadata( app_name: str, policyengine_version: str, ) -> BundleManifestMetadata: + """Describe the installed release bundle as served by ``app_name``.""" from policyengine_simulation_executor.spm import runtime_spm_capability capability = runtime_spm_capability() @@ -193,6 +206,12 @@ def validate_routing_state( *, required_policyengine_manifests: Iterable[str] = (), ) -> None: + """Raise ValueError unless ``state`` is a consistent schema-1 routing snapshot. + + Each required ``latest`` entry must have a route, and every stored bundle + manifest must match its policyengine route. A policyengine route may lack + a manifest unless its version is in ``required_policyengine_manifests``. + """ if state.get("schema_version") != 1: raise ValueError("Routing state schema_version must be 1") @@ -253,6 +272,12 @@ def build_next_routing_state( uk_version: str, force_latest: bool = False, ) -> RoutingState: + """Return ``current_state`` with routes and the bundle manifest for a new app added. + + ``latest`` advances only to a newer policyengine version, unless + ``force_latest`` (an intentional rollback). Raises ValueError when the US + or UK version differs from the installed bundle's model version. + """ manifest = build_bundle_manifest_metadata( app_name=app_name, policyengine_version=policyengine_version, @@ -867,6 +892,12 @@ def build_legacy_seed_routing_state( uk_versions: dict, app_release_bundles: dict | None = None, ) -> RoutingState: + """Build a routing state from snapshots of the legacy per-kind version Dicts. + + Policyengine routes missing from the legacy Dict are inferred from + country app names, and a missing policyengine ``latest`` is inferred from + the countries' latest apps, falling back to the newest route. + """ state = _empty_routing_state() state["generation"] = "legacy-seed" state["routes"]["us"] = _version_routes(us_versions) @@ -937,6 +968,11 @@ def _merge_legacy_seed_into_current( def seed_active_routing_state_from_legacy(*, environment: str) -> RoutingState: + """Seed the active routing state from the legacy Dicts and store it. + + Routes and ``latest`` entries already in the active state win over the + legacy seed. + """ _require_serialized_routing_workflow() store = modal.Dict.from_name( ROUTING_STATE_DICT_NAME, @@ -979,6 +1015,7 @@ def publish_routing_state( uk_version: str, force_latest: bool = False, ) -> RoutingState: + """Add a deployed app to the active routing state and store it in one write.""" _require_serialized_routing_workflow() store = modal.Dict.from_name( ROUTING_STATE_DICT_NAME, @@ -1031,6 +1068,7 @@ def restore_routing_state( def main() -> None: + """Publish a deployed app's routes, or seed the state from legacy Dicts.""" parser = argparse.ArgumentParser( formatter_class=argparse.RawDescriptionHelpFormatter, description=( diff --git a/projects/policyengine-simulation-executor/src/modal/utils/validate_hf_access.py b/projects/policyengine-simulation-executor/src/modal/utils/validate_hf_access.py index 42ce91ab0..895fc9549 100644 --- a/projects/policyengine-simulation-executor/src/modal/utils/validate_hf_access.py +++ b/projects/policyengine-simulation-executor/src/modal/utils/validate_hf_access.py @@ -10,11 +10,11 @@ def validate_local_access() -> HFCredentialAudit: """Validate the credential supplied directly to the CI process.""" - return validate_configured_uk_private_hf_access() def main() -> None: + """Validate the CI credential and print its audit record as JSON.""" print(validate_local_access().model_dump_json()) diff --git a/projects/policyengine-simulation-executor/src/modal/utils/validate_v2_worker_release.py b/projects/policyengine-simulation-executor/src/modal/utils/validate_v2_worker_release.py index fb96e8c07..b430a7f97 100644 --- a/projects/policyengine-simulation-executor/src/modal/utils/validate_v2_worker_release.py +++ b/projects/policyengine-simulation-executor/src/modal/utils/validate_v2_worker_release.py @@ -41,6 +41,11 @@ def _spawn_validation( def validate_release(*, environment: str) -> V2WorkerValidation: + """Invoke each deployed country validation function and return the release record. + + Raises RuntimeError when a function does not report success, returns no + invocation id, or reports bundle values that differ from the local bundle. + """ resolved = load_stage12_bundle() application_name = v2_application_name(resolved.bundle.policyengine_version) invocation_ids: dict[CountryId, ManifestText] = {} @@ -84,6 +89,7 @@ def validate_release(*, environment: str) -> V2WorkerValidation: def main() -> None: + """Write the release validation record to ``--output`` and print its id.""" parser = argparse.ArgumentParser( description="Validate the deployed Stage 12 v2 worker release", ) diff --git a/projects/policyengine-simulation-executor/src/modal/v2_app.py b/projects/policyengine-simulation-executor/src/modal/v2_app.py index dab372b87..9a15e1583 100644 --- a/projects/policyengine-simulation-executor/src/modal/v2_app.py +++ b/projects/policyengine-simulation-executor/src/modal/v2_app.py @@ -97,6 +97,11 @@ def _country_bundle(country: CountryId) -> Stage12CountryBundle: def build_v2_image(countries: tuple[CountryId, ...]) -> modal.Image: + """Build a Stage 12 image with the bundle's certified data for ``countries``. + + The bundle manifest digest and per-country bundle values are baked into + the image environment, which worker validation checks at run time. + """ country_values = { country: _country_bundle(country).model_dump(mode="json") for country in countries @@ -167,6 +172,7 @@ def _validate(country: CountryId) -> dict: secrets=worker_secrets, ) def validate_worker_us() -> dict: + """Validate the US worker image without serving and return its bundle values.""" return _validate("us") @@ -179,6 +185,7 @@ def validate_worker_us() -> dict: secrets=worker_secrets, ) def validate_worker_uk() -> dict: + """Validate the UK worker image without serving and return its bundle values.""" return _validate("uk") @@ -197,6 +204,7 @@ def run_single_simulation_us( *, observability_context: dict | None = None, ) -> dict: + """Calculate one planned US simulation and return its artifact descriptor.""" from policyengine_simulation_executor.stage12_runtime import ( run_single_simulation, ) @@ -242,6 +250,7 @@ def run_single_simulation_uk( *, observability_context: dict | None = None, ) -> dict: + """Calculate one planned UK simulation and return its artifact descriptor.""" from policyengine_simulation_executor.stage12_runtime import ( run_single_simulation, ) @@ -290,6 +299,12 @@ def coordinate_report( *, observability_context: dict | None = None, ) -> dict: + """Run one Stage 12 report: fan out both simulations, aggregate, and persist. + + A ``production_function_call_id`` in ``context`` marks a shadow run whose + result is compared with that production job's. A report another + invocation already claimed returns a ``deduplicated`` marker instead. + """ from policyengine_simulation_executor.stage12_runtime import ( coordinate_report as run_report_coordinator, ) diff --git a/projects/policyengine-simulation-executor/src/modal/v2_static_runtime_files_smoke.py b/projects/policyengine-simulation-executor/src/modal/v2_static_runtime_files_smoke.py index 4e5f5ac9a..a86d94a38 100644 --- a/projects/policyengine-simulation-executor/src/modal/v2_static_runtime_files_smoke.py +++ b/projects/policyengine-simulation-executor/src/modal/v2_static_runtime_files_smoke.py @@ -19,6 +19,10 @@ def validate_installed_uk_local_authority_dataset() -> str: + """Return the local authority boundary version of the installed UK dataset. + + Raises if the image has no installed certified UK dataset. + """ dataset_path = resolve_local_bundle_dataset_path("uk", None) if dataset_path is None: raise RuntimeError("installed certified UK dataset is unavailable") @@ -29,7 +33,7 @@ def validate_installed_uk_local_authority_dataset() -> str: @app.function(image=uk_worker_image, timeout=600, memory=512) def smoke_stage12_static_runtime_files() -> dict[str, int | str]: - + """Count the UK local authority files inside the UK worker image.""" resources = load_uk_local_authority_resources() return { "names": len(resources.names), @@ -47,6 +51,7 @@ def smoke_stage12_static_runtime_files() -> dict[str, int | str]: @app.local_entrypoint() def main() -> None: + """Run the smoke remotely and raise unless counts and boundary version match.""" observed = smoke_stage12_static_runtime_files.remote() expected_counts = { "names": 378, diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/__init__.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/__init__.py index 433343dff..b30653d84 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/__init__.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/__init__.py @@ -1,3 +1,5 @@ +"""PolicyEngine simulation executor: the economy-comparison API and its runtime.""" + from fastapi import FastAPI from policyengine_observability import ObservabilityRuntime @@ -10,8 +12,5 @@ def initialize(app: FastAPI, runtime: ObservabilityRuntime) -> None: - """ - attach all routes to the app and configure them to use the provided SQLModel engine - and jwt settings. - """ + """Mount the simulation routes on ``app``, reporting telemetry via ``runtime``.""" app.include_router(create_router(runtime)) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/artifact_keys.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/artifact_keys.py index 9eafa91ab..645f19026 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/artifact_keys.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/artifact_keys.py @@ -136,7 +136,7 @@ def baseline_key( def baseline_simulation_id(baseline_digest: str) -> str: - """The deterministic ``Simulation.id`` for a baseline artifact. + """Return the deterministic ``Simulation.id`` for a baseline artifact. The id doubles as the artifact filename stem (``{id}.h5`` beside the dataset — the path ``Simulation.ensure()`` already loads from), so it @@ -160,16 +160,19 @@ def baseline_simulation_id(baseline_digest: str) -> str: def dataset_artifact_filename(stem: str, year: int) -> str: + """Return the single-year dataset basename the runtime's existence check expects.""" return f"{stem}_year_{year}.h5" def baseline_artifact_filename(simulation_id: str) -> str: + """Return the ``{id}.h5`` basename ``Simulation.ensure()`` loads a baseline from.""" return f"{simulation_id}.h5" def dataset_artifact_path( country: str, dataset_digest: str, stem: str, year: int ) -> str: + """Return the store key ``datasets///_year_.h5``.""" return ( f"datasets/{country.lower()}/{dataset_digest}/" f"{dataset_artifact_filename(stem, year)}" @@ -179,6 +182,7 @@ def dataset_artifact_path( def baseline_artifact_path( country: str, baseline_digest: str, simulation_id: str ) -> str: + """Return the store key ``baselines///.h5``.""" return ( f"baselines/{country.lower()}/{baseline_digest}/" f"{baseline_artifact_filename(simulation_id)}" @@ -186,10 +190,12 @@ def baseline_artifact_path( def manifest_path(manifest_digest: str) -> str: + """Return the store key of a deploy manifest, addressed by its own digest.""" return f"manifests/{manifest_digest}.json" def deployed_marker_path(environment: str) -> str: + """Return the key of the marker naming the manifest deployed to ``environment``.""" return f"deployed/{environment}.json" @@ -214,6 +220,7 @@ class DatasetArtifactIdentity: @property def digest(self) -> str: + """Content key over every input that can change this dataset's bytes.""" return dataset_key( country=self.country, spm=self.spm, @@ -229,10 +236,12 @@ def digest(self) -> str: @property def filename(self) -> str: + """Basename the runtime resolves under ``POLICYENGINE_DATA_FOLDER``.""" return dataset_artifact_filename(self.stem, self.year) @property def store_path(self) -> str: + """Store key of this dataset artifact.""" return dataset_artifact_path(self.country, self.digest, self.stem, self.year) @@ -248,6 +257,7 @@ class BaselineArtifactIdentity: @property def digest(self) -> str: + """Content key; nests the dataset digest, so a dataset change rotates it.""" return baseline_key( country=self.country, region=self.region, @@ -260,16 +270,23 @@ def digest(self) -> str: @property def simulation_id(self) -> str: + """``Simulation.id`` the runtime assigns to this baseline.""" return baseline_simulation_id(self.digest) @property def storage_id(self) -> str: + """Artifact filename stem: the simulation id, plus the SPM digest when set. + + The suffix keeps SPM and non-SPM baselines that share a simulation id + from overwriting each other. + """ if self.spm is None: return self.simulation_id return f"{self.simulation_id}-spm-{canonical_digest(self.spm)}" @property def store_path(self) -> str: + """Store key of this baseline artifact, named by ``storage_id``.""" return baseline_artifact_path(self.country, self.digest, self.storage_id) @@ -350,6 +367,11 @@ def collect_baseline_identity( scope_key: str | None, spm: dict | None = None, ) -> BaselineArtifactIdentity: + """Return the identity of this bundle's baseline for one region and year. + + ``spm`` is normalized against the installed image's capability first, so + a request that omits it still keys to the image's default SPM selection. + """ from policyengine_simulation_executor.spm import normalize_runtime_spm selection = normalize_runtime_spm( diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/artifact_store.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/artifact_store.py index 9298fd237..12ef86cd0 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/artifact_store.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/artifact_store.py @@ -37,6 +37,10 @@ def resolve_bucket_name(explicit: str | None = None) -> str: + """Return ``explicit`` or the ``POLICYENGINE_ARTIFACT_BUCKET`` value. + + Raises ``RuntimeError`` when neither names a bucket. + """ bucket = explicit or os.environ.get(ARTIFACT_BUCKET_ENV) if not bucket: raise RuntimeError( @@ -59,14 +63,18 @@ def _default_client() -> storage.Client: class ArtifactStore: + """Artifact, manifest and deployed-marker storage in one GCS bucket.""" + def __init__( self, bucket_name: str | None = None, *, client: storage.Client | None = None ) -> None: + """Resolve the bucket now; the real GCS client is built lazily on first use.""" self.bucket_name = resolve_bucket_name(bucket_name) self._client = client @property def client(self) -> storage.Client: + """GCS client: the injected one, or a real client built on first access.""" if self._client is None: self._client = _default_client() return self._client @@ -75,12 +83,16 @@ def _blob(self, path: str) -> storage.Blob: return self.client.bucket(self.bucket_name).blob(path) def exists(self, path: str) -> bool: + """Report whether an object is stored at ``path``.""" return bool(self._blob(path).exists()) def upload_file(self, path: str, local_path: str | Path) -> bool: - """Write-once upload. Returns True if this call uploaded the object, - False if an identical-keyed object already existed (see module - docstring for why that is success).""" + """Upload a local file write-once; return whether this call created the object. + + False means an object already existed at that key, which counts as + success because keys digest the full input closure (see the module + docstring). + """ try: self._blob(path).upload_from_filename( str(local_path), if_generation_match=0 @@ -108,12 +120,17 @@ def upload_bytes( return True def read_bytes(self, path: str) -> bytes | None: + """Return the object's bytes, or None when nothing is stored at ``path``.""" try: return self._blob(path).download_as_bytes() except NotFound: return None def download_file(self, path: str, local_path: str | Path) -> None: + """Download ``path`` to ``local_path`` atomically, creating parent directories. + + A failed download never leaves a partial file at ``local_path``. + """ destination = Path(local_path) destination.parent.mkdir(parents=True, exist_ok=True) # Download to a sibling temp name and rename: callers guard on @@ -145,6 +162,7 @@ def _upload_json( return True def read_json(self, path: str) -> Mapping[str, Any] | None: + """Return the JSON object at ``path``; None if missing or not an object.""" try: payload = json.loads(self._blob(path).download_as_bytes()) except NotFound: @@ -160,6 +178,7 @@ def write_manifest(self, payload: Mapping[str, Any]) -> str: return digest def read_manifest(self, digest: str) -> Mapping[str, Any] | None: + """Return the manifest stored under ``digest``, or None if absent.""" return self.read_json(artifact_keys.manifest_path(digest)) # --- Deployed-environment markers (last-writer-wins) ------------------ @@ -167,9 +186,11 @@ def read_manifest(self, digest: str) -> Mapping[str, Any] | None: def write_deployed_marker( self, environment: str, payload: Mapping[str, Any] ) -> None: + """Overwrite ``environment``'s deployed marker with ``payload``.""" self._upload_json( artifact_keys.deployed_marker_path(environment), payload, overwrite=True ) def read_deployed_marker(self, environment: str) -> Mapping[str, Any] | None: + """Return ``environment``'s deployed marker, or None if none was written.""" return self.read_json(artifact_keys.deployed_marker_path(environment)) 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 ffc8d7091..cbc291ac1 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 @@ -111,11 +111,13 @@ def deterministic_baseline_id( scoping_strategy: ScopingStrategy | None, year: int, ) -> str | None: - """The runtime wrapper: id for a qualifying baseline, else None. + """Return the deterministic simulation id for a qualifying baseline, else None. + This is the runtime reader's wrapper around ``qualifying_baseline_identity``. Unlike the writer, the read path must never fail a request over identity collection (a manifest/receipt hiccup) — it degrades to a - random id, meaning no artifact reuse for that request. + random id, meaning no artifact reuse for that request. SPM selection + errors still propagate, since they mean the request itself is invalid. """ try: identity = qualifying_baseline_identity( @@ -154,6 +156,7 @@ def artifact_outcome(self) -> str | None: return self._artifact_outcome def run(self) -> None: + """Run the model, noting that this process computed the output itself.""" self._computed_this_process = True super().run() @@ -172,6 +175,13 @@ def _record_outcome(self, outcome: str) -> None: self._artifact_outcome = outcome def ensure(self) -> None: + """Load or compute the output, recomputing when a loaded artifact is incomplete. + + A load counts only if it has every requested output column and, under + an SPM selection, a valid receipt for the dataset year; otherwise the + model runs, the result is saved and the in-process cache entry is + replaced. Records the hit/miss/incomplete outcome on first call. + """ # The wrapper restores selection metadata on load/cache hits. Keep the # request's selection so an invalid cached entry cannot replace it. selection = getattr(self, "spm_config", None) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/generate_openapi.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/generate_openapi.py index 5d43c0925..3bf906d7e 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/generate_openapi.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/generate_openapi.py @@ -1,3 +1,5 @@ +"""Write the executor app's OpenAPI spec to ``artifacts/openapi.json`` when imported.""" + import json from pathlib import Path diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/hf_access_validation.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/hf_access_validation.py index 2565dc13e..76db81c62 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/hf_access_validation.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/hf_access_validation.py @@ -59,12 +59,16 @@ class HFAccessTokenIdentity(BaseModel): class HFAuthIdentity(BaseModel): + """The ``auth`` block of a whoami-v2 response: the token that made the call.""" + model_config = ConfigDict(extra="ignore", frozen=True, populate_by_name=True) access_token: HFAccessTokenIdentity = Field(alias="accessToken") class HFWhoAmIIdentity(BaseModel): + """The whoami-v2 fields the audit reads: account name and calling token.""" + model_config = ConfigDict(extra="ignore", frozen=True) name: str @@ -72,6 +76,8 @@ class HFWhoAmIIdentity(BaseModel): class HFArtifactAccess(BaseModel): + """Proof that the token can read one target: its resolved file metadata.""" + model_config = ConfigDict(extra="forbid", frozen=True) target: HFAccessTarget @@ -97,7 +103,6 @@ class HFCredentialAudit(BaseModel): def build_uk_bundle_access_plan() -> tuple[HFAccessTarget, ...]: """Derive private Hugging Face artifacts from the installed .py bundle.""" - bundle = get_country_release_bundle("uk") targets: dict[tuple[str, str, str, str], HFAccessTarget] = {} for dataset_name, dataset_uri in bundle.dataset_uris.items(): @@ -163,7 +168,6 @@ def validate_uk_private_hf_access( metadata_loader: MetadataLoader = _load_metadata, ) -> HFCredentialAudit: """Validate identity, read-only permissions, and every planned artifact.""" - token = environment.get(HF_RUNTIME_ENV_NAME) if not token: raise RuntimeError(f"{HF_RUNTIME_ENV_NAME} is required") @@ -201,5 +205,4 @@ def validate_uk_private_hf_access( def validate_configured_uk_private_hf_access() -> HFCredentialAudit: """Validate the configured token against every artifact in the UK bundle.""" - return validate_uk_private_hf_access(build_uk_bundle_access_plan()) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/main.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/main.py index 9e8391583..2e674a885 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/main.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/main.py @@ -1,3 +1,5 @@ +"""Local FastAPI entrypoint serving the executor routes plus ping and health probes.""" + import logging from collections.abc import AsyncIterator from contextlib import asynccontextmanager @@ -27,6 +29,7 @@ @asynccontextmanager async def lifespan(app: FastAPI) -> AsyncIterator[None]: + """Run the app inside the exit handler and flush observability on shutdown.""" async with exit.lifespan(): try: yield diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/national_partition.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/national_partition.py index f2a27f84b..c3bd64b9a 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/national_partition.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/national_partition.py @@ -38,7 +38,7 @@ def national_region_groups(country: str) -> list[list[str]] | None: - """The region-group partition for a country's segmented national run. + """Return the region-group partition for a country's segmented national run. Returns None when the country has no partition (only the US has one), in which case national requests run monolithically. diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py index 42f3c47be..77a765289 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py @@ -79,11 +79,12 @@ def manifest_digest_line(digest: str) -> str: + """Format the ``MANIFEST_DIGEST=`` line the CI deploy job parses.""" return f"{MANIFEST_DIGEST_PREFIX}{digest}" def cohort_params(year: int, group: list[str]) -> dict: - """The synthetic request a segmented-national child would receive.""" + """Return the synthetic request a segmented-national child would receive.""" return { "country": PRECOMPUTE_COUNTRY, "scope": "macro", @@ -155,7 +156,7 @@ def select_work(plan: PrecomputePlan, *, force: bool) -> WorkSelection: def build_manifest(plan: PrecomputePlan) -> ArtifactManifest: - """The deploy manifest: receipt + every artifact the image must fetch.""" + """Build the deploy manifest: the receipt and each artifact the image must fetch.""" artifacts = [ ManifestArtifact( type="dataset", @@ -448,6 +449,7 @@ def verify_determinism_impl( def publish_manifest_impl(bucket: str, plan: PrecomputePlan) -> str: + """Store the plan's manifest write-once and return its content digest.""" from policyengine_simulation_executor.artifact_store import ArtifactStore manifest = build_manifest(plan) @@ -465,7 +467,7 @@ def run_precompute( publish_manifest: RemoteFunction, echo: Callable[[str], None] = print, ) -> str: - """The local-entrypoint orchestration: plan, spawn misses, gate, publish. + """Orchestrate the local entrypoint: plan, spawn misses, gate, publish. Takes the five Modal function handles as parameters (each offering ``.remote``/``.spawn``) so the flow is testable with plain fakes. diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute_models.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute_models.py index 47dc14793..1807b7244 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute_models.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute_models.py @@ -109,7 +109,7 @@ class ArtifactManifest(_StrictModel): artifacts: list[ManifestArtifact] def canonical_payload(self) -> dict[str, Any]: - """The exact dict shape that gets digested and stored.""" + """Return the exact dict that gets digested and stored, keyed by wire alias.""" return self.model_dump(by_alias=True) @@ -137,6 +137,8 @@ class DeployedMarker(_StrictModel): class DatasetBuildResult(_StrictModel): + """Result of one dataset build; ``uploaded`` is False if the key existed.""" + year: int path: str uploaded: bool @@ -145,6 +147,12 @@ class DatasetBuildResult(_StrictModel): class BaselineComputeResult(_StrictModel): + """Outcome of one cohort baseline computation and its write-once upload. + + ``outcome`` is the baseline's artifact outcome (hit/incomplete/miss) from + its ``ensure()`` call; seconds are wall-clock, rounded to 0.1. + """ + year: int group: list[str] simulation_id: str @@ -155,6 +163,11 @@ class BaselineComputeResult(_StrictModel): class DeterminismVerdict(_StrictModel): + """Whether a stored baseline matches a fresh recompute, column by column. + + ``differences`` lists at most the first 50 mismatches. + """ + equal: bool differences: list[str] @@ -162,6 +175,10 @@ class DeterminismVerdict(_StrictModel): class RemoteFunction(Protocol): """The slice of a Modal function handle the orchestration uses.""" - def remote(self, *args: Any) -> Any: ... # noqa: ANN401 # mirrors modal.Function.remote; payloads and results vary per function + def remote(self, *args: Any) -> Any: # noqa: ANN401 # mirrors modal.Function.remote; payloads and results vary per function + """Run the function in its container and block until it returns.""" + ... - def spawn(self, *args: Any) -> Any: ... # noqa: ANN401 # mirrors modal.Function.spawn; payloads vary and the returned FunctionCall is untyped here + def spawn(self, *args: Any) -> Any: # noqa: ANN401 # mirrors modal.Function.spawn; payloads vary and the returned FunctionCall is untyped here + """Start the call without waiting; the handle's ``get()`` returns the result.""" + ... diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/release_bundle.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/release_bundle.py index 522b6df3e..e13aab3aa 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/release_bundle.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/release_bundle.py @@ -36,6 +36,13 @@ @dataclass(frozen=True) class CountryReleaseBundle: + """Package versions and dataset locations one country's installed bundle pins. + + ``dataset_uris`` maps dataset names to versioned URIs and + ``dataset_repo_types`` maps them to their Hugging Face repository type + (``"model"`` or ``"dataset"``). + """ + country: str policyengine_version: str model_package_name: str @@ -255,7 +262,6 @@ def _bundle_country_metadata( def get_bundled_package_version(package: str) -> str: """Return a runtime package version from policyengine.py's bundle manifest.""" - bundle = _current_policyengine_bundle() if not isinstance(bundle, Mapping): raise TypeError("PolicyEngine.py bundle metadata is unavailable") @@ -284,7 +290,6 @@ def _derive_data_package_version( manifest_data_package_version: str, ) -> str: """Return the data package version certified by policyengine.py.""" - version = bundled_data_package.get("version") or manifest_data_package_version if not isinstance(version, str) or not version.strip(): raise ValueError("PolicyEngine.py bundle has no data package version") @@ -294,7 +299,6 @@ def _derive_data_package_version( @lru_cache def get_country_release_bundle(country: str) -> CountryReleaseBundle: """Return package and dataset versions from policyengine.py metadata.""" - country = _normalise_country(country) from policyengine.provenance.manifest import get_release_manifest @@ -387,6 +391,7 @@ def get_country_release_bundle(country: str) -> CountryReleaseBundle: def get_bundled_country_model_version(country: str) -> str: + """Return the country model package version the installed bundle pins.""" return get_country_release_bundle(country).model_version @@ -400,6 +405,13 @@ def _split_requested_revision(requested_data: str) -> tuple[str, str | None]: def resolve_bundle_dataset_name(country: str, requested_data: str | None) -> str: + """Resolve a requested dataset to a bundle dataset name or a full URI. + + None selects the bundle's default dataset and a full URI passes through. + ``name@revision`` becomes the bundle URI for ``name`` re-pinned to that + Hugging Face revision without checking the revision exists; it raises + ``ValueError`` when the bundle has no dataset called ``name``. + """ bundle = get_country_release_bundle(country) if requested_data is None: return bundle.default_dataset @@ -420,6 +432,7 @@ def resolve_bundle_dataset_name(country: str, requested_data: str | None) -> str def resolve_bundle_dataset_uri(country: str, requested_data: str | None) -> str: + """Return the URI for a requested dataset; names the bundle lacks pass through.""" bundle = get_country_release_bundle(country) dataset_name = resolve_bundle_dataset_name(country, requested_data) if "://" in dataset_name: @@ -436,6 +449,11 @@ def _receipt_path() -> Path: def get_installed_bundle_receipt() -> Mapping | None: + """Return the image's bundle receipt, or None when it is absent or unreadable. + + The receipt is read from ``POLICYENGINE_BUNDLE_RECEIPT`` if set, else from + the data folder, on every call. + """ path = _receipt_path() if not path.exists(): return None @@ -485,7 +503,6 @@ def resolve_local_bundle_dataset_path( requested_data_version: str | None = None, ) -> str | None: """Return a local certified dataset path when the request uses the default.""" - country = _normalise_country(country) if not _is_default_bundle_dataset(country, requested_data, requested_data_version): return None @@ -511,7 +528,6 @@ def resolve_runtime_bundle_dataset_uri( prefer_local: bool = True, ) -> str: """Resolve a request dataset to the reference the worker should load.""" - if prefer_local: local_path = resolve_local_bundle_dataset_path( country, diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/segmented_national_reduce.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/segmented_national_reduce.py index 3fc6de14f..d3a87244f 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/segmented_national_reduce.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/segmented_national_reduce.py @@ -129,7 +129,7 @@ def build_national_output( resolved_data_version: str | None = None, runtime: ObservabilityRuntime | None = None, ) -> dict[str, Any]: - """The reduce: children's microdata -> full national macro output dict. + """Reduce the children's microdata into the full national macro output dict. Runs the existing ``SimulationOutputBuilder`` over stand-in simulations with ``resolved_region_code=country`` so the national congressional diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/settings.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/settings.py index 24bfb5547..e69c34794 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/settings.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/settings.py @@ -1,3 +1,5 @@ +"""Environment-driven settings for the executor app (read from ``.env``).""" + from enum import Enum from functools import lru_cache @@ -6,16 +8,21 @@ class Environment(Enum): + """Deployment environment the app runs in.""" + DESKTOP = "desktop" PRODUCTION = "production" class AppSettings(BaseSettings): + """Executor settings loaded from environment variables and a local ``.env`` file.""" + environment: Environment = Environment.DESKTOP @field_validator("environment", mode="before") @classmethod def strip_environment(cls, v: object) -> object: + """Strip whitespace from a string ``environment`` before enum parsing.""" if isinstance(v, str): return v.strip() return v @@ -34,4 +41,5 @@ def strip_environment(cls, v: object) -> object: @lru_cache def get_settings() -> AppSettings: + """Return the process-wide settings, loaded once and cached.""" return AppSettings() diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation.py index 4016de5ec..57203d8fb 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation.py @@ -1,3 +1,5 @@ +"""FastAPI router for the economy-comparison endpoint.""" + import logging from fastapi import APIRouter @@ -15,6 +17,10 @@ def create_router(runtime: ObservabilityRuntime) -> APIRouter: + """Build the router whose comparison endpoint runs simulations through ``runtime``. + + SPM selection errors become a 400 failure payload; other ``ValueError``s propagate. + """ router = APIRouter() @router.post("/simulate/economy/comparison", response_model=EconomyComparison) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_microdata.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_microdata.py index c51a28fc2..063eeed1b 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_microdata.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_microdata.py @@ -29,11 +29,11 @@ def extract_output_microdata( baseline: Simulation, reform: Simulation ) -> dict[str, Any]: - """Columnar dump of each simulation's computed ``output_dataset.data`` per - entity, for baseline and reform, plus per-column source dtypes. + """Dump each simulation's computed output columns per entity, with source dtypes. - The entity tables come from the model's own ``YearData.entity_data`` - property, so there is no hardcoded per-country entity list to drift. + Covers ``output_dataset.data`` for both baseline and reform. The entity + tables come from the model's own ``YearData.entity_data`` property, so + there is no hardcoded per-country entity list to drift. ``output_dataset.data`` is already populated (the output builder ran), so this only reads arrays — no model recompute. All columns (including the ``*_id`` link columns) are emitted so the coordinator's ``map_to_entity`` @@ -63,11 +63,12 @@ def dump( def rebuild_entity_frame( columns: dict[str, list], dtypes: dict[str, str] | None ) -> pd.DataFrame: - """Rebuild one transported entity table, casting each column back to its - source dtype (float32 stays float32 — see module docstring). A missing - or uncastable dtype falls back to pandas' inferred dtype rather than - failing the reduce, but LOUDLY: that fallback breaks the bit-exactness - guarantee (float64 widening, ~1e-7 aggregate drift). + """Rebuild one transported entity table with each column cast to its source dtype. + + Float32 stays float32 (see module docstring). A missing or uncastable + dtype falls back to pandas' inferred dtype rather than failing the + reduce, but LOUDLY: that fallback breaks the bit-exactness guarantee + (float64 widening, ~1e-7 aggregate drift). """ df = pd.DataFrame(columns) if dtypes is None: diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_budget.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_budget.py index 7d9a1e95e..1a3b17c64 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_budget.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_budget.py @@ -22,6 +22,11 @@ def build_detailed_budget(analysis: object) -> DetailedBudgetOutput: + """Return per-program baseline, reform and difference totals from the analysis. + + Rows without a program name are skipped, and a missing or non-finite + ``change`` falls back to reform minus baseline. + """ collection = getattr(analysis, "program_statistics", None) if isinstance(collection, DetailedBudgetOutput): return collection @@ -43,6 +48,12 @@ def build_detailed_budget(analysis: object) -> DetailedBudgetOutput: def build_budgetary_impact( country: str, baseline: Simulation, reform: Simulation ) -> BudgetaryImpact: + """Compute headline budget figures from household tax and benefit changes. + + ``budgetary_impact`` is the tax revenue change minus the benefit spending + change. State income tax is only modelled for the US; elsewhere its impact + is 0.0. + """ tax_revenue_impact = _change_output_variable( baseline, reform, "household_tax", entity="household" ) 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 3207f0fa9..44b714edc 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 @@ -45,6 +45,12 @@ @dataclass class SimulationOutputBuilder: + """Assemble one year's macro output from a computed baseline and reform. + + Each output segment is built inside its own observability span when a + ``runtime`` is given. + """ + country: str simulation_params: dict[str, Any] country_module: Any @@ -58,6 +64,7 @@ class SimulationOutputBuilder: _analysis: Any = field(default=None, init=False) def __post_init__(self) -> None: + """Normalize ``country`` to lower case.""" self.country = self.country.lower() def _span(self, name: str) -> AbstractContextManager[object]: @@ -65,6 +72,11 @@ def _span(self, name: str) -> AbstractContextManager[object]: @property def analysis(self) -> object: + """Country economic impact analysis, computed on first access and reused. + + Cliff impacts are included only when the request sets + ``include_cliffs`` to True. + """ if self._analysis is None: with self._span(ANNUAL_IMPACT_STAGES.name(Stage.ECONOMIC_IMPACT_ANALYSIS)): self._analysis = self.country_module.economic_impact_analysis( @@ -78,6 +90,7 @@ def _include_cliff_impacts(self) -> bool: return self.simulation_params.get("include_cliffs") is True def build(self) -> SingleYearMacroOutput: + """Build every output segment into the typed single-year macro output.""" with self._span(ANNUAL_IMPACT_STAGES.name(Stage.SIMULATION_OUTPUT_BUILD)): poverty_outputs = self._build_poverty_outputs() wealth_decile = getattr(self.analysis, "wealth_decile_impacts", None) @@ -110,6 +123,7 @@ def build(self) -> SingleYearMacroOutput: ) def serialize(self) -> dict[str, Any]: + """Build the output and dump it to JSON-compatible data.""" with self._span(ANNUAL_IMPACT_STAGES.name(Stage.CALCULATION)): output = self.build() with self._span(ANNUAL_IMPACT_STAGES.name(Stage.RESPONSE_SERIALIZATION)): diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_cliff.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_cliff.py index 746d1b8f2..e09a2db74 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_cliff.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_cliff.py @@ -12,6 +12,7 @@ def build_cliff_impact(analysis: object) -> CliffImpactOutput | None: + """Return the baseline and reform cliff impact, or None if it was not computed.""" cliff_impact = getattr(analysis, "cliff_impact", None) if isinstance(cliff_impact, CliffImpactOutput): return cliff_impact diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_distribution.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_distribution.py index 32032e308..50bd794eb 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_distribution.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_distribution.py @@ -27,10 +27,15 @@ def build_decile(analysis: object) -> DecileOutput: + """Return the income-decile changes from the analysis's ``decile_impacts``.""" return build_decile_output(getattr(analysis, "decile_impacts", None)) def build_decile_output(collection: object) -> DecileOutput: + """Key absolute and relative changes by decile number (``"1"``..``"10"``). + + Rows with a decile of 0 or below are dropped; missing values become 0.0. + """ if isinstance(collection, DecileOutput): return collection average: dict[str, float] = {} @@ -51,6 +56,11 @@ def build_decile_output(collection: object) -> DecileOutput: def build_intra_decile_output( baseline: Simulation, reform: Simulation ) -> IntraDecileOutput: + """Compute household net income winners and losers within each income decile. + + A calculation failure is logged and yields empty deciles with zero overall + shares rather than failing the request; SPM selection errors propagate. + """ from policyengine.outputs.intra_decile_impact import ( compute_intra_decile_impacts, ) @@ -67,6 +77,7 @@ def build_intra_decile_output( def build_wealth_decile(country: str, wealth_decile: object) -> DecileOutput | None: + """Return wealth-decile changes for the UK; None for other countries.""" if country != "uk": return None return build_decile_output(wealth_decile) @@ -75,6 +86,7 @@ def build_wealth_decile(country: str, wealth_decile: object) -> DecileOutput | N def build_intra_wealth_decile( country: str, intra_wealth_decile: object ) -> IntraDecileOutput | None: + """Return within-wealth-decile winners and losers for the UK; None elsewhere.""" if country != "uk": return None return build_intra_decile_output_from_collection(intra_wealth_decile) @@ -83,6 +95,10 @@ def build_intra_wealth_decile( def build_intra_decile_output_from_collection( collection: object, ) -> IntraDecileOutput: + """Reshape intra-decile rows into per-category decile lists plus overall shares. + + The ``all`` share for each category is the unweighted mean across deciles. + """ if isinstance(collection, IntraDecileOutput): return collection deciles: dict[str, list[float]] = {label: [] for label in INTRA_DECILE_COLUMNS} diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_geographic.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_geographic.py index 2bf92c1b9..281e00107 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_geographic.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_geographic.py @@ -31,7 +31,6 @@ def _required_uk_geography_lookup_csv_path(spec: UKGeographyAssetSpec) -> str: """Resolve one required lookup under the runtime GCP identity.""" - from policyengine.outputs.uk_geography_impact import ( resolve_uk_geography_lookup_csv_path, ) @@ -60,7 +59,6 @@ def _complete_uk_geography_output( name_field: str, ) -> GeographicImpactOutput: """Require every UK geography record to contain lookup-owned metadata.""" - output = build_geographic_impact_output(value) if output is None or not output.root: raise ValueError("UK geography output did not contain result records") @@ -91,6 +89,10 @@ def _policyengine_us_district_metadata() -> tuple[dict[int, str], frozenset[str] def build_geographic_impact_output(value: object) -> GeographicImpactOutput | None: + """Validate a list of per-area records (or a model dumping to one) into the output. + + Non-mapping items are dropped; None means ``value`` held no record list. + """ if isinstance(value, GeographicImpactOutput): return value records = _output_model_dump(value) @@ -108,6 +110,11 @@ def build_geographic_impact_output(value: object) -> GeographicImpactOutput | No def build_congressional_district_impact_output( value: object, ) -> CongressionalDistrictImpactOutput | None: + """Convert per-district records into public ``"ST-NN"`` district rows. + + Returns None when ``value`` holds no record list; raises ``ValueError`` when + a record's state FIPS or district number cannot be determined. + """ if isinstance(value, CongressionalDistrictImpactOutput): return value records = _output_model_dump(value) @@ -194,6 +201,11 @@ def build_congressional_district_impact( *, region_code: str | None = None, ) -> CongressionalDistrictImpactOutput | None: + """Compute US congressional district impacts for national, state or group runs. + + None for other countries or regions, and when the calculation fails (the + failure is logged; SPM selection errors still propagate). + """ if country != "us": return None if not _should_build_us_congressional_district_impact(region_code): @@ -225,7 +237,6 @@ def _require_uk_area_column( none of its own, and one built from a routed local-area dataset is a separate run. """ - if column not in household.columns: raise ValueError( f"UK {breakdown} breakdowns need {column} on the simulated dataset, " @@ -237,6 +248,11 @@ def _require_uk_area_column( def build_uk_constituency_impact( country: str, baseline: Simulation, reform: Simulation ) -> GeographicImpactOutput | None: + """Compute UK parliamentary constituency impacts; None for other countries. + + Downloads the constituency lookup if needed and raises when the dataset + lacks constituency codes or a record lacks lookup name or coordinates. + """ if country != "uk": return None _require_uk_area_column( @@ -266,6 +282,11 @@ def build_uk_local_authority_impact( *, uk_local_authority_metadata: UKLocalAuthorityMetadata | None = None, ) -> GeographicImpactOutput | None: + """Compute UK local authority impacts, named and placed for the detected boundaries. + + None for other countries. For the UK, ``uk_local_authority_metadata`` is + required, and a code outside its boundary version raises ``ValueError``. + """ if country != "uk": return None diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_inequality.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_inequality.py index 7bd1366a5..d5f7dff55 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_inequality.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_inequality.py @@ -10,6 +10,10 @@ def build_inequality(analysis: object) -> InequalityOutput: + """Pair baseline and reform Gini and top 10% and top 1% income shares. + + A metric the analysis lacks (or reports as NaN) is emitted as 0.0. + """ baseline = getattr(analysis, "baseline_inequality", None) reform = getattr(analysis, "reform_inequality", None) if isinstance(baseline, InequalityOutput): diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_labor.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_labor.py index 7033fe5ab..ab140b6ad 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_labor.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_labor.py @@ -9,6 +9,7 @@ def build_labor_supply_response(analysis: object) -> LaborSupplyResponseOutput | None: + """Return the analysis's labor supply response, or None if it computed none.""" labor_supply_response = getattr(analysis, "labor_supply_response", None) if isinstance(labor_supply_response, LaborSupplyResponseOutput): return labor_supply_response diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_poverty.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_poverty.py index fb260ce52..532699306 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_poverty.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_poverty.py @@ -115,6 +115,12 @@ def _race_poverty_output(values: dict[str, dict[str, float]]) -> RacePovertyOutp def build_poverty_outputs( country: str, baseline: Simulation, reform: Simulation, analysis: object ) -> PovertyModuleOutputs: + """Compute poverty rates with age, gender and (US only) race breakdowns. + + "Poverty" and "deep poverty" mean SPM and deep SPM poverty in the US, and + relative and absolute poverty before housing costs in the UK. A breakdown + whose calculation fails is logged and reported as zero rates. + """ prefix = "us" if country == "us" else "uk" baseline_poverty_by_age = _try_compute_output( "baseline poverty by age", @@ -181,6 +187,10 @@ def build_poverty_output( baseline_by_age: object, reform_by_age: object, ) -> PovertyOutput: + """Merge overall and by-age rates into child, adult, senior and all blocks. + + Rates the inputs do not report are 0.0. + """ if isinstance(baseline, PovertyOutput): return baseline result = { @@ -213,6 +223,7 @@ def build_poverty_by_gender_output( baseline_by_gender: object, reform_by_gender: object, ) -> PovertyByGenderOutput: + """Split poverty and deep poverty rates by gender; unreported rates are 0.0.""" if isinstance(baseline_by_gender, PovertyByGenderOutput): return baseline_by_gender result = { @@ -237,6 +248,10 @@ def build_poverty_by_race_output( baseline_by_race: object, reform_by_race: object, ) -> PovertyByRaceOutput: + """Split US SPM poverty rates by race; there is no deep poverty breakdown. + + Rates the inputs do not report are 0.0. + """ if isinstance(baseline_by_race, PovertyByRaceOutput): return baseline_by_race result = { 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 871656ae5..c04994263 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 @@ -1,5 +1,4 @@ -""" -Simulation implementation - pure logic with snapshotted imports. +"""Simulation implementation - pure logic with snapshotted imports. This module avoids importing policyengine at module level so the worker can load the requested country module without triggering cross-country imports. @@ -54,6 +53,12 @@ @dataclass(frozen=True) class RegionResolution: + """A request's normalized region code, the dataset to load, and how to scope it. + + ``scoping_strategy`` narrows the simulation to the region; None simulates + the whole dataset, as a national run does. + """ + code: str dataset_reference: str | None = None scoping_strategy: Any | None = None @@ -75,8 +80,8 @@ def _normalize_credentials_blob(creds_json: str) -> str: 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.""" - + clean blob. Only unwrap when the payload looks wrapped. + """ try: json.loads(creds_json) except json.JSONDecodeError: @@ -91,8 +96,7 @@ def _normalize_credentials_blob(creds_json: str) -> str: def setup_gcp_credentials( runtime: ObservabilityRuntime | None = None, ) -> Iterator[None]: - """ - Set up GCP credentials from environment variable. + """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 @@ -162,8 +166,7 @@ def run_simulation_impl( *, runtime: ObservabilityRuntime, ) -> dict: - """ - Execute economic simulation. + """Execute economic simulation. Pure implementation with no Modal dependencies. Accepts the gateway simulation payload and returns the legacy macro result dict. @@ -275,7 +278,7 @@ def _compile_interval_policy( def _runtime_policy( policy: dict[str, Any] | None, *, country_module: ModuleType, year: int ) -> dict[str, Any] | Policy | None: - """A request's policy as the wrapper must receive it. + """Convert a request's policy into the form the wrapper must receive. Normalises the period keys, then compiles a policy with interval keys into wrapper parameter values (see ``_compile_interval_policy``). Every @@ -376,7 +379,6 @@ def _require_certified_uk_weight_matrix( release replaces, is never downloaded (see ``_build_uk_weight_replacement_region``). """ - from policyengine.core.scoping_strategy import WeightReplacementStrategy if not isinstance(scoping_strategy, WeightReplacementStrategy): @@ -564,10 +566,9 @@ def _resolve_region_group( country: str, params: dict[str, Any], ) -> RegionResolution: - """Resolve a ``region_group`` (list of member region codes) to one - ``RegionGroupStrategy`` that scopes a single simulation to the union of the - members. + """Resolve a ``region_group`` list of region codes to one ``RegionGroupStrategy``. + The strategy scopes a single simulation to the union of the members. Members must be row-filter regions (states/CDs) so national weights are preserved (INV-5); the group filters the national dataset. """ @@ -613,7 +614,7 @@ def _country_module(country: str) -> ModuleType: def resolve_data_folder() -> str: - """The folder datasets and baseline artifacts are read from and saved to. + """Return the folder datasets and baseline artifacts are read from and saved to. One home for the default: the precompute downloads store artifacts into this folder for ``ensure_datasets``/``Simulation.ensure()`` to find, so diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/spm.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/spm.py index 8cb5696f8..240fb9674 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/spm.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/spm.py @@ -33,7 +33,10 @@ def _forecast(expected_sha256: str | None) -> SPMForecast: @lru_cache(maxsize=1) def runtime_spm_capability() -> SPMCapability | None: - """The installed image's certified SPM capability, or None. + """Return the installed image's certified SPM capability. + + None means the bundle configures no SPM measurement; a configured but + unusable installation raises ``SPMInputError`` instead. Memoized for the life of the process because every input is a property of the installed image -- the bundle configuration, the wrapper's model @@ -223,7 +226,6 @@ def compact_spm_provenance( reads only the scalar identity fields, executed-year keys, and explicit package versions that belong on the simulation API wire contract. """ - if not isinstance(country_receipt, Mapping): raise TypeError("SPM calculation receipt must be an object") resolved = SPMSelection.model_validate(selection) @@ -260,7 +262,6 @@ def simulation_spm_calculation_provenance( selection: object, ) -> SPMProvenance: """Read and compact one simulation's resolved SPM receipt.""" - resolved = SPMSelection.model_validate(selection) if getattr(simulation, "spm_config", None) != resolved.model_dump(mode="json"): raise SPMInputError( @@ -292,6 +293,12 @@ def simulation_spm_result( *, expected_year: int | None = None, ) -> dict[str, object]: + """Return the ``spm_provenance`` result fields for a baseline/reform pair. + + Empty when no SPM selection applies. Raises ``SPMInputError`` if either + simulation ignored the selection, has no valid receipt, or (when + ``expected_year`` is given) the receipts do not cover that year. + """ if selection is None: return {} baseline_provenance = simulation_spm_calculation_provenance(baseline, selection) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_artifacts.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_artifacts.py index 6d88745ce..272c0ff89 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_artifacts.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_artifacts.py @@ -37,6 +37,10 @@ def canonical_json_bytes(value: object) -> bytes: + """Encode ``value`` as compact, key-sorted UTF-8 JSON so equal values hash equally. + + NaN and infinity raise ``ValueError`` rather than producing non-standard JSON. + """ return json.dumps( value, sort_keys=True, @@ -52,6 +56,11 @@ def comparison_run_prefix( created_at: datetime, evaluation_id: UUID, ) -> str: + """Return ``stage-12-runs////`` for one run. + + Raises ``ValueError`` for an environment that could escape the prefix or a + naive ``created_at``. + """ if not environment or "/" in environment or ".." in environment: raise ValueError("invalid Stage 12 artifact environment") if created_at.tzinfo is None: @@ -62,18 +71,22 @@ def comparison_run_prefix( def input_path(*, prefix: str, role: str) -> str: + """Return the key of the planned execution input JSON for one simulation role.""" return f"{prefix}/inputs/{role}.json" def simulation_path(*, prefix: str, role: str) -> str: + """Return the key of the Parquet output for one simulation role.""" return f"{prefix}/simulations/{role}.parquet" def aggregate_path(*, prefix: str) -> str: + """Return the key of the run's aggregate report JSON.""" return f"{prefix}/reports/aggregate.json" def comparison_path(*, prefix: str) -> str: + """Return the key of the run's result-comparison report JSON.""" return f"{prefix}/reports/comparison.json" @@ -123,6 +136,12 @@ def serialize_simulation_frames( calculation_provenance: SPMProvenance | None = None, uk_local_authority_metadata: UKLocalAuthorityMetadata | None = None, ) -> tuple[bytes, RowIdentity]: + """Encode entity frames as one deterministic Parquet payload plus its row identity. + + Rows are ordered by each entity's ``_id`` column, which must exist + and be unique (``ValueError`` otherwise). Source dtypes and any SPM or UK + local authority metadata travel in the Parquet schema metadata. + """ combined, dtypes, identifier_values = _frame_payload(frames) identity_payload = canonical_json_bytes(identifier_values) row_identity = RowIdentity( @@ -166,6 +185,11 @@ def serialize_simulation_frames( def deserialize_simulation_frames(payload: bytes) -> dict[str, pd.DataFrame]: + """Decode a simulation payload back into per-entity frames with their source dtypes. + + Raises ``ValueError`` for an unsupported schema version or missing dtype + metadata. + """ table = pq.read_table(BytesIO(payload)) metadata = table.schema.metadata or {} if ( @@ -209,7 +233,6 @@ def deserialize_simulation_frames(payload: bytes) -> dict[str, pd.DataFrame]: def _parquet_schema_metadata(payload: bytes) -> Mapping[bytes, bytes]: """Read artifact metadata without materializing its simulation rows.""" - parquet_file = pq.ParquetFile(BytesIO(payload)) return parquet_file.schema_arrow.metadata or {} @@ -217,6 +240,7 @@ def _parquet_schema_metadata(payload: bytes) -> Mapping[bytes, bytes]: def deserialize_calculation_provenance( payload: bytes, ) -> SPMProvenance | None: + """Read a simulation artifact's SPM calculation receipt, or None if absent.""" raw = _parquet_schema_metadata(payload).get( PARQUET_CONTRACT.calculation_provenance_metadata_key.encode() ) @@ -229,7 +253,6 @@ def deserialize_uk_local_authority_metadata( payload: bytes, ) -> UKLocalAuthorityMetadata | None: """Read typed UK authority metadata from a Stage 12 simulation artifact.""" - raw = _parquet_schema_metadata(payload).get( PARQUET_CONTRACT.uk_local_authority_metadata_key.encode() ) @@ -239,7 +262,14 @@ def deserialize_uk_local_authority_metadata( class Stage12ArtifactStore: + """Write-once storage for one private bucket's Stage 12 run artifacts. + + A write to a path already holding identical bytes succeeds; different + bytes raise ``RuntimeError``. + """ + def __init__(self, bucket_name: str, *, store: ArtifactStore | None = None) -> None: + """Bind to ``bucket_name``, which must be non-empty (there is no default).""" if not bucket_name: raise ValueError( "Stage 12 requires an explicitly configured private bucket" @@ -279,6 +309,7 @@ def write_input( prefix: str, simulation: PlannedSimulationExecutionInput, ) -> ArtifactReference: + """Store a simulation's planned execution input as canonical JSON.""" return self._write_immutable( input_path(prefix=prefix, role=simulation.role.value), canonical_json_bytes(simulation.model_dump(mode="json")), @@ -294,6 +325,7 @@ def write_simulation( calculation_provenance: SPMProvenance | None = None, uk_local_authority_metadata: UKLocalAuthorityMetadata | None = None, ) -> SimulationArtifactDescriptor: + """Store a simulation's output frames and describe the artifact for reports.""" payload, row_identity = serialize_simulation_frames( frames, calculation_provenance=calculation_provenance, @@ -321,6 +353,7 @@ def write_aggregate( prefix: str, payload: Mapping[str, Any], ) -> ArtifactReference: + """Validate and store the run's aggregate report.""" validated = AggregateReportArtifactPayload.model_validate(payload) return self._write_immutable( aggregate_path(prefix=prefix), @@ -334,6 +367,7 @@ def write_comparison( prefix: str, payload: Mapping[str, Any], ) -> ArtifactReference: + """Validate and store the run's result-comparison report.""" validated = ResultComparisonArtifactPayload.model_validate(payload) return self._write_immutable( comparison_path(prefix=prefix), @@ -342,6 +376,11 @@ def write_comparison( ) def read(self, uri: str) -> bytes: + """Return the bytes at a ``gs://`` URI in this store's bucket. + + Raises ``ValueError`` for another bucket and ``FileNotFoundError`` when + the object is missing. + """ expected_prefix = f"gs://{self.bucket_name}/" if not uri.startswith(expected_prefix): raise ValueError("Stage 12 artifact URI names another bucket") diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_bundle/assertions.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_bundle/assertions.py index 01d090442..404b93593 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_bundle/assertions.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_bundle/assertions.py @@ -13,6 +13,10 @@ def assertion_values(bundle: Stage12BundleManifest) -> dict[str, str]: + """Flatten the bundle's assertable versions under dotted names. + + Per-country fields are named ``countries..``. + """ values = { "policyengine_version": bundle.policyengine_version, "policyengine_requirement": bundle.policyengine_requirement, @@ -41,7 +45,6 @@ def assert_expected_bundle_values( expected: Mapping[str, str], ) -> None: """Check external values without allowing them to select bundle content.""" - actual = assertion_values(bundle) for name, expected_value in expected.items(): if name not in actual: diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_bundle/manifest.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_bundle/manifest.py index 87adc695b..647c87faa 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_bundle/manifest.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_bundle/manifest.py @@ -87,7 +87,6 @@ def normalize_stage12_bundle( supported_countries: Sequence[CountryId] = SUPPORTED_COUNTRIES, ) -> Stage12BundleManifest: """Normalize only values certified by the packaged `.py` bundle.""" - try: return _normalize_stage12_bundle( raw_bundle, @@ -100,6 +99,7 @@ def normalize_stage12_bundle( def bundle_digest(bundle: Stage12BundleManifest) -> str: + """Return the sha256 of the manifest's canonical (key-sorted, compact) JSON.""" payload = json.dumps( bundle.model_dump(mode="json"), sort_keys=True, @@ -111,6 +111,10 @@ def bundle_digest(bundle: Stage12BundleManifest) -> str: def resolve_stage12_bundle(raw_bundle: object) -> ResolvedStage12Bundle: + """Normalize a raw bundle and pair it with its manifest digest. + + Raises ``Stage12BundleError`` when the bundle is incomplete or inconsistent. + """ bundle = normalize_stage12_bundle(raw_bundle) return ResolvedStage12Bundle( bundle=bundle, @@ -120,7 +124,6 @@ def resolve_stage12_bundle(raw_bundle: object) -> ResolvedStage12Bundle: def load_stage12_bundle() -> ResolvedStage12Bundle: """Read the installed PolicyEngine.py release's packaged bundle.""" - from policyengine.bundle import get_current_bundle return resolve_stage12_bundle(get_current_bundle()) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_bundle/parsing.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_bundle/parsing.py index bba9d4ee6..bb1e66e14 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_bundle/parsing.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_bundle/parsing.py @@ -14,12 +14,14 @@ def require_mapping(value: object, path: str) -> Mapping[str, object]: + """Return ``value`` as a string-keyed mapping, or raise naming ``path``.""" if not isinstance(value, Mapping): raise Stage12BundleError(f"PolicyEngine.py bundle field {path} is missing") return {str(key): item for key, item in value.items()} def require_text(value: object, path: str) -> str: + """Return ``value`` if a non-empty string without outer whitespace, else raise.""" if not isinstance(value, str) or not value.strip(): raise Stage12BundleError( f"PolicyEngine.py bundle field {path} must be a non-empty string" @@ -35,6 +37,11 @@ def require_exact_package( packages: Mapping[str, object], package_name: str, ) -> tuple[str, str]: + """Return a package's version and its exact ``name==version`` install requirement. + + Raises ``Stage12BundleError`` unless the entry names itself consistently and + its requirement pins exactly that version. + """ package = require_mapping(packages.get(package_name), f"packages.{package_name}") name = require_text(package.get("name"), f"packages.{package_name}.name") version = require_text( @@ -70,7 +77,6 @@ def normalize_country_bundle( data_releases: Mapping[str, object], ) -> Stage12CountryBundle: """Validate and normalize one country's model and dataset release.""" - country_metadata = require_mapping(countries.get(country), f"countries.{country}") country_package_name = require_text( country_metadata.get("model_package"), diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_parity.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_parity.py index 0a7e4984d..dd8db1d13 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_parity.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_parity.py @@ -12,10 +12,13 @@ @dataclass(frozen=True) class NumericalTolerance: + """Absolute and relative slack for one numeric field; the default is exact.""" + absolute: float = 0.0 relative: float = 0.0 def __post_init__(self) -> None: + """Reject negative tolerances with ``ValueError``.""" if self.absolute < 0 or self.relative < 0: raise ValueError("numerical tolerances must be non-negative") @@ -24,6 +27,7 @@ class ParityMismatch(ValueError): """A bounded diagnostic containing a code and schema path, never values.""" def __init__(self, code: str, path: str) -> None: + """Record the mismatch code and its path, truncated to 255 characters.""" self.code = code self.path = path[:255] super().__init__(f"{code} at {self.path}") @@ -32,6 +36,10 @@ def __init__(self, code: str, path: str) -> None: def canonicalize_simulation_frames( frames: Mapping[str, pd.DataFrame], ) -> dict[str, pd.DataFrame]: + """Sort each entity's rows by ``_id`` and its columns by name. + + Raises ``ParityMismatch`` if an entity lacks that id column or repeats an id. + """ canonical = {} for entity in sorted(frames): identifier = f"{entity}_id" @@ -54,6 +62,12 @@ def compare_simulation_frames( *, tolerances: Mapping[str, NumericalTolerance] | None = None, ) -> None: + """Raise ``ParityMismatch`` at the first difference between two simulation outputs. + + Entities, columns, row ids and dtypes must match exactly. Numeric columns + compare within the tolerance keyed by ``entity.column`` (exact by default, + NaN equal to NaN); other columns must be equal. + """ expected = canonicalize_simulation_frames(incumbent) actual = canonicalize_simulation_frames(candidate) if expected.keys() != actual.keys(): @@ -94,6 +108,12 @@ def compare_aggregate_reports( *, tolerances: Mapping[str, NumericalTolerance] | None = None, ) -> None: + """Raise ``ParityMismatch`` at the first difference between two aggregate reports. + + Mappings and lists must match in shape. Numbers (not booleans) compare + within the tolerance keyed by their path, such as ``budget.total`` or + ``deciles[3]``; any other value must be equal and of the same type. + """ allowed = tolerances or {} def compare(expected: object, actual: object, path: str) -> None: diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_qualification.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_qualification.py index 49c473f36..75c026831 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_qualification.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_qualification.py @@ -82,7 +82,6 @@ class Stage12ParityReceipt(BaseModel): def _local_qualification_runtime() -> ObservabilityRuntime: """Create a local runtime that cannot send qualification data remotely.""" - return ObservabilityRuntime( ObservabilityConfig( service=ServiceIdentity( @@ -106,7 +105,6 @@ def _validate_matching_input( report: ReportExecutionInput, ) -> None: """Reject a qualification that does not compare the same calculation.""" - baseline = report.baseline reform = report.reform bundle = baseline.bundle @@ -223,7 +221,6 @@ def qualify_report_parity( aggregate_tolerances: Mapping[str, NumericalTolerance] | None = None, ) -> Stage12ParityReceipt: """Run both implementations and fail on every unexplained difference.""" - report = ReportExecutionInput.model_validate(report_payload) _validate_matching_input(existing_request, report) output_plan = resolve_report_output_plan(report) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_result_comparison.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_result_comparison.py index 29d18a676..761deabba 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_result_comparison.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_result_comparison.py @@ -40,7 +40,6 @@ def _pointer_token(value: str) -> str: def _leaf_values(value: JsonValue, path: str = "") -> dict[str, JsonValue]: """Flatten JSON containers to scalar leaves using JSON Pointer paths.""" - if isinstance(value, dict): if not value: return {path: value} @@ -105,7 +104,6 @@ def compare_results( compared_at: datetime | None = None, ) -> ResultComparisonArtifactPayload: """Return exact digests and scalar differences outside the numeric tolerance.""" - production = _result_adapter.validate_python(production_result) stage12 = _result_adapter.validate_python(stage12_result) production_bytes = canonical_json_bytes(production) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/aggregation.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/aggregation.py index 4c8b05993..beda86dcf 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/aggregation.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/aggregation.py @@ -36,7 +36,6 @@ def validate_uk_local_authority_metadata( the report's region is national or nation-level and the bundle routes both constituency and local-authority regions to another dataset. """ - if country == "uk": if baseline is None and reform is None: from policyengine_simulation_executor.uk_local_authority_metadata import ( @@ -62,6 +61,11 @@ def validate_aligned_outputs( baseline: SimulationArtifactDescriptor, reform: SimulationArtifactDescriptor, ) -> None: + """Check that two simulation artifacts can be aggregated as this report's pair. + + Raises ``ValueError`` unless the roles, evaluation, bundle, output schema, + output plan and row identities all line up. + """ if baseline.role.value != "baseline" or reform.role.value != "reform": raise ValueError("simulation artifacts have incorrect report roles") if ( @@ -135,6 +139,11 @@ def build_aggregate_report( reform_descriptor: SimulationArtifactDescriptor, uk_local_authority_metadata: UKLocalAuthorityMetadata | None = None, ) -> dict[str, Any]: + """Build the aggregate report payload from retained baseline and reform frames. + + Runs the standard macro output builder over stand-in simulations that load + the stored frames, so no model is recomputed. + """ from policyengine_simulation_executor.segmented_national_reduce import ( PrecomputedSimulation, ) @@ -220,6 +229,11 @@ def build_spm_result( baseline_descriptor: SimulationArtifactDescriptor, reform_descriptor: SimulationArtifactDescriptor, ) -> dict[str, Any]: + """Return the ``spm_provenance`` fields for the report; empty without SPM. + + Raises ``ValueError`` when either artifact lacks an SPM receipt covering the + report year. + """ requested_spm = report.baseline.options.get("spm") if requested_spm is None: return {} diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/comparison.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/comparison.py index 33b0d0210..575217cf9 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/comparison.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/comparison.py @@ -39,7 +39,6 @@ def compare_completed_report( stage_plan: StagePlan | None = None, ) -> None: """Persist an isolated production/Stage 12 comparison after aggregation.""" - production_job_id = context.production_function_call_id if production_job_id is None: return diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py index 1e6aa584a..44ebdba45 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py @@ -96,7 +96,6 @@ def _validate_submitted_parent( parent: ComparisonReportRecord, ) -> None: """Reject dispatch metadata that does not describe this exact report.""" - bundle = report.baseline.bundle expected = { "evaluation_id": report.evaluation_id, @@ -148,7 +147,6 @@ def _claim_parent( store: ComparisonStore, ) -> tuple[ComparisonReportRecord, bool]: """Create the parent in Modal or claim a retry; reject duplicate execution.""" - coordinating_at = datetime.now(UTC) candidate = submitted.model_copy( update={ @@ -232,6 +230,16 @@ def coordinate_report( | None = None, runtime: ObservabilityRuntime | None = None, ) -> dict[str, Any]: + """Run one Stage 12 comparison report end to end; return its aggregate descriptor. + + Claims the parent report (one already running or finished returns a + ``deduplicated`` summary instead), starts or reattaches the baseline and + reform child simulations, and aggregates their digest-checked artifacts; + shadow runs are then compared with the production result. A dispatch + context that disagrees with the report raises ``ValueError`` before the + claim; any later failure marks the report FAILED and raises a + ``RuntimeError`` that carries no calculation detail. + """ report = ReportExecutionInput.model_validate(payload) context = Stage12InvocationContext.model_validate(context_payload) submitted_parent = ComparisonReportRecord.model_validate(parent_payload) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/dependencies.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/dependencies.py index f3d87885f..bbac965d2 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/dependencies.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/dependencies.py @@ -26,35 +26,49 @@ class ComparisonStore(Protocol): + """Report and simulation persistence; ``Stage12PersistenceStore`` implements it.""" + def create_or_resolve_report( self, record: ComparisonReportRecord, - ) -> ComparisonReportPersistenceResult: ... + ) -> ComparisonReportPersistenceResult: + """Insert the report, or return the stored one with the same identity.""" + ... - def get_report(self, evaluation_id: UUID) -> ComparisonReportRecord: ... + def get_report(self, evaluation_id: UUID) -> ComparisonReportRecord: + """Return the stored report; raise ``LookupError`` if there is none.""" + ... - def replace_report( - self, record: ComparisonReportRecord - ) -> ComparisonReportRecord: ... + def replace_report(self, record: ComparisonReportRecord) -> ComparisonReportRecord: + """Write the report's lifecycle and aggregation fields over the stored one.""" + ... def replace_report_result_comparison( self, record: ComparisonReportRecord - ) -> ComparisonReportRecord: ... + ) -> ComparisonReportRecord: + """Write the report's result-comparison fields over the stored one.""" + ... def create_or_resolve_simulation( self, record: ComparisonSimulationRecord, - ) -> ComparisonSimulationPersistenceResult: ... + ) -> ComparisonSimulationPersistenceResult: + """Insert the simulation, or return the stored one with the same identity.""" + ... def get_simulation( self, simulation_execution_id: UUID, - ) -> ComparisonSimulationRecord: ... + ) -> ComparisonSimulationRecord: + """Return the stored simulation; raise ``LookupError`` if there is none.""" + ... def replace_simulation( self, record: ComparisonSimulationRecord, - ) -> ComparisonSimulationRecord: ... + ) -> ComparisonSimulationRecord: + """Write the simulation's mutable fields over the stored one.""" + ... def attach_simulation_invocation( self, @@ -63,16 +77,27 @@ def attach_simulation_invocation( expected_placeholder: str, modal_invocation_id: str, updated_at: datetime, - ) -> ComparisonSimulationRecord: ... + ) -> ComparisonSimulationRecord: + """Replace ``expected_placeholder`` with the child's real Modal invocation id. + + Repeating an attach that already succeeded returns the stored record. + """ + ... class ChildCall(Protocol): + """Handle to a spawned child simulation, such as a ``modal.FunctionCall``.""" + object_id: str - def get(self, *, timeout: float | None = None) -> object: ... + def get(self, *, timeout: float | None = None) -> object: + """Wait up to ``timeout`` seconds (indefinitely if None) for the result.""" + ... class ChildInvoker(Protocol): + """Launcher for child simulation functions that can reattach to running calls.""" + def spawn( self, *, @@ -82,12 +107,18 @@ def spawn( simulation: dict[str, Any], context: dict[str, Any], observability_context: dict[str, Any] | None, - ) -> ChildCall: ... + ) -> ChildCall: + """Start the named deployed function on one simulation without waiting.""" + ... - def restore(self, invocation_id: str) -> ChildCall: ... + def restore(self, invocation_id: str) -> ChildCall: + """Return a handle to a call previously started, by its invocation id.""" + ... class ModalChildInvoker: + """``ChildInvoker`` backed by functions deployed to Modal.""" + def spawn( self, *, @@ -98,6 +129,7 @@ def spawn( context: dict[str, Any], observability_context: dict[str, Any] | None, ) -> ChildCall: + """Look up ``application_name.function_name`` in ``environment``; spawn it.""" from importlib import import_module modal: Any = import_module("modal") @@ -113,6 +145,7 @@ def spawn( ) def restore(self, invocation_id: str) -> ChildCall: + """Return the ``modal.FunctionCall`` with this id.""" from importlib import import_module modal: Any = import_module("modal") @@ -122,6 +155,11 @@ def restore(self, invocation_id: str) -> ChildCall: def runtime_store( runtime: ObservabilityRuntime | None = None, ) -> Stage12PersistenceStore: + """Build the persistence store from ``STAGE12_DATABASE_URL``. + + Raises ``ValueError`` when the variable is unset or not a complete + ``postgresql://`` URL. + """ # Every operation opens its own pooler session, so each one retries a # refused connection rather than failing the report during a burst. return Stage12PersistenceStore( @@ -131,4 +169,5 @@ def runtime_store( def artifact_store() -> Stage12ArtifactStore: + """Build the store for ``STAGE12_ARTIFACT_BUCKET``; ValueError if it is unset.""" return Stage12ArtifactStore(os.environ.get("STAGE12_ARTIFACT_BUCKET", "")) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/output_planning.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/output_planning.py index ed2731dca..b6e5f4c58 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/output_planning.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/output_planning.py @@ -38,7 +38,9 @@ class OutputVariableModel(Protocol): def resolve_entity_variables( self, simulation: Simulation, - ) -> dict[str, list[str]]: ... + ) -> dict[str, list[str]]: + """Return the output variables each entity must carry for ``simulation``.""" + ... StageScope = Callable[ @@ -76,7 +78,6 @@ def _planning_simulation( intentionally bypassed. This object must never call ``ensure`` or access a dataset. """ - return Simulation.model_construct( id=f"stage12-output-plan-{simulation.role.value}", policy=dict(simulation.policy), @@ -102,7 +103,6 @@ def _required_dataset_variables( aggregates: tuple[ReportAggregate, ...], ) -> dict[str, tuple[str, ...]]: """Return input-dataset columns copied into the materialized output.""" - if country == "uk" and ReportAggregate.GEOGRAPHIC in aggregates: return UK_GEOGRAPHIC_DATASET_VARIABLES return {} @@ -142,7 +142,6 @@ def resolve_report_output_plan( stage_scope: StageScope = _unmonitored_stage_scope, ) -> Stage12OutputPlan: """Resolve one country-owned output schema for both report simulations.""" - country = report.baseline.geography.country stage_attributes = {"country": country} with stage_scope(Stage.STAGE12_COUNTRY_MODEL_LOAD, stage_attributes): @@ -207,7 +206,6 @@ def plan_simulation_input( output_plan: Stage12OutputPlan, ) -> PlannedSimulationExecutionInput: """Attach a coordinator-resolved schema to one child input.""" - return PlannedSimulationExecutionInput.model_validate( { **simulation.model_dump(mode="json"), @@ -221,7 +219,6 @@ def apply_output_plan( output_plan: Stage12OutputPlan, ) -> None: """Install required extra variables before a live simulation is ensured.""" - extras = { entity: list(variables) for entity, variables in (simulation.extra_variables or {}).items() @@ -256,7 +253,6 @@ def validate_output_frames( output_plan: Stage12OutputPlan, ) -> None: """Require every planned entity and column in materialized output frames.""" - missing_entities = [ entity.entity for entity in output_plan.entities if entity.entity not in frames ] diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py index fd5601e24..5b50eea17 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py @@ -47,12 +47,15 @@ @dataclass(frozen=True) class SimulationCalculation: + """A calculator's entity output frames plus metadata stored alongside them.""" + frames: Mapping[str, pd.DataFrame] calculation_provenance: SPMProvenance | None = None uk_local_authority_metadata: UKLocalAuthorityMetadata | None = None def simulation_input_sha256(simulation: SimulationExecutionInput) -> str: + """Hash a simulation input ignoring its run ids, so equal inputs hash equally.""" normalized = simulation.model_dump( mode="json", exclude={"evaluation_id", "simulation_execution_id"}, @@ -115,7 +118,6 @@ def calculate_simulation_frames( runtime: ObservabilityRuntime | None = None, ) -> SimulationCalculation: """Run one policy and return the coordinator-planned entity output tables.""" - if simulation.population.kind != "dataset": raise ValueError("Stage 12 society-wide worker requires a dataset population") if simulation.population.artifact.uri != simulation.bundle.dataset.uri: @@ -250,6 +252,11 @@ def descriptor_from_record( child: ComparisonSimulationRecord, simulation: PlannedSimulationExecutionInput, ) -> SimulationArtifactDescriptor: + """Rebuild the artifact descriptor of a child that already succeeded. + + ``size_bytes`` is 0 because the record does not store it; raises + ``ValueError`` if the record lacks any output metadata. + """ output_uri = child.output_uri output_sha256 = child.output_sha256 output_schema_version = child.output_schema_version @@ -305,6 +312,12 @@ def run_single_simulation( ] = calculate_simulation_frames, runtime: ObservabilityRuntime | None = None, ) -> dict[str, Any]: + """Calculate one planned simulation, store its artifact, and return its descriptor. + + A child already marked succeeded returns its stored descriptor without + recalculating. A failure after the run is marked running marks it FAILED + and raises a ``RuntimeError`` that carries no calculation detail. + """ simulation = PlannedSimulationExecutionInput.model_validate(payload) context = Stage12InvocationContext.model_validate(context_payload) _require_context(simulation, context, required_country=required_country) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_worker_validation.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_worker_validation.py index 0175003cb..4fc038708 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_worker_validation.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_worker_validation.py @@ -111,7 +111,6 @@ def validate_country_worker( ), ) -> dict[str, Any]: """Validate imports, installed bundle, secrets, data, and calculation.""" - runtime_environment = environment if environment is not None else os.environ _require_secret_environment(runtime_environment) # Import the same typed persistence package and coordinator modules loaded diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/uk_local_authority_metadata.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/uk_local_authority_metadata.py index e30e3c6f2..b51f66918 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/uk_local_authority_metadata.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/uk_local_authority_metadata.py @@ -107,7 +107,6 @@ def metadata_for( boundary_version: UKLocalAuthorityBoundaryVersion, ) -> Mapping[str, LocalAuthorityDisplayMetadata]: """Combine common names with one boundary version's rendering coordinates.""" - return MappingProxyType( { code: LocalAuthorityDisplayMetadata( @@ -195,7 +194,6 @@ def _load_coordinates(filename: str) -> Mapping[str, LocalAuthorityCoordinate]: @lru_cache(maxsize=1) def load_uk_local_authority_resources() -> UKLocalAuthorityResources: """Load and validate every packaged local-authority resource.""" - names = _load_names() lad22 = _load_coordinates("coordinates_lad22.csv") lad23 = _load_coordinates("coordinates_lad23.csv") @@ -252,7 +250,6 @@ def detect_uk_local_authority_boundary_version( values: Iterable[object], ) -> UKLocalAuthorityMetadata: """Identify LAD22 or LAD23 from the unscoped dataset's authority codes.""" - observed_codes = frozenset( code for value in values @@ -298,7 +295,6 @@ def uk_area_regions_routed() -> bool: region types have one may the UK default, a national file, carry no area codes: their runs and breakdowns then belong to that local-area dataset. """ - from policyengine.provenance.manifest import get_release_manifest region_datasets = get_release_manifest("uk").region_datasets @@ -312,7 +308,6 @@ def uk_area_codes_required(region_code: str | None) -> bool: nation-level run may use a dataset without them only when the bundle routes both area region types to another dataset. """ - region_type = (region_code or "").split("/", maxsplit=1)[0] if region_type in UK_AREA_REGION_TYPES: return True @@ -331,7 +326,6 @@ def detect_uk_local_authority_metadata( that does not need area codes (see ``uk_area_codes_required``); a missing column is never read as "this dataset is national-only" on its own. """ - if country != "uk": return None data = getattr(dataset, "data", None) @@ -353,7 +347,6 @@ def detect_uk_local_authority_metadata( def uk_hdf_household_has_local_authority_codes(dataset_path: str) -> bool: """Whether an installed UK HDF dataset's household table has ``la_code_oa``.""" - with pd.HDFStore(dataset_path, mode="r") as store: if "household" not in store: raise ValueError("UK dataset contains no household table") @@ -367,7 +360,6 @@ def detect_uk_local_authority_metadata_from_hdf( dataset_path: str, ) -> UKLocalAuthorityMetadata: """Validate an installed UK HDF dataset against the packaged metadata.""" - observed_codes: set[object] = set() with pd.HDFStore(dataset_path, mode="r") as store: if "/household" not in store.keys(): diff --git a/projects/policyengine-simulation-executor/tests/integration/test_budget_window_ephemeral_modal.py b/projects/policyengine-simulation-executor/tests/integration/test_budget_window_ephemeral_modal.py index 13e0503b5..ec842b7f0 100644 --- a/projects/policyengine-simulation-executor/tests/integration/test_budget_window_ephemeral_modal.py +++ b/projects/policyengine-simulation-executor/tests/integration/test_budget_window_ephemeral_modal.py @@ -44,5 +44,4 @@ def test_budget_window_submission_reaches_complete( intentionally empty so the integration bucket exists as a real directory with a valid Python test module. """ - assert modal_integration_environment.startswith("http") diff --git a/projects/policyengine-simulation-executor/tests/test_app_redaction.py b/projects/policyengine-simulation-executor/tests/test_app_redaction.py index ffee1f9a5..ea1d9e68e 100644 --- a/projects/policyengine-simulation-executor/tests/test_app_redaction.py +++ b/projects/policyengine-simulation-executor/tests/test_app_redaction.py @@ -37,10 +37,12 @@ def test_redact_params_strips_signed_urls_and_reform_bodies(): def test_redact_params_strips_all_underscore_prefixed_keys(): - """Underscore-prefixed internal keys must never reach observability - attributes: the backend rejects attribute keys starting with '_', so a - leaked key crashes the span. Regression for the segmented-national - child failure where `_emit_microdata` reached operation() (#640).""" + """Underscore-prefixed internal keys must never reach observability attributes. + + The backend rejects attribute keys starting with '_', so a leaked key + crashes the span. Regression for the segmented-national child failure + where `_emit_microdata` reached operation() (#640). + """ params = { "country": "us", "scope": "macro", diff --git a/projects/policyengine-simulation-executor/tests/test_artifact_store.py b/projects/policyengine-simulation-executor/tests/test_artifact_store.py index d9329aae6..baa028cdc 100644 --- a/projects/policyengine-simulation-executor/tests/test_artifact_store.py +++ b/projects/policyengine-simulation-executor/tests/test_artifact_store.py @@ -115,8 +115,11 @@ def test_download_creates_parent_dirs(store, fake_client, tmp_path): def test_failed_download_leaves_no_file_at_the_final_name(store, fake_client, tmp_path): - """Callers guard on exists(): a truncated file surviving at the final - name would be trusted as complete by the next run in the container.""" + """A failed download must not leave a truncated file at the final name. + + Callers guard on exists(), so a truncated file surviving at the final + name would be trusted as complete by the next run in the container. + """ fake_client.objects["datasets/us/d/populace_year_2026.h5"] = b"content" fake_client.broken_downloads.add("datasets/us/d/populace_year_2026.h5") destination = tmp_path / "populace_year_2026.h5" diff --git a/projects/policyengine-simulation-executor/tests/test_baseline_artifacts.py b/projects/policyengine-simulation-executor/tests/test_baseline_artifacts.py index 8c2f06865..fd2e61e6e 100644 --- a/projects/policyengine-simulation-executor/tests/test_baseline_artifacts.py +++ b/projects/policyengine-simulation-executor/tests/test_baseline_artifacts.py @@ -279,11 +279,12 @@ def test_cache_hit_is_validated_too(self, fresh_cache): assert model.calls == ["run", "save"] def test_stale_cache_entry_is_replaced_after_recompute(self, fresh_cache): - """After a recompute that started from a cache hit on a DIFFERENT - (stale) object, the cache must hold the recomputed simulation. - LRUCache.add alone keeps the old value for an existing key, which - would make this request's second ensure() clobber the fresh output - and run the full model twice.""" + """A recompute after a stale cache hit must leave the fresh simulation cached. + + The stale hit is on a DIFFERENT object. LRUCache.add alone keeps the + old value for an existing key, which would make this request's second + ensure() clobber the fresh output and run the full model twice. + """ stale = _make_sim(FakeModelVersion()) stale.output_dataset = _output(_incomplete_frames()) fresh_cache.add("bl1-test", stale) @@ -324,9 +325,11 @@ def test_loaded_output_without_data_recomputes(self, fresh_cache): def test_second_ensure_keeps_a_genuine_miss(self, fresh_cache): """Every request ensures the baseline twice (analysis + deciles). + The second call self-hits the in-process cache with a complete column set — it must not overwrite the miss the request actually - experienced (the rollout metric depends on it).""" + experienced (the rollout metric depends on it). + """ model = FakeModelVersion(load_result="absent") sim = _make_sim(model) sim.ensure() @@ -564,9 +567,12 @@ def frame(entity, extra=None): class DiskModelVersion: - """Duck model version reusing the REAL load/save implementations, so - these tests exercise genuine h5 files on disk — including the exception - type a missing artifact raises — without loading the US tax system.""" + """Duck model version reusing the REAL load/save implementations. + + These tests therefore exercise genuine h5 files on disk — including the + exception type a missing artifact raises — without loading the US tax + system. + """ # PolicyEngine 6 reads this class-level value while restoring SPM metadata. # These disk tests are model-neutral, so keep the non-US branch selected. @@ -936,10 +942,13 @@ def test_receipt_for_another_year_recomputes(self, fresh_cache): assert model.calls == ["load", "run", "save"] def test_cached_entry_with_an_unusable_receipt_is_revalidated(self, fresh_cache): - """A cache hit skips the load entirely, so the entry's receipt is - the only thing standing between the request and someone else's - output. A tax-only run leaves a receipt with no years under the same - selection; the guard has to notice and recompute.""" + """A cache hit whose receipt is unusable must be recomputed. + + A cache hit skips the load entirely, so the entry's receipt is the + only thing standing between the request and someone else's output. A + tax-only run leaves a receipt with no years under the same selection; + the guard has to notice and recompute. + """ stale = _make_spm_sim(SPMModelVersion()) stale.output_dataset = _output(_complete_frames()) stale.spm_receipt = _spm_receipt(year="2024") @@ -953,11 +962,13 @@ def test_cached_entry_with_an_unusable_receipt_is_revalidated(self, fresh_cache) assert model.calls == ["run", "save"] def test_cached_selection_cannot_replace_the_request(self, fresh_cache): - """Defense in depth: an entry whose selection disagrees with the key - it is filed under. A correct wrapper never writes one — the key is - derived from the selection — but a cache hit restores the entry's - selection onto this request without revalidating, so the guard - keeps the request's own selection and recomputes under it.""" + """Defense in depth: an entry whose selection disagrees with its key. + + A correct wrapper never writes one — the key is derived from the + selection — but a cache hit restores the entry's selection onto this + request without revalidating, so the guard keeps the request's own + selection and recomputes under it. + """ other = {**SPM_SELECTION, "scenario": "zero_real"} stale = _make_spm_sim(SPMModelVersion(), spm=other) stale.output_dataset = _output(_complete_frames()) @@ -992,8 +1003,10 @@ def test_recompute_replaces_the_cache_with_this_request_s_selection( def test_missing_columns_are_decided_before_the_receipt(self, fresh_cache): """A tax-only artifact can legitimately carry no usable receipt. + Columns are checked first so the outcome is attributed to the gap - that actually exists, and the receipt is never consulted.""" + that actually exists, and the receipt is never consulted. + """ model = SPMModelVersion( load_result="incomplete", stored_config=deepcopy(SPM_SELECTION), @@ -1005,8 +1018,11 @@ def test_missing_columns_are_decided_before_the_receipt(self, fresh_cache): assert sim._provenance_reads == [] def test_a_request_without_a_selection_never_validates_receipts(self, fresh_cache): - """The legacy path: no selection, so a complete artifact is a hit - even though the wrapper surface exists and holds no receipt.""" + """Without a selection, a complete artifact is a hit even with no receipt. + + This is the legacy path: the wrapper surface exists but holds no + receipt. + """ model = SPMModelVersion(stored_config=None, stored_receipt=None) sim = _make_spm_sim(model, sim_id="bl1-legacy") sim.spm = None @@ -1050,9 +1066,11 @@ def test_the_cache_is_keyed_by_storage_id_not_simulation_id(self, fresh_cache): assert fresh_cache.get(first.id) is None def test_the_replaced_cache_entry_is_a_snapshot(self, fresh_cache): - """The replacement is a copy so that re-pointing this simulation at - another selection and running it again cannot overwrite the output - the earlier key names.""" + """The replacement cache entry is a copy, not the live simulation. + + Re-pointing this simulation at another selection and running it again + therefore cannot overwrite the output the earlier key names. + """ other = {**SPM_SELECTION, "scenario": "zero_real"} model = SPMModelVersion(stored_config=other, stored_receipt=_spm_receipt(other)) sim = _make_spm_sim(model) diff --git a/projects/policyengine-simulation-executor/tests/test_budget_window_batch.py b/projects/policyengine-simulation-executor/tests/test_budget_window_batch.py index e0642c76e..2731041ab 100644 --- a/projects/policyengine-simulation-executor/tests/test_budget_window_batch.py +++ b/projects/policyengine-simulation-executor/tests/test_budget_window_batch.py @@ -302,11 +302,12 @@ def test_run_budget_window_batch_impl_marks_failure( def test_scheduler_sleep_exponentially_backs_off_then_resets_on_progress( monkeypatch, mock_batch_modal, observability_runtime ): - """When every poll sees nothing resolve the runner should sleep with + """Back off exponentially while polls see no progress; reset on progress. + + When every poll sees nothing resolve, the runner should sleep with increasing intervals, capped at the configured max. Any progress in a subsequent poll should reset the cadence to the initial interval. """ - request, payload = _build_parent_payload(window_size=1) _seed_parent_batch(request, mock_batch_modal["parent_call_id"]) @@ -362,13 +363,13 @@ def test_scheduler_sleep_exponentially_backs_off_then_resets_on_progress( def test_scheduler_publishes_bounded_poll_aggregates( monkeypatch, mock_batch_modal, observability_runtime ): - """Poll/sleep telemetry must be published as overwriting aggregate - attributes, not one segment per probe: segments append nodes to the - operation's in-memory segment tree, so a near-timeout batch polling for - an hour would grow memory without bound and emit a final operation log - line large enough for Modal's log pipeline to truncate. - """ + """Publish poll/sleep telemetry as overwriting aggregates, not per-probe segments. + Segments append nodes to the operation's in-memory segment tree, so a + near-timeout batch polling for an hour would grow memory without bound + and emit a final operation log line large enough for Modal's log + pipeline to truncate. + """ request, payload = _build_parent_payload(window_size=1) _seed_parent_batch(request, mock_batch_modal["parent_call_id"]) diff --git a/projects/policyengine-simulation-executor/tests/test_budget_window_context.py b/projects/policyengine-simulation-executor/tests/test_budget_window_context.py index 1a1692dc0..afc312903 100644 --- a/projects/policyengine-simulation-executor/tests/test_budget_window_context.py +++ b/projects/policyengine-simulation-executor/tests/test_budget_window_context.py @@ -78,8 +78,11 @@ def test_build_child_simulation_request_removes_batch_only_fields(): def test_build_child_simulation_request_pins_children_monolithic(): - """A national budget-window child must not fan out again (years x 21 - containers); enabling nested segmentation is a deliberate follow-up.""" + """A national budget-window child must not fan out again. + + Fanning out would cost years x 21 containers; enabling nested + segmentation is a deliberate follow-up. + """ _, payload = _build_parent_payload() context = build_batch_context(payload, batch_job_id="parent-123") diff --git a/projects/policyengine-simulation-executor/tests/test_budget_window_results.py b/projects/policyengine-simulation-executor/tests/test_budget_window_results.py index cd8319c56..d17c416de 100644 --- a/projects/policyengine-simulation-executor/tests/test_budget_window_results.py +++ b/projects/policyengine-simulation-executor/tests/test_budget_window_results.py @@ -33,11 +33,13 @@ def test_extract_annual_impact_matches_v1_shape(): def test_extract_annual_impact_defaults_state_tax_for_uk_child_result(): - """UK worker results omit ``state_tax_revenue_impact`` because the UK - microsimulation has no devolved fiscal layer. The aggregator should - accept these results and treat the state component as zero instead of - failing with "missing numeric budget.state_tax_revenue_impact".""" + """Treat a UK child result's missing state tax revenue impact as zero. + UK worker results omit ``state_tax_revenue_impact`` because the UK + microsimulation has no devolved fiscal layer. The aggregator should + accept these results instead of failing with "missing numeric + budget.state_tax_revenue_impact". + """ annual = extract_annual_impact( simulation_year="2026", child_result={ diff --git a/projects/policyengine-simulation-executor/tests/test_budget_window_scheduler.py b/projects/policyengine-simulation-executor/tests/test_budget_window_scheduler.py index ca41b7b8b..e44192659 100644 --- a/projects/policyengine-simulation-executor/tests/test_budget_window_scheduler.py +++ b/projects/policyengine-simulation-executor/tests/test_budget_window_scheduler.py @@ -5,7 +5,8 @@ the ``modal`` module with in-memory fakes and run the scheduler against those fakes. Renamed to reflect the actual scope (#457). Real Modal integration tests live under ``tests/integration/`` and are skipped by -default.""" +default. +""" from __future__ import annotations @@ -38,7 +39,7 @@ def spm_child_result(runtime, simulation_year, *, receipt_year): - """A child result carrying a canonical SPM receipt for ``receipt_year``.""" + """Return a child result carrying a canonical SPM receipt for ``receipt_year``.""" receipt = { "schema_version": "canonical-spm-provenance-v2", "forecast_id": "test-only", diff --git a/projects/policyengine-simulation-executor/tests/test_image_setup_fetch.py b/projects/policyengine-simulation-executor/tests/test_image_setup_fetch.py index 33837f67a..dcd88b211 100644 --- a/projects/policyengine-simulation-executor/tests/test_image_setup_fetch.py +++ b/projects/policyengine-simulation-executor/tests/test_image_setup_fetch.py @@ -44,9 +44,11 @@ def bucket(self, name): def _manifest(): - """The wire shape fetch_artifacts consumes, derived from the schema - (the layer itself cannot import the model, so it speaks the canonical - dict — building the fake through ArtifactManifest keeps it honest).""" + """Return the wire shape fetch_artifacts consumes, derived from the schema. + + The layer itself cannot import the model, so it speaks the canonical + dict — building the fake through ArtifactManifest keeps it honest. + """ return ArtifactManifest.model_validate( { "schema": "mf1", @@ -123,8 +125,7 @@ def test_unresolved_manifest_fails_with_the_digest_env_var_named( assert fake_client.downloads == [] def test_overwrites_a_preexisting_same_named_file(self, fake_client, data_folder): - """No skip-if-exists: stale same-named bytes must never survive - into the image.""" + """No skip-if-exists: stale same-named bytes must never reach the image.""" data_folder.mkdir(parents=True) stale = data_folder / "populace_year_2026.h5" stale.write_bytes(b"stale") @@ -169,8 +170,7 @@ def test_passes_when_the_installed_data_version_matches( def test_proceeds_without_an_installed_receipt_file( self, monkeypatch, tmp_path, fake_client, data_folder ): - """Lenient by design: local fault-drill runs have no bundle - receipt file.""" + """Lenient by design: local fault-drill runs have no bundle receipt file.""" monkeypatch.setenv( "POLICYENGINE_BUNDLE_RECEIPT", str(tmp_path / "does-not-exist.json") ) @@ -179,10 +179,12 @@ def test_proceeds_without_an_installed_receipt_file( class TestCredentialShim: - """The client=None path: credentials are built in memory from the - secret blob via service_account.Credentials.from_service_account_info + """The client=None path: credentials are built in memory from the secret blob. + + They come from service_account.Credentials.from_service_account_info (never written to disk — the layer's filesystem is committed into the - image), falling back to ambient ADC when no blob is present.""" + image), falling back to ambient ADC when no blob is present. + """ @pytest.fixture def patched_google(self, monkeypatch, fake_client): @@ -249,16 +251,21 @@ def test_falls_back_to_ambient_adc_without_a_blob( def test_fetch_never_writes_credentials_to_disk(): - """A run_function layer commits its container filesystem into the - image, so the shim must never materialize the key as a file.""" + """The credential shim must never materialize the key as a file. + + A run_function layer commits its container filesystem into the image. + """ source = Path(image_setup.__file__).read_text(encoding="utf-8") assert "tempfile" not in source def test_image_setup_module_stays_self_contained(): - """The fetch layer runs before add_local_python_source, so the module - must have no module-level imports and must never import the executor - package, even lazily inside a function body.""" + """The image setup module must not import the executor package at all. + + The fetch layer runs before add_local_python_source, so the module must + have no module-level imports and must never import the executor package, + even lazily inside a function body. + """ source = Path(image_setup.__file__).read_text(encoding="utf-8") tree = ast.parse(source) diff --git a/projects/policyengine-simulation-executor/tests/test_modal_bundle_image.py b/projects/policyengine-simulation-executor/tests/test_modal_bundle_image.py index db42e8e02..8d57ad6d6 100644 --- a/projects/policyengine-simulation-executor/tests/test_modal_bundle_image.py +++ b/projects/policyengine-simulation-executor/tests/test_modal_bundle_image.py @@ -122,9 +122,12 @@ def test_modal_image_uses_policyengine_bundle_install(monkeypatch): def _fake_manifest(): - """A schema-valid store payload: app.py validates what it reads, so - the fake must satisfy ArtifactManifest — deriving it from the model - keeps the test from drifting off the wire shape.""" + """Return a schema-valid store manifest payload. + + app.py validates what it reads, so the fake must satisfy + ArtifactManifest — deriving it from the model keeps the test from + drifting off the wire shape. + """ from policyengine_simulation_executor.precompute_models import ArtifactManifest return ArtifactManifest.model_validate( @@ -243,11 +246,13 @@ def read_manifest(self, digest): def test_modal_image_uses_failing_sentinel_without_manifest_digest(monkeypatch): - """Without a digest the app must still import — the precompute and - smoke apps import it where no digest can exist — but the fetch layer - args must carry a None manifest, which makes fetch_artifacts fail the - build loudly. A digest-less deploy cannot silently ship an - artifact-less image.""" + """Without a digest the app imports, but its fetch layer fails the build. + + The precompute and smoke apps import the app where no digest can exist, + so the import must succeed. The fetch layer args must carry a None + manifest, which makes fetch_artifacts fail the build loudly: a + digest-less deploy cannot silently ship an artifact-less image. + """ install_fake_modal(monkeypatch) monkeypatch.setenv("POLICYENGINE_VERSION", "4.19.1") monkeypatch.setenv("POLICYENGINE_CORE_VERSION", "3.27.1") diff --git a/projects/policyengine-simulation-executor/tests/test_modal_scripts.py b/projects/policyengine-simulation-executor/tests/test_modal_scripts.py index 3cbe1ab1d..3b9b9f45f 100644 --- a/projects/policyengine-simulation-executor/tests/test_modal_scripts.py +++ b/projects/policyengine-simulation-executor/tests/test_modal_scripts.py @@ -14,7 +14,7 @@ class TestModalExtractVersions: - """Tests for modal-extract-versions.sh""" + """Tests for modal-extract-versions.sh.""" script = SCRIPTS_DIR / "modal-extract-versions.sh" @@ -91,7 +91,7 @@ def test_deploy_workflow_threads_force_latest_to_routing_publisher(self): class TestModalHealthCheck: - """Tests for modal-health-check.sh""" + """Tests for modal-health-check.sh.""" script = SCRIPTS_DIR / "modal-health-check.sh" @@ -258,7 +258,7 @@ def test_deployment_validates_before_syncing_the_exact_value(self): class TestSimulationDeploymentSummary: - """Tests for simulation-deployment-summary.sh""" + """Tests for simulation-deployment-summary.sh.""" script = SCRIPTS_DIR / "simulation-deployment-summary.sh" @@ -335,7 +335,7 @@ def test_generates_skipped_summary(self, temp_github_step_summary): class TestModalSyncSecrets: - """Tests for modal-sync-secrets.sh""" + """Tests for modal-sync-secrets.sh.""" script = SCRIPTS_DIR / "modal-sync-secrets.sh" @@ -513,7 +513,7 @@ def test_gcp_secret_sync_failure_stops_deployment(self, tmp_path): class TestModalPrecompute: - """Tests for modal-precompute.sh""" + """Tests for modal-precompute.sh.""" script = SCRIPTS_DIR / "modal-precompute.sh" @@ -707,8 +707,10 @@ def test_deploy_workflow_threads_force_recompute_to_script(self): ) def test_deploy_workflow_threads_manifest_and_store_credentials_to_deploy(self): - """The deploy step resolves the manifest from GCS on the runner, - so it needs the digest, the store bucket, and GCP credentials.""" + """The deploy step gets the digest, the store bucket, and GCP credentials. + + It resolves the manifest from GCS on the runner, so it needs all three. + """ reusable_workflow = ( REPO_ROOT / ".github" / "workflows" / "simulation-deploy.reusable.yml" ).read_text(encoding="utf-8") @@ -755,7 +757,7 @@ def test_precompute_is_gated_by_the_shared_secret_sync(self): class TestModalRecordDeployment: - """Tests for modal-record-deployment.sh""" + """Tests for modal-record-deployment.sh.""" script = SCRIPTS_DIR / "modal-record-deployment.sh" @@ -867,7 +869,7 @@ def test_deploy_workflow_records_marker(self): class TestModalGetUrl: - """Tests for modal-get-url.sh""" + """Tests for modal-get-url.sh.""" script = SCRIPTS_DIR / "modal-get-url.sh" @@ -895,7 +897,7 @@ def test_requires_modal_environment_argument(self): class TestModalSetupEnvironments: - """Tests for modal-setup-environments.sh""" + """Tests for modal-setup-environments.sh.""" script = SCRIPTS_DIR / "modal-setup-environments.sh" @@ -914,7 +916,7 @@ def test_script_syntax(self): class TestSimulationRunIntegTests: - """Tests for simulation-run-integ-tests.sh""" + """Tests for simulation-run-integ-tests.sh.""" script = SCRIPTS_DIR / "simulation-run-integ-tests.sh" diff --git a/projects/policyengine-simulation-executor/tests/test_national_partition.py b/projects/policyengine-simulation-executor/tests/test_national_partition.py index 948564d7b..0045acf4a 100644 --- a/projects/policyengine-simulation-executor/tests/test_national_partition.py +++ b/projects/policyengine-simulation-executor/tests/test_national_partition.py @@ -41,10 +41,13 @@ def test__groups_are_nonempty(self): class TestPartitionMatchesModelRegistry: def test__partition_covers_exactly_the_registry_state_regions(self): - """A policyengine pin bump that adds/renames a state-level region - must fail here, forcing a measured partition rebalance. (At runtime - the runner also assigns uncovered registry states to the last group - as a stopgap — this test is what makes the drift loud in CI.)""" + """The partition must cover exactly the model registry's state regions. + + A policyengine pin bump that adds/renames a state-level region must + fail here, forcing a measured partition rebalance. (At runtime the + runner also assigns uncovered registry states to the last group as a + stopgap — this test is what makes the drift loud in CI.) + """ from policyengine_simulation_executor.simulation_runtime import ( _country_module, ) diff --git a/projects/policyengine-simulation-executor/tests/test_precompute.py b/projects/policyengine-simulation-executor/tests/test_precompute.py index eba242ee9..9b2d97a88 100644 --- a/projects/policyengine-simulation-executor/tests/test_precompute.py +++ b/projects/policyengine-simulation-executor/tests/test_precompute.py @@ -1,10 +1,10 @@ -"""Precompute library: planning, manifest assembly, orchestration, and the -writer==reader identity contract. +"""Tests for the precompute library's planning, manifests and orchestration. -No Modal involved — the library takes function handles as parameters, so -the orchestration runs against plain fakes here. Everything structured is -a strict model from precompute_models; the fakes speak dicts, exactly -like Modal's serialization boundary does. +They also pin the writer==reader identity contract. No Modal involved — +the library takes function handles as parameters, so the orchestration runs +against plain fakes here. Everything structured is a strict model from +precompute_models; the fakes speak dicts, exactly like Modal's +serialization boundary does. """ import sys @@ -97,8 +97,11 @@ def test_missing_keys_are_rejected(self): PrecomputePlan.model_validate(payload) def test_artifact_outcome_mirrors_runtime_constants(self): - """The Literal must track baseline_artifacts.OUTCOME_*: drift would - surface as a ValidationError mid-precompute-run, not in review.""" + """The outcome Literal must track baseline_artifacts.OUTCOME_*. + + Drift would surface as a ValidationError mid-precompute-run, not in + review. + """ from typing import get_args from policyengine_simulation_executor import baseline_artifacts as ba @@ -126,9 +129,11 @@ def test_years_are_in_priority_order(self): assert precompute.PRECOMPUTE_YEARS == [2026, 2027, 2025] def test_precompute_years_cover_runtime_default_year(self): - """The image bakes exactly the precomputed years, so the runtime - default year must be one of them or every default request pays a - runtime dataset build.""" + """The runtime default year must be one of the precomputed years. + + The image bakes exactly the precomputed years, so otherwise every + default request pays a runtime dataset build. + """ from policyengine_simulation_executor.simulation_runtime import DEFAULT_YEAR assert DEFAULT_YEAR in precompute.PRECOMPUTE_YEARS @@ -152,9 +157,11 @@ def test_manifest_lists_every_artifact_with_runtime_filenames(self): ] def test_manifest_payload_is_wire_compatible(self): - """The manifest is content-addressed: the typed model must digest - identically to the raw dict shape already published to the store - (key names included — note "schema", not "manifest_schema").""" + """The typed manifest must digest identically to the published dict shape. + + The manifest is content-addressed, so key names matter too — note + "schema", not "manifest_schema". + """ raw = { "schema": "mf1", "country": "us", @@ -401,8 +408,11 @@ def test_precompute_identity_equals_runtime_id(self, identity_stubs, fake_region assert writer_identity.store_path.endswith(f"/{writer_identity.storage_id}.h5") def test_dataset_filename_matches_runtime_stem_lookup(self, monkeypatch): - """Real-manifest coverage: the artifact filename equals the exact - name the runtime's ensure_datasets existence check resolves.""" + """The dataset artifact filename is the name the runtime looks for. + + Real-manifest coverage: it equals the exact name the runtime's + ensure_datasets existence check resolves. + """ from policyengine.provenance.manifest import ( dataset_logical_name, resolve_dataset_reference, @@ -812,7 +822,7 @@ def test_refuses_a_mismatched_planned_storage_id(self, cohort_stubs): @staticmethod def _planned_identity(selection): - """A real ``BaselineArtifactIdentity``, not a restated expression. + """Return a real ``BaselineArtifactIdentity``, not a restated expression. The plan the deploy publishes comes from this object's properties, so keying the test off them is what makes the write path a @@ -896,9 +906,12 @@ def test_publishes_a_selection_scoped_artifact_under_the_planned_path( def test_refuses_a_selection_scoped_plan_the_container_does_not_share( self, cohort_stubs ): - """The abort the review worried about, under a real selection: a + """Abort when the container resolved a different SPM selection than the plan. + + This is the abort the review worried about, under a real selection: a container that resolved a different selection names a different - artifact, and nothing is published.""" + artifact, and nothing is published. + """ cohort_stubs.spm = {"scenario": "zero_real"} cohort_stubs.baseline = cohort_stubs.make_spm() entry = self._entry() @@ -925,8 +938,10 @@ def test_fails_loudly_when_ensure_leaves_no_artifact(self, cohort_stubs): class TestVerifyDeterminismImpl: - """The gate's comparator: the pipeline's only correctness check must - never fail open.""" + """The determinism gate's comparator must never fail open. + + It is the pipeline's only correctness check. + """ def _run( self, diff --git a/projects/policyengine-simulation-executor/tests/test_precompute_app.py b/projects/policyengine-simulation-executor/tests/test_precompute_app.py index 0ff348119..89a08973b 100644 --- a/projects/policyengine-simulation-executor/tests/test_precompute_app.py +++ b/projects/policyengine-simulation-executor/tests/test_precompute_app.py @@ -135,8 +135,10 @@ def _sample_plan(): def test_wrappers_validate_dicts_in_and_dump_models_out(precompute_module, monkeypatch): - """The wrappers ARE the Modal serialization boundary: dict in, - validated model to the impl, model back, dict out.""" + """The wrappers ARE the Modal serialization boundary. + + Dict in, validated model to the impl, model back, dict out. + """ from policyengine_simulation_executor import precompute as lib from policyengine_simulation_executor.precompute_models import ( BaselineComputeResult, diff --git a/projects/policyengine-simulation-executor/tests/test_release_bundle.py b/projects/policyengine-simulation-executor/tests/test_release_bundle.py index 9607ea254..12f7c5a06 100644 --- a/projects/policyengine-simulation-executor/tests/test_release_bundle.py +++ b/projects/policyengine-simulation-executor/tests/test_release_bundle.py @@ -40,7 +40,6 @@ def with_revision(dataset_uri, revision): @pytest.fixture def policyengine_uk_data_release(monkeypatch): """Expose the UK package-data release introduced by policyengine.py 5.0.4.""" - bundle = deepcopy(get_current_bundle()) bundle["bundle_version"] = "5.0.4" bundle["policyengine_version"] = "5.0.4" diff --git a/projects/policyengine-simulation-executor/tests/test_segmented_national.py b/projects/policyengine-simulation-executor/tests/test_segmented_national.py index 6155c8df8..df4536edf 100644 --- a/projects/policyengine-simulation-executor/tests/test_segmented_national.py +++ b/projects/policyengine-simulation-executor/tests/test_segmented_national.py @@ -17,7 +17,7 @@ def _bare_country_module(): - """A country module whose model has no loadable region registry.""" + """Return a country module whose model has no loadable region registry.""" return SimpleNamespace( model=SimpleNamespace(get_region=lambda code: None, region_registry=None) ) @@ -437,7 +437,7 @@ def _spm_child(index, *, selection=SPM_SELECTION, year="2026"): def _spm_child_missing(index, field): - """A child whose canonical receipt lost a resolved option.""" + """Return a child whose canonical receipt lost a resolved option.""" child = _spm_child(index) del child["spm_provenance"]["baseline"]["receipt"][field] return child @@ -445,8 +445,10 @@ def _spm_child_missing(index, field): @pytest.fixture def stub_reduce(monkeypatch): - """The microdata reduce itself is covered by - test_segmented_national_reduce; these tests pin the SPM merge.""" + """Stub out the microdata reduce so these tests pin only the SPM merge. + + The reduce itself is covered by test_segmented_national_reduce. + """ monkeypatch.setattr( sn, "build_national_output", lambda children, **kwargs: {"budget": {}} ) diff --git a/projects/policyengine-simulation-executor/tests/test_segmented_national_reduce.py b/projects/policyengine-simulation-executor/tests/test_segmented_national_reduce.py index 59d2e3fa6..2d0d57b4b 100644 --- a/projects/policyengine-simulation-executor/tests/test_segmented_national_reduce.py +++ b/projects/policyengine-simulation-executor/tests/test_segmented_national_reduce.py @@ -20,7 +20,7 @@ def _child(household_ids, incomes, **extra): - """A minimal child output dict with a _microdata payload.""" + """Return a minimal child output dict with a _microdata payload.""" frame = { "household_id": list(household_ids), "household_net_income": [float(x) for x in incomes], 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..124b90c30 100644 --- a/projects/policyengine-simulation-executor/tests/test_simulation_output_builder.py +++ b/projects/policyengine-simulation-executor/tests/test_simulation_output_builder.py @@ -633,8 +633,11 @@ def serialize(self): def test_run_simulation_impl_exports_baseline_artifact_outcome(monkeypatch): - """The hit/incomplete/miss attribute is the rollout metric for the - artifact pipeline; losing the export would blind that measurement.""" + """run_simulation_impl must export the baseline artifact outcome attribute. + + The hit/incomplete/miss attribute is the rollout metric for the artifact + pipeline; losing the export would blind that measurement. + """ dataset = object() country_module = SimpleNamespace(model=SimpleNamespace(version="1.715.2")) selection = installed_spm_selection() diff --git a/projects/policyengine-simulation-executor/tests/test_uk_weight_matrix_guard.py b/projects/policyengine-simulation-executor/tests/test_uk_weight_matrix_guard.py index dd1eb2f44..1c52bed2b 100644 --- a/projects/policyengine-simulation-executor/tests/test_uk_weight_matrix_guard.py +++ b/projects/policyengine-simulation-executor/tests/test_uk_weight_matrix_guard.py @@ -57,7 +57,6 @@ def _write_matrix(path: Path, households: int) -> None: def _assets(monkeypatch, tmp_path, *, households: int = 3, certified_dir=None): """Search UK geography assets in tmp_path and stub the certified download.""" - monkeypatch.setenv("POLICYENGINE_UK_GEOGRAPHY_DATA_DIR", str(tmp_path)) calls: list = [] From 526230481b2a61dc4e3c03fca19ca1c301486963 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 13 Oct 2026 16:56:47 +0400 Subject: [PATCH 08/18] Document the integration tests, repo scripts and package markers for ruff D Co-Authored-By: Claude Opus 5.5 --- .../scripts/update_policyengine_package.py | 5 +++++ projects/__init__.py | 1 + .../policyengine_api_full_integ/__init__.py | 1 + .../tests/simulation/test_budget_window.py | 6 +++--- .../tests/simulation/test_calculate.py | 19 +++++++++++-------- .../tests/simulation/test_ping.py | 6 +----- scripts/run_quality_guards.py | 4 ++++ 7 files changed, 26 insertions(+), 16 deletions(-) diff --git a/.github/scripts/update_policyengine_package.py b/.github/scripts/update_policyengine_package.py index 56197deb3..aab80691c 100644 --- a/.github/scripts/update_policyengine_package.py +++ b/.github/scripts/update_policyengine_package.py @@ -379,6 +379,11 @@ def _parser() -> argparse.ArgumentParser: def main() -> None: + """Run one subcommand for update-policyengine-package.sh. + + Results go to stdout for the shell script to capture; a failed check exits + non-zero through ``SystemExit`` with the reason as its message. + """ args = _parser().parse_args() if args.command == "current-version": print(current_version(args.pyproject, args.package)) diff --git a/projects/__init__.py b/projects/__init__.py index e69de29bb..28fda1efc 100644 --- a/projects/__init__.py +++ b/projects/__init__.py @@ -0,0 +1 @@ +"""Service projects of the simulation API; each subdirectory is its own uv project.""" diff --git a/projects/policyengine-apis-integ/policyengine_api_full_integ/__init__.py b/projects/policyengine-apis-integ/policyengine_api_full_integ/__init__.py index e69de29bb..2f52be592 100644 --- a/projects/policyengine-apis-integ/policyengine_api_full_integ/__init__.py +++ b/projects/policyengine-apis-integ/policyengine_api_full_integ/__init__.py @@ -0,0 +1 @@ +"""Empty package of the API integration-test project; the tests live in tests/.""" diff --git a/projects/policyengine-apis-integ/tests/simulation/test_budget_window.py b/projects/policyengine-apis-integ/tests/simulation/test_budget_window.py index b30b30083..19c4ae701 100644 --- a/projects/policyengine-apis-integ/tests/simulation/test_budget_window.py +++ b/projects/policyengine-apis-integ/tests/simulation/test_budget_window.py @@ -1,5 +1,4 @@ -""" -Integration tests for Modal-based budget-window batches. +"""Integration tests for Modal-based budget-window batches. These tests run against the staging Modal deployment and verify that the gateway can spawn the parent budget-window worker, the parent can spawn child @@ -26,7 +25,8 @@ def test_budget_window_multi_year_batch_completes( poll_budget_window_batch, us_model_version: str, ): - """ + """A two-year US budget-window batch returns both annual impacts and totals. + Given a two-year US budget-window request When the batch is submitted and polled to completion Then the response contains 2026 and 2027 annual impacts plus totals. diff --git a/projects/policyengine-apis-integ/tests/simulation/test_calculate.py b/projects/policyengine-apis-integ/tests/simulation/test_calculate.py index 89240f9ea..dbcdca26c 100644 --- a/projects/policyengine-apis-integ/tests/simulation/test_calculate.py +++ b/projects/policyengine-apis-integ/tests/simulation/test_calculate.py @@ -1,5 +1,4 @@ -""" -Integration tests for Modal-based simulation calculations. +"""Integration tests for Modal-based simulation calculations. These tests run against the staging Modal deployment and verify that economy-wide simulations complete successfully. @@ -71,8 +70,7 @@ def poll_for_completion( max_wait_seconds: float, poll_interval: float, ) -> JobStatusResponse: - """ - Poll for job completion. + """Poll for job completion. Args: client: The API client @@ -86,6 +84,7 @@ def poll_for_completion( Raises: TimeoutError: If job doesn't complete within max_wait_seconds AssertionError: If job fails + """ deadline = time.monotonic() + max_wait_seconds @@ -343,7 +342,8 @@ def test_calculate_default_model( max_wait_seconds: float, poll_interval: float, ): - """ + """A simulation on the default model version returns economic impact data. + Given a simulation request with default model version When the simulation is submitted and polled to completion Then the result contains expected economic impact data. @@ -387,7 +387,8 @@ def test_calculate_us_state_region_model( max_wait_seconds: float, poll_interval: float, ): - """ + """A US state-region simulation finds its dataset in the bundled manifest. + Given a US state-region simulation request When the simulation is submitted and polled to completion Then the worker can load the region dataset from its bundled manifest. @@ -435,7 +436,8 @@ def test_calculate_specific_model( max_wait_seconds: float, poll_interval: float, ): - """ + """A simulation pinned to a specific US model version returns impact data. + Given a simulation request with a specific model version When the simulation is submitted and polled to completion Then the result contains expected economic impact data. @@ -485,7 +487,8 @@ def test_calculate_uk_model( max_wait_seconds: float, poll_interval: float, ): - """ + """A UK simulation returns economic impact data. + Given a UK simulation request When the simulation is submitted and polled to completion Then the result contains expected economic impact data. diff --git a/projects/policyengine-apis-integ/tests/simulation/test_ping.py b/projects/policyengine-apis-integ/tests/simulation/test_ping.py index 71e9ae0b7..b68d74496 100644 --- a/projects/policyengine-apis-integ/tests/simulation/test_ping.py +++ b/projects/policyengine-apis-integ/tests/simulation/test_ping.py @@ -13,11 +13,7 @@ giveup=lambda e: getattr(e, "response", {}).get("status_code", 0) != 503, ) def test_ping(client: Client | AuthenticatedClient): - """ - Given a running simulation API - When a ping request is sent - Then the response contains the incremented value. - """ + """Given a running simulation API, when pinged, then it returns value + 1.""" response = ping_ping_post.sync(client=client, body=PingRequest(value=12)) assert isinstance(response, PingResponse) assert response.incremented == 13 diff --git a/scripts/run_quality_guards.py b/scripts/run_quality_guards.py index c3d8526ff..c6d7f4cc1 100644 --- a/scripts/run_quality_guards.py +++ b/scripts/run_quality_guards.py @@ -22,6 +22,10 @@ def main() -> int: + """Run every guard, print its violations, and return 1 if any guard failed. + + All guards run even after one fails, so a single run reports every problem. + """ failed = False for name, check in GUARDS: violations = check() From 67ccd40d098f316e77d05b1f74f23b2ed5a4430b Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 13 Oct 2026 16:56:47 +0400 Subject: [PATCH 09/18] Document the Stage 12 schema-DDL proof fixtures for ruff D The schema-authority proof reads the fixtures as text; the new docstrings avoid DDL phrases, so the bad fixture is still rejected and the good one accepted. Co-Authored-By: Claude Opus 5.5 --- rules/fixtures/sim.stage12.no-schema-ddl/bad_stage12_runtime.py | 1 + rules/fixtures/sim.stage12.no-schema-ddl/good_stage12_runtime.py | 1 + 2 files changed, 2 insertions(+) diff --git a/rules/fixtures/sim.stage12.no-schema-ddl/bad_stage12_runtime.py b/rules/fixtures/sim.stage12.no-schema-ddl/bad_stage12_runtime.py index 8e6ea09d6..1db8aeaff 100644 --- a/rules/fixtures/sim.stage12.no-schema-ddl/bad_stage12_runtime.py +++ b/rules/fixtures/sim.stage12.no-schema-ddl/bad_stage12_runtime.py @@ -11,4 +11,5 @@ def ensure_comparison_table(connection: Connection) -> None: + """Make the reports table from runtime code, the DDL the rule must catch.""" connection.execute(text("CREATE TABLE stage12_evaluation_reports (id uuid)")) diff --git a/rules/fixtures/sim.stage12.no-schema-ddl/good_stage12_runtime.py b/rules/fixtures/sim.stage12.no-schema-ddl/good_stage12_runtime.py index b8cf2de0c..bb03ec6fd 100644 --- a/rules/fixtures/sim.stage12.no-schema-ddl/good_stage12_runtime.py +++ b/rules/fixtures/sim.stage12.no-schema-ddl/good_stage12_runtime.py @@ -11,4 +11,5 @@ def record_report(connection: Connection, reports_table: Table, report_id: str) -> None: + """Insert one report row into a table that policyengine-api owns.""" connection.execute(insert(reports_table).values(id=report_id)) From d267fff930c4a4dc18d9bfbf1cae2abe78384db3 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 13 Oct 2026 16:56:47 +0400 Subject: [PATCH 10/18] Enable ruff group 4 (docstrings) Select pydocstyle (D) and ignore D203 and D213, the incompatible halves of the D203/D211 and D212/D213 pairs. Tests stay exempt from D1 only. Co-Authored-By: Claude Opus 5.5 --- ruff.toml | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/ruff.toml b/ruff.toml index 2aa803cb6..be05dff52 100644 --- a/ruff.toml +++ b/ruff.toml @@ -15,7 +15,12 @@ select = [ "TC", # Group 3: annotations. "ANN", + # Group 4: docstrings. + "D", ] +# Mutually exclusive pydocstyle pairs: keep D211 (no blank line before a class +# docstring) and D212 (summary on the first line). +ignore = ["D203", "D213"] [lint.per-file-ignores] # Tests may use pytest's bare assert, and need no docstrings, annotations or From 0826b3a6cef30364a83582cdb3e95a660e239cd1 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 13 Oct 2026 16:57:42 +0400 Subject: [PATCH 11/18] Drop redundant return None from documented no-op methods With a docstring as the body, the bare return None lines are redundant; remove them instead of suppressing RET501. Co-Authored-By: Claude Opus 5.5 --- .../src/policyengine_fastapi/observability/emitters.py | 8 -------- 1 file changed, 8 deletions(-) diff --git a/libs/policyengine-fastapi/src/policyengine_fastapi/observability/emitters.py b/libs/policyengine-fastapi/src/policyengine_fastapi/observability/emitters.py index ffb05c79b..5a65a8e94 100644 --- a/libs/policyengine-fastapi/src/policyengine_fastapi/observability/emitters.py +++ b/libs/policyengine-fastapi/src/policyengine_fastapi/observability/emitters.py @@ -29,15 +29,12 @@ def __exit__( traceback: TracebackType | None, ) -> None: """Do nothing, so any exception raised inside the span propagates.""" - return None # noqa: RET501 # explicit no-op body def set_attribute(self, key: str, value: object) -> None: """Discard the attribute.""" - return None # noqa: RET501 # explicit no-op body def add_event(self, name: str, attributes: Mapping[str, Any] | None = None) -> None: """Discard the event.""" - return None # noqa: RET501 # explicit no-op body class Observability(Protocol): @@ -88,7 +85,6 @@ class NoOpObservability: def emit_lifecycle_event(self, event: SimulationLifecycleEvent) -> None: """Discard the event.""" - return None # noqa: RET501 # explicit no-op body def emit_counter( self, @@ -97,7 +93,6 @@ def emit_counter( attributes: Mapping[str, str] | None = None, ) -> None: """Discard the measurement.""" - return None # noqa: RET501 # explicit no-op body def emit_histogram( self, @@ -106,11 +101,9 @@ def emit_histogram( attributes: Mapping[str, str] | None = None, ) -> None: """Discard the measurement.""" - return None # noqa: RET501 # explicit no-op body def record_artifact_manifest(self, manifest: TracerArtifactManifest) -> None: """Discard the manifest.""" - return None # noqa: RET501 # explicit no-op body def span(self, name: str, attributes: Mapping[str, Any] | None = None) -> NoOpSpan: """Return a span that records nothing.""" @@ -118,4 +111,3 @@ def span(self, name: str, attributes: Mapping[str, Any] | None = None) -> NoOpSp def flush(self) -> None: """Do nothing; there is never anything buffered.""" - return None # noqa: RET501 # explicit no-op body From 2a7aaf151aa75b709f7e853aa2210a0621d7dbcc Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 13 Oct 2026 17:11:36 +0400 Subject: [PATCH 12/18] Document published contract models and routes Replace the 61 D101 suppressions on contract, ping and health models with class docstrings written for API consumers, and fix the /ping/alive and gateway /ping route docstrings instead of suppressing D205. Pydantic and FastAPI copy these into schema and operation descriptions, so the entry, gateway and executor OpenAPI documents gain descriptions; paths, operations, properties, types and required lists are unchanged. The gateway golden spec is regenerated with generate_openapi. Co-Authored-By: Claude Opus 5.5 --- .../policyengine_fastapi/health/__init__.py | 12 +- .../src/policyengine_fastapi/ping/__init__.py | 19 ++- .../macro_output.py | 160 +++++++++++++++--- .../policyengine_simulation_contract/spm.py | 17 +- .../stage12_bundle.py | 19 ++- .../stage12_execution.py | 124 +++++++++++--- .../stage12_manifest.py | 32 +++- .../endpoints.py | 7 +- .../tests/golden/openapi.json | 75 +++++--- 9 files changed, 371 insertions(+), 94 deletions(-) diff --git a/libs/policyengine-fastapi/src/policyengine_fastapi/health/__init__.py b/libs/policyengine-fastapi/src/policyengine_fastapi/health/__init__.py index 4db20718f..ecf94141d 100644 --- a/libs/policyengine-fastapi/src/policyengine_fastapi/health/__init__.py +++ b/libs/policyengine-fastapi/src/policyengine_fastapi/health/__init__.py @@ -5,19 +5,25 @@ from pydantic import BaseModel -class ProbeStatus(BaseModel): # noqa: D101 # docstring would change the published OpenAPI description +class ProbeStatus(BaseModel): + """Result of one health probe, with an optional explanatory message.""" + name: str healthy: bool message: str | None = None -class SystemStatus(BaseModel): # noqa: D101 # docstring would change the published OpenAPI description +class SystemStatus(BaseModel): + """Health of one system; it is healthy only if every probe in ``detail`` is.""" + name: str healthy: bool detail: list[ProbeStatus] -class HealthStatus(BaseModel): # noqa: D101 # docstring would change the published OpenAPI description +class HealthStatus(BaseModel): + """Health of the service; it is healthy only if every system is.""" + healthy: bool systems: list[SystemStatus] diff --git a/libs/policyengine-fastapi/src/policyengine_fastapi/ping/__init__.py b/libs/policyengine-fastapi/src/policyengine_fastapi/ping/__init__.py index 9cc0c622d..867bf0c3c 100644 --- a/libs/policyengine-fastapi/src/policyengine_fastapi/ping/__init__.py +++ b/libs/policyengine-fastapi/src/policyengine_fastapi/ping/__init__.py @@ -7,11 +7,15 @@ from policyengine_fastapi.health import HealthRegistry, HealthStatus -class PingRequest(BaseModel): # noqa: D101 # docstring would change the published OpenAPI description +class PingRequest(BaseModel): + """Ping payload: an integer for the service to increment.""" + value: int -class PingResponse(BaseModel): # noqa: D101 # docstring would change the published OpenAPI description +class PingResponse(BaseModel): + """Ping reply carrying the request's ``value`` plus one.""" + incremented: int @@ -21,7 +25,7 @@ def create_router(health_registry: HealthRegistry) -> APIRouter: @router.post("/ping") async def ping(request: PingRequest) -> PingResponse: - """Verify the api is able to recieve and process unauthenticated requests.""" + """Verify the API can receive and process unauthenticated requests.""" return PingResponse(incremented=request.value + 1) @router.get("/ping/started", response_model=str) @@ -31,10 +35,11 @@ async def started() -> str: @router.get("/ping/alive", response_model=HealthStatus) async def alive() -> JSONResponse: - """Check if the service is healthy. This will always return - data, but will return a 503 response code if the service is - unhealthy. - """ # noqa: D205 # docstring would change the published OpenAPI description + """Report whether the service is healthy. + + The health report is always returned, with status 200 when the service + is healthy and 503 when it is not. + """ content = health_registry.report() return JSONResponse( content=content.model_dump(exclude_none=True), diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/macro_output.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/macro_output.py index 32fd1f886..a45cddec4 100644 --- a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/macro_output.py +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/macro_output.py @@ -21,7 +21,16 @@ class MacroRootModel(RootModel[T], Generic[T]): """Base model for named map and list output schemas.""" -class BudgetaryImpact(MacroOutputModel): # noqa: D101 # docstring would change the published OpenAPI description +class BudgetaryImpact(MacroOutputModel): + """Headline budget effect of the reform for the simulated year. + + Impacts are reform minus baseline totals in the country's currency. + ``budgetary_impact`` is ``tax_revenue_impact`` minus + ``benefit_spending_impact``; ``state_tax_revenue_impact`` covers US state + income tax and is 0 for the UK. ``households`` is the weighted number of + households and ``baseline_net_income`` their total baseline net income. + """ + tax_revenue_impact: float state_tax_revenue_impact: float benefit_spending_impact: float @@ -33,7 +42,12 @@ class BudgetaryImpact(MacroOutputModel): # noqa: D101 # docstring would change BudgetaryOutput = BudgetaryImpact -class DetailedBudgetProgramOutput(MacroOutputModel): # noqa: D101 # docstring would change the published OpenAPI description +class DetailedBudgetProgramOutput(MacroOutputModel): + """Baseline and reform totals for one program. + + ``difference`` is reform minus baseline. + """ + baseline: float reform: float difference: float @@ -43,82 +57,142 @@ class DetailedBudgetOutput(MacroRootModel[dict[str, DetailedBudgetProgramOutput] """Program name mapped to its baseline and reform totals.""" -class DecileOutput(MacroOutputModel): # noqa: D101 # docstring would change the published OpenAPI description +class DecileOutput(MacroOutputModel): + """Change in average household net income in each decile. + + Keys are decile numbers from ``"1"`` (lowest) to ``"10"`` (highest); the same + shape is used for income deciles and UK wealth deciles. ``average`` is the + absolute change and ``relative`` the change relative to the decile's + baseline average. + """ + average: dict[str, float] relative: dict[str, float] -class IntraDecileOutput(MacroOutputModel): # noqa: D101 # docstring would change the published OpenAPI description +class IntraDecileOutput(MacroOutputModel): + """Share of people in each decile who gain or lose, grouped by size of change. + + The categories are gaining or losing more or less than 5% of household net + income, or no change. ``deciles`` lists each category's share for deciles 1 + to 10 in order, and ``all`` gives each category's average across the deciles. + """ + deciles: dict[str, list[float]] all: dict[str, float] -class BaselineReformValue(MacroOutputModel): # noqa: D101 # docstring would change the published OpenAPI description +class BaselineReformValue(MacroOutputModel): + """A metric's value under the baseline and under the reform.""" + baseline: float reform: float -class AgePovertyOutput(MacroOutputModel): # noqa: D101 # docstring would change the published OpenAPI description +class AgePovertyOutput(MacroOutputModel): + """Baseline and reform poverty rates for children, adults, seniors and everyone.""" + child: BaselineReformValue adult: BaselineReformValue senior: BaselineReformValue all: BaselineReformValue -class GenderPovertyOutput(MacroOutputModel): # noqa: D101 # docstring would change the published OpenAPI description +class GenderPovertyOutput(MacroOutputModel): + """Baseline and reform poverty rates for males and females.""" + male: BaselineReformValue female: BaselineReformValue -class RacePovertyOutput(MacroOutputModel): # noqa: D101 # docstring would change the published OpenAPI description +class RacePovertyOutput(MacroOutputModel): + """Baseline and reform poverty rates by race and ethnicity (US only).""" + white: BaselineReformValue black: BaselineReformValue hispanic: BaselineReformValue other: BaselineReformValue -class PovertyOutput(MacroOutputModel): # noqa: D101 # docstring would change the published OpenAPI description +class PovertyOutput(MacroOutputModel): + """Poverty and deep poverty rates by age group. + + In the US these are the SPM poverty and deep SPM poverty rates. In the UK, + ``poverty`` is relative and ``deep_poverty`` absolute poverty, both measured + before housing costs. + """ + poverty: AgePovertyOutput deep_poverty: AgePovertyOutput -class PovertyByGenderOutput(MacroOutputModel): # noqa: D101 # docstring would change the published OpenAPI description +class PovertyByGenderOutput(MacroOutputModel): + """Poverty and deep poverty rates by gender, as defined for ``PovertyOutput``.""" + poverty: GenderPovertyOutput deep_poverty: GenderPovertyOutput -class PovertyByRaceOutput(MacroOutputModel): # noqa: D101 # docstring would change the published OpenAPI description +class PovertyByRaceOutput(MacroOutputModel): + """US poverty rates by race and ethnicity; there is no deep poverty breakdown.""" + poverty: RacePovertyOutput -class PovertyModuleOutputs(MacroOutputModel): # noqa: D101 # docstring would change the published JSON schema description +class PovertyModuleOutputs(MacroOutputModel): + """Overall, by-gender and by-race poverty results for one comparison. + + ``poverty_by_race`` is null outside the US. + """ + poverty: PovertyOutput poverty_by_gender: PovertyByGenderOutput poverty_by_race: PovertyByRaceOutput | None -class InequalityOutput(MacroOutputModel): # noqa: D101 # docstring would change the published OpenAPI description +class InequalityOutput(MacroOutputModel): + """Baseline and reform income inequality. + + Covers the Gini coefficient and the income shares of the top 10% and top 1%. + """ + gini: BaselineReformValue top_10_pct_share: BaselineReformValue top_1_pct_share: BaselineReformValue -class LaborSupplyRelativeResponse(MacroOutputModel): # noqa: D101 # docstring would change the published OpenAPI description +class LaborSupplyRelativeResponse(MacroOutputModel): + """Income- and substitution-effect responses relative to baseline earnings.""" + income: float substitution: float -class LaborSupplyDecileMetric(MacroOutputModel): # noqa: D101 # docstring would change the published OpenAPI description +class LaborSupplyDecileMetric(MacroOutputModel): + """Income- and substitution-effect responses keyed by income decile number.""" + income: dict[int, float] substitution: dict[int, float] -class LaborSupplyDecileOutput(MacroOutputModel): # noqa: D101 # docstring would change the published OpenAPI description +class LaborSupplyDecileOutput(MacroOutputModel): + """Labor supply responses by income decile. + + ``average`` is the response per household and ``relative`` the response + relative to the decile's baseline earnings. + """ + average: LaborSupplyDecileMetric relative: LaborSupplyDecileMetric -class LaborSupplyHoursOutput(MacroOutputModel): # noqa: D101 # docstring would change the published OpenAPI description +class LaborSupplyHoursOutput(MacroOutputModel): + """Total weekly hours worked under the baseline and reform, and how they change. + + ``change`` is reform minus baseline, split into its income and substitution + effects. + """ + baseline: float reform: float change: float @@ -126,7 +200,13 @@ class LaborSupplyHoursOutput(MacroOutputModel): # noqa: D101 # docstring would substitution_effect: float -class LaborSupplyResponseOutput(MacroOutputModel): # noqa: D101 # docstring would change the published OpenAPI description +class LaborSupplyResponseOutput(MacroOutputModel): + """Behavioral labor supply response to the reform. + + ``income_lsr`` and ``substitution_lsr`` are the total earnings changes from + the income and substitution effects, and ``total_change`` is their sum. + """ + substitution_lsr: float income_lsr: float relative_lsr: LaborSupplyRelativeResponse @@ -136,17 +216,32 @@ class LaborSupplyResponseOutput(MacroOutputModel): # noqa: D101 # docstring wo hours: LaborSupplyHoursOutput -class CliffImpactInSimulation(MacroOutputModel): # noqa: D101 # docstring would change the published OpenAPI description +class CliffImpactInSimulation(MacroOutputModel): + """Exposure to tax-benefit cliffs in one simulation. + + ``cliff_gap`` is the total net income people would lose from a small rise in + their earnings, and ``cliff_share`` the share of adults on a cliff. + """ + cliff_gap: float cliff_share: float -class CliffImpactOutput(MacroOutputModel): # noqa: D101 # docstring would change the published OpenAPI description +class CliffImpactOutput(MacroOutputModel): + """Exposure to tax-benefit cliffs under the baseline and under the reform.""" + baseline: CliffImpactInSimulation reform: CliffImpactInSimulation -class ConstituencyImpactRecord(MacroOutputModel): # noqa: D101 # docstring would change the published OpenAPI description +class ConstituencyImpactRecord(MacroOutputModel): + """Reform impact on household net income in one UK parliamentary constituency. + + ``average_household_income_change`` is the average change per household and + ``relative_household_income_change`` the relative change in the area's total + household net income. ``x`` and ``y`` are integer map-rendering coordinates. + """ + constituency_code: str constituency_name: str x: int | None @@ -156,7 +251,12 @@ class ConstituencyImpactRecord(MacroOutputModel): # noqa: D101 # docstring wou population: float -class LocalAuthorityImpactRecord(MacroOutputModel): # noqa: D101 # docstring would change the published OpenAPI description +class LocalAuthorityImpactRecord(MacroOutputModel): + """Reform impact on household net income in one UK local authority. + + Fields mean the same as in ``ConstituencyImpactRecord``. + """ + local_authority_code: str local_authority_name: str x: int | None @@ -172,7 +272,17 @@ class GeographicImpactOutput( """UK constituency or local-authority impact records.""" -class CongressionalDistrictImpactRecord(MacroOutputModel): # noqa: D101 # docstring would change the published OpenAPI description +class CongressionalDistrictImpactRecord(MacroOutputModel): + """Reform impact on household net income in one US congressional district. + + ``district`` is the state abbreviation and two-digit district number, such + as ``"CA-12"``; at-large districts are numbered ``01``. The income changes + are the average change per household and the relative change in the + district's total household net income. The three percentage fields are the + shares of people whose household net income rises, falls, or stays + essentially unchanged. + """ + district: str average_household_income_change: float relative_household_income_change: float @@ -182,7 +292,9 @@ class CongressionalDistrictImpactRecord(MacroOutputModel): # noqa: D101 # docs population: float -class CongressionalDistrictImpactOutput(MacroOutputModel): # noqa: D101 # docstring would change the published OpenAPI description +class CongressionalDistrictImpactOutput(MacroOutputModel): + """Reform impact for each US congressional district in the simulated region.""" + districts: list[CongressionalDistrictImpactRecord] diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/spm.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/spm.py index 7cdd7b4b8..10e41246e 100644 --- a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/spm.py +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/spm.py @@ -323,7 +323,14 @@ def validate_spm_calculation_provenance( ) -class SPMErrorDetail(BaseModel): # noqa: D101 # docstring would change the published OpenAPI description +class SPMErrorDetail(BaseModel): + """A typed SPM input error. + + ``code`` is one of a fixed set of stable values, such as + ``SPM_GEOGRAPHY_REQUIRED``, that clients can branch on; ``message`` explains + the problem to a person. + """ + code: str message: str @@ -359,7 +366,13 @@ def spm_error_detail(exc: BaseException) -> SPMErrorDetail | None: return None -class SPMCapability(BaseModel): # noqa: D101 # docstring would change the published OpenAPI description +class SPMCapability(BaseModel): + """Canonical SPM support offered by a simulation version. + + ``defaults`` is the pinned SPM selection, including the forecast artifact + hash and scenario, that applies to any option a request leaves out. + """ + model_config = ConfigDict(extra="forbid") contract_version: Literal["canonical-spm-v1"] = SPM_CONTRACT_VERSION defaults: SPMSelection diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_bundle.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_bundle.py index 2b0b929ce..5c406f742 100644 --- a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_bundle.py +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_bundle.py @@ -12,11 +12,18 @@ SUPPORTED_COUNTRIES: tuple[CountryId, ...] = ("us", "uk") -class StrictBundleModel(BaseModel): # noqa: D101 # docstring would change the published JSON schema description +class StrictBundleModel(BaseModel): + """Base for normalized bundle models; unknown fields are rejected.""" + model_config = ConfigDict(extra="forbid", frozen=True) -class Stage12Dataset(StrictBundleModel): # noqa: D101 # docstring would change the published JSON schema description +class Stage12Dataset(StrictBundleModel): + """A dataset available in a country bundle. + + Gives its identity, URI, artifact revision, SHA-256 and repository type. + """ + identity: NonEmptyText uri: NonEmptyText artifact_revision: NonEmptyText @@ -24,7 +31,9 @@ class Stage12Dataset(StrictBundleModel): # noqa: D101 # docstring would change repo_type: NonEmptyText -class Stage12CountryBundle(StrictBundleModel): # noqa: D101 # docstring would change the published JSON schema description +class Stage12CountryBundle(StrictBundleModel): + """One country's pinned model package, data release and datasets within a bundle.""" + country: CountryId country_package_name: NonEmptyText country_package_version: NonEmptyText @@ -51,6 +60,8 @@ class Stage12BundleManifest(StrictBundleModel): countries: tuple[Stage12CountryBundle, ...] -class ResolvedStage12Bundle(StrictBundleModel): # noqa: D101 # docstring would change the published JSON schema description +class ResolvedStage12Bundle(StrictBundleModel): + """A normalized bundle manifest together with its SHA-256 digest.""" + bundle: Stage12BundleManifest bundle_manifest_sha256: Sha256Digest diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_execution.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_execution.py index aff9f6b68..a125451d7 100644 --- a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_execution.py +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_execution.py @@ -53,17 +53,23 @@ ] -class StrictContractModel(BaseModel): # noqa: D101 # docstring would change the published JSON schema description +class StrictContractModel(BaseModel): + """Base for Stage 12 contracts; rejects unknown fields and non-finite numbers.""" + model_config = ConfigDict(extra="forbid", allow_inf_nan=False, frozen=True) -class SimulationRole(StrEnum): # noqa: D101 # docstring would change the published JSON schema description +class SimulationRole(StrEnum): + """Role of a simulation: baseline or reform of a comparison, or a standalone run.""" + BASELINE = "baseline" REFORM = "reform" STANDALONE = "standalone" -class ArtifactMediaType(StrEnum): # noqa: D101 # docstring would change the published JSON schema description +class ArtifactMediaType(StrEnum): + """Media type of a Stage 12 output artifact.""" + JSON = "application/json" PARQUET = "application/vnd.apache.parquet" @@ -104,13 +110,17 @@ class SimulationParquetPayloadContract(StrictContractModel): SIMULATION_PARQUET_PAYLOAD_CONTRACT = SimulationParquetPayloadContract() -class DatasetArtifactMediaType(StrEnum): # noqa: D101 # docstring would change the published JSON schema description +class DatasetArtifactMediaType(StrEnum): + """Media type of an input dataset artifact.""" + HDF5 = "application/x-hdf5" JSON = "application/json" PARQUET = "application/vnd.apache.parquet" -class DatasetProvenance(StrictContractModel): # noqa: D101 # docstring would change the published JSON schema description +class DatasetProvenance(StrictContractModel): + """The dataset a simulation used: identity, URI, revision and data package.""" + identity: ContractText uri: Annotated[str, Field(min_length=1, max_length=2048)] artifact_revision: ContractText @@ -118,7 +128,13 @@ class DatasetProvenance(StrictContractModel): # noqa: D101 # docstring would c data_package_version: ContractText -class BundleProvenance(StrictContractModel): # noqa: D101 # docstring would change the published JSON schema description +class BundleProvenance(StrictContractModel): + """The pinned software and data a simulation ran with. + + Covers the PolicyEngine.py version, the country model package, the dataset, + and the SHA-256 of the bundle manifest they came from. + """ + policyengine_version: ContractText country_package_name: ContractText country_package_version: ContractText @@ -126,26 +142,37 @@ class BundleProvenance(StrictContractModel): # noqa: D101 # docstring would ch bundle_manifest_sha256: Sha256Digest -class ArtifactReference(StrictContractModel): # noqa: D101 # docstring would change the published JSON schema description +class ArtifactReference(StrictContractModel): + """An artifact in Cloud Storage, with its media type, SHA-256 and size in bytes.""" + uri: StorageUri media_type: ArtifactMediaType content_sha256: Sha256Digest size_bytes: Annotated[int, Field(ge=0)] -class DatasetArtifactReference(StrictContractModel): # noqa: D101 # docstring would change the published JSON schema description +class DatasetArtifactReference(StrictContractModel): + """A dataset artifact in Cloud Storage (``gs://``) or on Hugging Face (``hf://``). + + ``size_bytes`` may be omitted when the size is not known. + """ + uri: DatasetArtifactUri media_type: DatasetArtifactMediaType content_sha256: Sha256Digest size_bytes: Annotated[int, Field(ge=0)] | None = None -class DatasetPopulationInput(StrictContractModel): # noqa: D101 # docstring would change the published JSON schema description +class DatasetPopulationInput(StrictContractModel): + """Population input that simulates the households in a dataset artifact.""" + kind: Literal["dataset"] = "dataset" artifact: DatasetArtifactReference -class HouseholdPopulationInput(StrictContractModel): # noqa: D101 # docstring would change the published JSON schema description +class HouseholdPopulationInput(StrictContractModel): + """Population input that simulates one household from a situation document.""" + kind: Literal["household"] = "household" document: dict[str, JsonValue] @@ -156,7 +183,13 @@ class HouseholdPopulationInput(StrictContractModel): # noqa: D101 # docstring ] -class GeographySelection(StrictContractModel): # noqa: D101 # docstring would change the published JSON schema description +class GeographySelection(StrictContractModel): + """Country and region to simulate, optionally filtered on one dataset field. + + ``filter_field`` and ``filter_value`` must be given together, and + ``filter_strategy`` requires both. + """ + country: CountryId region: ContractText filter_field: ContractText | None = None @@ -173,7 +206,15 @@ def require_complete_filter(self) -> GeographySelection: return self -class SimulationExecutionInput(StrictContractModel): # noqa: D101 # docstring would change the published JSON schema description +class SimulationExecutionInput(StrictContractModel): + """Everything needed to run one simulation of a comparison run. + + ``evaluation_id`` identifies the comparison run and + ``simulation_execution_id`` this simulation within it. The input fixes the + policy, population, year, geography, options and pinned bundle; a baseline + and reform pair is described by ``ReportExecutionInput`` instead. + """ + contract_version: Literal[1] = 1 evaluation_id: UUID simulation_execution_id: UUID @@ -200,7 +241,9 @@ def reject_combined_policy_input(cls, value: object) -> object: return value -class ReportAggregate(StrEnum): # noqa: D101 # docstring would change the published JSON schema description +class ReportAggregate(StrEnum): + """A family of aggregate results that a comparison report computes.""" + BUDGET = "budget" POVERTY = "poverty" INEQUALITY = "inequality" @@ -328,7 +371,13 @@ def require_output_plan_country(self) -> PlannedSimulationExecutionInput: return self -class RowIdentity(StrictContractModel): # noqa: D101 # docstring would change the published JSON schema description +class RowIdentity(StrictContractModel): + """Fingerprint of the rows in a simulation artifact. + + Names the identifier columns and gives the row count and a SHA-256 over the + row identifiers, so artifacts can be checked to cover the same rows. + """ + schema_version: Literal[1] = 1 identifier_columns: Annotated[tuple[ContractText, ...], Field(min_length=1)] row_count: Annotated[int, Field(ge=0)] @@ -343,7 +392,14 @@ def require_unique_columns(cls, value: tuple[str, ...]) -> tuple[str, ...]: return value -class SimulationArtifactDescriptor(StrictContractModel): # noqa: D101 # docstring would change the published JSON schema description +class SimulationArtifactDescriptor(StrictContractModel): + """The stored output of one simulation in a comparison run. + + Points to the artifact and records the digest of the output plan it follows, + its row identity, its bundle provenance and, for SPM calculations, the + calculation receipt. + """ + contract_version: Literal[1] = 1 evaluation_id: UUID simulation_execution_id: UUID @@ -356,7 +412,13 @@ class SimulationArtifactDescriptor(StrictContractModel): # noqa: D101 # docstr calculation_provenance: SPMProvenance | None = None -class ReportExecutionInput(StrictContractModel): # noqa: D101 # docstring would change the published JSON schema description +class ReportExecutionInput(StrictContractModel): + """A baseline and reform simulation pair to run and aggregate into one report. + + Both simulations must belong to this report, carry the baseline and reform + roles and distinct execution IDs, and differ only in their policy. + """ + contract_version: Literal[1] = 1 evaluation_id: UUID baseline: SimulationExecutionInput @@ -396,7 +458,13 @@ def require_aligned_simulations(self) -> ReportExecutionInput: return self -class AggregateReportArtifactDescriptor(StrictContractModel): # noqa: D101 # docstring would change the published JSON schema description +class AggregateReportArtifactDescriptor(StrictContractModel): + """The stored aggregate results of a comparison run. + + Records the digests of the baseline and reform artifacts the results were + computed from, and the bundle provenance. + """ + contract_version: Literal[1] = 1 evaluation_id: UUID artifact: ArtifactReference @@ -501,7 +569,15 @@ def validate_differences(self) -> ResultComparisonArtifactPayload: return self -class Stage12InvocationContext(StrictContractModel): # noqa: D101 # docstring would change the published JSON schema description +class Stage12InvocationContext(StrictContractModel): + """Deployment and retention context for one comparison-run invocation. + + Records the originating request, the environments, the worker version and + callable, the manifest digests and the artifact prefix. Automatic runs also + name the production function call they are compared with. Artifacts are kept + for at most 30 days after ``created_at``. + """ + contract_version: Literal[1] = 1 request_id: ContractText # The logical environment labels durable records and artifact paths. @@ -544,7 +620,9 @@ def _validate_retention(created_at: datetime, retention_expires_at: datetime) -> raise ValueError("comparison retention must not exceed 30 days") -class ComparisonRunLifecycleStatus(StrEnum): # noqa: D101 # docstring would change the published JSON schema description +class ComparisonRunLifecycleStatus(StrEnum): + """Overall status of a comparison run.""" + PENDING = "pending" RUNNING = "running" SUCCEEDED = "succeeded" @@ -553,14 +631,18 @@ class ComparisonRunLifecycleStatus(StrEnum): # noqa: D101 # docstring would ch SKIPPED = "skipped" -class ComparisonRunAggregationStatus(StrEnum): # noqa: D101 # docstring would change the published JSON schema description +class ComparisonRunAggregationStatus(StrEnum): + """Status of aggregating a comparison run's simulations into a report.""" + NOT_STARTED = "not_started" RUNNING = "running" SUCCEEDED = "succeeded" FAILED = "failed" -class ResultComparisonStatus(StrEnum): # noqa: D101 # docstring would change the published JSON schema description +class ResultComparisonStatus(StrEnum): + """Outcome of checking a comparison run's aggregate result against production.""" + NOT_REQUESTED = "not_requested" PENDING = "pending" RUNNING = "running" diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_manifest.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_manifest.py index 996e7db92..9301ec51e 100644 --- a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_manifest.py +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_manifest.py @@ -30,16 +30,26 @@ CallableName = Annotated[str, Field(pattern=_CALLABLE_PATTERN.pattern)] -class StrictManifestModel(BaseModel): # noqa: D101 # docstring would change the published JSON schema description +class StrictManifestModel(BaseModel): + """Base for v2 version manifest models; unknown fields are rejected.""" + model_config = ConfigDict(extra="forbid", frozen=True) -class V2CountryWorker(StrictManifestModel): # noqa: D101 # docstring would change the published JSON schema description +class V2CountryWorker(StrictManifestModel): + """A country served by a v2 worker, and the callable that runs its simulations.""" + country: CountryId single_simulation_callable: CallableName -class V2WorkerValidation(StrictManifestModel): # noqa: D101 # docstring would change the published JSON schema description +class V2WorkerValidation(StrictManifestModel): + """Receipt for a successful validation of a deployed v2 worker application. + + Names the application and bundle digest that were validated, when, and the + validation invocation IDs overall and for each country. + """ + validated: Literal[True] application_name: ManifestText bundle_manifest_sha256: Sha256Digest @@ -48,7 +58,13 @@ class V2WorkerValidation(StrictManifestModel): # noqa: D101 # docstring would country_validation_invocation_ids: dict[CountryId, ManifestText] -class V2WorkerVersion(StrictManifestModel): # noqa: D101 # docstring would change the published JSON schema description +class V2WorkerVersion(StrictManifestModel): + """One deployed v2 worker application and the bundle it runs. + + Lists its report coordinator callable, its per-country workers and its + validation receipt; the countries must match the bundle's, in order. + """ + application_name: ManifestText report_coordinator_callable: CallableName countries: tuple[V2CountryWorker, ...] @@ -87,7 +103,13 @@ def validate_worker(self) -> V2WorkerVersion: return self -class V2VersionManifest(StrictManifestModel): # noqa: D101 # docstring would change the published JSON schema description +class V2VersionManifest(StrictManifestModel): + """Registry of deployed v2 worker versions, keyed by PolicyEngine.py version. + + ``default_version`` is the version used when a request does not pin one, and + ``generation`` identifies the publish that produced this manifest. + """ + schema_version: Literal[1] = 1 generation: ManifestText default_version: ManifestText 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 5d6a62fa6..3fffdf1bf 100644 --- a/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/endpoints.py +++ b/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/endpoints.py @@ -1265,7 +1265,8 @@ async def health() -> HealthResponse: @router.post("/ping", response_model=PingResponse) async def ping(request: PingRequest) -> PingResponse: - """Verify the API is able to receive and process requests. - Matches the policyengine_fastapi.ping endpoint for test compatibility. - """ # noqa: D205 # docstring would change the published OpenAPI description + """Verify the API can receive and process requests. + + Responds with the request's ``value`` plus one as ``incremented``. + """ return PingResponse(incremented=request.value + 1) diff --git a/projects/policyengine-simulation-gateway/tests/golden/openapi.json b/projects/policyengine-simulation-gateway/tests/golden/openapi.json index bffd7ec05..8ce5d92cf 100644 --- a/projects/policyengine-simulation-gateway/tests/golden/openapi.json +++ b/projects/policyengine-simulation-gateway/tests/golden/openapi.json @@ -331,7 +331,8 @@ "senior", "all" ], - "title": "AgePovertyOutput" + "title": "AgePovertyOutput", + "description": "Baseline and reform poverty rates for children, adults, seniors and everyone." }, "BaselineReformValue": { "properties": { @@ -350,7 +351,8 @@ "baseline", "reform" ], - "title": "BaselineReformValue" + "title": "BaselineReformValue", + "description": "A metric's value under the baseline and under the reform." }, "BatchChildJobStatus": { "properties": { @@ -925,7 +927,8 @@ "households", "baseline_net_income" ], - "title": "BudgetaryImpact" + "title": "BudgetaryImpact", + "description": "Headline budget effect of the reform for the simulated year.\n\nImpacts are reform minus baseline totals in the country's currency.\n``budgetary_impact`` is ``tax_revenue_impact`` minus\n``benefit_spending_impact``; ``state_tax_revenue_impact`` covers US state\nincome tax and is 0 for the UK. ``households`` is the weighted number of\nhouseholds and ``baseline_net_income`` their total baseline net income." }, "CliffImpactInSimulation": { "properties": { @@ -944,7 +947,8 @@ "cliff_gap", "cliff_share" ], - "title": "CliffImpactInSimulation" + "title": "CliffImpactInSimulation", + "description": "Exposure to tax-benefit cliffs in one simulation.\n\n``cliff_gap`` is the total net income people would lose from a small rise in\ntheir earnings, and ``cliff_share`` the share of adults on a cliff." }, "CliffImpactOutput": { "properties": { @@ -961,7 +965,8 @@ "baseline", "reform" ], - "title": "CliffImpactOutput" + "title": "CliffImpactOutput", + "description": "Exposure to tax-benefit cliffs under the baseline and under the reform." }, "CongressionalDistrictImpactOutput": { "properties": { @@ -978,7 +983,8 @@ "required": [ "districts" ], - "title": "CongressionalDistrictImpactOutput" + "title": "CongressionalDistrictImpactOutput", + "description": "Reform impact for each US congressional district in the simulated region." }, "CongressionalDistrictImpactRecord": { "properties": { @@ -1022,7 +1028,8 @@ "no_change_percentage", "population" ], - "title": "CongressionalDistrictImpactRecord" + "title": "CongressionalDistrictImpactRecord", + "description": "Reform impact on household net income in one US congressional district.\n\n``district`` is the state abbreviation and two-digit district number, such\nas ``\"CA-12\"``; at-large districts are numbered ``01``. The income changes\nare the average change per household and the relative change in the\ndistrict's total household net income. The three percentage fields are the\nshares of people whose household net income rises, falls, or stays\nessentially unchanged." }, "ConstituencyImpactRecord": { "properties": { @@ -1080,7 +1087,8 @@ "relative_household_income_change", "population" ], - "title": "ConstituencyImpactRecord" + "title": "ConstituencyImpactRecord", + "description": "Reform impact on household net income in one UK parliamentary constituency.\n\n``average_household_income_change`` is the average change per household and\n``relative_household_income_change`` the relative change in the area's total\nhousehold net income. ``x`` and ``y`` are integer map-rendering coordinates." }, "DecileOutput": { "properties": { @@ -1105,7 +1113,8 @@ "average", "relative" ], - "title": "DecileOutput" + "title": "DecileOutput", + "description": "Change in average household net income in each decile.\n\nKeys are decile numbers from ``\"1\"`` (lowest) to ``\"10\"`` (highest); the same\nshape is used for income deciles and UK wealth deciles. ``average`` is the\nabsolute change and ``relative`` the change relative to the decile's\nbaseline average." }, "DetailedBudgetOutput": { "additionalProperties": { @@ -1137,7 +1146,8 @@ "reform", "difference" ], - "title": "DetailedBudgetProgramOutput" + "title": "DetailedBudgetProgramOutput", + "description": "Baseline and reform totals for one program.\n\n``difference`` is reform minus baseline." }, "GenderPovertyOutput": { "properties": { @@ -1154,7 +1164,8 @@ "male", "female" ], - "title": "GenderPovertyOutput" + "title": "GenderPovertyOutput", + "description": "Baseline and reform poverty rates for males and females." }, "GeographicImpactOutput": { "items": { @@ -1216,7 +1227,8 @@ "top_10_pct_share", "top_1_pct_share" ], - "title": "InequalityOutput" + "title": "InequalityOutput", + "description": "Baseline and reform income inequality.\n\nCovers the Gini coefficient and the income shares of the top 10% and top 1%." }, "IntraDecileOutput": { "properties": { @@ -1244,7 +1256,8 @@ "deciles", "all" ], - "title": "IntraDecileOutput" + "title": "IntraDecileOutput", + "description": "Share of people in each decile who gain or lose, grouped by size of change.\n\nThe categories are gaining or losing more or less than 5% of household net\nincome, or no change. ``deciles`` lists each category's share for deciles 1\nto 10 in order, and ``all`` gives each category's average across the deciles." }, "JobStatusResponse": { "properties": { @@ -1395,7 +1408,8 @@ "income", "substitution" ], - "title": "LaborSupplyDecileMetric" + "title": "LaborSupplyDecileMetric", + "description": "Income- and substitution-effect responses keyed by income decile number." }, "LaborSupplyDecileOutput": { "properties": { @@ -1412,7 +1426,8 @@ "average", "relative" ], - "title": "LaborSupplyDecileOutput" + "title": "LaborSupplyDecileOutput", + "description": "Labor supply responses by income decile.\n\n``average`` is the response per household and ``relative`` the response\nrelative to the decile's baseline earnings." }, "LaborSupplyHoursOutput": { "properties": { @@ -1446,7 +1461,8 @@ "income_effect", "substitution_effect" ], - "title": "LaborSupplyHoursOutput" + "title": "LaborSupplyHoursOutput", + "description": "Total weekly hours worked under the baseline and reform, and how they change.\n\n``change`` is reform minus baseline, split into its income and substitution\neffects." }, "LaborSupplyRelativeResponse": { "properties": { @@ -1465,7 +1481,8 @@ "income", "substitution" ], - "title": "LaborSupplyRelativeResponse" + "title": "LaborSupplyRelativeResponse", + "description": "Income- and substitution-effect responses relative to baseline earnings." }, "LaborSupplyResponseOutput": { "properties": { @@ -1506,7 +1523,8 @@ "decile", "hours" ], - "title": "LaborSupplyResponseOutput" + "title": "LaborSupplyResponseOutput", + "description": "Behavioral labor supply response to the reform.\n\n``income_lsr`` and ``substitution_lsr`` are the total earnings changes from\nthe income and substitution effects, and ``total_change`` is their sum." }, "LocalAuthorityImpactRecord": { "properties": { @@ -1564,7 +1582,8 @@ "relative_household_income_change", "population" ], - "title": "LocalAuthorityImpactRecord" + "title": "LocalAuthorityImpactRecord", + "description": "Reform impact on household net income in one UK local authority.\n\nFields mean the same as in ``ConstituencyImpactRecord``." }, "PingRequest": { "properties": { @@ -1684,7 +1703,8 @@ "poverty", "deep_poverty" ], - "title": "PovertyByGenderOutput" + "title": "PovertyByGenderOutput", + "description": "Poverty and deep poverty rates by gender, as defined for ``PovertyOutput``." }, "PovertyByRaceOutput": { "properties": { @@ -1697,7 +1717,8 @@ "required": [ "poverty" ], - "title": "PovertyByRaceOutput" + "title": "PovertyByRaceOutput", + "description": "US poverty rates by race and ethnicity; there is no deep poverty breakdown." }, "PovertyOutput": { "properties": { @@ -1714,7 +1735,8 @@ "poverty", "deep_poverty" ], - "title": "PovertyOutput" + "title": "PovertyOutput", + "description": "Poverty and deep poverty rates by age group.\n\nIn the US these are the SPM poverty and deep SPM poverty rates. In the UK,\n``poverty`` is relative and ``deep_poverty`` absolute poverty, both measured\nbefore housing costs." }, "RacePovertyOutput": { "properties": { @@ -1739,7 +1761,8 @@ "hispanic", "other" ], - "title": "RacePovertyOutput" + "title": "RacePovertyOutput", + "description": "Baseline and reform poverty rates by race and ethnicity (US only)." }, "SPMCapability": { "properties": { @@ -1758,7 +1781,8 @@ "required": [ "defaults" ], - "title": "SPMCapability" + "title": "SPMCapability", + "description": "Canonical SPM support offered by a simulation version.\n\n``defaults`` is the pinned SPM selection, including the forecast artifact\nhash and scenario, that applies to any option a request leaves out." }, "SPMComparisonProvenance": { "properties": { @@ -1800,7 +1824,8 @@ "code", "message" ], - "title": "SPMErrorDetail" + "title": "SPMErrorDetail", + "description": "A typed SPM input error.\n\n``code`` is one of a fixed set of stable values, such as\n``SPM_GEOGRAPHY_REQUIRED``, that clients can branch on; ``message`` explains\nthe problem to a person." }, "SPMExecutionProvenance": { "properties": { From cfc112aa072581209b6272b5badb5185a7e45733 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 13 Oct 2026 17:25:59 +0400 Subject: [PATCH 13/18] Use the Public API's docstrings in the shared rules tool Both repos documented scripts/rules independently for ruff D, so the copies drifted apart. Take the Public API's wording so the tool stays byte-identical in both repos; only docstrings differ. Co-Authored-By: Claude Opus 5.5 --- scripts/rules/adapters.py | 4 ++++ scripts/rules/baseline.py | 6 +++++- scripts/rules/ci_scripts.py | 20 +++++++++++++++++--- scripts/rules/cli.py | 4 ++++ scripts/rules/discover.py | 1 + scripts/rules/guard.py | 7 +++++++ scripts/rules/load.py | 2 ++ scripts/rules/model.py | 30 ++++++++++++++++++++++++++++++ scripts/rules/render.py | 10 +++++++++- scripts/rules/resolve.py | 31 ++++++++++++++++++++++--------- scripts/rules/validate.py | 5 +++++ scripts/rules/workflows.py | 3 +++ 12 files changed, 109 insertions(+), 14 deletions(-) diff --git a/scripts/rules/adapters.py b/scripts/rules/adapters.py index b22bd112a..19a7ee979 100644 --- a/scripts/rules/adapters.py +++ b/scripts/rules/adapters.py @@ -18,6 +18,10 @@ def check_adapters(agents: Agents, root: Path) -> list[Violation]: + """Return a violation for each way an adapter stops being a short pointer. + + Missing adapter files are skipped rather than reported. + """ violations = [] instructions = Path(agents.instructions).name for adapter in agents.adapters: diff --git a/scripts/rules/baseline.py b/scripts/rules/baseline.py index f1a8aeb4b..5bc3feb90 100644 --- a/scripts/rules/baseline.py +++ b/scripts/rules/baseline.py @@ -18,12 +18,16 @@ class BaseUnavailable(Exception): def default_base_ref() -> str | None: + """Return ``origin/`` on a pull-request run, otherwise None.""" base = os.environ.get("GITHUB_BASE_REF") return f"origin/{base}" if base else None def load_base_registry(root: Path, ref: str) -> Registry | None: - """The registry as it is on ``ref``, or None if the base has no registry yet.""" + """Return the registry as it is on ``ref``, or None if the base has none yet. + + Raise BaseUnavailable when ``ref`` does not name a commit. + """ verify = subprocess.run( ["git", "rev-parse", "--verify", "--quiet", f"{ref}^{{commit}}"], cwd=root, diff --git a/scripts/rules/ci_scripts.py b/scripts/rules/ci_scripts.py index 5a0a6a066..0a0edecd2 100644 --- a/scripts/rules/ci_scripts.py +++ b/scripts/rules/ci_scripts.py @@ -36,6 +36,8 @@ @dataclass class Findings: + """What one ci-scripts run found, as baseline-style entries for each check.""" + multiline: set[str] = field(default_factory=set) untested: set[str] = field(default_factory=set) shellcheck: set[str] | None = None # None when shellcheck could not run @@ -43,8 +45,11 @@ class Findings: def multiline_runs(config: CiScripts, root: Path) -> set[str]: - """``workflow :: job :: step`` for every run: block with more than one line of - code, or with commands chained by ``&&``, ``||`` or ``;``.""" + """Return ``workflow :: job :: step`` for each run: block that is not one command. + + A block fails when it has more than one line of code, or chains commands with + ``&&``, ``||`` or ``;``. + """ found = set() for workflow in sorted(root.glob(config.workflows)): relative = workflow.relative_to(root).as_posix() @@ -66,7 +71,7 @@ def called_scripts(config: CiScripts, root: Path) -> set[str]: def untested_scripts(config: CiScripts, root: Path) -> set[str]: - """Called scripts whose file name appears in no test file.""" + """Return the called scripts whose file name appears in no test file.""" tests = [ path.read_text(errors="replace") for pattern in config.test_globs @@ -101,6 +106,7 @@ def shellcheck_failures(config: CiScripts, root: Path) -> set[str] | None: def find(config: CiScripts, root: Path) -> Findings: + """Run the three checks on freshly parsed workflows, noting a skipped shellcheck.""" clear_caches() findings = Findings( multiline=multiline_runs(config, root), @@ -138,6 +144,10 @@ def _rule_for(registry: Registry, baseline: str, check: str) -> str | None: def compare(registry: Registry, findings: Findings, root: Path) -> list[Violation]: + """Report findings missing from their baseline and baseline entries now fixed. + + A check that could not run is skipped rather than reported. + """ config = registry.ci_scripts assert config is not None checks = ( @@ -204,6 +214,10 @@ def write_baselines(registry: Registry, findings: Findings, root: Path) -> None: def run_check(root: Path = REPO_ROOT) -> tuple[list[Violation], list[str]]: + """Load the registry and return the ci-scripts violations and notes. + + A missing shellcheck is a violation on GitHub Actions and only a note elsewhere. + """ registry, violations = load_registry(root) if registry is None or registry.ci_scripts is None: return violations, [ diff --git a/scripts/rules/cli.py b/scripts/rules/cli.py index 9322bb086..a4134cacb 100644 --- a/scripts/rules/cli.py +++ b/scripts/rules/cli.py @@ -24,6 +24,10 @@ def _print(violations: list, notes: list[str] | None = None) -> int: def main(argv: list[str] | None = None) -> int: + """Run one rules command and return 1 if it found violations. + + ``render`` refuses to write anything while the registry fails validation. + """ parser = argparse.ArgumentParser( prog="python -m scripts.rules", description=__doc__ ) diff --git a/scripts/rules/discover.py b/scripts/rules/discover.py index 323425618..a352c05ef 100644 --- a/scripts/rules/discover.py +++ b/scripts/rules/discover.py @@ -60,6 +60,7 @@ def _discovered_tests(registry: Registry, root: Path) -> set[str]: def discover_hidden(registry: Registry, root: Path) -> list[Violation]: + """Return a hidden-rule violation for each check no rule references or ignores.""" violations = [] rules = registry.rules diff --git a/scripts/rules/guard.py b/scripts/rules/guard.py index c68fce02b..3efa3d4f5 100644 --- a/scripts/rules/guard.py +++ b/scripts/rules/guard.py @@ -38,6 +38,8 @@ @dataclass class GuardResult: + """Guard violations, which fail the run, and notes, which only inform.""" + violations: list[Violation] = field(default_factory=list) notes: list[str] = field(default_factory=list) @@ -75,6 +77,11 @@ def run_guard( today: date | None = None, compare_base: bool = True, ) -> GuardResult: + """Run every registry check, then compare baselines with the base ref. + + Stop after the schema errors if the registry does not load. An unreadable + base ref is a violation on GitHub Actions and only a note elsewhere. + """ clear_caches() result = GuardResult() registry, schema_violations = load_registry(root) diff --git a/scripts/rules/load.py b/scripts/rules/load.py index 83b455e2e..936a3a9f4 100644 --- a/scripts/rules/load.py +++ b/scripts/rules/load.py @@ -389,6 +389,7 @@ def parse(raw: dict[str, Any]) -> tuple[Registry | None, list[Violation]]: def load_text(text: str) -> tuple[Registry | None, list[Violation]]: + """Parse registry TOML text, or return None and one violation if it is not TOML.""" try: raw = tomllib.loads(text) except tomllib.TOMLDecodeError as error: @@ -404,6 +405,7 @@ def load_text(text: str) -> tuple[Registry | None, list[Violation]]: def load_registry(root: Path) -> tuple[Registry | None, list[Violation]]: + """Load the registry under ``root``, or return None and a violation if missing.""" path = root / REGISTRY_PATH if not path.exists(): return None, [ diff --git a/scripts/rules/model.py b/scripts/rules/model.py index 6d8f9c365..401becaaa 100644 --- a/scripts/rules/model.py +++ b/scripts/rules/model.py @@ -12,12 +12,16 @@ class Tier(StrEnum): + """How mechanically a rule can be checked; judgment-tier rules never block.""" + MECHANICAL = "mechanical" METRIC = "metric" JUDGMENT = "judgment" class Status(StrEnum): + """How a rule is enforced today; only enforced and ratcheted rules can block.""" + ENFORCED = "enforced" RATCHETED = "ratcheted" ADVISORY = "advisory" @@ -25,6 +29,8 @@ class Status(StrEnum): class Kind(StrEnum): + """The checker behind a rule; it decides which ``[rule.enforcer]`` fields apply.""" + PYTEST = "pytest" QUALITY_GUARD = "quality-guard" RUFF = "ruff" @@ -47,12 +53,15 @@ class Violation: fix: str def __str__(self) -> str: + """Format as ``[rule] check: message. Fix: fix.``; rule defaults to registry.""" subject = self.rule_id or "registry" return f"[{subject}] {self.check}: {self.message}. Fix: {self.fix}." @dataclass(frozen=True) class Category: + """A ``[[category]]`` grouping rules, in declared order, in the rendered views.""" + id: str title: str @@ -67,6 +76,8 @@ class Ignored: @dataclass(frozen=True) class Discovery: + """Where the guard looks for unregistered checks, and which ones it may skip.""" + workflows: tuple[str, ...] = (".github/workflows/pr.yml",) guard_test_globs: tuple[str, ...] = () extra_guard_tests: tuple[str, ...] = () @@ -97,6 +108,8 @@ class CiScripts: @dataclass(frozen=True) class Enforcer: + """What checks a rule: its ``kind`` plus the raw ``[rule.enforcer]`` fields.""" + kind: Kind fields: Mapping[str, Any] = field(default_factory=dict) @@ -105,6 +118,7 @@ def get( key: str, default: object = None, ) -> Any: # noqa: ANN401 # a raw TOML value that callers read by its shape + """Return a raw enforcer field, or ``default`` when the registry omits it.""" return self.fields.get(key, default) @property @@ -153,6 +167,8 @@ def expects_proof(self) -> bool: @dataclass(frozen=True) class CiRef: + """A ``[[rule.ci]]`` reference: a workflow job whose run text has ``invocation``.""" + workflow: str job: str invocation: str @@ -160,17 +176,26 @@ class CiRef: @dataclass(frozen=True) class Proof: + """The raw ``[rule.proof]`` fields; which keys are required depends on the tier.""" + fields: Mapping[str, Any] @dataclass(frozen=True) class Baseline: + """A ratcheted rule's file of tolerated violations and the number it lists. + + The count must match the file and may never grow relative to the base branch. + """ + file: str count: int @dataclass(frozen=True) class RuleException: + """A temporary carve-out of ``path`` from a rule; the guard fails once expired.""" + path: str reason: str expires: date @@ -178,6 +203,8 @@ class RuleException: @dataclass(frozen=True) class Rule: + """One ``[[rule]]`` registry entry; rules-registry.md explains each field.""" + id: str category: str statement: str @@ -205,6 +232,8 @@ def blocking(self) -> bool: @dataclass(frozen=True) class Registry: + """The parsed contents of rules/registry.toml.""" + schema_version: int repo: str categories: tuple[Category, ...] @@ -214,6 +243,7 @@ class Registry: ci_scripts: CiScripts | None = None def rule(self, rule_id: str) -> Rule | None: + """Return the rule with ``rule_id``, or None if the registry has none.""" for rule in self.rules: if rule.id == rule_id: return rule diff --git a/scripts/rules/render.py b/scripts/rules/render.py index 97815bf30..2b70361e4 100644 --- a/scripts/rules/render.py +++ b/scripts/rules/render.py @@ -52,7 +52,7 @@ def _proof_label(rule: Rule) -> str: def _enforced_by(rule: Rule) -> str: - """What enforces a rule, naming the CI job when the job itself is the check.""" + """Describe what enforces a rule, naming the CI job when the job is the check.""" if rule.enforcer.kind is Kind.WORKFLOW_JOB: jobs = ", ".join(f"`{ref.job}`" for ref in rule.ci) return f"CI job {jobs}" @@ -121,6 +121,7 @@ def _bullets(items: list[str]) -> list[str]: def render_rules_md(registry: Registry) -> str: + """Render rules.md: the tier summary, the gaps, then each category's rules.""" rules = sorted(registry.rules, key=lambda r: r.id) out = [ "# Rules", @@ -211,6 +212,11 @@ def _agents_line(rule: Rule) -> str: def render_agents_block(registry: Registry) -> str: + """Render the AGENTS.md rules block, markers included. + + Past ``AGENTS_COLLAPSE_AT`` rules, each category lists its checked rules on + one line and spells out only the doc-only ones. + """ out = [ BEGIN, "", @@ -245,6 +251,7 @@ def splice_block(text: str, block: str) -> str: def extract_block(text: str) -> str | None: + """Return the generated rules block in ``text``, markers included, or None.""" start, end = text.find(BEGIN), text.find(END) if start == -1 or end == -1 or end < start: return None @@ -252,6 +259,7 @@ def extract_block(text: str) -> str | None: def write_outputs(registry: Registry, root: Path) -> list[Path]: + """Write rules.md and splice the rules block into AGENTS.md; return both paths.""" rules_md = root / RULES_MD rules_md.parent.mkdir(parents=True, exist_ok=True) rules_md.write_text(render_rules_md(registry)) diff --git a/scripts/rules/resolve.py b/scripts/rules/resolve.py index cddb6536d..ed2bba559 100644 --- a/scripts/rules/resolve.py +++ b/scripts/rules/resolve.py @@ -34,6 +34,7 @@ def _read(path: Path) -> str | None: def clear_caches() -> None: + """Forget cached file reads and parses, so the next check sees the files afresh.""" _read.cache_clear() workflow_jobs.cache_clear() test_names.cache_clear() @@ -41,6 +42,7 @@ def clear_caches() -> None: @cache def workflow_jobs(path: Path) -> tuple[dict[str, JobBlock], tuple[str, ...]]: + """Return a workflow's jobs by id and its parse problems; cached per path.""" text = _read(path) if text is None: return {}, (f"{path.name} is missing",) @@ -97,7 +99,10 @@ def node_exists(root: Path, node: str) -> str | None: def guard_names(root: Path, runner: str = DEFAULT_RUNNER) -> list[str] | None: - """String names in the ``GUARDS`` tuple of the quality-guard runner.""" + """Return the string names in the runner's ``GUARDS`` tuple. + + Return None if the runner is missing, and an empty list if it has no ``GUARDS``. + """ text = _read(root / runner) if text is None: return None @@ -131,7 +136,7 @@ def _toml(path: Path) -> dict[str, Any]: def ruff_lint_config(root: Path) -> dict[str, Any]: - """The ``lint`` table from ruff.toml, or ``[tool.ruff.lint]`` in pyproject.""" + """Return ruff.toml's ``lint`` table, else ``[tool.ruff.lint]`` in pyproject.""" ruff_toml = _toml(root / "ruff.toml") if ruff_toml: return ruff_toml.get("lint", {}) @@ -141,12 +146,14 @@ def ruff_lint_config(root: Path) -> dict[str, Any]: def ruff_selected(root: Path) -> tuple[str, ...]: + """Return ruff's ``select`` (its default when unset) plus ``extend-select``.""" lint = ruff_lint_config(root) selected = tuple(lint.get("select", RUFF_DEFAULT_SELECT)) return selected + tuple(lint.get("extend-select", ())) def mypy_files(root: Path) -> list[str]: + """Return ``[tool.mypy].files`` as a list, even when it is a single string.""" files = ( _toml(root / "pyproject.toml").get("tool", {}).get("mypy", {}).get("files", []) ) @@ -323,6 +330,7 @@ def _resolve_workflow_job(rule: Rule, root: Path) -> list[Violation]: def resolve_enforcer(rule: Rule, root: Path) -> list[Violation]: + """Check that the rule's enforcer exists where its kind says it should.""" return RESOLVERS[rule.enforcer.kind](rule, root) @@ -332,7 +340,7 @@ def resolve_enforcer(rule: Rule, root: Path) -> list[Violation]: def _job_for( rule_id: str, ref: CiRef, root: Path, check: str, field: str ) -> tuple[JobBlock | None, list[Violation]]: - """The job a CI reference names, if it exists and runs the invocation.""" + """Return the job a CI reference names if it runs the invocation, plus problems.""" def v(message: str, fix: str) -> Violation: return Violation(rule_id, check, message, fix) @@ -386,12 +394,13 @@ def resolve_ci(rule: Rule, root: Path) -> list[Violation]: def resolve_proof_ci(rule: Rule, root: Path) -> list[Violation]: - """A proof test must run in CI, or nothing shows the checker still works. + """Check that a CI job collects the rule's proof test. - The proof test itself asserts that the bad fixture is rejected and the good - one accepted; the guard checks that some CI job collects it. A proof runs - in the rule's own CI job when the rule is pytest-enforced; otherwise the - proof names its job with ``workflow``, ``job`` and ``invocation``. + Without that, nothing shows the checker still works. The proof test itself + asserts that the bad fixture is rejected and the good one accepted; the + guard checks that some CI job collects it. A proof runs in the rule's own + CI job when the rule is pytest-enforced; otherwise the proof names its job + with ``workflow``, ``job`` and ``invocation``. """ if rule.proof is None or not isinstance(rule.proof.fields.get("test"), str): return [] @@ -434,7 +443,11 @@ def _test_coverage( root: Path, check: str, ) -> list[Violation]: - """The job's test command actually collects these root-relative test nodes.""" + """Check that the job's test command collects these root-relative test nodes. + + Only ``make`` targets and ``unittest discover`` commands are inspected; any + other command is trusted. + """ violations = [] collected: list[str] | None = None described = "" diff --git a/scripts/rules/validate.py b/scripts/rules/validate.py index f7d117d65..406b40ef6 100644 --- a/scripts/rules/validate.py +++ b/scripts/rules/validate.py @@ -55,6 +55,7 @@ def baseline_entries(path: Path) -> int: + """Count the entries in a baseline file, skipping blank and ``#`` comment lines.""" return sum( 1 for line in path.read_text().splitlines() @@ -242,6 +243,10 @@ def _validate_proof(rule: Rule, root: Path) -> list[Violation]: def validate(registry: Registry, root: Path, today: date) -> list[Violation]: + """Check the registry's fields and file references against the tree at ``root``. + + ``today`` decides whether review dates and exceptions have expired. + """ violations: list[Violation] = [] def v(rule_id: str | None, message: str, fix: str, check: str = "schema") -> None: diff --git a/scripts/rules/workflows.py b/scripts/rules/workflows.py index a432455fd..5c7c039b4 100644 --- a/scripts/rules/workflows.py +++ b/scripts/rules/workflows.py @@ -32,11 +32,14 @@ class RunBlock: @property def single_line(self) -> bool: + """Whether the block is at most one line of code with no chained commands.""" return self.code_lines <= 1 and not self.chained @dataclass(frozen=True) class JobBlock: + """One workflow job: its run text, step names, ``uses:`` references and matrix.""" + id: str name: str run_text: str From de9ce11d0857f18ca062bce3e3af20b9b570ded3 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 13 Oct 2026 17:25:59 +0400 Subject: [PATCH 14/18] Register D in the ruff-lint rule sim.quality.ruff-lint now lists D, and the rendered rules page and AGENTS.md Rules block are regenerated for the new tool hash. Co-Authored-By: Claude Opus 5.5 --- AGENTS.md | 2 +- changelog_entry.yaml | 1 + docs/generated/rules.md | 8 ++++---- rules/registry.toml | 2 +- 4 files changed, 7 insertions(+), 6 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 0bae03c1c..2ff683839 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -138,7 +138,7 @@ To add or change a rule, read `docs/engineering/skills/rules-registry.md`. - `sim.quality.every-project-tested`: Every project and library with unit tests runs them in the PR test matrix. (doc-only, not checked: follow it yourself) - `sim.quality.ruff-format-src`: Source code under projects/*/src and libs/*/src is formatted with ruff format. (blocking: ruff format) - `sim.quality.ruff-format-tests`: Test code is formatted with ruff format too; CI currently checks only src directories. (doc-only, not checked: follow it yourself) -- `sim.quality.ruff-lint`: All Python code passes ruff check with every rule family selected in ruff.toml, with no ratchet; families are added one themed group at a time until the selection is ALL. (blocking: ruff check (ANN, E4, E7, E9, F, FLY, FURB, I, PIE, PYI, RET, RSE, TC, UP, W)) +- `sim.quality.ruff-lint`: All Python code passes ruff check with every rule family selected in ruff.toml, with no ratchet; families are added one themed group at a time until the selection is ALL. (blocking: ruff check (ANN, D, E4, E7, E9, F, FLY, FURB, I, PIE, PYI, RET, RSE, TC, UP, W)) - `sim.quality.tests-pass`: Each project's unit tests pass in its own locked environment. (blocking: CI job `test`) - `sim.quality.tests-ruff-exemptions`: Test files may use pytest's bare assert and need no docstrings, annotations or named constants for magic numbers (S101, D1, ANN, PLR2004); every other ruff rule applies to tests. (blocking: ruff per-file ignore of ANN, D1, PLR2004, S101 for `**/tests/**`) diff --git a/changelog_entry.yaml b/changelog_entry.yaml index 7ed733326..d6c9b77b6 100644 --- a/changelog_entry.yaml +++ b/changelog_entry.yaml @@ -4,6 +4,7 @@ - Made CLAUDE.md and the Copilot instructions thin pointers to AGENTS.md, required workflow logic longer than one line to live in tested and shellchecked scripts, and registered the tracking-issue rule - Added a rules registry (rules/registry.toml) with a standard-library tool that validates it, checks every rule is enforced where it claims, and renders docs/generated/rules.md and the Rules block in AGENTS.md changed: + - Enabled ruff's docstring rules (D) and documented every module, class and function outside tests; published contract models and routes now carry these descriptions in their JSON schema and OpenAPI, with no structural schema change - Enabled ruff's annotation rules (ANN) and annotated every function outside tests; test files are exempt from ANN, missing docstrings, magic numbers and bare-assert rules - Enabled ruff's typing-only import rules (TC) and moved annotation-only imports under TYPE_CHECKING, keeping runtime imports for Pydantic models and FastAPI endpoints - Ran ruff check in CI with the first group of rule families (imports, pyupgrade and other safe auto-fixes) and fixed every finding diff --git a/docs/generated/rules.md b/docs/generated/rules.md index 1d979899d..f3db22c6c 100644 --- a/docs/generated/rules.md +++ b/docs/generated/rules.md @@ -245,7 +245,7 @@ flowchart LR | `sim.quality.every-project-tested` | Every project and library with unit tests runs them in the PR test matrix. | mechanical | doc-only | no | nothing | — | — | | `sim.quality.ruff-format-src` | Source code under projects/*/src and libs/*/src is formatted with ruff format. | mechanical | enforced | yes | ruff format | `pr.yml:lint` | — | | `sim.quality.ruff-format-tests` | Test code is formatted with ruff format too; CI currently checks only src directories. | mechanical | doc-only | no | nothing | — | — | -| `sim.quality.ruff-lint` | All Python code passes ruff check with every rule family selected in ruff.toml, with no ratchet; families are added one themed group at a time until the selection is ALL. | mechanical | enforced | yes | ruff check (ANN, E4, E7, E9, F, FLY, FURB, I, PIE, PYI, RET, RSE, TC, UP, W) | `pr.yml:lint` | — | +| `sim.quality.ruff-lint` | All Python code passes ruff check with every rule family selected in ruff.toml, with no ratchet; families are added one themed group at a time until the selection is ALL. | mechanical | enforced | yes | ruff check (ANN, D, E4, E7, E9, F, FLY, FURB, I, PIE, PYI, RET, RSE, TC, UP, W) | `pr.yml:lint` | — | | `sim.quality.tests-pass` | Each project's unit tests pass in its own locked environment. | mechanical | enforced | yes | CI job `test` | `pr.yml:test` | — | | `sim.quality.tests-ruff-exemptions` | Test files may use pytest's bare assert and need no docstrings, annotations or named constants for magic numbers (S101, D1, ANN, PLR2004); every other ruff rule applies to tests. | mechanical | enforced | yes | ruff per-file ignore of ANN, D1, PLR2004, S101 for `**/tests/**` | `pr.yml:lint` | — | @@ -258,7 +258,7 @@ flowchart LR j_CI_pr_yml_lint["CI pr.yml:lint"] r_sim_quality_ruff_format_tests("sim.quality.ruff-format-tests (doc-only)") r_sim_quality_ruff_lint("sim.quality.ruff-lint (enforced)") - e_ruff_check__ANN__E4__E7__E9__F__FLY__FURB__I__PIE__PYI__RET__RSE__TC__UP__W_["ruff check (ANN, E4, E7, E9, F, FLY, FURB, I, PIE, PYI, RET, RSE, TC, UP, W)"] + e_ruff_check__ANN__D__E4__E7__E9__F__FLY__FURB__I__PIE__PYI__RET__RSE__TC__UP__W_["ruff check (ANN, D, E4, E7, E9, F, FLY, FURB, I, PIE, PYI, RET, RSE, TC, UP, W)"] r_sim_quality_tests_pass("sim.quality.tests-pass (enforced)") e_CI_job_test["CI job test"] j_CI_pr_yml_test["CI pr.yml:test"] @@ -268,8 +268,8 @@ flowchart LR r_sim_quality_ruff_format_src --> e_ruff_format e_ruff_format --> j_CI_pr_yml_lint r_sim_quality_ruff_format_tests -.-> gap_unenforced - r_sim_quality_ruff_lint --> e_ruff_check__ANN__E4__E7__E9__F__FLY__FURB__I__PIE__PYI__RET__RSE__TC__UP__W_ - e_ruff_check__ANN__E4__E7__E9__F__FLY__FURB__I__PIE__PYI__RET__RSE__TC__UP__W_ --> j_CI_pr_yml_lint + r_sim_quality_ruff_lint --> e_ruff_check__ANN__D__E4__E7__E9__F__FLY__FURB__I__PIE__PYI__RET__RSE__TC__UP__W_ + e_ruff_check__ANN__D__E4__E7__E9__F__FLY__FURB__I__PIE__PYI__RET__RSE__TC__UP__W_ --> j_CI_pr_yml_lint r_sim_quality_tests_pass --> e_CI_job_test e_CI_job_test --> j_CI_pr_yml_test r_sim_quality_tests_ruff_exemptions --> e_ruff_per_file_ignore_of_ANN__D1__PLR2004__S101_for____tests___ diff --git a/rules/registry.toml b/rules/registry.toml index ae72bb9f1..7f57a58c8 100644 --- a/rules/registry.toml +++ b/rules/registry.toml @@ -710,7 +710,7 @@ autofix = true [rule.enforcer] kind = "ruff" mode = "lint" -codes = ["ANN", "E4", "E7", "E9", "F", "FLY", "FURB", "I", "PIE", "PYI", "RET", "RSE", "TC", "UP", "W"] +codes = ["ANN", "D", "E4", "E7", "E9", "F", "FLY", "FURB", "I", "PIE", "PYI", "RET", "RSE", "TC", "UP", "W"] [[rule.ci]] workflow = ".github/workflows/pr.yml" job = "lint" From ca508fef32ab2c47b491b74c4b3bf066a1195662 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 13 Oct 2026 21:07:16 +0400 Subject: [PATCH 15/18] Put the docstrings #767 added in the imperative mood #767 landed on main after this stack branched. D401, which this PR enables, wants each docstring's summary line in the imperative mood. Co-Authored-By: Claude Opus 5.5 --- .../simulation_output_common.py | 2 +- .../tests/test_simulation_output_budget.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_common.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_common.py index e2fd0d7f2..f14b24a6e 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_common.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_common.py @@ -58,7 +58,7 @@ def _entity_data(simulation: Simulation, entity: str) -> pd.DataFrame: def _count_output_entities(simulation: Simulation, entity: str, weight: str) -> float: - """The weighted number of entities: the plain sum of their weights. + """Return the weighted number of entities: the plain sum of their weights. Output frames are microdf ``MicroDataFrame``s weighted by the entity's weight, so a column's ``.sum()`` is already a weighted sum. Summing the diff --git a/projects/policyengine-simulation-executor/tests/test_simulation_output_budget.py b/projects/policyengine-simulation-executor/tests/test_simulation_output_budget.py index 9774d7597..42dca109d 100644 --- a/projects/policyengine-simulation-executor/tests/test_simulation_output_budget.py +++ b/projects/policyengine-simulation-executor/tests/test_simulation_output_budget.py @@ -12,7 +12,7 @@ def _simulation(columns: dict[str, list[float]]): - """A simulation whose household output is a MicroDataFrame, as policyengine builds it.""" + """Build a simulation whose household output is a MicroDataFrame, as policyengine does.""" household = MicroDataFrame(pd.DataFrame(columns), weights="household_weight") return SimpleNamespace( output_dataset=SimpleNamespace(data=SimpleNamespace(household=household)) From d752603a0830bfb6fc6f5fc34cfeb7d1ec454c76 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 14 Oct 2026 20:36:31 +0400 Subject: [PATCH 16/18] Put the docstrings #743 and #744 added in the D style #743 and #744 landed on main after this stack branched. D, which this PR enables, flags their new code: blank lines after docstrings (D202), a two-line summary (D205, D209), summaries that are not in the imperative mood (D401), an escaped backslash in the sentinel's module docstring (D301), and three public functions without docstrings (D103). Co-Authored-By: Claude Opus 5.5 --- .../scripts/update_policyengine_package.py | 1 - .../tests/test_deployment_assets.py | 9 -------- .../tests/test_routing_workflow_invariants.py | 2 -- .../tests/test_simulation_rollback.py | 5 ---- .../src/modal/utils/release_sentinel.py | 15 +++++------- .../modal/utils/update_version_registry.py | 23 ++++--------------- .../release_bundle.py | 6 ++--- ...est_policyengine_package_update_scripts.py | 5 ---- .../tests/test_release_sentinel.py | 5 ---- .../tests/test_update_version_registry.py | 8 ------- .../endpoints.py | 4 +--- .../tests/test_endpoints.py | 1 - .../tests/test_spm_routes.py | 6 +++-- .../tests/test_versions_fixture.py | 3 --- 14 files changed, 18 insertions(+), 75 deletions(-) diff --git a/.github/scripts/update_policyengine_package.py b/.github/scripts/update_policyengine_package.py index aab80691c..c1bb21f2a 100644 --- a/.github/scripts/update_policyengine_package.py +++ b/.github/scripts/update_policyengine_package.py @@ -285,7 +285,6 @@ def verify_dispatch() -> None: promotable is refused. ``dispatch_verified`` is true only when at least one field was compared. Prints digests for the PR body. """ - raw = os.environ.get("RELEASE_PAYLOAD", "").strip() payload = json.loads(raw) if raw and raw != "null" else {} installed = {} diff --git a/projects/policyengine-simulation-entry/tests/test_deployment_assets.py b/projects/policyengine-simulation-entry/tests/test_deployment_assets.py index 5c7b0c080..d978edba6 100644 --- a/projects/policyengine-simulation-entry/tests/test_deployment_assets.py +++ b/projects/policyengine-simulation-entry/tests/test_deployment_assets.py @@ -1121,7 +1121,6 @@ def test_finalizer_always_resolves_a_publication_with_the_release_gate_outcome( @pytest.mark.parametrize("job", ["integration", "promote_entrypoint"]) def test_a_failing_sentinel_fails_its_step(tmp_path, job): """#744 review: `|| true` on the sentinel command must not pass a failure.""" - step = _sentinel_check_step(_workflow_jobs(REUSABLE_DEPLOY_WORKFLOW)[job]) bin_dir = tmp_path / "bin" bin_dir.mkdir() @@ -1151,7 +1150,6 @@ def test_a_failing_sentinel_fails_its_step(tmp_path, job): def test_stage12_only_promotion_runs_whatever_the_skipped_stack_jobs(): """#743 review: the explicit status function must stay active.""" - promote = _workflow_jobs(REUSABLE_DEPLOY_WORKFLOW)["promote_entrypoint"] context = { "inputs": {"promote_entrypoint": True, "deploy_existing_stack": False}, @@ -1174,7 +1172,6 @@ def test_stage12_only_promotion_runs_whatever_the_skipped_stack_jobs(): def test_smoke_rejects_a_latest_wrapper_that_lags_alone(tmp_path: Path): """#743 review: the wrapper's latest is checked, not only the model's.""" - import copy versions = copy.deepcopy(VERSIONS_BODY) @@ -1190,7 +1187,6 @@ def test_smoke_rejects_a_latest_wrapper_that_lags_alone(tmp_path: Path): def _run_routing_step_with_a_failing_command(tmp_path, step, failed_command, **env): """Run the real body with Actions' shell and a failing CLI seam.""" - bin_dir = tmp_path / "bin" bin_dir.mkdir() (bin_dir / "uv").write_text( @@ -1867,7 +1863,6 @@ def test_smoke_accepts_an_unrecorded_digest_only_when_none_is_served( tmp_path: Path, ): """A release published before digests were recorded serves none.""" - import copy versions = copy.deepcopy(VERSIONS_BODY) @@ -1885,7 +1880,6 @@ def test_smoke_accepts_an_unrecorded_digest_only_when_none_is_served( def test_smoke_rejects_worker_routes_to_an_old_app(tmp_path: Path): """#743 round-2 review: the version and its model route to the new app.""" - import copy versions = copy.deepcopy(VERSIONS_BODY) @@ -1913,7 +1907,6 @@ def test_smoke_rejects_worker_routes_to_an_old_app(tmp_path: Path): @pytest.mark.parametrize("how", ["unset", "empty"]) def test_smoke_refuses_an_incomplete_identity(tmp_path: Path, missing, how): """#743 round-2 review: an absent expectation never weakens the gate.""" - expected = dict(IDENTITY) if how == "unset": del expected[missing] @@ -1950,7 +1943,6 @@ def test_smoke_refuses_malformed_expectations(tmp_path: Path, name, value): def test_smoke_rejects_a_bundle_recorded_for_another_wrapper(tmp_path: Path): """#743 round-2 review: the bundle must be the expected version's own.""" - import copy versions = copy.deepcopy(VERSIONS_BODY) @@ -1970,7 +1962,6 @@ def _workflow_jobs(path: Path) -> dict: def test_deploy_gates_pass_the_release_identity_to_the_smoke(): """Review M2: the active workflow (not comments) wires the whole identity.""" - jobs = _workflow_jobs(REUSABLE_DEPLOY_WORKFLOW) assert ( jobs["prepare"]["outputs"]["uk_dataset_sha256"] diff --git a/projects/policyengine-simulation-entry/tests/test_routing_workflow_invariants.py b/projects/policyengine-simulation-entry/tests/test_routing_workflow_invariants.py index e80226ab1..72fdfc29c 100644 --- a/projects/policyengine-simulation-entry/tests/test_routing_workflow_invariants.py +++ b/projects/policyengine-simulation-entry/tests/test_routing_workflow_invariants.py @@ -105,7 +105,6 @@ def _routing_commands(script: str): Importing the library or using shell wrappers needs a deliberate extension here, rather than silently escaping discovery or the ordering proof. """ - script = re.sub(r"\\\n\s*", " ", script) script = "\n".join( line for line in script.splitlines() if not line.lstrip().startswith("#") @@ -205,7 +204,6 @@ def _routing_commands(script: str): def _command_options(arguments: tuple[str, ...]) -> set[str]: """Parse option positions; an option-looking value is not a CLI mode.""" - options = set() index = 0 while index < len(arguments): diff --git a/projects/policyengine-simulation-entry/tests/test_simulation_rollback.py b/projects/policyengine-simulation-entry/tests/test_simulation_rollback.py index ebf9a9990..972e43b3d 100644 --- a/projects/policyengine-simulation-entry/tests/test_simulation_rollback.py +++ b/projects/policyengine-simulation-entry/tests/test_simulation_rollback.py @@ -205,7 +205,6 @@ def test_operator_recovery_prints_the_journal_then_restores_without_artifacts(): def test_rollback_verifies_the_whole_restored_identity(): """#743/#744: the smoke's identity is all or nothing, so pass all of it.""" - verify = _rollback_steps()[ _step_index(_rollback_steps(), "cloud-run-simulation-entry-smoke.sh") ] @@ -253,7 +252,6 @@ def _fake_gh( inventory: dict[int, list[dict]] | None = None, ) -> None: """Return API JSON and execute the workflow's real --jq filter with jq.""" - payloads = {} for attempt in range(1, attempts + 1): jobs = [_inventory_job() for _ in range(kept.get(attempt, 0))] @@ -326,7 +324,6 @@ def _run_restore( def test_rollback_passes_every_attempts_plan_oldest_first(tmp_path): """#744 review: a later attempt must not hide an earlier publication.""" - attempts = [1, 10, 2, 9] result = _run_restore( tmp_path, 10, {attempt: 1 for attempt in attempts}, present=attempts @@ -343,7 +340,6 @@ def test_rollback_passes_every_attempts_plan_oldest_first(tmp_path): def test_rollback_refuses_when_an_attempts_plan_is_missing(tmp_path): """#744 review: attempt 1 published but its plan is gone; attempt 2's is not.""" - result = _run_restore(tmp_path, 2, {1: 1, 2: 1}, present=[2]) assert result.returncode == 1 assert "Attempt 1 of run 1 kept a prod routing plan that cannot be found" in ( @@ -381,7 +377,6 @@ def test_inventory_jq_selects_only_successfully_kept_target_plans( tmp_path, ignored_job ): """The workflow's job/step/conclusion jq selectors all run, not a stub.""" - result = _run_restore( tmp_path, 2, diff --git a/projects/policyengine-simulation-executor/src/modal/utils/release_sentinel.py b/projects/policyengine-simulation-executor/src/modal/utils/release_sentinel.py index 486326510..fc6bcf76c 100644 --- a/projects/policyengine-simulation-executor/src/modal/utils/release_sentinel.py +++ b/projects/policyengine-simulation-executor/src/modal/utils/release_sentinel.py @@ -1,4 +1,4 @@ -"""Numerical release sentinel: the deployed API computes what the bundle computes. +r"""Numerical release sentinel: the deployed API computes what the bundle computes. ``expected`` runs each sentinel request in this process through ``run_simulation_impl``, the pure implementation the Modal simulation worker @@ -17,7 +17,7 @@ Usage:: python -m src.modal.utils.release_sentinel expected --output sentinel.json - python -m src.modal.utils.release_sentinel check --base-url URL \\ + python -m src.modal.utils.release_sentinel check --base-url URL \ --expected sentinel.json """ @@ -79,8 +79,7 @@ class SentinelError(RuntimeError): def worker_payload(request: Mapping[str, Any]) -> dict[str, Any]: - """The payload the gateway forwards to the worker for ``request``.""" - + """Build the payload the gateway forwards to the worker for ``request``.""" from policyengine_simulation_contract.gateway_models import SimulationRequest return SimulationRequest.model_validate(dict(request)).model_dump( @@ -91,6 +90,7 @@ def worker_payload(request: Mapping[str, Any]) -> dict[str, Any]: def value_at(result: Mapping[str, Any], path: str) -> float: + """Return the finite number at the dotted ``path`` in ``result``.""" node: Any = result for key in path.split("."): if not isinstance(node, Mapping) or key not in node: @@ -105,7 +105,6 @@ def value_at(result: Mapping[str, Any], path: str) -> float: def agree(expected: float, actual: float) -> bool: """Equal up to floating-point reduction order (symmetric).""" - return abs(expected - actual) <= max( ABS_TOL, REL_TOL * max(abs(expected), abs(actual)) ) @@ -146,7 +145,6 @@ def compute_expected( cases: tuple[dict[str, Any], ...] = SENTINEL_CASES, ) -> dict[str, Any]: """Run every case in this process and record its values.""" - runner = runner or _local_runner() identity = dict(identity or _bundle_identity()) if not identity["uk"].get("dataset_sha256"): @@ -201,8 +199,7 @@ def _http( def access_token(environ: Mapping[str, str] = os.environ) -> str | None: - """A client-credentials token when the test auth variables are configured.""" - + """Return a client-credentials token when the test auth variables are configured.""" present = [name for name in _AUTH_VARS if environ.get(name)] if not present: if _truthy(environ.get("SIMULATION_TEST_AUTH_REQUIRED")): @@ -249,7 +246,6 @@ def check_deployment( clock: Callable[[], float] = time.monotonic, ) -> list[str]: """Submit each expected case to ``base_url`` and return what disagrees.""" - if expected.get("schema") != SCHEMA: raise SentinelError(f"expected file schema is {expected.get('schema')!r}") http = http or _http @@ -320,6 +316,7 @@ def check_deployment( def main(argv: list[str] | None = None) -> int: + """Run the ``expected`` or ``check`` command and return its exit status.""" parser = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) commands = parser.add_subparsers(dest="command", required=True) expected_parser = commands.add_parser("expected", help="compute in-process values") 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 ae24844be..96127ee20 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 @@ -340,7 +340,6 @@ def build_restored_routing_state( the routing plan its deployment recorded (:func:`restore_routing_plan`): which app served a model version before cannot be read from the state. """ - state = _coerce_routing_state(current_state) current_latest = state["latest"].get("policyengine") if expected_latest is not None and current_latest != expected_latest: @@ -397,7 +396,6 @@ class RoutingConflict(RuntimeError): def routing_state_digest(state: object) -> str: """SHA-256 of a routing state's canonical JSON (``null`` when absent).""" - text = json.dumps(state, sort_keys=True, separators=(",", ":"), ensure_ascii=False) return hashlib.sha256(text.encode("utf-8")).hexdigest() @@ -416,7 +414,7 @@ def plan_routing_publication( force_latest: bool = False, publication_id: str | None = None, ) -> dict: - """The routing state before and after publishing a release, with digests. + """Plan the routing state before and after publishing a release, with digests. The published state's ``generation`` names this publication uniquely (``publication_id``, by default a random one), so a later publication of @@ -424,7 +422,6 @@ def plan_routing_publication( restored over it. ``advances_latest`` says whether unpinned requests route to this release once the plan is applied. """ - publication_id = publication_id or uuid.uuid4().hex before = _json_copy(current_state) if current_state is not None else None after = _json_copy( @@ -452,8 +449,7 @@ def plan_routing_publication( def validate_routing_plan(plan: object) -> dict: - """A plan whose digests match its states, or ``ValueError``.""" - + """Return a plan whose digests match its states, or raise ``ValueError``.""" if not isinstance(plan, dict) or plan.get("schema") != ROUTING_PLAN_SCHEMA: raise ValueError(f"not a {ROUTING_PLAN_SCHEMA} routing plan") for side in ("before", "after"): @@ -491,7 +487,6 @@ def _routing_store(environment: str, *, create_if_missing: bool) -> "modal.Dict def _require_serialized_routing_workflow() -> None: """Refuse routing writes outside the workflows that serialize them.""" - # GitHub associates a reusable workflow's github context with its caller, # so GITHUB_WORKFLOW_REF identifies the caller, not the reusable file: # https://docs.github.com/en/actions/reference/workflows-and-actions/reusing-workflow-configurations#github-context @@ -569,7 +564,6 @@ def record_routing_plan(*, environment: str, plan: dict) -> None: file or uploading its artifact fails. Operator recovery can restore a plan if the finalizer itself is lost. """ - _require_serialized_routing_workflow() validate_routing_plan(plan) store = _routing_store(environment, create_if_missing=True) @@ -584,7 +578,6 @@ def record_routing_plan(*, environment: str, plan: dict) -> None: def complete_routing_plan(*, environment: str, plan: dict) -> None: """Acknowledge the owning release only after all its gates have passed.""" - _require_serialized_routing_workflow() validate_routing_plan(plan) store = _routing_store(environment, create_if_missing=False) @@ -600,7 +593,6 @@ def complete_routing_plan(*, environment: str, plan: dict) -> None: def apply_routing_plan(*, environment: str, plan: dict) -> None: """Publish ``plan``'s state if the live state is still its starting point.""" - _require_serialized_routing_workflow() validate_routing_plan(plan) store = _routing_store(environment, create_if_missing=True) @@ -626,7 +618,6 @@ def restore_routing_plans(*, environment: str, plans: list[dict]) -> str: starting state no plan replaced (nothing was applied, or it was already undone). Any other live state is a conflict and nothing is changed. """ - _require_serialized_routing_workflow() for plan in plans: validate_routing_plan(plan) @@ -674,7 +665,6 @@ def restore_routing_plans(*, environment: str, plans: list[dict]) -> str: def restore_routing_plan(*, environment: str, plan: dict) -> str: """Undo one plan: ``restored``, ``unchanged`` (never applied) or a conflict.""" - _require_serialized_routing_workflow() return restore_routing_plans(environment=environment, plans=[plan]) @@ -683,7 +673,6 @@ def _publication_owned_by_release( plan: dict, *, publication_run: str | None, release_environment: str ) -> bool: """Match a workflow run across attempts, or the operator's selected store.""" - if publication_run is None: return True @@ -713,7 +702,6 @@ def finalize_routing_publication( Already finished records owned by this run return their recorded outcome without writing, so retrying a finalizer after output loss is safe. """ - _require_serialized_routing_workflow() if release_environment not in {"beta", "prod"}: raise ValueError("release_environment must be beta or prod") @@ -757,8 +745,7 @@ def finalize_routing_publication( def latest_release_identity(state: object) -> dict[str, str]: - """What unpinned requests are routed to: version, app, UK model and dataset.""" - + """Describe what unpinned requests route to: version, app, UK model and dataset.""" routing = _coerce_routing_state(state) version = routing["latest"].get("policyengine") or "" manifest = routing["bundles"].get(version) or {} @@ -1036,8 +1023,7 @@ def publish_routing_state( def read_latest_policyengine_version(*, environment: str) -> str | None: - """The policyengine.py version unpinned requests are routed to now.""" - + """Return the policyengine.py version unpinned requests are routed to now.""" snapshot = _modal_dict_snapshot( name=ROUTING_STATE_DICT_NAME, environment=environment ) @@ -1051,6 +1037,7 @@ def restore_routing_state( policyengine_version: str, expected_latest: str | None = None, ) -> RoutingState: + """Point ``latest`` back at a published policyengine.py version and store it.""" _require_serialized_routing_workflow() store = modal.Dict.from_name( ROUTING_STATE_DICT_NAME, diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/release_bundle.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/release_bundle.py index e13aab3aa..42daa350f 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/release_bundle.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/release_bundle.py @@ -62,12 +62,11 @@ class CountryReleaseBundle: def _sha256_or_none(value: object) -> str | None: - """A SHA-256 as 64 lowercase hex digits, or ``None`` when malformed. + """Return a SHA-256 as 64 lowercase hex digits, or ``None`` when malformed. Only hexadecimal digits count: ``int(value, 16)`` would also accept a sign, underscores or surrounding whitespace. """ - if isinstance(value, str) and _SHA256_HEX.fullmatch(value): return value.lower() return None @@ -79,12 +78,11 @@ def _default_dataset_sha256( data_release: Mapping, manifest: CountryReleaseManifest, ) -> str | None: - """The certified SHA-256 of the default dataset, from the bundle. + """Return the certified SHA-256 of the default dataset, from the bundle. The certified data artifact's hash counts only when that artifact is the default dataset; otherwise the default dataset's own artifact pin is used. """ - certified = _mapping(data_release.get("certified_data_artifact")) if certified.get("dataset") == default_dataset: digest = _sha256_or_none(certified.get("sha256")) diff --git a/projects/policyengine-simulation-executor/tests/test_policyengine_package_update_scripts.py b/projects/policyengine-simulation-executor/tests/test_policyengine_package_update_scripts.py index ba5f0522f..a627dd5a8 100644 --- a/projects/policyengine-simulation-executor/tests/test_policyengine_package_update_scripts.py +++ b/projects/policyengine-simulation-executor/tests/test_policyengine_package_update_scripts.py @@ -217,7 +217,6 @@ def test_a_recovered_branch_gets_no_pr_without_verification( message: str, ) -> None: """#743 round-2 review: branch recovery may not bypass validation.""" - gh_log = tmp_path / "gh.log" install_fake_git( fake_bin, @@ -470,7 +469,6 @@ def test_verify_dispatch_rejects_another_bundle() -> None: def test_verify_dispatch_compares_the_certificate_when_installed_reports_one() -> None: """#743 round-2 review: removing the certificate comparison must fail a test.""" - import os installed = INSTALLED.replace( @@ -495,7 +493,6 @@ def verify(payload: str) -> subprocess.CompletedProcess[str]: def test_verify_dispatch_fails_closed_when_the_installed_bundle_cannot_answer() -> None: """Review M3: a named dataset digest needs an installed one to compare.""" - import os installed = INSTALLED.replace("uk_dataset_sha256=" + "a" * 64, "uk_dataset_sha256=") @@ -585,7 +582,6 @@ def test_bundle_versions_matches_policyengine_release_identity() -> None: certified_data_artifact.sha256, packages.policyengine-core.version), so a real dispatch for the locked wrapper verifies. """ - import json import os @@ -631,7 +627,6 @@ def test_bundle_versions_reports_the_receipt_certificate_digest(tmp_path) -> Non the digest it computes from the UK record must reach the output unchanged, and a dispatch naming another certificate must then be refused. """ - import json import os diff --git a/projects/policyengine-simulation-executor/tests/test_release_sentinel.py b/projects/policyengine-simulation-executor/tests/test_release_sentinel.py index 2cb153d4a..3d0c125a5 100644 --- a/projects/policyengine-simulation-executor/tests/test_release_sentinel.py +++ b/projects/policyengine-simulation-executor/tests/test_release_sentinel.py @@ -254,7 +254,6 @@ def test_the_sentinel_compares_the_budget_and_net_income(): @pytest.mark.parametrize("path", COMPARED) def test_every_compared_value_must_agree(path): """#744 review: removing any one comparison must fail a test.""" - section, key = path.split(".") result = {**RESULT, section: {**RESULT[section], key: RESULT[section][key] * 1.01}} problems, _ = _check(FakeAPI(result=result)) @@ -269,7 +268,6 @@ def test_every_compared_value_must_agree(path): @pytest.mark.parametrize("value", [True, False, "6.1e9", None, [6.1e9]]) def test_a_value_that_is_not_a_number_fails(value): """#744 review: a boolean is not a number, whatever Python says.""" - result = {**RESULT, "budget": {**RESULT["budget"], "budgetary_impact": value}} problems, _ = _check(FakeAPI(result=result)) assert any("not a number" in problem for problem in problems), problems @@ -301,7 +299,6 @@ def test_a_job_that_did_not_complete_fails_even_with_the_right_numbers(status): ) def test_a_complete_job_without_an_object_result_fails(job): """#744 review: a worker returning None must not pass the sentinel.""" - problems, _ = _check(FakeAPI(job=job)) assert len(problems) == 1 assert "result, not an object" in problems[0] @@ -312,7 +309,6 @@ def test_a_complete_job_without_an_object_result_fails(job): ) def test_the_tolerance_is_floating_point_reordering_only(relative, agrees): """#744 review: a looser REL_TOL (a 1e-3 mismatch is £4m) must fail a test.""" - value = 4e9 assert sentinel.agree(value, value * (1 + relative)) is agrees @@ -323,7 +319,6 @@ def test_the_tolerance_is_floating_point_reordering_only(relative, agrees): ) def test_the_absolute_tolerance_is_one_penny(actual, agrees): """Near zero, numerical agreement admits up to and including £0.01 of error.""" - assert sentinel.agree(0, actual) is agrees assert sentinel.agree(actual, 0) is agrees 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..0ec5e9af5 100644 --- a/projects/policyengine-simulation-executor/tests/test_update_version_registry.py +++ b/projects/policyengine-simulation-executor/tests/test_update_version_registry.py @@ -612,7 +612,6 @@ def _store(patched_modal, state): @given(releases=RELEASES, legacy=st.booleans()) def test_publish_then_restore_returns_exactly_the_previous_state(releases, legacy): """#744 review: first, repeat and older releases are all undone exactly.""" - monkeypatch = pytest.MonkeyPatch() stores: dict = {} try: @@ -664,7 +663,6 @@ def from_name(name, environment_name, create_if_missing): def test_the_review_counterexample_restores_every_pin(monkeypatch, patched_modal): """Legacy 6.0/model1, then 6.1/model2, then a failed 6.2/model1.""" - state = registry._coerce_routing_state(None) state["routes"]["policyengine"]["6.0.0"] = _app("6.0.0") state["routes"]["uk"]["2.1.0"] = _app("6.0.0") @@ -868,7 +866,6 @@ def test_latest_identity_of_an_empty_or_legacy_state_is_blank(): def test_a_failed_retry_attempt_cannot_hide_the_publication(monkeypatch, patched_modal): """#744 review: attempt 1 published A to B; attempt 2 republished B.""" - a = _publish(monkeypatch, None, "6.0.0", "1.1.0", "2.1.0") store = _store(patched_modal, a) first = _plan(monkeypatch, a, "6.1.0", "1.1.0", "2.2.0") @@ -888,7 +885,6 @@ def test_an_old_plan_cannot_withdraw_a_healthy_republication( monkeypatch, patched_modal ): """#744 review: A to B failed and was restored; a healthy run republished B.""" - a = _publish(monkeypatch, None, "6.0.0", "1.1.0", "2.1.0") store = _store(patched_modal, a) failed = _plan(monkeypatch, a, "6.1.0", "1.1.0", "2.2.0") @@ -1292,7 +1288,6 @@ def pop(self, key): def test_runner_loss_cannot_be_adopted_by_the_next_release(monkeypatch, patched_modal): """Failed A cannot become successful B's starting state or rollback target.""" - original = _publish(monkeypatch, None, "6.0.0", "1.1.0", "2.1.0") store = _store(patched_modal, original) failed_a = _plan(monkeypatch, original, "6.1.0", "1.1.0", "2.2.0") @@ -1342,7 +1337,6 @@ def test_a_failed_same_version_redeploy_is_resolved_by_its_restore( A plan's published state carries its own publication id, so even a same-version redeploy differs from what it replaces and has a restore. """ - live = _publish(monkeypatch, None, "6.0.0", "1.1.0", "2.1.0") store = _store(patched_modal, live) redeploy = _plan(monkeypatch, live, "6.0.0", "1.1.0", "2.1.0") @@ -1452,7 +1446,6 @@ def test_publication_ownership_compares_the_entire_plan( monkeypatch, patched_modal, operation ): """A reused publication id never grants permission to a changed plan.""" - before = _publish(monkeypatch, None, "6.0.0", "1.1.0", "2.1.0") store = _store(patched_modal, before) pending = _plan(monkeypatch, before, "6.1.0", "1.1.0", "2.2.0") @@ -1497,7 +1490,6 @@ def test_finalization_finishes_each_release_and_allows_the_next_one( first_publication, ): """Upload loss and attempt 2 use the full journal, with no artifact reads.""" - before = ( None if first_publication 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 3fffdf1bf..c9a9b2028 100644 --- a/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/endpoints.py +++ b/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/endpoints.py @@ -509,13 +509,12 @@ def _resolve_from_legacy_dicts( def _certified_dataset_sha256(country_bundle: dict) -> str | None: - """The certified SHA-256 of the app's default dataset, if recorded. + """Return the certified SHA-256 of the app's default dataset, if recorded. ``dataset`` in a bundle response is always the app's default dataset URI (``_resolve_dataset_uri_from_app_bundle``), so the digest describes it. Malformed or missing values are omitted rather than repeated. """ - digest = country_bundle.get("default_dataset_sha256") # fullmatch: ``$`` would also accept a trailing newline. if isinstance(digest, str) and _SHA256_HEX.fullmatch(digest): @@ -572,7 +571,6 @@ def _bundle_identities(state: dict) -> dict[str, dict[str, PolicyEngineBundle]]: dataset cannot be resolved is reported without a dataset rather than failing the whole response. """ - identities: dict[str, dict[str, PolicyEngineBundle]] = {} for policyengine_version, app_bundle in _routing_state_bundles(state).items(): if not isinstance(app_bundle, dict): diff --git a/projects/policyengine-simulation-gateway/tests/test_endpoints.py b/projects/policyengine-simulation-gateway/tests/test_endpoints.py index 60c12d5ab..0fa500987 100644 --- a/projects/policyengine-simulation-gateway/tests/test_endpoints.py +++ b/projects/policyengine-simulation-gateway/tests/test_endpoints.py @@ -1106,7 +1106,6 @@ def test__given_malformed_dataset_sha__then_versions_omit_it( self, mock_modal, client: TestClient, malformed ): """#743 round-2 review: a digest with a trailing newline gave a 500.""" - state = deepcopy(TEST_ROUTING_STATE) state["bundles"]["4.10.0"]["uk"] = { **state["bundles"]["4.10.0"]["uk"], diff --git a/projects/policyengine-simulation-gateway/tests/test_spm_routes.py b/projects/policyengine-simulation-gateway/tests/test_spm_routes.py index 8ff06fd43..df9500900 100644 --- a/projects/policyengine-simulation-gateway/tests/test_spm_routes.py +++ b/projects/policyengine-simulation-gateway/tests/test_spm_routes.py @@ -622,8 +622,10 @@ def test_a_bundle_that_states_no_model_version_contradicts_nothing( def test_poll_bodies_carry_the_dataset_sha256_only_when_recorded(mock_modal, client): - """A worker bundle with a certified dataset digest reports it in raw - 202 bodies; one without keeps the legacy key set.""" + """Report a recorded dataset digest in raw 202 bodies. + + A worker bundle without a certified dataset digest keeps the legacy key set. + """ from copy import deepcopy from fixtures.gateway_endpoints import TEST_ROUTING_STATE diff --git a/projects/policyengine-simulation-gateway/tests/test_versions_fixture.py b/projects/policyengine-simulation-gateway/tests/test_versions_fixture.py index 352961eea..790800bea 100644 --- a/projects/policyengine-simulation-gateway/tests/test_versions_fixture.py +++ b/projects/policyengine-simulation-gateway/tests/test_versions_fixture.py @@ -51,7 +51,6 @@ def _app(version: str) -> str: def routing_state() -> dict: """Two releases, published in order; 6.2.7 is latest.""" - releases = { "6.2.6": ("1.700.0", "2.125.1", "5" * 64), "6.2.7": ("1.700.0", "2.126.8", "a" * 64), @@ -92,7 +91,6 @@ def test_smoke_fixture_is_the_gateway_versions_response(mock_modal, client): def test_versions_survives_an_unresolvable_historical_dataset(mock_modal, client): """Review M1: one bad entry is reported without a dataset, not a 500.""" - state = routing_state() state["bundles"]["6.2.6"]["uk"]["default_dataset_uri"] = ( "gs://policyengine-uk-data-private/enhanced_frs_2024_25.h5@1.39.0" @@ -115,7 +113,6 @@ def test_versions_skips_country_entries_without_a_model_version( mock_modal, client, model_version ): """#743 review: a blank or non-string model is no identity at all.""" - state = routing_state() if model_version == "missing": del state["bundles"]["6.2.6"]["uk"]["model_version"] From bf536f8f9330211a5a10be60dbc846c461299a12 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 14 Oct 2026 21:19:17 +0400 Subject: [PATCH 17/18] Put the docstrings #789 added in the D style #789 landed on main after this stack branched. D, which this PR enables, flags blank lines after its new docstrings (D202) and the undocumented __call__ of its ConnectionOpener protocol (D102). Co-Authored-By: Claude Opus 5.5 --- .../src/policyengine_stage12_persistence/store.py | 5 +++-- libs/policyengine-stage12-persistence/tests/test_store.py | 1 - .../stage12_runtime/database_retry.py | 1 - .../stage12_runtime/failure_recording.py | 2 -- 4 files changed, 3 insertions(+), 6 deletions(-) diff --git a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/store.py b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/store.py index 9936cb36f..87996a455 100644 --- a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/store.py +++ b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/store.py @@ -142,7 +142,9 @@ def __call__( connect: Callable[[], Connection], *, operation_name: str, - ) -> Connection: ... + ) -> Connection: + """Return the connection ``connect`` opens for ``operation_name``.""" + ... def open_connection_once( @@ -151,7 +153,6 @@ def open_connection_once( operation_name: str, ) -> Connection: """Open a connection once, so a refused connection fails the operation.""" - return connect() diff --git a/libs/policyengine-stage12-persistence/tests/test_store.py b/libs/policyengine-stage12-persistence/tests/test_store.py index 9df7938b2..96dad3b0e 100644 --- a/libs/policyengine-stage12-persistence/tests/test_store.py +++ b/libs/policyengine-stage12-persistence/tests/test_store.py @@ -199,7 +199,6 @@ def __call__(self, connect, *, operation_name): def _sqlalchemy_commit_failure(monkeypatch) -> OperationalError: """Return the error SQLAlchemy raises when a COMMIT's reply is lost.""" - engine = create_engine("sqlite://") def commit_then_lose_the_reply(dbapi_connection): diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/database_retry.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/database_retry.py index cbaed0477..0869b4fc2 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/database_retry.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/database_retry.py @@ -46,7 +46,6 @@ def open_connection_with_retry( Errors other than ``OperationalError`` propagate immediately, and retries stop once the retry budget is spent. """ - started = monotonic() attempt = 1 delay_ceiling = INITIAL_RETRY_DELAY_SECONDS diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/failure_recording.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/failure_recording.py index 2da7527ab..d0da8beeb 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/failure_recording.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/failure_recording.py @@ -27,7 +27,6 @@ def report_original_failure( session pooler has no free connections. Telemetry therefore receives the original cause first, so it survives a failed or hung database write. """ - _emit_error( "Stage 12 execution failed", error=error, @@ -49,7 +48,6 @@ def persist_failure_record( report and simulation handlers raise their stable ``RuntimeError``. The original cause then remains visible only in telemetry. """ - try: write_failure_record() except Exception as persistence_error: # noqa: BLE001 From 05d4fb2106fe79658523aa7ae3607838d2fda1f2 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 15 Oct 2026 00:25:56 +0400 Subject: [PATCH 18/18] Document the shared object store fake #787 added #787 landed on main after this stack branched. D, which this PR enables, flags the undocumented __init__ (D107), upload_bytes and read_bytes (D102) of the in-memory object store it moved into the executor's fixtures. Co-Authored-By: Claude Opus 5.5 --- .../fixtures/fake_object_store.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/projects/policyengine-simulation-executor/fixtures/fake_object_store.py b/projects/policyengine-simulation-executor/fixtures/fake_object_store.py index 23cbca3e9..de476cb50 100644 --- a/projects/policyengine-simulation-executor/fixtures/fake_object_store.py +++ b/projects/policyengine-simulation-executor/fixtures/fake_object_store.py @@ -10,9 +10,11 @@ class InMemoryObjectStore: """Keep uploaded objects by path, refusing to replace an existing one.""" def __init__(self) -> None: + """Start with no stored objects.""" self.objects: dict[str, bytes] = {} def upload_bytes(self, path: str, payload: bytes, *, content_type: str) -> bool: + """Store ``payload`` at ``path`` and report whether this call created it.""" # The real store uploads with ``if_generation_match=0``: an existing # object is left in place and the call reports that it created nothing. if path in self.objects: @@ -21,4 +23,5 @@ def upload_bytes(self, path: str, payload: bytes, *, content_type: str) -> bool: return True def read_bytes(self, path: str) -> bytes | None: + """Return the bytes stored at ``path``, or None when nothing is there.""" return self.objects.get(path)