Stop src/modal shadowing the modal library in the executor's Pyright check

docs/generated/rules.md

@@ -68,7 +68,7 @@ Every rule that governs `policyengine-sim-api` is listed here with how it is enf
 - `sim.ci.multiline-run-is-script`: 13 remaining in `rules/baselines/ci-multiline-run.txt`
 - `sim.microdata.weight-stripping-confined`: 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`: 124 remaining in `rules/baselines/pyright-executor.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

projects/policyengine-simulation-executor/pyproject.toml

@@ -97,3 +97,12 @@ markers = [
 [tool.pyright]
 # The pr.yml pyright job checks these paths: python -m scripts.rules pyright.
 include = ["src"]
+# Without extraPaths, Pyright adds src as an import root, and `import modal`
+# in the Modal app modules under src/modal resolves to that directory instead
+# of the installed modal library. Where the apps run (modal deploy and
+# modal run, and the Modal containers), src is not an import root: the package
+# is installed from src, and the app modules import each other as
+# src.modal.<module> from the project directory. The unit tests are the
+# exception: pytest's pythonpath puts src first, so there `import modal` can
+# find src/modal, and the tests replace the modal module where they need it.
+extraPaths = []

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

@@ -32,6 +32,9 @@ def run_budget_window_batch_impl(
     if selection is not None:
         params = {**params, "spm": selection}
     batch_job_id = modal.current_function_call_id()
+    if batch_job_id is None:
+        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)):
         context = build_batch_context(

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

@@ -197,9 +197,10 @@ def main(*, force: bool = False) -> None:
     run_precompute(
         resolve_bucket_name(),
         force=force,
-        # Pyright doesn't resolve `modal` here to the Modal package (it reports
-        # `modal.App` as unknown), so it types each decorated function as the
-        # plain function it wraps. At run time each is a Modal function.
+        # Pyright types each decorated function as a modal.Function whose
+        # `remote` and `spawn` take that function's own parameters, which it
+        # doesn't accept for RemoteFunction's `*args: Any`. At run time each is
+        # a Modal function, called with its own arguments.
         functions=PrecomputeFunctions(
             plan_artifacts=cast("RemoteFunction", plan_artifacts),
             build_dataset=cast("RemoteFunction", build_dataset),

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

@@ -46,6 +46,7 @@
     NationalReduceInputs,
     build_national_output,
 )
+from policyengine_simulation_observability import scopes
 from policyengine_simulation_observability.errors import log_and_redact_exception
 from policyengine_simulation_observability.stages import (
     SEGMENTED_NATIONAL_STAGES,
@@ -329,7 +330,8 @@ def _extend_last_group_with_uncovered_states(
     def _spawn_child(self, group: list[str], *, attempt: int) -> modal.FunctionCall:
         payload = build_group_child_payload(self.params, group)
         observability_context = self.runtime.capture_context()
-        with self.runtime.span(
+        with scopes.span(
+            self.runtime,
             SEGMENTED_NATIONAL_STAGES.name(Stage.SEGMENTED_NATIONAL_CHILD_SPAWN),
             attributes={"region_group": "+".join(group), "child_attempt": attempt},
         ):

projects/policyengine-simulation-executor/src/modal/utils/update_version_registry.py

@@ -439,6 +439,13 @@ def build_restored_routing_state(
 ROUTING_PLAN_SCHEMA = "policyengine-simulation.routing-plan/v1"
 
 
+# Raised, as a RuntimeError, when a routing write or a plan's completion finds
+# no routing-state Dict to write to.
+UNOPENED_ROUTING_STORE = (
+    "the routing state Modal Dict could not be opened; nothing was written"
+)
+
+
 class RoutingConflictError(RuntimeError):
     """The live routing state is not the one a plan expects."""
 
@@ -509,18 +516,27 @@ def validate_routing_plan(plan: object) -> dict:
 
 
 def _routing_store(environment: str, *, create_if_missing: bool) -> "modal.Dict | None":
+    """Return the routing-state Dict, or None when it doesn't exist.
+
+    modal's Dict.from_name returns a handle without contacting Modal, so a
+    missing Dict would surface as NotFoundError at the first read. When the
+    Dict must already exist, look it up now, so a missing one is reported here.
+    """
     try:
-        return modal.Dict.from_name(
+        store = modal.Dict.from_name(
             ROUTING_STATE_DICT_NAME,
             environment_name=environment,
             create_if_missing=create_if_missing,
         )
+        if not create_if_missing:
+            store.hydrate()
     except KeyError:
         return None
     except Exception as exc:
         if exc.__class__.__name__ == "NotFoundError":
             return None
         raise
+    return store
 
 
 CANONICAL_GITHUB_REPOSITORY = "PolicyEngine/policyengine-sim-api"
@@ -617,6 +633,8 @@ def record_routing_plan(*, environment: str, plan: dict) -> None:
     _require_serialized_routing_workflow()
     validate_routing_plan(plan)
     store = _routing_store(environment, create_if_missing=True)
+    if store is None:
+        raise RuntimeError(UNOPENED_ROUTING_STORE)
     _require_finished_publication(store)
     if (
         routing_state_digest(store.get(ROUTING_STATE_ACTIVE_KEY))
@@ -632,6 +650,8 @@ def complete_routing_plan(*, environment: str, plan: dict) -> None:
     _require_serialized_routing_workflow()
     validate_routing_plan(plan)
     store = _routing_store(environment, create_if_missing=False)
+    if store is None:
+        raise RuntimeError(UNOPENED_ROUTING_STORE)
     if _unfinished_publication(store) != plan:
         msg = "no matching unfinished routing publication to complete"
         raise RoutingConflictError(msg)
@@ -649,6 +669,8 @@ def apply_routing_plan(*, environment: str, plan: dict) -> None:
     _require_serialized_routing_workflow()
     validate_routing_plan(plan)
     store = _routing_store(environment, create_if_missing=True)
+    if store is None:
+        raise RuntimeError(UNOPENED_ROUTING_STORE)
     _require_finished_publication(store, allowed_plans=[plan])
     current = store.get(ROUTING_STATE_ACTIVE_KEY)
     if routing_state_digest(current) != plan["before_sha256"]:

projects/policyengine-simulation-executor/tests/test_budget_window_batch.py

@@ -255,6 +255,19 @@ def test_run_budget_window_batch_impl_completes_and_respects_max_parallel(
     assert state.result.totals.budgetaryImpact == 51
 
 
+def test_run_budget_window_batch_impl_refuses_to_run_outside_a_modal_call(
+    mock_batch_modal, observability_runtime, monkeypatch
+):
+    """Without a Modal call id there is no batch job id to key the batch by."""
+    _, payload = _build_parent_payload()
+    monkeypatch.setattr(batch_module.modal, "current_function_call_id", lambda: None)
+
+    with pytest.raises(RuntimeError, match="must run inside a Modal function call"):
+        run_budget_window_batch_impl(payload, runtime=observability_runtime)
+
+    assert mock_batch_modal["call_registry"] == {}
+
+
 def test_run_budget_window_batch_impl_marks_failure(
     mock_batch_modal, observability_runtime
 ):

projects/policyengine-simulation-executor/tests/test_update_version_registry.py

@@ -42,6 +42,9 @@ def pop(self, key):
     def snapshot(self) -> dict:
         return dict(self._data)
 
+    def hydrate(self):
+        return self
+
 
 @pytest.fixture
 def serialized_workflow(monkeypatch):
@@ -1336,6 +1339,92 @@ def pop(self, key):
     assert store.get("active") == plan["after"]
 
 
+class _UnopenableDict:
+    """A Modal Dict lookup that finds nothing, even when asked to create the Dict."""
+
+    @staticmethod
+    def from_name(name: str, **lookup_options: object):
+        missing_dict = f"{lookup_options['environment_name']}/{name}"
+        raise KeyError(missing_dict)
+
+
+def test_routing_writes_refuse_a_routing_store_that_cannot_be_opened(
+    monkeypatch, serialized_workflow
+):
+    """Recording or applying a plan stops with a RuntimeError, not an AttributeError."""
+    monkeypatch.setattr(registry, "modal", SimpleNamespace(Dict=_UnopenableDict))
+    original = _publish(monkeypatch, None, "6.0.0", "1.1.0", "2.1.0")
+    plan = _plan(monkeypatch, original, ("6.1.0", "1.1.0", "2.2.0"))
+
+    for operation in (registry.record_routing_plan, registry.apply_routing_plan):
+        with pytest.raises(
+            RuntimeError, match="the routing state Modal Dict could not be opened"
+        ) as raised:
+            operation(environment="staging", plan=plan)
+        # Not a RoutingConflictError: no plan disagrees with the live state.
+        assert type(raised.value) is RuntimeError
+
+
+class NotFoundError(Exception):
+    """Named like modal.exception.NotFoundError, which the utility matches by name."""
+
+
+class _LazyMissingDict:
+    """A missing Dict as modal 1.x reports it.
+
+    Dict.from_name returns a handle without contacting Modal; looking the Dict
+    up, or the first read, raises NotFoundError.
+    """
+
+    @staticmethod
+    def from_name(name: str, **lookup_options: object):
+        del lookup_options
+        return _LazyMissingDict.Handle(name)
+
+    class Handle:
+        def __init__(self, name: str):
+            self.name = name
+
+        def hydrate(self):
+            raise NotFoundError(self.name)
+
+        def get(self, key, default=None):
+            raise NotFoundError(self.name)
+
+
+@pytest.mark.parametrize(
+    "missing_dict", [_UnopenableDict, _LazyMissingDict], ids=["eager", "lazy"]
+)
+def test_completing_a_plan_refuses_a_missing_routing_store(
+    monkeypatch, serialized_workflow, missing_dict
+):
+    monkeypatch.setattr(registry, "modal", SimpleNamespace(Dict=missing_dict))
+    original = _publish(monkeypatch, None, "6.0.0", "1.1.0", "2.1.0")
+    plan = _plan(monkeypatch, original, ("6.1.0", "1.1.0", "2.2.0"))
+
+    with pytest.raises(RuntimeError, match=registry.UNOPENED_ROUTING_STORE) as raised:
+        registry.complete_routing_plan(environment="staging", plan=plan)
+    assert type(raised.value) is RuntimeError
+
+
+def test_restoring_with_a_lazily_missing_routing_store(
+    monkeypatch, serialized_workflow
+):
+    """With modal 1.x a missing Dict is reported at lookup; restore sees no store."""
+    monkeypatch.setattr(registry, "modal", SimpleNamespace(Dict=_LazyMissingDict))
+    first = _plan(monkeypatch, None, ("6.1.0", "1.1.0", "2.2.0"))
+    assert (
+        registry.restore_routing_plans(environment="staging", plans=[first])
+        == "unchanged"
+    )
+    later = _plan(monkeypatch, first["after"], ("6.2.0", "1.1.0", "2.3.0"))
+    with pytest.raises(
+        registry.RoutingConflictError,
+        match=r"^there is no routing state to restore into$",
+    ):
+        registry.restore_routing_plans(environment="staging", plans=[first, later])
+
+
 def test_runner_loss_cannot_be_adopted_by_the_next_release(monkeypatch, patched_modal):
     """Failed A cannot become successful B's starting state or rollback target."""
     original = _publish(monkeypatch, None, "6.0.0", "1.1.0", "2.1.0")