Fix the executor's remaining Pyright errors

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:
 

docs/engineering/skills/type-checking.md

@@ -114,8 +114,14 @@ Baselines are under `rules/`, so changing one needs code-owner approval
   same line: `# pyright: ignore[reportAttributeAccessIssue]  # <reason>`.
 - 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
+  `<simulation>.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.

docs/generated/rules.md

@@ -67,7 +67,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`: 0 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`: 66 remaining in `rules/baselines/pyright-executor.txt`
+- `sim.quality.pyright-executor`: 0 remaining in `rules/baselines/pyright-executor.txt`
 - `sim.quality.pyright-gateway`: 0 remaining in `rules/baselines/pyright-gateway.txt`
 
 ### Exceptions and review dates

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)

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)

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,
@@ -366,7 +367,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,
@@ -438,7 +440,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,
@@ -478,7 +481,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,

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,

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(

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)