diff --git a/changelog_entry.yaml b/changelog_entry.yaml index c044351cc..41efb7fd9 100644 --- a/changelog_entry.yaml +++ b/changelog_entry.yaml @@ -16,6 +16,7 @@ - Ran Stage 12 US national simulations as 20 region-group calls each, and built an independent, immutable Stage 12 cache of three annual US datasets and 60 current-law baselines before each deployment, publishing it only after a recorded determinism check and recording its release in the v2 version manifest under a new key, while the existing key keeps a copy that earlier readers accept - Made the budget refuse a household count outside its country's plausible range (above 200 million in the US or 60 million in the UK, negative or not finite), raising ImplausibleEntityCountError instead of publishing a count like #767's changed: + - Required policyengine-observability 3.0.4 (>=3.0.4,<4) in the observability library, the entrypoint, the gateway and the executor, and locked it in the six locks that include it, where it is the only package that moves; 3.0.4 annotates its span and operation scopes' __exit__ and __aexit__ as returning Literal[False], so the services now open spans and operations with runtime.span and runtime.operation directly, and the policyengine_simulation_observability.scopes module, which only retyped those scopes for Pyright, is removed with its tests; a new observability library test fails if an older release is installed, Pyright still reports no errors in any package, and 3.0.3's fix, which keeps the recorded exception on the request.completed record and server span of a request that ends with an error outcome but passes no exception to end_request, is included - Enforced the file-length rule sim.quality.python-file-length, now that every tracked Python file is within its limit: its empty baseline and its downgrade section are gone, and its tier is metric rather than mechanical, because it measures a number against a threshold, with the at-the-limit and one-over fixtures as its below and above proof - Moved the entrypoint's backend calls, response telemetry and comparison dispatch out of app.py (925 lines, now 583) into entry_runtime.py, which app.py re-exports and which still logs under the policyengine_simulation_entry.app logger; the routes stay in app.py, so FastAPI and ruff read their annotations as before; the 2,251-line deployment asset test file is split into seven test files and a support module named test_deployment_assets_support.py, the rules registry names each moved test by its new file, and its guard-test discovery now covers tests/test_deployment_assets*.py - Moved the worker's credential setup, request year and policy parsing, and region resolution out of simulation_runtime.py (1,111 lines, now 651) into gcp_credentials.py, request_policy.py and region_resolution.py; simulation_runtime.py keeps dataset selection and loading, simulation building, local-area breakdowns and the run itself, which tests replace by name there, and re-exports the moved names, so every caller still imports them from simulation_runtime; the shared helpers _parse_year, _normalise_policy, _compile_interval_policy, _runtime_policy (now wrapper_policy), _resolve_region, _resolve_dataset_reference and _require_local_area_breakdown are now public, and the credential messages still log under the simulation_runtime logger; the 3,048-line output builder test file is split into seven test files and a module of shared fakes, all named test_simulation_output_*.py so the weighted-frame rule still checks them diff --git a/docs/engineering/skills/observability.md b/docs/engineering/skills/observability.md index d0d08d303..c7ad2ac63 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 -`scopes.operation(runtime, ..., remote_context=context)`. The operation scope must +`runtime.operation(..., 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,11 +79,10 @@ 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))`, 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 +Open a span with `runtime.span(plan.name(Stage.VALUE))`, and an operation with +`runtime.operation(...)`. policyengine-observability 3.0.4 types both scopes +correctly: their `__exit__` and `__aexit__` return `Literal[False]`, so 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 1a202c579..56a444b58 100644 --- a/docs/engineering/skills/type-checking.md +++ b/docs/engineering/skills/type-checking.md @@ -173,12 +173,14 @@ section is still in the registry fails the guard on every pull request. 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 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. +- Open observability spans and operations with `runtime.span(...)` and + `runtime.operation(...)`. policyengine-observability 3.0.4 types them + correctly: their `__exit__` and `__aexit__` return `Literal[False]`, so + Pyright knows they never suppress an exception and doesn't report a variable + the block assigns as possibly unbound after it. Every package that depends + on policyengine-observability requires 3.0.4 or later, and + `libs/policyengine-simulation-observability/tests/test_scope_exit_typing.py` + fails if an older release is installed. - Before reading a simulation's output frames, narrow `.output_dataset.data` with the `EntityFrames` protocol in `policyengine_simulation_executor.output_frames`, and keep reading through diff --git a/libs/policyengine-simulation-contract/uv.lock b/libs/policyengine-simulation-contract/uv.lock index a1b23fb94..8915800ea 100644 --- a/libs/policyengine-simulation-contract/uv.lock +++ b/libs/policyengine-simulation-contract/uv.lock @@ -1069,11 +1069,11 @@ wheels = [ [[package]] name = "policyengine-observability" -version = "3.0.2" +version = "3.0.4" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/74/2a/0c261c8ca693bb3a13d7fe837d3f89f7fdda245496d4b3e92eefed3c972e/policyengine_observability-3.0.2.tar.gz", hash = "sha256:554f907e43d8eeb274c1298fbe990b1626ac45f194c49995ddde091ddcdc2a64", size = 125718, upload-time = "2026-09-30T21:36:02.028Z" } +sdist = { url = "https://files.pythonhosted.org/packages/c3/d0/09b91aca37fb3efd493c164082097e85705fcac2062d2b5ef8a8c4e2be04/policyengine_observability-3.0.4.tar.gz", hash = "sha256:7e43a5d0e9b9e5f683623643efca7c47e651baaad7b557a50ec8b344602dc61f", size = 127972, upload-time = "2026-10-27T17:23:24.343Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/a5/9c/b48fb64c1bebf83bd5fdeebe914188806c3513cf2de0fc95ee1d95082fc4/policyengine_observability-3.0.2-py3-none-any.whl", hash = "sha256:d84c7564922f6a3294268e89549065924edde932c31adf8b7c714558e107cb31", size = 42029, upload-time = "2026-09-30T21:36:00.841Z" }, + { url = "https://files.pythonhosted.org/packages/00/00/d9c4850a12b4d0bef72034316909aac97228de302384af8936a958b7dc8e/policyengine_observability-3.0.4-py3-none-any.whl", hash = "sha256:1b3ebbfd00c138abcc2ec9580b13aa55679b91745783c5567cc73f8b2bbaaaf7", size = 42204, upload-time = "2026-10-27T17:23:22.641Z" }, ] [package.optional-dependencies] @@ -1145,7 +1145,7 @@ requires-dist = [ { name = "black", marker = "extra == 'build'", specifier = ">=25.1.0" }, { name = "fastapi", specifier = ">=0.115.0" }, { name = "httpx2", marker = "extra == 'test'" }, - { name = "policyengine-observability", extras = ["fastapi", "google", "otlp-grpc"], specifier = ">=3.0.2,<4" }, + { name = "policyengine-observability", extras = ["fastapi", "google", "otlp-grpc"], specifier = ">=3.0.4,<4" }, { name = "pydantic", specifier = ">=2.0" }, { name = "pyright", marker = "extra == 'build'", specifier = ">=1.1.401" }, { name = "pytest", marker = "extra == 'test'", specifier = ">=8.3.4" }, diff --git a/libs/policyengine-simulation-observability/pyproject.toml b/libs/policyengine-simulation-observability/pyproject.toml index 15f8048ca..d7cf62251 100644 --- a/libs/policyengine-simulation-observability/pyproject.toml +++ b/libs/policyengine-simulation-observability/pyproject.toml @@ -10,7 +10,7 @@ requires-python = ">=3.13" dependencies = [ "fastapi>=0.115.0", "pydantic>=2.0", - "policyengine-observability[fastapi,google,otlp-grpc]>=3.0.2,<4", + "policyengine-observability[fastapi,google,otlp-grpc]>=3.0.4,<4", ] [project.optional-dependencies] diff --git a/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/scopes.py b/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/scopes.py deleted file mode 100644 index d1a31400b..000000000 --- a/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/scopes.py +++ /dev/null @@ -1,54 +0,0 @@ -"""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 -``bool`` result as "this context manager may suppress the exception raised in -its block", so after ``with runtime.span(...)`` Pyright reports every variable -the block assigned as possibly unbound, and a function that returns from inside -the block as missing a return. - -``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 - -from typing import TYPE_CHECKING, Any - -if TYPE_CHECKING: - from collections.abc import Mapping - from contextlib import AbstractContextManager - - from policyengine_observability import ObservabilityRuntime - - -def span( - runtime: ObservabilityRuntime, - name: str, - *, - attributes: Mapping[str, Any] | None = None, -) -> AbstractContextManager[object]: - """Return ``runtime.span(name, attributes=attributes)``. - - The scope times its block as a child span, records an exception raised in - 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_scope_exit_typing.py b/libs/policyengine-simulation-observability/tests/test_scope_exit_typing.py new file mode 100644 index 000000000..cd54ad5ca --- /dev/null +++ b/libs/policyengine-simulation-observability/tests/test_scope_exit_typing.py @@ -0,0 +1,46 @@ +"""The installed policyengine-observability types its scopes as never suppressing. + +The services open spans and operations with ``runtime.span(...)`` and +``runtime.operation(...)``. A type checker reads an ``__exit__`` or +``__aexit__`` annotated as returning ``bool`` as "this context manager may +suppress the exception raised in its block", so Pyright would then report every +variable such a block assigns as possibly unbound after it, and every function +that returns from inside the block as possibly not returning. +policyengine-observability 3.0.4 annotates both methods as returning +``Literal[False]``. These tests fail if an older release is installed. +""" + +from typing import Literal, get_type_hints + +import pytest +from policyengine_simulation_observability.observability import build_runtime + + +@pytest.fixture +def runtime(monkeypatch): + monkeypatch.setenv("OTEL_SDK_DISABLED", "true") + real_runtime = build_runtime( + service_name="policyengine-simulation-scope-typing-test", + service_role="test", + platform="local", + environment="test", + ) + yield real_runtime + real_runtime.shutdown() + + +@pytest.mark.parametrize("exit_method_name", ["__exit__", "__aexit__"]) +@pytest.mark.parametrize( + "open_scope", + [ + pytest.param(lambda runtime: runtime.span("scope_typing"), id="span"), + pytest.param(lambda runtime: runtime.operation("scope_typing"), id="operation"), + ], +) +def test_scope_exit_is_typed_as_never_suppressing_an_exception( + runtime, open_scope, exit_method_name +): + scope_manager = open_scope(runtime) + exit_method = getattr(type(scope_manager), exit_method_name) + + assert get_type_hints(exit_method)["return"] == Literal[False] diff --git a/libs/policyengine-simulation-observability/tests/test_scopes.py b/libs/policyengine-simulation-observability/tests/test_scopes.py deleted file mode 100644 index 7b524bc45..000000000 --- a/libs/policyengine-simulation-observability/tests/test_scopes.py +++ /dev/null @@ -1,103 +0,0 @@ -"""Tests for the typed span and operation scopes of the observability runtime.""" - -from contextlib import nullcontext - -import pytest -from policyengine_simulation_observability import scopes -from policyengine_simulation_observability.observability import build_runtime - - -class RecordingRuntime: - """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): - monkeypatch.setenv("OTEL_SDK_DISABLED", "true") - real_runtime = build_runtime( - service_name="policyengine-simulation-scopes-test", - service_role="test", - platform="local", - environment="test", - ) - yield real_runtime - real_runtime.shutdown() - - -def test_span_opens_the_runtime_span_with_its_name_and_attributes(): - recording_runtime = RecordingRuntime() - - with scopes.span( - recording_runtime, "simulation.parse", attributes={"country": "us"} - ) as entered_state: - assert recording_runtime.opened_spans == [ - ("simulation.parse", {"country": "us"}) - ] - - assert entered_state == "entered span state" - - -def test_span_defaults_to_no_attributes(): - recording_runtime = RecordingRuntime() - - with scopes.span(recording_runtime, "simulation.parse"): - pass - - assert recording_runtime.opened_spans == [("simulation.parse", None)] - - -def test_span_lets_an_exception_raised_inside_it_propagate(runtime): - message = "invalid simulation request" - - with ( - pytest.raises(ValueError, match=message), - 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/libs/policyengine-simulation-observability/uv.lock b/libs/policyengine-simulation-observability/uv.lock index dc1e298e0..8e971eb08 100644 --- a/libs/policyengine-simulation-observability/uv.lock +++ b/libs/policyengine-simulation-observability/uv.lock @@ -720,11 +720,11 @@ wheels = [ [[package]] name = "policyengine-observability" -version = "3.0.2" +version = "3.0.4" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/74/2a/0c261c8ca693bb3a13d7fe837d3f89f7fdda245496d4b3e92eefed3c972e/policyengine_observability-3.0.2.tar.gz", hash = "sha256:554f907e43d8eeb274c1298fbe990b1626ac45f194c49995ddde091ddcdc2a64", size = 125718, upload-time = "2026-09-30T21:36:02.028Z" } +sdist = { url = "https://files.pythonhosted.org/packages/c3/d0/09b91aca37fb3efd493c164082097e85705fcac2062d2b5ef8a8c4e2be04/policyengine_observability-3.0.4.tar.gz", hash = "sha256:7e43a5d0e9b9e5f683623643efca7c47e651baaad7b557a50ec8b344602dc61f", size = 127972, upload-time = "2026-10-27T17:23:24.343Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/a5/9c/b48fb64c1bebf83bd5fdeebe914188806c3513cf2de0fc95ee1d95082fc4/policyengine_observability-3.0.2-py3-none-any.whl", hash = "sha256:d84c7564922f6a3294268e89549065924edde932c31adf8b7c714558e107cb31", size = 42029, upload-time = "2026-09-30T21:36:00.841Z" }, + { url = "https://files.pythonhosted.org/packages/00/00/d9c4850a12b4d0bef72034316909aac97228de302384af8936a958b7dc8e/policyengine_observability-3.0.4-py3-none-any.whl", hash = "sha256:1b3ebbfd00c138abcc2ec9580b13aa55679b91745783c5567cc73f8b2bbaaaf7", size = 42204, upload-time = "2026-10-27T17:23:22.641Z" }, ] [package.optional-dependencies] @@ -768,7 +768,7 @@ requires-dist = [ { name = "black", marker = "extra == 'build'", specifier = ">=25.1.0" }, { name = "fastapi", specifier = ">=0.115.0" }, { name = "httpx2", marker = "extra == 'test'" }, - { name = "policyengine-observability", extras = ["fastapi", "google", "otlp-grpc"], specifier = ">=3.0.2,<4" }, + { name = "policyengine-observability", extras = ["fastapi", "google", "otlp-grpc"], specifier = ">=3.0.4,<4" }, { name = "pydantic", specifier = ">=2.0" }, { name = "pyright", marker = "extra == 'build'", specifier = ">=1.1.401" }, { name = "pytest", marker = "extra == 'test'", specifier = ">=8.3.4" }, diff --git a/libs/policyengine-stage12-persistence/uv.lock b/libs/policyengine-stage12-persistence/uv.lock index e8d8bbe36..5e3c5dca8 100644 --- a/libs/policyengine-stage12-persistence/uv.lock +++ b/libs/policyengine-stage12-persistence/uv.lock @@ -566,11 +566,11 @@ wheels = [ [[package]] name = "policyengine-observability" -version = "3.0.2" +version = "3.0.4" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/74/2a/0c261c8ca693bb3a13d7fe837d3f89f7fdda245496d4b3e92eefed3c972e/policyengine_observability-3.0.2.tar.gz", hash = "sha256:554f907e43d8eeb274c1298fbe990b1626ac45f194c49995ddde091ddcdc2a64", size = 125718, upload-time = "2026-09-30T21:36:02.028Z" } +sdist = { url = "https://files.pythonhosted.org/packages/c3/d0/09b91aca37fb3efd493c164082097e85705fcac2062d2b5ef8a8c4e2be04/policyengine_observability-3.0.4.tar.gz", hash = "sha256:7e43a5d0e9b9e5f683623643efca7c47e651baaad7b557a50ec8b344602dc61f", size = 127972, upload-time = "2026-10-27T17:23:24.343Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/a5/9c/b48fb64c1bebf83bd5fdeebe914188806c3513cf2de0fc95ee1d95082fc4/policyengine_observability-3.0.2-py3-none-any.whl", hash = "sha256:d84c7564922f6a3294268e89549065924edde932c31adf8b7c714558e107cb31", size = 42029, upload-time = "2026-09-30T21:36:00.841Z" }, + { url = "https://files.pythonhosted.org/packages/00/00/d9c4850a12b4d0bef72034316909aac97228de302384af8936a958b7dc8e/policyengine_observability-3.0.4-py3-none-any.whl", hash = "sha256:1b3ebbfd00c138abcc2ec9580b13aa55679b91745783c5567cc73f8b2bbaaaf7", size = 42204, upload-time = "2026-10-27T17:23:22.641Z" }, ] [package.optional-dependencies] @@ -627,7 +627,7 @@ requires-dist = [ { name = "black", marker = "extra == 'build'", specifier = ">=25.1.0" }, { name = "fastapi", specifier = ">=0.115.0" }, { name = "httpx2", marker = "extra == 'test'" }, - { name = "policyengine-observability", extras = ["fastapi", "google", "otlp-grpc"], specifier = ">=3.0.2,<4" }, + { name = "policyengine-observability", extras = ["fastapi", "google", "otlp-grpc"], specifier = ">=3.0.4,<4" }, { name = "pydantic", specifier = ">=2.0" }, { name = "pyright", marker = "extra == 'build'", specifier = ">=1.1.401" }, { name = "pytest", marker = "extra == 'test'", specifier = ">=8.3.4" }, diff --git a/projects/policyengine-simulation-entry/pyproject.toml b/projects/policyengine-simulation-entry/pyproject.toml index bc419e070..1b43b4f1d 100644 --- a/projects/policyengine-simulation-entry/pyproject.toml +++ b/projects/policyengine-simulation-entry/pyproject.toml @@ -16,7 +16,7 @@ dependencies = [ "fastapi>=0.115.0,<1", "httpx>=0.28.0,<1", "modal>=1.4,<2", - "policyengine-observability[fastapi,google,httpx,otlp-grpc]>=3.0.2,<4", + "policyengine-observability[fastapi,google,httpx,otlp-grpc]>=3.0.4,<4", "policyengine-simulation-contract", "policyengine-simulation-observability", "policyengine-stage12-persistence", diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/stage12_backend.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/stage12_backend.py index a1d232941..29dad7567 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/stage12_backend.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/stage12_backend.py @@ -24,7 +24,6 @@ V2ManifestLoader, V2WorkerVersion, ) -from policyengine_simulation_observability import scopes from policyengine_simulation_observability.identifiers import ( normalize_observability_id, ) @@ -493,8 +492,7 @@ async def _spawn_report( else STAGE12_CANONICAL_REPORT_STAGES ) try: - with scopes.span( - self._runtime, + with self._runtime.span( stage_plan.name(Stage.STAGE12_ENTRY_DISPATCH), attributes={ "evaluation_id": str(parent.evaluation_id), diff --git a/projects/policyengine-simulation-entry/uv.lock b/projects/policyengine-simulation-entry/uv.lock index 943b788d5..c3bfa7d53 100644 --- a/projects/policyengine-simulation-entry/uv.lock +++ b/projects/policyengine-simulation-entry/uv.lock @@ -840,11 +840,11 @@ wheels = [ [[package]] name = "policyengine-observability" -version = "3.0.2" +version = "3.0.4" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/74/2a/0c261c8ca693bb3a13d7fe837d3f89f7fdda245496d4b3e92eefed3c972e/policyengine_observability-3.0.2.tar.gz", hash = "sha256:554f907e43d8eeb274c1298fbe990b1626ac45f194c49995ddde091ddcdc2a64", size = 125718, upload-time = "2026-09-30T21:36:02.028Z" } +sdist = { url = "https://files.pythonhosted.org/packages/c3/d0/09b91aca37fb3efd493c164082097e85705fcac2062d2b5ef8a8c4e2be04/policyengine_observability-3.0.4.tar.gz", hash = "sha256:7e43a5d0e9b9e5f683623643efca7c47e651baaad7b557a50ec8b344602dc61f", size = 127972, upload-time = "2026-10-27T17:23:24.343Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/a5/9c/b48fb64c1bebf83bd5fdeebe914188806c3513cf2de0fc95ee1d95082fc4/policyengine_observability-3.0.2-py3-none-any.whl", hash = "sha256:d84c7564922f6a3294268e89549065924edde932c31adf8b7c714558e107cb31", size = 42029, upload-time = "2026-09-30T21:36:00.841Z" }, + { url = "https://files.pythonhosted.org/packages/00/00/d9c4850a12b4d0bef72034316909aac97228de302384af8936a958b7dc8e/policyengine_observability-3.0.4-py3-none-any.whl", hash = "sha256:1b3ebbfd00c138abcc2ec9580b13aa55679b91745783c5567cc73f8b2bbaaaf7", size = 42204, upload-time = "2026-10-27T17:23:22.641Z" }, ] [package.optional-dependencies] @@ -925,7 +925,7 @@ requires-dist = [ { name = "httpx", specifier = ">=0.28.0,<1" }, { name = "modal", specifier = ">=1.4,<2" }, { name = "openapi-python-client", marker = "extra == 'build'", specifier = ">=0.21.6" }, - { name = "policyengine-observability", extras = ["fastapi", "google", "httpx", "otlp-grpc"], specifier = ">=3.0.2,<4" }, + { name = "policyengine-observability", extras = ["fastapi", "google", "httpx", "otlp-grpc"], specifier = ">=3.0.4,<4" }, { name = "policyengine-simulation-contract", editable = "../../libs/policyengine-simulation-contract" }, { name = "policyengine-simulation-observability", editable = "../../libs/policyengine-simulation-observability" }, { name = "policyengine-stage12-persistence", editable = "../../libs/policyengine-stage12-persistence" }, @@ -954,7 +954,7 @@ requires-dist = [ { name = "black", marker = "extra == 'build'", specifier = ">=25.1.0" }, { name = "fastapi", specifier = ">=0.115.0" }, { name = "httpx2", marker = "extra == 'test'" }, - { name = "policyengine-observability", extras = ["fastapi", "google", "otlp-grpc"], specifier = ">=3.0.2,<4" }, + { name = "policyengine-observability", extras = ["fastapi", "google", "otlp-grpc"], specifier = ">=3.0.4,<4" }, { name = "pydantic", specifier = ">=2.0" }, { name = "pyright", marker = "extra == 'build'", specifier = ">=1.1.401" }, { name = "pytest", marker = "extra == 'test'", specifier = ">=8.3.4" }, diff --git a/projects/policyengine-simulation-executor/pyproject.toml b/projects/policyengine-simulation-executor/pyproject.toml index a2147ca9f..6a3b7f131 100644 --- a/projects/policyengine-simulation-executor/pyproject.toml +++ b/projects/policyengine-simulation-executor/pyproject.toml @@ -24,7 +24,7 @@ dependencies = [ "policyengine[models]==6.10.0", "tables>=3.10.2", "modal>=0.73.0", - "policyengine-observability[fastapi,google,otlp-grpc]>=3.0.2,<4", + "policyengine-observability[fastapi,google,otlp-grpc]>=3.0.4,<4", # Imported directly by the artifact store client; do not rely on the # transitive pin via policyengine. "google-cloud-storage>=2", @@ -49,7 +49,7 @@ modal-simulation-image = [ "policyengine[models]==6.10.0", "fastapi>=0.115.0", "tables>=3.10.2", - "policyengine-observability[fastapi,google,otlp-grpc]>=3.0.2,<4", + "policyengine-observability[fastapi,google,otlp-grpc]>=3.0.4,<4", # The artifact fetch layer and store client read GCS. The lib is also a # transitive dependency of policyengine, but the image must not depend # on an upstream package keeping it. diff --git a/projects/policyengine-simulation-executor/src/modal/app.py b/projects/policyengine-simulation-executor/src/modal/app.py index 0e5fcd1f3..3ee231c03 100644 --- a/projects/policyengine-simulation-executor/src/modal/app.py +++ b/projects/policyengine-simulation-executor/src/modal/app.py @@ -24,7 +24,6 @@ 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, @@ -367,8 +366,7 @@ def run_simulation( } try: propagated = normalize_observability_context(observability_context) - with scopes.operation( - runtime, + with runtime.operation( ANNUAL_IMPACT_STAGES.name(Stage.ANNUAL_EXECUTION), attributes=redacted_params, remote_context=propagated, @@ -440,8 +438,7 @@ def run_simulation_segment( } try: propagated = normalize_observability_context(observability_context) - with scopes.operation( - runtime, + with runtime.operation( SEGMENTED_NATIONAL_STAGES.name(Stage.SEGMENTED_NATIONAL_EXECUTION), attributes=redacted_params, remote_context=propagated, @@ -481,8 +478,7 @@ def run_budget_window_batch( } try: propagated = normalize_observability_context(observability_context) - with scopes.operation( - runtime, + with runtime.operation( 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 d03a4c065..90ec90208 100644 --- a/projects/policyengine-simulation-executor/src/modal/budget_window_batch.py +++ b/projects/policyengine-simulation-executor/src/modal/budget_window_batch.py @@ -4,7 +4,6 @@ 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 @@ -37,7 +36,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 scopes.span(runtime, BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_CONTEXT)): + with runtime.span(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_scheduler.py b/projects/policyengine-simulation-executor/src/modal/budget_window_scheduler.py index 3213ab1c5..f19c5796d 100644 --- a/projects/policyengine-simulation-executor/src/modal/budget_window_scheduler.py +++ b/projects/policyengine-simulation-executor/src/modal/budget_window_scheduler.py @@ -20,7 +20,6 @@ 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 ( @@ -72,8 +71,8 @@ def serialize_batch_status( runtime: ObservabilityRuntime, ) -> dict[str, Any]: """Render ``state`` as the JSON-ready batch status response.""" - with scopes.span( - runtime, BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_STATUS_SERIALIZATION) + with runtime.span( + BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_STATUS_SERIALIZATION) ): return build_batch_status_response(state).model_dump(mode="json") @@ -84,9 +83,7 @@ def load_or_create_batch_state( runtime: ObservabilityRuntime, ) -> BudgetWindowBatchState: """Return the batch's submit-time seed state, creating and storing it if absent.""" - with scopes.span( - runtime, BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_STATE_LOAD) - ): + with runtime.span(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( @@ -97,9 +94,7 @@ def load_or_create_batch_state( ), bundle=context.bundle, ) - with scopes.span( - runtime, BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_STATE_WRITE) - ): + with runtime.span(BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_STATE_WRITE)): put_batch_job_seed(state) return state @@ -173,8 +168,8 @@ def run(self) -> dict[str, Any]: failure fails the whole batch. """ mark_batch_running(self.state) - with scopes.span( - self.runtime, BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_STATE_WRITE) + with self.runtime.span( + BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_STATE_WRITE) ): put_batch_job_state(self.state) @@ -214,8 +209,7 @@ def spawn_until_capacity(self) -> None: and self.state.queued_years ): simulation_year = self.state.queued_years[0] - with scopes.span( - self.runtime, + with self.runtime.span( BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_CHILD_REQUEST_BUILD), attributes={"simulation_year": simulation_year}, ): @@ -224,8 +218,7 @@ def spawn_until_capacity(self) -> None: simulation_year=simulation_year, ) observability_context = self.runtime.capture_context() - with scopes.span( - self.runtime, + with self.runtime.span( BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_CHILD_SPAWN), attributes={"simulation_year": simulation_year}, ): @@ -243,8 +236,8 @@ def spawn_until_capacity(self) -> None: year=simulation_year, child_job_id=call.object_id, ) - with scopes.span( - self.runtime, BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_STATE_WRITE) + with self.runtime.span( + BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_STATE_WRITE) ): put_batch_job_state(self.state) @@ -294,8 +287,7 @@ def poll_running_children_once(self) -> bool: self._publish_poll_stats() try: - with scopes.span( - self.runtime, + with self.runtime.span( BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_RESULT_PARSE), attributes={"simulation_year": simulation_year}, ): @@ -334,8 +326,8 @@ def poll_running_children_once(self) -> bool: year=simulation_year, annual_impact=annual_impact, ) - with scopes.span( - self.runtime, BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_STATE_WRITE) + with self.runtime.span( + BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_STATE_WRITE) ): put_batch_job_state(self.state) progress_made = True @@ -378,15 +370,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 scopes.span( - self.runtime, BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_STATE_WRITE) + with self.runtime.span( + 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 scopes.span( - self.runtime, BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_AGGREGATION) + with self.runtime.span( + BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_AGGREGATION) ): annual_impacts = [ self.state.partial_annual_impacts[simulation_year] @@ -399,8 +391,8 @@ def complete_batch(self) -> dict[str, Any]: annual_impacts=annual_impacts, ) mark_batch_complete(self.state, result=result) - with scopes.span( - self.runtime, BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_STATE_WRITE) + with self.runtime.span( + 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 73d3a82a9..ea2a3ed04 100644 --- a/projects/policyengine-simulation-executor/src/modal/segmented_national.py +++ b/projects/policyengine-simulation-executor/src/modal/segmented_national.py @@ -46,7 +46,6 @@ 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, @@ -330,8 +329,7 @@ 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 scopes.span( - self.runtime, + with self.runtime.span( SEGMENTED_NATIONAL_STAGES.name(Stage.SEGMENTED_NATIONAL_CHILD_SPAWN), attributes={"region_group": "+".join(group), "child_attempt": attempt}, ): @@ -587,8 +585,7 @@ 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 scopes.span( - self.runtime, + with self.runtime.span( SEGMENTED_NATIONAL_STAGES.name(Stage.SEGMENTED_NATIONAL_REDUCE), ): output = build_national_output( diff --git a/projects/policyengine-simulation-executor/src/modal/v2_app.py b/projects/policyengine-simulation-executor/src/modal/v2_app.py index cee879432..92642d673 100644 --- a/projects/policyengine-simulation-executor/src/modal/v2_app.py +++ b/projects/policyengine-simulation-executor/src/modal/v2_app.py @@ -30,7 +30,6 @@ 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, @@ -281,8 +280,7 @@ def run_single_simulation_segment_us( try: propagated = normalize_observability_context(observability_context) result = None - with scopes.operation( - runtime, + with runtime.operation( STAGE12_SEGMENT_STAGES.name(Stage.STAGE12_SEGMENT_EXECUTION), attributes={ "runner_name": "stage12", @@ -339,8 +337,7 @@ def segmented_calculator( try: propagated = normalize_observability_context(observability_context) result = None - with scopes.operation( - runtime, + with runtime.operation( STAGE12_SIMULATION_STAGES.name(Stage.STAGE12_SIMULATION_EXECUTION), attributes={ "runner_name": "stage12", @@ -390,8 +387,7 @@ def run_single_simulation_uk( try: propagated = normalize_observability_context(observability_context) result = None - with scopes.operation( - runtime, + with runtime.operation( STAGE12_SIMULATION_STAGES.name(Stage.STAGE12_SIMULATION_EXECUTION), attributes={ "runner_name": "stage12", @@ -457,8 +453,7 @@ def coordinate_report( try: propagated = normalize_observability_context(observability_context) result = None - with scopes.operation( - runtime, + with runtime.operation( stage_plan.name(Stage.STAGE12_COORDINATOR_EXECUTION), attributes={ "runner_name": "stage12", diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/gcp_credentials.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/gcp_credentials.py index 8d614524e..b4989df5e 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/gcp_credentials.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/gcp_credentials.py @@ -9,7 +9,6 @@ import tempfile from typing import TYPE_CHECKING -from policyengine_simulation_observability import scopes from policyengine_simulation_observability.stages import ANNUAL_IMPACT_STAGES, Stage if TYPE_CHECKING: @@ -62,7 +61,7 @@ def setup_gcp_credentials( with contextlib.ExitStack() as credentials_file_stack: try: credential_span = ( - scopes.span(runtime, ANNUAL_IMPACT_STAGES.name(Stage.CREDENTIAL_SETUP)) + runtime.span(ANNUAL_IMPACT_STAGES.name(Stage.CREDENTIAL_SETUP)) if runtime is not None else contextlib.nullcontext() ) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_provenance.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_provenance.py index 2316cf119..bd90ca638 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_provenance.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_provenance.py @@ -5,7 +5,6 @@ from contextlib import AbstractContextManager, nullcontext from typing import TYPE_CHECKING, Any -from policyengine_simulation_observability import scopes from policyengine_simulation_observability.stages import ( ANNUAL_IMPACT_STAGES, Stage, @@ -26,7 +25,7 @@ def output_span( """Open a span named ``name``, or nothing when no runtime is given.""" if runtime is None: return nullcontext() - return scopes.span(runtime, name) + return runtime.span(name) def output_model_version( 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 5088f1066..f6a81a6d1 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 @@ -18,7 +18,6 @@ from policyengine_simulation_contract.dataset_uri import runtime_dataset_uri from policyengine_simulation_contract.job_progress import JobStage from policyengine_simulation_contract.spm import SPMInputError, spm_error_detail -from policyengine_simulation_observability import scopes from policyengine_simulation_observability.stages import ANNUAL_IMPACT_STAGES, Stage from policyengine_simulation_observability.telemetry import split_internal_payload @@ -352,8 +351,7 @@ def _uk_area_simulations( ) area_role = {"dataset_role": "uk_area"} - with scopes.span( - runtime, + with runtime.span( ANNUAL_IMPACT_STAGES.name(Stage.DATASET_RESOLUTION), attributes=area_role, ): @@ -363,8 +361,8 @@ def _uk_area_simulations( code=region.code, dataset_reference=routed_uk_area_dataset_name() ), ) - with scopes.span( - runtime, ANNUAL_IMPACT_STAGES.name(Stage.DATASET_LOAD), attributes=area_role + with runtime.span( + ANNUAL_IMPACT_STAGES.name(Stage.DATASET_LOAD), attributes=area_role ): dataset = _load_dataset( params, selection=selection, country_module=country_module @@ -376,8 +374,7 @@ def _uk_area_simulations( for kind, policy in zip( ("uk_area_baseline", "uk_area_reform"), policies, strict=True ): - with scopes.span( - runtime, + with runtime.span( ANNUAL_IMPACT_STAGES.name(Stage.SIMULATION_BUILD), attributes={"simulation_kind": kind}, ): @@ -414,16 +411,16 @@ def _run_local_area_breakdown_core( impacts only (``LocalAreaBreakdownOutput``). """ require_local_area_breakdown(simulation_params, country=country) - with scopes.span(runtime, ANNUAL_IMPACT_STAGES.name(Stage.COUNTRY_MODULE_LOAD)): + with runtime.span(ANNUAL_IMPACT_STAGES.name(Stage.COUNTRY_MODULE_LOAD)): country_module = _country_module(country) - with scopes.span(runtime, ANNUAL_IMPACT_STAGES.name(Stage.REGION_RESOLUTION)): + with runtime.span(ANNUAL_IMPACT_STAGES.name(Stage.REGION_RESOLUTION)): region_resolution = resolve_region( country_module=country_module, country=country, params=simulation_params, ) runtime.set_context(region=region_resolution.code) - with scopes.span(runtime, ANNUAL_IMPACT_STAGES.name(Stage.POLICY_NORMALIZATION)): + with runtime.span(ANNUAL_IMPACT_STAGES.name(Stage.POLICY_NORMALIZATION)): policies = ( normalise_policy(simulation_params.get("baseline")), normalise_policy(simulation_params.get("reform")), @@ -463,7 +460,7 @@ def _run_simulation_impl_core( runtime: ObservabilityRuntime, progress: JobProgressReporter = NULL_JOB_PROGRESS, ) -> dict: - with scopes.span(runtime, ANNUAL_IMPACT_STAGES.name(Stage.REQUEST_PARSE)): + with runtime.span(ANNUAL_IMPACT_STAGES.name(Stage.REQUEST_PARSE)): simulation_params, telemetry, metadata = split_internal_payload(params) metadata = metadata or {} from policyengine_simulation_executor.spm import ( @@ -490,20 +487,20 @@ def _run_simulation_impl_core( simulation_params, runtime=runtime, progress=progress, country=country ) - with scopes.span(runtime, ANNUAL_IMPACT_STAGES.name(Stage.COUNTRY_MODULE_LOAD)): + with runtime.span(ANNUAL_IMPACT_STAGES.name(Stage.COUNTRY_MODULE_LOAD)): country_module = _country_module(country) - with scopes.span(runtime, ANNUAL_IMPACT_STAGES.name(Stage.REGION_RESOLUTION)): + with runtime.span(ANNUAL_IMPACT_STAGES.name(Stage.REGION_RESOLUTION)): region_resolution = resolve_region( country_module=country_module, country=country, params=simulation_params, ) runtime.set_context(region=region_resolution.code) - with scopes.span(runtime, ANNUAL_IMPACT_STAGES.name(Stage.DATASET_RESOLUTION)): + with runtime.span(ANNUAL_IMPACT_STAGES.name(Stage.DATASET_RESOLUTION)): dataset_selection = _resolve_dataset_selection( simulation_params, region_resolution=region_resolution ) - with scopes.span(runtime, ANNUAL_IMPACT_STAGES.name(Stage.DATASET_LOAD)): + with runtime.span(ANNUAL_IMPACT_STAGES.name(Stage.DATASET_LOAD)): baked_folder = state_datasets.baked_state_data_folder( country=country, selection=dataset_selection, @@ -536,13 +533,12 @@ def _run_simulation_impl_core( region_code=region_resolution.code, dataset_name=dataset_selection.name, ) - with scopes.span(runtime, ANNUAL_IMPACT_STAGES.name(Stage.POLICY_NORMALIZATION)): + with runtime.span(ANNUAL_IMPACT_STAGES.name(Stage.POLICY_NORMALIZATION)): baseline_policy = normalise_policy(simulation_params.get("baseline")) reform_policy = normalise_policy(simulation_params.get("reform")) logger.info("Initialising baseline and reform simulations") - with scopes.span( - runtime, + with runtime.span( ANNUAL_IMPACT_STAGES.name(Stage.SIMULATION_BUILD), attributes={"simulation_kind": "baseline"}, ): @@ -553,8 +549,7 @@ def _run_simulation_impl_core( policy=baseline_policy, region=region_resolution, ) - with scopes.span( - runtime, + with runtime.span( ANNUAL_IMPACT_STAGES.name(Stage.SIMULATION_BUILD), attributes={"simulation_kind": "reform"}, ): 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 7524c2525..fefc8281b 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,7 +13,6 @@ 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 @@ -60,7 +59,7 @@ def compare_completed_report( return try: comparison_span = ( - scopes.span(runtime, stage_plan.name(Stage.STAGE12_RESULT_COMPARISON)) + runtime.span(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 8a60a4d19..27aac2e99 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,7 +35,6 @@ 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 @@ -187,9 +186,9 @@ def tracked_stage( if runtime is None: scope = nullcontext() elif operation: - scope = scopes.operation(runtime, stage_name, attributes=attributes) + scope = runtime.operation(stage_name, attributes=attributes) else: - scope = scopes.span(runtime, stage_name, attributes=attributes) + scope = runtime.span(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 147c5549a..83194b9f5 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,7 +8,6 @@ from typing import TYPE_CHECKING, Any import modal -from policyengine_simulation_observability import scopes from policyengine_simulation_observability.stages import ( STAGE12_SIMULATION_STAGES, Stage, @@ -132,8 +131,7 @@ def run(self) -> SimulationCalculation: try: for segment in segments: call = None - with scopes.span( - self.runtime, + with self.runtime.span( STAGE12_SIMULATION_STAGES.name(Stage.STAGE12_SEGMENT_DISPATCH), attributes={"segment_index": segment.segment_index}, ): @@ -147,16 +145,16 @@ def run(self) -> SimulationCalculation: raise payloads: list[object] | None = None - with scopes.span( - self.runtime, STAGE12_SIMULATION_STAGES.name(Stage.STAGE12_SEGMENT_WAIT) + with self.runtime.span( + 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 scopes.span( - self.runtime, STAGE12_SIMULATION_STAGES.name(Stage.STAGE12_SEGMENT_MERGE) + with self.runtime.span( + STAGE12_SIMULATION_STAGES.name(Stage.STAGE12_SEGMENT_MERGE) ): merged = merge_segment_results( self.simulation, diff --git a/projects/policyengine-simulation-executor/uv.lock b/projects/policyengine-simulation-executor/uv.lock index 7453cde31..b68fcbe9e 100644 --- a/projects/policyengine-simulation-executor/uv.lock +++ b/projects/policyengine-simulation-executor/uv.lock @@ -1840,11 +1840,11 @@ provides-extras = ["test", "build"] [[package]] name = "policyengine-observability" -version = "3.0.2" +version = "3.0.4" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/74/2a/0c261c8ca693bb3a13d7fe837d3f89f7fdda245496d4b3e92eefed3c972e/policyengine_observability-3.0.2.tar.gz", hash = "sha256:554f907e43d8eeb274c1298fbe990b1626ac45f194c49995ddde091ddcdc2a64", size = 125718, upload-time = "2026-09-30T21:36:02.028Z" } +sdist = { url = "https://files.pythonhosted.org/packages/c3/d0/09b91aca37fb3efd493c164082097e85705fcac2062d2b5ef8a8c4e2be04/policyengine_observability-3.0.4.tar.gz", hash = "sha256:7e43a5d0e9b9e5f683623643efca7c47e651baaad7b557a50ec8b344602dc61f", size = 127972, upload-time = "2026-10-27T17:23:24.343Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/a5/9c/b48fb64c1bebf83bd5fdeebe914188806c3513cf2de0fc95ee1d95082fc4/policyengine_observability-3.0.2-py3-none-any.whl", hash = "sha256:d84c7564922f6a3294268e89549065924edde932c31adf8b7c714558e107cb31", size = 42029, upload-time = "2026-09-30T21:36:00.841Z" }, + { url = "https://files.pythonhosted.org/packages/00/00/d9c4850a12b4d0bef72034316909aac97228de302384af8936a958b7dc8e/policyengine_observability-3.0.4-py3-none-any.whl", hash = "sha256:1b3ebbfd00c138abcc2ec9580b13aa55679b91745783c5567cc73f8b2bbaaaf7", size = 42204, upload-time = "2026-10-27T17:23:22.641Z" }, ] [package.optional-dependencies] @@ -1953,7 +1953,7 @@ requires-dist = [ { name = "opentelemetry-instrumentation-sqlalchemy", specifier = ">=0.65b0,<0.66" }, { name = "policyengine", extras = ["models"], specifier = "==6.10.0" }, { name = "policyengine-fastapi", editable = "../../libs/policyengine-fastapi" }, - { name = "policyengine-observability", extras = ["fastapi", "google", "otlp-grpc"], specifier = ">=3.0.2,<4" }, + { name = "policyengine-observability", extras = ["fastapi", "google", "otlp-grpc"], specifier = ">=3.0.4,<4" }, { name = "policyengine-simulation-contract", extras = ["modal"], editable = "../../libs/policyengine-simulation-contract" }, { name = "policyengine-simulation-observability", editable = "../../libs/policyengine-simulation-observability" }, { name = "policyengine-stage12-persistence", editable = "../../libs/policyengine-stage12-persistence" }, @@ -1974,7 +1974,7 @@ modal-simulation-image = [ { name = "google-cloud-storage", specifier = ">=2" }, { name = "huggingface-hub", specifier = ">=1,<2" }, { name = "policyengine", extras = ["models"], specifier = "==6.10.0" }, - { name = "policyengine-observability", extras = ["fastapi", "google", "otlp-grpc"], specifier = ">=3.0.2,<4" }, + { name = "policyengine-observability", extras = ["fastapi", "google", "otlp-grpc"], specifier = ">=3.0.4,<4" }, { name = "psycopg", extras = ["binary"], specifier = ">=3.2,<4" }, { name = "pyarrow", specifier = ">=20,<24" }, { name = "sqlalchemy", specifier = ">=2,<3" }, @@ -2004,7 +2004,7 @@ requires-dist = [ { name = "modal", specifier = ">=1.4,<2" }, { name = "openapi-python-client", marker = "extra == 'build'", specifier = ">=0.21.6" }, { name = "packaging", specifier = ">=24" }, - { name = "policyengine-observability", extras = ["fastapi", "google", "otlp-grpc"], specifier = ">=3.0.2,<4" }, + { name = "policyengine-observability", extras = ["fastapi", "google", "otlp-grpc"], specifier = ">=3.0.4,<4" }, { name = "pydantic", specifier = ">=2.0" }, { name = "pyjwt", specifier = ">=2.10.1,<3.0.0" }, { name = "pyright", marker = "extra == 'build'", specifier = ">=1.1.401" }, @@ -2036,7 +2036,7 @@ requires-dist = [ { name = "black", marker = "extra == 'build'", specifier = ">=25.1.0" }, { name = "fastapi", specifier = ">=0.115.0" }, { name = "httpx2", marker = "extra == 'test'" }, - { name = "policyengine-observability", extras = ["fastapi", "google", "otlp-grpc"], specifier = ">=3.0.2,<4" }, + { name = "policyengine-observability", extras = ["fastapi", "google", "otlp-grpc"], specifier = ">=3.0.4,<4" }, { name = "pydantic", specifier = ">=2.0" }, { name = "pyright", marker = "extra == 'build'", specifier = ">=1.1.401" }, { name = "pytest", marker = "extra == 'test'", specifier = ">=8.3.4" }, diff --git a/projects/policyengine-simulation-gateway/pyproject.toml b/projects/policyengine-simulation-gateway/pyproject.toml index d72967079..4518faa1e 100644 --- a/projects/policyengine-simulation-gateway/pyproject.toml +++ b/projects/policyengine-simulation-gateway/pyproject.toml @@ -25,7 +25,7 @@ dependencies = [ "pyjwt>=2.10.1,<3.0.0", # JWTDecoder (policyengine-fastapi lib) needs cryptography at runtime. "cryptography>=41.0.0", - "policyengine-observability[fastapi,google,otlp-grpc]>=3.0.2,<4", + "policyengine-observability[fastapi,google,otlp-grpc]>=3.0.4,<4", "modal>=1.4,<2", ] diff --git a/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/job_creation_routes.py b/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/job_creation_routes.py index 174a16138..36bbc4ed3 100644 --- a/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/job_creation_routes.py +++ b/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/job_creation_routes.py @@ -35,7 +35,6 @@ resolve_spm_selection, spm_error_detail, ) -from policyengine_simulation_observability import scopes from policyengine_simulation_observability.identifiers import ( OBSERVABILITY_ID_HEADER, resolve_observability_id, @@ -219,7 +218,7 @@ async def submit_simulation( breakdown=request.breakdown, ) try: - with scopes.span(runtime, ANNUAL_IMPACT_STAGES.name(Stage.ROUTE_RESOLUTION)): + with runtime.span(ANNUAL_IMPACT_STAGES.name(Stage.ROUTE_RESOLUTION)): route = resolve_route( request.country, request.version, @@ -236,7 +235,7 @@ async def submit_simulation( policyengine_version=route.policyengine_version, ) - with scopes.span(runtime, ANNUAL_IMPACT_STAGES.name(Stage.REQUEST_PARSE)): + with runtime.span(ANNUAL_IMPACT_STAGES.name(Stage.REQUEST_PARSE)): payload = request.model_dump( exclude={"version", "policyengine_version", "telemetry"}, mode="json", @@ -246,7 +245,7 @@ async def submit_simulation( payload["_telemetry"] = request.telemetry.model_dump(mode="json") try: - with scopes.span(runtime, ANNUAL_IMPACT_STAGES.name(Stage.POLICYENGINE_BUNDLE)): + with runtime.span(ANNUAL_IMPACT_STAGES.name(Stage.POLICYENGINE_BUNDLE)): bundle = build_policyengine_bundle( request.country, route, @@ -279,7 +278,7 @@ async def submit_simulation( # handle — the control-plane RPC (hydration) happens inside ``spawn`` — # so both live under the spawn segment to time the real network cost. submitted_at = datetime.now(UTC) - with scopes.span(runtime, ANNUAL_IMPACT_STAGES.name(Stage.MODAL_FUNCTION_SPAWN)): + with runtime.span(ANNUAL_IMPACT_STAGES.name(Stage.MODAL_FUNCTION_SPAWN)): sim_func = modal.Function.from_name(route.app_name, "run_simulation") call = sim_func.spawn( payload, @@ -295,12 +294,10 @@ async def submit_simulation( submitted_at=submitted_at, job_progress=_route_job_progress_capability(route), ) - with scopes.span( - runtime, ANNUAL_IMPACT_STAGES.name(Stage.MODAL_JOB_METADATA_WRITE) - ): + with runtime.span(ANNUAL_IMPACT_STAGES.name(Stage.MODAL_JOB_METADATA_WRITE)): job_metadata_store()[call.object_id] = job_metadata - with scopes.span(runtime, ANNUAL_IMPACT_STAGES.name(Stage.RESPONSE_SERIALIZATION)): + with runtime.span(ANNUAL_IMPACT_STAGES.name(Stage.RESPONSE_SERIALIZATION)): return JobSubmitResponse( job_id=call.object_id, status="submitted", @@ -332,7 +329,7 @@ async def submit_budget_window_batch( observability_id=observability_id, ) try: - with scopes.span(runtime, BUDGET_WINDOW_STAGES.name(Stage.ROUTE_RESOLUTION)): + with runtime.span(BUDGET_WINDOW_STAGES.name(Stage.ROUTE_RESOLUTION)): route = resolve_route( request.country, request.version, @@ -349,7 +346,7 @@ async def submit_budget_window_batch( ) try: - with scopes.span(runtime, BUDGET_WINDOW_STAGES.name(Stage.POLICYENGINE_BUNDLE)): + with runtime.span(BUDGET_WINDOW_STAGES.name(Stage.POLICYENGINE_BUNDLE)): bundle = build_policyengine_bundle( request.country, route, @@ -365,7 +362,7 @@ async def submit_budget_window_batch( ) runtime.record_exception(exc, handled=True, status_code=400) raise HTTPException(status_code=400, detail=str(exc)) from exc - with scopes.span(runtime, BUDGET_WINDOW_STAGES.name(Stage.REQUEST_PARSE)): + with runtime.span(BUDGET_WINDOW_STAGES.name(Stage.REQUEST_PARSE)): payload = _build_budget_window_parent_payload( request, resolved_version=route.response_version, @@ -375,7 +372,7 @@ async def submit_budget_window_batch( # Lazy handle + spawn together: the RPC cost lands in ``spawn``. observability_context = runtime.capture_context() - with scopes.span(runtime, BUDGET_WINDOW_STAGES.name(Stage.MODAL_FUNCTION_SPAWN)): + with runtime.span(BUDGET_WINDOW_STAGES.name(Stage.MODAL_FUNCTION_SPAWN)): batch_func = modal.Function.from_name(route.app_name, "run_budget_window_batch") call = batch_func.spawn( payload, @@ -391,12 +388,10 @@ async def submit_budget_window_batch( bundle=bundle, observability_id=observability_id, ) - with scopes.span( - runtime, BUDGET_WINDOW_STAGES.name(Stage.MODAL_JOB_METADATA_WRITE) - ): + with runtime.span(BUDGET_WINDOW_STAGES.name(Stage.MODAL_JOB_METADATA_WRITE)): put_batch_job_seed(seed_state) - with scopes.span(runtime, BUDGET_WINDOW_STAGES.name(Stage.RESPONSE_SERIALIZATION)): + with runtime.span(BUDGET_WINDOW_STAGES.name(Stage.RESPONSE_SERIALIZATION)): return BudgetWindowBatchSubmitResponse( batch_job_id=batch_job_id, # A new batch's seed state is always submitted. diff --git a/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/job_status_routes.py b/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/job_status_routes.py index fea6f7838..986440248 100644 --- a/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/job_status_routes.py +++ b/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/job_status_routes.py @@ -23,7 +23,6 @@ ) from policyengine_simulation_contract.job_progress_state import get_job_progress 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.identifiers import OBSERVABILITY_ID_HEADER from policyengine_simulation_observability.stages import ( @@ -85,9 +84,7 @@ def _job_progress_payload( record = None if reports_stages: try: - with scopes.span( - runtime, ANNUAL_IMPACT_STAGES.name(Stage.MODAL_JOB_PROGRESS_READ) - ): + with runtime.span(ANNUAL_IMPACT_STAGES.name(Stage.MODAL_JOB_PROGRESS_READ)): record = get_job_progress(job_id) except Exception as exc: # noqa: BLE001 # stage reporting is optional logger.warning( @@ -122,7 +119,7 @@ async def get_job_status( """ runtime = observability_runtime(request) runtime.set_context(job_id=job_id) - with scopes.span(runtime, ANNUAL_IMPACT_STAGES.name(Stage.MODAL_JOB_METADATA_READ)): + with runtime.span(ANNUAL_IMPACT_STAGES.name(Stage.MODAL_JOB_METADATA_READ)): job_metadata = job_metadata_store().get(job_id) if job_metadata is None: record_not_found(runtime, f"Job not found: {job_id}") @@ -138,9 +135,7 @@ async def get_job_status( response_headers = observability_headers(observability_id) try: - with scopes.span( - runtime, ANNUAL_IMPACT_STAGES.name(Stage.MODAL_JOB_STATUS_POLL) - ): + with runtime.span(ANNUAL_IMPACT_STAGES.name(Stage.MODAL_JOB_STATUS_POLL)): call = modal.FunctionCall.from_id(job_id) except Exception as exc: if is_modal_job_not_found(exc): @@ -158,22 +153,16 @@ def with_progress(status: JobStatus) -> dict: } try: - with scopes.span( - runtime, ANNUAL_IMPACT_STAGES.name(Stage.MODAL_JOB_STATUS_POLL) - ): + with runtime.span(ANNUAL_IMPACT_STAGES.name(Stage.MODAL_JOB_STATUS_POLL)): result = call.get(timeout=0) complete_metadata = with_progress("complete") - with scopes.span( - runtime, ANNUAL_IMPACT_STAGES.name(Stage.RESPONSE_SERIALIZATION) - ): + with runtime.span(ANNUAL_IMPACT_STAGES.name(Stage.RESPONSE_SERIALIZATION)): return JobStatusResponse( status="complete", result=result, **complete_metadata ) except TimeoutError: running_metadata = with_progress("running") - with scopes.span( - runtime, ANNUAL_IMPACT_STAGES.name(Stage.RESPONSE_SERIALIZATION) - ): + with runtime.span(ANNUAL_IMPACT_STAGES.name(Stage.RESPONSE_SERIALIZATION)): return running_job_response(running_metadata, headers=response_headers) # call.get re-raises whatever the worker raised; every failure becomes a # failed-job response, logged with its traceback by log_and_redact_exception. @@ -198,9 +187,7 @@ def with_progress(status: JobStatus) -> dict: context={"job_id": job_id}, ) failed_metadata = with_progress("failed") - with scopes.span( - runtime, ANNUAL_IMPACT_STAGES.name(Stage.RESPONSE_SERIALIZATION) - ): + with runtime.span(ANNUAL_IMPACT_STAGES.name(Stage.RESPONSE_SERIALIZATION)): return failed_job_response( error=redacted, job_metadata=failed_metadata, @@ -222,25 +209,21 @@ async def get_budget_window_job_status( """Poll for budget-window batch status.""" runtime = observability_runtime(request) runtime.set_context(batch_job_id=batch_job_id) - with scopes.span( - runtime, BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_STATE_LOAD) - ): + with runtime.span(BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_STATE_LOAD)): state = get_batch_job_state(batch_job_id) if state is not None: bind_observability_id(request, runtime, state.observability_id) response_headers = observability_headers(state.observability_id) if response_headers is not None: response.headers.update(response_headers) - with scopes.span( - runtime, BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_STATUS_SERIALIZATION) + with runtime.span( + BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_STATUS_SERIALIZATION) ): return batch_status_response( build_batch_status_response(state), headers=response_headers ) - with scopes.span( - runtime, BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_STATE_LOAD) - ): + with runtime.span(BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_STATE_LOAD)): seed_state = get_batch_job_seed(batch_job_id) if seed_state is None: record_not_found(runtime, f"Budget-window job not found: {batch_job_id}") @@ -253,9 +236,7 @@ async def get_budget_window_job_status( response.headers.update(response_headers) try: - with scopes.span( - runtime, BUDGET_WINDOW_STAGES.name(Stage.MODAL_JOB_STATUS_POLL) - ): + with runtime.span(BUDGET_WINDOW_STAGES.name(Stage.MODAL_JOB_STATUS_POLL)): call = modal.FunctionCall.from_id(batch_job_id) except Exception as exc: # noqa: BLE001 # any lookup failure degrades to the seed state; logged below with its traceback at warning level # The endpoint degrades gracefully to the seed state (a successful @@ -276,21 +257,19 @@ async def get_budget_window_job_status( batch_job_id, exc_info=exc, ) - with scopes.span( - runtime, BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_STATUS_SERIALIZATION) + with runtime.span( + BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_STATUS_SERIALIZATION) ): return batch_status_response( build_batch_status_response(seed_state), headers=response_headers ) try: - with scopes.span( - runtime, BUDGET_WINDOW_STAGES.name(Stage.MODAL_JOB_STATUS_POLL) - ): + with runtime.span(BUDGET_WINDOW_STAGES.name(Stage.MODAL_JOB_STATUS_POLL)): result = call.get(timeout=0) except TimeoutError: - with scopes.span( - runtime, BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_STATUS_SERIALIZATION) + with runtime.span( + BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_STATUS_SERIALIZATION) ): return batch_status_response( build_batch_status_response(seed_state), headers=response_headers @@ -314,23 +293,19 @@ async def get_budget_window_job_status( seed_state.status = "failed" seed_state.errors = [detail] if detail else None seed_state.error = message - with scopes.span( - runtime, BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_STATE_WRITE) - ): + with runtime.span(BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_STATE_WRITE)): put_batch_job_state(seed_state) put_batch_job_seed(seed_state) - with scopes.span( - runtime, BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_STATUS_SERIALIZATION) + with runtime.span( + BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_STATUS_SERIALIZATION) ): return batch_status_response( build_batch_status_response(seed_state), headers=response_headers ) - with scopes.span( - runtime, BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_RESULT_PARSE) - ): + with runtime.span(BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_RESULT_PARSE)): status_response = BudgetWindowBatchStatusResponse.model_validate(result) - with scopes.span( - runtime, BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_STATUS_SERIALIZATION) + with runtime.span( + BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_STATUS_SERIALIZATION) ): return batch_status_response(status_response, headers=response_headers) diff --git a/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/version_routes.py b/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/version_routes.py index 53d9b2c36..e3b64f1d5 100644 --- a/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/version_routes.py +++ b/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/version_routes.py @@ -10,7 +10,6 @@ VersionsResponse, ) from policyengine_simulation_contract.spm import SPMCapability -from policyengine_simulation_observability import scopes from policyengine_simulation_observability.stages import ANNUAL_IMPACT_STAGES, Stage from policyengine_simulation_gateway.bundle_payloads import bundle_identities @@ -37,7 +36,7 @@ async def list_versions(request: Request) -> VersionsResponse: """List all available routing versions.""" runtime = observability_runtime(request) - with scopes.span(runtime, ANNUAL_IMPACT_STAGES.name(Stage.ROUTE_RESOLUTION)): + with runtime.span(ANNUAL_IMPACT_STAGES.name(Stage.ROUTE_RESOLUTION)): state = active_routing_state() if state: capabilities = {} @@ -58,7 +57,7 @@ async def list_versions(request: Request) -> VersionsResponse: # ``Dict.from_name`` is a lazy handle; the RPCs fire during the # ``dict(...)`` iterations, so those are what the segment must cover. - with scopes.span(runtime, ANNUAL_IMPACT_STAGES.name(Stage.MODAL_DICT_READ)): + with runtime.span(ANNUAL_IMPACT_STAGES.name(Stage.MODAL_DICT_READ)): policyengine_dict = optional_modal_dict(POLICYENGINE_VERSION_DICT_NAME) us_dict = modal.Dict.from_name("simulation-api-us-versions") uk_dict = modal.Dict.from_name("simulation-api-uk-versions") @@ -92,7 +91,7 @@ async def get_country_versions(kind: str, request: Request) -> VersionMap: record_not_found(runtime, f"Unknown version kind: {kind}") raise HTTPException(status_code=404, detail=f"Unknown version kind: {kind}") - with scopes.span(runtime, ANNUAL_IMPACT_STAGES.name(Stage.ROUTE_RESOLUTION)): + with runtime.span(ANNUAL_IMPACT_STAGES.name(Stage.ROUTE_RESOLUTION)): state = active_routing_state() if state: return _version_map_from_state(state, kind_lower) @@ -100,11 +99,11 @@ async def get_country_versions(kind: str, request: Request) -> VersionMap: # The ``dict(...)`` iteration is where the Modal Dict RPCs happen; the # from_name handle itself is lazy and free. if kind_lower == "policyengine": - with scopes.span(runtime, ANNUAL_IMPACT_STAGES.name(Stage.MODAL_DICT_READ)): + with runtime.span(ANNUAL_IMPACT_STAGES.name(Stage.MODAL_DICT_READ)): version_dict = optional_modal_dict(POLICYENGINE_VERSION_DICT_NAME) return VersionMap(dict(version_dict) if version_dict is not None else {}) - with scopes.span(runtime, ANNUAL_IMPACT_STAGES.name(Stage.MODAL_DICT_READ)): + with runtime.span(ANNUAL_IMPACT_STAGES.name(Stage.MODAL_DICT_READ)): version_dict = modal.Dict.from_name(f"simulation-api-{kind_lower}-versions") return VersionMap(dict(version_dict)) diff --git a/projects/policyengine-simulation-gateway/uv.lock b/projects/policyengine-simulation-gateway/uv.lock index a17de18fd..082898381 100644 --- a/projects/policyengine-simulation-gateway/uv.lock +++ b/projects/policyengine-simulation-gateway/uv.lock @@ -1261,11 +1261,11 @@ provides-extras = ["test", "build"] [[package]] name = "policyengine-observability" -version = "3.0.2" +version = "3.0.4" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/74/2a/0c261c8ca693bb3a13d7fe837d3f89f7fdda245496d4b3e92eefed3c972e/policyengine_observability-3.0.2.tar.gz", hash = "sha256:554f907e43d8eeb274c1298fbe990b1626ac45f194c49995ddde091ddcdc2a64", size = 125718, upload-time = "2026-09-30T21:36:02.028Z" } +sdist = { url = "https://files.pythonhosted.org/packages/c3/d0/09b91aca37fb3efd493c164082097e85705fcac2062d2b5ef8a8c4e2be04/policyengine_observability-3.0.4.tar.gz", hash = "sha256:7e43a5d0e9b9e5f683623643efca7c47e651baaad7b557a50ec8b344602dc61f", size = 127972, upload-time = "2026-10-27T17:23:24.343Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/a5/9c/b48fb64c1bebf83bd5fdeebe914188806c3513cf2de0fc95ee1d95082fc4/policyengine_observability-3.0.2-py3-none-any.whl", hash = "sha256:d84c7564922f6a3294268e89549065924edde932c31adf8b7c714558e107cb31", size = 42029, upload-time = "2026-09-30T21:36:00.841Z" }, + { url = "https://files.pythonhosted.org/packages/00/00/d9c4850a12b4d0bef72034316909aac97228de302384af8936a958b7dc8e/policyengine_observability-3.0.4-py3-none-any.whl", hash = "sha256:1b3ebbfd00c138abcc2ec9580b13aa55679b91745783c5567cc73f8b2bbaaaf7", size = 42204, upload-time = "2026-10-27T17:23:22.641Z" }, ] [package.optional-dependencies] @@ -1355,7 +1355,7 @@ requires-dist = [ { name = "modal", specifier = ">=1.4,<2" }, { name = "openapi-python-client", marker = "extra == 'build'", specifier = ">=0.21.6" }, { name = "packaging", specifier = ">=24" }, - { name = "policyengine-observability", extras = ["fastapi", "google", "otlp-grpc"], specifier = ">=3.0.2,<4" }, + { name = "policyengine-observability", extras = ["fastapi", "google", "otlp-grpc"], specifier = ">=3.0.4,<4" }, { name = "pydantic", specifier = ">=2.0" }, { name = "pyjwt", specifier = ">=2.10.1,<3.0.0" }, { name = "pyright", marker = "extra == 'build'", specifier = ">=1.1.401" }, @@ -1387,7 +1387,7 @@ requires-dist = [ { name = "black", marker = "extra == 'build'", specifier = ">=25.1.0" }, { name = "fastapi", specifier = ">=0.115.0" }, { name = "httpx2", marker = "extra == 'test'" }, - { name = "policyengine-observability", extras = ["fastapi", "google", "otlp-grpc"], specifier = ">=3.0.2,<4" }, + { name = "policyengine-observability", extras = ["fastapi", "google", "otlp-grpc"], specifier = ">=3.0.4,<4" }, { name = "pydantic", specifier = ">=2.0" }, { name = "pyright", marker = "extra == 'build'", specifier = ">=1.1.401" }, { name = "pytest", marker = "extra == 'test'", specifier = ">=8.3.4" },