From 94d801debde3ea30deec9c66a3f501919a0c7068 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 13 Oct 2026 23:20:05 +0400 Subject: [PATCH 1/5] Record a redacted Stage 12 failure cause instead of only the exception class Stage 12 coordinator, worker and result-comparison failures stored only type(error).__name__, so a database session limit, a UK aggregate validation defect and an invalid reform path all appeared as RuntimeError or ValidationError behind "Stage 12 report coordination failed". Failure records now store a bounded summary of the form " in|after (): ", capped at the 512-character error_summary limit. Messages are redacted first: URL credentials, signed URL query strings, secret-looking assignments, bearer and known token formats, this process's secret environment values, Pydantic input values and SQL statement parameters are removed. Stage tracking follows the existing registered span names, so the summary names the stage that failed or the last stage that finished before a failure between spans. The coordinator no longer overwrites a cause that the failed worker already persisted, and copies that cause onto the report record. The redacted cause, with the original stack, is also recorded on the active operation before the failure is persisted. The exceptions returned to Modal are unchanged. Co-Authored-By: Claude Opus 5.5 --- changelog_entry.yaml | 1 + .../stage12_runtime/comparison.py | 3 +- .../stage12_runtime/coordination.py | 130 +++++---- .../stage12_runtime/failure_causes.py | 209 +++++++++++++++ .../stage12_runtime/failure_recording.py | 15 +- .../stage12_runtime/simulation.py | 138 +++++----- .../tests/test_stage12_failure_causes.py | 248 ++++++++++++++++++ .../tests/test_stage12_failure_recording.py | 15 ++ .../tests/test_stage12_runtime.py | 219 +++++++++++++++- 9 files changed, 827 insertions(+), 151 deletions(-) create mode 100644 projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/failure_causes.py create mode 100644 projects/policyengine-simulation-executor/tests/test_stage12_failure_causes.py diff --git a/changelog_entry.yaml b/changelog_entry.yaml index f5595afdd..c2e911bf5 100644 --- a/changelog_entry.yaml +++ b/changelog_entry.yaml @@ -19,3 +19,4 @@ - Bound UK constituency and local-authority weight matrices to the certified data release of the selected dataset instead of the unversioned bucket copy, and stopped runs before they start when the matrix lacks the run year or household dimension - Released each Stage 12 database session after its operation instead of holding pooled sessions through calculations, retried opening each Stage 12 coordinator and worker database connection that the pooler refused without ever retrying a statement or commit, and reported Stage 12 failures to telemetry before recording them in the database so a failed write cannot hide the original error - Accepted a null SPM provenance in Stage 12 aggregate artifacts, so UK and other reports without SPM calculations are saved, while every supplied provenance record is still validated + - Recorded a redacted, bounded Stage 12 failure cause naming the exception, failing stage, simulation role and message instead of only the exception class, and kept a failed worker's recorded cause on its simulation and report 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 99055a349..1f76aa306 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 @@ -19,6 +19,7 @@ from policyengine_simulation_executor.stage12_result_comparison import compare_results from .dependencies import ChildInvoker, ComparisonStore +from .failure_causes import failure_summary PRODUCTION_RESULT_WAIT_TIMEOUT_SECONDS = 900 logger = logging.getLogger(__name__) @@ -122,7 +123,7 @@ def compare_completed_report( "comparison_schema_version": None, "comparison_completed_at": failed_at, "comparison_error_code": "result_comparison_failed", - "comparison_error_summary": type(error).__name__, + "comparison_error_summary": failure_summary(error), "updated_at": failed_at, } ) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py index 5ccf33c70..d3dc6ddb4 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py @@ -2,8 +2,8 @@ from __future__ import annotations -from collections.abc import Callable, Mapping -from contextlib import nullcontext +from collections.abc import Callable, Iterator, Mapping +from contextlib import contextmanager, nullcontext from datetime import UTC, datetime from functools import partial from hashlib import sha256 @@ -52,6 +52,7 @@ artifact_store, runtime_store, ) +from .failure_causes import StageProgress, failure_summary, record_failure_cause from .failure_recording import persist_failure_record, report_original_failure from .output_planning import ( plan_simulation_input, @@ -63,6 +64,11 @@ SIMULATION_WAIT_TIMEOUT_SECONDS = 3_000 +def _simulation_role(attributes: Mapping[str, object] | None) -> str | None: + role = attributes.get("simulation_role") if attributes is not None else None + return str(role) if role is not None else None + + def _child_record( *, simulation: PlannedSimulationExecutionInput, @@ -195,10 +201,22 @@ def _mark_child_failed( *, error_code: str, error_summary: str, -) -> None: +) -> str | None: + """Mark an unfinished child failed and return the cause its record keeps. + + A worker that failed has already persisted its own redacted cause before + Modal returned the worker's stable exception, so that cause is kept rather + than overwritten. A succeeded child is left unchanged and has no cause. + """ + latest = persistence.get_simulation(simulation_execution_id) if latest.status is ComparisonRunLifecycleStatus.SUCCEEDED: - return + return None + if ( + latest.status is ComparisonRunLifecycleStatus.FAILED + and latest.error_summary is not None + ): + return latest.error_summary failed_at = datetime.now(UTC) persistence.replace_simulation( latest.model_copy( @@ -211,6 +229,7 @@ def _mark_child_failed( } ) ) + return error_summary def coordinate_report( @@ -269,21 +288,45 @@ def coordinate_report( } descriptors: dict[str, SimulationArtifactDescriptor] = {} calls: dict[str, ChildCall] = {} + progress = StageProgress() # A child failure is reported where it occurs and then propagates to the # report handler below, which must not report the same error twice. reported_child_failure: Exception | None = None + # A failed child records its own cause; the report keeps that cause + # rather than the generic exception Modal returns to the coordinator. + child_failure_summary: str | None = None try: + @contextmanager def stage_scope( stage: Stage, attributes: Mapping[str, object] | None = None, - ): - if runtime is None: - return nullcontext() - return runtime.operation( - stage_plan.name(stage), - attributes=attributes, + ) -> Iterator[None]: + stage_name = stage_plan.name(stage) + operation = ( + runtime.operation(stage_name, attributes=attributes) + if runtime is not None + else nullcontext() ) + with ( + progress.stage(stage_name, role=_simulation_role(attributes)), + operation, + ): + yield + + @contextmanager + def stage_span( + stage: Stage, + attributes: Mapping[str, object] | None = None, + ) -> Iterator[None]: + stage_name = stage_plan.name(stage) + span = ( + runtime.span(stage_name, attributes=attributes) + if runtime is not None + else nullcontext() + ) + with progress.stage(stage_name, role=_simulation_role(attributes)), span: + yield with stage_scope(Stage.STAGE12_COORDINATOR_PREPARATION): context = context.model_copy( @@ -317,12 +360,7 @@ def stage_scope( ): simulations.append(plan_simulation_input(simulation, output_plan)) planned_simulations = tuple(simulations) - child_state_span = ( - runtime.span(stage_plan.name(Stage.STAGE12_CHILD_STATE_CREATE)) - if runtime is not None - else nullcontext() - ) - with child_state_span: + with stage_span(Stage.STAGE12_CHILD_STATE_CREATE): children = { simulation.role: persistence.create_or_resolve_simulation( _child_record( @@ -377,15 +415,10 @@ def stage_scope( ) ) try: - dispatch_span = ( - runtime.span( - stage_plan.name(Stage.STAGE12_CHILD_DISPATCH), - attributes={"simulation_role": simulation.role.value}, - ) - if runtime is not None - else nullcontext() - ) - with dispatch_span: + with stage_span( + Stage.STAGE12_CHILD_DISPATCH, + {"simulation_role": simulation.role.value}, + ): child_observability_context = ( runtime.capture_context() if runtime is not None else None ) @@ -411,7 +444,7 @@ def stage_scope( persistence, simulation.simulation_execution_id, error_code="simulation_dispatch_failed", - error_summary=type(error).__name__, + error_summary=failure_summary(error, progress=progress), ), runtime=runtime, failure_scope="child_dispatch", @@ -427,15 +460,7 @@ def stage_scope( # Every required call has been started before the coordinator waits. for role, call in calls.items(): try: - wait_span = ( - runtime.span( - stage_plan.name(Stage.STAGE12_CHILD_WAIT), - attributes={"simulation_role": role}, - ) - if runtime is not None - else nullcontext() - ) - with wait_span: + with stage_span(Stage.STAGE12_CHILD_WAIT, {"simulation_role": role}): descriptors[role] = SimulationArtifactDescriptor.model_validate( call.get(timeout=SIMULATION_WAIT_TIMEOUT_SECONDS) ) @@ -450,13 +475,13 @@ def stage_scope( simulation = next( item for item in simulations if item.role.value == role ) - persist_failure_record( + child_failure_summary = persist_failure_record( partial( _mark_child_failed, persistence, simulation.simulation_execution_id, error_code="simulation_invocation_failed", - error_summary=type(error).__name__, + error_summary=failure_summary(error, progress=progress), ), runtime=runtime, failure_scope="child_wait", @@ -465,12 +490,7 @@ def stage_scope( baseline = descriptors["baseline"] reform = descriptors["reform"] validate_aligned_outputs(report, output_plan, baseline, reform) - artifact_read_span = ( - runtime.span(stage_plan.name(Stage.STAGE12_ARTIFACT_READ)) - if runtime is not None - else nullcontext() - ) - with artifact_read_span: + with stage_span(Stage.STAGE12_ARTIFACT_READ): baseline_payload = artifact_storage.read(baseline.artifact.uri) reform_payload = artifact_storage.read(reform.artifact.uri) if sha256(baseline_payload).hexdigest() != baseline.artifact.content_sha256: @@ -487,12 +507,7 @@ def stage_scope( ) validate_output_frames(baseline_frames, output_plan) validate_output_frames(reform_frames, output_plan) - aggregation_span = ( - runtime.span(stage_plan.name(Stage.STAGE12_AGGREGATION)) - if runtime is not None - else nullcontext() - ) - with aggregation_span: + with stage_span(Stage.STAGE12_AGGREGATION): aggregate = aggregator( report=report, baseline_frames=baseline_frames, @@ -501,12 +516,7 @@ def stage_scope( reform_descriptor=reform, uk_local_authority_metadata=uk_local_authority_metadata, ) - aggregate_write_span = ( - runtime.span(stage_plan.name(Stage.STAGE12_AGGREGATE_ARTIFACT_WRITE)) - if runtime is not None - else nullcontext() - ) - with aggregate_write_span: + with stage_span(Stage.STAGE12_AGGREGATE_ARTIFACT_WRITE): aggregate_artifact = artifact_storage.write_aggregate( prefix=context.artifact_prefix, payload=aggregate, @@ -543,8 +553,8 @@ def stage_scope( stage_plan=stage_plan, ) return descriptor.model_dump(mode="json") - # Persist any coordinator, child-call, aggregation, or artifact failure - # before returning a stable exception to Modal. + # Persist a redacted cause for any coordinator, child-call, aggregation, or + # artifact failure before returning a stable exception to Modal. except Exception as error: # noqa: BLE001 if error is not reported_child_failure: report_original_failure( @@ -552,6 +562,10 @@ def stage_scope( runtime=runtime, failure_scope="report_coordination", ) + report_failure_summary = child_failure_summary or failure_summary( + error, progress=progress + ) + record_failure_cause(runtime, error, report_failure_summary) failed_at = datetime.now(UTC) persist_failure_record( partial( @@ -561,7 +575,7 @@ def stage_scope( "status": ComparisonRunLifecycleStatus.FAILED, "aggregation_status": ComparisonRunAggregationStatus.FAILED, "error_code": "report_coordination_failed", - "error_summary": type(error).__name__, + "error_summary": report_failure_summary, "updated_at": failed_at, "completed_at": failed_at, } 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 new file mode 100644 index 000000000..d5362aaf5 --- /dev/null +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/failure_causes.py @@ -0,0 +1,209 @@ +"""Bounded, redacted failure causes for durable Stage 12 records. + +A Stage 12 failure record names the exception type, the stage the process +was in (or had last finished), the simulation role where one applies, and a +redacted message. Without the message an operator cannot tell a database +session limit from a validation defect, because both surface as a generic +exception class. + +Messages are redacted before they are stored or exported. Redaction removes +credentials embedded in URLs, signed URL query strings, secret-looking +key/value pairs, bearer and access tokens, the values of this process's +secret environment variables, Pydantic validation input values and SQL +statement parameters. The summary fits the 512-character ``error_summary`` +column and contract field. +""" + +from __future__ import annotations + +import logging +import os +import re +from collections.abc import Iterator +from contextlib import contextmanager +from urllib.parse import urlsplit + +from policyengine_observability import ObservabilityRuntime +from pydantic import ValidationError + +from policyengine_simulation_executor.hf_access_validation import HF_RUNTIME_ENV_NAME + +MAX_FAILURE_SUMMARY_LENGTH = 512 +REDACTED = "[REDACTED]" +TRUNCATION_MARKER = "..." +SECRET_ENVIRONMENT_VARIABLES = ("STAGE12_DATABASE_URL", HF_RUNTIME_ENV_NAME) +MIN_SECRET_VALUE_LENGTH = 8 +logger = logging.getLogger(__name__) + +_PRIVATE_KEY_BLOCK = re.compile( + r"-----BEGIN [A-Z ]*PRIVATE KEY-----.*?(?:-----END [A-Z ]*PRIVATE KEY-----|$)", + re.DOTALL, +) +_URL_CREDENTIALS = re.compile(r"(?P\b[a-zA-Z][a-zA-Z0-9+.-]*://)[^\s/@'\"]+@") +_URL_QUERY = re.compile( + r"(?P
\b[a-zA-Z][a-zA-Z0-9+.-]*://[^\s?'\"]+)\?[^\s'\"]+" +) +_BEARER_TOKEN = re.compile(r"(?i)\bbearer\s+[A-Za-z0-9._~+/=-]+") +_SECRET_ASSIGNMENT = re.compile( + r"(?i)\b(?P[a-z0-9_-]*(?:password|passwd|pwd|secret|token|api[_-]?key" + r"|access[_-]?key|private[_-]?key|authorization|credentials?))" + r"(?P\s*[=:]\s*)(?P\"[^\"]*\"|'[^']*'|[^\s,;)\]}]+)" +) +_KNOWN_TOKEN_FORMATS = re.compile( + r"\b(?:hf_[A-Za-z0-9]{20,}" + r"|gh[pousr]_[A-Za-z0-9]{20,}" + r"|AIza[0-9A-Za-z_-]{35}" + r"|ya29\.[0-9A-Za-z_-]+" + r"|eyJ[A-Za-z0-9_-]+\.[A-Za-z0-9_-]+\.[A-Za-z0-9_-]+)" +) +_PYDANTIC_INPUT_VALUE = re.compile(r"input_value=.*?(?=, input_type=)", re.DOTALL) +_SQL_PARAMETERS = re.compile(r"\[parameters: .*?\](?=\s*(?:\(Background|$))", re.DOTALL) +_SQLALCHEMY_BACKGROUND = re.compile(r"\(Background on this error at: [^)]*\)") +_WHITESPACE = re.compile(r"\s+") + + +class StageProgress: + """Remember the last Stage 12 stage a process entered and whether it ended. + + A failure inside a stage is reported as ``in ``. A failure between + stages is reported as ``after ``, naming the last stage that + completed, so code between two spans is not attributed to either of them. + """ + + def __init__(self, *, role: str | None = None) -> None: + self.process_role = role + self.stage_name: str | None = None + self.stage_role: str | None = role + self.stage_finished = False + + @contextmanager + def stage(self, stage_name: str, *, role: str | None = None) -> Iterator[None]: + stage_role = role if role is not None else self.process_role + self.stage_name = stage_name + self.stage_role = stage_role + self.stage_finished = False + yield + # An exception skips this, leaving the failing stage recorded. + if self.stage_name == stage_name and self.stage_role == stage_role: + self.stage_finished = True + + +class Stage12FailureCause(Exception): + """A redacted stand-in for a failure, safe to export to telemetry.""" + + +def redact_failure_message(message: str) -> str: + """Remove credentials, tokens and input values from an exception message.""" + + redacted = _PRIVATE_KEY_BLOCK.sub(REDACTED, message) + for secret in _secret_environment_values(): + redacted = redacted.replace(secret, REDACTED) + redacted = _URL_CREDENTIALS.sub(rf"\g{REDACTED}@", redacted) + redacted = _URL_QUERY.sub(rf"\g
?{REDACTED}", redacted) + redacted = _BEARER_TOKEN.sub(f"Bearer {REDACTED}", redacted) + redacted = _SECRET_ASSIGNMENT.sub( + rf"\g\g{REDACTED}", + redacted, + ) + redacted = _KNOWN_TOKEN_FORMATS.sub(REDACTED, redacted) + redacted = _PYDANTIC_INPUT_VALUE.sub(f"input_value={REDACTED}", redacted) + redacted = _SQL_PARAMETERS.sub(f"[parameters: {REDACTED}]", redacted) + redacted = _SQLALCHEMY_BACKGROUND.sub("", redacted) + return _WHITESPACE.sub(" ", redacted).strip() + + +def failure_summary( + error: BaseException, + *, + progress: StageProgress | None = None, +) -> str: + """Return a redacted `` in (): `` summary.""" + + error_type = type(error).__name__ + try: + heading = f"{error_type}{_location(progress)}" + message = redact_failure_message(_failure_message(error)) + return _bounded(f"{heading}: {message}" if message else heading) + except Exception: # noqa: BLE001 - a summary must never mask the failure + return _bounded(error_type) + + +def record_failure_cause( + runtime: ObservabilityRuntime | None, + error: BaseException, + summary: str, +) -> None: + """Attach the redacted cause, with the original stack, to the operation.""" + + if runtime is None: + return + try: + cause = Stage12FailureCause(summary).with_traceback(error.__traceback__) + runtime.record_exception(cause, handled=True) + except Exception: # noqa: BLE001 - telemetry must not change the outcome + logger.warning("stage12_failure_cause_not_recorded") + + +def _location(progress: StageProgress | None) -> str: + if progress is None: + return "" + location = "" + if progress.stage_name is not None: + position = "after" if progress.stage_finished else "in" + location = f" {position} {progress.stage_name}" + if progress.stage_role is not None: + location = f"{location} ({progress.stage_role})" + return location + + +def _failure_message(error: BaseException) -> str: + if isinstance(error, ValidationError): + return _validation_message(error) + database_error = getattr(error, "orig", None) + if isinstance(database_error, BaseException) and hasattr(error, "statement"): + # SQLAlchemy wraps the driver error with the statement and parameters. + # The driver message carries the cause without stored values. + return f"{type(database_error).__name__}: {database_error}" + return str(error) + + +def _validation_message(error: ValidationError) -> str: + details = [] + for detail in error.errors( + include_url=False, + include_context=False, + include_input=False, + ): + description = f"{detail['msg']} [{detail['type']}]" + location = ".".join(str(part) for part in detail["loc"]) + details.append(f"{location}: {description}" if location else description) + count = error.error_count() + noun = "error" if count == 1 else "errors" + return f"{count} validation {noun} for {error.title}: {'; '.join(details)}" + + +def _secret_environment_values() -> tuple[str, ...]: + values: list[str] = [] + for name in SECRET_ENVIRONMENT_VARIABLES: + value = os.environ.get(name, "") + values.append(value) + values.append(_url_password(value) or "") + # Very short values would redact ordinary words; URL credentials are + # removed by pattern regardless of length. + secrets = {value for value in values if len(value) >= MIN_SECRET_VALUE_LENGTH} + # Replace longer values first so a password inside a URL is not split. + return tuple(sorted(secrets, key=len, reverse=True)) + + +def _url_password(value: str) -> str | None: + try: + return urlsplit(value).password + except ValueError: + return None + + +def _bounded(summary: str) -> str: + if len(summary) <= MAX_FAILURE_SUMMARY_LENGTH: + return summary + kept = MAX_FAILURE_SUMMARY_LENGTH - len(TRUNCATION_MARKER) + return f"{summary[:kept]}{TRUNCATION_MARKER}" diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/failure_recording.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/failure_recording.py index 342c315f4..b86c0679c 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/failure_recording.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/failure_recording.py @@ -33,22 +33,24 @@ def report_original_failure( ) -def persist_failure_record( - write_failure_record: Callable[[], object], +def persist_failure_record[WriteResult]( + write_failure_record: Callable[[], WriteResult], *, runtime: ObservabilityRuntime | None, failure_scope: str, -) -> None: +) -> WriteResult | None: """Attempt a failure write without letting its error replace the original. - The caller raises the same failure after this returns, whether or not the + Returns what the write returned, or ``None`` when the write failed. The + caller raises the same failure after this returns, whether or not the record was written: a child handler re-raises the original error, and the report and simulation handlers raise their stable ``RuntimeError``. The - original cause then remains visible only in telemetry. + original exception then remains visible only in telemetry; the record + holds its redacted summary. """ try: - write_failure_record() + return write_failure_record() except Exception as persistence_error: # noqa: BLE001 _emit_error( "Stage 12 failure record could not be persisted", @@ -56,6 +58,7 @@ def persist_failure_record( runtime=runtime, attributes={"stage12_failure_scope": failure_scope}, ) + return None def _emit_error( diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py index 139ce06cd..f593c796e 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py @@ -2,8 +2,8 @@ from __future__ import annotations -from collections.abc import Callable, Mapping -from contextlib import nullcontext +from collections.abc import Callable, Iterator, Mapping +from contextlib import contextmanager, nullcontext from dataclasses import dataclass from datetime import UTC, datetime from functools import partial @@ -39,6 +39,7 @@ ) from .dependencies import ComparisonStore, artifact_store, runtime_store +from .failure_causes import StageProgress, failure_summary, record_failure_cause from .failure_recording import persist_failure_record, report_original_failure from .output_planning import apply_output_plan, validate_output_frames @@ -107,13 +108,29 @@ def _require_installed_bundle(simulation: SimulationExecutionInput) -> None: raise RuntimeError("simulation provenance differs from the installed bundle") +@contextmanager +def _simulation_stage( + stage: Stage, + *, + runtime: ObservabilityRuntime | None, + progress: StageProgress, +) -> Iterator[None]: + stage_name = STAGE12_SIMULATION_STAGES.name(stage) + span = runtime.span(stage_name) if runtime is not None else nullcontext() + with progress.stage(stage_name), span: + yield + + def calculate_simulation_frames( simulation: PlannedSimulationExecutionInput, *, runtime: ObservabilityRuntime | None = None, + progress: StageProgress | None = None, ) -> SimulationCalculation: """Run one policy and return the coordinator-planned entity output tables.""" + if progress is None: + progress = StageProgress(role=simulation.role.value) if simulation.population.kind != "dataset": raise ValueError("Stage 12 society-wide worker requires a dataset population") if simulation.population.artifact.uri != simulation.bundle.dataset.uri: @@ -138,46 +155,30 @@ def calculate_simulation_frames( setup_gcp_credentials, ) - credential_span = ( - runtime.span(STAGE12_SIMULATION_STAGES.name(Stage.CREDENTIAL_SETUP)) - if runtime is not None - else nullcontext() - ) - with credential_span, setup_gcp_credentials(): - country_span = ( - runtime.span(STAGE12_SIMULATION_STAGES.name(Stage.COUNTRY_MODULE_LOAD)) - if runtime is not None - else nullcontext() - ) - with country_span: + with ( + _simulation_stage(Stage.CREDENTIAL_SETUP, runtime=runtime, progress=progress), + setup_gcp_credentials(), + ): + with _simulation_stage( + Stage.COUNTRY_MODULE_LOAD, runtime=runtime, progress=progress + ): country_module = _country_module(country) - region_span = ( - runtime.span(STAGE12_SIMULATION_STAGES.name(Stage.REGION_RESOLUTION)) - if runtime is not None - else nullcontext() - ) - with region_span: + with _simulation_stage( + Stage.REGION_RESOLUTION, runtime=runtime, progress=progress + ): region = _resolve_region( country_module=country_module, country=country, params=params, ) - dataset_resolution_span = ( - runtime.span(STAGE12_SIMULATION_STAGES.name(Stage.DATASET_RESOLUTION)) - if runtime is not None - else nullcontext() - ) - with dataset_resolution_span: + with _simulation_stage( + Stage.DATASET_RESOLUTION, runtime=runtime, progress=progress + ): dataset_selection = _resolve_dataset_selection( params, region_resolution=region, ) - dataset_span = ( - runtime.span(STAGE12_SIMULATION_STAGES.name(Stage.DATASET_LOAD)) - if runtime is not None - else nullcontext() - ) - with dataset_span: + with _simulation_stage(Stage.DATASET_LOAD, runtime=runtime, progress=progress): dataset = _load_dataset( params, selection=dataset_selection, @@ -191,19 +192,13 @@ def calculate_simulation_frames( _require_certified_uk_weight_matrix( region.scoping_strategy, dataset, dataset_selection ) - policy_span = ( - runtime.span(STAGE12_SIMULATION_STAGES.name(Stage.POLICY_NORMALIZATION)) - if runtime is not None - else nullcontext() - ) - with policy_span: + with _simulation_stage( + Stage.POLICY_NORMALIZATION, runtime=runtime, progress=progress + ): policy = _normalise_policy(simulation.policy) - build_span = ( - runtime.span(STAGE12_SIMULATION_STAGES.name(Stage.SIMULATION_BUILD)) - if runtime is not None - else nullcontext() - ) - with build_span: + with _simulation_stage( + Stage.SIMULATION_BUILD, runtime=runtime, progress=progress + ): model = _build_simulation( params, dataset=dataset, @@ -212,12 +207,9 @@ def calculate_simulation_frames( scoping_strategy=region.scoping_strategy, region_code=region.code, ) - calculation_span = ( - runtime.span(STAGE12_SIMULATION_STAGES.name(Stage.STAGE12_CALCULATION)) - if runtime is not None - else nullcontext() - ) - with calculation_span: + with _simulation_stage( + Stage.STAGE12_CALCULATION, runtime=runtime, progress=progress + ): apply_output_plan(model, simulation.output_plan) model.ensure() output_data = getattr(getattr(model, "output_dataset", None), "data", None) @@ -322,26 +314,25 @@ def run_single_simulation( } ) persistence.replace_simulation(running) + progress = StageProgress(role=simulation.role.value) try: - input_span = ( - runtime.span(STAGE12_SIMULATION_STAGES.name(Stage.STAGE12_INPUT_WRITE)) - if runtime is not None - else nullcontext() - ) - with input_span: + with _simulation_stage( + Stage.STAGE12_INPUT_WRITE, runtime=runtime, progress=progress + ): artifact_storage.write_input( prefix=context.artifact_prefix, simulation=simulation, ) if calculator is calculate_simulation_frames: - calculated = calculate_simulation_frames(simulation, runtime=runtime) - else: - calculation_span = ( - runtime.span(STAGE12_SIMULATION_STAGES.name(Stage.STAGE12_CALCULATION)) - if runtime is not None - else nullcontext() + calculated = calculate_simulation_frames( + simulation, + runtime=runtime, + progress=progress, ) - with calculation_span: + else: + with _simulation_stage( + Stage.STAGE12_CALCULATION, runtime=runtime, progress=progress + ): calculated = calculator(simulation) if isinstance(calculated, SimulationCalculation): frames = calculated.frames @@ -352,14 +343,11 @@ def run_single_simulation( calculation_provenance = None uk_local_authority_metadata = None validate_output_frames(frames, simulation.output_plan) - artifact_span = ( - runtime.span( - STAGE12_SIMULATION_STAGES.name(Stage.STAGE12_SIMULATION_ARTIFACT_WRITE) - ) - if runtime is not None - else nullcontext() - ) - with artifact_span: + with _simulation_stage( + Stage.STAGE12_SIMULATION_ARTIFACT_WRITE, + runtime=runtime, + progress=progress, + ): descriptor = artifact_storage.write_simulation( prefix=context.artifact_prefix, simulation=simulation, @@ -384,8 +372,8 @@ def run_single_simulation( ) ) return descriptor.model_dump(mode="json") - # Persist any country-package calculation failure before returning a - # stable exception to Modal. + # Persist a redacted cause for any country-package calculation failure + # before returning a stable exception to Modal. except Exception as error: # noqa: BLE001 report_original_failure( error, @@ -393,6 +381,8 @@ def run_single_simulation( failure_scope="simulation_execution", attributes={"simulation_role": simulation.role.value}, ) + execution_failure_summary = failure_summary(error, progress=progress) + record_failure_cause(runtime, error, execution_failure_summary) failed_at = datetime.now(UTC) persist_failure_record( partial( @@ -401,7 +391,7 @@ def run_single_simulation( update={ "status": ComparisonRunLifecycleStatus.FAILED, "error_code": "simulation_execution_failed", - "error_summary": type(error).__name__, + "error_summary": execution_failure_summary, "updated_at": failed_at, "completed_at": failed_at, } diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_failure_causes.py b/projects/policyengine-simulation-executor/tests/test_stage12_failure_causes.py new file mode 100644 index 000000000..bd4071e60 --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_stage12_failure_causes.py @@ -0,0 +1,248 @@ +"""Tests for the redacted, bounded causes persisted on Stage 12 failures.""" + +from __future__ import annotations + +import pytest +from policyengine_simulation_contract.spm import SPMComparisonProvenance +from pydantic import BaseModel, ValidationError +from sqlalchemy.exc import OperationalError + +from policyengine_simulation_executor.stage12_runtime.failure_causes import ( + MAX_FAILURE_SUMMARY_LENGTH, + Stage12FailureCause, + StageProgress, + failure_summary, + record_failure_cause, +) + +DATABASE_PASSWORD = "pooler-password-value" +DATABASE_URL = ( + f"postgresql://postgres.project:{DATABASE_PASSWORD}" + "@pooler.example.com:5432/postgres" +) +SESSION_LIMIT_MESSAGE = ( + 'connection to server at "pooler.example.com" (10.0.0.1), port 5432 failed: ' + "FATAL: (EMAXCONNSESSION) max clients reached in session mode - " + "max clients are limited to pool_size: 15" +) + + +class DriverOperationalError(Exception): + """Stands in for the database driver's own connection error.""" + + +class RecordingRuntime: + def __init__(self) -> None: + self.recorded_errors: list[tuple[Exception, bool]] = [] + + def record_exception(self, error, *, handled, status_code=None) -> None: + self.recorded_errors.append((error, handled)) + + +def _raised(error: Exception) -> Exception: + try: + raise error + except Exception as caught: # noqa: BLE001 + return caught + + +def test_summary_names_type_stage_role_and_message() -> None: + progress = StageProgress(role="reform") + + with pytest.raises(ValueError) as caught, progress.stage("stage12_calculation"): + raise ValueError("Variable 'country' not found in household data") + + assert failure_summary(caught.value, progress=progress) == ( + "ValueError in stage12_calculation (reform): " + "Variable 'country' not found in household data" + ) + + +def test_summary_names_the_last_finished_stage_for_a_failure_between_stages() -> None: + progress = StageProgress(role="baseline") + with progress.stage("stage12_dataset_load"): + pass + + summary = failure_summary(KeyError("region"), progress=progress) + + assert summary == "KeyError after stage12_dataset_load (baseline): 'region'" + + +def test_outer_stage_does_not_hide_the_finished_inner_stage() -> None: + progress = StageProgress() + with ( + progress.stage("stage12_credential_setup"), + progress.stage("stage12_dataset_load"), + ): + pass + + assert progress.stage_name == "stage12_dataset_load" + assert progress.stage_finished + + +def test_database_url_credentials_are_redacted_from_the_summary() -> None: + error = RuntimeError(f"could not connect using {DATABASE_URL}") + + summary = failure_summary(error) + + assert DATABASE_PASSWORD not in summary + assert "postgres.project" not in summary + assert "postgresql://[REDACTED]@pooler.example.com:5432/postgres" in summary + + +def test_connection_string_password_assignment_is_redacted() -> None: + error = RuntimeError( + f"invalid dsn: host=pooler.example.com user=postgres password={DATABASE_PASSWORD}" + ) + + summary = failure_summary(error) + + assert DATABASE_PASSWORD not in summary + assert "password=[REDACTED]" in summary + + +def test_secret_environment_values_are_redacted(monkeypatch) -> None: + monkeypatch.setenv("STAGE12_DATABASE_URL", DATABASE_URL) + monkeypatch.setenv("HUGGING_FACE_TOKEN", "private-dataset-read-credential") + error = RuntimeError( + f"authentication failed for {DATABASE_PASSWORD} " + "using private-dataset-read-credential" + ) + + summary = failure_summary(error) + + assert DATABASE_PASSWORD not in summary + assert "private-dataset-read-credential" not in summary + assert summary == ( + "RuntimeError: authentication failed for [REDACTED] using [REDACTED]" + ) + + +@pytest.mark.parametrize( + ("message", "secret"), + [ + ("Authorization: Bearer abc.def-ghi", "abc.def-ghi"), + ( + "download https://storage.example.com/data.h5?X-Goog-Signature=signed123 failed", + "signed123", + ), + ("token rejected: hf_abcdefghijklmnopqrstuvwxyz0123", "hf_abcdefghij"), + ("api_key=very-private-key rejected", "very-private-key"), + ], +) +def test_tokens_and_signed_urls_are_redacted(message, secret) -> None: + summary = failure_summary(RuntimeError(message)) + + assert secret not in summary + assert "[REDACTED]" in summary + + +def test_validation_input_values_are_not_persisted() -> None: + class HouseholdIncome(BaseModel): + employment_income: float + + with pytest.raises(ValidationError) as caught: + HouseholdIncome.model_validate({"employment_income": "household-value-123"}) + + summary = failure_summary(caught.value) + + assert "household-value-123" not in summary + assert summary.startswith("ValidationError: 1 validation error for HouseholdIncome") + assert "employment_income: Input should be a valid number" in summary + + +def test_null_spm_provenance_validation_failure_is_specific() -> None: + progress = StageProgress() + with ( + pytest.raises(ValidationError) as caught, + progress.stage("stage12_aggregate_artifact_write"), + ): + SPMComparisonProvenance.model_validate(None) + + summary = failure_summary(caught.value, progress=progress) + + assert summary == ( + "ValidationError in stage12_aggregate_artifact_write: " + "1 validation error for SPMComparisonProvenance: " + "Input should be a valid dictionary or instance " + "of SPMComparisonProvenance [model_type]" + ) + + +def test_stringified_validation_input_values_are_redacted() -> None: + error = RuntimeError( + "1 validation error for Simulation\n Value error [type=value_error, " + "input_value={'income': 52000.0}, input_type=dict]" + ) + + summary = failure_summary(error) + + assert "52000" not in summary + assert "input_value=[REDACTED], input_type=dict" in summary + + +def test_database_errors_report_the_driver_cause_without_parameters() -> None: + error = OperationalError( + "INSERT INTO stage12_eval_reports (error_summary) VALUES (%(summary)s)", + {"summary": "stored-value-that-must-not-leak"}, + DriverOperationalError(SESSION_LIMIT_MESSAGE), + ) + + summary = failure_summary( + error, progress=_progress_in("stage12_child_state_create") + ) + + assert "stored-value-that-must-not-leak" not in summary + assert "Background on this error" not in summary + assert summary == ( + "OperationalError in stage12_child_state_create: " + f"DriverOperationalError: {SESSION_LIMIT_MESSAGE}" + ) + + +def test_summary_is_bounded_to_the_persisted_column_length() -> None: + summary = failure_summary( + ValueError("x" * 2_000), progress=StageProgress(role="reform") + ) + + assert len(summary) == MAX_FAILURE_SUMMARY_LENGTH + assert summary.endswith("...") + + +def test_unprintable_errors_still_produce_a_summary() -> None: + class UnprintableError(Exception): + def __str__(self) -> str: + raise RuntimeError("cannot render") + + assert failure_summary(UnprintableError()) == "UnprintableError" + + +def test_recorded_cause_is_redacted_and_keeps_the_original_stack() -> None: + runtime = RecordingRuntime() + error = _raised(RuntimeError(f"could not connect using {DATABASE_URL}")) + summary = failure_summary(error) + + record_failure_cause(runtime, error, summary) + + [(recorded, handled)] = runtime.recorded_errors + assert isinstance(recorded, Stage12FailureCause) + assert handled + assert str(recorded) == summary + assert DATABASE_PASSWORD not in str(recorded) + assert recorded.__traceback__ is error.__traceback__ + + +def test_recording_failures_never_escape() -> None: + class FailingRuntime: + def record_exception(self, *_args, **_kwargs) -> None: + raise RuntimeError("exporter unavailable") + + record_failure_cause(FailingRuntime(), ValueError("cause"), "ValueError: cause") + record_failure_cause(None, ValueError("cause"), "ValueError: cause") + + +def _progress_in(stage_name: str) -> StageProgress: + progress = StageProgress() + with pytest.raises(RuntimeError), progress.stage(stage_name): + raise RuntimeError("stage failed") + return progress diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_failure_recording.py b/projects/policyengine-simulation-executor/tests/test_stage12_failure_recording.py index 7ca71d018..20f041cca 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_failure_recording.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_failure_recording.py @@ -51,3 +51,18 @@ def fail_to_write_record() -> None: "Stage 12 failure record could not be persisted (report_coordination)" ) assert record.exc_info[1] is persistence_error + + +def test_written_failure_record_returns_what_the_write_returned() -> None: + worker_cause = "ValueError in stage12_calculation (baseline): country missing" + + def keep_worker_cause() -> str: + return worker_cause + + kept_cause = persist_failure_record( + keep_worker_cause, + runtime=None, + failure_scope="child_wait", + ) + + assert kept_cause == worker_cause diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py index cddfa82dd..5d53ee98e 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py @@ -54,6 +54,7 @@ stage12_output_plan_sha256, ) from policyengine_simulation_contract.spm import ( + SPMComparisonProvenance, SPMRuntimeVersions, SPMSelection, build_spm_provenance, @@ -230,6 +231,7 @@ def coordinate_report(*args, **kwargs): class RecordingRuntime: def __init__(self) -> None: self.operations: list[tuple[str, dict[str, object]]] = [] + self.recorded_exceptions: list[Exception] = [] @contextmanager def operation(self, name, *, attributes=None, **_kwargs): @@ -242,6 +244,10 @@ def span(self, *_args, **_kwargs): def capture_context(self): return {} + def record_exception(self, error, *, handled, status_code=None): + assert handled + self.recorded_exceptions.append(error) + def _context() -> Stage12InvocationContext: return Stage12InvocationContext( @@ -728,7 +734,66 @@ def fail(_simulation): child = store.children[simulation.simulation_execution_id] assert child.status is ComparisonRunLifecycleStatus.FAILED assert child.error_code == "simulation_execution_failed" - assert child.error_summary == "RuntimeError" + assert child.error_summary == ( + "RuntimeError in stage12_calculation (baseline): " + "contains sensitive calculation values" + ) + + +def test_single_worker_persists_and_records_a_redacted_cause() -> None: + store = FakeStore() + simulation = _planned_simulation(SimulationRole.REFORM) + store.children[simulation.simulation_execution_id] = _child(simulation) + runtime = RecordingRuntime() + + def fail(_simulation): + raise RuntimeError( + "could not reach postgresql://runtime:database-secret@db.example/postgres" + ) + + with pytest.raises(RuntimeError, match="Stage 12 simulation execution failed"): + run_single_simulation( + simulation.model_dump(mode="json"), + _context().model_dump(mode="json"), + required_country="us", + store=store, + artifacts=FakeArtifacts(), + calculator=fail, + runtime=runtime, + ) + + child = store.children[simulation.simulation_execution_id] + assert child.error_summary == ( + "RuntimeError in stage12_calculation (reform): " + "could not reach postgresql://[REDACTED]@db.example/postgres" + ) + assert "database-secret" not in child.error_summary + [recorded_cause] = runtime.recorded_exceptions + assert str(recorded_cause) == child.error_summary + + +def test_single_worker_names_the_failing_artifact_write_stage() -> None: + class FailingSimulationArtifacts(FakeArtifacts): + def write_simulation(self, **_kwargs): + raise OSError("bucket write rejected") + + store = FakeStore() + simulation = _planned_simulation(SimulationRole.BASELINE) + store.children[simulation.simulation_execution_id] = _child(simulation) + + with pytest.raises(RuntimeError, match="Stage 12 simulation execution failed"): + run_single_simulation( + simulation.model_dump(mode="json"), + _context().model_dump(mode="json"), + required_country="us", + store=store, + artifacts=FailingSimulationArtifacts(), + calculator=lambda _: _frames(), + ) + + assert store.children[simulation.simulation_execution_id].error_summary == ( + "OSError in stage12_simulation_artifact_write (baseline): bucket write rejected" + ) def _connection_limit_error() -> OperationalError: @@ -816,7 +881,11 @@ def fail(_simulation): "stage12_failure_scope": "simulation_execution", "simulation_role": "baseline", } - assert failure_write == ("failure_write", "ValueError") + assert failure_write == ( + "failure_write", + "ValueError in stage12_calculation (baseline): " + "country column missing from household data", + ) assert persistence_report[0] == "telemetry" assert persistence_report[1][0] == ( "Stage 12 failure record could not be persisted" @@ -1556,14 +1625,14 @@ def test_coordinator_never_writes_partial_aggregate_when_a_child_fails() -> None assert "sensitive" not in str(error.value) + expected_summary = "RuntimeError in stage12_child_wait (reform): child failed with sensitive values" assert artifacts.aggregate_writes == [] assert store.parent.status is ComparisonRunLifecycleStatus.FAILED - assert store.parent.error_summary == "RuntimeError" - assert "sensitive" not in store.parent.error_summary + assert store.parent.error_summary == expected_summary failed_child = store.children[REFORM_ID] assert failed_child.status is ComparisonRunLifecycleStatus.FAILED assert failed_child.error_code == "simulation_invocation_failed" - assert failed_child.error_summary == "RuntimeError" + assert failed_child.error_summary == expected_summary @pytest.mark.parametrize( @@ -1617,7 +1686,119 @@ def create_or_resolve_simulation(self, record): assert store.parent.status is ComparisonRunLifecycleStatus.FAILED assert store.parent.error_code == "report_coordination_failed" - assert store.parent.error_summary == "RuntimeError" + assert store.parent.error_summary == ( + "RuntimeError in stage12_child_state_create: sensitive database detail" + ) + + +class WorkerRecordedFailureCall: + """A child whose worker persisted its cause before Modal raised.""" + + def __init__(self, object_id, store, simulation_execution_id, worker_summary): + self.object_id = object_id + self.store = store + self.simulation_execution_id = simulation_execution_id + self.worker_summary = worker_summary + + def get(self, *, timeout=None): + child = self.store.children[self.simulation_execution_id] + self.store.children[self.simulation_execution_id] = child.model_copy( + update={ + "status": ComparisonRunLifecycleStatus.FAILED, + "error_code": "simulation_execution_failed", + "error_summary": self.worker_summary, + "completed_at": NOW, + } + ) + raise RuntimeError("Stage 12 simulation execution failed") + + +class UnreachedCall: + def __init__(self, object_id): + self.object_id = object_id + + def get(self, *, timeout=None): + raise AssertionError("the coordinator waits for baseline first") + + +class WorkerRecordedFailureInvoker: + def __init__(self, store, worker_summary): + self.store = store + self.worker_summary = worker_summary + + def spawn(self, *, simulation, **_): + planned = PlannedSimulationExecutionInput.model_validate(simulation) + if planned.role is SimulationRole.BASELINE: + return WorkerRecordedFailureCall( + "call-baseline", + self.store, + planned.simulation_execution_id, + self.worker_summary, + ) + return UnreachedCall("call-reform") + + def restore(self, invocation_id): + raise AssertionError(f"unexpected restored invocation {invocation_id}") + + +def test_coordinator_keeps_the_cause_recorded_by_a_failed_worker() -> None: + store = FakeStore() + worker_summary = ( + "ValueError in stage12_calculation (baseline): " + "Variable 'country' not found in household data" + ) + runtime = RecordingRuntime() + + with pytest.raises(RuntimeError, match="Stage 12 report coordination failed"): + coordinate_report( + _report().model_dump(mode="json"), + _context().model_dump(mode="json"), + _parent().model_dump(mode="json"), + application_name=_context().modal_application, + coordinator_invocation_id="coordinator-1", + store=store, + artifacts=FakeArtifacts(), + invoker=WorkerRecordedFailureInvoker(store, worker_summary), + aggregator=lambda **_: {"must": "not run"}, + runtime=runtime, + ) + + failed_child = store.children[BASELINE_ID] + assert failed_child.error_code == "simulation_execution_failed" + assert failed_child.error_summary == worker_summary + assert store.parent.error_code == "report_coordination_failed" + assert store.parent.error_summary == worker_summary + [recorded_cause] = runtime.recorded_exceptions + assert str(recorded_cause) == worker_summary + + +def test_coordinator_names_the_stage_that_rejected_the_aggregate() -> None: + class RejectingAggregateArtifacts(FakeArtifacts): + def write_aggregate(self, *, prefix, payload): + SPMComparisonProvenance.model_validate(payload["result"]["spm_provenance"]) + + store = FakeStore() + artifacts = RejectingAggregateArtifacts() + + with pytest.raises(RuntimeError, match="Stage 12 report coordination failed"): + coordinate_report( + _report().model_dump(mode="json"), + _context().model_dump(mode="json"), + _parent().model_dump(mode="json"), + application_name=_context().modal_application, + coordinator_invocation_id="coordinator-1", + store=store, + artifacts=artifacts, + invoker=ConcurrentInvoker(artifacts), + aggregator=lambda **_: {"result": {"spm_provenance": None}}, + ) + + assert store.parent.error_summary == ( + "ValidationError in stage12_aggregate_artifact_write: " + "1 validation error for SPMComparisonProvenance: " + "Input should be a valid dictionary or instance " + "of SPMComparisonProvenance [model_type]" + ) def test_coordinator_builds_its_store_with_the_coordinator_runtime( @@ -1850,7 +2031,10 @@ def fail_aggregation(**_): original_report, failure_write, persistence_report = ordered_events assert original_report[1][1] == {"stage12_failure_scope": "report_coordination"} assert original_report[1][2] is aggregation_error - assert failure_write == ("failure_write", "ValueError") + assert failure_write == ( + "failure_write", + "ValueError in stage12_aggregation: aggregate failed validation", + ) assert isinstance(persistence_report[1][2], OperationalError) assert store.parent.status is ComparisonRunLifecycleStatus.RUNNING @@ -1894,7 +2078,10 @@ def get_simulation(self, simulation_execution_id): assert ordered_events[1] == ("child_read", str(BASELINE_ID)) # The report record is still written after the child record could not be. assert store.parent.status is ComparisonRunLifecycleStatus.FAILED - assert store.parent.error_summary == "RuntimeError" + assert store.parent.error_summary == ( + "RuntimeError in stage12_child_wait (baseline): " + "child failed with sensitive values" + ) class RefusingDispatchInvoker: @@ -1946,9 +2133,13 @@ def test_coordinator_reports_a_child_dispatch_failure_once() -> None: baseline_child = store.children[BASELINE_ID] assert baseline_child.status is ComparisonRunLifecycleStatus.FAILED assert baseline_child.error_code == "simulation_dispatch_failed" - assert baseline_child.error_summary == "ConnectionError" + expected_summary = ( + "ConnectionError in stage12_child_dispatch (baseline): " + "Modal refused the dispatch" + ) + assert baseline_child.error_summary == expected_summary assert store.parent.status is ComparisonRunLifecycleStatus.FAILED - assert store.parent.error_summary == "ConnectionError" + assert store.parent.error_summary == expected_summary def test_coordinator_rejects_incompatible_row_identity() -> None: @@ -2081,8 +2272,12 @@ def test_coordinator_records_bounded_timeout_without_an_aggregate() -> None: assert invoker.events == ["spawn:baseline", "spawn:reform"] assert artifacts.aggregate_writes == [] + expected_summary = ( + "TimeoutError in stage12_child_wait (baseline): " + "contains sensitive timeout details" + ) assert store.parent.status is ComparisonRunLifecycleStatus.FAILED - assert store.parent.error_summary == "TimeoutError" + assert store.parent.error_summary == expected_summary timed_out_child = store.children[BASELINE_ID] assert timed_out_child.status is ComparisonRunLifecycleStatus.FAILED - assert timed_out_child.error_summary == "TimeoutError" + assert timed_out_child.error_summary == expected_summary From 5219b58277a1532ef66aa8fdacee72f04936eb81 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 14 Oct 2026 18:50:38 +0400 Subject: [PATCH 2/5] Redact more credential formats in Stage 12 failure causes Redact quoted JSON and dictionary keys, non-Bearer Authorization schemes and cookie headers, signed query strings without a URL scheme, and the row values that PostgreSQL constraint violations echo in their detail, from both the stored summary and the exported cause. Patterns that could start inside a run of word characters now start only where that run starts, so redaction is linear in the message length. Redaction also reads at most 4,096 characters: the word that limit cuts is dropped, and quoted values, SQL parameters and validation inputs whose closing delimiter was cut are redacted to the end. Take the summary length from the contract's ErrorSummary, detect SQLAlchemy errors by type, and state in the module docstring that span exception events are outside this redaction. Co-Authored-By: Claude Opus 5.5 --- .../stage12_runtime/failure_causes.py | 136 ++++++++--- .../tests/test_stage12_failure_causes.py | 220 +++++++++++++++++- 2 files changed, 324 insertions(+), 32 deletions(-) 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 d5362aaf5..901eb2864 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 @@ -6,12 +6,21 @@ session limit from a validation defect, because both surface as a generic exception class. -Messages are redacted before they are stored or exported. Redaction removes -credentials embedded in URLs, signed URL query strings, secret-looking -key/value pairs, bearer and access tokens, the values of this process's -secret environment variables, Pydantic validation input values and SQL -statement parameters. The summary fits the 512-character ``error_summary`` -column and contract field. +The summary stored in ``error_summary`` and the ``Stage12FailureCause`` +recorded to telemetry are redacted. Redaction removes credentials embedded in +URLs, the query strings of URLs and paths, secret-looking key/value pairs +(including quoted JSON and dictionary keys), ``Authorization`` and cookie +header values, bearer and other known access tokens, the values of this +process's secret environment variables, Pydantic validation input values, SQL +statement parameters and the row values that database constraint errors +echo. Redaction reads a bounded prefix of the message, and the summary fits +the contract's ``ErrorSummary`` length. + +Redaction covers only this summary and the cause recorded with it. The +exception events that ``runtime.span`` and ``runtime.operation`` scopes +export when an exception leaves them carry the original message, filtered +only by the observability package's own redaction; this module does not +change them. """ from __future__ import annotations @@ -24,11 +33,17 @@ from urllib.parse import urlsplit from policyengine_observability import ObservabilityRuntime -from pydantic import ValidationError +from policyengine_simulation_contract.stage12_execution import ErrorSummary +from pydantic import TypeAdapter, ValidationError +from sqlalchemy.exc import StatementError from policyengine_simulation_executor.hf_access_validation import HF_RUNTIME_ENV_NAME -MAX_FAILURE_SUMMARY_LENGTH = 512 +# The longest ``error_summary`` the contract accepts. +MAX_FAILURE_SUMMARY_LENGTH: int = TypeAdapter(ErrorSummary).json_schema()["maxLength"] +# Redaction reads at most this much of a message. The cap bounds the cost of +# redacting an arbitrarily long message and is far above the summary length. +MAX_REDACTED_MESSAGE_LENGTH = 4_096 REDACTED = "[REDACTED]" TRUNCATION_MARKER = "..." SECRET_ENVIRONMENT_VARIABLES = ("STAGE12_DATABASE_URL", HF_RUNTIME_ENV_NAME) @@ -39,25 +54,62 @@ r"-----BEGIN [A-Z ]*PRIVATE KEY-----.*?(?:-----END [A-Z ]*PRIVATE KEY-----|$)", re.DOTALL, ) -_URL_CREDENTIALS = re.compile(r"(?P\b[a-zA-Z][a-zA-Z0-9+.-]*://)[^\s/@'\"]+@") -_URL_QUERY = re.compile( - r"(?P
\b[a-zA-Z][a-zA-Z0-9+.-]*://[^\s?'\"]+)\?[^\s'\"]+" +# A pattern that could start anywhere in a run of word characters only starts +# where the run starts, so matching stays linear in the length of the message. +_URL_SCHEME = r"(?{_URL_SCHEME})[^\s/'\"]+@") +# Every URL matches, with or without a query, so a failed match never +# rescans the rest of a long address. +_URL_ADDRESS = re.compile( + rf"(?P
{_URL_SCHEME}[^\s?'\"]+)(?P\?[^\s'\"]+)?" +) +# A query string without a scheme, such as on a signed object path. +_QUERY_STRING = re.compile(r"\?(?=[\w.~%+-]++=)[^\s'\"]+") +# A quoted value may contain escaped quotes. A value whose closing quote was +# cut off runs to the end of the message. +_QUOTED_VALUE = r"\"(?:[^\"\\]|\\.)*(?:\"|\Z)|'(?:[^'\\]|\\.)*(?:'|\Z)" +# An unquoted header value can contain spaces, as in "Basic " or +# a Digest parameter list, so it runs to the end of the line. +_CREDENTIAL_HEADER = re.compile( + r"(?[\w-]*(?:authorization|cookie))" + r"(?P[\"']?\s*[=:]\s*)" + rf"(?P{_QUOTED_VALUE}|[^\r\n]*)", + re.IGNORECASE | re.DOTALL, ) _BEARER_TOKEN = re.compile(r"(?i)\bbearer\s+[A-Za-z0-9._~+/=-]+") +_SECRET_KEY = ( + r"[\w-]*(?:password|passwd|pwd|secret|token|api[_-]?key|access[_-]?key" + r"|private[_-]?key|credentials?|signature|sig)" +) _SECRET_ASSIGNMENT = re.compile( - r"(?i)\b(?P[a-z0-9_-]*(?:password|passwd|pwd|secret|token|api[_-]?key" - r"|access[_-]?key|private[_-]?key|authorization|credentials?))" - r"(?P\s*[=:]\s*)(?P\"[^\"]*\"|'[^']*'|[^\s,;)\]}]+)" + rf"(?{_SECRET_KEY})(?P[\"']?\s*[=:]\s*)" + rf"(?P{re.escape(REDACTED)}|{_QUOTED_VALUE}|[^\s,;)\]}}]+)", + re.IGNORECASE | re.DOTALL, ) _KNOWN_TOKEN_FORMATS = re.compile( r"\b(?:hf_[A-Za-z0-9]{20,}" r"|gh[pousr]_[A-Za-z0-9]{20,}" r"|AIza[0-9A-Za-z_-]{35}" r"|ya29\.[0-9A-Za-z_-]+" - r"|eyJ[A-Za-z0-9_-]+\.[A-Za-z0-9_-]+\.[A-Za-z0-9_-]+)" + r"|(?[^()\r\n]*(?:\([^()\r\n]*\)[^()\r\n]*)*)\)=\(.*", + re.DOTALL, ) -_PYDANTIC_INPUT_VALUE = re.compile(r"input_value=.*?(?=, input_type=)", re.DOTALL) -_SQL_PARAMETERS = re.compile(r"\[parameters: .*?\](?=\s*(?:\(Background|$))", re.DOTALL) +_FAILING_ROW_VALUES = re.compile(r"\bFailing row contains \(.*", re.DOTALL) _SQLALCHEMY_BACKGROUND = re.compile(r"\(Background on this error at: [^)]*\)") _WHITESPACE = re.compile(r"\s+") @@ -93,23 +145,36 @@ class Stage12FailureCause(Exception): def redact_failure_message(message: str) -> str: - """Remove credentials, tokens and input values from an exception message.""" + """Remove credentials, tokens and stored values from an exception message. - redacted = _PRIVATE_KEY_BLOCK.sub(REDACTED, message) + Only the first ``MAX_REDACTED_MESSAGE_LENGTH`` characters are redacted and + returned. The word that limit cuts is dropped, because part of a secret + may no longer match its pattern. + """ + + truncated = len(message) > MAX_REDACTED_MESSAGE_LENGTH + redacted = ( + _before_last_word(message[:MAX_REDACTED_MESSAGE_LENGTH]) + if truncated + else message + ) + redacted = _PRIVATE_KEY_BLOCK.sub(REDACTED, redacted) for secret in _secret_environment_values(): redacted = redacted.replace(secret, REDACTED) redacted = _URL_CREDENTIALS.sub(rf"\g{REDACTED}@", redacted) - redacted = _URL_QUERY.sub(rf"\g
?{REDACTED}", redacted) + redacted = _URL_ADDRESS.sub(_without_url_query, redacted) + redacted = _QUERY_STRING.sub(f"?{REDACTED}", redacted) + redacted = _CREDENTIAL_HEADER.sub(rf"\g\g{REDACTED}", redacted) redacted = _BEARER_TOKEN.sub(f"Bearer {REDACTED}", redacted) - redacted = _SECRET_ASSIGNMENT.sub( - rf"\g\g{REDACTED}", - redacted, - ) + redacted = _SECRET_ASSIGNMENT.sub(rf"\g\g{REDACTED}", redacted) redacted = _KNOWN_TOKEN_FORMATS.sub(REDACTED, redacted) redacted = _PYDANTIC_INPUT_VALUE.sub(f"input_value={REDACTED}", redacted) redacted = _SQL_PARAMETERS.sub(f"[parameters: {REDACTED}]", redacted) + redacted = _CONSTRAINT_KEY_VALUES.sub(rf"Key (\g)={REDACTED}", redacted) + redacted = _FAILING_ROW_VALUES.sub(f"Failing row contains {REDACTED}", redacted) redacted = _SQLALCHEMY_BACKGROUND.sub("", redacted) - return _WHITESPACE.sub(" ", redacted).strip() + redacted = _WHITESPACE.sub(" ", redacted).strip() + return f"{redacted}{TRUNCATION_MARKER}" if truncated else redacted def failure_summary( @@ -159,11 +224,11 @@ def _location(progress: StageProgress | None) -> str: def _failure_message(error: BaseException) -> str: if isinstance(error, ValidationError): return _validation_message(error) - database_error = getattr(error, "orig", None) - if isinstance(database_error, BaseException) and hasattr(error, "statement"): - # SQLAlchemy wraps the driver error with the statement and parameters. - # The driver message carries the cause without stored values. - return f"{type(database_error).__name__}: {database_error}" + if isinstance(error, StatementError) and error.orig is not None: + # SQLAlchemy wraps the driver error with the statement and its + # parameters, so only the driver error is reported. Its detail can + # still echo row values, which redaction removes. + return f"{type(error.orig).__name__}: {error.orig}" return str(error) @@ -182,6 +247,17 @@ def _validation_message(error: ValidationError) -> str: return f"{count} validation {noun} for {error.title}: {'; '.join(details)}" +def _without_url_query(url_match: re.Match[str]) -> str: + if url_match["query"] is None: + return url_match[0] + return f"{url_match['address']}?{REDACTED}" + + +def _before_last_word(text: str) -> str: + words = text.rsplit(maxsplit=1) + return words[0] if len(words) == 2 else "" + + def _secret_environment_values() -> tuple[str, ...]: values: list[str] = [] for name in SECRET_ENVIRONMENT_VARIABLES: diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_failure_causes.py b/projects/policyengine-simulation-executor/tests/test_stage12_failure_causes.py index bd4071e60..bc8fe2458 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_failure_causes.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_failure_causes.py @@ -2,17 +2,23 @@ from __future__ import annotations +import time + import pytest from policyengine_simulation_contract.spm import SPMComparisonProvenance -from pydantic import BaseModel, ValidationError -from sqlalchemy.exc import OperationalError +from policyengine_simulation_contract.stage12_execution import ErrorSummary +from pydantic import BaseModel, TypeAdapter, ValidationError +from sqlalchemy.exc import IntegrityError, OperationalError +from policyengine_simulation_executor.stage12_runtime import failure_causes from policyengine_simulation_executor.stage12_runtime.failure_causes import ( MAX_FAILURE_SUMMARY_LENGTH, + MAX_REDACTED_MESSAGE_LENGTH, Stage12FailureCause, StageProgress, failure_summary, record_failure_cause, + redact_failure_message, ) DATABASE_PASSWORD = "pooler-password-value" @@ -25,12 +31,22 @@ "FATAL: (EMAXCONNSESSION) max clients reached in session mode - " "max clients are limited to pool_size: 15" ) +# Shapes that made the earlier patterns backtrack quadratically. +PATHOLOGICAL_MESSAGE_UNITS = ("a-", "a.", "eyJ-", "a://", "s://a@") +# Redacting this many pathological characters took several seconds before the +# patterns were made linear. +UNCAPPED_PATHOLOGICAL_LENGTH = 20_000 +MAX_PATHOLOGICAL_SECONDS = 1.0 class DriverOperationalError(Exception): """Stands in for the database driver's own connection error.""" +class DriverIntegrityError(Exception): + """Stands in for the database driver's constraint violation.""" + + class RecordingRuntime: def __init__(self) -> None: self.recorded_errors: list[tuple[Exception, bool]] = [] @@ -137,6 +153,100 @@ def test_tokens_and_signed_urls_are_redacted(message, secret) -> None: assert "[REDACTED]" in summary +@pytest.mark.parametrize( + ("message", "expected"), + [ + ( + 'request body {"password": "json-password", "user": "runtime"}', + 'request body {"password": [REDACTED], "user": "runtime"}', + ), + ( + 'request body {"password":"escaped \\"quote\\" password"}', + 'request body {"password":[REDACTED]}', + ), + ("options {'token': 'dict-token'}", "options {'token': [REDACTED]}"), + ("options {'api_key': 'dict-api-key'}", "options {'api_key': [REDACTED]}"), + ( + "connect(password='two word password') failed", + "connect(password=[REDACTED]) failed", + ), + ], +) +def test_quoted_secret_keys_are_redacted(message, expected) -> None: + assert redact_failure_message(message) == expected + + +@pytest.mark.parametrize( + ("message", "expected"), + [ + ("Authorization: Basic YWxpY2U6c2VjcmV0", "Authorization: [REDACTED]"), + ("Authorization: Token opaque-api-token", "Authorization: [REDACTED]"), + ( + 'Authorization: Digest username="runtime", response="6629fae4"', + "Authorization: [REDACTED]", + ), + ("Authorization: opaque-header-value", "Authorization: [REDACTED]"), + ( + "Proxy-Authorization: Basic YWxpY2U6c2VjcmV0", + "Proxy-Authorization: [REDACTED]", + ), + ( + "headers {'Authorization': 'Basic YWxpY2U6c2VjcmV0', 'Accept': 'json'}", + "headers {'Authorization': [REDACTED], 'Accept': 'json'}", + ), + ("Cookie: session=opaque-session", "Cookie: [REDACTED]"), + ], +) +def test_authorization_header_values_are_redacted(message, expected) -> None: + assert redact_failure_message(message) == expected + + +@pytest.mark.parametrize( + ("message", "expected"), + [ + ( + "GET /bucket/data.h5?X-Goog-Algorithm=GOOG4&X-Goog-Signature=abc failed", + "GET /bucket/data.h5?[REDACTED] failed", + ), + ( + "read bucket/data.h5?X-Amz-Signature=abc rejected", + "read bucket/data.h5?[REDACTED] rejected", + ), + ( + "blob container/data.h5?sv=2024-01-01&sig=abc expired", + "blob container/data.h5?[REDACTED] expired", + ), + ( + "canonical query X-Goog-Signature=abc mismatch", + "canonical query X-Goog-Signature=[REDACTED] mismatch", + ), + ("request sig=abc rejected", "request sig=[REDACTED] rejected"), + ], +) +def test_signed_query_parameters_without_a_scheme_are_redacted( + message, expected +) -> None: + assert redact_failure_message(message) == expected + + +def test_url_password_containing_an_at_sign_is_redacted_whole() -> None: + message = "could not connect to postgresql://runtime:pass@word@db.example/postgres" + + assert redact_failure_message(message) == ( + "could not connect to postgresql://[REDACTED]@db.example/postgres" + ) + + +def test_secret_environment_value_in_an_assignment_is_redacted_once( + monkeypatch, +) -> None: + monkeypatch.setenv("STAGE12_DATABASE_URL", DATABASE_URL) + + summary = failure_summary(RuntimeError(f"password={DATABASE_PASSWORD} rejected")) + + assert summary == "RuntimeError: password=[REDACTED] rejected" + + def test_validation_input_values_are_not_persisted() -> None: class HouseholdIncome(BaseModel): employment_income: float @@ -200,6 +310,111 @@ def test_database_errors_report_the_driver_cause_without_parameters() -> None: ) +@pytest.mark.parametrize( + ("driver_message", "expected_message"), + [ + ( + ( + 'duplicate key value violates unique constraint "reports_pkey"\n' + "DETAIL: Key (evaluation_id)=(stored-identifier) already exists." + ), + ( + 'duplicate key value violates unique constraint "reports_pkey" ' + "DETAIL: Key (evaluation_id)=[REDACTED]" + ), + ), + ( + ( + 'duplicate key value violates unique constraint "users_email_key"\n' + "DETAIL: Key (lower(email))=(stored@example.com) already exists." + ), + ( + 'duplicate key value violates unique constraint "users_email_key" ' + "DETAIL: Key (lower(email))=[REDACTED]" + ), + ), + ( + ( + 'null value in column "status" violates not-null constraint\n' + "DETAIL: Failing row contains (stored-identifier, stored\nvalue, null)." + ), + ( + 'null value in column "status" violates not-null constraint ' + "DETAIL: Failing row contains [REDACTED]" + ), + ), + ], +) +def test_constraint_violations_keep_the_cause_without_row_values( + driver_message, expected_message +) -> None: + error = IntegrityError( + "INSERT INTO stage12_eval_reports (evaluation_id) VALUES (%(id)s)", + {"id": "stored-identifier"}, + DriverIntegrityError(driver_message), + ) + + summary = failure_summary(error) + + assert "stored" not in summary + assert summary == f"IntegrityError: DriverIntegrityError: {expected_message}" + + +@pytest.mark.parametrize( + ("secret_construct", "kept_length", "secret_fragment"), + [ + ( + "postgresql://runtime:cutpassword@db.example/postgres", + 30, + "cutpass", + ), + ("{'password': 'open sesame word'}", 28, "sesame"), + ("[parameters: {'summary': 'stored row value'}]", 36, "stored"), + ], +) +def test_a_secret_cut_by_the_redaction_limit_is_not_exposed( + secret_construct, kept_length, secret_fragment +) -> None: + padding_length = MAX_REDACTED_MESSAGE_LENGTH - kept_length + padding = ("padding " * padding_length)[:padding_length] + message = f"{padding}{secret_construct} and more" + kept_text = message[:MAX_REDACTED_MESSAGE_LENGTH] + assert secret_fragment in kept_text + + redacted = redact_failure_message(message) + + assert secret_fragment not in redacted + assert redacted.endswith("...") + + +@pytest.mark.parametrize("message_unit", PATHOLOGICAL_MESSAGE_UNITS) +def test_long_pathological_messages_are_summarised_quickly(message_unit) -> None: + message = message_unit * (1_000_000 // len(message_unit)) + + started = time.perf_counter() + summary = failure_summary(RuntimeError(message)) + elapsed = time.perf_counter() - started + + assert elapsed < MAX_PATHOLOGICAL_SECONDS + assert len(summary) <= MAX_FAILURE_SUMMARY_LENGTH + + +@pytest.mark.parametrize("message_unit", PATHOLOGICAL_MESSAGE_UNITS) +def test_redaction_patterns_are_linear_without_the_length_cap( + monkeypatch, message_unit +) -> None: + monkeypatch.setattr( + failure_causes, "MAX_REDACTED_MESSAGE_LENGTH", UNCAPPED_PATHOLOGICAL_LENGTH + ) + message = message_unit * (UNCAPPED_PATHOLOGICAL_LENGTH // len(message_unit)) + + started = time.perf_counter() + redact_failure_message(message) + elapsed = time.perf_counter() - started + + assert elapsed < MAX_PATHOLOGICAL_SECONDS + + def test_summary_is_bounded_to_the_persisted_column_length() -> None: summary = failure_summary( ValueError("x" * 2_000), progress=StageProgress(role="reform") @@ -207,6 +422,7 @@ def test_summary_is_bounded_to_the_persisted_column_length() -> None: assert len(summary) == MAX_FAILURE_SUMMARY_LENGTH assert summary.endswith("...") + assert TypeAdapter(ErrorSummary).validate_python(summary) == summary def test_unprintable_errors_still_produce_a_summary() -> None: From 71193731f82762b68c7922d0205ef0d72d8592eb Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 14 Oct 2026 18:51:25 +0400 Subject: [PATCH 3/5] Share one stage helper for Stage 12 spans and failure progress Replace the coordinator's stage_scope and stage_span closures and the worker's _simulation_stage with tracked_stage, which opens a stage's span or operation and records it as the process's progress. The coordinator claim span uses it too. Co-Authored-By: Claude Opus 5.5 --- .../stage12_runtime/coordination.py | 61 ++++----------- .../stage12_runtime/failure_causes.py | 46 +++++++++-- .../stage12_runtime/simulation.py | 77 ++++++++----------- 3 files changed, 86 insertions(+), 98 deletions(-) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py index d3dc6ddb4..ae1824909 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py @@ -2,8 +2,7 @@ from __future__ import annotations -from collections.abc import Callable, Iterator, Mapping -from contextlib import contextmanager, nullcontext +from collections.abc import Callable from datetime import UTC, datetime from functools import partial from hashlib import sha256 @@ -52,7 +51,12 @@ artifact_store, runtime_store, ) -from .failure_causes import StageProgress, failure_summary, record_failure_cause +from .failure_causes import ( + StageProgress, + failure_summary, + record_failure_cause, + tracked_stage, +) from .failure_recording import persist_failure_record, report_original_failure from .output_planning import ( plan_simulation_input, @@ -64,11 +68,6 @@ SIMULATION_WAIT_TIMEOUT_SECONDS = 3_000 -def _simulation_role(attributes: Mapping[str, object] | None) -> str | None: - role = attributes.get("simulation_role") if attributes is not None else None - return str(role) if role is not None else None - - def _child_record( *, simulation: PlannedSimulationExecutionInput, @@ -269,12 +268,15 @@ def coordinate_report( if context.production_function_call_id is not None else STAGE12_CANONICAL_REPORT_STAGES ) - claim_span = ( - runtime.span(stage_plan.name(Stage.STAGE12_COORDINATOR_CLAIM)) - if runtime is not None - else nullcontext() + progress = StageProgress() + stage_span = partial( + tracked_stage, + stage_plan=stage_plan, + progress=progress, + runtime=runtime, ) - with claim_span: + stage_scope = partial(stage_span, operation=True) + with stage_span(Stage.STAGE12_COORDINATOR_CLAIM): parent, should_execute = _claim_parent( submitted=submitted_parent, coordinator_invocation_id=coordinator_invocation_id, @@ -288,7 +290,6 @@ def coordinate_report( } descriptors: dict[str, SimulationArtifactDescriptor] = {} calls: dict[str, ChildCall] = {} - progress = StageProgress() # A child failure is reported where it occurs and then propagates to the # report handler below, which must not report the same error twice. reported_child_failure: Exception | None = None @@ -296,38 +297,6 @@ def coordinate_report( # rather than the generic exception Modal returns to the coordinator. child_failure_summary: str | None = None try: - - @contextmanager - def stage_scope( - stage: Stage, - attributes: Mapping[str, object] | None = None, - ) -> Iterator[None]: - stage_name = stage_plan.name(stage) - operation = ( - runtime.operation(stage_name, attributes=attributes) - if runtime is not None - else nullcontext() - ) - with ( - progress.stage(stage_name, role=_simulation_role(attributes)), - operation, - ): - yield - - @contextmanager - def stage_span( - stage: Stage, - attributes: Mapping[str, object] | None = None, - ) -> Iterator[None]: - stage_name = stage_plan.name(stage) - span = ( - runtime.span(stage_name, attributes=attributes) - if runtime is not None - else nullcontext() - ) - with progress.stage(stage_name, role=_simulation_role(attributes)), span: - yield - with stage_scope(Stage.STAGE12_COORDINATOR_PREPARATION): context = context.model_copy( update={ 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 901eb2864..b6e20b191 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 @@ -18,9 +18,10 @@ Redaction covers only this summary and the cause recorded with it. The exception events that ``runtime.span`` and ``runtime.operation`` scopes -export when an exception leaves them carry the original message, filtered -only by the observability package's own redaction; this module does not -change them. +export when an exception leaves them, and the original-failure record that +``failure_recording`` logs before a failure write, carry the original +message, filtered only by the observability package's own redaction; this +module does not change them. """ from __future__ import annotations @@ -28,12 +29,13 @@ import logging import os import re -from collections.abc import Iterator -from contextlib import contextmanager +from collections.abc import Iterator, Mapping +from contextlib import AbstractContextManager, contextmanager, nullcontext from urllib.parse import urlsplit from policyengine_observability import ObservabilityRuntime from policyengine_simulation_contract.stage12_execution import ErrorSummary +from policyengine_simulation_observability.stages import Stage, StagePlan from pydantic import TypeAdapter, ValidationError from sqlalchemy.exc import StatementError @@ -140,6 +142,35 @@ def stage(self, stage_name: str, *, role: str | None = None) -> Iterator[None]: self.stage_finished = True +@contextmanager +def tracked_stage( + stage: Stage, + attributes: Mapping[str, object] | None = None, + *, + stage_plan: StagePlan, + progress: StageProgress, + runtime: ObservabilityRuntime | None, + operation: bool = False, +) -> Iterator[None]: + """Run a block as one Stage 12 stage for telemetry and failure causes. + + The block runs inside the stage's span, or its operation when + ``operation`` is true, and ``progress`` records the stage so that a + failure names it. A ``simulation_role`` attribute sets the stage's role. + """ + + stage_name = stage_plan.name(stage) + scope: AbstractContextManager[object] + if runtime is None: + scope = nullcontext() + elif operation: + scope = runtime.operation(stage_name, attributes=attributes) + else: + scope = runtime.span(stage_name, attributes=attributes) + with progress.stage(stage_name, role=_simulation_role(attributes)), scope: + yield + + class Stage12FailureCause(Exception): """A redacted stand-in for a failure, safe to export to telemetry.""" @@ -209,6 +240,11 @@ def record_failure_cause( logger.warning("stage12_failure_cause_not_recorded") +def _simulation_role(attributes: Mapping[str, object] | None) -> str | None: + role = attributes.get("simulation_role") if attributes is not None else None + return str(role) if role is not None else None + + def _location(progress: StageProgress | None) -> str: if progress is None: return "" diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py index f593c796e..1ba2f57cb 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py @@ -2,8 +2,7 @@ from __future__ import annotations -from collections.abc import Callable, Iterator, Mapping -from contextlib import contextmanager, nullcontext +from collections.abc import Callable, Mapping from dataclasses import dataclass from datetime import UTC, datetime from functools import partial @@ -39,7 +38,12 @@ ) from .dependencies import ComparisonStore, artifact_store, runtime_store -from .failure_causes import StageProgress, failure_summary, record_failure_cause +from .failure_causes import ( + StageProgress, + failure_summary, + record_failure_cause, + tracked_stage, +) from .failure_recording import persist_failure_record, report_original_failure from .output_planning import apply_output_plan, validate_output_frames @@ -108,19 +112,6 @@ def _require_installed_bundle(simulation: SimulationExecutionInput) -> None: raise RuntimeError("simulation provenance differs from the installed bundle") -@contextmanager -def _simulation_stage( - stage: Stage, - *, - runtime: ObservabilityRuntime | None, - progress: StageProgress, -) -> Iterator[None]: - stage_name = STAGE12_SIMULATION_STAGES.name(stage) - span = runtime.span(stage_name) if runtime is not None else nullcontext() - with progress.stage(stage_name), span: - yield - - def calculate_simulation_frames( simulation: PlannedSimulationExecutionInput, *, @@ -131,6 +122,12 @@ def calculate_simulation_frames( if progress is None: progress = StageProgress(role=simulation.role.value) + stage_span = partial( + tracked_stage, + stage_plan=STAGE12_SIMULATION_STAGES, + progress=progress, + runtime=runtime, + ) if simulation.population.kind != "dataset": raise ValueError("Stage 12 society-wide worker requires a dataset population") if simulation.population.artifact.uri != simulation.bundle.dataset.uri: @@ -156,29 +153,23 @@ def calculate_simulation_frames( ) with ( - _simulation_stage(Stage.CREDENTIAL_SETUP, runtime=runtime, progress=progress), + stage_span(Stage.CREDENTIAL_SETUP), setup_gcp_credentials(), ): - with _simulation_stage( - Stage.COUNTRY_MODULE_LOAD, runtime=runtime, progress=progress - ): + with stage_span(Stage.COUNTRY_MODULE_LOAD): country_module = _country_module(country) - with _simulation_stage( - Stage.REGION_RESOLUTION, runtime=runtime, progress=progress - ): + with stage_span(Stage.REGION_RESOLUTION): region = _resolve_region( country_module=country_module, country=country, params=params, ) - with _simulation_stage( - Stage.DATASET_RESOLUTION, runtime=runtime, progress=progress - ): + with stage_span(Stage.DATASET_RESOLUTION): dataset_selection = _resolve_dataset_selection( params, region_resolution=region, ) - with _simulation_stage(Stage.DATASET_LOAD, runtime=runtime, progress=progress): + with stage_span(Stage.DATASET_LOAD): dataset = _load_dataset( params, selection=dataset_selection, @@ -192,13 +183,9 @@ def calculate_simulation_frames( _require_certified_uk_weight_matrix( region.scoping_strategy, dataset, dataset_selection ) - with _simulation_stage( - Stage.POLICY_NORMALIZATION, runtime=runtime, progress=progress - ): + with stage_span(Stage.POLICY_NORMALIZATION): policy = _normalise_policy(simulation.policy) - with _simulation_stage( - Stage.SIMULATION_BUILD, runtime=runtime, progress=progress - ): + with stage_span(Stage.SIMULATION_BUILD): model = _build_simulation( params, dataset=dataset, @@ -207,9 +194,7 @@ def calculate_simulation_frames( scoping_strategy=region.scoping_strategy, region_code=region.code, ) - with _simulation_stage( - Stage.STAGE12_CALCULATION, runtime=runtime, progress=progress - ): + with stage_span(Stage.STAGE12_CALCULATION): apply_output_plan(model, simulation.output_plan) model.ensure() output_data = getattr(getattr(model, "output_dataset", None), "data", None) @@ -315,10 +300,14 @@ def run_single_simulation( ) persistence.replace_simulation(running) progress = StageProgress(role=simulation.role.value) + stage_span = partial( + tracked_stage, + stage_plan=STAGE12_SIMULATION_STAGES, + progress=progress, + runtime=runtime, + ) try: - with _simulation_stage( - Stage.STAGE12_INPUT_WRITE, runtime=runtime, progress=progress - ): + with stage_span(Stage.STAGE12_INPUT_WRITE): artifact_storage.write_input( prefix=context.artifact_prefix, simulation=simulation, @@ -330,9 +319,7 @@ def run_single_simulation( progress=progress, ) else: - with _simulation_stage( - Stage.STAGE12_CALCULATION, runtime=runtime, progress=progress - ): + with stage_span(Stage.STAGE12_CALCULATION): calculated = calculator(simulation) if isinstance(calculated, SimulationCalculation): frames = calculated.frames @@ -343,11 +330,7 @@ def run_single_simulation( calculation_provenance = None uk_local_authority_metadata = None validate_output_frames(frames, simulation.output_plan) - with _simulation_stage( - Stage.STAGE12_SIMULATION_ARTIFACT_WRITE, - runtime=runtime, - progress=progress, - ): + with stage_span(Stage.STAGE12_SIMULATION_ARTIFACT_WRITE): descriptor = artifact_storage.write_simulation( prefix=context.artifact_prefix, simulation=simulation, From 9e1f2cc084a3899daeb075ea37ebae01253fe0fc Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 14 Oct 2026 18:51:32 +0400 Subject: [PATCH 4/5] Test dispatch and comparison failure summaries in Stage 12 Assert the redacted summary stored on a child and its report when Modal rejects a dispatch, and the comparison_error_summary stored when the production result cannot be read. Reword fixture failure messages that described themselves as sensitive, since the tests now assert those messages are stored. Co-Authored-By: Claude Opus 5.5 --- .../tests/test_stage12_runtime.py | 138 ++++++++++++++---- 1 file changed, 110 insertions(+), 28 deletions(-) diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py index 5d53ee98e..399afea31 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py @@ -716,7 +716,7 @@ def test_single_worker_exposes_only_a_bounded_failure() -> None: store.children[simulation.simulation_execution_id] = _child(simulation) def fail(_simulation): - raise RuntimeError("contains sensitive calculation values") + raise RuntimeError("country package calculation failed") with pytest.raises( RuntimeError, match="Stage 12 simulation execution failed" @@ -730,13 +730,13 @@ def fail(_simulation): calculator=fail, ) - assert "sensitive" not in str(error.value) + assert "country package calculation failed" not in str(error.value) child = store.children[simulation.simulation_execution_id] assert child.status is ComparisonRunLifecycleStatus.FAILED assert child.error_code == "simulation_execution_failed" assert child.error_summary == ( "RuntimeError in stage12_calculation (baseline): " - "contains sensitive calculation values" + "country package calculation failed" ) @@ -1254,7 +1254,7 @@ def run(): started = time.monotonic() time.sleep(0.05) if role == self.fail_role: - raise RuntimeError("child failed with sensitive values") + raise RuntimeError("child calculation failed") frames = _frames(100.0 if role == "baseline" else 120.0) if self.missing_column and role == "reform": frames["household"] = frames["household"].drop( @@ -1551,6 +1551,50 @@ def test_comparison_failure_does_not_change_successful_stage12_report() -> None: assert artifacts.comparison_writes == [] +class UnavailableProductionCall: + object_id = "production-job-1" + + def get(self, *, timeout=None): + raise ConnectionError( + "production result unavailable at " + "https://storage.example.com/result.json?X-Goog-Signature=signed-value" + ) + + +class UnavailableProductionInvoker(ConcurrentInvoker): + def restore(self, invocation_id): + if invocation_id == "production-job-1": + return UnavailableProductionCall() + return super().restore(invocation_id) + + +def test_comparison_failure_stores_a_redacted_cause() -> None: + store = FakeStore() + artifacts = FakeArtifacts() + context = _context().model_copy( + update={"production_function_call_id": "production-job-1"} + ) + + coordinate_report( + _report().model_dump(mode="json"), + context.model_dump(mode="json"), + _automatic_parent().model_dump(mode="json"), + application_name=context.modal_application, + coordinator_invocation_id="coordinator-1", + store=store, + artifacts=artifacts, + invoker=UnavailableProductionInvoker(artifacts), + aggregator=lambda **_: {"result": {"budget": 100}}, + ) + + assert store.parent.status is ComparisonRunLifecycleStatus.SUCCEEDED + assert store.parent.comparison_status is ResultComparisonStatus.FAILED + assert store.parent.comparison_error_summary == ( + "ConnectionError: production result unavailable at " + "https://storage.example.com/result.json?[REDACTED]" + ) + + def test_repeated_submission_does_not_retry_failed_comparison() -> None: store = FakeStore() parent = _automatic_parent() @@ -1623,9 +1667,11 @@ def test_coordinator_never_writes_partial_aggregate_when_a_child_fails() -> None aggregator=lambda **_: _aggregate_payload({"must": "not run"}), ) - assert "sensitive" not in str(error.value) + assert "child calculation failed" not in str(error.value) - expected_summary = "RuntimeError in stage12_child_wait (reform): child failed with sensitive values" + expected_summary = ( + "RuntimeError in stage12_child_wait (reform): child calculation failed" + ) assert artifacts.aggregate_writes == [] assert store.parent.status is ComparisonRunLifecycleStatus.FAILED assert store.parent.error_summary == expected_summary @@ -1668,7 +1714,7 @@ def test_coordinator_rejects_artifacts_that_do_not_satisfy_the_shared_plan( def test_coordinator_records_failure_when_child_persistence_fails() -> None: class FailingChildStore(FakeStore): def create_or_resolve_simulation(self, record): - raise RuntimeError("sensitive database detail") + raise RuntimeError("database rejected the child record") store = FailingChildStore() @@ -1687,10 +1733,62 @@ def create_or_resolve_simulation(self, record): assert store.parent.status is ComparisonRunLifecycleStatus.FAILED assert store.parent.error_code == "report_coordination_failed" assert store.parent.error_summary == ( - "RuntimeError in stage12_child_state_create: sensitive database detail" + "RuntimeError in stage12_child_state_create: database rejected the child record" ) +class RefusingDispatchInvoker: + def __init__(self, dispatch_error): + self.dispatch_error = dispatch_error + self.events = [] + + def spawn(self, *, simulation, **_): + role = PlannedSimulationExecutionInput.model_validate(simulation).role.value + self.events.append(f"spawn:{role}") + raise self.dispatch_error + + def restore(self, invocation_id): + raise AssertionError(f"unexpected restored invocation {invocation_id}") + + +def test_coordinator_records_a_redacted_dispatch_failure_cause() -> None: + store = FakeStore() + artifacts = FakeArtifacts() + invoker = RefusingDispatchInvoker( + ConnectionError("Modal rejected the call with token=opaque-token") + ) + runtime = RecordingRuntime() + + with pytest.raises(RuntimeError, match="Stage 12 report coordination failed"): + coordinate_report( + _report().model_dump(mode="json"), + _context().model_dump(mode="json"), + _parent().model_dump(mode="json"), + application_name=_context().modal_application, + coordinator_invocation_id="coordinator-1", + store=store, + artifacts=artifacts, + invoker=invoker, + aggregator=lambda **_: {"must": "not run"}, + runtime=runtime, + ) + + expected_summary = ( + "ConnectionError in stage12_child_dispatch (baseline): " + "Modal rejected the call with token=[REDACTED]" + ) + assert invoker.events == ["spawn:baseline"] + assert artifacts.aggregate_writes == [] + failed_child = store.children[BASELINE_ID] + assert failed_child.status is ComparisonRunLifecycleStatus.FAILED + assert failed_child.error_code == "simulation_dispatch_failed" + assert failed_child.error_summary == expected_summary + assert store.parent.error_code == "report_coordination_failed" + assert store.parent.error_summary == expected_summary + [recorded_cause] = runtime.recorded_exceptions + assert str(recorded_cause) == expected_summary + + class WorkerRecordedFailureCall: """A child whose worker persisted its cause before Modal raised.""" @@ -2079,25 +2177,10 @@ def get_simulation(self, simulation_execution_id): # The report record is still written after the child record could not be. assert store.parent.status is ComparisonRunLifecycleStatus.FAILED assert store.parent.error_summary == ( - "RuntimeError in stage12_child_wait (baseline): " - "child failed with sensitive values" + "RuntimeError in stage12_child_wait (baseline): child calculation failed" ) -class RefusingDispatchInvoker: - def __init__(self, dispatch_error): - self.dispatch_error = dispatch_error - self.events = [] - - def spawn(self, *, simulation, **_): - role = PlannedSimulationExecutionInput.model_validate(simulation).role.value - self.events.append(f"spawn:{role}") - raise self.dispatch_error - - def restore(self, invocation_id): - raise AssertionError(f"unexpected restored invocation {invocation_id}") - - def test_coordinator_reports_a_child_dispatch_failure_once() -> None: ordered_events: list[tuple[str, object]] = [] dispatch_error = ConnectionError("Modal refused the dispatch") @@ -2232,7 +2315,7 @@ def __init__(self, object_id): self.object_id = object_id def get(self, *, timeout=None): - raise TimeoutError("contains sensitive timeout details") + raise TimeoutError("child result was not ready") class TimeoutInvoker: @@ -2268,13 +2351,12 @@ def test_coordinator_records_bounded_timeout_without_an_aggregate() -> None: aggregator=lambda **_: _aggregate_payload({"must": "not run"}), ) - assert "sensitive" not in str(error.value) + assert "child result was not ready" not in str(error.value) assert invoker.events == ["spawn:baseline", "spawn:reform"] assert artifacts.aggregate_writes == [] expected_summary = ( - "TimeoutError in stage12_child_wait (baseline): " - "contains sensitive timeout details" + "TimeoutError in stage12_child_wait (baseline): child result was not ready" ) assert store.parent.status is ComparisonRunLifecycleStatus.FAILED assert store.parent.error_summary == expected_summary From 160e16b245b0500213e77a66d91f2a96257d2d54 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 15 Oct 2026 00:23:56 +0400 Subject: [PATCH 5/5] Return valid aggregates from the Stage 12 test aggregators #787 made the coordinator tests' artifact double write aggregates through the real Stage12ArtifactStore, which rejects anything that is not a valid AggregateReportArtifactPayload. Use _aggregate_payload for this PR's aggregator stubs, as #787 did for main's. The redacted comparison-failure test needs a valid aggregate for the report to succeed. An aggregator that must not run now shows up as a written aggregate if it does run, instead of as a rejection hidden by the failure the test expects. Co-Authored-By: Claude Opus 5.5 --- .../tests/test_stage12_runtime.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py index 399afea31..3816d9513 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py @@ -1584,7 +1584,7 @@ def test_comparison_failure_stores_a_redacted_cause() -> None: store=store, artifacts=artifacts, invoker=UnavailableProductionInvoker(artifacts), - aggregator=lambda **_: {"result": {"budget": 100}}, + aggregator=lambda **_: _aggregate_payload({"budget": 100}), ) assert store.parent.status is ComparisonRunLifecycleStatus.SUCCEEDED @@ -1769,7 +1769,7 @@ def test_coordinator_records_a_redacted_dispatch_failure_cause() -> None: store=store, artifacts=artifacts, invoker=invoker, - aggregator=lambda **_: {"must": "not run"}, + aggregator=lambda **_: _aggregate_payload({"must": "not run"}), runtime=runtime, ) @@ -1857,7 +1857,7 @@ def test_coordinator_keeps_the_cause_recorded_by_a_failed_worker() -> None: store=store, artifacts=FakeArtifacts(), invoker=WorkerRecordedFailureInvoker(store, worker_summary), - aggregator=lambda **_: {"must": "not run"}, + aggregator=lambda **_: _aggregate_payload({"must": "not run"}), runtime=runtime, )