[anth-volk]

Fixes #872

Same-Repository Draft PR Check

• This PR is a draft.
• This PR is from a branch in PolicyEngine/policyengine-sim-api, not a personal fork.

Summary

Fifth in the Pyright stack, on top of #871. Fixes the executor's remaining 72 Pyright errors, so it has none: rules/baselines/pyright-executor.txt goes from 72 entries to 0. The stack's top PR deletes it and makes sim.quality.pyright-executor enforced.

Commit	Errors fixed	Baseline after
Add a typed operation scope next to the span scope	(helper)	72
Open the executor's remaining spans and operations with typed scopes	17	55
Type the version-registry utility's routing state as it is stored	8	47
Check the shape of a simulation's output data before reading it	16	31
Narrow the output helpers' column sums and aggregate results	7	24
Type the baseline artifact outcome and the precompute manifest	6	18
Build the Stage 12 simulation's two types in separate branches	7	11
Narrow the Stage 12 parity check's report values and columns	7	4
Fix the executor's last four Pyright errors	4	0
Fail the batch, as before, when a child's call handle is missing	(review)	0
Raise TypeError wherever a simulation's output data is missing	(review)	0
Test the output-frame and model-version protocols on real US types	(review)	0
Point the observability and type-checking guides at scopes.operation	(review)	0
Open tracked_stage's scopes through the typed scopes helpers	(review)	0
Test the UK weight-matrix check's TypeError for a dataset with no entity data	(review)	0

What changed

• scopes.operation (observability library): runtime.operation(...) typed like scopes.span, taking the attributes and remote_context the Modal functions pass. The Modal functions' operations (app.py, v2_app.py) and the remaining spans (budget-window batch and scheduler, segmented national reduce, Stage 12 fan-out and comparison, output builder) use the typed scopes. stage12_runtime.failure_causes.tracked_stage, which records a Stage 12 stage for failure causes and times it, also opens its span or operation with scopes.span / scopes.operation instead of calling runtime.span / runtime.operation directly. No other code in the executor's src calls those runtime methods directly: the remaining .span(...) calls in stage12_runtime/coordination.py are Stage12Coordination.span, a method already typed AbstractContextManager[object] that calls tracked_stage.

• Version-registry utility (update_version_registry.py, run by the deploy workflows):

  • LatestVersions is non-total: a kind has no latest entry until one is routed, and readers already used .get. The few direct reads now use .get too. They are the plan's advances_latest flag and the post-publish and post-restore reports, which run after validate_routing_state has required those entries, plus a merge that copied a latest only when present.
  • RoutingState.bundles holds Mapping[str, Any], because stored manifests are read back unvalidated.
  • _json_copy is generic.
  • The finalize paths check for a missing Dict where they already reached the same outcome.

• Output data reads. output_frames adds two runtime-checkable protocols, EntityFrames and HouseholdFrame. Each read of <simulation>.output_dataset.data.<frames> first checks isinstance(<simulation>.output_dataset.data, ...) and raises TypeError if the frames are missing (before, AttributeError on None). Pyright narrows that same expression, so the reads are unchanged, and that matters for the weighted-microdata checks:

  • they recognise weighted frames by these expressions;
  • their stripping baseline lists pd.DataFrame(simulation.output_dataset.data.household) verbatim;
  • a helper returning the frames would hide them from the checks.

  The protocol checks need only the attribute, so the tests' stand-in datasets pass. Frames are typed as pandas DataFrames (a MicroDataFrame is one), because microdf types a column read as MicroSeries | MicroDataFrame.

• Output helpers' sums. _sum_output_variable and _change_output_variable require the selected column to be a Series before summing. A MicroSeries is one, so sums stay weighted (checked: a weighted column sums to its weighted total). They raise TypeError for a repeated column name, where float() raised TypeError before. Aggregate and ChangeAggregate results are checked for None. _output_model_dump reads model_dump with getattr.

• Precompute and artifact outcome:

  • The manifest is constructed with its alias schema=. Pydantic accepts both names (populate_by_name), but Pyright's Pydantic support knows only the alias.
  • ArtifactBaselineSimulation records an ArtifactOutcome Literal instead of str.
  • The determinism check reads the outcome with getattr.
  • _missing_output_columns checks the model version against a runtime-checkable EntityVariableResolver protocol, since policyengine defines resolve_entity_variables only on MicrosimulationModelVersion; the tests' stand-in model versions pass.

• Stage 12 simulation builder. It builds Simulation and Stage12CachedSimulation in separate branches instead of unpacking {"id": ...}, which Pyright matched against every keyword.

• Stage 12 parity:

  • _is_report_list and _is_report_number return TypeGuards. TypeGuard, unlike TypeIs, doesn't narrow on False, which matters because a str is a Sequence the list check rejects.
  • Frame columns must be Series; a repeated column name raises TypeError, where the comparison failed with AttributeError on .dtype before.

• Last four:

  • the budget-window totals are passed by name;
  • the scheduler checks the call handle that ChildSimulationHandle types as optional;
  • release_sentinel reads __doc__ or "" (__doc__ is None under -OO);
  • the SPM capability uses SPMCapability.model_validate.

Behaviour

Behaviour is unchanged on every path production reaches. The changes that differ do so only in paths that already failed, and they now fail with a clearer error:

