diff --git a/docs/engineering/skills/observability.md b/docs/engineering/skills/observability.md index 2134d896f..d0d08d303 100644 --- a/docs/engineering/skills/observability.md +++ b/docs/engineering/skills/observability.md @@ -52,7 +52,7 @@ add `_observability_context` to it. At the receiving function, validate the separate argument with `normalize_observability_context` and pass the result to -`runtime.operation(..., remote_context=context)`. The operation scope must +`scopes.operation(runtime, ..., remote_context=context)`. The operation scope must cover the worker call. Child dispatches made inside that scope call `runtime.capture_context()` again and pass the resulting context through their own keyword-only argument. @@ -79,10 +79,12 @@ All span names for supported run configurations are defined in appropriate `StagePlan` and call `plan.name(Stage.VALUE)`. Do not add span name string literals in services or workers. -Open a span with `scopes.span(runtime, plan.name(Stage.VALUE))` from -`policyengine_simulation_observability.scopes` rather than `runtime.span(...)`. -It returns the same scope, typed so that Pyright knows an exception raised in -the block propagates; see `docs/engineering/skills/type-checking.md`. +Open a span with `scopes.span(runtime, plan.name(Stage.VALUE))`, and an +operation with `scopes.operation(runtime, ...)`, from +`policyengine_simulation_observability.scopes` rather than `runtime.span(...)` +or `runtime.operation(...)`. Each returns the runtime's own scope, typed so +that Pyright knows an exception raised in the block propagates; see +`docs/engineering/skills/type-checking.md`. When adding a run configuration or runtime stage: diff --git a/docs/engineering/skills/type-checking.md b/docs/engineering/skills/type-checking.md index bef001d5b..5d2b83cf6 100644 --- a/docs/engineering/skills/type-checking.md +++ b/docs/engineering/skills/type-checking.md @@ -9,12 +9,12 @@ package. Every uv package is type-checked with Pyright on every pull request, by the `pyright` job in `.github/workflows/pr.yml`: -- `sim.quality.pyright-libraries-and-entry`: the contract, observability, - Stage 12 persistence and fastapi libraries and the simulation entry have no - Pyright errors. -- `sim.quality.pyright-gateway`, `sim.quality.pyright-executor` and - `sim.quality.pyright-apis-integ`: these packages may have only the errors - their baselines list (see [Baselines](#baselines)). +- `sim.quality.pyright-libraries-and-entry`, `sim.quality.pyright-gateway` and + `sim.quality.pyright-executor`: the contract, observability, Stage 12 + persistence and fastapi libraries, the simulation entry, the gateway and the + executor have no Pyright errors. +- `sim.quality.pyright-apis-integ`: the integration tests may have only the + errors their baseline lists (see [Baselines](#baselines)). - `sim.quality.every-package-type-checked`: a Pyright rule lists every uv package under `libs/` and `projects/`. When you add a package, give it a `[tool.pyright] include`, add it to a Pyright rule's `enforcer.projects`, and @@ -114,8 +114,14 @@ Baselines are under `rules/`, so changing one needs code-owner approval same line: `# pyright: ignore[reportAttributeAccessIssue] # `. - If an error shows a real bug, first write a test that demonstrates it, then fix it. -- Open observability spans with `scopes.span(runtime, ...)` from - `policyengine_simulation_observability.scopes`, not `runtime.span(...)`. - policyengine-observability types its scope managers as possibly suppressing - exceptions, so after `with runtime.span(...)` Pyright reports each variable - the block assigned as possibly unbound. +- Open observability spans and operations with `scopes.span(runtime, ...)` and + `scopes.operation(runtime, ...)` from + `policyengine_simulation_observability.scopes`, not `runtime.span(...)` or + `runtime.operation(...)`. policyengine-observability types its scope + managers as possibly suppressing exceptions, so after `with runtime.span(...)` + Pyright reports each variable the block assigned as possibly unbound. +- Before reading a simulation's output frames, narrow + `.output_dataset.data` with the `EntityFrames` or + `HouseholdFrame` protocols in `policyengine_simulation_executor.output_frames`, + and keep reading through that expression: the microdata weight checks + recognise weighted frames by it. diff --git a/docs/generated/rules.md b/docs/generated/rules.md index f394f4683..8137e8ce7 100644 --- a/docs/generated/rules.md +++ b/docs/generated/rules.md @@ -8,7 +8,7 @@ Every rule that governs `policyengine-sim-api` is listed here with how it is enf | Tier | enforced | ratcheted | advisory | doc-only | Total | | --- | ---: | ---: | ---: | ---: | ---: | -| mechanical | 45 | 5 | 0 | 15 | 65 | +| mechanical | 47 | 3 | 0 | 15 | 65 | | metric | 0 | 0 | 0 | 0 | 0 | | judgment | 0 | 0 | 0 | 0 | 0 | @@ -68,8 +68,6 @@ Every rule that governs `policyengine-sim-api` is listed here with how it is enf - `sim.ci.multiline-run-is-script`: 13 remaining in `rules/baselines/ci-multiline-run.txt` - `sim.microdata.weight-stripping-confined`: 1 remaining in `rules/baselines/sim-microdata-weight-stripping.txt` - `sim.quality.pyright-apis-integ`: 4 remaining in `rules/baselines/pyright-apis-integ.txt` -- `sim.quality.pyright-executor`: 77 remaining in `rules/baselines/pyright-executor.txt` -- `sim.quality.pyright-gateway`: 0 remaining in `rules/baselines/pyright-gateway.txt` ### Exceptions and review dates @@ -288,8 +286,8 @@ flowchart LR | `sim.quality.every-package-type-checked` | Every uv package under libs/ and projects/ (a directory with a pyproject.toml) is listed by a Pyright rule, so the pyright CI job type-checks it. | mechanical | enforced | yes | quality guard `pyright-coverage` | `pr.yml:quality-guards` | yes | | `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.pyright-apis-integ` | The generated-client integration-test project, apis-integ, has no Pyright errors beyond those its baseline lists, and the baseline lists only errors that still occur, so it only shrinks; checked with the pinned Pyright version in the project's own locked environment. | mechanical | ratcheted | yes | pyright in `projects/policyengine-apis-integ` | `pr.yml:pyright` | yes | -| `sim.quality.pyright-executor` | The simulation executor has no Pyright errors beyond those its baseline lists, and the baseline lists only errors that still occur, so it only shrinks; checked with the pinned Pyright version in the executor's own locked environment. | mechanical | ratcheted | yes | pyright in `projects/policyengine-simulation-executor` | `pr.yml:pyright` | yes | -| `sim.quality.pyright-gateway` | The simulation gateway has no Pyright errors beyond those its baseline lists, and the baseline lists only errors that still occur, so it only shrinks; checked with the pinned Pyright version in the gateway's own locked environment. | mechanical | ratcheted | yes | pyright in `projects/policyengine-simulation-gateway` | `pr.yml:pyright` | yes | +| `sim.quality.pyright-executor` | The simulation executor has no Pyright errors, checked with the pinned Pyright version in the executor's own locked environment. | mechanical | enforced | yes | pyright in `projects/policyengine-simulation-executor` | `pr.yml:pyright` | yes | +| `sim.quality.pyright-gateway` | The simulation gateway has no Pyright errors, checked with the pinned Pyright version in the gateway's own locked environment. | mechanical | enforced | yes | pyright in `projects/policyengine-simulation-gateway` | `pr.yml:pyright` | yes | | `sim.quality.pyright-libraries-and-entry` | The contract, observability, Stage 12 persistence and fastapi libraries and the simulation entry have no Pyright errors, checked with the pinned Pyright version in each package's own locked environment. | mechanical | enforced | yes | pyright in `libs/policyengine-simulation-contract`, `libs/policyengine-simulation-observability`, `libs/policyengine-stage12-persistence`, `libs/policyengine-fastapi`, `projects/policyengine-simulation-entry` | `pr.yml:pyright` | yes | | `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 | — | — | @@ -317,9 +315,9 @@ flowchart LR r_sim_quality_pyright_apis_integ("sim.quality.pyright-apis-integ (ratcheted)") e_pyright_in_projects_policyengine_apis_integ["pyright in projects/policyengine-apis-integ"] j_CI_pr_yml_pyright["CI pr.yml:pyright"] - r_sim_quality_pyright_executor("sim.quality.pyright-executor (ratcheted)") + r_sim_quality_pyright_executor("sim.quality.pyright-executor (enforced)") e_pyright_in_projects_policyengine_simulation_executor["pyright in projects/policyengine-simulation-executor"] - r_sim_quality_pyright_gateway("sim.quality.pyright-gateway (ratcheted)") + r_sim_quality_pyright_gateway("sim.quality.pyright-gateway (enforced)") e_pyright_in_projects_policyengine_simulation_gateway["pyright in projects/policyengine-simulation-gateway"] r_sim_quality_pyright_libraries_and_entry("sim.quality.pyright-libraries-and-entry (enforced)") e_pyright_in_libs_policyengine_simulation_contract__libs_policyengine_simulation_observability__libs_policyengine_stage12_persistence__libs_policyengine_fastapi__projects_policyengine_simulation_entry["pyright in libs/policyengine-simulation-contract, libs/policyengine-simulation-observability, libs/policyengine-stage12-persistence, libs/policyengine-fastapi, projects/policyengine-simulation-entry"] diff --git a/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/scopes.py b/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/scopes.py index c79bc552c..d1a31400b 100644 --- a/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/scopes.py +++ b/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/scopes.py @@ -1,4 +1,4 @@ -"""The observability runtime's span scope, typed as never suppressing an exception. +"""The observability runtime's scopes, typed as never suppressing an exception. policyengine-observability annotates its scope managers' ``__exit__`` as returning ``bool``, although it always returns ``False``. A type checker reads a @@ -7,9 +7,10 @@ the block assigned as possibly unbound, and a function that returns from inside the block as missing a return. -``span`` returns the runtime's own scope unchanged, typed as a context manager -whose ``__exit__`` returns ``bool | None``: the type ``contextlib.contextmanager`` -scopes have, which type checkers don't treat as suppressing exceptions. +``span`` and ``operation`` return the runtime's own scope unchanged, typed as a +context manager whose ``__exit__`` returns ``bool | None``: the type +``contextlib.contextmanager`` scopes have, which type checkers don't treat as +suppressing exceptions. """ from __future__ import annotations @@ -35,3 +36,19 @@ def span( the block, and lets that exception propagate. """ return runtime.span(name, attributes=attributes) + + +def operation( + runtime: ObservabilityRuntime, + name: str, + *, + attributes: Mapping[str, Any] | None = None, + remote_context: Mapping[str, Any] | None = None, +) -> AbstractContextManager[object]: + """Return ``runtime.operation(name, ...)`` with these attributes and context. + + The scope times its block as an operation, continuing the trace in + ``remote_context`` when there is one, records an exception raised in the + block, and lets that exception propagate. + """ + return runtime.operation(name, attributes=attributes, remote_context=remote_context) diff --git a/libs/policyengine-simulation-observability/tests/test_scopes.py b/libs/policyengine-simulation-observability/tests/test_scopes.py index c86cad316..7b524bc45 100644 --- a/libs/policyengine-simulation-observability/tests/test_scopes.py +++ b/libs/policyengine-simulation-observability/tests/test_scopes.py @@ -1,4 +1,4 @@ -"""Tests for the typed span scope of the observability runtime.""" +"""Tests for the typed span and operation scopes of the observability runtime.""" from contextlib import nullcontext @@ -8,15 +8,20 @@ class RecordingRuntime: - """Records the spans it is asked to open, and enters each with a marker.""" + """Records the scopes it is asked to open, and enters each with a marker.""" def __init__(self) -> None: self.opened_spans: list[tuple[str, object]] = [] + self.opened_operations: list[tuple[str, object, object]] = [] def span(self, name, *, attributes=None): self.opened_spans.append((name, attributes)) return nullcontext("entered span state") + def operation(self, name, *, attributes=None, remote_context=None): + self.opened_operations.append((name, attributes, remote_context)) + return nullcontext("entered operation state") + @pytest.fixture def runtime(monkeypatch): @@ -61,3 +66,38 @@ def test_span_lets_an_exception_raised_inside_it_propagate(runtime): scopes.span(runtime, "simulation.parse"), ): raise ValueError(message) + + +def test_operation_opens_the_runtime_operation_with_its_context(): + recording_runtime = RecordingRuntime() + + with scopes.operation( + recording_runtime, + "simulation.execution", + attributes={"country": "us"}, + remote_context={"traceparent": "00-trace"}, + ) as entered_state: + assert recording_runtime.opened_operations == [ + ("simulation.execution", {"country": "us"}, {"traceparent": "00-trace"}) + ] + + assert entered_state == "entered operation state" + + +def test_operation_defaults_to_no_attributes_or_remote_context(): + recording_runtime = RecordingRuntime() + + with scopes.operation(recording_runtime, "simulation.execution"): + pass + + assert recording_runtime.opened_operations == [("simulation.execution", None, None)] + + +def test_operation_lets_an_exception_raised_inside_it_propagate(runtime): + message = "simulation failed" + + with ( + pytest.raises(RuntimeError, match=message), + scopes.operation(runtime, "simulation.execution"), + ): + raise RuntimeError(message) diff --git a/projects/policyengine-simulation-executor/src/modal/app.py b/projects/policyengine-simulation-executor/src/modal/app.py index eb52bab63..b43760f32 100644 --- a/projects/policyengine-simulation-executor/src/modal/app.py +++ b/projects/policyengine-simulation-executor/src/modal/app.py @@ -24,6 +24,7 @@ get_bundled_country_model_version, get_bundled_package_version, ) +from policyengine_simulation_observability import scopes from policyengine_simulation_observability.observability import ( init_process_observability, modal_image_environment, @@ -361,7 +362,8 @@ def run_simulation( } try: propagated = normalize_observability_context(observability_context) - with runtime.operation( + with scopes.operation( + runtime, ANNUAL_IMPACT_STAGES.name(Stage.ANNUAL_EXECUTION), attributes=redacted_params, remote_context=propagated, @@ -433,7 +435,8 @@ def run_simulation_segment( } try: propagated = normalize_observability_context(observability_context) - with runtime.operation( + with scopes.operation( + runtime, SEGMENTED_NATIONAL_STAGES.name(Stage.SEGMENTED_NATIONAL_EXECUTION), attributes=redacted_params, remote_context=propagated, @@ -473,7 +476,8 @@ def run_budget_window_batch( } try: propagated = normalize_observability_context(observability_context) - with runtime.operation( + with scopes.operation( + runtime, BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_EXECUTION), attributes=redacted_params, remote_context=propagated, 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 90ec90208..d03a4c065 100644 --- a/projects/policyengine-simulation-executor/src/modal/budget_window_batch.py +++ b/projects/policyengine-simulation-executor/src/modal/budget_window_batch.py @@ -4,6 +4,7 @@ from typing import TYPE_CHECKING, Any +from policyengine_simulation_observability import scopes from policyengine_simulation_observability.stages import BUDGET_WINDOW_STAGES, Stage from src.modal.budget_window_context import build_batch_context from src.modal.budget_window_scheduler import BudgetWindowBatchRunner @@ -36,7 +37,7 @@ def run_budget_window_batch_impl( msg = "run_budget_window_batch must run inside a Modal function call" raise RuntimeError(msg) runtime.set_context(batch_job_id=batch_job_id) - with runtime.span(BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_CONTEXT)): + with scopes.span(runtime, BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_CONTEXT)): context = build_batch_context( params, batch_job_id=batch_job_id, 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 4c1c06ae1..c2b082ae3 100644 --- a/projects/policyengine-simulation-executor/src/modal/budget_window_results.py +++ b/projects/policyengine-simulation-executor/src/modal/budget_window_results.py @@ -115,7 +115,13 @@ def sum_annual_impacts( ) totals["budgetaryImpact"] += _as_decimal(annual_impact.budgetaryImpact) - return BudgetWindowTotals(**{key: float(value) for key, value in totals.items()}) + return BudgetWindowTotals( + taxRevenueImpact=float(totals["taxRevenueImpact"]), + federalTaxRevenueImpact=float(totals["federalTaxRevenueImpact"]), + stateTaxRevenueImpact=float(totals["stateTaxRevenueImpact"]), + benefitSpendingImpact=float(totals["benefitSpendingImpact"]), + budgetaryImpact=float(totals["budgetaryImpact"]), + ) def build_budget_window_result( 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 3478c46b2..3213ab1c5 100644 --- a/projects/policyengine-simulation-executor/src/modal/budget_window_scheduler.py +++ b/projects/policyengine-simulation-executor/src/modal/budget_window_scheduler.py @@ -20,6 +20,7 @@ put_batch_job_state, ) from policyengine_simulation_contract.spm import spm_error_detail +from policyengine_simulation_observability import scopes from policyengine_simulation_observability.errors import log_and_redact_exception from policyengine_simulation_observability.stages import BUDGET_WINDOW_STAGES, Stage from src.modal.budget_window_context import ( @@ -71,8 +72,8 @@ 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) + with scopes.span( + runtime, BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_STATUS_SERIALIZATION) ): return build_batch_status_response(state).model_dump(mode="json") @@ -83,7 +84,9 @@ 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)): + with scopes.span( + runtime, BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_STATE_LOAD) + ): state = get_batch_job_seed(context.batch_job_id) if state is None: state = create_initial_batch_state( @@ -94,11 +97,26 @@ def load_or_create_batch_state( ), bundle=context.bundle, ) - with runtime.span(BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_STATE_WRITE)): + with scopes.span( + runtime, BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_STATE_WRITE) + ): put_batch_job_seed(state) return state +def _child_call(handle: ChildSimulationHandle) -> modal.FunctionCall[Any]: + """Return the child's Modal call, raising if the handle has none. + + resolve_child_handle always returns a handle with its call. The poll calls + this inside its try, so a missing call would fail the batch like any other + poll error. + """ + if handle.call is None: + msg = f"the {handle.simulation_year} child has no call handle" + raise RuntimeError(msg) + return handle.call + + class BudgetWindowBatchRunner: """Runs a parent budget-window batch job to completion.""" @@ -155,8 +173,8 @@ def run(self) -> dict[str, Any]: failure fails the whole batch. """ mark_batch_running(self.state) - with self.runtime.span( - BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_STATE_WRITE) + with scopes.span( + self.runtime, BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_STATE_WRITE) ): put_batch_job_state(self.state) @@ -196,7 +214,8 @@ def spawn_until_capacity(self) -> None: and self.state.queued_years ): simulation_year = self.state.queued_years[0] - with self.runtime.span( + with scopes.span( + self.runtime, BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_CHILD_REQUEST_BUILD), attributes={"simulation_year": simulation_year}, ): @@ -205,7 +224,8 @@ def spawn_until_capacity(self) -> None: simulation_year=simulation_year, ) observability_context = self.runtime.capture_context() - with self.runtime.span( + with scopes.span( + self.runtime, BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_CHILD_SPAWN), attributes={"simulation_year": simulation_year}, ): @@ -223,8 +243,8 @@ def spawn_until_capacity(self) -> None: year=simulation_year, child_job_id=call.object_id, ) - with self.runtime.span( - BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_STATE_WRITE) + with scopes.span( + self.runtime, BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_STATE_WRITE) ): put_batch_job_state(self.state) @@ -241,7 +261,7 @@ def poll_running_children_once(self) -> bool: poll_started = time.perf_counter() try: - child_result = handle.call.get(timeout=0) + child_result = _child_call(handle).get(timeout=0) except TimeoutError: continue except Exception as exc: # noqa: BLE001 # call.get re-raises whatever the child worker raised; every failure fails the batch, logged with its traceback by log_and_redact_exception @@ -274,7 +294,8 @@ def poll_running_children_once(self) -> bool: self._publish_poll_stats() try: - with self.runtime.span( + with scopes.span( + self.runtime, BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_RESULT_PARSE), attributes={"simulation_year": simulation_year}, ): @@ -313,8 +334,8 @@ def poll_running_children_once(self) -> bool: year=simulation_year, annual_impact=annual_impact, ) - with self.runtime.span( - BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_STATE_WRITE) + with scopes.span( + self.runtime, BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_STATE_WRITE) ): put_batch_job_state(self.state) progress_made = True @@ -357,15 +378,15 @@ def fail_batch_for_child_error( """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( - BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_STATE_WRITE) + with scopes.span( + self.runtime, BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_STATE_WRITE) ): 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) + with scopes.span( + self.runtime, BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_AGGREGATION) ): annual_impacts = [ self.state.partial_annual_impacts[simulation_year] @@ -378,8 +399,8 @@ def complete_batch(self) -> dict[str, Any]: annual_impacts=annual_impacts, ) mark_batch_complete(self.state, result=result) - with self.runtime.span( - BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_STATE_WRITE) + with scopes.span( + self.runtime, BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_STATE_WRITE) ): put_batch_job_state(self.state) return serialize_batch_status(self.state, runtime=self.runtime) diff --git a/projects/policyengine-simulation-executor/src/modal/segmented_national.py b/projects/policyengine-simulation-executor/src/modal/segmented_national.py index 0a291da68..955b42d20 100644 --- a/projects/policyengine-simulation-executor/src/modal/segmented_national.py +++ b/projects/policyengine-simulation-executor/src/modal/segmented_national.py @@ -587,8 +587,9 @@ def _reduce( # The canonical internal-key stripper, plus the opt-out knob. simulation_params, _, _ = split_internal_payload(self.params) simulation_params.pop("segmented", None) - with self.runtime.span( - SEGMENTED_NATIONAL_STAGES.name(Stage.SEGMENTED_NATIONAL_REDUCE) + with scopes.span( + self.runtime, + SEGMENTED_NATIONAL_STAGES.name(Stage.SEGMENTED_NATIONAL_REDUCE), ): output = build_national_output( child_results, 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 6b4eb35fd..55e93f3f1 100644 --- a/projects/policyengine-simulation-executor/src/modal/utils/release_sentinel.py +++ b/projects/policyengine-simulation-executor/src/modal/utils/release_sentinel.py @@ -411,7 +411,8 @@ def _report(line: str, *, error: bool = False) -> None: 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]) + # __doc__ is None when Python runs with -OO, which strips docstrings. + parser = argparse.ArgumentParser(description=(__doc__ or "").split("\n\n")[0]) commands = parser.add_subparsers(dest="command", required=True) expected_parser = commands.add_parser("expected", help="compute in-process values") expected_parser.add_argument("--output", required=True) 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 b11acff9b..0c2b5fd9b 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 @@ -11,11 +11,11 @@ import json import os import uuid -from collections.abc import Callable, Iterable +from collections.abc import Callable, Iterable, Mapping from copy import deepcopy from dataclasses import dataclass from pathlib import Path -from typing import NotRequired, TypedDict +from typing import Any, NotRequired, TypedDict import modal from packaging.version import InvalidVersion, Version @@ -66,8 +66,11 @@ class BundleManifestMetadata(TypedDict): job_progress: NotRequired[int] -class LatestVersions(TypedDict): - """The versions that unpinned requests route to, per route kind.""" +class LatestVersions(TypedDict, total=False): + """The versions that unpinned requests route to, per route kind. + + A kind has no entry until a release of that kind is routed. + """ policyengine: str us: str @@ -89,7 +92,11 @@ class RoutingState(TypedDict): generation: str latest: LatestVersions routes: RouteMaps - bundles: dict[str, BundleManifestMetadata] + # Each policyengine version's bundle manifest, as publishing writes it + # (BundleManifestMetadata). A stored state is read back unvalidated, so + # its manifests are only known to be mappings until validate_routing_state + # checks them against their routes. + bundles: dict[str, Mapping[str, Any]] def _is_newer_version(candidate: str, current: str | None) -> bool: @@ -456,7 +463,8 @@ def routing_state_digest(state: object) -> str: return hashlib.sha256(text.encode("utf-8")).hexdigest() -def _json_copy(value: object) -> object: +def _json_copy[JsonValue](value: JsonValue) -> JsonValue: + """Return a deep copy of JSON-safe ``value``; the round trip keeps its shape.""" return json.loads(json.dumps(value)) @@ -496,7 +504,7 @@ def plan_routing_publication( "after": after, "before_sha256": routing_state_digest(before), "after_sha256": routing_state_digest(after), - "advances_latest": after["latest"]["policyengine"] == policyengine_version, + "advances_latest": after["latest"].get("policyengine") == policyengine_version, } @@ -810,9 +818,12 @@ def finalize_routing_publication( msg = "operator recovery requires gates_passed=false" raise ValueError(msg) store = _routing_store(environment, create_if_missing=False) + if store is None: + # No routing-state Dict, so nothing was ever published or recorded. + return "no-pending" pending = _unfinished_publication(store) if pending is None: - record = store.get(ROUTING_PUBLICATION_KEY) if store is not None else None + record = store.get(ROUTING_PUBLICATION_KEY) if ( record is not None and publication_run is not None @@ -1048,8 +1059,9 @@ def _merge_legacy_seed_into_current( for kind in ("policyengine", "us", "uk"): merged["routes"][kind].update(current["routes"][kind]) - if kind in current["latest"]: - merged["latest"][kind] = current["latest"][kind] + current_latest = current["latest"].get(kind) + if current_latest is not None: + merged["latest"][kind] = current_latest merged["bundles"].update(current["bundles"]) validate_routing_state(merged) @@ -1372,7 +1384,7 @@ def _finalize_command(args: argparse.Namespace) -> None: raise _UsageError(msg) store = _routing_store(args.environment, create_if_missing=False) pending = _unfinished_publication(store) - if pending is None: + if store is None or pending is None: _report("No pending routing publication; routing is unchanged.") else: _report("Pending routing publication:") @@ -1437,9 +1449,9 @@ def _restore_latest_command(args: argparse.Namespace) -> None: expected_latest=args.expected_latest, ) _report(f"Restored routing state in Modal environment: {args.environment}") - _report(f" latest.policyengine: {state['latest']['policyengine']}") - _report(f" latest.us: {state['latest']['us']}") - _report(f" latest.uk: {state['latest']['uk']}") + _report(f" latest.policyengine: {state['latest'].get('policyengine')}") + _report(f" latest.us: {state['latest'].get('us')}") + _report(f" latest.uk: {state['latest'].get('uk')}") def _expected_latest_without_restore(_args: argparse.Namespace) -> None: @@ -1518,9 +1530,9 @@ def _publish_command(args: argparse.Namespace) -> None: ) _report() _report(f" {ROUTING_STATE_DICT_NAME}[{ROUTING_STATE_ACTIVE_KEY}]: updated") - _report(f" latest.policyengine: {state['latest']['policyengine']}") - _report(f" latest.us: {state['latest']['us']}") - _report(f" latest.uk: {state['latest']['uk']}") + _report(f" latest.policyengine: {state['latest'].get('policyengine')}") + _report(f" latest.us: {state['latest'].get('us')}") + _report(f" latest.uk: {state['latest'].get('uk')}") _report() _report("Routing state published successfully.") diff --git a/projects/policyengine-simulation-executor/src/modal/v2_app.py b/projects/policyengine-simulation-executor/src/modal/v2_app.py index ff6f7f09a..669c95a61 100644 --- a/projects/policyengine-simulation-executor/src/modal/v2_app.py +++ b/projects/policyengine-simulation-executor/src/modal/v2_app.py @@ -30,6 +30,7 @@ UK_SIMULATION_TIMEOUT_SECONDS, US_SIMULATION_TIMEOUT_SECONDS, ) +from policyengine_simulation_observability import scopes from policyengine_simulation_observability.observability import ( init_process_observability, modal_image_environment, @@ -273,7 +274,8 @@ def run_single_simulation_segment_us( try: propagated = normalize_observability_context(observability_context) result = None - with runtime.operation( + with scopes.operation( + runtime, STAGE12_SEGMENT_STAGES.name(Stage.STAGE12_SEGMENT_EXECUTION), attributes={ "runner_name": "stage12", @@ -330,7 +332,8 @@ def segmented_calculator( try: propagated = normalize_observability_context(observability_context) result = None - with runtime.operation( + with scopes.operation( + runtime, STAGE12_SIMULATION_STAGES.name(Stage.STAGE12_SIMULATION_EXECUTION), attributes={ "runner_name": "stage12", @@ -380,7 +383,8 @@ def run_single_simulation_uk( try: propagated = normalize_observability_context(observability_context) result = None - with runtime.operation( + with scopes.operation( + runtime, STAGE12_SIMULATION_STAGES.name(Stage.STAGE12_SIMULATION_EXECUTION), attributes={ "runner_name": "stage12", @@ -446,7 +450,8 @@ def coordinate_report( try: propagated = normalize_observability_context(observability_context) result = None - with runtime.operation( + with scopes.operation( + runtime, stage_plan.name(Stage.STAGE12_COORDINATOR_EXECUTION), attributes={ "runner_name": "stage12", 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 7f0b82cfb..c0e54e47d 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 @@ -22,18 +22,20 @@ import logging from dataclasses import dataclass -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Final, Protocol, runtime_checkable from policyengine.core.scoping_strategy import RegionGroupStrategy from policyengine_simulation_contract.spm import SPMInputError, spm_error_detail from pydantic import PrivateAttr from policyengine_simulation_executor import artifact_keys +from policyengine_simulation_executor.precompute_models import ArtifactOutcome from policyengine_simulation_executor.stage_reporting_simulation import ( StageReportingSimulation, ) if TYPE_CHECKING: + from policyengine.core import Simulation from policyengine.core.scoping_strategy import ScopingStrategy logger = logging.getLogger(__name__) @@ -43,9 +45,23 @@ # "incomplete": loaded but missing requested columns or a valid SPM receipt # -> recomputed; # "miss": no artifact -> computed (today's behavior). -OUTCOME_HIT = "hit" -OUTCOME_INCOMPLETE = "incomplete" -OUTCOME_MISS = "miss" +OUTCOME_HIT: Final = "hit" +OUTCOME_INCOMPLETE: Final = "incomplete" +OUTCOME_MISS: Final = "miss" + + +@runtime_checkable +class EntityVariableResolver(Protocol): + """A model version that lists the output variables a simulation must carry. + + policyengine defines ``resolve_entity_variables`` on + MicrosimulationModelVersion, not on the TaxBenefitModelVersion a + Simulation declares. + """ + + def resolve_entity_variables(self, simulation: Simulation) -> dict[str, list[str]]: + """Return each entity's required output variables.""" + ... @dataclass(frozen=True) @@ -154,10 +170,10 @@ class ArtifactBaselineSimulation(StageReportingSimulation): """A Simulation whose ``ensure()`` validates what it loads.""" _computed_this_process: bool = PrivateAttr(default=False) - _artifact_outcome: str | None = PrivateAttr(default=None) + _artifact_outcome: ArtifactOutcome | None = PrivateAttr(default=None) @property - def artifact_outcome(self) -> str | None: + def artifact_outcome(self) -> ArtifactOutcome | None: """Set by ``ensure()``; None until then.""" return self._artifact_outcome @@ -166,7 +182,7 @@ def run(self) -> None: self._computed_this_process = True super().run() - def _record_outcome(self, outcome: str) -> None: + def _record_outcome(self, outcome: ArtifactOutcome) -> None: """First outcome wins for the life of this instance. Every request ensures its baseline more than once (the economic @@ -247,7 +263,11 @@ def ensure(self) -> None: def _missing_output_columns(self) -> list[tuple[str, str]]: data = getattr(self.output_dataset, "data", None) frames = getattr(data, "entity_data", None) - resolved = self.tax_benefit_model_version.resolve_entity_variables(self) + model_version = self.tax_benefit_model_version + if not isinstance(model_version, EntityVariableResolver): + msg = "the baseline's model version can't list its output variables" + raise TypeError(msg) + resolved = model_version.resolve_entity_variables(self) if not frames: return [(entity, "") for entity in resolved] missing: list[tuple[str, str]] = [] diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/output_frames.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/output_frames.py new file mode 100644 index 000000000..81d63116c --- /dev/null +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/output_frames.py @@ -0,0 +1,46 @@ +"""The shapes of a simulation's output data, as types code can check. + +policyengine types ``Dataset.data`` as an optional ``BaseModel``, but each +country's year data holds weighted entity frames. Code that reads them checks +the shape first, for example:: + + if not isinstance(simulation.output_dataset.data, EntityFrames): + raise TypeError(...) + frames = simulation.output_dataset.data.entity_data + +Pyright then knows the type of that same expression. Reading through the +expression itself, rather than a helper's return value, keeps it one the +microdata weight checks recognise as a weighted frame. + +The protocols are runtime-checkable, so the check only needs the attribute: +real year data and the tests' stand-ins both pass. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Protocol, runtime_checkable + +if TYPE_CHECKING: + from collections.abc import Mapping + + import pandas as pd + + +@runtime_checkable +class EntityFrames(Protocol): + """Year data that holds each entity's weighted frame, by entity name.""" + + @property + def entity_data(self) -> Mapping[str, pd.DataFrame]: + """Return each entity's weighted frame (a microdf MicroDataFrame).""" + ... + + +@runtime_checkable +class HouseholdFrame(Protocol): + """Year data with a weighted household frame.""" + + @property + def household(self) -> pd.DataFrame: + """Return the weighted household frame (a microdf MicroDataFrame).""" + ... 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 0e87f9dc8..50d0561f6 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py @@ -44,6 +44,7 @@ BaselineArtifactIdentity, baseline_artifact_filename, ) +from policyengine_simulation_executor.output_frames import EntityFrames from policyengine_simulation_executor.precompute_models import ( ArtifactManifest, BaselineComputeResult, @@ -284,7 +285,9 @@ def build_manifest(plan: PrecomputePlan) -> ArtifactManifest: for entry in plan.state_datasets ] return ArtifactManifest( - manifest_schema=MANIFEST_SCHEMA, + # The field's alias, which Pydantic accepts (and which its type + # checker support knows) at construction. + schema=MANIFEST_SCHEMA, country=PRECOMPUTE_COUNTRY, receipt=plan.receipt, artifacts=artifacts, @@ -760,10 +763,13 @@ def verify_determinism_impl( if not artifact_file.exists(): store.download_file(expected.path, artifact_file) baseline.ensure() - if baseline.artifact_outcome != OUTCOME_HIT: + # A cohort baseline is an ArtifactBaselineSimulation, which records the + # outcome; the attribute isn't on the Simulation type it is declared as. + outcome = getattr(baseline, "artifact_outcome", None) + if outcome != OUTCOME_HIT: msg = ( "Determinism gate could not load the artifact it verifies " - f"(outcome={baseline.artifact_outcome})" + f"(outcome={outcome})" ) raise RuntimeError(msg) @@ -778,6 +784,15 @@ def verify_determinism_impl( ) fresh.run() + missing_frames = "Determinism gate found a simulation with no output entity frames" + if baseline.output_dataset is None or not isinstance( + baseline.output_dataset.data, EntityFrames + ): + raise TypeError(missing_frames) + if fresh.output_dataset is None or not isinstance( + fresh.output_dataset.data, EntityFrames + ): + raise TypeError(missing_frames) differences = [] loaded = baseline.output_dataset.data.entity_data recomputed = fresh.output_dataset.data.entity_data 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 063eeed1b..d2acb0ba6 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 @@ -20,6 +20,8 @@ import pandas as pd +from policyengine_simulation_executor.output_frames import EntityFrames + if TYPE_CHECKING: from policyengine.core import Simulation @@ -45,6 +47,11 @@ def dump( ) -> tuple[dict[str, dict[str, list]], dict[str, dict[str, str]]]: columns: dict[str, dict[str, list]] = {} dtypes: dict[str, dict[str, str]] = {} + if simulation.output_dataset is None or not isinstance( + simulation.output_dataset.data, EntityFrames + ): + msg = "the simulation has no output entity frames; run it first" + raise TypeError(msg) for entity, frame in simulation.output_dataset.data.entity_data.items(): df = pd.DataFrame(frame) columns[entity] = df.to_dict("list") 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 5a86a0482..9e851922c 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 @@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, Any from policyengine_simulation_contract.job_progress import JobStage +from policyengine_simulation_observability import scopes from policyengine_simulation_observability.stages import ( ANNUAL_IMPACT_STAGES, Stage, @@ -74,7 +75,11 @@ def __post_init__(self) -> None: self.country = self.country.lower() def _span(self, name: str) -> AbstractContextManager[object]: - return self.runtime.span(name) if self.runtime is not None else nullcontext() + return ( + scopes.span(self.runtime, name) + if self.runtime is not None + else nullcontext() + ) @property def analysis(self) -> object: 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 4ce0d2feb..8bc282fab 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 @@ -11,6 +11,8 @@ import pandas as pd from policyengine_simulation_contract.spm import spm_error_detail +from policyengine_simulation_executor.output_frames import HouseholdFrame + if TYPE_CHECKING: from _typeshed import ConvertibleToFloat from policyengine.core import Simulation @@ -50,8 +52,9 @@ def _collection_records(collection: object) -> list[dict[str, Any]]: def _output_model_dump(value: object) -> object: if value is None: return None - if hasattr(value, "model_dump"): - return value.model_dump(mode="json") + model_dump = getattr(value, "model_dump", None) + if model_dump is not None: + return model_dump(mode="json") if isinstance(value, Mapping): return dict(value) return None @@ -60,6 +63,9 @@ def _output_model_dump(value: object) -> object: def _entity_data(simulation: Simulation, entity: str) -> pd.DataFrame: if simulation.output_dataset is None or simulation.output_dataset.data is None: simulation.ensure() + if simulation.output_dataset is None or simulation.output_dataset.data is None: + msg = "the simulation has no output data after ensure()" + raise TypeError(msg) return getattr(simulation.output_dataset.data, entity) @@ -91,6 +97,11 @@ def _unweighted_household_frame(simulation: Simulation) -> pd.DataFrame: shrink. It goes once the breakdown hands policyengine the simulations, as the constituency and US congressional district breakdowns already do. """ + if simulation.output_dataset is None or not isinstance( + simulation.output_dataset.data, HouseholdFrame + ): + msg = "the simulation has no household output frame; run it first" + raise TypeError(msg) return pd.DataFrame(simulation.output_dataset.data.household) @@ -132,7 +143,11 @@ def _warn_on_mismatched_weights( def _sum_output_variable(simulation: Simulation, variable: str, entity: str) -> float: data = _entity_data(simulation, entity) if variable in data.columns: - return float(data[variable].sum()) + column = data[variable] + if not isinstance(column, pd.Series): + msg = f"{entity} output has more than one {variable!r} column" + raise TypeError(msg) + return float(column.sum()) from policyengine.outputs import ( # noqa: PLC0415 # policyengine stays out of this module's import-time imports Aggregate, @@ -146,6 +161,9 @@ def _sum_output_variable(simulation: Simulation, variable: str, entity: str) -> aggregate_type=AggregateType.SUM, ) output.run() + if output.result is None: + msg = f"Aggregate of {entity}.{variable} produced no result" + raise RuntimeError(msg) return float(output.result) @@ -156,7 +174,14 @@ def _change_output_variable( reform_data = _entity_data(reform, entity) _warn_on_mismatched_weights(baseline_data, reform_data, entity) if variable in baseline_data.columns and variable in reform_data.columns: - return float((reform_data[variable] - baseline_data[variable]).sum()) + reform_column = reform_data[variable] + baseline_column = baseline_data[variable] + if not isinstance(reform_column, pd.Series) or not isinstance( + baseline_column, pd.Series + ): + msg = f"{entity} output has more than one {variable!r} column" + raise TypeError(msg) + return float((reform_column - baseline_column).sum()) from policyengine.outputs import ( # noqa: PLC0415 # policyengine stays out of this module's import-time imports ChangeAggregate, @@ -171,6 +196,9 @@ def _change_output_variable( aggregate_type=ChangeAggregateType.SUM, ) output.run() + if output.result is None: + msg = f"ChangeAggregate of {entity}.{variable} produced no result" + raise RuntimeError(msg) return float(output.result) 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 849c6c7a1..b19a555b2 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 @@ -10,6 +10,7 @@ CONSTITUENCY_ASSET_SPEC, ) +from policyengine_simulation_executor.output_frames import HouseholdFrame from policyengine_simulation_executor.simulation_macro_output import ( CongressionalDistrictImpactOutput, CongressionalDistrictImpactRecord, @@ -264,6 +265,11 @@ def build_uk_constituency_impact( """ if country != "uk": return None + if baseline.output_dataset is None or not isinstance( + baseline.output_dataset.data, HouseholdFrame + ): + msg = "the baseline has no household output frame; run it first" + raise TypeError(msg) _require_uk_area_column( baseline.output_dataset.data.household, "constituency_code_oa", "constituency" ) 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 154f13002..128ec1716 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 @@ -35,6 +35,7 @@ NULL_JOB_PROGRESS, JobProgressReporter, ) +from policyengine_simulation_executor.output_frames import EntityFrames from policyengine_simulation_executor.release_bundle import ( get_country_release_bundle, resolve_bundle_dataset_name, @@ -133,7 +134,7 @@ def setup_gcp_credentials( with contextlib.ExitStack() as credentials_file_stack: try: credential_span = ( - runtime.span(ANNUAL_IMPACT_STAGES.name(Stage.CREDENTIAL_SETUP)) + scopes.span(runtime, ANNUAL_IMPACT_STAGES.name(Stage.CREDENTIAL_SETUP)) if runtime is not None else contextlib.nullcontext() ) @@ -290,6 +291,9 @@ def _compile_interval_policy( ) compiled = compile_reform(policy, year=year, model_version=country_module.model) + if compiled is None: + # compile_reform returns None only for an empty policy, returned above. + return policy parameter_values = [] for path, periods in compiled.items(): dated_values = [ @@ -363,7 +367,9 @@ def _build_uk_weight_replacement_region(region_code: str) -> Region | None: return None prefix, value = region_code.split("/", maxsplit=1) - if prefix not in {"constituency", "local_authority"}: + # A tuple, unlike a set display, keeps the literal types, so the check + # narrows prefix to the two UK region types it allows. + if prefix not in ("constituency", "local_authority"): return None from policyengine.core.region import ( # noqa: PLC0415 # this module keeps policyengine out of its module-level imports (see its docstring) @@ -509,6 +515,9 @@ def release(reference: ArtifactPathReference) -> tuple[str, str]: msg = f"{matrix_name} entry {year} is not a weight matrix" raise TypeError(msg) matrix_households = weights.shape[-1] + if not isinstance(dataset.data, EntityFrames): + msg = f"the {dataset_selection.name!r} dataset has no entity data loaded" + raise TypeError(msg) households = len(pd.DataFrame(dataset.data.entity_data["household"])) if households != matrix_households: msg = ( @@ -768,6 +777,11 @@ def _build_simulation( policy: dict[str, Any] | None, region: RegionResolution, ) -> StageReportingSimulation: + """Build a simulation that reports when its calculation starts. + + A run with a deterministic baseline id gets an ArtifactBaselineSimulation, + which is a StageReportingSimulation that can load a precomputed baseline. + """ from policyengine_simulation_executor.baseline_artifacts import ( ArtifactBaselineSimulation, BaselineIdentityRequest, @@ -785,7 +799,7 @@ def _build_simulation( spm_kwargs: dict[str, Any] = {"spm": selection} if selection is not None else {} country = params.get("country", "us") country_module = _country_module(country) - policy = _compile_interval_policy( + runtime_policy = _compile_interval_policy( policy, country_module=country_module, year=_parse_year(params) ) scoping_strategy = region.scoping_strategy @@ -794,7 +808,7 @@ def _build_simulation( params, country=country, dataset_is_default=dataset_selection.is_default, - policy=policy, + policy=runtime_policy, region_code=region.code, scoping_strategy=scoping_strategy, year=_parse_year(params), @@ -809,14 +823,14 @@ def _build_simulation( id=simulation_id, dataset=dataset, tax_benefit_model_version=country_module.model, - policy=policy, + policy=runtime_policy, scoping_strategy=scoping_strategy, ) return StageReportingSimulation( **spm_kwargs, dataset=dataset, tax_benefit_model_version=country_module.model, - policy=policy, + policy=runtime_policy, scoping_strategy=scoping_strategy, ) 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 b5c2490a7..b4b4ade8e 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/spm.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/spm.py @@ -103,7 +103,7 @@ def _certified_us_spm_capability(configured: object) -> SPMCapability: ): msg = "The installed US model does not support canonical SPM" raise ValueError(msg) - capability = SPMCapability(defaults=configured) + capability = SPMCapability.model_validate({"defaults": configured}) forecast = _forecast(capability.defaults.forecast_content_sha256) forecast.entry( forecast.years[0], 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 2e68cebb2..a4da63563 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 @@ -5,6 +5,7 @@ import math from collections.abc import Mapping, Sequence from dataclasses import dataclass +from typing import TypeGuard import numpy as np import pandas as pd @@ -100,10 +101,17 @@ def _compare_entity_frames( raise ParityMismatchError(error_code, entity) for column in expected_frame.columns: path = f"{entity}.{column}" + expected_values = expected_frame[column] + actual_values = actual_frame[column] + if not isinstance(expected_values, pd.Series) or not isinstance( + actual_values, pd.Series + ): + msg = f"{entity} has more than one {column!r} column" + raise TypeError(msg) _compare_column( path, - expected_frame[column], - actual_frame[column], + expected_values, + actual_values, tolerances.get(path, NumericalTolerance()), ) @@ -148,14 +156,14 @@ def compare_aggregate_reports( _compare_report_values(incumbent, candidate, "", tolerances or {}) -def _is_report_list(value: object) -> bool: +def _is_report_list(value: object) -> TypeGuard[Sequence[object]]: """Whether a report value is a list-like sequence rather than text or bytes.""" return isinstance(value, Sequence) and not isinstance( value, (str, bytes, bytearray) ) -def _is_report_number(value: object) -> bool: +def _is_report_number(value: object) -> TypeGuard[int | float]: """Whether a report value is an int or float; booleans are not numbers here.""" return isinstance(value, (int, float)) and not isinstance(value, bool) 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 fefc8281b..7524c2525 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 @@ -13,6 +13,7 @@ ResultComparisonStatus, Stage12InvocationContext, ) +from policyengine_simulation_observability import scopes from policyengine_simulation_observability.stages import Stage, StagePlan from policyengine_simulation_executor.stage12_result_comparison import compare_results @@ -59,7 +60,7 @@ def compare_completed_report( return try: comparison_span = ( - runtime.span(stage_plan.name(Stage.STAGE12_RESULT_COMPARISON)) + scopes.span(runtime, stage_plan.name(Stage.STAGE12_RESULT_COMPARISON)) if runtime is not None and stage_plan is not None else nullcontext() ) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/failure_causes.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/failure_causes.py index 27aac2e99..8a60a4d19 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/failure_causes.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/failure_causes.py @@ -35,6 +35,7 @@ from urllib.parse import urlsplit from policyengine_simulation_contract.stage12_execution import ErrorSummary +from policyengine_simulation_observability import scopes from pydantic import TypeAdapter, ValidationError from sqlalchemy.exc import StatementError @@ -186,9 +187,9 @@ def tracked_stage( if runtime is None: scope = nullcontext() elif operation: - scope = runtime.operation(stage_name, attributes=attributes) + scope = scopes.operation(runtime, stage_name, attributes=attributes) else: - scope = runtime.span(stage_name, attributes=attributes) + scope = scopes.span(runtime, stage_name, attributes=attributes) with ( tracking.progress.stage(stage_name, role=_simulation_role(attributes)), scope, diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/fanout.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/fanout.py index 83194b9f5..147c5549a 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/fanout.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/fanout.py @@ -8,6 +8,7 @@ from typing import TYPE_CHECKING, Any import modal +from policyengine_simulation_observability import scopes from policyengine_simulation_observability.stages import ( STAGE12_SIMULATION_STAGES, Stage, @@ -131,7 +132,8 @@ def run(self) -> SimulationCalculation: try: for segment in segments: call = None - with self.runtime.span( + with scopes.span( + self.runtime, STAGE12_SIMULATION_STAGES.name(Stage.STAGE12_SEGMENT_DISPATCH), attributes={"segment_index": segment.segment_index}, ): @@ -145,16 +147,16 @@ def run(self) -> SimulationCalculation: raise payloads: list[object] | None = None - with self.runtime.span( - STAGE12_SIMULATION_STAGES.name(Stage.STAGE12_SEGMENT_WAIT) + with scopes.span( + self.runtime, STAGE12_SIMULATION_STAGES.name(Stage.STAGE12_SEGMENT_WAIT) ): payloads = self._collect(handles) if payloads is None: msg = "Stage 12 segment wait returned no payloads" raise RuntimeError(msg) merged: SimulationCalculation | None = None - with self.runtime.span( - STAGE12_SIMULATION_STAGES.name(Stage.STAGE12_SEGMENT_MERGE) + with scopes.span( + self.runtime, STAGE12_SIMULATION_STAGES.name(Stage.STAGE12_SEGMENT_MERGE) ): merged = merge_segment_results( self.simulation, 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 b1a5695c9..7be2fc154 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 @@ -115,17 +115,24 @@ def build_stage12_simulation( scoping_strategy=resolved.region.scoping_strategy, spm=selection, ) - simulation_type = Stage12CachedSimulation if identity is not None else Simulation - model = simulation_type( - **({"spm": selection} if selection is not None else {}), - **({"id": identity.simulation_id} if identity is not None else {}), + spm_kwargs: dict[str, Any] = {"spm": selection} if selection is not None else {} + if identity is None: + return Simulation( + **spm_kwargs, + dataset=resolved.dataset, + tax_benefit_model_version=resolved.country_module.model, + policy=policy, + scoping_strategy=resolved.region.scoping_strategy, + ) + model = Stage12CachedSimulation( + **spm_kwargs, + id=identity.simulation_id, dataset=resolved.dataset, tax_benefit_model_version=resolved.country_module.model, policy=policy, scoping_strategy=resolved.region.scoping_strategy, ) - if isinstance(model, Stage12CachedSimulation): - model.configure_stage12_cache(execution.output_plan) + model.configure_stage12_cache(execution.output_plan) return 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 72ede3566..dba29ed02 100644 --- a/projects/policyengine-simulation-executor/tests/test_budget_window_batch.py +++ b/projects/policyengine-simulation-executor/tests/test_budget_window_batch.py @@ -13,6 +13,7 @@ PolicyEngineBundle, ) from src.modal.budget_window_batch import run_budget_window_batch_impl +from src.modal.budget_window_context import ChildSimulationHandle from src.modal.fanout import PollBackoff @@ -315,6 +316,35 @@ def test_run_budget_window_batch_impl_marks_failure( assert state.child_jobs["2027"].status == "cancelled" +def test_a_child_without_a_call_handle_fails_the_batch( + monkeypatch, mock_batch_modal, observability_runtime +): + """A missing call handle fails the batch like any other child poll error.""" + request, payload = _build_parent_payload(window_size=1) + _seed_parent_batch(request, mock_batch_modal["parent_call_id"]) + run_simulation = MockRunSimulationFunction( + tracker=SpawnTracker(), + results_by_year={"2026": [TimeoutError()]}, + call_registry=mock_batch_modal["call_registry"], + ) + mock_batch_modal["functions"][ + ("policyengine-simulation-py4-10-0", "run_simulation") + ] = run_simulation + monkeypatch.setattr( + scheduler_module.BudgetWindowBatchRunner, + "resolve_child_handle", + lambda _runner, simulation_year: ChildSimulationHandle( + simulation_year=simulation_year, job_id="child-2026", call=None + ), + ) + + result = run_budget_window_batch_impl(payload, runtime=observability_runtime) + + assert result["status"] == "failed" + assert result["failed_years"] == ["2026"] + assert result["error"].startswith("Simulation failed") + + def test_scheduler_sleep_exponentially_backs_off_then_resets_on_progress( monkeypatch, mock_batch_modal, observability_runtime ): diff --git a/projects/policyengine-simulation-executor/tests/test_output_frames.py b/projects/policyengine-simulation-executor/tests/test_output_frames.py new file mode 100644 index 000000000..84fdfc3fb --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_output_frames.py @@ -0,0 +1,96 @@ +"""The checkable shapes of a simulation's output data.""" + +from types import SimpleNamespace + +import pandas as pd +import pytest +from microdf import MicroDataFrame +from policyengine.tax_benefit_models.common.model_version import ( + MicrosimulationModelVersion, +) +from policyengine.tax_benefit_models.uk.datasets import UKYearData +from policyengine.tax_benefit_models.us.datasets import USYearData +from policyengine_simulation_executor.baseline_artifacts import EntityVariableResolver +from policyengine_simulation_executor.output_frames import EntityFrames, HouseholdFrame +from policyengine_simulation_executor.simulation_output_common import _entity_data + + +def _frame(entity: str) -> MicroDataFrame: + return MicroDataFrame( + pd.DataFrame({f"{entity}_id": [1], f"{entity}_weight": [2.0]}), + weights=f"{entity}_weight", + ) + + +def _uk_year_data() -> UKYearData: + return UKYearData( + person=_frame("person"), + benunit=_frame("benunit"), + household=_frame("household"), + ) + + +def _us_year_data() -> USYearData: + return USYearData( + **{ + entity: _frame(entity) + for entity in ( + "person", + "marital_unit", + "family", + "spm_unit", + "tax_unit", + "household", + ) + } + ) + + +@pytest.mark.parametrize( + "year_data_factory", [_uk_year_data, _us_year_data], ids=["UK", "US"] +) +def test_country_year_data_has_entity_frames_and_a_household_frame( + year_data_factory, +): + year_data = year_data_factory() + + assert isinstance(year_data, EntityFrames) + assert isinstance(year_data, HouseholdFrame) + + +def test_country_model_versions_resolve_entity_variables(): + # policyengine's US and UK model versions both derive from this class. + assert issubclass(MicrosimulationModelVersion, EntityVariableResolver) + + +def test_data_with_frames_by_attribute_matches_the_shapes(): + assert isinstance(SimpleNamespace(entity_data={}), EntityFrames) + assert isinstance(SimpleNamespace(household=None), HouseholdFrame) + + +def test_missing_output_data_matches_neither_shape(): + assert not isinstance(None, EntityFrames) + assert not isinstance(None, HouseholdFrame) + assert not isinstance(SimpleNamespace(person=None), HouseholdFrame) + + +class _SimulationWithoutOutput: + """A simulation whose ensure() leaves the output it was given.""" + + def __init__(self, output_dataset): + self.output_dataset = output_dataset + + def ensure(self): + pass + + +@pytest.mark.parametrize( + "output_dataset", + [None, SimpleNamespace(data=None)], + ids=["no output dataset", "no output data"], +) +def test_reading_an_entity_frame_without_output_data_raises_type_error( + output_dataset, +): + with pytest.raises(TypeError, match="no output data"): + _entity_data(_SimulationWithoutOutput(output_dataset), "household") diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_failure_causes.py b/projects/policyengine-simulation-executor/tests/test_stage12_failure_causes.py index fb651aa53..b908bbc84 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_failure_causes.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_failure_causes.py @@ -491,7 +491,10 @@ def span(self, name, *, attributes=None): yield @contextmanager - def operation(self, name, *, attributes=None): + def operation(self, name, *, attributes=None, remote_context=None): + # Mirrors ObservabilityRuntime.operation; a Stage 12 stage continues + # no remote trace. + assert remote_context is None self.scopes.append(("operation", name, attributes)) yield 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 7242afaa2..0bc66cca7 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 @@ -172,3 +172,22 @@ def test_guard_rejects_a_year_the_matrix_does_not_cover(monkeypatch, tmp_path) - with pytest.raises(ValueError, match=r"has no weights for 2026 \(it covers 2025\)"): sr._require_certified_uk_weight_matrix(_STRATEGY, _dataset(3, 2026), _SELECTION) + + +@pytest.mark.parametrize( + "unloaded_data", + [None, SimpleNamespace(tables={})], + ids=["no-data", "no-entity-data"], +) +def test_guard_rejects_a_dataset_with_no_entity_data( + monkeypatch, tmp_path, unloaded_data +) -> None: + _bundle(monkeypatch) + _assets(monkeypatch, tmp_path) + dataset = SimpleNamespace(year=2025, data=unloaded_data) + + with pytest.raises( + TypeError, + match=r"the 'enhanced_frs_2024_25' dataset has no entity data loaded", + ): + sr._require_certified_uk_weight_matrix(_STRATEGY, dataset, _SELECTION) diff --git a/rules/baselines/pyright-executor.txt b/rules/baselines/pyright-executor.txt deleted file mode 100644 index 77651b94b..000000000 --- a/rules/baselines/pyright-executor.txt +++ /dev/null @@ -1,86 +0,0 @@ -# Tolerated Pyright errors in projects/policyengine-simulation-executor. -# This list may only shrink. -# -# Rule sim.quality.pyright-executor. Each entry is -# :: :: , one line per -# error, sorted. After fixing errors, remove the entries that no longer -# occur with -# python -m scripts.rules pyright projects/policyengine-simulation-executor --write-baseline -# then set the rule's [rule.baseline] count in rules/registry.toml to match. -projects/policyengine-simulation-executor/src/modal/app.py :: reportReturnType :: Function with declared return type "dict[Unknown, Unknown]" must return value on all code paths -projects/policyengine-simulation-executor/src/modal/app.py :: reportReturnType :: Function with declared return type "dict[Unknown, Unknown]" must return value on all code paths -projects/policyengine-simulation-executor/src/modal/app.py :: reportReturnType :: Function with declared return type "dict[Unknown, Unknown]" must return value on all code paths -projects/policyengine-simulation-executor/src/modal/budget_window_batch.py :: reportPossiblyUnboundVariable :: "context" is possibly unbound -projects/policyengine-simulation-executor/src/modal/budget_window_results.py :: reportArgumentType :: Argument of type "float" cannot be assigned to parameter "year" of type "Literal['Total']" in function "__init__" -projects/policyengine-simulation-executor/src/modal/budget_window_scheduler.py :: reportOptionalMemberAccess :: "get" is not a known attribute of "None" -projects/policyengine-simulation-executor/src/modal/budget_window_scheduler.py :: reportPossiblyUnboundVariable :: "annual_impact" is possibly unbound -projects/policyengine-simulation-executor/src/modal/budget_window_scheduler.py :: reportPossiblyUnboundVariable :: "call" is possibly unbound -projects/policyengine-simulation-executor/src/modal/budget_window_scheduler.py :: reportPossiblyUnboundVariable :: "call" is possibly unbound -projects/policyengine-simulation-executor/src/modal/budget_window_scheduler.py :: reportPossiblyUnboundVariable :: "call" is possibly unbound -projects/policyengine-simulation-executor/src/modal/budget_window_scheduler.py :: reportPossiblyUnboundVariable :: "child_request" is possibly unbound -projects/policyengine-simulation-executor/src/modal/budget_window_scheduler.py :: reportPossiblyUnboundVariable :: "result" is possibly unbound -projects/policyengine-simulation-executor/src/modal/budget_window_scheduler.py :: reportPossiblyUnboundVariable :: "state" is possibly unbound -projects/policyengine-simulation-executor/src/modal/budget_window_scheduler.py :: reportPossiblyUnboundVariable :: "state" is possibly unbound -projects/policyengine-simulation-executor/src/modal/budget_window_scheduler.py :: reportReturnType :: Function with declared return type "dict[str, Any]" must return value on all code paths -projects/policyengine-simulation-executor/src/modal/segmented_national.py :: reportPossiblyUnboundVariable :: "output" is possibly unbound -projects/policyengine-simulation-executor/src/modal/segmented_national.py :: reportPossiblyUnboundVariable :: "output" is possibly unbound -projects/policyengine-simulation-executor/src/modal/segmented_national.py :: reportPossiblyUnboundVariable :: "output" is possibly unbound -projects/policyengine-simulation-executor/src/modal/segmented_national.py :: reportPossiblyUnboundVariable :: "output" is possibly unbound -projects/policyengine-simulation-executor/src/modal/utils/release_sentinel.py :: reportOptionalMemberAccess :: "split" is not a known attribute of "None" -projects/policyengine-simulation-executor/src/modal/utils/update_version_registry.py :: reportArgumentType :: Argument of type "dict[Unknown, Unknown]" cannot be assigned to parameter "value" of type "BundleManifestMetadata" in function "__setitem__" -projects/policyengine-simulation-executor/src/modal/utils/update_version_registry.py :: reportGeneralTypeIssues :: Could not assign item in TypedDict -projects/policyengine-simulation-executor/src/modal/utils/update_version_registry.py :: reportIndexIssue :: "__getitem__" method not defined on type "object" -projects/policyengine-simulation-executor/src/modal/utils/update_version_registry.py :: reportIndexIssue :: "__setitem__" method not defined on type "object" -projects/policyengine-simulation-executor/src/modal/utils/update_version_registry.py :: reportOptionalMemberAccess :: "get" is not a known attribute of "None" -projects/policyengine-simulation-executor/src/modal/utils/update_version_registry.py :: reportOptionalMemberAccess :: "get" is not a known attribute of "None" -projects/policyengine-simulation-executor/src/modal/utils/update_version_registry.py :: reportOptionalMemberAccess :: "get" is not a known attribute of "None" -projects/policyengine-simulation-executor/src/modal/utils/update_version_registry.py :: reportReturnType :: Type "dict[str, int | str | RouteMaps | dict[str, BundleManifestMetadata]]" is not assignable to return type "RoutingState" -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/baseline_artifacts.py :: reportAttributeAccessIssue :: Cannot access attribute "resolve_entity_variables" for class "TaxBenefitModelVersion" -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py :: reportArgumentType :: Argument of type "str | None" cannot be assigned to parameter "outcome" of type "ArtifactOutcome | None" in function "__init__" -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py :: reportAttributeAccessIssue :: Cannot access attribute "artifact_outcome" for class "Simulation" -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py :: reportAttributeAccessIssue :: Cannot access attribute "artifact_outcome" for class "Simulation" -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py :: reportAttributeAccessIssue :: Cannot access attribute "entity_data" for class "BaseModel" -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py :: reportAttributeAccessIssue :: Cannot access attribute "entity_data" for class "BaseModel" -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py :: reportCallIssue :: Argument missing for parameter "schema" -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py :: reportCallIssue :: No parameter named "manifest_schema" -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py :: reportOptionalMemberAccess :: "data" is not a known attribute of "None" -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py :: reportOptionalMemberAccess :: "data" is not a known attribute of "None" -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py :: reportOptionalMemberAccess :: "entity_data" is not a known attribute of "None" -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py :: reportOptionalMemberAccess :: "entity_data" is not a known attribute of "None" -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_microdata.py :: reportAttributeAccessIssue :: Cannot access attribute "entity_data" for class "BaseModel" -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_microdata.py :: reportOptionalMemberAccess :: "data" is not a known attribute of "None" -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_microdata.py :: reportOptionalMemberAccess :: "entity_data" is not a known attribute of "None" -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_common.py :: reportArgumentType :: Argument of type "Any | None" cannot be assigned to parameter "x" of type "ConvertibleToFloat" in function "__new__" -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_common.py :: reportArgumentType :: Argument of type "Any | None" cannot be assigned to parameter "x" of type "ConvertibleToFloat" in function "__new__" -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_common.py :: reportArgumentType :: Argument of type "Any | None" cannot be assigned to parameter "x" of type "ConvertibleToFloat" in function "__new__" -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_common.py :: reportArgumentType :: Argument of type "Any | None" cannot be assigned to parameter "x" of type "ConvertibleToFloat" in function "__new__" -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_common.py :: reportArgumentType :: Argument of type "Unknown | Any | Series" cannot be assigned to parameter "x" of type "ConvertibleToFloat" in function "__new__" -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_common.py :: reportArgumentType :: Argument of type "Unknown | Any | Series" cannot be assigned to parameter "x" of type "ConvertibleToFloat" in function "__new__" -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_common.py :: reportAttributeAccessIssue :: Cannot access attribute "household" for class "BaseModel" -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_common.py :: reportAttributeAccessIssue :: Cannot access attribute "model_dump" for class "object" -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_common.py :: reportOptionalMemberAccess :: "data" is not a known attribute of "None" -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_common.py :: reportOptionalMemberAccess :: "data" is not a known attribute of "None" -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_common.py :: reportOptionalMemberAccess :: "household" is not a known attribute of "None" -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_geographic.py :: reportAttributeAccessIssue :: Cannot access attribute "household" for class "BaseModel" -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_geographic.py :: reportOptionalMemberAccess :: "data" is not a known attribute of "None" -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_geographic.py :: reportOptionalMemberAccess :: "household" is not a known attribute of "None" -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py :: reportArgumentType :: Argument of type "str" cannot be assigned to parameter "region_type" of type "RegionType" in function "__init__" -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py :: reportAssignmentType :: Type "dict[str, Any] | Policy | None" is not assignable to declared type "dict[str, Any] | None" -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py :: reportAttributeAccessIssue :: Cannot access attribute "entity_data" for class "BaseModel" -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py :: reportOptionalMemberAccess :: "entity_data" is not a known attribute of "None" -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py :: reportOptionalMemberAccess :: "items" is not a known attribute of "None" -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/spm.py :: reportArgumentType :: Argument of type "object" cannot be assigned to parameter "defaults" of type "SPMSelection" in function "__init__" -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_parity.py :: reportArgumentType :: Argument of type "Series | Unknown | Self@DataFrame" cannot be assigned to parameter "actual_values" of type "Series" in function "_compare_column" -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_parity.py :: reportArgumentType :: Argument of type "Series | Unknown | Self@DataFrame" cannot be assigned to parameter "expected_values" of type "Series" in function "_compare_column" -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_parity.py :: reportArgumentType :: Argument of type "object" cannot be assigned to parameter "expected" of type "Sequence[Unknown]" in function "_compare_report_lists" -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_parity.py :: reportArgumentType :: Argument of type "object" cannot be assigned to parameter "obj" of type "Sized" in function "len" -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_parity.py :: reportArgumentType :: Argument of type "object" cannot be assigned to parameter "x" of type "ConvertibleToFloat" in function "__new__" -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_parity.py :: reportArgumentType :: Argument of type "object" cannot be assigned to parameter "x" of type "ConvertibleToFloat" in function "__new__" -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_parity.py :: reportIndexIssue :: "__getitem__" method not defined on type "object" -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py :: reportArgumentType :: Argument of type "str" cannot be assigned to parameter "created_at" of type "datetime" in function "__init__" -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py :: reportArgumentType :: Argument of type "str" cannot be assigned to parameter "dynamic" of type "Dynamic | dict[str, Any] | None" in function "__init__" -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py :: reportArgumentType :: Argument of type "str" cannot be assigned to parameter "extra_variables" of type "dict[str, list[str]]" in function "__init__" -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py :: reportArgumentType :: Argument of type "str" cannot be assigned to parameter "output_dataset" of type "Dataset | None" in function "__init__" -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py :: reportArgumentType :: Argument of type "str" cannot be assigned to parameter "spm" of type "SPMSelection | None" in function "__init__" -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py :: reportArgumentType :: Argument of type "str" cannot be assigned to parameter "spm_receipt" of type "SPMProvenance | None" in function "__init__" -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py :: reportArgumentType :: Argument of type "str" cannot be assigned to parameter "updated_at" of type "datetime" in function "__init__" diff --git a/rules/baselines/pyright-gateway.txt b/rules/baselines/pyright-gateway.txt deleted file mode 100644 index b407bea2d..000000000 --- a/rules/baselines/pyright-gateway.txt +++ /dev/null @@ -1,9 +0,0 @@ -# Tolerated Pyright errors in projects/policyengine-simulation-gateway. -# This list may only shrink. -# -# Rule sim.quality.pyright-gateway. Each entry is -# :: :: , one line per -# error, sorted. After fixing errors, remove the entries that no longer -# occur with -# python -m scripts.rules pyright projects/policyengine-simulation-gateway --write-baseline -# then set the rule's [rule.baseline] count in rules/registry.toml to match. diff --git a/rules/fixtures/sim.quality.pyright-executor/bad/baseline.txt b/rules/fixtures/sim.quality.pyright-executor/bad/baseline.txt deleted file mode 100644 index 099021df0..000000000 --- a/rules/fixtures/sim.quality.pyright-executor/bad/baseline.txt +++ /dev/null @@ -1,5 +0,0 @@ -# Proof fixture: the baseline omits the "Secret" error in pyright.json, and -# still lists an "Image" error that no longer occurs. -projects/policyengine-simulation-executor/src/modal/app.py :: reportAttributeAccessIssue :: "App" is not a known attribute of module "modal" -projects/policyengine-simulation-executor/src/modal/app.py :: reportAttributeAccessIssue :: "Image" is not a known attribute of module "modal" -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py :: reportPossiblyUnboundVariable :: "metadata" is possibly unbound diff --git a/rules/fixtures/sim.quality.pyright-executor/good/baseline.txt b/rules/fixtures/sim.quality.pyright-executor/good/baseline.txt deleted file mode 100644 index 541795061..000000000 --- a/rules/fixtures/sim.quality.pyright-executor/good/baseline.txt +++ /dev/null @@ -1,4 +0,0 @@ -# Proof fixture: the baseline lists exactly the errors in pyright.json. -projects/policyengine-simulation-executor/src/modal/app.py :: reportAttributeAccessIssue :: "App" is not a known attribute of module "modal" -projects/policyengine-simulation-executor/src/modal/app.py :: reportAttributeAccessIssue :: "Secret" is not a known attribute of module "modal" -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py :: reportPossiblyUnboundVariable :: "metadata" is possibly unbound diff --git a/rules/fixtures/sim.quality.pyright-executor/good/pyright.json b/rules/fixtures/sim.quality.pyright-executor/good/pyright.json index 668fe4ddb..4d419d935 100644 --- a/rules/fixtures/sim.quality.pyright-executor/good/pyright.json +++ b/rules/fixtures/sim.quality.pyright-executor/good/pyright.json @@ -1,58 +1,9 @@ { "version": "1.1.414", - "generalDiagnostics": [ - { - "file": "/repository/projects/policyengine-simulation-executor/src/modal/app.py", - "severity": "error", - "message": "\"App\" is not a known attribute of module \"modal\"", - "range": { - "start": { - "line": 118, - "character": 12 - }, - "end": { - "line": 118, - "character": 15 - } - }, - "rule": "reportAttributeAccessIssue" - }, - { - "file": "/repository/projects/policyengine-simulation-executor/src/modal/app.py", - "severity": "error", - "message": "\"Secret\" is not a known attribute of module \"modal\"", - "range": { - "start": { - "line": 122, - "character": 19 - }, - "end": { - "line": 122, - "character": 25 - } - }, - "rule": "reportAttributeAccessIssue" - }, - { - "file": "/repository/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py", - "severity": "error", - "message": "\"metadata\" is possibly unbound", - "range": { - "start": { - "line": 823, - "character": 15 - }, - "end": { - "line": 823, - "character": 23 - } - }, - "rule": "reportPossiblyUnboundVariable" - } - ], + "generalDiagnostics": [], "summary": { - "filesAnalyzed": 90, - "errorCount": 3, + "filesAnalyzed": 91, + "errorCount": 0, "warningCount": 0, "informationCount": 0 } diff --git a/rules/fixtures/sim.quality.pyright-gateway/bad/baseline.txt b/rules/fixtures/sim.quality.pyright-gateway/bad/baseline.txt deleted file mode 100644 index 673fa595f..000000000 --- a/rules/fixtures/sim.quality.pyright-gateway/bad/baseline.txt +++ /dev/null @@ -1,5 +0,0 @@ -# Proof fixture: pyright.json has two "route" errors and the baseline lists -# one, and the baseline still lists a "bundle" error that no longer occurs. -projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/endpoints.py :: reportPossiblyUnboundVariable :: "route" is possibly unbound -projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/endpoints.py :: reportPossiblyUnboundVariable :: "payload" is possibly unbound -projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/endpoints.py :: reportPossiblyUnboundVariable :: "bundle" is possibly unbound diff --git a/rules/fixtures/sim.quality.pyright-gateway/good/baseline.txt b/rules/fixtures/sim.quality.pyright-gateway/good/baseline.txt deleted file mode 100644 index 3a7b94654..000000000 --- a/rules/fixtures/sim.quality.pyright-gateway/good/baseline.txt +++ /dev/null @@ -1,4 +0,0 @@ -# Proof fixture: the baseline lists exactly the errors in pyright.json. -projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/endpoints.py :: reportPossiblyUnboundVariable :: "route" is possibly unbound -projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/endpoints.py :: reportPossiblyUnboundVariable :: "route" is possibly unbound -projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/endpoints.py :: reportPossiblyUnboundVariable :: "payload" is possibly unbound diff --git a/rules/fixtures/sim.quality.pyright-gateway/good/pyright.json b/rules/fixtures/sim.quality.pyright-gateway/good/pyright.json index a902297dd..cd2a79ab8 100644 --- a/rules/fixtures/sim.quality.pyright-gateway/good/pyright.json +++ b/rules/fixtures/sim.quality.pyright-gateway/good/pyright.json @@ -1,58 +1,9 @@ { "version": "1.1.414", - "generalDiagnostics": [ - { - "file": "/repository/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/endpoints.py", - "severity": "error", - "message": "\"route\" is possibly unbound", - "range": { - "start": { - "line": 922, - "character": 26 - }, - "end": { - "line": 922, - "character": 31 - } - }, - "rule": "reportPossiblyUnboundVariable" - }, - { - "file": "/repository/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/endpoints.py", - "severity": "error", - "message": "\"route\" is possibly unbound", - "range": { - "start": { - "line": 923, - "character": 25 - }, - "end": { - "line": 923, - "character": 30 - } - }, - "rule": "reportPossiblyUnboundVariable" - }, - { - "file": "/repository/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/endpoints.py", - "severity": "error", - "message": "\"payload\" is possibly unbound", - "range": { - "start": { - "line": 934, - "character": 8 - }, - "end": { - "line": 934, - "character": 15 - } - }, - "rule": "reportPossiblyUnboundVariable" - } - ], + "generalDiagnostics": [], "summary": { "filesAnalyzed": 9, - "errorCount": 3, + "errorCount": 0, "warningCount": 0, "informationCount": 0 } diff --git a/rules/registry.toml b/rules/registry.toml index 44257ba8c..2a33ccb11 100644 --- a/rules/registry.toml +++ b/rules/registry.toml @@ -1097,11 +1097,11 @@ invocation = "unittest discover -s scripts/rules/tests" [[rule]] id = "sim.quality.pyright-gateway" category = "quality" -statement = "The simulation gateway has no Pyright errors beyond those its baseline lists, and the baseline lists only errors that still occur, so it only shrinks; checked with the pinned Pyright version in the gateway's own locked environment." -rationale = "docs/engineering/skills/type-checking.md#baselines" +statement = "The simulation gateway has no Pyright errors, checked with the pinned Pyright version in the gateway's own locked environment." +rationale = "docs/engineering/skills/type-checking.md#rules" scope = ["projects/policyengine-simulation-gateway/src/**/*.py"] tier = "mechanical" -status = "ratcheted" +status = "enforced" autofix = false [rule.enforcer] kind = "pyright" @@ -1110,9 +1110,6 @@ projects = ["projects/policyengine-simulation-gateway"] workflow = ".github/workflows/pr.yml" job = "pyright" invocation = "python -m scripts.rules pyright" -[rule.baseline] -file = "rules/baselines/pyright-gateway.txt" -count = 0 [rule.proof] bad_fixture = "rules/fixtures/sim.quality.pyright-gateway/bad" good_fixture = "rules/fixtures/sim.quality.pyright-gateway/good" @@ -1124,11 +1121,11 @@ invocation = "unittest discover -s scripts/rules/tests" [[rule]] id = "sim.quality.pyright-executor" category = "quality" -statement = "The simulation executor has no Pyright errors beyond those its baseline lists, and the baseline lists only errors that still occur, so it only shrinks; checked with the pinned Pyright version in the executor's own locked environment." -rationale = "docs/engineering/skills/type-checking.md#baselines" +statement = "The simulation executor has no Pyright errors, checked with the pinned Pyright version in the executor's own locked environment." +rationale = "docs/engineering/skills/type-checking.md#rules" scope = ["projects/policyengine-simulation-executor/src/**/*.py"] tier = "mechanical" -status = "ratcheted" +status = "enforced" autofix = false [rule.enforcer] kind = "pyright" @@ -1137,9 +1134,6 @@ projects = ["projects/policyengine-simulation-executor"] workflow = ".github/workflows/pr.yml" job = "pyright" invocation = "python -m scripts.rules pyright" -[rule.baseline] -file = "rules/baselines/pyright-executor.txt" -count = 77 [rule.proof] bad_fixture = "rules/fixtures/sim.quality.pyright-executor/bad" good_fixture = "rules/fixtures/sim.quality.pyright-executor/good" diff --git a/scripts/rules/tests/test_pyright.py b/scripts/rules/tests/test_pyright.py index 378de485d..7fa74dc8c 100644 --- a/scripts/rules/tests/test_pyright.py +++ b/scripts/rules/tests/test_pyright.py @@ -190,35 +190,37 @@ def test_rule_proof_pyright_libraries_and_entry(self) -> None: def test_rule_proof_pyright_gateway(self) -> None: rule_id = "sim.quality.pyright-gateway" - # The baseline lists one of two "route" errors, and an error that no - # longer occurs. + # Three of the gateway's 81 real errors before its fixes: variables + # assigned inside a span, which Pyright took to swallow exceptions. bad = proof_comparison(rule_id, "bad") self.assertEqual( [(error.line, error.summary) for error in bad.new_errors], - [(924, '"route" is possibly unbound')], - ) - self.assertEqual( - bad.stale_entries, - (f'{GATEWAY_ENDPOINTS} :: {UNBOUND} :: "bundle" is possibly unbound',), + [ + (923, '"route" is possibly unbound'), + (924, '"route" is possibly unbound'), + (935, '"payload" is possibly unbound'), + ], ) + self.assertEqual(bad.stale_entries, ()) self.assertTrue(proof_comparison(rule_id, "good").passed) def test_rule_proof_pyright_executor(self) -> None: rule_id = "sim.quality.pyright-executor" + # Three of the executor's 178 real errors before its fixes: modal + # attributes unknown while src/modal shadowed the library, and a + # variable assigned inside a span. bad = proof_comparison(rule_id, "bad") self.assertEqual( [(error.line, error.summary) for error in bad.new_errors], - [(123, '"Secret" is not a known attribute of module "modal"')], - ) - self.assertEqual( - bad.stale_entries, - ( - f"{EXECUTOR_APP} :: reportAttributeAccessIssue :: " - '"Image" is not a known attribute of module "modal"', - ), + [ + (119, '"App" is not a known attribute of module "modal"'), + (123, '"Secret" is not a known attribute of module "modal"'), + (824, '"metadata" is possibly unbound'), + ], ) + self.assertEqual(bad.stale_entries, ()) self.assertTrue(proof_comparison(rule_id, "good").passed) def test_rule_proof_pyright_apis_integ(self) -> None: