[anth-volk]

Fixes #868

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

Third in the Pyright stack, on top of #866. Pyright now resolves import modal in the executor's Modal app modules to the installed modal library instead of the local src/modal directory, and the errors that the shadowing hid are fixed. The executor's baseline (rules/baselines/pyright-executor.txt) goes from 178 entries to 126 and grows in no commit.

Commit	Baseline after
Refuse a routing write when the routing-state Dict can't be opened	173
Refuse to run a budget-window batch outside a Modal function call	173
Make Pyright resolve import modal to the modal library	126
Raise RuntimeError, not a routing conflict, for an unopened store	126
Correct the executor's extraPaths comment about the unit tests	126
Report a missing routing-state Dict however modal reports it	126

The shadowing fix: configuration only

Without extraPaths, Pyright adds src as an import root, so import modal inside src/modal/*.py resolved to src/modal/__init__.py. At runtime src is never an import root: the package is installed from src (an editable install adds src to sys.path after site-packages), and the app modules import each other as src.modal.<module> from the project directory. extraPaths = [] in the executor's [tool.pyright] turns the automatic root off. Pyright still resolves src.modal.* from the project root and policyengine_simulation_executor from the environment. Only Pyright reads this setting: no rename or move, and no change to Modal deployment entrypoints.

On its own it removes 47 errors and exposes 9 that the shadowing hid. To keep the baseline shrinking in every commit, the 3 exposed errors that need a behaviour change in an unreachable path are fixed first, each with a test written before the fix. The configuration commit then fixes the other 6.

What changed

• Routing-state Dict (2 exposed errors, plus 5 baselined None errors in the same functions). _routing_store returns None when Modal finds no routing-state Dict. record_routing_plan and apply_routing_plan ask it to create the Dict, so None should never come back. If it did, they failed with AttributeError: 'NoneType' object has no attribute 'get'; they now raise RuntimeError ("could not be opened; nothing was written"). complete_routing_plan now raises the same RuntimeError for a missing store, before its check that a matching plan is pending. The command exits non-zero either way, as before, and the utility's main handles neither exception.
• Finding a missing routing-state Dict (last commit). In modal 1.4.2 (the executor's locked version), modal.Dict.from_name returns a handle without contacting Modal; a Dict that doesn't exist raises modal.exception.NotFoundError at the first read. So _routing_store never returned None for a missing Dict, and every store is None branch was unreachable: callers failed with a NotFoundError traceback instead. _routing_store now calls the handle's hydrate() (synchronous; it looks the Dict up) when the caller passes create_if_missing=False, so a missing Dict returns None whether modal reports it at lookup or at first read. The None branches are kept, and are now reachable:
  • complete_routing_plan raises RuntimeError with UNOPENED_ROUTING_STORE ("the routing state Modal Dict could not be opened; nothing was written"), the message the routing writes use.
  • restore_routing_plans keeps its existing, tested contract for a missing Dict: "unchanged" when no plan had a prior state, otherwise RoutingConflictError ("there is no routing state to restore into"). I did not change it to the RuntimeError: a first restore with nothing to restore is a valid no-op that the existing tests require, and the conflict already names the missing state.
  • finalize and the pending-plan report reach their existing "nothing pending" outcomes. In the deploy workflow, the finalize step then ends with outcome no-pending instead of a NotFoundError; .github/scripts/simulation-deployment-ready.sh still fails the deployment, because it requires FINALIZE_OUTCOME=completed. In simulation-rollback.yml, the step that resolves the unfinished routing publication (--print-pending-plan) prints null instead of failing.
  • The routing writes pass create_if_missing=True, so they don't call hydrate() and are unchanged.
• Budget-window batch call id (1 exposed error). modal.current_function_call_id() returns str | None. The batch keys its children and state by it, and its docstring says it must run inside a Modal call. Outside one, it went on until building its state failed with a Pydantic ValidationError. It now raises RuntimeError first, as the Stage 12 coordinator in v2_app.py already does. Nothing changes inside a Modal call.
• RemoteFunction (5 exposed errors, previously baselined with other wording). Each Modal function's remote and spawn take that function's own parameters, so a protocol method declared (*args: Any) didn't accept them. The protocol now declares (*args: Any, **kwargs: Any), the gradual signature type checkers treat as "any arguments". This changes only the lines' annotations; their existing noqa: ANN401 comments already cover Any.
• Segmented national child spawn (1 exposed error). The span around the spawn that _spawn_child returns now uses scopes.span, so Pyright no longer reports the method as missing a return.

Behaviour changes

Both are in paths that production doesn't reach, and the code failed in both before. Tests written first:

Change	Test	Red	Green
Routing writes with no store raise RuntimeError	test_routing_writes_refuse_a_routing_store_that_cannot_be_opened	AttributeError: 'NoneType' object has no attribute 'get'	passes
The batch refuses to run outside a Modal call	test_run_budget_window_batch_impl_refuses_to_run_outside_a_modal_call	Pydantic ValidationError on batch_job_id	passes

The last commit changes what a missing routing-state Dict produces. With modal 1.4.2 these paths failed with NotFoundError before:

Change	Test	Red	Green
complete_routing_plan raises RuntimeError(UNOPENED_ROUTING_STORE) for a missing Dict, reported at lookup or at first read	test_completing_a_plan_refuses_a_missing_routing_store (parametrised over a stand-in that reports at lookup and one that reports at first read, like modal 1.x)	RoutingConflictError and NotFoundError	passes
restore_routing_plans returns its tested outcome for a Dict modal reports missing at first read	test_restoring_with_a_lazily_missing_routing_store	NotFoundError	passes

Commits (red, then green)

1. Routing store. Red: the test above. Green: 275 version-registry tests pass; the check reports 5 entries that no longer occur, removed (178 to 173).
2. Batch call id. Red: the test above. Green: batch and scheduler tests pass (13); baseline unchanged at 173.
3. Configuration. Red: with extraPaths = [], the check reported 6 errors the baseline didn't list and 47 entries that no longer occur. Green: 0 new errors after the protocol and span fixes; baseline 126.
4. Plain RuntimeError for the routing writes (review finding 2). Red: the test pinning the exact type got RoutingConflictError. Green: passes.
5. extraPaths comment (review finding 1). Comment only.
6. Missing routing-state Dict (review finding 4). Red: the two tests above. Green: 277 version-registry tests pass; baseline unchanged at 126.

Suppressions added

None.

Real bugs found

None in production paths. The two behaviour changes above replace a confusing failure with an explicit one in paths that production doesn't reach.

Self-review

review-code on this PR's diff, run by a separate agent, found no bugs. It reproduced the 47 removed and 9 exposed errors, and showed both new tests failing on the parent revision. Code-quality findings:

1. The extraPaths comment said src is never an import root at runtime, but pytest's pythonpath puts it first in the unit tests. Fixed: the comment now says where the apps run and names the test exception.
2. RoutingConflictError means the live state isn't the one a plan expects, which a Dict that can't be opened isn't. Fixed: the two writers raise a plain RuntimeError with the same message. The test pins the exact type, and was red against RoutingConflictError. The reviewer also noted that modal.Dict.from_name returns a handle without contacting Modal, so this path needs a lookup error and is unreachable in a real deployment; the comment now says so.
3. Only one span in segmented_national is converted here. Declined: the plan converts the rest in Fix the executor's remaining Pyright errors #873.

A later review-code pass over the whole stack found one more:

4. Because Dict.from_name is lazy, the store is None branches could not run, and complete_routing_plan reported a missing Dict as NotFoundError (or, with an eager stand-in, as a routing conflict). Fixed in the last commit, as described above.

A review-code pass over the last commit found no bugs. It confirmed in the installed modal 1.4.2 that from_name is lazy, that hydrate() is synchronous, and that a missing Dict raises NotFoundError, and it ran the new tests red against the parent commit. One minor finding, not changed: finalize and --print-pending-plan with a Dict that modal reports missing at first read are checked by hand, not by a test.

Environment variables

update_version_registry.py is a deployment utility that the deploy workflows run. Its environment-variable reads are unchanged, including in the last commit, which changes only _routing_store and complete_routing_plan: GITHUB_ACTIONS, GITHUB_REPOSITORY and GITHUB_WORKFLOW_REF, which GitHub Actions provides to every step, are read in _require_serialized_routing_workflow. budget_window_batch.py and segmented_national.py run inside Modal and read no new variables. segmented_national.py now imports policyengine_simulation_observability.scopes (#864), part of the observability library that the executor images already install. No workflow, image or Modal app definition changed.

Rebase onto #879 (2026-10-22)

Rebased onto main after #879 merged. Where this section and the text above differ, this section is current.

What conflicted. rules/registry.toml, docs/generated/rules.md and rules/baselines/pyright-executor.txt. #879 removed 54 executor baseline entries (178 to 124): 49 in simulation_runtime.py, and the 5 precompute_app.py errors for the precompute functions passed to PrecomputeFunctions, which #879 fixed by casting its six Modal functions (it adds build_state_datasets) to RemoteFunction.

How it was resolved.

• Each commit's baseline was regenerated with python -m scripts.rules pyright projects/policyengine-simulation-executor --write-baseline, starting from its parent's baseline, and the count set to match. Pyright passes at every commit.
• The RemoteFunction protocol change (*args: Any, **kwargs: Any) is dropped, and Bake per-state ACS-local year files into the main path's worker image #879's casts are kept. With the casts, Pyright reports nothing at that call with or without the protocol change. Without either, it reports 6 errors once modal resolves.
• The comment above Bake per-state ACS-local year files into the main path's worker image #879's casts said Pyright doesn't resolve modal in precompute_app.py, which this PR makes untrue. It now says why the casts are still needed: each modal.Function's remote and spawn take that function's own parameters, which Pyright doesn't accept for RemoteFunction's *args: Any. No code changed there.
• extraPaths = [] now removes 42 errors and exposes 4: the 3 fixed by the first two commits, and the segmented national span.

Commit	Baseline after
Refuse a routing write when the routing-state Dict can't be opened	119
Refuse to run a budget-window batch outside a Modal function call	119
Make Pyright resolve import modal to the modal library	77
Raise RuntimeError, not a routing conflict, for an unopened store	77
Correct the executor's extraPaths comment about the unit tests	77
Report a missing routing-state Dict however modal reports it	77

Baseline at this PR's head: 77 entries (from 124 on main).

Checks after the rebase. Pyright passes with 77. Executor unit tests: 1645 passed, 22 skipped, 3 deselected. Ruff 0.15.22 format and check are clean. The quality guards and python -m scripts.rules guard --base-ref origin/main pass. render leaves no changes. No new noqa.

Testing

• python -m scripts.rules pyright projects/policyengine-simulation-executor: passes with 126.
• tests/test_update_version_registry.py: 277 passed.
• Executor unit tests: 1555 passed, 22 skipped, before the last commit, which adds 2 tests to the version-registry file.
• 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