• a simulation with no output frames, or no output data after ensure(): TypeError instead of AttributeError on None (or, in _entity_data, RuntimeError in this PR's first version);
• a repeated column name: TypeError with the column's name, instead of TypeError or AttributeError from pandas;
• an Aggregate with no result: RuntimeError instead of TypeError from float(None);
• a missing child call handle (unreachable): the batch fails, as before, with a RuntimeError in place of the AttributeError.

The last two commits change no behaviour. scopes.span(runtime, name, attributes=...) and scopes.operation(runtime, name, attributes=...) call runtime.span(name, attributes=...) and runtime.operation(name, attributes=..., remote_context=None) and return the runtime's scope unchanged, so tracked_stage opens the same scopes as before.

The microdata weight checks (sim.microdata.*) and their baseline are unchanged and pass.

Commits (red, then green)

Each type-fix commit's red is the Pyright check run on the code before the fix, with that commit's entries removed from the baseline. The check reported exactly those entries as new errors (17, 8, 16, 7, 6, 7, 7, 4). Green is the check passing with the shrunk baseline, plus the executor tests. Two commits add tests first:

• scopes.operation: red module ... has no attribute 'operation'; green, 38 library tests.
• The output-frame protocols: red No module named 'policyengine_simulation_executor.output_frames'; green, 3 tests covering real UK year data, stand-ins and None.

The last two commits answer a later review pass:

• tracked_stage with the typed scopes: a refactor, so there is no red. test_tracked_stage_scopes_and_records_the_stage_with_its_role records the scopes opened and passes before and after; its stand-in runtime's operation now accepts remote_context, as ObservabilityRuntime.operation does, and asserts it is None. Green: 56 failure-cause tests; Pyright 0 errors.
• The UK weight-matrix check's TypeError: test_guard_rejects_a_dataset_with_no_entity_data passes data=None and a stand-in with no entity_data. Red, with the isinstance(dataset.data, EntityFrames) check removed: both cases failed with AttributeError ('NoneType' object has no attribute 'entity_data', 'types.SimpleNamespace' object has no attribute 'entity_data'). Green: 11 weight-matrix tests pass.

Suppressions added

None.

Real bugs found

None. precompute.build_manifest's manifest_schema= looked like a wrong keyword, but ArtifactManifest sets populate_by_name=True, so it worked.

Self-review

review-code on this PR's diff, run by a separate agent, found no bugs. It reproduced each commit's baseline, confirmed that weighted sums stay weighted and that the microdata checks still see every read, and checked the Pydantic and deployment-utility changes. Code-quality findings:

1. The guides still prescribed runtime.operation(..., remote_context=...) and only scopes.span. Fixed: both guides name scopes.operation, and the type-checking guide says how to narrow output data without hiding the read from the microdata checks.
2. Missing output data was handled unevenly. _entity_data raised RuntimeError or AttributeError where the new reads raise TypeError, and the UK household count used its own getattr check. Fixed: _entity_data raises TypeError in both cases (red: a two-case test got RuntimeError and AttributeError), and the household count uses EntityFrames.
3. The scheduler's missing-call check sat before the poll's try, so it would have left the batch running instead of failed. Fixed: the check moved into _child_call, called inside the try, which restores the earlier handling. Red: a test with a handle lacking its call got the RuntimeError out of the scheduler. Green: the batch is marked failed.
4. The protocol tests covered UK year data only. Added: US year data, and MicrosimulationModelVersion against EntityVariableResolver. Both pass on the existing code.
5. The determinism check reads the outcome with getattr where the compute path uses isinstance(..., ArtifactBaselineSimulation). Declined: the determinism test passes a stand-in baseline that isn't an ArtifactBaselineSimulation. The error message still reports the outcome it read.

A later review-code pass over the whole stack found two more, both fixed:

6. The UK weight-matrix check's new TypeError for a dataset with no entity data had no test. Added (the last commit).
7. tracked_stage still called runtime.span / runtime.operation directly. Converted to the typed scopes, with no behaviour change; it was the only direct call left in src.

A review-code pass over the last two commits found no bugs. It confirmed that both helpers pass the same arguments as the direct calls and return the runtime's own scope unwrapped. It also confirmed that the new weight-matrix test fails with AttributeError without the check.

Environment variables

update_version_registry.py and release_sentinel.py are deployment utilities that the deploy workflows run. Their environment-variable reads are unchanged: GITHUB_ACTIONS, GITHUB_REPOSITORY and GITHUB_WORKFLOW_REF, which GitHub Actions provides, are read in _require_serialized_routing_workflow. No workflow, image or Modal app definition changed. app.py and v2_app.py now import policyengine_simulation_observability.scopes, part of the observability library the images already install, and the Modal image smoke covers them when this PR reaches the bottom of the stack.

Testing

• python -m scripts.rules pyright projects/policyengine-simulation-executor: passes with 0 errors.
• Executor unit tests: 1567 passed, 22 skipped. Observability library tests: 38 passed.
• Quality guards and rendered rules views: pass. Ruff 0.15.22 format and check: clean.

This PR is above the bottom of the stack, so its CI runs only CI complete, which fails by design until the PRs below merge.

🤖 Generated with Claude Code