Skip to content

Commit ac7f8fe

Browse files
authored
Merge pull request #965 from PolicyEngine/remove-span-none-presets
Remove the None presets and checks around Stage 12 span and operation blocks
2 parents 2a9212f + 8f87202 commit ac7f8fe

6 files changed

Lines changed: 173 additions & 57 deletions

File tree

‎changelog_entry.yaml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
- 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
1818
changed:
1919
- 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
20+
- Removed the None presets and None checks around span and operation blocks that only satisfied Pyright before policyengine-observability 3.0.4; the Stage 12 Modal functions run_single_simulation_segment_us, run_single_simulation_us, run_single_simulation_uk and coordinate_report return their result from inside the operation block, and the Stage 12 segment runner no longer sets the dispatched call, the collected payloads and the merged calculation to None before their spans or checks them afterwards, and returns the merge from inside its span; 3.0.4's scopes never suppress an exception, and each value comes from a function that never returns None, so the removed RuntimeErrors (returned no result, no call, no payloads or no calculation) could not be raised; new tests run each of those blocks through the scopes of the runtime the services build
2021
- 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
2122
- 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
2223
- 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

‎projects/policyengine-simulation-executor/src/modal/v2_app.py‎

Lines changed: 4 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -279,7 +279,6 @@ def run_single_simulation_segment_us(
279279

280280
try:
281281
propagated = normalize_observability_context(observability_context)
282-
result = None
283282
with runtime.operation(
284283
STAGE12_SEGMENT_STAGES.name(Stage.STAGE12_SEGMENT_EXECUTION),
285284
attributes={
@@ -289,13 +288,9 @@ def run_single_simulation_segment_us(
289288
},
290289
remote_context=propagated,
291290
):
292-
result = stage12_runtime.calculate_segment(
291+
return stage12_runtime.calculate_segment(
293292
payload, runtime=runtime
294293
).model_dump(mode="python")
295-
if result is None:
296-
msg = "Stage 12 segment calculation returned no result"
297-
raise RuntimeError(msg)
298-
return result
299294
finally:
300295
runtime.shutdown()
301296

@@ -336,7 +331,6 @@ def segmented_calculator(
336331

337332
try:
338333
propagated = normalize_observability_context(observability_context)
339-
result = None
340334
with runtime.operation(
341335
STAGE12_SIMULATION_STAGES.name(Stage.STAGE12_SIMULATION_EXECUTION),
342336
attributes={
@@ -345,7 +339,7 @@ def segmented_calculator(
345339
},
346340
remote_context=propagated,
347341
):
348-
result = stage12_runtime.run_single_simulation(
342+
return stage12_runtime.run_single_simulation(
349343
payload,
350344
context,
351345
required_country="us",
@@ -354,10 +348,6 @@ def segmented_calculator(
354348
runtime=runtime,
355349
),
356350
)
357-
if result is None:
358-
msg = "Stage 12 US simulation returned no result"
359-
raise RuntimeError(msg)
360-
return result
361351
finally:
362352
runtime.shutdown()
363353

@@ -386,7 +376,6 @@ def run_single_simulation_uk(
386376
)
387377
try:
388378
propagated = normalize_observability_context(observability_context)
389-
result = None
390379
with runtime.operation(
391380
STAGE12_SIMULATION_STAGES.name(Stage.STAGE12_SIMULATION_EXECUTION),
392381
attributes={
@@ -395,18 +384,14 @@ def run_single_simulation_uk(
395384
},
396385
remote_context=propagated,
397386
):
398-
result = stage12_runtime.run_single_simulation(
387+
return stage12_runtime.run_single_simulation(
399388
payload,
400389
context,
401390
required_country="uk",
402391
dependencies=stage12_runtime.SimulationWorkerDependencies(
403392
runtime=runtime
404393
),
405394
)
406-
if result is None:
407-
msg = "Stage 12 UK simulation returned no result"
408-
raise RuntimeError(msg)
409-
return result
410395
finally:
411396
runtime.shutdown()
412397

@@ -452,7 +437,6 @@ def coordinate_report(
452437
)
453438
try:
454439
propagated = normalize_observability_context(observability_context)
455-
result = None
456440
with runtime.operation(
457441
stage_plan.name(Stage.STAGE12_COORDINATOR_EXECUTION),
458442
attributes={
@@ -465,7 +449,7 @@ def coordinate_report(
465449
},
466450
remote_context=propagated,
467451
):
468-
result = stage12_runtime.coordinate_report(
452+
return stage12_runtime.coordinate_report(
469453
payload,
470454
context,
471455
parent,
@@ -475,9 +459,5 @@ def coordinate_report(
475459
),
476460
dependencies=stage12_runtime.CoordinatorDependencies(runtime=runtime),
477461
)
478-
if result is None:
479-
msg = "Stage 12 report coordinator returned no result"
480-
raise RuntimeError(msg)
481-
return result
482462
finally:
483463
runtime.shutdown()

‎projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/fanout.py‎

Lines changed: 2 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -42,14 +42,6 @@
4242
CONSECUTIVE_POLL_ERROR_LIMIT = 2
4343

4444

45-
def _require_call(call: modal.FunctionCall | None) -> modal.FunctionCall:
46-
"""Return ``call``, raising ``RuntimeError`` when a dispatch returned none."""
47-
if call is None:
48-
msg = "Stage 12 segment dispatch returned no call"
49-
raise RuntimeError(msg)
50-
return call
51-
52-
5345
def _next_backoff(current: float, *, maximum: float) -> float:
5446
return min(current * POLL_INTERVAL_BACKOFF_FACTOR, maximum)
5547

@@ -130,7 +122,6 @@ def run(self) -> SimulationCalculation:
130122
handles: list[tuple[int, Any]] = []
131123
try:
132124
for segment in segments:
133-
call = None
134125
with self.runtime.span(
135126
STAGE12_SIMULATION_STAGES.name(Stage.STAGE12_SEGMENT_DISPATCH),
136127
attributes={"segment_index": segment.segment_index},
@@ -139,32 +130,23 @@ def run(self) -> SimulationCalculation:
139130
segment.model_dump(mode="python"),
140131
observability_context=self.runtime.capture_context(),
141132
)
142-
handles.append((segment.segment_index, _require_call(call)))
133+
handles.append((segment.segment_index, call))
143134
except Exception:
144135
self._cancel_all(handles)
145136
raise
146137

147-
payloads: list[object] | None = None
148138
with self.runtime.span(
149139
STAGE12_SIMULATION_STAGES.name(Stage.STAGE12_SEGMENT_WAIT)
150140
):
151141
payloads = self._collect(handles)
152-
if payloads is None:
153-
msg = "Stage 12 segment wait returned no payloads"
154-
raise RuntimeError(msg)
155-
merged: SimulationCalculation | None = None
156142
with self.runtime.span(
157143
STAGE12_SIMULATION_STAGES.name(Stage.STAGE12_SEGMENT_MERGE)
158144
):
159-
merged = merge_segment_results(
145+
return merge_segment_results(
160146
self.simulation,
161147
payloads,
162148
groups=groups,
163149
)
164-
if merged is None:
165-
msg = "Stage 12 segment merge returned no calculation"
166-
raise RuntimeError(msg)
167-
return merged
168150

169151
def _collect(self, handles: list[tuple[int, Any]]) -> list[object]:
170152
results: list[object | None] = [None] * len(handles)

‎projects/policyengine-simulation-executor/tests/conftest.py‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,29 @@ def observability_runtime():
5050
return NoOpObservabilityRuntime()
5151

5252

53+
@pytest.fixture
54+
def real_observability_runtime(monkeypatch):
55+
"""Build the observability runtime the services build, exporting nothing.
56+
57+
Its span and operation scopes are the ones deployed code opens, so a test
58+
can check what a block inside them returns or raises. OpenTelemetry is
59+
disabled and Google Cloud logging is not configured.
60+
"""
61+
from policyengine_simulation_observability.observability import build_runtime
62+
63+
monkeypatch.setenv("OTEL_SDK_DISABLED", "true")
64+
monkeypatch.delenv("OBSERVABILITY_LOGGING_PROJECT_ID", raising=False)
65+
monkeypatch.delenv("OBSERVABILITY_LOG_NAME", raising=False)
66+
runtime = build_runtime(
67+
service_name="policyengine-simulation-executor-test",
68+
service_role="test",
69+
platform="local",
70+
environment="test",
71+
)
72+
yield runtime
73+
runtime.shutdown()
74+
75+
5376
@pytest.fixture(autouse=True)
5477
def _forget_spm_runtime_caches():
5578
"""Keep the SPM runtime memoization out of every other test's way.

‎projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py‎

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -398,6 +398,57 @@ def fail(*_, **__):
398398
assert runtime.shutdown_calls == 1
399399

400400

401+
@pytest.mark.parametrize(
402+
"function_name",
403+
[
404+
"run_single_simulation_us",
405+
"run_single_simulation_uk",
406+
"run_single_simulation_segment_us",
407+
"coordinate_report",
408+
],
409+
)
410+
def test_stage12_functions_return_the_result_from_the_real_operation_scope(
411+
monkeypatch,
412+
real_observability_runtime,
413+
function_name,
414+
) -> None:
415+
module = _load(monkeypatch)
416+
monkeypatch.setattr(
417+
module, "init_process_observability", lambda **_: real_observability_runtime
418+
)
419+
420+
result = _invoke(monkeypatch, module, function_name, lambda *_, **__: {"ok": True})
421+
422+
assert result == {"ok": True}
423+
424+
425+
@pytest.mark.parametrize(
426+
"function_name",
427+
[
428+
"run_single_simulation_us",
429+
"run_single_simulation_uk",
430+
"run_single_simulation_segment_us",
431+
"coordinate_report",
432+
],
433+
)
434+
def test_stage12_functions_raise_the_application_error_from_the_real_operation_scope(
435+
monkeypatch,
436+
real_observability_runtime,
437+
function_name,
438+
) -> None:
439+
module = _load(monkeypatch)
440+
monkeypatch.setattr(
441+
module, "init_process_observability", lambda **_: real_observability_runtime
442+
)
443+
444+
def fail(*_, **__):
445+
msg = "application failure"
446+
raise LookupError(msg)
447+
448+
with pytest.raises(LookupError, match="application failure"):
449+
_invoke(monkeypatch, module, function_name, fail)
450+
451+
401452
def test_external_package_values_are_assertions_only(monkeypatch) -> None:
402453
install_fake_modal(monkeypatch)
403454
monkeypatch.setenv("STAGE12_EXPECT_US_VERSION", "not-the-bundled-version")

0 commit comments

Comments
 (0)