From 385ca404376ebceeed74edeb87e1beb074d7283a Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 22 Sep 2026 03:57:58 +0400 Subject: [PATCH 01/10] Centralize API v1 observability --- .github/scripts/deploy_cloud_run_candidate.sh | 7 + changelog.d/3847.changed.md | 3 + policyengine_api/api.py | 6 +- policyengine_api/asgi.py | 2 + policyengine_api/gcp_logging.py | 88 ++++++------- .../libs/simulation_entrypoint.py | 3 + policyengine_api/migration_logging.py | 84 ++++++++---- policyengine_api/observability.py | 110 ++++++++++++++++ policyengine_api/services/economy_service.py | 18 --- pyproject.toml | 4 + tests/unit/libs/test_simulation_entrypoint.py | 14 ++ .../routes/test_migration_context_logging.py | 33 ++++- tests/unit/test_cloud_run_deploy_scripts.py | 15 +++ tests/unit/test_gcp_logging.py | 80 ++++++------ uv.lock | 123 ++++++++++++++---- 15 files changed, 427 insertions(+), 163 deletions(-) create mode 100644 changelog.d/3847.changed.md create mode 100644 policyengine_api/observability.py diff --git a/.github/scripts/deploy_cloud_run_candidate.sh b/.github/scripts/deploy_cloud_run_candidate.sh index cb19b9bf3..e0898e64c 100755 --- a/.github/scripts/deploy_cloud_run_candidate.sh +++ b/.github/scripts/deploy_cloud_run_candidate.sh @@ -8,6 +8,7 @@ cloud_run_set_defaults bash .github/scripts/validate_cloud_run_deploy_env.sh env_vars=( + "APP_ENVIRONMENT=${DEPLOYMENT_ENVIRONMENT}" "POLICYENGINE_DB_INSTANCE_CONNECTION_NAME=${POLICYENGINE_DB_INSTANCE_CONNECTION_NAME}" "POLICYENGINE_DB_USER=${POLICYENGINE_DB_USER:-policyengine}" "POLICYENGINE_DB_NAME=${POLICYENGINE_DB_NAME:-policyengine}" @@ -32,6 +33,12 @@ env_vars=( "RUNTIME_CACHE_MODE=deployed" "RUNTIME_CACHE_ENVIRONMENT=${CLOUD_RUN_RUNTIME_CACHE_ENVIRONMENT}" "RUNTIME_CACHE_SERVICE=api" + "OTEL_EXPORTER_OTLP_ENDPOINT=https://policyengine-api-v1-otel-collector-790230211054.us-central1.run.app" + "OTEL_EXPORTER_OTLP_PROTOCOL=grpc" + "OTEL_TRACES_EXPORTER=otlp" + "OTEL_METRICS_EXPORTER=otlp" + "OTEL_TRACES_SAMPLER_ARG=1.0" + "POLICYENGINE_OTEL_GOOGLE_AUDIENCE=https://policyengine-api-v1-otel-collector-790230211054.us-central1.run.app" "V2_SUPABASE_PROJECT_REF=${V2_SUPABASE_PROJECT_REF}" "V2_SUPABASE_ENVIRONMENT=${V2_SUPABASE_ENVIRONMENT}" "V2_RUNTIME_DATABASE_URL_SECRET_RESOURCE=${V2_RUNTIME_DATABASE_URL_SECRET_RESOURCE}" diff --git a/changelog.d/3847.changed.md b/changelog.d/3847.changed.md new file mode 100644 index 000000000..dba636323 --- /dev/null +++ b/changelog.d/3847.changed.md @@ -0,0 +1,3 @@ +Route API v1 structured logs, traces, and metrics through the explicit +policyengine-observability version 2 runtime and propagate request context to +the simulation entry service. diff --git a/policyengine_api/api.py b/policyengine_api/api.py index 8c470596b..c4bba167d 100644 --- a/policyengine_api/api.py +++ b/policyengine_api/api.py @@ -23,7 +23,9 @@ def log_timing(message): from policyengine_api.extensions import cache from policyengine_api.migration_logging import register_migration_request_logging +from policyengine_api.observability import runtime as observability_runtime from policyengine_api.runtime_cache.settings import load_runtime_cache_settings +from policyengine_observability import instrument_flask log_timing("Caching utilities import completed") @@ -60,6 +62,8 @@ def log_timing(message): app = application = flask.Flask(__name__) log_timing("Flask app created") +instrument_flask(app, observability_runtime) +log_timing("Observability initialised") runtime_cache_settings = load_runtime_cache_settings() if runtime_cache_settings.enabled: @@ -102,7 +106,7 @@ def log_timing(message): CORS(app) log_timing("CORS initialised") -register_migration_request_logging(app) +register_migration_request_logging(app, runtime=observability_runtime) log_timing("Migration request logging initialised") app.register_blueprint(error_bp) diff --git a/policyengine_api/asgi.py b/policyengine_api/asgi.py index d49ec1405..1f981c56a 100644 --- a/policyengine_api/asgi.py +++ b/policyengine_api/asgi.py @@ -7,6 +7,7 @@ from policyengine_api.api import app as flask_app from policyengine_api.asgi_factory import create_asgi_app from policyengine_api.data.orm import close_v1_engines +from policyengine_api.observability import get_runtime from policyengine_api.readiness import mark_not_ready, mark_ready from policyengine_api.runtime_cache.client import close_runtime_cache_clients from policyengine_api.warmup import run_startup_warmup @@ -15,6 +16,7 @@ def _close_runtime_resources() -> None: close_v1_engines() close_runtime_cache_clients() + get_runtime().shutdown() app = application = create_asgi_app( diff --git a/policyengine_api/gcp_logging.py b/policyengine_api/gcp_logging.py index be3c96e1b..6c3b55b7e 100644 --- a/policyengine_api/gcp_logging.py +++ b/policyengine_api/gcp_logging.py @@ -1,62 +1,50 @@ -import logging -import os -from typing import Optional - - -class _LazyGoogleLogger: - """Lazily initialize Google Cloud Logging and fall back to stderr.""" - - def __init__(self, logger_name: str): - self._logger_name = logger_name - self._google_logger = None - self._initialization_failed = False - self._fallback_logger = logging.getLogger(logger_name) - - def _get_google_logger(self): - if not os.environ.get("K_SERVICE"): - self._initialization_failed = True - return None - if self._google_logger is not None: - return self._google_logger - if self._initialization_failed: - return None - try: - from google.cloud.logging import Client +"""Compatibility facade for application-owned structured logging. - self._google_logger = Client().logger(self._logger_name) - return self._google_logger - except Exception: - self._initialization_failed = True - return None +Existing API modules call ``logger.log_struct``. The facade keeps that small +surface while sending records through the explicitly owned v2 runtime. Cloud +Run captures the resulting JSON from standard output, so request threads never +call the Cloud Logging API. +""" + +from __future__ import annotations + +from collections.abc import Mapping +from typing import Any + +from policyengine_api.observability import get_runtime + + +class _RuntimeLogger: + """Adapt the former ``log_struct`` call shape to the v2 runtime.""" def log_struct( self, - info: dict, + info: Mapping[str, Any], severity: str = "INFO", *, - labels: Optional[dict] = None, + labels: Mapping[str, Any] | None = None, ) -> None: - """Record structured diagnostics without changing caller behavior.""" - - google_logger = self._get_google_logger() - if google_logger is not None: - try: - google_logger.log_struct(info, severity=severity, labels=labels) - return - except Exception: - # Observability must never invalidate a successful request or - # cache operation. Cloud Run collects stderr as a fallback - # when the structured logging API is unavailable. - self._google_logger = None - self._initialization_failed = True - - level = getattr(logging, severity.upper(), logging.INFO) + """Record an allowlisted structured message without affecting callers.""" + try: - self._fallback_logger.log(level, "%s", info) + message = str(info.get("message") or "API event") + attributes = { + key: value + for key, value in info.items() + if key not in {"message", "migration", "response_text"} + } + migration = info.get("migration") + if isinstance(migration, Mapping): + attributes.update(migration) + if labels: + attributes.update(labels) + get_runtime().log( + message, + severity=severity, + attributes=attributes, + ) except Exception: - # Logging is diagnostic only. A broken local handler must not - # change the result of the operation that attempted to log. pass -logger = _LazyGoogleLogger("policyengine-api") +logger = _RuntimeLogger() diff --git a/policyengine_api/libs/simulation_entrypoint.py b/policyengine_api/libs/simulation_entrypoint.py index 7d4b2f19d..05b88a6ba 100644 --- a/policyengine_api/libs/simulation_entrypoint.py +++ b/policyengine_api/libs/simulation_entrypoint.py @@ -7,6 +7,7 @@ from urllib.parse import urlparse import httpx +from policyengine_observability import instrument_httpx from policyengine_api.gcp_logging import logger from policyengine_api.libs.gateway_auth import ( GatewayAuthError, @@ -16,6 +17,7 @@ gateway_auth_required, ) from policyengine_api.migration_flags import get_sim_entrypoint +from policyengine_api.observability import get_runtime from policyengine_api.request_context import ( REQUEST_ID_HEADER, current_request_id, @@ -142,6 +144,7 @@ def __init__(self, entrypoint: str | None = None): auth=auth, event_hooks={"request": [_attach_current_request_id]}, ) + instrument_httpx(self.client, get_runtime()) def _normalize_submission_payload(self, payload: dict) -> dict: if "data" in payload or "data_version" in payload: diff --git a/policyengine_api/migration_logging.py b/policyengine_api/migration_logging.py index 75b1ed447..8f9f6cbe0 100644 --- a/policyengine_api/migration_logging.py +++ b/policyengine_api/migration_logging.py @@ -5,6 +5,7 @@ import time import flask +from policyengine_observability import ObservabilityRuntime from policyengine_api.gcp_logging import logger from policyengine_api.migration_flags import ( RouteImplementation, @@ -66,13 +67,18 @@ def _is_v2_household_resource(method: str, path: str) -> bool: ) -def register_migration_request_logging(app: flask.Flask) -> None: +def register_migration_request_logging( + app: flask.Flask, + *, + runtime: ObservabilityRuntime | None = None, +) -> None: """Register request IDs and migration logging for Flask.""" @app.before_request def set_request_migration_context(): flask.g.request_started_at = time.time() - flask.g.request_id = ( + captured = runtime.capture_context() if runtime is not None else {} + flask.g.request_id = captured.get("request_id") or ( flask.request.headers.get(REQUEST_ID_HEADER) or generate_request_id() ) @@ -82,17 +88,30 @@ def log_request_migration_context(response): if request_id is not None: response.headers[REQUEST_ID_HEADER] = request_id try: - log_migration_request( - request_id=request_id, - method=flask.request.method, - path=flask.request.path, - status_code=response.status_code, - started_at=getattr(flask.g, "request_started_at", None), - country_id=flask.request.view_args.get("country_id") + country_id = ( + flask.request.view_args.get("country_id") if flask.request.view_args - else None, - route_impl=RouteImplementation.FLASK_FALLBACK, + else None ) + if runtime is not None: + runtime.set_context( + country_id=country_id, + **_migration_context( + method=flask.request.method, + path=flask.request.path, + route_impl=RouteImplementation.FLASK_FALLBACK, + ), + ) + else: + log_migration_request( + request_id=request_id, + method=flask.request.method, + path=flask.request.path, + status_code=response.status_code, + started_at=getattr(flask.g, "request_started_at", None), + country_id=country_id, + route_impl=RouteImplementation.FLASK_FALLBACK, + ) except Exception: try: app.logger.exception("Failed to log migration request context") @@ -117,6 +136,33 @@ def log_migration_request( if started_at is not None: elapsed_ms = round((time.time() - started_at) * 1000, 2) + migration_context = _migration_context( + method=method, + path=path, + route_impl=route_impl, + ) + + logger.log_struct( + { + "message": "API request served", + "request_id": request_id, + "method": method, + "path": path, + "status_code": status_code, + "latency_ms": elapsed_ms, + "country_id": country_id, + "migration": migration_context, + }, + severity="INFO" if status_code < 500 else "ERROR", + ) + + +def _migration_context( + *, + method: str, + path: str, + route_impl: RouteImplementation | None, +) -> dict[str, str | None]: route_group = infer_route_group(path) is_v2_metadata_read = _is_v2_metadata_resource_read(method, path) is_v2_policy_resource = _is_v2_policy_resource(method, path) @@ -124,7 +170,7 @@ def log_migration_request( uses_explicit_v2_source = ( is_v2_metadata_read or is_v2_policy_resource or is_v2_household_resource ) - migration_context = get_migration_log_context( + return get_migration_log_context( route_group, route_impl=route_impl, use_configured_db_sources=not uses_explicit_v2_source, @@ -142,17 +188,3 @@ def log_migration_request( else None ), ) - - logger.log_struct( - { - "message": "API request served", - "request_id": request_id, - "method": method, - "path": path, - "status_code": status_code, - "latency_ms": elapsed_ms, - "country_id": country_id, - "migration": migration_context, - }, - severity="INFO" if status_code < 500 else "ERROR", - ) diff --git a/policyengine_api/observability.py b/policyengine_api/observability.py new file mode 100644 index 000000000..f7ef90a80 --- /dev/null +++ b/policyengine_api/observability.py @@ -0,0 +1,110 @@ +"""Explicit API v1 observability runtime ownership.""" + +from __future__ import annotations + +import os +from importlib.metadata import PackageNotFoundError, version + +from policyengine_observability import ( + DeploymentIdentity, + LoggingConfig, + ObservabilityConfig, + ObservabilityRuntime, + ServiceIdentity, + configure, +) + + +GOOGLE_CLOUD_PROJECT = "policyengine-observability" +SERVICE_NAMESPACE = "policyengine.api-v1" +APPLICATION_ATTRIBUTE_KEYS = frozenset( + { + "backend", + "configured_write_source", + "baseline_policy_id", + "batch_job_id", + "cache_backend", + "cache_event", + "country_id", + "data_version", + "db_entity", + "db_read", + "db_read_source", + "db_write", + "db_write_source", + "elapsed_ms", + "error_code", + "error_type", + "execution_id", + "failure_category", + "http_status", + "job_id", + "latency_ms", + "max_parallel", + "method", + "metric_name", + "metric_value", + "migration_flag_error", + "model_version", + "path", + "policy_id", + "policyengine_version", + "process_id", + "requested_through_revision", + "resource", + "request_id", + "resolved_app_name", + "route_group", + "route_impl", + "run_id", + "simulation_year", + "sim_compute", + "sim_entrypoint", + "sim_flow", + "source_revision", + "start_year", + "status", + "status_code", + "window_size", + } +) + + +def _package_version() -> str: + try: + return version("policyengine-api") + except PackageNotFoundError: + return "4.1.0" + + +def _build_runtime() -> ObservabilityRuntime: + environment = os.getenv("APP_ENVIRONMENT", "local").strip() or "local" + return configure( + ObservabilityConfig.from_env( + service=ServiceIdentity( + name="policyengine-api", + namespace=SERVICE_NAMESPACE, + version=_package_version(), + role="api", + ), + deployment=DeploymentIdentity( + environment=environment, + platform="google_cloud_run", + region=os.getenv("CLOUD_RUN_REGION") or "us-central1", + instance_id=os.getenv("K_REVISION"), + ), + google_cloud_project_id=GOOGLE_CLOUD_PROJECT, + logging=LoggingConfig( + stdout_enabled=True, + capture_standard_library=True, + ), + application_attribute_keys=APPLICATION_ATTRIBUTE_KEYS, + ) + ) + + +runtime = _build_runtime() + + +def get_runtime() -> ObservabilityRuntime: + return runtime diff --git a/policyengine_api/services/economy_service.py b/policyengine_api/services/economy_service.py index 9f3112655..f1b564c37 100644 --- a/policyengine_api/services/economy_service.py +++ b/policyengine_api/services/economy_service.py @@ -1444,7 +1444,6 @@ def _build_simulation_telemetry( return { "run_id": str(uuid.uuid4()), "process_id": setup_options.process_id, - "traceparent": self._get_current_traceparent(), "requested_at": datetime.datetime.now(datetime.UTC).isoformat(), "simulation_kind": simulation_kind, "geography_code": geography_code, @@ -1479,23 +1478,6 @@ def _stable_config_hash(self, payload: dict[str, Any]) -> str: ).encode("utf-8") return f"sha256:{hashlib.sha256(encoded).hexdigest()}" - def _get_current_traceparent(self) -> str | None: - try: - from opentelemetry import trace - except Exception: - return None - - span = trace.get_current_span() - span_context = span.get_span_context() - if not getattr(span_context, "is_valid", False): - return None - - trace_flags = int(getattr(span_context, "trace_flags", 0)) - return ( - f"00-{span_context.trace_id:032x}-" - f"{span_context.span_id:016x}-{trace_flags:02x}" - ) - # Note: The following methods that interface with the ReformImpactsService # are written separately because the service relies upon mutating an original # 'computing' record to 'ok' or 'error' status, rather than creating a new record. diff --git a/pyproject.toml b/pyproject.toml index 141159f22..bf0d4fa32 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -39,6 +39,7 @@ dependencies = [ "microdf_python>=1.0.0", "openai", "packaging>=24,<27", + "policyengine-observability[flask,google,httpx,otlp-grpc]>=2,<3", "policyengine_canada==0.96.3", "policyengine-ng==0.5.1", "policyengine-il==0.1.0", @@ -74,6 +75,9 @@ dev = [ [tool.hatch.build.targets.wheel] packages = ["policyengine_api"] +[tool.uv.sources] +policyengine-observability = { git = "https://github.com/PolicyEngine/policyengine-observability.git", rev = "a25e63c147064918ad20b2e396cac84900727df5" } + [tool.ruff.lint.per-file-ignores] "migrations/*/versions/*.py" = ["F401"] diff --git a/tests/unit/libs/test_simulation_entrypoint.py b/tests/unit/libs/test_simulation_entrypoint.py index ddf51e16b..517d8429e 100644 --- a/tests/unit/libs/test_simulation_entrypoint.py +++ b/tests/unit/libs/test_simulation_entrypoint.py @@ -377,6 +377,20 @@ def test__given_client_initialized__then_installs_one_request_id_hook( assert list(kwargs["event_hooks"]) == ["request"] assert len(kwargs["event_hooks"]["request"]) == 1 + def test__given_client_initialized__then_instruments_explicit_httpx_client( + self, mock_httpx_client + ): + from policyengine_api.libs import simulation_entrypoint as module + + runtime = object() + with ( + patch.object(module, "get_runtime", return_value=runtime), + patch.object(module, "instrument_httpx") as instrument, + ): + client = SimulationAPIModal() + + instrument.assert_called_once_with(client.client, runtime) + def test__given_flask_request__then_hook_uses_current_request_id( self, mock_httpx_client ): diff --git a/tests/unit/routes/test_migration_context_logging.py b/tests/unit/routes/test_migration_context_logging.py index 89dff5b14..c2ce9be38 100644 --- a/tests/unit/routes/test_migration_context_logging.py +++ b/tests/unit/routes/test_migration_context_logging.py @@ -1,5 +1,5 @@ from types import SimpleNamespace -from unittest.mock import patch +from unittest.mock import Mock, patch from fastapi.testclient import TestClient from flask import Flask, Response @@ -94,6 +94,37 @@ def test_request_logging_includes_migration_context(): assert log_payload["migration"]["route_impl"] == "flask_fallback" +def test_instrumented_flask_request_enriches_single_adapter_record(): + app = Flask(__name__) + app.config["TESTING"] = True + runtime = Mock() + runtime.capture_context.return_value = {"request_id": "request-123"} + + @app.route("//metadata") + def metadata(country_id): + return Response(country_id, status=200, mimetype="text/plain") + + register_migration_request_logging(app, runtime=runtime) + + with patch("policyengine_api.migration_logging.logger") as mock_logger: + response = app.test_client().get("/us/metadata") + + assert response.status_code == 200 + assert response.headers[REQUEST_ID_HEADER] == "request-123" + runtime.set_context.assert_called_once_with( + country_id="us", + route_group="metadata", + route_impl="flask_fallback", + db_entity="metadata", + db_write="cloud_sql", + db_read="cloud_sql", + sim_flow=None, + sim_entrypoint="old_gateway_direct", + sim_compute=None, + ) + mock_logger.log_struct.assert_not_called() + + def test_flask_preserves_policyengine_request_id_in_context_log_and_response(): with patch("policyengine_api.migration_logging.logger") as mock_logger: response = ( diff --git a/tests/unit/test_cloud_run_deploy_scripts.py b/tests/unit/test_cloud_run_deploy_scripts.py index 601011366..b63ce7d79 100644 --- a/tests/unit/test_cloud_run_deploy_scripts.py +++ b/tests/unit/test_cloud_run_deploy_scripts.py @@ -988,6 +988,21 @@ def test_deploy_cloud_run_candidate_dry_run_preserves_access_and_traffic(): assert "RUNTIME_CACHE_MODE=deployed" in result.stdout assert "RUNTIME_CACHE_ENVIRONMENT=production" in result.stdout assert "RUNTIME_CACHE_SERVICE=api" in result.stdout + assert "APP_ENVIRONMENT=production" in result.stdout + assert ( + "OTEL_EXPORTER_OTLP_ENDPOINT=https://" + "policyengine-api-v1-otel-collector-790230211054.us-central1.run.app" + in result.stdout + ) + assert "OTEL_EXPORTER_OTLP_PROTOCOL=grpc" in result.stdout + assert "OTEL_TRACES_EXPORTER=otlp" in result.stdout + assert "OTEL_METRICS_EXPORTER=otlp" in result.stdout + assert "OTEL_TRACES_SAMPLER_ARG=1.0" in result.stdout + assert ( + "POLICYENGINE_OTEL_GOOGLE_AUDIENCE=https://" + "policyengine-api-v1-otel-collector-790230211054.us-central1.run.app" + in result.stdout + ) assert ( "RUNTIME_CACHE_URL=policyengine-api-prod-runtime-cache-url:latest" in result.stdout diff --git a/tests/unit/test_gcp_logging.py b/tests/unit/test_gcp_logging.py index 2edceadea..af939614f 100644 --- a/tests/unit/test_gcp_logging.py +++ b/tests/unit/test_gcp_logging.py @@ -1,49 +1,47 @@ from unittest.mock import Mock -from policyengine_api.gcp_logging import _LazyGoogleLogger - - -def test_local_logging_uses_stderr_without_initializing_google(monkeypatch): - monkeypatch.delenv("K_SERVICE", raising=False) - logger = _LazyGoogleLogger("test-local") - logger._fallback_logger = Mock() - payload = {"message": "cache miss"} - - logger.log_struct(payload, severity="WARNING", labels={"cache": "analysis"}) - - assert logger._initialization_failed is True - assert logger._google_logger is None - logger._fallback_logger.log.assert_called_once_with(30, "%s", payload) - - -def test_remote_logging_failure_falls_back_and_disables_retries(monkeypatch): - monkeypatch.setenv("K_SERVICE", "policyengine-api") - remote_logger = Mock() - remote_logger.log_struct.side_effect = ConnectionError("logging unavailable") - logger = _LazyGoogleLogger("test-deployed") - logger._google_logger = remote_logger - logger._fallback_logger = Mock() - payload = {"message": "cache write"} - - logger.log_struct(payload, severity="INFO", labels={"cache": "household"}) - logger.log_struct(payload, severity="INFO", labels={"cache": "household"}) +from policyengine_api import gcp_logging +from policyengine_api.gcp_logging import _RuntimeLogger + + +def test_runtime_logger_flattens_migration_context_and_omits_response_text( + monkeypatch, +): + runtime = Mock() + monkeypatch.setattr(gcp_logging, "get_runtime", lambda: runtime) + logger = _RuntimeLogger() + + logger.log_struct( + { + "message": "API request served", + "request_id": "request-1", + "response_text": "must not be recorded", + "migration": { + "route_impl": "flask_fallback", + "db_read_source": "cloud_sql", + }, + }, + severity="WARNING", + labels={"backend": "simulation_entry"}, + ) - remote_logger.log_struct.assert_called_once_with( - payload, - severity="INFO", - labels={"cache": "household"}, + runtime.log.assert_called_once_with( + "API request served", + severity="WARNING", + attributes={ + "request_id": "request-1", + "route_impl": "flask_fallback", + "db_read_source": "cloud_sql", + "backend": "simulation_entry", + }, ) - assert logger._initialization_failed is True - assert logger._google_logger is None - assert logger._fallback_logger.log.call_count == 2 -def test_fallback_logging_failure_does_not_escape(monkeypatch): - monkeypatch.delenv("K_SERVICE", raising=False) - logger = _LazyGoogleLogger("test-broken-fallback") - logger._fallback_logger = Mock() - logger._fallback_logger.log.side_effect = RuntimeError("logging unavailable") +def test_runtime_logging_failure_does_not_escape(monkeypatch): + runtime = Mock() + runtime.log.side_effect = RuntimeError("logging unavailable") + monkeypatch.setattr(gcp_logging, "get_runtime", lambda: runtime) - logger.log_struct({"message": "operation succeeded"}, severity="INFO") + _RuntimeLogger().log_struct({"message": "operation succeeded"}, severity="INFO") - logger._fallback_logger.log.assert_called_once() + runtime.log.assert_called_once() diff --git a/uv.lock b/uv.lock index 64d96a293..b0375d6f6 100644 --- a/uv.lock +++ b/uv.lock @@ -1444,18 +1444,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/5d/13/ad7d7ca3808a898b4612b6fe93cde56b53f3034dcde235acb1f0e1df24c6/idna-3.13-py3-none-any.whl", hash = "sha256:892ea0cde124a99ce773decba204c5552b69c3c67ffd5f232eb7696135bc8bb3", size = 68629, upload-time = "2026-04-22T16:42:40.909Z" }, ] -[[package]] -name = "importlib-metadata" -version = "8.7.1" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "zipp" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/f3/49/3b30cad09e7771a4982d9975a8cbf64f00d4a1ececb53297f1d9a7be1b10/importlib_metadata-8.7.1.tar.gz", hash = "sha256:49fef1ae6440c182052f407c8d34a68f72efc36db9ca90dc0113398f2fdde8bb", size = 57107, upload-time = "2025-12-21T10:00:19.278Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/fa/5e/f8e9a1d23b9c20a551a8a02ea3637b4642e22c2626e3a13a9a29cdea99eb/importlib_metadata-8.7.1-py3-none-any.whl", hash = "sha256:5a1f80bf1daa489495071efbb095d75a634cf28a8bc299581244063b53176151", size = 27865, upload-time = "2025-12-21T10:00:18.329Z" }, -] - [[package]] name = "iniconfig" version = "2.3.0" @@ -2448,15 +2436,83 @@ wheels = [ [[package]] name = "opentelemetry-api" -version = "1.41.1" +version = "1.44.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/ee/8b/aa9e2d8b8dfa7c946f7dec5d1f8f6ba8eca062f43509a06bdb5ce93d26c0/opentelemetry_api-1.44.0.tar.gz", hash = "sha256:67647e5e9566edcf421166fdf022b3537f818635daa852b289e34604dc6fb33a", size = 72406, upload-time = "2026-07-16T15:25:32.678Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ca/6f/a04e900f465ff3221ccc395522503e2d10e79fa21f2723c8e177aae1e0d1/opentelemetry_api-1.44.0-py3-none-any.whl", hash = "sha256:94b98c893a91b88657eaac1e3ba89618cdb85be6918196705354f34728b2cdef", size = 60018, upload-time = "2026-07-16T15:25:11.657Z" }, +] + +[[package]] +name = "opentelemetry-exporter-otlp-proto-common" +version = "1.44.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "opentelemetry-proto" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/61/09/4d717852c1cf3f854b76c7110a5d00883bc3c99288b9b0dbcbeb9e306eb6/opentelemetry_exporter_otlp_proto_common-1.44.0.tar.gz", hash = "sha256:dc87a5a5bc58f149a56d1547e4691588fa12994cdc3bc039a694ccb3375862ac", size = 20202, upload-time = "2026-07-16T15:25:37.658Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/5e/71/65fd9d54c10b860f87c045ccee1264cab7011268895d3528818a29c1172a/opentelemetry_exporter_otlp_proto_common-1.44.0-py3-none-any.whl", hash = "sha256:9a9fe61bba73d802904bc989f1d6b4a7b1ee40f06c40e98d6f85af65aaebb694", size = 17045, upload-time = "2026-07-16T15:25:18.201Z" }, +] + +[[package]] +name = "opentelemetry-exporter-otlp-proto-grpc" +version = "1.44.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "googleapis-common-protos" }, + { name = "grpcio" }, + { name = "opentelemetry-api" }, + { name = "opentelemetry-exporter-otlp-proto-common" }, + { name = "opentelemetry-proto" }, + { name = "opentelemetry-sdk" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/1f/47/80d9e9d468dc5de3af5096f5ccdb065fa4dd1470f74495cc53e59e397f47/opentelemetry_exporter_otlp_proto_grpc-1.44.0.tar.gz", hash = "sha256:40d1ae9e03fcc36de3cbac610cc99f35894938bff9cfd90fc4ec68bd85448463", size = 27225, upload-time = "2026-07-16T15:25:38.308Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/54/29/6ae42ba32b153ae0a44ae125f0caff2188bbe62d99c82d1768da30864e72/opentelemetry_exporter_otlp_proto_grpc-1.44.0-py3-none-any.whl", hash = "sha256:6a1a645ea182a2f59440c51fa8301d309f3324a8f9d65f8395584b064b67ee4e", size = 19624, upload-time = "2026-07-16T15:25:19.096Z" }, +] + +[[package]] +name = "opentelemetry-proto" +version = "1.44.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "protobuf" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/64/01/40ac4ae9a149263cc52c2cee200ddd80cb6d8db1a4610abf8eabce0fe771/opentelemetry_proto-1.44.0.tar.gz", hash = "sha256:c547a79c2f8c0c515d31509154682e5921c7cfd5ca67b70e1f9266e2c3e103f3", size = 46488, upload-time = "2026-07-16T15:25:45.34Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d1/7c/8be563d68e93bbefa5c8affb82ddcff91b3ad858ce49957ba7b16fd3e0ab/opentelemetry_proto-1.44.0-py3-none-any.whl", hash = "sha256:898b155a0e1557afd867478fb6158e8122a46329ca0bb8dc53cc55e98f017f56", size = 72483, upload-time = "2026-07-16T15:25:28.429Z" }, +] + +[[package]] +name = "opentelemetry-sdk" +version = "1.44.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "opentelemetry-api" }, + { name = "opentelemetry-semantic-conventions" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/5d/77/a6592cbc7c8d9bcc9d6757a9df45e04a7c585e3e6e7a13456da522b21109/opentelemetry_sdk-1.44.0.tar.gz", hash = "sha256:cebe7f65dc12f26ead75c6064de12fd2a9052e5060c0272d402cfa203aae123b", size = 208624, upload-time = "2026-07-16T15:25:46.078Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e7/23/ff077e61886ee020a17ce9c8b6fa11c601c8d8345b09ea24f605445df62a/opentelemetry_sdk-1.44.0-py3-none-any.whl", hash = "sha256:df081c4c6bcfdb1211e3e86140376792643128a25f8d72d1d27675936e7e96ad", size = 137221, upload-time = "2026-07-16T15:25:29.534Z" }, +] + +[[package]] +name = "opentelemetry-semantic-conventions" +version = "0.65b0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "importlib-metadata" }, + { name = "opentelemetry-api" }, { name = "typing-extensions" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/fa/fc/b7564cbef36601aef0d6c9bc01f7badb64be8e862c2e1c3c5c3b43b53e4f/opentelemetry_api-1.41.1.tar.gz", hash = "sha256:0ad1814d73b875f84494387dae86ce0b12c68556331ce6ce8fe789197c949621", size = 71416, upload-time = "2026-04-24T13:15:38.262Z" } +sdist = { url = "https://files.pythonhosted.org/packages/8f/73/0cbdebcb4cf545fdd328da14f5137e37d0770c3f26185e478b0d15d94f50/opentelemetry_semantic_conventions-0.65b0.tar.gz", hash = "sha256:f9b2b81e9d5b64f11bc952075e7e9c7fb0aab075c7fd1c46d597f1b919852d60", size = 148774, upload-time = "2026-07-16T15:25:46.902Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/29/59/3e7118ed140f76b0982ba4321bdaed1997a0473f9720de2d10788a577033/opentelemetry_api-1.41.1-py3-none-any.whl", hash = "sha256:a22df900e75c76dc08440710e51f52f1aa6b451b429298896023e60db5b3139f", size = 69007, upload-time = "2026-04-24T13:15:15.662Z" }, + { url = "https://files.pythonhosted.org/packages/a6/0e/49df70d9b81fb5cbae4bbf2a49d865b09bcbcbc4eb53f5851b1027738d78/opentelemetry_semantic_conventions-0.65b0-py3-none-any.whl", hash = "sha256:1cacde7b0ad306f84c5ef08c3dbe1bbaf20165bba6f8bff43b670e555a086bcb", size = 204645, upload-time = "2026-07-16T15:25:30.688Z" }, ] [[package]] @@ -2713,7 +2769,7 @@ models = [ [[package]] name = "policyengine-api" -version = "3.56.1" +version = "4.1.0" source = { editable = "." } dependencies = [ { name = "a2wsgi" }, @@ -2738,6 +2794,7 @@ dependencies = [ { name = "policyengine-canada" }, { name = "policyengine-il" }, { name = "policyengine-ng" }, + { name = "policyengine-observability", extra = ["flask", "google", "httpx", "otlp-grpc"] }, { name = "psycopg", extra = ["binary"] }, { name = "pydantic" }, { name = "pymysql" }, @@ -2791,6 +2848,7 @@ requires-dist = [ { name = "policyengine-canada", specifier = "==0.96.3" }, { name = "policyengine-il", specifier = "==0.1.0" }, { name = "policyengine-ng", specifier = "==0.5.1" }, + { name = "policyengine-observability", extras = ["flask", "google", "httpx", "otlp-grpc"], git = "https://github.com/PolicyEngine/policyengine-observability.git?rev=a25e63c147064918ad20b2e396cac84900727df5" }, { name = "psycopg", extras = ["binary"], specifier = ">=3.3,<4" }, { name = "pydantic" }, { name = "pymysql" }, @@ -2892,6 +2950,28 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/03/9e/1040e63f72f3857d540e52c0f99115f8db04da5e94231959d4245a119bef/policyengine_ng-0.5.1-py3-none-any.whl", hash = "sha256:21fad6aae8d80a156142ac876cf1b7679e036c1640ca6cb375661701f10b9920", size = 31074, upload-time = "2023-04-19T13:14:28.242Z" }, ] +[[package]] +name = "policyengine-observability" +version = "1.4.1" +source = { git = "https://github.com/PolicyEngine/policyengine-observability.git?rev=a25e63c147064918ad20b2e396cac84900727df5#a25e63c147064918ad20b2e396cac84900727df5" } + +[package.optional-dependencies] +flask = [ + { name = "flask" }, +] +google = [ + { name = "google-auth" }, + { name = "google-cloud-logging" }, +] +httpx = [ + { name = "httpx" }, +] +otlp-grpc = [ + { name = "opentelemetry-api" }, + { name = "opentelemetry-exporter-otlp-proto-grpc" }, + { name = "opentelemetry-sdk" }, +] + [[package]] name = "policyengine-uk" version = "2.90.2" @@ -4567,15 +4647,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/69/68/c8739671f5699c7dc470580a4f821ef37c32c4cb0b047ce223a7f115757f/yarl-1.23.0-py3-none-any.whl", hash = "sha256:a2df6afe50dea8ae15fa34c9f824a3ee958d785fd5d089063d960bae1daa0a3f", size = 48288, upload-time = "2026-03-01T22:07:51.388Z" }, ] -[[package]] -name = "zipp" -version = "3.23.1" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/30/21/093488dfc7cc8964ded15ab726fad40f25fd3d788fd741cc1c5a17d78ee8/zipp-3.23.1.tar.gz", hash = "sha256:32120e378d32cd9714ad503c1d024619063ec28aad2248dc6672ad13edfa5110", size = 25965, upload-time = "2026-04-13T23:21:46.6Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/08/8a/0861bec20485572fbddf3dfba2910e38fe249796cb73ecdeb74e07eeb8d3/zipp-3.23.1-py3-none-any.whl", hash = "sha256:0b3596c50a5c700c9cb40ba8d86d9f2cc4807e9bedb06bcdf7fac85633e444dc", size = 10378, upload-time = "2026-04-13T23:21:45.386Z" }, -] - [[package]] name = "zope-interface" version = "8.4" From cefbb9dc721073638d3293094f21b190ab46736f Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 22 Sep 2026 03:59:38 +0400 Subject: [PATCH 02/10] Make draft dependency available to pip --- pyproject.toml | 5 +---- uv.lock | 38 ++++++++++++++++++++++++++++++++++++-- 2 files changed, 37 insertions(+), 6 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index bf0d4fa32..48951c0f6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -39,7 +39,7 @@ dependencies = [ "microdf_python>=1.0.0", "openai", "packaging>=24,<27", - "policyengine-observability[flask,google,httpx,otlp-grpc]>=2,<3", + "policyengine-observability[flask,google,httpx,otlp-grpc] @ https://github.com/PolicyEngine/policyengine-observability/archive/a25e63c147064918ad20b2e396cac84900727df5.tar.gz", "policyengine_canada==0.96.3", "policyengine-ng==0.5.1", "policyengine-il==0.1.0", @@ -75,9 +75,6 @@ dev = [ [tool.hatch.build.targets.wheel] packages = ["policyengine_api"] -[tool.uv.sources] -policyengine-observability = { git = "https://github.com/PolicyEngine/policyengine-observability.git", rev = "a25e63c147064918ad20b2e396cac84900727df5" } - [tool.ruff.lint.per-file-ignores] "migrations/*/versions/*.py" = ["F401"] diff --git a/uv.lock b/uv.lock index b0375d6f6..7d59431d7 100644 --- a/uv.lock +++ b/uv.lock @@ -2848,7 +2848,7 @@ requires-dist = [ { name = "policyengine-canada", specifier = "==0.96.3" }, { name = "policyengine-il", specifier = "==0.1.0" }, { name = "policyengine-ng", specifier = "==0.5.1" }, - { name = "policyengine-observability", extras = ["flask", "google", "httpx", "otlp-grpc"], git = "https://github.com/PolicyEngine/policyengine-observability.git?rev=a25e63c147064918ad20b2e396cac84900727df5" }, + { name = "policyengine-observability", extras = ["flask", "google", "httpx", "otlp-grpc"], url = "https://github.com/PolicyEngine/policyengine-observability/archive/a25e63c147064918ad20b2e396cac84900727df5.tar.gz" }, { name = "psycopg", extras = ["binary"], specifier = ">=3.3,<4" }, { name = "pydantic" }, { name = "pymysql" }, @@ -2953,7 +2953,8 @@ wheels = [ [[package]] name = "policyengine-observability" version = "1.4.1" -source = { git = "https://github.com/PolicyEngine/policyengine-observability.git?rev=a25e63c147064918ad20b2e396cac84900727df5#a25e63c147064918ad20b2e396cac84900727df5" } +source = { url = "https://github.com/PolicyEngine/policyengine-observability/archive/a25e63c147064918ad20b2e396cac84900727df5.tar.gz" } +sdist = { hash = "sha256:6504945c1768f555087a53b7b26199c394f6482196c9c1828684d57d823c22ae" } [package.optional-dependencies] flask = [ @@ -2972,6 +2973,39 @@ otlp-grpc = [ { name = "opentelemetry-sdk" }, ] +[package.metadata] +requires-dist = [ + { name = "build", marker = "extra == 'dev'" }, + { name = "coverage", marker = "extra == 'dev'" }, + { name = "fastapi", marker = "extra == 'all'" }, + { name = "fastapi", marker = "extra == 'fastapi'" }, + { name = "flask", marker = "extra == 'all'", specifier = ">=2.2" }, + { name = "flask", marker = "extra == 'flask'", specifier = ">=2.2" }, + { name = "google-auth", marker = "extra == 'all'", specifier = ">=2.38.0" }, + { name = "google-auth", marker = "extra == 'google'", specifier = ">=2.38.0" }, + { name = "google-cloud-logging", marker = "extra == 'all'", specifier = ">=3.15.0" }, + { name = "google-cloud-logging", marker = "extra == 'google'", specifier = ">=3.15.0" }, + { name = "httpx", marker = "extra == 'all'" }, + { name = "httpx", marker = "extra == 'httpx'" }, + { name = "opentelemetry-api", marker = "extra == 'all'", specifier = ">=1.43.0" }, + { name = "opentelemetry-api", marker = "extra == 'otel'", specifier = ">=1.43.0" }, + { name = "opentelemetry-api", marker = "extra == 'otlp-grpc'", specifier = ">=1.43.0" }, + { name = "opentelemetry-api", marker = "extra == 'otlp-http'", specifier = ">=1.43.0" }, + { name = "opentelemetry-exporter-otlp-proto-grpc", marker = "extra == 'all'", specifier = ">=1.43.0" }, + { name = "opentelemetry-exporter-otlp-proto-grpc", marker = "extra == 'otlp-grpc'", specifier = ">=1.43.0" }, + { name = "opentelemetry-exporter-otlp-proto-http", marker = "extra == 'all'", specifier = ">=1.43.0" }, + { name = "opentelemetry-exporter-otlp-proto-http", marker = "extra == 'otlp-http'", specifier = ">=1.43.0" }, + { name = "opentelemetry-sdk", marker = "extra == 'all'", specifier = ">=1.43.0" }, + { name = "opentelemetry-sdk", marker = "extra == 'otel'", specifier = ">=1.43.0" }, + { name = "opentelemetry-sdk", marker = "extra == 'otlp-grpc'", specifier = ">=1.43.0" }, + { name = "opentelemetry-sdk", marker = "extra == 'otlp-http'", specifier = ">=1.43.0" }, + { name = "pyright", marker = "extra == 'dev'", specifier = ">=1.1.405" }, + { name = "pytest", marker = "extra == 'dev'" }, + { name = "ruff", marker = "extra == 'dev'", specifier = ">=0.9.0" }, + { name = "towncrier", marker = "extra == 'dev'", specifier = ">=24.8.0" }, +] +provides-extras = ["otel", "otlp-grpc", "otlp-http", "flask", "fastapi", "httpx", "google", "all", "dev"] + [[package]] name = "policyengine-uk" version = "2.90.2" From b0374292a317e7ee017af7205ef12be4fabfd85f Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 22 Sep 2026 04:01:58 +0400 Subject: [PATCH 03/10] Allow pinned draft dependency source --- pyproject.toml | 3 +++ 1 file changed, 3 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index 48951c0f6..9b6276a62 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -75,6 +75,9 @@ dev = [ [tool.hatch.build.targets.wheel] packages = ["policyengine_api"] +[tool.hatch.metadata] +allow-direct-references = true + [tool.ruff.lint.per-file-ignores] "migrations/*/versions/*.py" = ["F401"] From d4a25fa6a04801275089abba515607183c7c5f43 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 22 Sep 2026 20:12:46 +0400 Subject: [PATCH 04/10] Own API v1 observability deployment and instrumentation --- .github/scripts/deploy_cloud_run_candidate.sh | 6 +- .../scripts/validate_cloud_run_deploy_env.sh | 4 + .github/workflows/push.yml | 8 + docs/operations/api-v1-observability.md | 153 ++++++++++ gcp/observability/README.md | 271 ++++++++++++++++++ gcp/observability/alerts.template.yaml | 42 +++ gcp/observability/collector/Dockerfile | 6 + gcp/observability/collector/config.yaml | 71 +++++ .../collector/service.template.yaml | 48 ++++ gcp/observability/create_alerts.py | 246 ++++++++++++++++ gcp/observability/dashboard.template.json | 183 ++++++++++++ gcp/observability/iam.template.yaml | 36 +++ gcp/observability/log-routing.template.yaml | 37 +++ gcp/observability/render_deployment.py | 88 ++++++ gcp/observability/verify.sh | 63 ++++ gcp/observability/verify_modal_wif.py | 221 ++++++++++++++ gcp/observability/verify_otel.py | 252 ++++++++++++++++ .../workload-inventory.template.yaml | 70 +++++ policyengine_api/observability.py | 14 +- policyengine_api/services/economy_service.py | 33 +++ pyproject.toml | 2 +- tests/unit/libs/test_simulation_entrypoint.py | 8 +- tests/unit/test_cloud_run_deploy_scripts.py | 27 +- .../test_observability_deployment_assets.py | 141 +++++++++ tests/unit/test_observability_runtime.py | 23 ++ uv.lock | 8 +- 26 files changed, 2032 insertions(+), 29 deletions(-) create mode 100644 docs/operations/api-v1-observability.md create mode 100644 gcp/observability/README.md create mode 100644 gcp/observability/alerts.template.yaml create mode 100644 gcp/observability/collector/Dockerfile create mode 100644 gcp/observability/collector/config.yaml create mode 100644 gcp/observability/collector/service.template.yaml create mode 100644 gcp/observability/create_alerts.py create mode 100644 gcp/observability/dashboard.template.json create mode 100644 gcp/observability/iam.template.yaml create mode 100644 gcp/observability/log-routing.template.yaml create mode 100644 gcp/observability/render_deployment.py create mode 100644 gcp/observability/verify.sh create mode 100644 gcp/observability/verify_modal_wif.py create mode 100644 gcp/observability/verify_otel.py create mode 100644 gcp/observability/workload-inventory.template.yaml create mode 100644 tests/unit/test_observability_deployment_assets.py create mode 100644 tests/unit/test_observability_runtime.py diff --git a/.github/scripts/deploy_cloud_run_candidate.sh b/.github/scripts/deploy_cloud_run_candidate.sh index e0898e64c..2e5d80773 100755 --- a/.github/scripts/deploy_cloud_run_candidate.sh +++ b/.github/scripts/deploy_cloud_run_candidate.sh @@ -33,12 +33,14 @@ env_vars=( "RUNTIME_CACHE_MODE=deployed" "RUNTIME_CACHE_ENVIRONMENT=${CLOUD_RUN_RUNTIME_CACHE_ENVIRONMENT}" "RUNTIME_CACHE_SERVICE=api" - "OTEL_EXPORTER_OTLP_ENDPOINT=https://policyengine-api-v1-otel-collector-790230211054.us-central1.run.app" + "OBSERVABILITY_SERVICE_NAMESPACE=${OBSERVABILITY_SERVICE_NAMESPACE}" + "OBSERVABILITY_TRACE_PROJECT_ID=${OBSERVABILITY_TRACE_PROJECT_ID}" + "OTEL_EXPORTER_OTLP_ENDPOINT=${OTEL_EXPORTER_OTLP_ENDPOINT}" "OTEL_EXPORTER_OTLP_PROTOCOL=grpc" "OTEL_TRACES_EXPORTER=otlp" "OTEL_METRICS_EXPORTER=otlp" "OTEL_TRACES_SAMPLER_ARG=1.0" - "POLICYENGINE_OTEL_GOOGLE_AUDIENCE=https://policyengine-api-v1-otel-collector-790230211054.us-central1.run.app" + "POLICYENGINE_OTEL_GOOGLE_AUDIENCE=${POLICYENGINE_OTEL_GOOGLE_AUDIENCE}" "V2_SUPABASE_PROJECT_REF=${V2_SUPABASE_PROJECT_REF}" "V2_SUPABASE_ENVIRONMENT=${V2_SUPABASE_ENVIRONMENT}" "V2_RUNTIME_DATABASE_URL_SECRET_RESOURCE=${V2_RUNTIME_DATABASE_URL_SECRET_RESOURCE}" diff --git a/.github/scripts/validate_cloud_run_deploy_env.sh b/.github/scripts/validate_cloud_run_deploy_env.sh index 732b607d9..8647f1235 100755 --- a/.github/scripts/validate_cloud_run_deploy_env.sh +++ b/.github/scripts/validate_cloud_run_deploy_env.sh @@ -36,6 +36,10 @@ cloud_run_require_env \ CLOUD_RUN_VPC_NETWORK \ CLOUD_RUN_VPC_SUBNET \ CLOUD_RUN_VPC_EGRESS \ + OBSERVABILITY_SERVICE_NAMESPACE \ + OBSERVABILITY_TRACE_PROJECT_ID \ + OTEL_EXPORTER_OTLP_ENDPOINT \ + POLICYENGINE_OTEL_GOOGLE_AUDIENCE \ V2_SUPABASE_PROJECT_REF \ V2_SUPABASE_ENVIRONMENT \ V2_RUNTIME_DATABASE_URL_SECRET_RESOURCE \ diff --git a/.github/workflows/push.yml b/.github/workflows/push.yml index 1390bb8cb..6dffa5e2b 100644 --- a/.github/workflows/push.yml +++ b/.github/workflows/push.yml @@ -273,6 +273,10 @@ jobs: GATEWAY_AUTH_AUDIENCE: ${{ secrets.GATEWAY_AUTH_AUDIENCE }} GATEWAY_AUTH_CLIENT_ID: ${{ secrets.GATEWAY_AUTH_CLIENT_ID }} GATEWAY_AUTH_CLIENT_SECRET_RESOURCE: ${{ secrets.GATEWAY_AUTH_CLIENT_SECRET_RESOURCE }} + OBSERVABILITY_SERVICE_NAMESPACE: ${{ vars.OBSERVABILITY_SERVICE_NAMESPACE }} + OBSERVABILITY_TRACE_PROJECT_ID: ${{ vars.OBSERVABILITY_TRACE_PROJECT_ID }} + OTEL_EXPORTER_OTLP_ENDPOINT: ${{ vars.OBSERVABILITY_OTLP_ENDPOINT }} + POLICYENGINE_OTEL_GOOGLE_AUDIENCE: ${{ vars.OBSERVABILITY_OTLP_GOOGLE_AUDIENCE }} - name: Resolve exact Cloud Run staging candidate id: candidate run: bash .github/scripts/resolve_cloud_run_candidate_state.sh >> "$GITHUB_OUTPUT" @@ -497,6 +501,10 @@ jobs: GATEWAY_AUTH_AUDIENCE: ${{ secrets.GATEWAY_AUTH_AUDIENCE }} GATEWAY_AUTH_CLIENT_ID: ${{ secrets.GATEWAY_AUTH_CLIENT_ID }} GATEWAY_AUTH_CLIENT_SECRET_RESOURCE: ${{ secrets.GATEWAY_AUTH_CLIENT_SECRET_RESOURCE }} + OBSERVABILITY_SERVICE_NAMESPACE: ${{ vars.OBSERVABILITY_SERVICE_NAMESPACE }} + OBSERVABILITY_TRACE_PROJECT_ID: ${{ vars.OBSERVABILITY_TRACE_PROJECT_ID }} + OTEL_EXPORTER_OTLP_ENDPOINT: ${{ vars.OBSERVABILITY_OTLP_ENDPOINT }} + POLICYENGINE_OTEL_GOOGLE_AUDIENCE: ${{ vars.OBSERVABILITY_OTLP_GOOGLE_AUDIENCE }} - name: Resolve exact Cloud Run production candidate id: candidate run: bash .github/scripts/resolve_cloud_run_candidate_state.sh >> "$GITHUB_OUTPUT" diff --git a/docs/operations/api-v1-observability.md b/docs/operations/api-v1-observability.md new file mode 100644 index 000000000..15441ae64 --- /dev/null +++ b/docs/operations/api-v1-observability.md @@ -0,0 +1,153 @@ +# API v1 observability operating policy + +## Scope + +The machine-readable workload inventory is +[`gcp/observability/workload-inventory.template.yaml`](../../gcp/observability/workload-inventory.template.yaml). +Only the listed `policyengine-api`, simulation entry, simulation gateway, and +versioned simulation executor workloads participate. + +Cloud Run candidate, canary, and tagged revisions use the identity of their +containing service and are included. Modal smoke, precompute, and ephemeral +applications are excluded. `policyengine-household-api` and +`policyengine-uk-chat` remain unchanged and receive no migration or +service-specific validation in this work. + +The inventory is the configuration source for log sink filters, collector +invocation permissions, and the Modal Workload Identity Federation condition. +Telemetry attributes such as `service.namespace` do not grant access. + +## Storage decision + +Production and nonproduction application logs use the existing +`policyengine-observability` bucket in the `policyengine-observability` Google +Cloud project. The bucket is global, analytics enabled, and retains records for +30 days. `deployment.environment.name` distinguishes production and staging. + +A single bucket is appropriate for the initial rollout because the same small +operator group requires both environments, the current bucket already exists, +and a shared analytics surface simplifies request investigations. Access is +controlled at the bucket and project level rather than by environment. This +decision must be revisited before an environment requires different readers, +retention, residency, or deletion policy. + +Cloud Run writes structured JSON to standard output. Exact-service sinks in +the source projects route selected records into this bucket. Modal writes the +same records to standard output and uses the package's bounded asynchronous +Cloud Logging destination under the `policyengine-api-v1-modal` log ID. A +central exclusion prevents a directly ingested record from also being retained +in `_Default`. + +Traces and metrics use Cloud Trace and Cloud Monitoring in the same project. +They are correlated with logs by resource identity, trace ID, request ID, and +job ID; they are not stored in the log bucket. + +## Initial trace sampling + +The initial rollout uses these head-sampling settings: + +| Environment | Services | Parent-based ratio | +| --- | --- | ---: | +| Production | API v1, simulation entry, gateway, executors | 1.0 | +| Staging | API v1, simulation entry, gateway, executors | 1.0 | +| Local development | All | No remote exporter | + +Sampling every initial production trace provides a complete baseline for +volume, cost, errors, and slow operations. After at least one representative +week, operators may lower the normal-request ratio only after recording a cost +and coverage review. A lower head-sampling ratio cannot retroactively retain a +request after its outcome or duration becomes known. Retaining all errors or +slow requests with a lower normal ratio therefore requires a reviewed, +bounded collector tail-sampling policy. + +The checked-in collector configuration implements error and 30-second latency +policies and initially retains 100% of all remaining traces. Any later +reduction applies to the collector's general tail policy while SDK head +sampling remains at 100%, allowing the collector to evaluate completed traces. + +Parent sampling decisions are preserved. Society-wide simulation dispatches +must not override a sampled parent with an unsampled child. + +## Asynchronous trace relationships + +Captured dispatch context contains its UTC capture time. A worker uses the +dispatch span as its parent only when all of these conditions hold: + +- The work is a direct continuation of one dispatch. +- The worker starts no more than five minutes after capture. +- The invocation is not an independent retry. +- The invocation does not aggregate multiple dispatches. + +Otherwise the worker starts a new trace and links the dispatch span. Request +and job identifiers remain the same in either representation. Malformed or +expired remote context is ignored without rejecting the job. + +## Data policy + +### Required log fields + +- `schema_version` +- `timestamp` +- `severity` +- `message` or `event.name` +- `service.name`, `service.namespace`, `service.version`, `service.role` +- `deployment.environment.name` and `cloud.platform` +- Request, operation, trace, span, duration, outcome, and bounded error fields + when applicable + +Application attributes are stored below `attributes`. The initial allowlist is +limited to bounded operational values such as country, model version, backend, +requested version, resolved channel, authentication outcome, job type, and +simulation year. Attribute strings are truncated at 1,024 characters and one +record contains at most 32 application attributes. + +### Trace attributes + +Traces may contain the standard service resource fields, HTTP route templates, +HTTP methods, status codes, operation names, bounded deployment identifiers, +request IDs, job IDs, and explicitly approved operational attributes. Raw URLs, +query values, request bodies, response bodies, and arbitrary baggage are not +recorded. + +### Metric labels + +Metrics use only these bounded labels: + +- `service.name` +- `service.role` +- `deployment.environment.name` +- `cloud.platform` +- `http.route` +- `http.request.method` +- `http.response.status_code_class` +- `operation.name` +- `operation.kind` +- `outcome` + +Request IDs, trace IDs, job IDs, simulation IDs, raw paths, error messages, +user-provided values, unrestricted geography values, and unrestricted version +values are prohibited metric labels. + +### Prohibited telemetry data + +Logs, traces, metrics, and dispatch context must not contain: + +- Authorization headers, cookies, credentials, tokens, or secret values +- Request or response bodies +- Household situations, entity records, or person-level values +- Reform definitions or parameter payloads +- Raw IP addresses +- Prompts, model inputs, or model responses +- Exception local variables +- Function arguments or return values captured automatically + +Exception messages and stacks are truncated and passed through configured +secret-value redaction before remote delivery. + +## Operational limits + +Remote application delivery is best effort. Every application queue, exporter, +retry, network request, flush, and shutdown action has a finite bound. Queue +overflow drops the new record and increments a local counter. Internal +diagnostics are rate limited and written directly to standard error so they do +not recurse through a failing exporter. diff --git a/gcp/observability/README.md b/gcp/observability/README.md new file mode 100644 index 000000000..f023e7271 --- /dev/null +++ b/gcp/observability/README.md @@ -0,0 +1,271 @@ +# Google Cloud deployment plan + +This directory defines the centralized API v1 observability resources in the +project selected by `OBSERVABILITY_PROJECT_ID`. The workload boundary is +defined in +[`workload-inventory.template.yaml`](workload-inventory.template.yaml). +Applications absent from that inventory receive no credentials or destination +permissions. + +These files are a reviewable deployment plan. Applying them changes live IAM, +Cloud Logging routing, Cloud Run, and monitoring resources and therefore +requires an operator-approved deployment window. + +## Resources + +| File | Resource | +| --- | --- | +| `collector/config.yaml` | OTLP gRPC receiver, bounded processors, and Google Telemetry API exporter for traces and metrics | +| `collector/Dockerfile` | Google-built OTel Collector 0.160.0 plus the reviewed configuration | +| `collector/service.template.yaml` | Authenticated Cloud Run collector with fixed CPU, memory, concurrency, health checks, and scaling bounds | +| `log-routing.template.yaml` | Exact Cloud Run source sinks, restricted Modal direct-log sink, and `_Default` duplicate exclusion | +| `iam.template.yaml` | Collector and Modal service accounts, Cloud Run invokers, and a dedicated Modal API v1 identity provider | +| `dashboard.template.json` | Initial request, latency, error, dropped-item, and exporter-failure dashboard | +| `alerts.template.yaml` | Initial alert policy inputs | +| `render_deployment.py` | Validates deployment variables and writes private rendered files under `rendered/` | +| `verify.sh` | Read-only resource and routing checks after deployment | + +The collector accepts traces and metrics. Application logs do not enter the +collector. Cloud Run JSON output uses source-project sinks, while authorized +Modal processes use the package's bounded Cloud Logging writer. + +The initial collector tail policy retains 100% of traces. Separate error and +30-second latency policies are evaluated before the general policy. If the +general percentage is reduced after the volume review, those two policies keep +error and slow traces. SDK head sampling must remain at 100% for the collector +to receive spans needed for this decision. + +## Deployment order + +### 1. Configure and render deployment values + +Configure these GitHub Actions repository variables: + +- `OBSERVABILITY_PROJECT_ID` +- `OBSERVABILITY_PROJECT_NUMBER` +- `API_PROJECT_ID` +- `SIMULATION_ENTRY_PROJECT_ID` + +Configure `MODAL_WORKSPACE_ID` as a GitHub Actions repository secret. It is +private deployment metadata and must not be printed by workflows. + +Workflows that render or apply the deployment must map the values explicitly: + +```yaml +env: + OBSERVABILITY_PROJECT_ID: ${{ vars.OBSERVABILITY_PROJECT_ID }} + OBSERVABILITY_PROJECT_NUMBER: ${{ vars.OBSERVABILITY_PROJECT_NUMBER }} + API_PROJECT_ID: ${{ vars.API_PROJECT_ID }} + SIMULATION_ENTRY_PROJECT_ID: ${{ vars.SIMULATION_ENTRY_PROJECT_ID }} + MODAL_WORKSPACE_ID: ${{ secrets.MODAL_WORKSPACE_ID }} +``` + +The API and simulation repositories own their runtime destination settings. +Configure these GitHub Actions variables in both repositories: + +- `OBSERVABILITY_SERVICE_NAMESPACE` +- `OBSERVABILITY_TRACE_PROJECT_ID` +- `OBSERVABILITY_OTLP_ENDPOINT` +- `OBSERVABILITY_OTLP_GOOGLE_AUDIENCE` + +The simulation repository also configures direct Modal log delivery and its +Google identity with: + +- `OBSERVABILITY_LOGGING_PROJECT_ID` +- `OBSERVABILITY_LOG_NAME` +- `OBSERVABILITY_GOOGLE_WORKLOAD_IDENTITY_PROVIDER` +- `OBSERVABILITY_GOOGLE_SERVICE_ACCOUNT_EMAIL` + +The API Cloud Run service writes logs to standard output, so its source-project +sink selects the central log destination. It does not need direct Cloud +Logging credentials. + +For an operator-run deployment, set the same five values in the local process +without writing them to a tracked file, then render the templates: + +```bash +.venv/bin/python gcp/observability/render_deployment.py +``` + +The renderer validates every value, reports only variable names, and writes +files with owner-only permissions under the ignored `gcp/observability/rendered/` +directory. + +### 2. Enable services + +```bash +gcloud services enable \ + artifactregistry.googleapis.com \ + cloudbuild.googleapis.com \ + cloudresourcemanager.googleapis.com \ + iamcredentials.googleapis.com \ + logging.googleapis.com \ + monitoring.googleapis.com \ + run.googleapis.com \ + sts.googleapis.com \ + telemetry.googleapis.com \ + tracing.googleapis.com \ + --project="${OBSERVABILITY_PROJECT_ID}" +``` + +### 3. Create identities + +Create these service accounts in the central project: + +```text +policyengine-otel-collector@${OBSERVABILITY_PROJECT_ID}.iam.gserviceaccount.com +policyengine-api-v1-modal@${OBSERVABILITY_PROJECT_ID}.iam.gserviceaccount.com +``` + +Grant only the roles listed in the rendered `iam.yaml`. The collector receives +`roles/telemetry.writer` and `roles/serviceusage.serviceUsageConsumer`. The +Modal identity receives `roles/logging.logWriter` and collector invocation +permission. The four existing Cloud Run identities in the inventory receive +collector invocation permission on the collector service only. + +Remove project-level `roles/logging.logWriter` bindings from every identity +outside this inventory. Source-project logging service agents use conditional +`roles/logging.bucketWriter` access to the named analytics bucket and do not +receive project-level log write access. `verify.sh` fails when another +project-level log writer is present. + +Create a separate `modal-api-v1` workload identity pool and provider using the +issuer, audience, mappings, workspace, environment, and application condition +in the rendered `iam.yaml`. Do not modify the existing `modal/modal` provider +during this deployment; it belongs to applications excluded from this change. +Grant the new provider permission to impersonate only the API v1 Modal service +account. + +Before enabling the provider, decode one production and one staging Modal +identity token locally and confirm that `workspace_id`, `environment_name`, +and `app_name` exactly match the reviewed condition. + +### 4. Build and deploy the collector + +```bash +gcloud artifacts repositories create observability \ + --repository-format=docker \ + --location=us-central1 \ + --immutable-tags \ + --project="${OBSERVABILITY_PROJECT_ID}" + +gcloud builds submit gcp/observability/collector \ + --tag="us-central1-docker.pkg.dev/${OBSERVABILITY_PROJECT_ID}/observability/otel-collector:0.160.0-api-v1-1" \ + --project="${OBSERVABILITY_PROJECT_ID}" + +gcloud run services replace gcp/observability/rendered/collector/service.yaml \ + --region=us-central1 \ + --project="${OBSERVABILITY_PROJECT_ID}" +``` + +Apply `roles/run.invoker` bindings for the five identities listed in the +rendered `iam.yaml`. Do not grant unauthenticated invocation. Record the HTTPS +service URL as both `OTEL_EXPORTER_OTLP_ENDPOINT` and +`POLICYENGINE_OTEL_GOOGLE_AUDIENCE` in participating service configuration. + +### 5. Configure log routing + +Create one aggregated sink in each source project using the exact service and +schema filters in the rendered `log-routing.yaml`. Grant each generated sink +writer identity `roles/logging.bucketWriter` on the central log bucket. + +Update `policyengine-observability-app-logs` to the listed direct-log filter. +Add the listed exclusion to `_Default`; this prevents a direct Modal log from +being stored in both `_Default` and the analytics bucket. Preserve the existing +Cloud Audit Log exclusions. + +After routing one synthetic record per participating service, confirm each +`insertId` exists exactly once in the central project. + +### 6. Create dashboard and alerts + +```bash +gcloud monitoring dashboards create \ + --config-from-file=gcp/observability/rendered/dashboard.json \ + --project="${OBSERVABILITY_PROJECT_ID}" +``` + +Create the API-ready alert policies with: + +```bash +.venv/bin/python gcp/observability/create_alerts.py +``` + +The script is idempotent by policy display name. It leaves notification-channel +configuration empty when the project has no channel; add operator-owned channel +identifiers after creating the relevant email, Slack, or paging destination. + +### 7. Verify before consumer deployment + +```bash +bash gcp/observability/verify.sh +``` + +Then use an approved workload identity to send one trace and metric. Attempt +the same request with a synthetic Modal token whose application name is not in +the inventory; token exchange or collector invocation must return permission +denial. Do not invoke an excluded application to perform this check. + +For an operator-run Cloud Run identity check, temporarily grant the operator +`roles/iam.serviceAccountTokenCreator` on one inventoried runtime identity, run: + +```bash +.venv/bin/python gcp/observability/verify_otel.py \ + --endpoint="${POLICYENGINE_OTEL_GOOGLE_AUDIENCE}" \ + --service-account="sim-entry-beta-runtime@${SIMULATION_ENTRY_PROJECT_ID}.iam.gserviceaccount.com" +``` + +Remove the temporary operator binding immediately after the check. The script +requires an authenticated `gcloud` session, sends one trace and metric, verifies +both Google Cloud stores, and confirms that the collector rejects OTLP logs. + +Use [`verify_modal_wif.py`](verify_modal_wif.py) with the Modal CLI to run an +allowed app name and a synthetic denied app name. The remote function exchanges +its automatically injected Modal OIDC token, verifies service-account access, +invokes the collector, and writes one routing record without exposing any +token. Run an allowed app name from a temporary non-allowlisted environment to +verify the environment restriction, then delete that environment. + +## Rollback + +1. Remove the OTel endpoint from participating service configuration. Local + structured logging continues and no remote OTel exporter is created. +2. Remove Modal remote logging configuration. Modal JSON output continues. +3. Revert each participating service to its previous package version and + deployment revision. +4. Remove the new source sinks and restore the prior central direct-log sink + filter and `_Default` exclusion state. +5. Remove invoker bindings, disable the `modal-api-v1` provider, and disable or + delete the collector service. +6. Keep the central bucket during the retention period unless the stored data + itself caused the incident. + +Rollback does not modify the existing `modal/modal` provider or any excluded +application deployment. + +## Deployment record + +The infrastructure portion of this runbook was applied and verified on +2026-09-22: + +- the global central log bucket retains records for 30 + days and has log analytics enabled; +- exact source-project sinks route the two API services and the two simulation + entry services to that bucket; +- the authenticated collector runs in `us-central1` as + `policyengine-api-v1-otel-collector`; +- the dedicated `modal-api-v1` identity provider is active with the workspace, + environment, and application conditions in the rendered `iam.yaml`; +- the only project-level `roles/logging.logWriter` identity is the API v1 + Modal service account; +- the dashboard and six alert policies are present and enabled; and +- the project currently has no alert notification channel, so the policies + record incidents without sending email, Slack, or paging notifications. + +The package and consumer service rollout remains pending until the three draft +pull requests are reviewed, the package is published as version 2.0, and the +temporary Git source pins in both consumer repositories are replaced with the +published version. Run the synthetic cross-service request, volume and cost +measurement, and destination comparison after those deployments. Record the +deployed revisions and the observation interval here before declaring the +consumer rollout complete. diff --git a/gcp/observability/alerts.template.yaml b/gcp/observability/alerts.template.yaml new file mode 100644 index 000000000..19fdbac27 --- /dev/null +++ b/gcp/observability/alerts.template.yaml @@ -0,0 +1,42 @@ +policies: + - display_name: API v1 exporter failures + condition_type: prometheus_query_language + query: 'sum(increase({"policyengine.telemetry.exporter.failure"}[5m])) > 0' + duration: 0s + evaluation_interval: 60s + disable_metric_validation: true + - display_name: API v1 dropped telemetry + condition_type: prometheus_query_language + query: 'sum(increase({"policyengine.telemetry.dropped"}[5m])) > 0' + duration: 0s + evaluation_interval: 60s + disable_metric_validation: true + - display_name: API v1 elevated error rate + condition_type: prometheus_query_language + query: 'sum(increase({"policyengine.error.count"}[10m])) > 1' + duration: 0s + evaluation_interval: 60s + disable_metric_validation: true + - display_name: API v1 high request latency + condition_type: prometheus_query_language + query: >- + histogram_quantile(0.99, sum by (le) + (rate({"policyengine.request.duration_bucket"}[5m]))) > 30 + duration: 300s + evaluation_interval: 60s + disable_metric_validation: true + - display_name: API v1 collector unavailable + condition_type: metric_threshold + metric_type: run.googleapis.com/service_health_count + resource_label_service_name: policyengine-api-v1-otel-collector + metric_label_service_health: UNHEALTHY + comparison: COMPARISON_GT + threshold: 0 + duration: 300s + - display_name: API v1 monthly log ingestion above 10 GiB + condition_type: metric_threshold + metric_type: logging.googleapis.com/billing/log_bucket_monthly_bytes_ingested + metric_label_log_bucket_id: ${OBSERVABILITY_PROJECT_ID} + comparison: COMPARISON_GT + threshold: 10737418240 + duration: 0s diff --git a/gcp/observability/collector/Dockerfile b/gcp/observability/collector/Dockerfile new file mode 100644 index 000000000..ae3f63207 --- /dev/null +++ b/gcp/observability/collector/Dockerfile @@ -0,0 +1,6 @@ +FROM us-docker.pkg.dev/cloud-ops-agents-artifacts/google-cloud-opentelemetry-collector/otelcol-google:0.160.0 + +COPY config.yaml /etc/otelcol-google/config.yaml + +ENTRYPOINT ["/otelcol-google"] +CMD ["--config=/etc/otelcol-google/config.yaml"] diff --git a/gcp/observability/collector/config.yaml b/gcp/observability/collector/config.yaml new file mode 100644 index 000000000..2fa3ced55 --- /dev/null +++ b/gcp/observability/collector/config.yaml @@ -0,0 +1,71 @@ +receivers: + otlp: + protocols: + grpc: + endpoint: 0.0.0.0:8080 + +processors: + memory_limiter: + check_interval: 1s + limit_mib: 384 + spike_limit_mib: 64 + resource/destination: + attributes: + - key: gcp.project_id + value: ${env:OBSERVABILITY_PROJECT_ID} + action: upsert + tail_sampling: + decision_wait: 30s + num_traces: 50000 + expected_new_traces_per_sec: 500 + policies: + - name: retain-errors + type: status_code + status_code: + status_codes: [ERROR] + - name: retain-slow-operations + type: latency + latency: + threshold_ms: 30000 + - name: initial-full-sample + type: probabilistic + probabilistic: + sampling_percentage: 100 + batch: + send_batch_size: 200 + send_batch_max_size: 1000 + timeout: 5s + +exporters: + otlp_grpc: + endpoint: telemetry.googleapis.com:443 + auth: + authenticator: googleclientauth + sending_queue: + enabled: true + queue_size: 2000 + retry_on_failure: + enabled: true + initial_interval: 1s + max_interval: 5s + max_elapsed_time: 30s + +extensions: + googleclientauth: + health_check: + endpoint: 0.0.0.0:13133 + +service: + extensions: [googleclientauth, health_check] + pipelines: + traces: + receivers: [otlp] + processors: [memory_limiter, resource/destination, tail_sampling, batch] + exporters: [otlp_grpc] + metrics: + receivers: [otlp] + processors: [memory_limiter, resource/destination, batch] + exporters: [otlp_grpc] + telemetry: + logs: + level: info diff --git a/gcp/observability/collector/service.template.yaml b/gcp/observability/collector/service.template.yaml new file mode 100644 index 000000000..8cde2e4aa --- /dev/null +++ b/gcp/observability/collector/service.template.yaml @@ -0,0 +1,48 @@ +apiVersion: serving.knative.dev/v1 +kind: Service +metadata: + name: policyengine-api-v1-otel-collector + namespace: "${OBSERVABILITY_PROJECT_NUMBER}" + labels: + cloud.googleapis.com/location: us-central1 + annotations: + run.googleapis.com/ingress: all +spec: + template: + metadata: + annotations: + autoscaling.knative.dev/minScale: "1" + autoscaling.knative.dev/maxScale: "10" + run.googleapis.com/cpu-throttling: "false" + run.googleapis.com/startup-cpu-boost: "true" + spec: + containerConcurrency: 100 + timeoutSeconds: 30 + serviceAccountName: policyengine-otel-collector@${OBSERVABILITY_PROJECT_ID}.iam.gserviceaccount.com + containers: + - name: collector + image: us-central1-docker.pkg.dev/${OBSERVABILITY_PROJECT_ID}/observability/otel-collector:0.160.0-api-v1-1 + env: + - name: OBSERVABILITY_PROJECT_ID + value: ${OBSERVABILITY_PROJECT_ID} + ports: + - name: h2c + containerPort: 8080 + resources: + limits: + cpu: "1" + memory: 512Mi + startupProbe: + httpGet: + path: / + port: 13133 + periodSeconds: 2 + timeoutSeconds: 1 + failureThreshold: 30 + livenessProbe: + httpGet: + path: / + port: 13133 + periodSeconds: 30 + timeoutSeconds: 2 + failureThreshold: 3 diff --git a/gcp/observability/create_alerts.py b/gcp/observability/create_alerts.py new file mode 100644 index 000000000..1f9637867 --- /dev/null +++ b/gcp/observability/create_alerts.py @@ -0,0 +1,246 @@ +"""Create the initial API v1 Cloud Monitoring alert policies idempotently.""" + +from __future__ import annotations + +import json +import os +import subprocess +import urllib.parse +import urllib.request +from typing import Any + +PROJECT = os.environ.get("OBSERVABILITY_PROJECT_ID", "") +API_ROOT = f"https://monitoring.googleapis.com/v3/projects/{PROJECT}" + + +def _threshold_condition( + *, + display_name: str, + filter_value: str, + comparison: str, + threshold: float, + duration: str, + alignment_period: str, + aligner: str, + reducer: str, +) -> dict[str, Any]: + return { + "displayName": display_name, + "conditionThreshold": { + "filter": filter_value, + "comparison": comparison, + "thresholdValue": threshold, + "duration": duration, + "aggregations": [ + { + "alignmentPeriod": alignment_period, + "perSeriesAligner": aligner, + "crossSeriesReducer": reducer, + } + ], + "trigger": {"count": 1}, + }, + } + + +def _promql_condition( + *, + display_name: str, + query: str, + duration: str, + alert_rule: str, +) -> dict[str, Any]: + return { + "displayName": display_name, + "conditionPrometheusQueryLanguage": { + "query": query, + "duration": duration, + "evaluationInterval": "60s", + "alertRule": alert_rule, + "ruleGroup": "policyengine_api_v1", + "disableMetricValidation": True, + }, + } + + +def _policy( + display_name: str, + condition: dict[str, Any], + documentation: str, +) -> dict[str, Any]: + return { + "displayName": display_name, + "combiner": "OR", + "enabled": True, + "notificationChannels": [], + "documentation": { + "content": documentation, + "mimeType": "text/markdown", + }, + "alertStrategy": {"autoClose": "1800s"}, + "conditions": [condition], + } + + +POLICIES = [ + _policy( + "API v1 exporter failures", + _promql_condition( + display_name="Exporter failure counter increased", + query=( + 'sum(increase({"policyengine.telemetry.exporter.failure"}[5m])) > 0' + ), + duration="0s", + alert_rule="ExporterFailure", + ), + "The API v1 runtime reported at least one telemetry exporter failure.", + ), + _policy( + "API v1 dropped telemetry", + _promql_condition( + display_name="Dropped telemetry counter increased", + query=('sum(increase({"policyengine.telemetry.dropped"}[5m])) > 0'), + duration="0s", + alert_rule="DroppedTelemetry", + ), + "An API v1 bounded telemetry queue dropped at least one item.", + ), + _policy( + "API v1 elevated error rate", + _promql_condition( + display_name="More than one application error in ten minutes", + query='sum(increase({"policyengine.error.count"}[10m])) > 1', + duration="0s", + alert_rule="ElevatedErrorRate", + ), + "API v1 application error counters increased more than once in ten minutes.", + ), + _policy( + "API v1 high request latency", + _promql_condition( + display_name="P99 request duration above 30 seconds", + query=( + "histogram_quantile(0.99, sum by (le) " + '(rate({"policyengine.request.duration_bucket"}[5m]))) > 30' + ), + duration="300s", + alert_rule="HighRequestLatency", + ), + "API v1 P99 request duration exceeded 30 seconds for five minutes.", + ), + _policy( + "API v1 collector unavailable", + _threshold_condition( + display_name="Collector Cloud Run service is unhealthy", + filter_value=( + 'resource.type="cloud_run_revision" AND ' + 'metric.type="run.googleapis.com/service_health_count" AND ' + 'resource.label."service_name"=' + '"policyengine-api-v1-otel-collector" AND ' + 'metric.label."service_health"="UNHEALTHY"' + ), + comparison="COMPARISON_GT", + threshold=0, + duration="300s", + alignment_period="60s", + aligner="ALIGN_MEAN", + reducer="REDUCE_MAX", + ), + "The authenticated API v1 collector reported an unhealthy revision.", + ), + _policy( + "API v1 monthly log ingestion above 10 GiB", + _threshold_condition( + display_name="Central analytics bucket exceeds 10 GiB this month", + filter_value=( + 'resource.type="global" AND ' + 'metric.type="logging.googleapis.com/billing/' + 'log_bucket_monthly_bytes_ingested" AND ' + f'metric.label."log_bucket_id"="{PROJECT}"' + ), + comparison="COMPARISON_GT", + threshold=10_737_418_240, + duration="0s", + alignment_period="1800s", + aligner="ALIGN_MAX", + reducer="REDUCE_SUM", + ), + "The central API v1 analytics bucket exceeded 10 GiB of month-to-date ingestion.", + ), +] + + +def _token() -> str: + return subprocess.check_output( + ["gcloud", "auth", "print-access-token"], + text=True, + ).strip() + + +def _request( + url: str, + *, + token: str, + payload: dict[str, Any] | None = None, +) -> dict[str, Any]: + request = urllib.request.Request( + url, + data=json.dumps(payload).encode() if payload is not None else None, + headers={ + "Authorization": f"Bearer {token}", + "Content-Type": "application/json", + }, + method="POST" if payload is not None else "GET", + ) + with urllib.request.urlopen(request, timeout=30) as response: + return json.loads(response.read()) + + +def main() -> None: + if not PROJECT: + raise SystemExit("Missing deployment variable: OBSERVABILITY_PROJECT_ID") + token = _token() + existing: dict[str, str] = {} + page_token = "" + while True: + query = urllib.parse.urlencode({"pageToken": page_token}) if page_token else "" + suffix = f"?{query}" if query else "" + response = _request( + f"{API_ROOT}/alertPolicies{suffix}", + token=token, + ) + for policy in response.get("alertPolicies", []): + existing[str(policy["displayName"])] = str(policy["name"]) + page_token = str(response.get("nextPageToken", "")) + if not page_token: + break + + results = [] + for policy in POLICIES: + display_name = str(policy["displayName"]) + if display_name in existing: + results.append( + { + "displayName": display_name, + "name": existing[display_name], + "status": "existing", + } + ) + continue + created = _request( + f"{API_ROOT}/alertPolicies", + token=token, + payload=policy, + ) + results.append( + { + "displayName": display_name, + "name": created["name"], + "status": "created", + } + ) + print(json.dumps(results, indent=2, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/gcp/observability/dashboard.template.json b/gcp/observability/dashboard.template.json new file mode 100644 index 000000000..2fa07b296 --- /dev/null +++ b/gcp/observability/dashboard.template.json @@ -0,0 +1,183 @@ +{ + "displayName": "PolicyEngine API v1 observability", + "mosaicLayout": { + "columns": 12, + "tiles": [ + { + "xPos": 0, + "yPos": 0, + "width": 6, + "height": 4, + "widget": { + "title": "Request rate by service", + "xyChart": { + "dataSets": [{ + "timeSeriesQuery": { + "timeSeriesFilter": { + "filter": "metric.type=\"prometheus.googleapis.com/policyengine.request.count/counter\"", + "aggregation": { + "alignmentPeriod": "60s", + "perSeriesAligner": "ALIGN_RATE", + "crossSeriesReducer": "REDUCE_SUM", + "groupByFields": ["resource.label.\"job\""] + } + } + }, + "plotType": "LINE" + }] + } + } + }, + { + "xPos": 6, + "yPos": 0, + "width": 6, + "height": 4, + "widget": { + "title": "Request duration", + "xyChart": { + "dataSets": [{ + "timeSeriesQuery": { + "timeSeriesFilter": { + "filter": "metric.type=\"prometheus.googleapis.com/policyengine.request.duration/histogram\"", + "aggregation": { + "alignmentPeriod": "60s", + "perSeriesAligner": "ALIGN_PERCENTILE_99" + } + } + }, + "plotType": "LINE" + }] + } + } + }, + { + "xPos": 0, + "yPos": 4, + "width": 6, + "height": 4, + "widget": { + "title": "Application errors", + "xyChart": { + "dataSets": [{ + "timeSeriesQuery": { + "timeSeriesFilter": { + "filter": "metric.type=\"prometheus.googleapis.com/policyengine.error.count/counter\"", + "aggregation": { + "alignmentPeriod": "60s", + "perSeriesAligner": "ALIGN_RATE", + "crossSeriesReducer": "REDUCE_SUM" + } + } + }, + "plotType": "STACKED_AREA" + }] + } + } + }, + { + "xPos": 6, + "yPos": 4, + "width": 6, + "height": 4, + "widget": { + "title": "Dropped telemetry and exporter failures", + "xyChart": { + "dataSets": [ + { + "timeSeriesQuery": { + "timeSeriesFilter": { + "filter": "metric.type=\"prometheus.googleapis.com/policyengine.telemetry.dropped/counter\"", + "aggregation": { + "alignmentPeriod": "60s", + "perSeriesAligner": "ALIGN_RATE", + "crossSeriesReducer": "REDUCE_SUM" + } + } + }, + "plotType": "LINE" + }, + { + "timeSeriesQuery": { + "timeSeriesFilter": { + "filter": "metric.type=\"prometheus.googleapis.com/policyengine.telemetry.exporter.failure/counter\"", + "aggregation": { + "alignmentPeriod": "60s", + "perSeriesAligner": "ALIGN_RATE", + "crossSeriesReducer": "REDUCE_SUM" + } + } + }, + "plotType": "LINE" + } + ] + } + } + }, + { + "xPos": 0, + "yPos": 8, + "width": 6, + "height": 4, + "widget": { + "title": "Collector health", + "xyChart": { + "dataSets": [{ + "timeSeriesQuery": { + "timeSeriesFilter": { + "filter": "resource.type=\"cloud_run_revision\" AND metric.type=\"run.googleapis.com/service_health_count\" AND resource.label.\"service_name\"=\"policyengine-api-v1-otel-collector\"", + "aggregation": { + "alignmentPeriod": "60s", + "perSeriesAligner": "ALIGN_MEAN", + "crossSeriesReducer": "REDUCE_SUM", + "groupByFields": ["metric.label.\"service_health\""] + } + } + }, + "plotType": "LINE" + }] + } + } + }, + { + "xPos": 6, + "yPos": 8, + "width": 6, + "height": 4, + "widget": { + "title": "Telemetry ingestion volume", + "xyChart": { + "dataSets": [ + { + "timeSeriesQuery": { + "timeSeriesFilter": { + "filter": "resource.type=\"global\" AND metric.type=\"logging.googleapis.com/billing/log_bucket_bytes_ingested\" AND metric.label.\"log_bucket_id\"=\"${OBSERVABILITY_PROJECT_ID}\"", + "aggregation": { + "alignmentPeriod": "3600s", + "perSeriesAligner": "ALIGN_SUM", + "crossSeriesReducer": "REDUCE_SUM" + } + } + }, + "plotType": "LINE" + }, + { + "timeSeriesQuery": { + "timeSeriesFilter": { + "filter": "resource.type=\"global\" AND metric.type=\"monitoring.googleapis.com/billing/samples_ingested\" AND metric.label.\"metric_domain\"=\"prometheus.googleapis.com\"", + "aggregation": { + "alignmentPeriod": "3600s", + "perSeriesAligner": "ALIGN_SUM", + "crossSeriesReducer": "REDUCE_SUM" + } + } + }, + "plotType": "LINE" + } + ] + } + } + } + ] + } +} diff --git a/gcp/observability/iam.template.yaml b/gcp/observability/iam.template.yaml new file mode 100644 index 000000000..062e90bbd --- /dev/null +++ b/gcp/observability/iam.template.yaml @@ -0,0 +1,36 @@ +service_accounts: + collector: + email: policyengine-otel-collector@${OBSERVABILITY_PROJECT_ID}.iam.gserviceaccount.com + project_roles: + - roles/telemetry.writer + - roles/serviceusage.serviceUsageConsumer + modal: + email: policyengine-api-v1-modal@${OBSERVABILITY_PROJECT_ID}.iam.gserviceaccount.com + project_roles: + - roles/logging.logWriter + collector_roles: + - roles/run.invoker + +collector_invokers: + - policyengine-api-cr-runtime@${API_PROJECT_ID}.iam.gserviceaccount.com + - policyengine-api-cr-staging@${API_PROJECT_ID}.iam.gserviceaccount.com + - sim-entry-prod-runtime@${SIMULATION_ENTRY_PROJECT_ID}.iam.gserviceaccount.com + - sim-entry-beta-runtime@${SIMULATION_ENTRY_PROJECT_ID}.iam.gserviceaccount.com + - policyengine-api-v1-modal@${OBSERVABILITY_PROJECT_ID}.iam.gserviceaccount.com + +modal_workload_identity: + pool_id: modal-api-v1 + provider_id: modal-api-v1 + issuer: https://oidc.modal.com + allowed_audience: oidc.modal.com + workspace_id: ${MODAL_WORKSPACE_ID} + attribute_mapping: + google.subject: assertion.app_name + attribute.app_name: assertion.app_name + attribute.environment_name: assertion.environment_name + attribute.workspace_id: assertion.workspace_id + attribute_condition: >- + assertion.workspace_id == "${MODAL_WORKSPACE_ID}" && + (assertion.environment_name == "main" || assertion.environment_name == "staging") && + (assertion.app_name == "policyengine-simulation-gateway" || + assertion.app_name.matches("^policyengine-simulation-py[0-9]+-[0-9]+-[0-9]+$")) diff --git a/gcp/observability/log-routing.template.yaml b/gcp/observability/log-routing.template.yaml new file mode 100644 index 000000000..8d84ed8a6 --- /dev/null +++ b/gcp/observability/log-routing.template.yaml @@ -0,0 +1,37 @@ +destination: + project: ${OBSERVABILITY_PROJECT_ID} + location: global + bucket: ${OBSERVABILITY_PROJECT_ID} + direct_log_name: policyengine-api-v1-modal + +source_sinks: + - project: ${API_PROJECT_ID} + name: api-v1-central-observability + filter: >- + resource.type="cloud_run_revision" AND + (resource.labels.service_name="policyengine-api" OR + resource.labels.service_name="policyengine-api-staging") AND + jsonPayload.schema_version="policyengine.observability.v2" + - project: ${SIMULATION_ENTRY_PROJECT_ID} + name: api-v1-central-observability + filter: >- + resource.type="cloud_run_revision" AND + (resource.labels.service_name="policyengine-simulation-entry" OR + resource.labels.service_name="policyengine-simulation-entry-staging") AND + jsonPayload.schema_version="policyengine.observability.v2" + +central_direct_sink: + name: policyengine-observability-app-logs + filter: >- + LOG_ID("policyengine-api-v1-modal") AND + jsonPayload.schema_version="policyengine.observability.v2" AND + jsonPayload."service.namespace"="policyengine.api-v1" + +default_bucket_exclusion: + name: exclude-api-v1-direct-observability + filter: >- + LOG_ID("policyengine-api-v1-modal") AND + jsonPayload.schema_version="policyengine.observability.v2" AND + jsonPayload."service.namespace"="policyengine.api-v1" AND + (jsonPayload."service.name"="policyengine-simulation-gateway" OR + jsonPayload."service.name"=~"^policyengine-simulation-py[0-9]+-[0-9]+-[0-9]+$") diff --git a/gcp/observability/render_deployment.py b/gcp/observability/render_deployment.py new file mode 100644 index 000000000..a95474852 --- /dev/null +++ b/gcp/observability/render_deployment.py @@ -0,0 +1,88 @@ +"""Render Google Cloud deployment templates from validated environment values.""" + +from __future__ import annotations + +import argparse +import os +import re +from pathlib import Path + +ROOT = Path(__file__).resolve().parent +DEFAULT_OUTPUT = ROOT / "rendered" +PLACEHOLDER = re.compile(r"\$\{([A-Z][A-Z0-9_]*)\}") +TEMPLATES = ( + ("iam.template.yaml", "iam.yaml"), + ("workload-inventory.template.yaml", "workload-inventory.yaml"), + ("log-routing.template.yaml", "log-routing.yaml"), + ("alerts.template.yaml", "alerts.yaml"), + ("dashboard.template.json", "dashboard.json"), + ("collector/service.template.yaml", "collector/service.yaml"), +) +VALIDATORS = { + "OBSERVABILITY_PROJECT_ID": re.compile(r"[a-z][a-z0-9-]{4,28}[a-z0-9]"), + "OBSERVABILITY_PROJECT_NUMBER": re.compile(r"[1-9][0-9]{5,29}"), + "API_PROJECT_ID": re.compile(r"[a-z][a-z0-9-]{4,28}[a-z0-9]"), + "SIMULATION_ENTRY_PROJECT_ID": re.compile(r"[a-z][a-z0-9-]{4,28}[a-z0-9]"), + "MODAL_WORKSPACE_ID": re.compile(r"ac-[A-Za-z0-9_-]+"), +} + + +def _deployment_values() -> dict[str, str]: + missing = sorted(name for name in VALIDATORS if not os.getenv(name)) + if missing: + raise SystemExit("Missing deployment variables: " + ", ".join(missing)) + + values = {name: os.environ[name] for name in VALIDATORS} + invalid = sorted( + name + for name, pattern in VALIDATORS.items() + if pattern.fullmatch(values[name]) is None + ) + if invalid: + raise SystemExit("Invalid deployment variables: " + ", ".join(invalid)) + return values + + +def _render(content: str, values: dict[str, str]) -> str: + referenced = set(PLACEHOLDER.findall(content)) + unknown = sorted(referenced - values.keys()) + if unknown: + raise SystemExit("Unknown deployment variables: " + ", ".join(unknown)) + rendered = PLACEHOLDER.sub(lambda match: values[match.group(1)], content) + unresolved = sorted(set(PLACEHOLDER.findall(rendered))) + if unresolved: + raise SystemExit("Unresolved deployment variables: " + ", ".join(unresolved)) + return rendered + + +def render_deployment(output_directory: Path) -> None: + values = _deployment_values() + rendered_files = [] + for source_name, output_name in TEMPLATES: + source = ROOT / source_name + rendered_files.append((output_name, _render(source.read_text(), values))) + + for output_name, content in rendered_files: + destination = output_directory / output_name + destination.parent.mkdir(parents=True, exist_ok=True) + destination.write_text(content) + destination.chmod(0o600) + + variable_names = ", ".join(sorted(values)) + print(f"Rendered {len(TEMPLATES)} deployment files using: {variable_names}") + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument( + "--output-dir", + type=Path, + default=DEFAULT_OUTPUT, + help="Destination for rendered files (default: gcp/observability/rendered)", + ) + arguments = parser.parse_args() + render_deployment(arguments.output_dir) + + +if __name__ == "__main__": + main() diff --git a/gcp/observability/verify.sh b/gcp/observability/verify.sh new file mode 100644 index 000000000..a14188a21 --- /dev/null +++ b/gcp/observability/verify.sh @@ -0,0 +1,63 @@ +#!/usr/bin/env bash +set -euo pipefail + +: "${OBSERVABILITY_PROJECT_ID:?Missing OBSERVABILITY_PROJECT_ID}" +: "${API_PROJECT_ID:?Missing API_PROJECT_ID}" +: "${SIMULATION_ENTRY_PROJECT_ID:?Missing SIMULATION_ENTRY_PROJECT_ID}" + +project="${OBSERVABILITY_PROJECT_ID}" +region="us-central1" +collector="policyengine-api-v1-otel-collector" + +gcloud logging buckets describe "${OBSERVABILITY_PROJECT_ID}" \ + --location=global \ + --project="${project}" \ + --format='value(name,retentionDays,analyticsEnabled)' + +gcloud run services describe "${collector}" \ + --region="${region}" \ + --project="${project}" \ + --format='value(status.url,spec.template.spec.serviceAccountName)' + +gcloud projects get-iam-policy "${project}" \ + --flatten='bindings[].members' \ + --filter='bindings.role:roles/telemetry.writer OR bindings.role:roles/serviceusage.serviceUsageConsumer OR bindings.role:roles/logging.logWriter' \ + --format='table(bindings.role,bindings.members)' + +log_writers="$( + gcloud projects get-iam-policy "${project}" \ + --flatten='bindings[].members' \ + --filter='bindings.role=roles/logging.logWriter' \ + --format='value(bindings.members)' +)" +expected_log_writer="serviceAccount:policyengine-api-v1-modal@${OBSERVABILITY_PROJECT_ID}.iam.gserviceaccount.com" +if [[ "${log_writers}" != "${expected_log_writer}" ]]; then + echo "Unexpected project-level Cloud Logging writers: ${log_writers}" >&2 + exit 1 +fi + +gcloud logging views get-iam-policy _AllLogs \ + --bucket="${OBSERVABILITY_PROJECT_ID}" \ + --location=global \ + --project="${project}" \ + --format=json + +gcloud iam workload-identity-pools providers describe modal-api-v1 \ + --workload-identity-pool=modal-api-v1 \ + --location=global \ + --project="${project}" \ + --format='yaml(state,attributeCondition,attributeMapping,oidc)' + +for source_project in "${API_PROJECT_ID}" "${SIMULATION_ENTRY_PROJECT_ID}"; do + gcloud logging sinks describe api-v1-central-observability \ + --project="${source_project}" \ + --format='yaml(destination,filter,writerIdentity)' +done + +gcloud logging sinks describe policyengine-observability-app-logs \ + --project="${project}" \ + --format='yaml(destination,filter)' + +gcloud logging sinks describe _Default \ + --project="${project}" \ + --format='yaml(filter,exclusions)' diff --git a/gcp/observability/verify_modal_wif.py b/gcp/observability/verify_modal_wif.py new file mode 100644 index 000000000..ebe9f7a2f --- /dev/null +++ b/gcp/observability/verify_modal_wif.py @@ -0,0 +1,221 @@ +"""Verify allowed and denied Modal identities against Google WIF. + +Set ``POLICYENGINE_WIF_TEST_APP`` in the local process before ``modal run``. +The script never returns or prints the Modal identity token or Google tokens. +""" + +# pyright: reportMissingImports=false + +from __future__ import annotations + +import base64 +import hashlib +import json +import os +import time +import urllib.error +import urllib.parse +import urllib.request + +import modal + +POOL = "modal-api-v1" +PROVIDER = "modal-api-v1" +APP_NAME = os.environ.get( + "POLICYENGINE_WIF_TEST_APP", + "policyengine-observability-wif-denied-test", +) + +app = modal.App(APP_NAME) + + +def _required_environment(name: str) -> str: + value = os.getenv(name) + if not value: + raise SystemExit(f"Missing deployment variable: {name}") + return value + + +def _jwt_claims(token: str) -> dict[str, object]: + payload = token.split(".")[1] + payload += "=" * (-len(payload) % 4) + return json.loads(base64.urlsafe_b64decode(payload)) + + +def _post_form(url: str, values: dict[str, str]) -> dict[str, object]: + request = urllib.request.Request( + url, + data=urllib.parse.urlencode(values).encode(), + headers={"Content-Type": "application/x-www-form-urlencoded"}, + ) + with urllib.request.urlopen(request, timeout=15) as response: + return json.loads(response.read()) + + +def _post_json( + url: str, + payload: dict[str, object], + *, + bearer_token: str, +) -> tuple[int, dict[str, object]]: + request = urllib.request.Request( + url, + data=json.dumps(payload).encode(), + headers={ + "Authorization": f"Bearer {bearer_token}", + "Content-Type": "application/json", + }, + ) + with urllib.request.urlopen(request, timeout=15) as response: + body = response.read() + return response.status, json.loads(body) if body else {} + + +def _error_result(error: urllib.error.HTTPError) -> dict[str, object]: + try: + payload = json.loads(error.read()) + except (json.JSONDecodeError, UnicodeDecodeError): + payload = {} + error_payload = payload.get("error", {}) + if isinstance(error_payload, dict): + error_name = error_payload.get("status") or error_payload.get("error") + else: + error_name = error_payload + return { + "http_status": error.code, + "error": error_name, + } + + +@app.function(timeout=60) +def verify_identity( + project: str, + project_number: str, + workspace_id_digest: str, +) -> dict[str, object]: + service_account = f"policyengine-api-v1-modal@{project}.iam.gserviceaccount.com" + collector = ( + "https://policyengine-api-v1-otel-collector-" + f"{project_number}.us-central1.run.app" + ) + identity_token = os.environ["MODAL_IDENTITY_TOKEN"] + claims = _jwt_claims(identity_token) + safe_claims = { + key: claims.get(key) + for key in ( + "environment_name", + "app_name", + "function_name", + "aud", + "iss", + ) + } + workspace_id = str(claims.get("workspace_id", "")) + workspace_id_matches = ( + hashlib.sha256(workspace_id.encode()).hexdigest() == workspace_id_digest + ) + audience = ( + "//iam.googleapis.com/projects/" + f"{project_number}/locations/global/workloadIdentityPools/{POOL}" + f"/providers/{PROVIDER}" + ) + try: + sts_payload = _post_form( + "https://sts.googleapis.com/v1/token", + { + "audience": audience, + "grant_type": "urn:ietf:params:oauth:grant-type:token-exchange", + "requested_token_type": ( + "urn:ietf:params:oauth:token-type:access_token" + ), + "scope": "https://www.googleapis.com/auth/cloud-platform", + "subject_token": identity_token, + "subject_token_type": "urn:ietf:params:oauth:token-type:jwt", + }, + ) + except urllib.error.HTTPError as error: + return { + "claims": safe_claims, + "workspace_id_matches": workspace_id_matches, + "token_exchange": _error_result(error), + } + + federated_token = str(sts_payload["access_token"]) + service_account_url = ( + "https://iamcredentials.googleapis.com/v1/projects/-/serviceAccounts/" + f"{service_account}" + ) + access_status, access_payload = _post_json( + f"{service_account_url}:generateAccessToken", + { + "scope": ["https://www.googleapis.com/auth/cloud-platform"], + "lifetime": "300s", + }, + bearer_token=federated_token, + ) + service_access_token = str(access_payload["accessToken"]) + + identity_status, identity_payload = _post_json( + f"{service_account_url}:generateIdToken", + {"audience": collector, "includeEmail": True}, + bearer_token=federated_token, + ) + collector_identity_token = str(identity_payload["token"]) + collector_request = urllib.request.Request( + collector, + headers={"Authorization": f"Bearer {collector_identity_token}"}, + ) + try: + with urllib.request.urlopen(collector_request, timeout=15) as response: + collector_status = response.status + except urllib.error.HTTPError as error: + collector_status = error.code + + verification_id = f"modal-wif-{int(time.time())}" + logging_status, _ = _post_json( + "https://logging.googleapis.com/v2/entries:write", + { + "logName": f"projects/{project}/logs/policyengine-api-v1-modal", + "resource": { + "type": "global", + "labels": {"project_id": project}, + }, + "entries": [ + { + "insertId": verification_id, + "jsonPayload": { + "schema_version": "policyengine.observability.v2", + "verification_id": verification_id, + "service.namespace": "policyengine.api-v1", + "service.name": safe_claims["app_name"], + "message": "Modal WIF verification", + }, + } + ], + }, + bearer_token=service_access_token, + ) + return { + "claims": safe_claims, + "workspace_id_matches": workspace_id_matches, + "token_exchange": {"http_status": 200}, + "service_account_access": {"http_status": access_status}, + "service_account_identity": {"http_status": identity_status}, + "collector_http_status": collector_status, + "logging_http_status": logging_status, + "verification_id": verification_id, + } + + +@app.local_entrypoint() +def main() -> None: + project = _required_environment("OBSERVABILITY_PROJECT_ID") + project_number = _required_environment("OBSERVABILITY_PROJECT_NUMBER") + workspace_id = _required_environment("MODAL_WORKSPACE_ID") + workspace_id_digest = hashlib.sha256(workspace_id.encode()).hexdigest() + result = getattr(verify_identity, "remote")( + project, + project_number, + workspace_id_digest, + ) + print(json.dumps(result, sort_keys=True)) diff --git a/gcp/observability/verify_otel.py b/gcp/observability/verify_otel.py new file mode 100644 index 000000000..ab30a7485 --- /dev/null +++ b/gcp/observability/verify_otel.py @@ -0,0 +1,252 @@ +"""Send and verify synthetic OTLP signals through the Cloud Run collector.""" + +from __future__ import annotations + +import argparse +import json +import os +import secrets +import subprocess +import time +import urllib.error +import urllib.parse +import urllib.request + +import grpc +from opentelemetry.exporter.otlp.proto.grpc.metric_exporter import ( + OTLPMetricExporter, +) +from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import ( + OTLPSpanExporter, +) +from opentelemetry.proto.collector.logs.v1.logs_service_pb2 import ( + ExportLogsServiceRequest, +) +from opentelemetry.proto.collector.logs.v1.logs_service_pb2_grpc import ( + LogsServiceStub, +) +from opentelemetry.proto.common.v1.common_pb2 import AnyValue, KeyValue +from opentelemetry.proto.logs.v1.logs_pb2 import ( + LogRecord, + ResourceLogs, + ScopeLogs, +) +from opentelemetry.proto.resource.v1.resource_pb2 import ( + Resource as ProtoResource, +) +from opentelemetry.sdk.metrics import MeterProvider +from opentelemetry.sdk.metrics.export import PeriodicExportingMetricReader +from opentelemetry.sdk.resources import Resource +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor + +PROJECT = os.environ.get("OBSERVABILITY_PROJECT_ID", "") +METRIC_NAME = "policyengine.verification.counter" +METRIC_TYPE = f"prometheus.googleapis.com/{METRIC_NAME}/counter" + + +def _gcloud_output(*arguments: str) -> str: + return subprocess.check_output( + ["gcloud", *arguments], + text=True, + ).strip() + + +def _identity_token(service_account: str, audience: str) -> str: + return _gcloud_output( + "auth", + "print-identity-token", + f"--impersonate-service-account={service_account}", + f"--audiences={audience}", + ) + + +def _access_token() -> str: + return _gcloud_output("auth", "print-access-token") + + +def _authorized_get(url: str, token: str) -> dict[str, object]: + request = urllib.request.Request( + url, + headers={"Authorization": f"Bearer {token}"}, + ) + with urllib.request.urlopen(request, timeout=15) as response: + return json.loads(response.read()) + + +def _send_signals(endpoint: str, service_account: str) -> dict[str, object]: + identity_token = _identity_token(service_account, endpoint) + headers = (("authorization", f"Bearer {identity_token}"),) + credentials = grpc.ssl_channel_credentials() + resource = Resource.create( + { + "service.name": "policyengine-observability-verification", + "service.namespace": "policyengine.api-v1", + "service.version": "2.0.0-verification", + "service.role": "verification", + "deployment.environment.name": "staging", + "cloud.platform": "gcp_cloud_run", + "cloud.region": "us-central1", + } + ) + + trace_exporter = OTLPSpanExporter( + endpoint=endpoint, + credentials=credentials, + headers=headers, + timeout=10, + ) + trace_provider = TracerProvider(resource=resource) + trace_provider.add_span_processor(SimpleSpanProcessor(trace_exporter)) + tracer = trace_provider.get_tracer("policyengine.observability.verification") + with tracer.start_as_current_span("policyengine.observability.verify") as span: + span.set_attribute("policyengine.verification", True) + trace_id = f"{span.get_span_context().trace_id:032x}" + trace_provider.force_flush(timeout_millis=15_000) + trace_provider.shutdown() + + metric_exporter = OTLPMetricExporter( + endpoint=endpoint, + credentials=credentials, + headers=headers, + timeout=10, + ) + metric_reader = PeriodicExportingMetricReader( + metric_exporter, + export_interval_millis=60_000, + export_timeout_millis=10_000, + ) + metric_provider = MeterProvider( + resource=resource, + metric_readers=[metric_reader], + ) + meter = metric_provider.get_meter("policyengine.observability.verification") + counter = meter.create_counter(METRIC_NAME) + counter.add(1, {"outcome": "success"}) + metric_provider.shutdown() + + host = urllib.parse.urlparse(endpoint).netloc + channel = grpc.secure_channel(host, credentials) + logs_stub = LogsServiceStub(channel) + log_request = ExportLogsServiceRequest( + resource_logs=[ + ResourceLogs( + resource=ProtoResource( + attributes=[ + KeyValue( + key="service.name", + value=AnyValue( + string_value="policyengine-observability-verification" + ), + ) + ] + ), + scope_logs=[ + ScopeLogs( + log_records=[ + LogRecord( + time_unix_nano=time.time_ns(), + severity_text="INFO", + body=AnyValue( + string_value="collector log rejection verification" + ), + ) + ] + ) + ], + ) + ] + ) + log_status = "accepted" + try: + logs_stub.Export(log_request, timeout=10, metadata=headers) + except grpc.RpcError as error: + log_status = error.code().name + finally: + channel.close() + + return { + "trace_id": trace_id, + "metric_type": METRIC_TYPE, + "log_export_status": log_status, + } + + +def _wait_for_storage( + *, trace_id: str, metric_type: str, timeout_seconds: int +) -> dict[str, object]: + access_token = _access_token() + deadline = time.monotonic() + timeout_seconds + start = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime(time.time() - 300)) + end = time.strftime( + "%Y-%m-%dT%H:%M:%SZ", time.gmtime(time.time() + timeout_seconds + 60) + ) + metric_filter = urllib.parse.quote(f'metric.type = "{metric_type}"') + metric_url = ( + f"https://monitoring.googleapis.com/v3/projects/{PROJECT}/timeSeries" + f"?filter={metric_filter}&interval.startTime={start}" + f"&interval.endTime={end}&view=HEADERS" + ) + trace_url = ( + f"https://cloudtrace.googleapis.com/v1/projects/{PROJECT}/traces/{trace_id}" + ) + trace_found = False + metric_found = False + while time.monotonic() < deadline and not (trace_found and metric_found): + if not trace_found: + try: + trace_payload = _authorized_get(trace_url, access_token) + trace_found = bool(trace_payload.get("spans")) + except urllib.error.HTTPError as error: + if error.code != 404: + raise + if not metric_found: + try: + metric_payload = _authorized_get(metric_url, access_token) + metric_found = bool(metric_payload.get("timeSeries")) + except urllib.error.HTTPError as error: + if error.code != 404: + raise + if not (trace_found and metric_found): + time.sleep(5) + return { + "trace_stored": trace_found, + "metric_stored": metric_found, + } + + +def main() -> int: + if not PROJECT: + raise SystemExit("Missing deployment variable: OBSERVABILITY_PROJECT_ID") + parser = argparse.ArgumentParser() + parser.add_argument("--endpoint", required=True) + parser.add_argument("--service-account", required=True) + parser.add_argument("--timeout-seconds", type=int, default=120) + arguments = parser.parse_args() + + # Add a run identifier without placing high-cardinality values on the metric. + run_id = secrets.token_hex(4) + result = _send_signals(arguments.endpoint, arguments.service_account) + result.update( + _wait_for_storage( + trace_id=str(result["trace_id"]), + metric_type=str(result["metric_type"]), + timeout_seconds=arguments.timeout_seconds, + ) + ) + result["run_id"] = run_id + print(json.dumps(result, sort_keys=True)) + + return ( + 0 + if ( + result["trace_stored"] + and result["metric_stored"] + and result["log_export_status"] != "accepted" + ) + else 1 + ) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/gcp/observability/workload-inventory.template.yaml b/gcp/observability/workload-inventory.template.yaml new file mode 100644 index 000000000..f74db600e --- /dev/null +++ b/gcp/observability/workload-inventory.template.yaml @@ -0,0 +1,70 @@ +schema_version: policyengine.observability.workloads.v1 +destination_project: + project_id: ${OBSERVABILITY_PROJECT_ID} + project_number: "${OBSERVABILITY_PROJECT_NUMBER}" + log_bucket: ${OBSERVABILITY_PROJECT_ID} + log_bucket_location: global + log_retention_days: 30 + +cloud_run: + - service_name: policyengine-api + source_project: ${API_PROJECT_ID} + region: us-central1 + environment: production + service_account: policyengine-api-cr-runtime@${API_PROJECT_ID}.iam.gserviceaccount.com + signals: [logs, traces, metrics] + include_tagged_revisions: true + - service_name: policyengine-api-staging + source_project: ${API_PROJECT_ID} + region: us-central1 + environment: staging + service_account: policyengine-api-cr-staging@${API_PROJECT_ID}.iam.gserviceaccount.com + signals: [logs, traces, metrics] + include_tagged_revisions: true + - service_name: policyengine-simulation-entry + source_project: ${SIMULATION_ENTRY_PROJECT_ID} + region: us-central1 + environment: production + service_account: sim-entry-prod-runtime@${SIMULATION_ENTRY_PROJECT_ID}.iam.gserviceaccount.com + signals: [logs, traces, metrics] + include_tagged_revisions: true + - service_name: policyengine-simulation-entry-staging + source_project: ${SIMULATION_ENTRY_PROJECT_ID} + region: us-central1 + environment: staging + service_account: sim-entry-beta-runtime@${SIMULATION_ENTRY_PROJECT_ID}.iam.gserviceaccount.com + signals: [logs, traces, metrics] + include_tagged_revisions: true + +modal: + workspace_name: policyengine + workspace_id: ${MODAL_WORKSPACE_ID} + environments: + - name: main + deployment_environment: production + - name: staging + deployment_environment: staging + applications: + - match: exact + value: policyengine-simulation-gateway + signals: [logs, traces, metrics] + - match: prefix + value: policyengine-simulation-py + name_constraint: '^policyengine-simulation-py[0-9]+-[0-9]+-[0-9]+$' + signals: [logs, traces, metrics] + +excluded: + modal_exact: + - policyengine-simulation-precompute + - policyengine-simulation-executor-smoke + - policyengine-simulation-gateway-smoke + modal_prefixes: + - policyengine-simulation-gateway-ephemeral- + - policyengine-simulation-py4-12-0-ephemeral- + - policyengine-household-api- + - policyengine-uk-chat + - peukchat- + policy: >- + Every Cloud Run service and Modal application absent from the allowlist is + excluded. Excluded application repositories, deployments, and test suites + are outside this change. diff --git a/policyengine_api/observability.py b/policyengine_api/observability.py index f7ef90a80..bcddc8554 100644 --- a/policyengine_api/observability.py +++ b/policyengine_api/observability.py @@ -7,16 +7,16 @@ from policyengine_observability import ( DeploymentIdentity, + GoogleCloudLogFormatter, LoggingConfig, ObservabilityConfig, ObservabilityRuntime, ServiceIdentity, + StdoutLogDestination, configure, ) -GOOGLE_CLOUD_PROJECT = "policyengine-observability" -SERVICE_NAMESPACE = "policyengine.api-v1" APPLICATION_ATTRIBUTE_KEYS = frozenset( { "backend", @@ -79,11 +79,16 @@ def _package_version() -> str: def _build_runtime() -> ObservabilityRuntime: environment = os.getenv("APP_ENVIRONMENT", "local").strip() or "local" + trace_project = os.getenv("OBSERVABILITY_TRACE_PROJECT_ID", "").strip() + formatter = GoogleCloudLogFormatter(trace_project) if trace_project else None return configure( ObservabilityConfig.from_env( service=ServiceIdentity( name="policyengine-api", - namespace=SERVICE_NAMESPACE, + namespace=os.getenv( + "OBSERVABILITY_SERVICE_NAMESPACE", + "policyengine.api-v1", + ), version=_package_version(), role="api", ), @@ -93,9 +98,8 @@ def _build_runtime() -> ObservabilityRuntime: region=os.getenv("CLOUD_RUN_REGION") or "us-central1", instance_id=os.getenv("K_REVISION"), ), - google_cloud_project_id=GOOGLE_CLOUD_PROJECT, logging=LoggingConfig( - stdout_enabled=True, + destinations=(StdoutLogDestination(formatter=formatter),), capture_standard_library=True, ), application_attribute_keys=APPLICATION_ATTRIBUTE_KEYS, diff --git a/policyengine_api/services/economy_service.py b/policyengine_api/services/economy_service.py index f1b564c37..85d062a6f 100644 --- a/policyengine_api/services/economy_service.py +++ b/policyengine_api/services/economy_service.py @@ -25,6 +25,7 @@ from policyengine_api.data.places import validate_place_code from policyengine_api.gcp_logging import logger from policyengine_api.libs.simulation_entrypoint import simulation_entrypoint +from policyengine_api.observability import runtime as observability_runtime from policyengine_api.services.budget_window_cache import BudgetWindowCache from policyengine_api.services.policy_service import PolicyService from policyengine_api.services.reform_impacts_service import ( @@ -273,6 +274,7 @@ def _budget_window_cache(self) -> BudgetWindowCache: def _simulation_gateway(self): return self._injected_simulation_entrypoint or simulation_entrypoint + @observability_runtime.span("economy.load_policies") def _get_policy_jsons( self, country_id: str, @@ -301,6 +303,7 @@ def _parse_json_object(value: dict[str, Any] | str) -> dict[str, Any]: raise TypeError("Expected a JSON object") return parsed + @observability_runtime.span("economy.request") def get_economic_impact( self, country_id: str, @@ -322,6 +325,12 @@ def get_economic_impact( the status is "computing" or "error". """ + observability_runtime.set_context( + country_id=country_id, + policy_id=policy_id, + baseline_policy_id=baseline_policy_id, + simulation_year=time_period, + ) try: # Normalize region early for US; this allows us to accommodate legacy # regions that don't contain a region prefix. @@ -348,6 +357,7 @@ def get_economic_impact( print(f"Error getting economic impact: {str(e)}") raise e + @observability_runtime.span("economy.budget_window_request") def get_budget_window_economic_impact( self, country_id: str, @@ -362,6 +372,13 @@ def get_budget_window_economic_impact( target: Literal["general", "cliff"] = "general", max_active_years: int = BUDGET_WINDOW_MAX_ACTIVE_YEARS, ) -> BudgetWindowEconomicImpactResult: + observability_runtime.set_context( + country_id=country_id, + policy_id=policy_id, + baseline_policy_id=baseline_policy_id, + start_year=start_year, + window_size=window_size, + ) try: if country_id == "us": region = normalize_us_region(region) @@ -519,6 +536,7 @@ def _build_budget_window_batch_payload( sim_params["target"] = setup_options.target return sim_params + @observability_runtime.span("economy.start_budget_window_batch") def _start_budget_window_batch( self, *, @@ -567,6 +585,7 @@ def _build_budget_window_submission_error_message( return str(error) + @observability_runtime.span("economy.poll_budget_window_batch") def _get_budget_window_result_from_batch_job_id( self, *, @@ -728,6 +747,7 @@ def _build_economic_impact_setup_options( } ) + @observability_runtime.span("economy.resolve_cached_or_new_impact") def _get_or_create_economic_impact( self, setup_options: EconomicImpactSetupOptions ) -> EconomicImpactResult: @@ -757,6 +777,10 @@ def _get_or_create_economic_impact( impact_action: ImpactAction = self._determine_impact_action( most_recent_impact=most_recent_impact ) + observability_runtime.event( + "economy.cache_decision", + attributes={"cache_event": impact_action.value}, + ) if impact_action == ImpactAction.COMPLETED: logger.log_struct( @@ -822,6 +846,7 @@ def _get_or_create_economic_impact( raise ValueError(f"Unexpected impact action: {impact_action}") + @observability_runtime.span("economy.resolve_runtime_bundle") def _resolve_runtime_bundle_for_setup_options( self, setup_options: EconomicImpactSetupOptions, @@ -976,6 +1001,7 @@ def _determine_impact_action( else: raise ValueError(f"Unknown impact status: {status}") + @observability_runtime.span("economy.handle_execution_state") def _handle_execution_state( self, setup_options: EconomicImpactSetupOptions, @@ -1044,6 +1070,7 @@ def _handle_execution_state( else: raise ValueError(f"Unexpected sim API execution state: {execution_state}") + @observability_runtime.span("economy.read_completed_impact") def _handle_completed_impact( self, setup_options: EconomicImpactSetupOptions, @@ -1103,6 +1130,7 @@ def _record_uncertifiable_stored_impact( except Exception: pass + @observability_runtime.span("economy.read_failed_impact") def _handle_failed_impact( self, most_recent_impact: ReformImpact, @@ -1123,6 +1151,7 @@ def _handle_failed_impact( ) ) + @observability_runtime.span("economy.poll_active_impact") def _handle_computing_impact( self, setup_options: EconomicImpactSetupOptions, @@ -1148,6 +1177,7 @@ def _handle_computing_impact( ) raise + @observability_runtime.span("economy.submit_impact") def _handle_create_impact( self, setup_options: EconomicImpactSetupOptions, @@ -1482,6 +1512,7 @@ def _stable_config_hash(self, payload: dict[str, Any]) -> str: # are written separately because the service relies upon mutating an original # 'computing' record to 'ok' or 'error' status, rather than creating a new record. # This should be addressed in the future. + @observability_runtime.span("economy.persist_computing_impact") def _set_reform_impact_computing( self, setup_options: EconomicImpactSetupOptions, @@ -1515,6 +1546,7 @@ def _set_reform_impact_computing( ) raise e + @observability_runtime.span("economy.persist_completed_impact") def _set_reform_impact_complete( self, setup_options: EconomicImpactSetupOptions, @@ -1545,6 +1577,7 @@ def _set_reform_impact_complete( ) raise e + @observability_runtime.span("economy.persist_failed_impact") def _set_reform_impact_error( self, setup_options: EconomicImpactSetupOptions, diff --git a/pyproject.toml b/pyproject.toml index 9b6276a62..d674ec0f3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -39,7 +39,7 @@ dependencies = [ "microdf_python>=1.0.0", "openai", "packaging>=24,<27", - "policyengine-observability[flask,google,httpx,otlp-grpc] @ https://github.com/PolicyEngine/policyengine-observability/archive/a25e63c147064918ad20b2e396cac84900727df5.tar.gz", + "policyengine-observability[flask,google,httpx,otlp-grpc] @ https://github.com/PolicyEngine/policyengine-observability/archive/4ea849342462eae314bd49978e5c672bf9dcb05b.tar.gz", "policyengine_canada==0.96.3", "policyengine-ng==0.5.1", "policyengine-il==0.1.0", diff --git a/tests/unit/libs/test_simulation_entrypoint.py b/tests/unit/libs/test_simulation_entrypoint.py index 517d8429e..21c04656c 100644 --- a/tests/unit/libs/test_simulation_entrypoint.py +++ b/tests/unit/libs/test_simulation_entrypoint.py @@ -6,18 +6,12 @@ """ import os -import sys -from types import SimpleNamespace -from unittest.mock import MagicMock, patch +from unittest.mock import patch import httpx import pytest from flask import Flask, g -sys.modules.setdefault( - "policyengine_api.gcp_logging", - SimpleNamespace(logger=MagicMock()), -) os.environ.setdefault("FLASK_DEBUG", "1") from policyengine_api.constants import ( # noqa: E402 diff --git a/tests/unit/test_cloud_run_deploy_scripts.py b/tests/unit/test_cloud_run_deploy_scripts.py index b63ce7d79..24cbe8c89 100644 --- a/tests/unit/test_cloud_run_deploy_scripts.py +++ b/tests/unit/test_cloud_run_deploy_scripts.py @@ -19,6 +19,7 @@ TEST_V2_RUNTIME_SECRET_RESOURCE = ( "projects/test-project/secrets/v2-runtime-database-url/versions/latest" ) +TEST_OTEL_ENDPOINT = "https://collector.example.test" CLOUD_RUN_SERVICE_SCRIPTS = ( "scripts/deploy_cloud_run_candidate.sh", "scripts/capture_cloud_run_service_state.sh", @@ -73,6 +74,15 @@ def _v2_target_env() -> dict[str, str]: } +def _observability_env() -> dict[str, str]: + return { + "OBSERVABILITY_SERVICE_NAMESPACE": "policyengine.api-v1", + "OBSERVABILITY_TRACE_PROJECT_ID": "central-observability", + "OTEL_EXPORTER_OTLP_ENDPOINT": TEST_OTEL_ENDPOINT, + "POLICYENGINE_OTEL_GOOGLE_AUDIENCE": TEST_OTEL_ENDPOINT, + } + + def _required_runtime_env() -> dict[str, str]: return { "DEPLOYMENT_ENVIRONMENT": "production", @@ -106,6 +116,7 @@ def _required_runtime_env() -> dict[str, str]: "DB_WRITE_POLICY": "cloud_sql", "DB_READ_HOUSEHOLD": "cloud_sql", "DB_WRITE_HOUSEHOLD": "cloud_sql", + **_observability_env(), **_v2_target_env(), **_gateway_auth_env(), } @@ -610,6 +621,7 @@ def test_validate_cloud_run_deploy_env_accepts_direct_mode_from_environment(): ), **_v2_target_env(), **_gateway_auth_env(), + **_observability_env(), ), ) @@ -766,6 +778,7 @@ def test_validate_cloud_run_deploy_env_requires_only_selected_url( ), **_v2_target_env(), **_gateway_auth_env(), + **_observability_env(), ) missing_result = _run_script( ".github/scripts/validate_cloud_run_deploy_env.sh", @@ -989,20 +1002,14 @@ def test_deploy_cloud_run_candidate_dry_run_preserves_access_and_traffic(): assert "RUNTIME_CACHE_ENVIRONMENT=production" in result.stdout assert "RUNTIME_CACHE_SERVICE=api" in result.stdout assert "APP_ENVIRONMENT=production" in result.stdout - assert ( - "OTEL_EXPORTER_OTLP_ENDPOINT=https://" - "policyengine-api-v1-otel-collector-790230211054.us-central1.run.app" - in result.stdout - ) + assert f"OTEL_EXPORTER_OTLP_ENDPOINT={TEST_OTEL_ENDPOINT}" in result.stdout + assert "OBSERVABILITY_SERVICE_NAMESPACE=policyengine.api-v1" in result.stdout + assert "OBSERVABILITY_TRACE_PROJECT_ID=central-observability" in result.stdout assert "OTEL_EXPORTER_OTLP_PROTOCOL=grpc" in result.stdout assert "OTEL_TRACES_EXPORTER=otlp" in result.stdout assert "OTEL_METRICS_EXPORTER=otlp" in result.stdout assert "OTEL_TRACES_SAMPLER_ARG=1.0" in result.stdout - assert ( - "POLICYENGINE_OTEL_GOOGLE_AUDIENCE=https://" - "policyengine-api-v1-otel-collector-790230211054.us-central1.run.app" - in result.stdout - ) + assert f"POLICYENGINE_OTEL_GOOGLE_AUDIENCE={TEST_OTEL_ENDPOINT}" in result.stdout assert ( "RUNTIME_CACHE_URL=policyengine-api-prod-runtime-cache-url:latest" in result.stdout diff --git a/tests/unit/test_observability_deployment_assets.py b/tests/unit/test_observability_deployment_assets.py new file mode 100644 index 000000000..5f048a608 --- /dev/null +++ b/tests/unit/test_observability_deployment_assets.py @@ -0,0 +1,141 @@ +from __future__ import annotations + +import json +import os +import stat +import subprocess +import sys +from pathlib import Path + +ROOT = Path(__file__).parents[2] +DEPLOY = ROOT / "gcp" / "observability" + + +def test_dashboard_is_valid_json_with_required_signals() -> None: + dashboard = json.loads((DEPLOY / "dashboard.template.json").read_text()) + serialized = json.dumps(dashboard) + assert "policyengine.request.count" in serialized + assert "policyengine.request.duration" in serialized + assert "policyengine.error.count" in serialized + assert "policyengine.telemetry.dropped" in serialized + assert "policyengine.telemetry.exporter.failure" in serialized + + +def test_collector_accepts_only_traces_and_metrics() -> None: + config = (DEPLOY / "collector" / "config.yaml").read_text() + assert "telemetry.googleapis.com:443" in config + assert "memory_limiter" in config + assert "googleclientauth" in config + assert " traces:" in config + assert " metrics:" in config + assert " logs:\n receivers:" not in config + + +def test_authorization_assets_exclude_unrelated_applications() -> None: + iam = (DEPLOY / "iam.template.yaml").read_text() + routing = (DEPLOY / "log-routing.template.yaml").read_text() + for excluded in ( + "policyengine-household-api", + "policyengine-uk-chat", + "peukchat", + "precompute", + "smoke", + "ephemeral", + ): + assert excluded not in iam + assert excluded not in routing + assert "policyengine-simulation-gateway" in iam + assert "policyengine-simulation-py" in iam + assert 'jsonPayload."service.namespace"' in routing + + +def test_deployment_templates_use_environment_placeholders() -> None: + templates = [ + DEPLOY / "iam.template.yaml", + DEPLOY / "workload-inventory.template.yaml", + DEPLOY / "log-routing.template.yaml", + DEPLOY / "alerts.template.yaml", + DEPLOY / "dashboard.template.json", + DEPLOY / "collector" / "service.template.yaml", + ] + content = "\n".join(path.read_text() for path in templates) + for variable in ( + "OBSERVABILITY_PROJECT_ID", + "OBSERVABILITY_PROJECT_NUMBER", + "API_PROJECT_ID", + "SIMULATION_ENTRY_PROJECT_ID", + "MODAL_WORKSPACE_ID", + ): + assert f"${{{variable}}}" in content + assert "workspace_id: ac-" not in content + + +def test_deployment_renderer_validates_and_does_not_print_values( + tmp_path: Path, +) -> None: + values = { + "OBSERVABILITY_PROJECT_ID": "central-observability", + "OBSERVABILITY_PROJECT_NUMBER": "123456789012", + "API_PROJECT_ID": "api-project", + "SIMULATION_ENTRY_PROJECT_ID": "simulation-entry-project", + "MODAL_WORKSPACE_ID": "ac-private-workspace", + } + environment = os.environ.copy() + environment.update(values) + result = subprocess.run( + [ + sys.executable, + str(DEPLOY / "render_deployment.py"), + "--output-dir", + str(tmp_path), + ], + check=True, + capture_output=True, + text=True, + env=environment, + ) + + assert all(value not in result.stdout for value in values.values()) + rendered_iam = (tmp_path / "iam.yaml").read_text() + assert "policyengine-otel-collector@central-observability" in rendered_iam + assert 'assertion.workspace_id == "ac-private-workspace"' in rendered_iam + assert stat.S_IMODE((tmp_path / "iam.yaml").stat().st_mode) == 0o600 + json.loads((tmp_path / "dashboard.json").read_text()) + + +def test_deployment_renderer_rejects_missing_values(tmp_path: Path) -> None: + environment = os.environ.copy() + for variable in ( + "OBSERVABILITY_PROJECT_ID", + "OBSERVABILITY_PROJECT_NUMBER", + "API_PROJECT_ID", + "SIMULATION_ENTRY_PROJECT_ID", + "MODAL_WORKSPACE_ID", + ): + environment.pop(variable, None) + + result = subprocess.run( + [ + sys.executable, + str(DEPLOY / "render_deployment.py"), + "--output-dir", + str(tmp_path), + ], + check=False, + capture_output=True, + text=True, + env=environment, + ) + + assert result.returncode != 0 + assert "Missing deployment variables:" in result.stderr + assert not list(tmp_path.iterdir()) + + +def test_verification_script_has_valid_shell_syntax() -> None: + subprocess.run( + ["bash", "-n", str(DEPLOY / "verify.sh")], + check=True, + capture_output=True, + text=True, + ) diff --git a/tests/unit/test_observability_runtime.py b/tests/unit/test_observability_runtime.py new file mode 100644 index 000000000..12f199ebc --- /dev/null +++ b/tests/unit/test_observability_runtime.py @@ -0,0 +1,23 @@ +from policyengine_observability import ( + GoogleCloudLogFormatter, + StdoutLogDestination, +) + +from policyengine_api.observability import _build_runtime + + +def test_runtime_uses_consumer_owned_identity_and_stdout(monkeypatch): + monkeypatch.setenv("OTEL_SDK_DISABLED", "true") + monkeypatch.setenv("OBSERVABILITY_SERVICE_NAMESPACE", "example.stack") + monkeypatch.setenv("OBSERVABILITY_TRACE_PROJECT_ID", "trace-project") + + runtime = _build_runtime() + try: + assert runtime.config.service.namespace == "example.stack" + assert len(runtime.config.logging.destinations) == 1 + destination = runtime.config.logging.destinations[0] + assert isinstance(destination, StdoutLogDestination) + assert isinstance(destination.formatter, GoogleCloudLogFormatter) + assert destination.formatter.project_id == "trace-project" + finally: + runtime.shutdown() diff --git a/uv.lock b/uv.lock index 7d59431d7..eb8ca1b39 100644 --- a/uv.lock +++ b/uv.lock @@ -2848,7 +2848,7 @@ requires-dist = [ { name = "policyengine-canada", specifier = "==0.96.3" }, { name = "policyengine-il", specifier = "==0.1.0" }, { name = "policyengine-ng", specifier = "==0.5.1" }, - { name = "policyengine-observability", extras = ["flask", "google", "httpx", "otlp-grpc"], url = "https://github.com/PolicyEngine/policyengine-observability/archive/a25e63c147064918ad20b2e396cac84900727df5.tar.gz" }, + { name = "policyengine-observability", extras = ["flask", "google", "httpx", "otlp-grpc"], url = "https://github.com/PolicyEngine/policyengine-observability/archive/4ea849342462eae314bd49978e5c672bf9dcb05b.tar.gz" }, { name = "psycopg", extras = ["binary"], specifier = ">=3.3,<4" }, { name = "pydantic" }, { name = "pymysql" }, @@ -2952,9 +2952,9 @@ wheels = [ [[package]] name = "policyengine-observability" -version = "1.4.1" -source = { url = "https://github.com/PolicyEngine/policyengine-observability/archive/a25e63c147064918ad20b2e396cac84900727df5.tar.gz" } -sdist = { hash = "sha256:6504945c1768f555087a53b7b26199c394f6482196c9c1828684d57d823c22ae" } +version = "2.0.0" +source = { url = "https://github.com/PolicyEngine/policyengine-observability/archive/4ea849342462eae314bd49978e5c672bf9dcb05b.tar.gz" } +sdist = { hash = "sha256:af0090a654745f32d8cdbfa0ddc1af295ba8ebb28aeb4eeaf8ebdf0c809562a4" } [package.optional-dependencies] flask = [ From aa7ac3a20b205e5d62de2e2670f87042e89f289f Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 22 Sep 2026 22:31:25 +0400 Subject: [PATCH 05/10] Require full observability trace sampling --- policyengine_api/observability.py | 48 +++++++++++++----------- tests/unit/test_observability_runtime.py | 2 + 2 files changed, 28 insertions(+), 22 deletions(-) diff --git a/policyengine_api/observability.py b/policyengine_api/observability.py index bcddc8554..25fbba0b1 100644 --- a/policyengine_api/observability.py +++ b/policyengine_api/observability.py @@ -3,6 +3,7 @@ from __future__ import annotations import os +from dataclasses import replace from importlib.metadata import PackageNotFoundError, version from policyengine_observability import ( @@ -81,30 +82,33 @@ def _build_runtime() -> ObservabilityRuntime: environment = os.getenv("APP_ENVIRONMENT", "local").strip() or "local" trace_project = os.getenv("OBSERVABILITY_TRACE_PROJECT_ID", "").strip() formatter = GoogleCloudLogFormatter(trace_project) if trace_project else None - return configure( - ObservabilityConfig.from_env( - service=ServiceIdentity( - name="policyengine-api", - namespace=os.getenv( - "OBSERVABILITY_SERVICE_NAMESPACE", - "policyengine.api-v1", - ), - version=_package_version(), - role="api", + config = ObservabilityConfig.from_env( + service=ServiceIdentity( + name="policyengine-api", + namespace=os.getenv( + "OBSERVABILITY_SERVICE_NAMESPACE", + "policyengine.api-v1", ), - deployment=DeploymentIdentity( - environment=environment, - platform="google_cloud_run", - region=os.getenv("CLOUD_RUN_REGION") or "us-central1", - instance_id=os.getenv("K_REVISION"), - ), - logging=LoggingConfig( - destinations=(StdoutLogDestination(formatter=formatter),), - capture_standard_library=True, - ), - application_attribute_keys=APPLICATION_ATTRIBUTE_KEYS, - ) + version=_package_version(), + role="api", + ), + deployment=DeploymentIdentity( + environment=environment, + platform="google_cloud_run", + region=os.getenv("CLOUD_RUN_REGION") or "us-central1", + instance_id=os.getenv("K_REVISION"), + ), + logging=LoggingConfig( + destinations=(StdoutLogDestination(formatter=formatter),), + capture_standard_library=True, + ), + application_attribute_keys=APPLICATION_ATTRIBUTE_KEYS, + ) + config = replace( + config, + otel=replace(config.otel, sampling_ratio=1.0), ) + return configure(config) runtime = _build_runtime() diff --git a/tests/unit/test_observability_runtime.py b/tests/unit/test_observability_runtime.py index 12f199ebc..958863686 100644 --- a/tests/unit/test_observability_runtime.py +++ b/tests/unit/test_observability_runtime.py @@ -8,12 +8,14 @@ def test_runtime_uses_consumer_owned_identity_and_stdout(monkeypatch): monkeypatch.setenv("OTEL_SDK_DISABLED", "true") + monkeypatch.setenv("OTEL_TRACES_SAMPLER_ARG", "0.01") monkeypatch.setenv("OBSERVABILITY_SERVICE_NAMESPACE", "example.stack") monkeypatch.setenv("OBSERVABILITY_TRACE_PROJECT_ID", "trace-project") runtime = _build_runtime() try: assert runtime.config.service.namespace == "example.stack" + assert runtime.config.otel.sampling_ratio == 1.0 assert len(runtime.config.logging.destinations) == 1 destination = runtime.config.logging.destinations[0] assert isinstance(destination, StdoutLogDestination) From cc68f331647b71d5902b6108f277f0edda4fc762 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 23 Sep 2026 16:26:57 +0400 Subject: [PATCH 06/10] Use released observability package --- changelog.d/3847.changed.md | 2 +- gcp/observability/README.md | 13 +++++------ pyproject.toml | 2 +- uv.lock | 44 ++++++------------------------------- 4 files changed, 15 insertions(+), 46 deletions(-) diff --git a/changelog.d/3847.changed.md b/changelog.d/3847.changed.md index dba636323..1cd7dafbd 100644 --- a/changelog.d/3847.changed.md +++ b/changelog.d/3847.changed.md @@ -1,3 +1,3 @@ Route API v1 structured logs, traces, and metrics through the explicit -policyengine-observability version 2 runtime and propagate request context to +policyengine-observability 3.x runtime and propagate request context to the simulation entry service. diff --git a/gcp/observability/README.md b/gcp/observability/README.md index f023e7271..e5f8c60d9 100644 --- a/gcp/observability/README.md +++ b/gcp/observability/README.md @@ -262,10 +262,9 @@ The infrastructure portion of this runbook was applied and verified on - the project currently has no alert notification channel, so the policies record incidents without sending email, Slack, or paging notifications. -The package and consumer service rollout remains pending until the three draft -pull requests are reviewed, the package is published as version 2.0, and the -temporary Git source pins in both consumer repositories are replaced with the -published version. Run the synthetic cross-service request, volume and cost -measurement, and destination comparison after those deployments. Record the -deployed revisions and the observation interval here before declaring the -consumer rollout complete. +The package was published as version 3.0.0 on 2026-09-23. The consumer service +rollout remains pending until the two consumer pull requests are reviewed and +deployed. Run the synthetic cross-service request, volume and cost measurement, +and destination comparison after those deployments. Record the deployed +revisions and the observation interval here before declaring the consumer +rollout complete. diff --git a/pyproject.toml b/pyproject.toml index d674ec0f3..9bc3a6458 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -39,7 +39,7 @@ dependencies = [ "microdf_python>=1.0.0", "openai", "packaging>=24,<27", - "policyengine-observability[flask,google,httpx,otlp-grpc] @ https://github.com/PolicyEngine/policyengine-observability/archive/4ea849342462eae314bd49978e5c672bf9dcb05b.tar.gz", + "policyengine-observability[flask,google,httpx,otlp-grpc]>=3,<4", "policyengine_canada==0.96.3", "policyengine-ng==0.5.1", "policyengine-il==0.1.0", diff --git a/uv.lock b/uv.lock index eb8ca1b39..b407e8151 100644 --- a/uv.lock +++ b/uv.lock @@ -2848,7 +2848,7 @@ requires-dist = [ { name = "policyengine-canada", specifier = "==0.96.3" }, { name = "policyengine-il", specifier = "==0.1.0" }, { name = "policyengine-ng", specifier = "==0.5.1" }, - { name = "policyengine-observability", extras = ["flask", "google", "httpx", "otlp-grpc"], url = "https://github.com/PolicyEngine/policyengine-observability/archive/4ea849342462eae314bd49978e5c672bf9dcb05b.tar.gz" }, + { name = "policyengine-observability", extras = ["flask", "google", "httpx", "otlp-grpc"], specifier = ">=3,<4" }, { name = "psycopg", extras = ["binary"], specifier = ">=3.3,<4" }, { name = "pydantic" }, { name = "pymysql" }, @@ -2952,9 +2952,12 @@ wheels = [ [[package]] name = "policyengine-observability" -version = "2.0.0" -source = { url = "https://github.com/PolicyEngine/policyengine-observability/archive/4ea849342462eae314bd49978e5c672bf9dcb05b.tar.gz" } -sdist = { hash = "sha256:af0090a654745f32d8cdbfa0ddc1af295ba8ebb28aeb4eeaf8ebdf0c809562a4" } +version = "3.0.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/44/69/963d1044a8519e2967b473db1ff1fe304966c9f93396660252bd3e339be5/policyengine_observability-3.0.0.tar.gz", hash = "sha256:bcf8027e2526c22ac490e0ef5f9e1fe89254e08b479611fa16a94142e54c34d9", size = 122222, upload-time = "2026-09-23T10:58:47.354Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/cb/f2/8174cbb0378573553fc6eea0073063fb4bcc197cbdbbc2170f0e73bba57a/policyengine_observability-3.0.0-py3-none-any.whl", hash = "sha256:8ce4cf186c81ba136fd824c6c81c78d4b637da177c6ac657d15ffb538fe88d75", size = 40200, upload-time = "2026-09-23T10:58:45.799Z" }, +] [package.optional-dependencies] flask = [ @@ -2973,39 +2976,6 @@ otlp-grpc = [ { name = "opentelemetry-sdk" }, ] -[package.metadata] -requires-dist = [ - { name = "build", marker = "extra == 'dev'" }, - { name = "coverage", marker = "extra == 'dev'" }, - { name = "fastapi", marker = "extra == 'all'" }, - { name = "fastapi", marker = "extra == 'fastapi'" }, - { name = "flask", marker = "extra == 'all'", specifier = ">=2.2" }, - { name = "flask", marker = "extra == 'flask'", specifier = ">=2.2" }, - { name = "google-auth", marker = "extra == 'all'", specifier = ">=2.38.0" }, - { name = "google-auth", marker = "extra == 'google'", specifier = ">=2.38.0" }, - { name = "google-cloud-logging", marker = "extra == 'all'", specifier = ">=3.15.0" }, - { name = "google-cloud-logging", marker = "extra == 'google'", specifier = ">=3.15.0" }, - { name = "httpx", marker = "extra == 'all'" }, - { name = "httpx", marker = "extra == 'httpx'" }, - { name = "opentelemetry-api", marker = "extra == 'all'", specifier = ">=1.43.0" }, - { name = "opentelemetry-api", marker = "extra == 'otel'", specifier = ">=1.43.0" }, - { name = "opentelemetry-api", marker = "extra == 'otlp-grpc'", specifier = ">=1.43.0" }, - { name = "opentelemetry-api", marker = "extra == 'otlp-http'", specifier = ">=1.43.0" }, - { name = "opentelemetry-exporter-otlp-proto-grpc", marker = "extra == 'all'", specifier = ">=1.43.0" }, - { name = "opentelemetry-exporter-otlp-proto-grpc", marker = "extra == 'otlp-grpc'", specifier = ">=1.43.0" }, - { name = "opentelemetry-exporter-otlp-proto-http", marker = "extra == 'all'", specifier = ">=1.43.0" }, - { name = "opentelemetry-exporter-otlp-proto-http", marker = "extra == 'otlp-http'", specifier = ">=1.43.0" }, - { name = "opentelemetry-sdk", marker = "extra == 'all'", specifier = ">=1.43.0" }, - { name = "opentelemetry-sdk", marker = "extra == 'otel'", specifier = ">=1.43.0" }, - { name = "opentelemetry-sdk", marker = "extra == 'otlp-grpc'", specifier = ">=1.43.0" }, - { name = "opentelemetry-sdk", marker = "extra == 'otlp-http'", specifier = ">=1.43.0" }, - { name = "pyright", marker = "extra == 'dev'", specifier = ">=1.1.405" }, - { name = "pytest", marker = "extra == 'dev'" }, - { name = "ruff", marker = "extra == 'dev'", specifier = ">=0.9.0" }, - { name = "towncrier", marker = "extra == 'dev'", specifier = ">=24.8.0" }, -] -provides-extras = ["otel", "otlp-grpc", "otlp-http", "flask", "fastapi", "httpx", "google", "all", "dev"] - [[package]] name = "policyengine-uk" version = "2.90.2" From 362a1572dc1e8afbb1b693e5c2f84c7453058112 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 23 Sep 2026 19:32:38 +0400 Subject: [PATCH 07/10] Add durable observability IDs and stage registry --- policyengine_api/api.py | 4 +- policyengine_api/asgi_factory.py | 32 ++++- .../libs/simulation_entrypoint.py | 28 +++- policyengine_api/migration_logging.py | 22 +++- policyengine_api/observability/__init__.py | 5 + policyengine_api/observability/identifiers.py | 25 ++++ .../runtime.py} | 4 +- policyengine_api/observability/stages.py | 103 +++++++++++++++ policyengine_api/request_context.py | 35 +++++ .../runtime_cache/reform_impacts.py | 2 + .../services/budget_window_cache.py | 31 +++++ policyengine_api/services/economy_service.py | 123 +++++++++++++----- .../services/household_calculation_service.py | 113 +++++++++------- .../services/reform_impacts_service.py | 2 + tests/fixtures/libs/simulation_entrypoint.py | 8 +- tests/fixtures/services/economy_service.py | 20 +-- tests/unit/libs/test_simulation_entrypoint.py | 19 +-- .../routes/test_migration_context_logging.py | 26 +++- tests/unit/services/test_economy_service.py | 106 ++++++++------- tests/unit/test_asgi_factory.py | 21 +++ .../unit/test_observability_stage_registry.py | 31 +++++ 21 files changed, 596 insertions(+), 164 deletions(-) create mode 100644 policyengine_api/observability/__init__.py create mode 100644 policyengine_api/observability/identifiers.py rename policyengine_api/{observability.py => observability/runtime.py} (98%) create mode 100644 policyengine_api/observability/stages.py create mode 100644 tests/unit/test_observability_stage_registry.py diff --git a/policyengine_api/api.py b/policyengine_api/api.py index c4bba167d..2c0d23a35 100644 --- a/policyengine_api/api.py +++ b/policyengine_api/api.py @@ -23,6 +23,8 @@ def log_timing(message): from policyengine_api.extensions import cache from policyengine_api.migration_logging import register_migration_request_logging +from policyengine_api.observability.identifiers import OBSERVABILITY_ID_HEADER +from policyengine_api.request_context import REQUEST_ID_HEADER from policyengine_api.observability import runtime as observability_runtime from policyengine_api.runtime_cache.settings import load_runtime_cache_settings from policyengine_observability import instrument_flask @@ -103,7 +105,7 @@ def log_timing(message): cache.init_app(app) log_timing("Caching initialised") -CORS(app) +CORS(app, expose_headers=[REQUEST_ID_HEADER, OBSERVABILITY_ID_HEADER]) log_timing("CORS initialised") register_migration_request_logging(app, runtime=observability_runtime) diff --git a/policyengine_api/asgi_factory.py b/policyengine_api/asgi_factory.py index bcc9802ed..cc1977ac7 100644 --- a/policyengine_api/asgi_factory.py +++ b/policyengine_api/asgi_factory.py @@ -31,9 +31,13 @@ from policyengine_api.migration_logging import log_migration_request from policyengine_api.request_context import ( REQUEST_ID_HEADER, + _asgi_observability_id, _asgi_request_id, generate_request_id, + resolve_observability_id, ) +from policyengine_api.observability.identifiers import OBSERVABILITY_ID_HEADER +from policyengine_api.request_context import current_observability_id from starlette.datastructures import MutableHeaders from starlette.middleware.cors import CORSMiddleware from starlette.middleware.gzip import GZipMiddleware @@ -48,6 +52,13 @@ def _apply_request_id_header( response.headers[REQUEST_ID_HEADER] = request_id +def _apply_observability_id_header( + response: Response, + observability_id: str, +) -> None: + response.headers[OBSERVABILITY_ID_HEADER] = observability_id + + def create_asgi_app( wsgi_app, *, @@ -97,6 +108,12 @@ async def add_headers_to_unhandled_errors( request.headers.get(REQUEST_ID_HEADER) or generate_request_id(), ) _apply_request_id_header(response, request_id) + observability_id = getattr( + request.state, + "policyengine_observability_id", + None, + ) or resolve_observability_id(request.headers.get(OBSERVABILITY_ID_HEADER)) + _apply_observability_id_header(response, observability_id) return response @app.exception_handler(RequestValidationError) @@ -119,9 +136,15 @@ async def oversized_v2_request( async def add_request_context_and_migration_logging(request, call_next): started_at = time.time() request_id = request.headers.get(REQUEST_ID_HEADER) or generate_request_id() + observability_id = resolve_observability_id( + request.headers.get(OBSERVABILITY_ID_HEADER) + ) MutableHeaders(scope=request.scope)[REQUEST_ID_HEADER] = request_id + MutableHeaders(scope=request.scope)[OBSERVABILITY_ID_HEADER] = observability_id request.state.policyengine_request_id = request_id + request.state.policyengine_observability_id = observability_id context_token = _asgi_request_id.set(request_id) + observability_context_token = _asgi_observability_id.set(observability_id) def log_native_route(status_code: int) -> None: if not isinstance(request.scope.get("route"), APIRoute): @@ -149,10 +172,17 @@ def log_native_route(status_code: int) -> None: log_native_route(500) raise _apply_request_id_header(response, request_id) + response_observability_id = ( + response.headers.get(OBSERVABILITY_ID_HEADER) + or current_observability_id() + or observability_id + ) + _apply_observability_id_header(response, response_observability_id) log_native_route(response.status_code) return response finally: _asgi_request_id.reset(context_token) + _asgi_observability_id.reset(observability_context_token) app.include_router(build_core_health_router(dependencies)) app.include_router(build_v2_router(dependencies)) @@ -171,7 +201,7 @@ def log_native_route(status_code: int) -> None: allow_origin_regex=".*", allow_methods=["DELETE", "GET", "HEAD", "OPTIONS", "PATCH", "POST", "PUT"], allow_headers=["*"], - expose_headers=[REQUEST_ID_HEADER], + expose_headers=[REQUEST_ID_HEADER, OBSERVABILITY_ID_HEADER], allow_credentials=False, max_age=600, ) diff --git a/policyengine_api/libs/simulation_entrypoint.py b/policyengine_api/libs/simulation_entrypoint.py index 05b88a6ba..1d9d2678b 100644 --- a/policyengine_api/libs/simulation_entrypoint.py +++ b/policyengine_api/libs/simulation_entrypoint.py @@ -20,8 +20,10 @@ from policyengine_api.observability import get_runtime from policyengine_api.request_context import ( REQUEST_ID_HEADER, + current_observability_id, current_request_id, ) +from policyengine_api.observability.identifiers import OBSERVABILITY_ID_HEADER from policyengine_api.worker_spm import validate_worker_spm, raise_worker_spm_error @@ -62,6 +64,9 @@ def _attach_current_request_id(request: httpx.Request) -> None: request_id = current_request_id() if request_id is not None: request.headers[REQUEST_ID_HEADER] = request_id + observability_id = current_observability_id() + if observability_id is not None: + request.headers[OBSERVABILITY_ID_HEADER] = observability_id @dataclass @@ -72,7 +77,7 @@ class ModalSimulationExecution: job_id: str status: str - run_id: Optional[str] = None + observability_id: Optional[str] = None result: Optional[dict] = None error: Optional[str] = None policyengine_bundle: Optional[dict] = None @@ -92,6 +97,7 @@ class ModalBudgetWindowBatchExecution: batch_job_id: str status: str + observability_id: Optional[str] = None progress: Optional[int] = None completed_years: list[str] = field(default_factory=list) running_years: list[str] = field(default_factory=list) @@ -221,7 +227,7 @@ def run(self, payload: dict) -> ModalSimulationExecution: { "message": "Simulation entrypoint job submitted", "job_id": data.get("job_id"), - "run_id": data.get("run_id"), + "observability_id": data.get("observability_id"), "status": data.get("status"), }, severity="INFO", @@ -232,14 +238,16 @@ def run(self, payload: dict) -> ModalSimulationExecution: status=data["status"], policyengine_bundle=data.get("policyengine_bundle"), resolved_app_name=data.get("resolved_app_name"), - run_id=data.get("run_id"), + observability_id=data.get("observability_id"), ) except httpx.HTTPStatusError as e: logger.log_struct( { "message": f"Simulation entrypoint HTTP error: {e.response.status_code}", - "run_id": (payload.get("_telemetry") or {}).get("run_id"), + "observability_id": (payload.get("_telemetry") or {}).get( + "observability_id" + ), "response_text": e.response.text[:500], }, severity="ERROR", @@ -250,7 +258,9 @@ def run(self, payload: dict) -> ModalSimulationExecution: logger.log_struct( { "message": f"Simulation entrypoint request error: {str(e)}", - "run_id": (payload.get("_telemetry") or {}).get("run_id"), + "observability_id": (payload.get("_telemetry") or {}).get( + "observability_id" + ), }, severity="ERROR", ) @@ -283,6 +293,7 @@ def run_budget_window_batch(self, payload: dict) -> ModalBudgetWindowBatchExecut return ModalBudgetWindowBatchExecution( batch_job_id=data["batch_job_id"], status=data["status"], + observability_id=data.get("observability_id"), ) except httpx.HTTPStatusError as e: @@ -299,7 +310,9 @@ def run_budget_window_batch(self, payload: dict) -> ModalBudgetWindowBatchExecut logger.log_struct( { "message": f"Simulation batch API request error: {str(e)}", - "run_id": (payload.get("_telemetry") or {}).get("run_id"), + "observability_id": (payload.get("_telemetry") or {}).get( + "observability_id" + ), }, severity="ERROR", ) @@ -434,7 +447,7 @@ def get_execution_by_id(self, job_id: str) -> ModalSimulationExecution: return ModalSimulationExecution( job_id=job_id, status=data["status"], - run_id=data.get("run_id"), + observability_id=data.get("observability_id"), result=data.get("result"), error=data.get("error"), policyengine_bundle=data.get("policyengine_bundle"), @@ -478,6 +491,7 @@ def get_budget_window_batch_by_id( return ModalBudgetWindowBatchExecution( batch_job_id=batch_job_id, status=data["status"], + observability_id=data.get("observability_id"), progress=data.get("progress"), completed_years=data.get("completed_years", []), running_years=data.get("running_years", []), diff --git a/policyengine_api/migration_logging.py b/policyengine_api/migration_logging.py index 8f9f6cbe0..f278367ba 100644 --- a/policyengine_api/migration_logging.py +++ b/policyengine_api/migration_logging.py @@ -15,7 +15,9 @@ from policyengine_api.request_context import ( REQUEST_ID_HEADER, generate_request_id, + resolve_observability_id, ) +from policyengine_api.observability.identifiers import OBSERVABILITY_ID_HEADER V2_METADATA_RESOURCE_SEGMENTS = frozenset( @@ -77,16 +79,34 @@ def register_migration_request_logging( @app.before_request def set_request_migration_context(): flask.g.request_started_at = time.time() - captured = runtime.capture_context() if runtime is not None else {} + try: + captured = runtime.capture_context() if runtime is not None else {} + except Exception: + captured = {} flask.g.request_id = captured.get("request_id") or ( flask.request.headers.get(REQUEST_ID_HEADER) or generate_request_id() ) + flask.g.observability_id = resolve_observability_id( + captured.get("observability_id") + or flask.request.headers.get(OBSERVABILITY_ID_HEADER) + ) + if runtime is not None: + try: + runtime.set_context( + request_id=flask.g.request_id, + observability_id=flask.g.observability_id, + ) + except Exception: + pass @app.after_request def log_request_migration_context(response): request_id = getattr(flask.g, "request_id", None) if request_id is not None: response.headers[REQUEST_ID_HEADER] = request_id + observability_id = getattr(flask.g, "observability_id", None) + if observability_id is not None: + response.headers[OBSERVABILITY_ID_HEADER] = observability_id try: country_id = ( flask.request.view_args.get("country_id") diff --git a/policyengine_api/observability/__init__.py b/policyengine_api/observability/__init__.py new file mode 100644 index 000000000..df2252a8a --- /dev/null +++ b/policyengine_api/observability/__init__.py @@ -0,0 +1,5 @@ +"""API observability runtime, identifiers, and registered stage plans.""" + +from .runtime import _build_runtime, get_runtime, runtime + +__all__ = ["_build_runtime", "get_runtime", "runtime"] diff --git a/policyengine_api/observability/identifiers.py b/policyengine_api/observability/identifiers.py new file mode 100644 index 000000000..e836fc887 --- /dev/null +++ b/policyengine_api/observability/identifiers.py @@ -0,0 +1,25 @@ +"""Diagnostic correlation identifiers for API requests and report work.""" + +from __future__ import annotations + +from typing import Any +from uuid import UUID, uuid4 + +OBSERVABILITY_ID_HEADER = "X-PolicyEngine-Observability-Id" + + +def generate_observability_id() -> str: + """Create an identifier used only to correlate observability records.""" + + return str(uuid4()) + + +def normalize_observability_id(value: Any) -> str | None: + """Return a canonical UUID string, or ``None`` for malformed input.""" + + if not isinstance(value, str): + return None + try: + return str(UUID(value)) + except (ValueError, AttributeError): + return None diff --git a/policyengine_api/observability.py b/policyengine_api/observability/runtime.py similarity index 98% rename from policyengine_api/observability.py rename to policyengine_api/observability/runtime.py index 25fbba0b1..a3c2da72a 100644 --- a/policyengine_api/observability.py +++ b/policyengine_api/observability/runtime.py @@ -50,14 +50,14 @@ "path", "policy_id", "policyengine_version", - "process_id", + "submission_claim_id", "requested_through_revision", "resource", "request_id", + "observability_id", "resolved_app_name", "route_group", "route_impl", - "run_id", "simulation_year", "sim_compute", "sim_entrypoint", diff --git a/policyengine_api/observability/stages.py b/policyengine_api/observability/stages.py new file mode 100644 index 000000000..eaba1b7d9 --- /dev/null +++ b/policyengine_api/observability/stages.py @@ -0,0 +1,103 @@ +"""Canonical stage registry for API calculation configurations.""" + +from __future__ import annotations + +from dataclasses import dataclass +from enum import StrEnum +from types import MappingProxyType +from typing import Mapping + + +class RunConfiguration(StrEnum): + HOUSEHOLD = "household" + ECONOMY_ANNUAL = "economy_annual" + ECONOMY_BUDGET_WINDOW = "economy_budget_window" + + +class Stage(StrEnum): + HOUSEHOLD_LOAD_INPUTS = "household.load_inputs" + HOUSEHOLD_CACHE_LOOKUP = "household.cache_lookup" + HOUSEHOLD_INPUT_NORMALIZATION = "household.input_normalization" + HOUSEHOLD_CALCULATION = "household.calculation" + HOUSEHOLD_CACHE_WRITE = "household.cache_write" + + ECONOMY_REQUEST = "economy.request" + ECONOMY_BUDGET_WINDOW_REQUEST = "economy.budget_window_request" + ECONOMY_LOAD_POLICIES = "economy.load_policies" + ECONOMY_RESOLVE_CACHED_OR_NEW = "economy.resolve_cached_or_new_impact" + ECONOMY_RESOLVE_RUNTIME_BUNDLE = "economy.resolve_runtime_bundle" + ECONOMY_SUBMIT = "economy.submit_impact" + ECONOMY_START_BUDGET_WINDOW = "economy.start_budget_window_batch" + ECONOMY_POLL_BUDGET_WINDOW = "economy.poll_budget_window_batch" + ECONOMY_HANDLE_EXECUTION_STATE = "economy.handle_execution_state" + ECONOMY_READ_COMPLETED = "economy.read_completed_impact" + ECONOMY_READ_FAILED = "economy.read_failed_impact" + ECONOMY_POLL_ACTIVE = "economy.poll_active_impact" + ECONOMY_PERSIST_COMPUTING = "economy.persist_computing_impact" + ECONOMY_PERSIST_COMPLETED = "economy.persist_completed_impact" + ECONOMY_PERSIST_FAILED = "economy.persist_failed_impact" + + +@dataclass(frozen=True) +class StagePlan: + configuration: RunConfiguration + stages: tuple[Stage, ...] + + def name(self, stage: Stage) -> str: + if stage not in self.stages: + raise ValueError( + f"{stage.value!r} is not registered for {self.configuration.value!r}" + ) + return stage.value + + +_ECONOMY_COMMON = ( + Stage.ECONOMY_LOAD_POLICIES, + Stage.ECONOMY_RESOLVE_CACHED_OR_NEW, + Stage.ECONOMY_RESOLVE_RUNTIME_BUNDLE, + Stage.ECONOMY_HANDLE_EXECUTION_STATE, + Stage.ECONOMY_READ_COMPLETED, + Stage.ECONOMY_READ_FAILED, + Stage.ECONOMY_POLL_ACTIVE, + Stage.ECONOMY_PERSIST_COMPUTING, + Stage.ECONOMY_PERSIST_COMPLETED, + Stage.ECONOMY_PERSIST_FAILED, +) + +RUN_STAGE_REGISTRY: Mapping[RunConfiguration, StagePlan] = MappingProxyType( + { + RunConfiguration.HOUSEHOLD: StagePlan( + RunConfiguration.HOUSEHOLD, + ( + Stage.HOUSEHOLD_LOAD_INPUTS, + Stage.HOUSEHOLD_CACHE_LOOKUP, + Stage.HOUSEHOLD_INPUT_NORMALIZATION, + Stage.HOUSEHOLD_CALCULATION, + Stage.HOUSEHOLD_CACHE_WRITE, + ), + ), + RunConfiguration.ECONOMY_ANNUAL: StagePlan( + RunConfiguration.ECONOMY_ANNUAL, + ( + Stage.ECONOMY_REQUEST, + *_ECONOMY_COMMON, + Stage.ECONOMY_SUBMIT, + ), + ), + RunConfiguration.ECONOMY_BUDGET_WINDOW: StagePlan( + RunConfiguration.ECONOMY_BUDGET_WINDOW, + ( + Stage.ECONOMY_BUDGET_WINDOW_REQUEST, + *_ECONOMY_COMMON, + Stage.ECONOMY_START_BUDGET_WINDOW, + Stage.ECONOMY_POLL_BUDGET_WINDOW, + ), + ), + } +) + +HOUSEHOLD_STAGES = RUN_STAGE_REGISTRY[RunConfiguration.HOUSEHOLD] +ECONOMY_ANNUAL_STAGES = RUN_STAGE_REGISTRY[RunConfiguration.ECONOMY_ANNUAL] +ECONOMY_BUDGET_WINDOW_STAGES = RUN_STAGE_REGISTRY[ + RunConfiguration.ECONOMY_BUDGET_WINDOW +] diff --git a/policyengine_api/request_context.py b/policyengine_api/request_context.py index feff6194f..4a279faed 100644 --- a/policyengine_api/request_context.py +++ b/policyengine_api/request_context.py @@ -6,6 +6,10 @@ from contextvars import ContextVar import flask +from policyengine_api.observability.identifiers import ( + generate_observability_id, + normalize_observability_id, +) REQUEST_ID_HEADER = "X-PolicyEngine-Request-Id" @@ -13,6 +17,10 @@ "policyengine_api_request_id", default=None, ) +_asgi_observability_id: ContextVar[str | None] = ContextVar( + "policyengine_api_observability_id", + default=None, +) def generate_request_id() -> str: @@ -27,3 +35,30 @@ def current_request_id() -> str | None: if flask.has_request_context(): return getattr(flask.g, "request_id", None) return _asgi_request_id.get() + + +def current_observability_id() -> str | None: + """Return the diagnostic correlation identifier for the current request.""" + + if flask.has_request_context(): + return getattr(flask.g, "observability_id", None) + return _asgi_observability_id.get() + + +def adopt_observability_id(value: object) -> str | None: + """Adopt a stored identifier without allowing malformed data to fail work.""" + + observability_id = normalize_observability_id(value) + if observability_id is None: + return current_observability_id() + if flask.has_request_context(): + flask.g.observability_id = observability_id + else: + _asgi_observability_id.set(observability_id) + return observability_id + + +def resolve_observability_id(value: object) -> str: + """Use a valid caller value or create a new diagnostic identifier.""" + + return normalize_observability_id(value) or generate_observability_id() diff --git a/policyengine_api/runtime_cache/reform_impacts.py b/policyengine_api/runtime_cache/reform_impacts.py index 920b4f422..578fa3add 100644 --- a/policyengine_api/runtime_cache/reform_impacts.py +++ b/policyengine_api/runtime_cache/reform_impacts.py @@ -43,6 +43,7 @@ class CachedReformImpact: end_time: datetime | None execution_id: str | None error_code: str | None = None + observability_id: str | None = None def _datetime_to_wire(value: datetime | None) -> str | None: @@ -97,6 +98,7 @@ def _impact_from_wire(payload: Any) -> CachedReformImpact | None: else None ), error_code=payload.get("error_code"), + observability_id=payload.get("observability_id"), ) except (KeyError, TypeError, ValueError): return None diff --git a/policyengine_api/services/budget_window_cache.py b/policyengine_api/services/budget_window_cache.py index ae2e0cf6b..ef5261051 100644 --- a/policyengine_api/services/budget_window_cache.py +++ b/policyengine_api/services/budget_window_cache.py @@ -82,6 +82,10 @@ def _error_key(cache_key: str) -> str: def _batch_key(cache_key: str) -> str: return f"{cache_key}:batch-job-id" + @staticmethod + def _observability_key(cache_key: str) -> str: + return f"{cache_key}:observability-id" + @staticmethod def _handle_cache_error( operation: str, @@ -218,6 +222,15 @@ def get_batch_job_id(self, cache_key: str) -> str | None: ) return value + def get_observability_id(self, cache_key: str) -> str | None: + """Read optional diagnostic state without affecting calculation flow.""" + + try: + value = self.client.get(self._observability_key(cache_key)) + except Exception: + return None + return value if isinstance(value, str) and value else None + def claim_batch_start(self, cache_key: str, claim_token: str) -> bool: try: return self._claims.acquire( @@ -261,6 +274,24 @@ def store_batch_job_id(self, cache_key: str, batch_job_id: str) -> None: started_at=started_at, ) + def store_observability_id( + self, + cache_key: str, + observability_id: str | None, + ) -> None: + """Store diagnostic state best effort; exporter state cannot block work.""" + + if not observability_id: + return + try: + self.client.set( + self._observability_key(cache_key), + observability_id, + ex=BUDGET_WINDOW_BATCH_TTL_SECONDS, + ) + except Exception: + return + def clear_starting_claim(self, cache_key: str, claim_token: str) -> None: try: self._claims.release( diff --git a/policyengine_api/services/economy_service.py b/policyengine_api/services/economy_service.py index 85d062a6f..7282ecffa 100644 --- a/policyengine_api/services/economy_service.py +++ b/policyengine_api/services/economy_service.py @@ -6,7 +6,6 @@ from typing import Any, Literal, Optional import httpx -import numpy as np from dotenv import load_dotenv from policyengine_api.constants import ( COUNTRY_PACKAGE_VERSIONS, @@ -26,6 +25,15 @@ from policyengine_api.gcp_logging import logger from policyengine_api.libs.simulation_entrypoint import simulation_entrypoint from policyengine_api.observability import runtime as observability_runtime +from policyengine_api.observability.stages import ( + ECONOMY_ANNUAL_STAGES, + ECONOMY_BUDGET_WINDOW_STAGES, + Stage, +) +from policyengine_api.request_context import ( + adopt_observability_id, + current_observability_id, +) from policyengine_api.services.budget_window_cache import BudgetWindowCache from policyengine_api.services.policy_service import PolicyService from policyengine_api.services.reform_impacts_service import ( @@ -85,7 +93,8 @@ class SimulationOptions(BaseModel): class EconomicImpactSetupOptions(BaseModel): - process_id: str + submission_claim_id: str + observability_id: str country_id: str reform_policy_id: int baseline_policy_id: int @@ -274,7 +283,7 @@ def _budget_window_cache(self) -> BudgetWindowCache: def _simulation_gateway(self): return self._injected_simulation_entrypoint or simulation_entrypoint - @observability_runtime.span("economy.load_policies") + @observability_runtime.span(ECONOMY_ANNUAL_STAGES.name(Stage.ECONOMY_LOAD_POLICIES)) def _get_policy_jsons( self, country_id: str, @@ -303,7 +312,7 @@ def _parse_json_object(value: dict[str, Any] | str) -> dict[str, Any]: raise TypeError("Expected a JSON object") return parsed - @observability_runtime.span("economy.request") + @observability_runtime.span(ECONOMY_ANNUAL_STAGES.name(Stage.ECONOMY_REQUEST)) def get_economic_impact( self, country_id: str, @@ -357,7 +366,9 @@ def get_economic_impact( print(f"Error getting economic impact: {str(e)}") raise e - @observability_runtime.span("economy.budget_window_request") + @observability_runtime.span( + ECONOMY_BUDGET_WINDOW_STAGES.name(Stage.ECONOMY_BUDGET_WINDOW_REQUEST) + ) def get_budget_window_economic_impact( self, country_id: str, @@ -433,6 +444,14 @@ def get_budget_window_economic_impact( batch_job_id = self._budget_window_cache.get_batch_job_id(cache_key) if batch_job_id: + stored_observability_id = adopt_observability_id( + self._budget_window_cache.get_observability_id(cache_key) + ) + if stored_observability_id is not None: + setup_options.observability_id = stored_observability_id + observability_runtime.set_context( + observability_id=stored_observability_id + ) return self._get_budget_window_result_from_batch_job_id( batch_job_id=batch_job_id, spm=setup_options.options.get("spm"), @@ -442,7 +461,7 @@ def get_budget_window_economic_impact( cache_status="batch-id-hit", ) - claim_token = setup_options.process_id + claim_token = setup_options.submission_claim_id cache_status = "starting-claim-hit" if self._budget_window_cache.claim_batch_start(cache_key, claim_token): cache_status = "miss" @@ -456,6 +475,11 @@ def get_budget_window_economic_impact( self._budget_window_cache.store_batch_job_id( cache_key, batch_execution.batch_job_id ) + self._budget_window_cache.store_observability_id( + cache_key, + batch_execution.observability_id + or setup_options.observability_id, + ) except httpx.HTTPStatusError as error: self._budget_window_cache.clear_starting_claim( cache_key, claim_token @@ -536,7 +560,9 @@ def _build_budget_window_batch_payload( sim_params["target"] = setup_options.target return sim_params - @observability_runtime.span("economy.start_budget_window_batch") + @observability_runtime.span( + ECONOMY_BUDGET_WINDOW_STAGES.name(Stage.ECONOMY_START_BUDGET_WINDOW) + ) def _start_budget_window_batch( self, *, @@ -585,7 +611,9 @@ def _build_budget_window_submission_error_message( return str(error) - @observability_runtime.span("economy.poll_budget_window_batch") + @observability_runtime.span( + ECONOMY_BUDGET_WINDOW_STAGES.name(Stage.ECONOMY_POLL_BUDGET_WINDOW) + ) def _get_budget_window_result_from_batch_job_id( self, *, @@ -711,7 +739,8 @@ def _build_economic_impact_setup_options( ) if resolved_spm is not None: options = {**options, "spm": resolved_spm} - process_id: str = self._create_process_id() + submission_claim_id = self._create_submission_claim_id() + observability_id = current_observability_id() or str(uuid.uuid4()) cache_version = get_economy_impact_cache_version(country_id, api_version) country_package_version = COUNTRY_PACKAGE_VERSIONS.get(country_id) resolved_dataset = "default" @@ -729,7 +758,8 @@ def _build_economic_impact_setup_options( return EconomicImpactSetupOptions.model_validate( { - "process_id": process_id, + "submission_claim_id": submission_claim_id, + "observability_id": observability_id, "country_id": country_id, "reform_policy_id": policy_id, "baseline_policy_id": baseline_policy_id, @@ -747,7 +777,9 @@ def _build_economic_impact_setup_options( } ) - @observability_runtime.span("economy.resolve_cached_or_new_impact") + @observability_runtime.span( + ECONOMY_ANNUAL_STAGES.name(Stage.ECONOMY_RESOLVE_CACHED_OR_NEW) + ) def _get_or_create_economic_impact( self, setup_options: EconomicImpactSetupOptions ) -> EconomicImpactResult: @@ -762,6 +794,15 @@ def _get_or_create_economic_impact( most_recent_impact: dict | None = self._get_most_recent_impact( setup_options=setup_options ) + if most_recent_impact is not None: + stored_observability_id = adopt_observability_id( + getattr(most_recent_impact, "observability_id", None) + ) + if stored_observability_id is not None: + setup_options.observability_id = stored_observability_id + observability_runtime.set_context( + observability_id=stored_observability_id + ) if most_recent_impact and self._should_refresh_cached_impact( setup_options=setup_options, @@ -846,7 +887,9 @@ def _get_or_create_economic_impact( raise ValueError(f"Unexpected impact action: {impact_action}") - @observability_runtime.span("economy.resolve_runtime_bundle") + @observability_runtime.span( + ECONOMY_ANNUAL_STAGES.name(Stage.ECONOMY_RESOLVE_RUNTIME_BUNDLE) + ) def _resolve_runtime_bundle_for_setup_options( self, setup_options: EconomicImpactSetupOptions, @@ -886,7 +929,7 @@ def _reform_impact_start_claim_arguments( "options_hash": setup_options.options_hash, "api_version": setup_options.api_version, "target": setup_options.target, - "claim_token": setup_options.process_id, + "claim_token": setup_options.submission_claim_id, } def _claim_reform_impact_start( @@ -1001,7 +1044,9 @@ def _determine_impact_action( else: raise ValueError(f"Unknown impact status: {status}") - @observability_runtime.span("economy.handle_execution_state") + @observability_runtime.span( + ECONOMY_ANNUAL_STAGES.name(Stage.ECONOMY_HANDLE_EXECUTION_STATE) + ) def _handle_execution_state( self, setup_options: EconomicImpactSetupOptions, @@ -1070,7 +1115,9 @@ def _handle_execution_state( else: raise ValueError(f"Unexpected sim API execution state: {execution_state}") - @observability_runtime.span("economy.read_completed_impact") + @observability_runtime.span( + ECONOMY_ANNUAL_STAGES.name(Stage.ECONOMY_READ_COMPLETED) + ) def _handle_completed_impact( self, setup_options: EconomicImpactSetupOptions, @@ -1130,7 +1177,7 @@ def _record_uncertifiable_stored_impact( except Exception: pass - @observability_runtime.span("economy.read_failed_impact") + @observability_runtime.span(ECONOMY_ANNUAL_STAGES.name(Stage.ECONOMY_READ_FAILED)) def _handle_failed_impact( self, most_recent_impact: ReformImpact, @@ -1151,7 +1198,7 @@ def _handle_failed_impact( ) ) - @observability_runtime.span("economy.poll_active_impact") + @observability_runtime.span(ECONOMY_ANNUAL_STAGES.name(Stage.ECONOMY_POLL_ACTIVE)) def _handle_computing_impact( self, setup_options: EconomicImpactSetupOptions, @@ -1177,7 +1224,7 @@ def _handle_computing_impact( ) raise - @observability_runtime.span("economy.submit_impact") + @observability_runtime.span(ECONOMY_ANNUAL_STAGES.name(Stage.ECONOMY_SUBMIT)) def _handle_create_impact( self, setup_options: EconomicImpactSetupOptions, @@ -1210,7 +1257,7 @@ def _handle_create_impact( logger.log_struct( { "message": "Setting up sim API job", - "run_id": telemetry["run_id"], + "observability_id": telemetry["observability_id"], **setup_options.model_dump(), } ) @@ -1219,7 +1266,7 @@ def _handle_create_impact( sim_params["_metadata"] = { "reform_policy_id": setup_options.reform_policy_id, "baseline_policy_id": setup_options.baseline_policy_id, - "process_id": setup_options.process_id, + "submission_claim_id": setup_options.submission_claim_id, "model_version": setup_options.model_version, "policyengine_version": setup_options.policyengine_version, "data_version": setup_options.data_version, @@ -1243,15 +1290,16 @@ def _handle_create_impact( entrypoint_execution ) - run_id = ( - getattr(entrypoint_execution, "run_id", None) or telemetry["run_id"] + observability_id = ( + getattr(entrypoint_execution, "observability_id", None) + or telemetry["observability_id"] ) progress_log = { **setup_options.model_dump(), "message": "Sim API job started", "execution_id": execution_id, - "run_id": run_id, + "observability_id": observability_id, } logger.log_struct(progress_log, severity="INFO") @@ -1472,8 +1520,8 @@ def _build_simulation_telemetry( ) return { - "run_id": str(uuid.uuid4()), - "process_id": setup_options.process_id, + "observability_id": setup_options.observability_id, + "submission_claim_id": setup_options.submission_claim_id, "requested_at": datetime.datetime.now(datetime.UTC).isoformat(), "simulation_kind": simulation_kind, "geography_code": geography_code, @@ -1512,7 +1560,9 @@ def _stable_config_hash(self, payload: dict[str, Any]) -> str: # are written separately because the service relies upon mutating an original # 'computing' record to 'ok' or 'error' status, rather than creating a new record. # This should be addressed in the future. - @observability_runtime.span("economy.persist_computing_impact") + @observability_runtime.span( + ECONOMY_ANNUAL_STAGES.name(Stage.ECONOMY_PERSIST_COMPUTING) + ) def _set_reform_impact_computing( self, setup_options: EconomicImpactSetupOptions, @@ -1536,6 +1586,7 @@ def _set_reform_impact_computing( reform_impact_json={}, start_time=datetime.datetime.now(), execution_id=execution_id, + observability_id=setup_options.observability_id, ) except Exception as e: logger.log_struct( @@ -1546,7 +1597,9 @@ def _set_reform_impact_computing( ) raise e - @observability_runtime.span("economy.persist_completed_impact") + @observability_runtime.span( + ECONOMY_ANNUAL_STAGES.name(Stage.ECONOMY_PERSIST_COMPLETED) + ) def _set_reform_impact_complete( self, setup_options: EconomicImpactSetupOptions, @@ -1577,7 +1630,9 @@ def _set_reform_impact_complete( ) raise e - @observability_runtime.span("economy.persist_failed_impact") + @observability_runtime.span( + ECONOMY_ANNUAL_STAGES.name(Stage.ECONOMY_PERSIST_FAILED) + ) def _set_reform_impact_error( self, setup_options: EconomicImpactSetupOptions, @@ -1610,11 +1665,7 @@ def _set_reform_impact_error( ) raise e - def _create_process_id(self) -> str: - """ - Generate a unique process ID based on the current timestamp and a random number. - This is used to track the process in the database and logs. - """ - timestamp = datetime.datetime.now().strftime("%Y%m%d%H%M%S") - random_number = np.random.randint(1000, 9999) - return f"job_{timestamp}_{random_number}" + def _create_submission_claim_id(self) -> str: + """Create an opaque token for one submission ownership claim.""" + + return str(uuid.uuid4()) diff --git a/policyengine_api/services/household_calculation_service.py b/policyengine_api/services/household_calculation_service.py index 8eb26e11f..9f324a0b6 100644 --- a/policyengine_api/services/household_calculation_service.py +++ b/policyengine_api/services/household_calculation_service.py @@ -10,6 +10,8 @@ from sqlalchemy.orm import Session, sessionmaker from policyengine_api.constants import COUNTRY_PACKAGE_VERSIONS, POLICYENGINE_VERSION +from policyengine_api.observability import runtime as observability_runtime +from policyengine_api.observability.stages import HOUSEHOLD_STAGES, Stage from policyengine_api.data.orm import get_v1_session_factory from policyengine_api.data.v1_models import ( Household, @@ -144,6 +146,7 @@ def _cache_identity( spm=spm, ) + @observability_runtime.span(HOUSEHOLD_STAGES.name(Stage.HOUSEHOLD_LOAD_INPUTS)) def _get_inputs( self, country_id: str, @@ -165,6 +168,7 @@ def _get_inputs( ) return household, policy + @observability_runtime.span(HOUSEHOLD_STAGES.name(Stage.HOUSEHOLD_CACHE_WRITE)) def _store_result( self, identity: HouseholdCalculationIdentity, @@ -205,7 +209,10 @@ def calculate_stored_household( api_version, spm, ) - cached = self._cache.get(cache_identity) + with observability_runtime.span( + HOUSEHOLD_STAGES.name(Stage.HOUSEHOLD_CACHE_LOOKUP) + ): + cached = self._cache.get(cache_identity) if cached is not None: return HouseholdCalculationResult( household=cached.household, @@ -215,34 +222,40 @@ def calculate_stored_household( spm_provenance=cached.spm_provenance, ) - countries = self._countries() - country = countries.get(country_id) - household_json = add_yearly_variables( - household_inputs, - country_id, - countries, - ) - deprecated_inputs = drop_deprecated_inputs(household_json) - household_json = deprecated_inputs.household - invalid_inputs = find_unrecognized_inputs( - household_json, - policy.policy_json, - country.metadata, - ) + with observability_runtime.span( + HOUSEHOLD_STAGES.name(Stage.HOUSEHOLD_INPUT_NORMALIZATION) + ): + countries = self._countries() + country = countries.get(country_id) + household_json = add_yearly_variables( + household_inputs, + country_id, + countries, + ) + deprecated_inputs = drop_deprecated_inputs(household_json) + household_json = deprecated_inputs.household + invalid_inputs = find_unrecognized_inputs( + household_json, + policy.policy_json, + country.metadata, + ) if invalid_inputs: raise InvalidHouseholdInputsError(invalid_inputs) calculation_started_at = time.perf_counter() try: - raw_calculation = country.calculate( - household_json, - policy.policy_json, - **( - {"spm": spm, "spm_requested": saved_spm is not None} - if spm is not None - else {} - ), - ) + with observability_runtime.span( + HOUSEHOLD_STAGES.name(Stage.HOUSEHOLD_CALCULATION) + ): + raw_calculation = country.calculate( + household_json, + policy.policy_json, + **( + {"spm": spm, "spm_requested": saved_spm is not None} + if spm is not None + else {} + ), + ) except Exception: record_cache_event( family="household-calculation", @@ -298,32 +311,42 @@ def calculate_household( spm_requested: bool = False, ) -> HouseholdCalculationResult: """Validate and calculate request-provided household and policy data.""" - countries = self._countries() - country = countries.get(country_id) - spm = normalize_spm_selection(country_id, spm) - household_json = deepcopy(household_json) - if add_missing: - household_json = add_yearly_variables( + with observability_runtime.span( + HOUSEHOLD_STAGES.name(Stage.HOUSEHOLD_INPUT_NORMALIZATION) + ): + countries = self._countries() + country = countries.get(country_id) + spm = normalize_spm_selection(country_id, spm) + household_json = deepcopy(household_json) + if add_missing: + household_json = add_yearly_variables( + household_json, + country_id, + countries, + ) + + deprecated_inputs = drop_deprecated_inputs(household_json) + household_json = deprecated_inputs.household + invalid_inputs = find_unrecognized_inputs( household_json, - country_id, - countries, + policy_json, + country.metadata, ) - - deprecated_inputs = drop_deprecated_inputs(household_json) - household_json = deprecated_inputs.household - invalid_inputs = find_unrecognized_inputs( - household_json, - policy_json, - country.metadata, - ) if invalid_inputs: raise InvalidHouseholdInputsError(invalid_inputs) - raw_calculation = country.calculate( - household_json, - policy_json, - **({"spm": spm, "spm_requested": spm_requested} if spm is not None else {}), - ) + with observability_runtime.span( + HOUSEHOLD_STAGES.name(Stage.HOUSEHOLD_CALCULATION) + ): + raw_calculation = country.calculate( + household_json, + policy_json, + **( + {"spm": spm, "spm_requested": spm_requested} + if spm is not None + else {} + ), + ) if isinstance(raw_calculation, dict): household = raw_calculation calculation_warnings = () diff --git a/policyengine_api/services/reform_impacts_service.py b/policyengine_api/services/reform_impacts_service.py index 81e66cad2..fb1962bea 100644 --- a/policyengine_api/services/reform_impacts_service.py +++ b/policyengine_api/services/reform_impacts_service.py @@ -153,6 +153,7 @@ def set_reform_impact( reform_impact_json: dict[str, Any], start_time, execution_id: str, + observability_id: str | None = None, ) -> CachedReformImpact: impact = CachedReformImpact( reform_impact_id=reform_impact_id(execution_id), @@ -171,6 +172,7 @@ def set_reform_impact( start_time=start_time, end_time=None, execution_id=execution_id, + observability_id=observability_id, ) if not self._cache.set(impact): raise ReformImpactHandoffError( diff --git a/tests/fixtures/libs/simulation_entrypoint.py b/tests/fixtures/libs/simulation_entrypoint.py index 9a11b4c90..5b8f12c92 100644 --- a/tests/fixtures/libs/simulation_entrypoint.py +++ b/tests/fixtures/libs/simulation_entrypoint.py @@ -18,7 +18,7 @@ # Mock data constants MOCK_MODAL_JOB_ID = "fc-abc123xyz" -MOCK_RUN_ID = "run-abc123xyz" +MOCK_OBSERVABILITY_ID = "run-abc123xyz" MOCK_BATCH_JOB_ID = "fc-batch123xyz" MOCK_MODAL_BASE_URL = "https://test-modal-api.modal.run" @@ -35,8 +35,8 @@ MOCK_SIMULATION_PAYLOAD_WITH_TELEMETRY = { **MOCK_SIMULATION_PAYLOAD, "_telemetry": { - "run_id": MOCK_RUN_ID, - "process_id": "job_20250626120000_1234", + "observability_id": MOCK_OBSERVABILITY_ID, + "submission_claim_id": "job_20250626120000_1234", "capture_mode": "disabled", }, } @@ -59,7 +59,7 @@ MOCK_SUBMIT_RESPONSE_SUCCESS = { "job_id": MOCK_MODAL_JOB_ID, - "run_id": MOCK_RUN_ID, + "observability_id": MOCK_OBSERVABILITY_ID, "status": MODAL_EXECUTION_STATUS_SUBMITTED, "poll_url": f"/jobs/{MOCK_MODAL_JOB_ID}", "country": "us", diff --git a/tests/fixtures/services/economy_service.py b/tests/fixtures/services/economy_service.py index cd40425b2..35cb466ec 100644 --- a/tests/fixtures/services/economy_service.py +++ b/tests/fixtures/services/economy_service.py @@ -1,6 +1,7 @@ import datetime import json from unittest.mock import MagicMock, patch +from uuid import UUID import pytest from policyengine_api.constants import ( @@ -35,8 +36,8 @@ ) MOCK_MODAL_JOB_ID = "fc-test123xyz" MOCK_EXECUTION_ID = MOCK_MODAL_JOB_ID # Alias for test compatibility -MOCK_RUN_ID = "run-test123xyz" -MOCK_PROCESS_ID = "job_20250626120000_1234" +MOCK_OBSERVABILITY_ID = "00000000-0000-4000-8000-000000000001" +MOCK_SUBMISSION_CLAIM_ID = "00000000-0000-4000-8000-000000000002" MOCK_MODEL_VERSION = "1.2.3" MOCK_POLICYENGINE_VERSION = "3.4.0" MOCK_RESOLVED_APP_NAME = "policyengine-simulation-us1-2-3-uk2-7-8" @@ -192,11 +193,11 @@ def mock_datetime(): @pytest.fixture -def mock_numpy_random(): - """Mock numpy random integer generation.""" +def mock_submission_claim_id(): + """Return one stable UUID for submission ownership claims.""" with patch( - "policyengine_api.services.economy_service.np.random.randint", - return_value=1234, + "policyengine_api.services.economy_service.uuid.uuid4", + return_value=UUID(MOCK_SUBMISSION_CLAIM_ID), ) as mock: yield mock @@ -223,7 +224,7 @@ def create_mock_reform_impact( }, } ) - return ReformImpact( + impact = ReformImpact( reform_impact_id=1, country_id=MOCK_COUNTRY_ID, reform_policy_id=MOCK_POLICY_ID, @@ -241,6 +242,8 @@ def create_mock_reform_impact( start_time=start_time or datetime.datetime(2025, 6, 26, 12, 0, 0), end_time=(datetime.datetime(2025, 6, 26, 12, 5, 0) if status == "ok" else None), ) + impact.observability_id = MOCK_OBSERVABILITY_ID + return impact def create_mock_modal_execution( @@ -271,7 +274,7 @@ def create_mock_modal_execution( """ mock_execution = MagicMock() mock_execution.job_id = job_id - mock_execution.run_id = MOCK_RUN_ID + mock_execution.observability_id = MOCK_OBSERVABILITY_ID mock_execution.name = job_id # Alias for compatibility mock_execution.status = status mock_execution.result = result @@ -297,6 +300,7 @@ def create_mock_budget_window_batch_execution( mock_execution.batch_job_id = batch_job_id mock_execution.name = batch_job_id mock_execution.status = status + mock_execution.observability_id = MOCK_OBSERVABILITY_ID mock_execution.progress = progress mock_execution.completed_years = completed_years or [] mock_execution.running_years = running_years or [] diff --git a/tests/unit/libs/test_simulation_entrypoint.py b/tests/unit/libs/test_simulation_entrypoint.py index 21c04656c..02980fe83 100644 --- a/tests/unit/libs/test_simulation_entrypoint.py +++ b/tests/unit/libs/test_simulation_entrypoint.py @@ -46,7 +46,7 @@ MOCK_POLL_RESPONSE_FAILED, MOCK_POLL_RESPONSE_RUNNING, MOCK_RESOLVED_APP_NAME, - MOCK_RUN_ID, + MOCK_OBSERVABILITY_ID, MOCK_SIMULATION_PAYLOAD, MOCK_SIMULATION_PAYLOAD_WITH_TELEMETRY, MOCK_SIMULATION_RESULT, @@ -528,7 +528,7 @@ def test__given_valid_payload__then_returns_execution_with_job_id( # Then assert execution.job_id == MOCK_MODAL_JOB_ID - assert execution.run_id == MOCK_RUN_ID + assert execution.observability_id == MOCK_OBSERVABILITY_ID assert execution.status == MODAL_EXECUTION_STATUS_SUBMITTED assert execution.policyengine_bundle == MOCK_POLICYENGINE_BUNDLE assert execution.resolved_app_name == MOCK_RESOLVED_APP_NAME @@ -568,7 +568,10 @@ def test__given_telemetry_payload__then_preserves_it_in_post_body( api.run(MOCK_SIMULATION_PAYLOAD_WITH_TELEMETRY) call_args = mock_httpx_client.post.call_args - assert call_args[1]["json"]["_telemetry"]["run_id"] == MOCK_RUN_ID + assert ( + call_args[1]["json"]["_telemetry"]["observability_id"] + == MOCK_OBSERVABILITY_ID + ) def test__given_model_and_bundle_versions__then_translates_payload_for_modal( self, @@ -614,7 +617,7 @@ def test__given_api_v1_default_bundle_payload__then_posts_gateway_contract_body( "model_version": "1.729.0", "policyengine_version": "4.18.3", "_metadata": { - "process_id": "job_20260629120000_1234", + "submission_claim_id": "job_20260629120000_1234", "model_version": "1.729.0", "policyengine_version": "4.18.3", "data_version": None, @@ -622,8 +625,8 @@ def test__given_api_v1_default_bundle_payload__then_posts_gateway_contract_body( "resolved_app_name": "policyengine-simulation-py4-18-3", }, "_telemetry": { - "run_id": "run_20260629120000_1234", - "process_id": "job_20260629120000_1234", + "observability_id": "run_20260629120000_1234", + "submission_claim_id": "job_20260629120000_1234", "capture_mode": "disabled", }, } @@ -686,7 +689,7 @@ def test__given_network_error__then_raises_exception( log_payload = mock_modal_logger.log_struct.call_args.args[0] assert "Simulation entrypoint request error" in log_payload["message"] - assert log_payload["run_id"] == MOCK_RUN_ID + assert log_payload["observability_id"] == MOCK_OBSERVABILITY_ID class TestResolveAppName: def test__given_country_and_version__then_returns_registered_app( @@ -823,7 +826,7 @@ def test__given_network_error__then_raises_exception( api.run_budget_window_batch(MOCK_SIMULATION_PAYLOAD_WITH_TELEMETRY) log_payload = mock_modal_logger.log_struct.call_args.args[0] - assert log_payload["run_id"] == MOCK_RUN_ID + assert log_payload["observability_id"] == MOCK_OBSERVABILITY_ID class TestGetExecutionById: def test__given_running_job__then_returns_running_status( diff --git a/tests/unit/routes/test_migration_context_logging.py b/tests/unit/routes/test_migration_context_logging.py index c2ce9be38..509da031a 100644 --- a/tests/unit/routes/test_migration_context_logging.py +++ b/tests/unit/routes/test_migration_context_logging.py @@ -111,7 +111,12 @@ def metadata(country_id): assert response.status_code == 200 assert response.headers[REQUEST_ID_HEADER] == "request-123" - runtime.set_context.assert_called_once_with( + assert runtime.set_context.call_count == 2 + runtime.set_context.assert_any_call( + request_id="request-123", + observability_id=response.headers["X-PolicyEngine-Observability-Id"], + ) + runtime.set_context.assert_any_call( country_id="us", route_group="metadata", route_impl="flask_fallback", @@ -125,6 +130,25 @@ def metadata(country_id): mock_logger.log_struct.assert_not_called() +def test_observability_runtime_failure_does_not_reject_flask_request(): + app = Flask(__name__) + runtime = Mock() + runtime.capture_context.side_effect = RuntimeError("runtime unavailable") + runtime.set_context.side_effect = RuntimeError("runtime unavailable") + register_migration_request_logging(app, runtime=runtime) + + @app.get("/health") + def health(): + return {"status": "ok"} + + response = app.test_client().get("/health") + + assert response.status_code == 200 + assert response.json == {"status": "ok"} + assert response.headers[REQUEST_ID_HEADER] + assert response.headers["X-PolicyEngine-Observability-Id"] + + def test_flask_preserves_policyengine_request_id_in_context_log_and_response(): with patch("policyengine_api.migration_logging.logger") as mock_logger: response = ( diff --git a/tests/unit/services/test_economy_service.py b/tests/unit/services/test_economy_service.py index ab0fc7353..2bf00e73c 100644 --- a/tests/unit/services/test_economy_service.py +++ b/tests/unit/services/test_economy_service.py @@ -32,12 +32,12 @@ MOCK_OPTIONS_HASH, MOCK_POLICY_ID, MOCK_POLICYENGINE_VERSION, - MOCK_PROCESS_ID, + MOCK_SUBMISSION_CLAIM_ID, MOCK_REFORM_IMPACT_DATA, MOCK_REGION, MOCK_RESOLVED_APP_NAME, MOCK_RESOLVED_DATASET, - MOCK_RUN_ID, + MOCK_OBSERVABILITY_ID, MOCK_TIME_PERIOD, create_mock_budget_window_batch_execution, create_mock_reform_impact, @@ -115,7 +115,7 @@ def test__given_completed_impact__returns_completed_result( mock_simulation_entrypoint, mock_logger, mock_datetime, - mock_numpy_random, + mock_submission_claim_id, ): completed_impact = create_mock_reform_impact(status="ok") mock_reform_impacts_service.get_all_reform_impacts_by_options_hash_prefix.return_value = [ @@ -151,7 +151,7 @@ def test__given_orm_decoded_completed_impact__returns_completed_result( mock_simulation_entrypoint, mock_logger, mock_datetime, - mock_numpy_random, + mock_submission_claim_id, ): completed_impact = create_mock_reform_impact(status="ok") completed_impact.reform_impact_json = json.loads( @@ -193,7 +193,7 @@ def test__given_cached_error_impact__returns_error_message( mock_simulation_entrypoint, mock_logger, mock_datetime, - mock_numpy_random, + mock_submission_claim_id, ): failed_impact = create_mock_reform_impact( status="error", @@ -222,7 +222,7 @@ def test__given_legacy_completed_impact__refreshes_cache( mock_simulation_entrypoint, mock_logger, mock_datetime, - mock_numpy_random, + mock_submission_claim_id, ): completed_impact = create_mock_reform_impact( status="ok", @@ -254,7 +254,7 @@ def test__given_computing_impact_with_succeeded_execution__returns_completed_res mock_simulation_entrypoint, mock_logger, mock_datetime, - mock_numpy_random, + mock_submission_claim_id, ): computing_impact = create_mock_reform_impact(status="computing") mock_reform_impacts_service.get_all_reform_impacts_by_options_hash_prefix.return_value = [ @@ -293,7 +293,7 @@ def test__given_computing_impact_with_failed_execution__returns_error_result( mock_simulation_entrypoint, mock_logger, mock_datetime, - mock_numpy_random, + mock_submission_claim_id, ): computing_impact = create_mock_reform_impact(status="computing") mock_reform_impacts_service.get_all_reform_impacts_by_options_hash_prefix.return_value = [ @@ -325,7 +325,7 @@ def test__given_computing_impact_with_active_execution__returns_computing_result mock_simulation_entrypoint, mock_logger, mock_datetime, - mock_numpy_random, + mock_submission_claim_id, ): computing_impact = create_mock_reform_impact(status="computing") mock_reform_impacts_service.get_all_reform_impacts_by_options_hash_prefix.return_value = [ @@ -349,7 +349,7 @@ def test__given_no_previous_impact__creates_new_simulation( mock_simulation_entrypoint, mock_logger, mock_datetime, - mock_numpy_random, + mock_submission_claim_id, ): mock_reform_impacts_service.get_all_reform_impacts_by_options_hash_prefix.return_value = [] @@ -386,7 +386,7 @@ def test__given_existing_start_claim__does_not_submit_duplicate_simulation( mock_simulation_entrypoint, mock_logger, mock_datetime, - mock_numpy_random, + mock_submission_claim_id, ): mock_reform_impacts_service.claim_reform_impact_start.return_value = False @@ -408,7 +408,7 @@ def test__given_start_claim_cache_failure__fails_before_submission( mock_simulation_entrypoint, mock_logger, mock_datetime, - mock_numpy_random, + mock_submission_claim_id, ): mock_reform_impacts_service.claim_reform_impact_start.side_effect = ( CacheCoordinationError("cache unavailable") @@ -431,7 +431,7 @@ def test__given_gateway_raises_before_returning_execution__releases_start_claim( mock_simulation_entrypoint, mock_logger, mock_datetime, - mock_numpy_random, + mock_submission_claim_id, ): mock_simulation_entrypoint.run.side_effect = RuntimeError( "submission failed" @@ -455,7 +455,7 @@ def test__given_submitted_simulation_handoff_failure__retains_start_claim( mock_simulation_entrypoint, mock_logger, mock_datetime, - mock_numpy_random, + mock_submission_claim_id, ): mock_reform_impacts_service.set_reform_impact.side_effect = ( ReformImpactHandoffError("cache unavailable") @@ -478,7 +478,7 @@ def test__given_submitted_simulation_without_execution_id__retains_start_claim( mock_simulation_entrypoint, mock_logger, mock_datetime, - mock_numpy_random, + mock_submission_claim_id, ): mock_simulation_entrypoint.get_execution_id.side_effect = RuntimeError( "missing execution identifier" @@ -519,7 +519,7 @@ def test__given_policies_created_through_orm__submits_decoded_json( MOCK_MODEL_VERSION, ) simulation_gateway.get_execution_id.return_value = "execution-1" - simulation_gateway.run.return_value.run_id = "run-1" + simulation_gateway.run.return_value.observability_id = "run-1" monkeypatch.setattr( "policyengine_api.services.economy_service.logger", MagicMock(), @@ -559,7 +559,7 @@ def test__given_no_previous_impact__includes_metadata_in_simulation_params( mock_simulation_entrypoint, mock_logger, mock_datetime, - mock_numpy_random, + mock_submission_claim_id, ): """Verify that _metadata with policy IDs is passed to simulation API.""" mock_reform_impacts_service.get_all_reform_impacts_by_options_hash_prefix.return_value = [] @@ -576,7 +576,10 @@ def test__given_no_previous_impact__includes_metadata_in_simulation_params( assert ( sim_params["_metadata"]["baseline_policy_id"] == MOCK_BASELINE_POLICY_ID ) - assert sim_params["_metadata"]["process_id"] == MOCK_PROCESS_ID + assert ( + sim_params["_metadata"]["submission_claim_id"] + == MOCK_SUBMISSION_CLAIM_ID + ) assert sim_params["_metadata"]["model_version"] == MOCK_MODEL_VERSION assert ( sim_params["_metadata"]["policyengine_version"] @@ -600,7 +603,7 @@ def test__given_no_previous_impact__includes_telemetry_in_simulation_params( mock_simulation_entrypoint, mock_logger, mock_datetime, - mock_numpy_random, + mock_submission_claim_id, ): mock_reform_impacts_service.get_all_reform_impacts.return_value = [] @@ -608,15 +611,18 @@ def test__given_no_previous_impact__includes_telemetry_in_simulation_params( sim_params = mock_simulation_entrypoint.run.call_args[0][0] - assert sim_params["_telemetry"]["run_id"] - assert sim_params["_telemetry"]["process_id"] == MOCK_PROCESS_ID + assert sim_params["_telemetry"]["observability_id"] + assert ( + sim_params["_telemetry"]["submission_claim_id"] + == MOCK_SUBMISSION_CLAIM_ID + ) assert sim_params["_telemetry"]["simulation_kind"] == "national" assert sim_params["_telemetry"]["geography_type"] == "national" assert sim_params["_telemetry"]["geography_code"] == MOCK_COUNTRY_ID assert sim_params["_telemetry"]["capture_mode"] == "disabled" assert sim_params["_telemetry"]["config_hash"].startswith("sha256:") progress_log = mock_logger.log_struct.call_args_list[-1].args[0] - assert progress_log["run_id"] == MOCK_RUN_ID + assert progress_log["observability_id"] == MOCK_OBSERVABILITY_ID assert ( mock_logger.log_struct.call_args_list[-1].kwargs["severity"] == "INFO" ) @@ -632,7 +638,7 @@ def test__given_runtime_cache_version__uses_versioned_economy_cache_key( mock_simulation_entrypoint, mock_logger, mock_datetime, - mock_numpy_random, + mock_submission_claim_id, monkeypatch, ): cache_version = "e1cache01" @@ -669,7 +675,7 @@ def test__given_default_dataset__queries_previous_impacts_with_resolved_bundle( mock_simulation_entrypoint, mock_logger, mock_datetime, - mock_numpy_random, + mock_submission_claim_id, ): mock_reform_impacts_service.get_all_reform_impacts_by_options_hash_prefix.return_value = [] @@ -696,7 +702,7 @@ def test__given_completed_impact__uses_resolved_runtime_bundle_for_cache_lookup( mock_simulation_entrypoint, mock_logger, mock_datetime, - mock_numpy_random, + mock_submission_claim_id, ): completed_impact = create_mock_reform_impact(status="ok") mock_reform_impacts_service.get_all_reform_impacts_by_options_hash_prefix.return_value = [ @@ -723,7 +729,7 @@ def test__given_cached_impact_and_runtime_lookup_fails__then_returns_cached_resu mock_simulation_entrypoint, mock_logger, mock_datetime, - mock_numpy_random, + mock_submission_claim_id, ): completed_impact = create_mock_reform_impact(status="ok") mock_reform_impacts_service.get_all_reform_impacts_by_options_hash_prefix.return_value = [ @@ -752,7 +758,7 @@ def test__given_legacy_cached_impact_without_resolved_app_name__then_refreshes_c mock_simulation_entrypoint, mock_logger, mock_datetime, - mock_numpy_random, + mock_submission_claim_id, ): completed_impact = create_mock_reform_impact( status="ok", @@ -784,7 +790,7 @@ def test__given_legacy_and_refreshed_cached_impacts__then_reuses_refreshed_entry mock_simulation_entrypoint, mock_logger, mock_datetime, - mock_numpy_random, + mock_submission_claim_id, ): legacy_impact = create_mock_reform_impact( status="ok", @@ -823,7 +829,7 @@ def test__given_legacy_cached_impact_and_runtime_lookup_fails__then_returns_cach mock_simulation_entrypoint, mock_logger, mock_datetime, - mock_numpy_random, + mock_submission_claim_id, ): completed_impact = create_mock_reform_impact( status="ok", @@ -854,7 +860,7 @@ def test__given_legacy_computing_impact_without_resolved_app_name__then_reuses_e mock_simulation_entrypoint, mock_logger, mock_datetime, - mock_numpy_random, + mock_submission_claim_id, ): computing_impact = create_mock_reform_impact( status="computing", @@ -882,7 +888,7 @@ def test__given_exception__raises_error( mock_simulation_entrypoint, mock_logger, mock_datetime, - mock_numpy_random, + mock_submission_claim_id, ): mock_reform_impacts_service.get_all_reform_impacts_by_options_hash_prefix.side_effect = Exception( "Database error" @@ -902,7 +908,7 @@ def test__given_uk_request__preserves_model_version_in_bundle( mock_simulation_entrypoint, mock_logger, mock_datetime, - mock_numpy_random, + mock_submission_claim_id, ): mock_country_package_versions["uk"] = "2.7.8" mock_reform_impacts_service.get_all_reform_impacts_by_options_hash_prefix.return_value = [] @@ -930,7 +936,7 @@ def economy_service( mock_policy_service, mock_logger, mock_datetime, - mock_numpy_random, + mock_submission_claim_id, ): return EconomyService() @@ -984,7 +990,7 @@ def test__given_no_cached_batch__submits_parent_batch_and_returns_queued_result( assert submitted_payload["target"] == "general" assert "time_period" not in submitted_payload mock_budget_window_cache.claim_batch_start.assert_called_once_with( - "budget-window-cache-key", MOCK_PROCESS_ID + "budget-window-cache-key", MOCK_SUBMISSION_CLAIM_ID ) mock_budget_window_cache.store_batch_job_id.assert_called_once_with( "budget-window-cache-key", "fc-budget-123" @@ -1277,7 +1283,7 @@ def test__given_gateway_raises_before_returning_batch__clears_start_claim( economy_service.get_budget_window_economic_impact(**base_params) mock_budget_window_cache.clear_starting_claim.assert_called_once_with( - "budget-window-cache-key", MOCK_PROCESS_ID + "budget-window-cache-key", MOCK_SUBMISSION_CLAIM_ID ) @pytest.mark.parametrize("status_code", [400, 422]) @@ -1312,7 +1318,7 @@ def test__given_modal_rejects_batch_submission_for_validation__returns_failed_re assert result.queued_years == ["2026", "2027", "2028"] assert result.cache_status == "miss" mock_budget_window_cache.clear_starting_claim.assert_called_once_with( - "budget-window-cache-key", MOCK_PROCESS_ID + "budget-window-cache-key", MOCK_SUBMISSION_CLAIM_ID ) mock_budget_window_cache.store_batch_job_id.assert_not_called() @@ -1333,7 +1339,7 @@ def test__given_modal_non_validation_error_on_batch_submission__raises( economy_service.get_budget_window_economic_impact(**base_params) mock_budget_window_cache.clear_starting_claim.assert_called_once_with( - "budget-window-cache-key", MOCK_PROCESS_ID + "budget-window-cache-key", MOCK_SUBMISSION_CLAIM_ID ) mock_budget_window_cache.store_batch_job_id.assert_not_called() @@ -1422,7 +1428,7 @@ def test__given_runtime_cache_version__uses_versioned_cache_key_for_budget_windo mock_budget_window_cache, mock_logger, mock_datetime, - mock_numpy_random, + mock_submission_claim_id, monkeypatch, ): cache_version = "e1cache01" @@ -1583,7 +1589,8 @@ def economy_service(self): @pytest.fixture def setup_options(self): return EconomicImpactSetupOptions( - process_id=MOCK_PROCESS_ID, + submission_claim_id=MOCK_SUBMISSION_CLAIM_ID, + observability_id=MOCK_OBSERVABILITY_ID, country_id=MOCK_COUNTRY_ID, reform_policy_id=MOCK_POLICY_ID, baseline_policy_id=MOCK_BASELINE_POLICY_ID, @@ -1683,7 +1690,8 @@ def economy_service(self): @pytest.fixture def setup_options(self): return EconomicImpactSetupOptions( - process_id=MOCK_PROCESS_ID, + submission_claim_id=MOCK_SUBMISSION_CLAIM_ID, + observability_id=MOCK_OBSERVABILITY_ID, country_id=MOCK_COUNTRY_ID, reform_policy_id=MOCK_POLICY_ID, baseline_policy_id=MOCK_BASELINE_POLICY_ID, @@ -1899,19 +1907,16 @@ def test__given_modal_submitted_state__then_returns_computing_result( assert result.status == ImpactStatus.COMPUTING assert result.data is None - class TestCreateProcessId: + class TestCreateSubmissionClaimId: @pytest.fixture def economy_service(self): return EconomyService() - def test_given_mocked_datetime_and_random_returns_expected_format( - self, economy_service, mock_datetime, mock_numpy_random - ): - result = economy_service._create_process_id() + def test_returns_uuid_string(self, economy_service, mock_submission_claim_id): + result = economy_service._create_submission_claim_id() - assert result == "job_20250626120000_1234" - mock_datetime.now.assert_called_once() - mock_numpy_random.assert_called_once_with(1000, 9999) + assert result == MOCK_SUBMISSION_CLAIM_ID + mock_submission_claim_id.assert_called_once_with() class TestEconomicImpactResult: @@ -1978,7 +1983,8 @@ def test__given_error__creates_correct_instance_and_logs(self): class TestEconomicImpactSetupOptions: def test__given_valid_data__creates_instance(self): options = EconomicImpactSetupOptions( - process_id=MOCK_PROCESS_ID, + submission_claim_id=MOCK_SUBMISSION_CLAIM_ID, + observability_id=MOCK_OBSERVABILITY_ID, country_id=MOCK_COUNTRY_ID, reform_policy_id=MOCK_POLICY_ID, baseline_policy_id=MOCK_BASELINE_POLICY_ID, @@ -1991,7 +1997,7 @@ def test__given_valid_data__creates_instance(self): options_hash=MOCK_OPTIONS_HASH, ) - assert options.process_id == MOCK_PROCESS_ID + assert options.submission_claim_id == MOCK_SUBMISSION_CLAIM_ID assert options.country_id == MOCK_COUNTRY_ID assert options.reform_policy_id == MOCK_POLICY_ID assert options.baseline_policy_id == MOCK_BASELINE_POLICY_ID diff --git a/tests/unit/test_asgi_factory.py b/tests/unit/test_asgi_factory.py index d4d1aceaa..ec911df4d 100644 --- a/tests/unit/test_asgi_factory.py +++ b/tests/unit/test_asgi_factory.py @@ -17,6 +17,7 @@ RouteImplementation, RouteImplementationSettings, ) +from policyengine_api.observability.identifiers import OBSERVABILITY_ID_HEADER from policyengine_api.request_context import ( REQUEST_ID_HEADER, current_request_id, @@ -49,6 +50,14 @@ def request_echo(): response.headers["X-Echo"] = "present" return response + @app.get("/stored-observability-id") + def stored_observability_id(): + response = make_response("stored", 200) + response.headers[OBSERVABILITY_ID_HEADER] = ( + "00000000-0000-4000-8000-000000000012" + ) + return response + @app.get("/readiness-check") def readiness_check(): return Response("OK", status=200, mimetype="text/plain") @@ -415,6 +424,18 @@ def test_flask_fallback_preserves_status_body_headers_and_cookies(): assert response.headers["content-type"].startswith("text/html") +def test_flask_fallback_preserves_a_stored_observability_id(): + client = TestClient(create_asgi_app(create_test_wsgi_app())) + + response = client.get("/stored-observability-id") + + assert response.status_code == 200 + assert ( + response.headers[OBSERVABILITY_ID_HEADER] + == "00000000-0000-4000-8000-000000000012" + ) + + def test_large_flask_fallback_response_supports_http_gzip(): client = TestClient(create_asgi_app(create_test_wsgi_app())) diff --git a/tests/unit/test_observability_stage_registry.py b/tests/unit/test_observability_stage_registry.py new file mode 100644 index 000000000..3698be2c9 --- /dev/null +++ b/tests/unit/test_observability_stage_registry.py @@ -0,0 +1,31 @@ +import pytest + +from policyengine_api.observability.stages import ( + RUN_STAGE_REGISTRY, + RunConfiguration, + Stage, +) + + +def test_registry_defines_every_run_configuration_and_stage() -> None: + assert set(RUN_STAGE_REGISTRY) == set(RunConfiguration) + registered = { + stage + for stage_plan in RUN_STAGE_REGISTRY.values() + for stage in stage_plan.stages + } + assert registered == set(Stage) + + +@pytest.mark.parametrize("stage_plan", RUN_STAGE_REGISTRY.values()) +def test_each_stage_plan_is_ordered_and_contains_no_duplicates(stage_plan) -> None: + assert stage_plan.stages + assert len(stage_plan.stages) == len(set(stage_plan.stages)) + assert all(stage_plan.name(stage) == stage.value for stage in stage_plan.stages) + + +def test_stage_plan_rejects_a_stage_from_another_configuration() -> None: + household = RUN_STAGE_REGISTRY[RunConfiguration.HOUSEHOLD] + + with pytest.raises(ValueError, match="not registered"): + household.name(Stage.ECONOMY_SUBMIT) From f5884d6eeef3aff6c526ac45b06160fbf9dee33f Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 24 Sep 2026 23:18:11 +0400 Subject: [PATCH 08/10] Use observability header as canonical transport --- .../libs/simulation_entrypoint.py | 37 +++++++++------- policyengine_api/services/economy_service.py | 7 ++- tests/fixtures/libs/simulation_entrypoint.py | 8 ++-- tests/unit/libs/test_simulation_entrypoint.py | 43 ++++++++++++++----- tests/unit/services/test_economy_service.py | 2 +- 5 files changed, 64 insertions(+), 33 deletions(-) diff --git a/policyengine_api/libs/simulation_entrypoint.py b/policyengine_api/libs/simulation_entrypoint.py index 1d9d2678b..be9f1c44d 100644 --- a/policyengine_api/libs/simulation_entrypoint.py +++ b/policyengine_api/libs/simulation_entrypoint.py @@ -23,7 +23,10 @@ current_observability_id, current_request_id, ) -from policyengine_api.observability.identifiers import OBSERVABILITY_ID_HEADER +from policyengine_api.observability.identifiers import ( + OBSERVABILITY_ID_HEADER, + normalize_observability_id, +) from policyengine_api.worker_spm import validate_worker_spm, raise_worker_spm_error @@ -69,6 +72,15 @@ def _attach_current_request_id(request: httpx.Request) -> None: request.headers[OBSERVABILITY_ID_HEADER] = observability_id +def _response_observability_id(response: httpx.Response) -> str | None: + """Read the diagnostic identifier from the canonical response header.""" + + return ( + normalize_observability_id(response.headers.get(OBSERVABILITY_ID_HEADER)) + or current_observability_id() + ) + + @dataclass class ModalSimulationExecution: """ @@ -222,12 +234,13 @@ def run(self, payload: dict) -> ModalSimulationExecution: raise_worker_spm_error(response) response.raise_for_status() data = response.json() + observability_id = _response_observability_id(response) logger.log_struct( { "message": "Simulation entrypoint job submitted", "job_id": data.get("job_id"), - "observability_id": data.get("observability_id"), + "observability_id": observability_id, "status": data.get("status"), }, severity="INFO", @@ -238,16 +251,14 @@ def run(self, payload: dict) -> ModalSimulationExecution: status=data["status"], policyengine_bundle=data.get("policyengine_bundle"), resolved_app_name=data.get("resolved_app_name"), - observability_id=data.get("observability_id"), + observability_id=observability_id, ) except httpx.HTTPStatusError as e: logger.log_struct( { "message": f"Simulation entrypoint HTTP error: {e.response.status_code}", - "observability_id": (payload.get("_telemetry") or {}).get( - "observability_id" - ), + "observability_id": current_observability_id(), "response_text": e.response.text[:500], }, severity="ERROR", @@ -258,9 +269,7 @@ def run(self, payload: dict) -> ModalSimulationExecution: logger.log_struct( { "message": f"Simulation entrypoint request error: {str(e)}", - "observability_id": (payload.get("_telemetry") or {}).get( - "observability_id" - ), + "observability_id": current_observability_id(), }, severity="ERROR", ) @@ -293,7 +302,7 @@ def run_budget_window_batch(self, payload: dict) -> ModalBudgetWindowBatchExecut return ModalBudgetWindowBatchExecution( batch_job_id=data["batch_job_id"], status=data["status"], - observability_id=data.get("observability_id"), + observability_id=_response_observability_id(response), ) except httpx.HTTPStatusError as e: @@ -310,9 +319,7 @@ def run_budget_window_batch(self, payload: dict) -> ModalBudgetWindowBatchExecut logger.log_struct( { "message": f"Simulation batch API request error: {str(e)}", - "observability_id": (payload.get("_telemetry") or {}).get( - "observability_id" - ), + "observability_id": current_observability_id(), }, severity="ERROR", ) @@ -447,7 +454,7 @@ def get_execution_by_id(self, job_id: str) -> ModalSimulationExecution: return ModalSimulationExecution( job_id=job_id, status=data["status"], - observability_id=data.get("observability_id"), + observability_id=_response_observability_id(response), result=data.get("result"), error=data.get("error"), policyengine_bundle=data.get("policyengine_bundle"), @@ -491,7 +498,7 @@ def get_budget_window_batch_by_id( return ModalBudgetWindowBatchExecution( batch_job_id=batch_job_id, status=data["status"], - observability_id=data.get("observability_id"), + observability_id=_response_observability_id(response), progress=data.get("progress"), completed_years=data.get("completed_years", []), running_years=data.get("running_years", []), diff --git a/policyengine_api/services/economy_service.py b/policyengine_api/services/economy_service.py index 7282ecffa..435441d67 100644 --- a/policyengine_api/services/economy_service.py +++ b/policyengine_api/services/economy_service.py @@ -1257,12 +1257,12 @@ def _handle_create_impact( logger.log_struct( { "message": "Setting up sim API job", - "observability_id": telemetry["observability_id"], + "observability_id": setup_options.observability_id, **setup_options.model_dump(), } ) - # Preserve both legacy metadata and the new telemetry envelope. + # Preserve execution metadata and non-identity simulation telemetry. sim_params["_metadata"] = { "reform_policy_id": setup_options.reform_policy_id, "baseline_policy_id": setup_options.baseline_policy_id, @@ -1292,7 +1292,7 @@ def _handle_create_impact( observability_id = ( getattr(entrypoint_execution, "observability_id", None) - or telemetry["observability_id"] + or setup_options.observability_id ) progress_log = { @@ -1520,7 +1520,6 @@ def _build_simulation_telemetry( ) return { - "observability_id": setup_options.observability_id, "submission_claim_id": setup_options.submission_claim_id, "requested_at": datetime.datetime.now(datetime.UTC).isoformat(), "simulation_kind": simulation_kind, diff --git a/tests/fixtures/libs/simulation_entrypoint.py b/tests/fixtures/libs/simulation_entrypoint.py index 5b8f12c92..47263c121 100644 --- a/tests/fixtures/libs/simulation_entrypoint.py +++ b/tests/fixtures/libs/simulation_entrypoint.py @@ -18,7 +18,7 @@ # Mock data constants MOCK_MODAL_JOB_ID = "fc-abc123xyz" -MOCK_OBSERVABILITY_ID = "run-abc123xyz" +MOCK_OBSERVABILITY_ID = "00000000-0000-4000-8000-000000000001" MOCK_BATCH_JOB_ID = "fc-batch123xyz" MOCK_MODAL_BASE_URL = "https://test-modal-api.modal.run" @@ -35,7 +35,6 @@ MOCK_SIMULATION_PAYLOAD_WITH_TELEMETRY = { **MOCK_SIMULATION_PAYLOAD, "_telemetry": { - "observability_id": MOCK_OBSERVABILITY_ID, "submission_claim_id": "job_20250626120000_1234", "capture_mode": "disabled", }, @@ -59,7 +58,6 @@ MOCK_SUBMIT_RESPONSE_SUCCESS = { "job_id": MOCK_MODAL_JOB_ID, - "observability_id": MOCK_OBSERVABILITY_ID, "status": MODAL_EXECUTION_STATUS_SUBMITTED, "poll_url": f"/jobs/{MOCK_MODAL_JOB_ID}", "country": "us", @@ -142,6 +140,7 @@ def create_mock_httpx_response( status_code: int = 200, json_data: dict = None, + headers: dict | None = None, ): """ Helper function to create a mock httpx response. @@ -160,6 +159,9 @@ def create_mock_httpx_response( """ mock_response = MagicMock() mock_response.status_code = status_code + mock_response.headers = headers or { + "X-PolicyEngine-Observability-Id": MOCK_OBSERVABILITY_ID + } mock_response.json.return_value = json_data or {} mock_response.text = json.dumps(json_data or {}) mock_response.raise_for_status = MagicMock() diff --git a/tests/unit/libs/test_simulation_entrypoint.py b/tests/unit/libs/test_simulation_entrypoint.py index 02980fe83..1a1fc9d33 100644 --- a/tests/unit/libs/test_simulation_entrypoint.py +++ b/tests/unit/libs/test_simulation_entrypoint.py @@ -29,8 +29,12 @@ ) from policyengine_api.request_context import ( # noqa: E402 REQUEST_ID_HEADER, + _asgi_observability_id, _asgi_request_id, ) +from policyengine_api.observability.identifiers import ( # noqa: E402 + OBSERVABILITY_ID_HEADER, +) from tests.fixtures.libs.simulation_entrypoint import ( # noqa: E402 MOCK_BATCH_JOB_ID, @@ -111,7 +115,12 @@ def _response(self, method, url, json=None): payload = MOCK_HEALTH_RESPONSE status_code = 200 - return httpx.Response(status_code, request=request, json=payload) + return httpx.Response( + status_code, + request=request, + json=payload, + headers={"X-PolicyEngine-Observability-Id": MOCK_OBSERVABILITY_ID}, + ) def post(self, url, json=None): return self._response("POST", url, json=json) @@ -405,7 +414,7 @@ def test__given_flask_request__then_hook_uses_current_request_id( assert request.headers[REQUEST_ID_HEADER] == "flask-request-id" - def test__given_asgi_request__then_hook_uses_current_request_id( + def test__given_asgi_request__then_hook_uses_current_correlation_ids( self, mock_httpx_client ): from policyengine_api.libs.simulation_entrypoint import httpx as modal_httpx @@ -414,12 +423,15 @@ def test__given_asgi_request__then_hook_uses_current_request_id( hook = modal_httpx.Client.call_args.kwargs["event_hooks"]["request"][0] request = httpx.Request("GET", MOCK_MODAL_BASE_URL) token = _asgi_request_id.set("asgi-request-id") + observability_token = _asgi_observability_id.set(MOCK_OBSERVABILITY_ID) try: hook(request) finally: _asgi_request_id.reset(token) + _asgi_observability_id.reset(observability_token) assert request.headers[REQUEST_ID_HEADER] == "asgi-request-id" + assert request.headers[OBSERVABILITY_ID_HEADER] == MOCK_OBSERVABILITY_ID def test__given_no_request_context__then_hook_omits_request_id( self, monkeypatch, mock_modal_logger @@ -554,7 +566,7 @@ def test__given_valid_payload__then_posts_to_correct_endpoint( assert "/simulate/economy/comparison" in call_args[0][0] assert call_args[1]["json"] == MOCK_SIMULATION_PAYLOAD - def test__given_telemetry_payload__then_preserves_it_in_post_body( + def test__given_telemetry_payload__then_preserves_non_identity_fields( self, mock_httpx_client, mock_modal_logger, @@ -568,10 +580,10 @@ def test__given_telemetry_payload__then_preserves_it_in_post_body( api.run(MOCK_SIMULATION_PAYLOAD_WITH_TELEMETRY) call_args = mock_httpx_client.post.call_args - assert ( - call_args[1]["json"]["_telemetry"]["observability_id"] - == MOCK_OBSERVABILITY_ID - ) + assert call_args[1]["json"]["_telemetry"] == { + "submission_claim_id": "job_20250626120000_1234", + "capture_mode": "disabled", + } def test__given_model_and_bundle_versions__then_translates_payload_for_modal( self, @@ -625,7 +637,6 @@ def test__given_api_v1_default_bundle_payload__then_posts_gateway_contract_body( "resolved_app_name": "policyengine-simulation-py4-18-3", }, "_telemetry": { - "observability_id": "run_20260629120000_1234", "submission_claim_id": "job_20260629120000_1234", "capture_mode": "disabled", }, @@ -684,7 +695,13 @@ def test__given_network_error__then_raises_exception( api = SimulationAPIModal() # When/Then - with pytest.raises(httpx.RequestError): + with ( + patch( + "policyengine_api.libs.simulation_entrypoint.current_observability_id", + return_value=MOCK_OBSERVABILITY_ID, + ), + pytest.raises(httpx.RequestError), + ): api.run(MOCK_SIMULATION_PAYLOAD_WITH_TELEMETRY) log_payload = mock_modal_logger.log_struct.call_args.args[0] @@ -822,7 +839,13 @@ def test__given_network_error__then_raises_exception( mock_httpx_client.post.side_effect = httpx.RequestError("Connection failed") api = SimulationAPIModal() - with pytest.raises(httpx.RequestError): + with ( + patch( + "policyengine_api.libs.simulation_entrypoint.current_observability_id", + return_value=MOCK_OBSERVABILITY_ID, + ), + pytest.raises(httpx.RequestError), + ): api.run_budget_window_batch(MOCK_SIMULATION_PAYLOAD_WITH_TELEMETRY) log_payload = mock_modal_logger.log_struct.call_args.args[0] diff --git a/tests/unit/services/test_economy_service.py b/tests/unit/services/test_economy_service.py index 2bf00e73c..0e65d9ee1 100644 --- a/tests/unit/services/test_economy_service.py +++ b/tests/unit/services/test_economy_service.py @@ -611,7 +611,7 @@ def test__given_no_previous_impact__includes_telemetry_in_simulation_params( sim_params = mock_simulation_entrypoint.run.call_args[0][0] - assert sim_params["_telemetry"]["observability_id"] + assert "observability_id" not in sim_params["_telemetry"] assert ( sim_params["_telemetry"]["submission_claim_id"] == MOCK_SUBMISSION_CLAIM_ID From da8d7511e0f76cf421edeacc4cd43de8aa014e29 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 24 Sep 2026 23:58:24 +0400 Subject: [PATCH 09/10] Complete native request observability and Stage 12 IAM --- gcp/observability/iam.template.yaml | 3 +- gcp/observability/log-routing.template.yaml | 3 +- .../workload-inventory.template.yaml | 4 + policyengine_api/asgi_factory.py | 80 ++++++++++++++++++- tests/unit/test_asgi_factory.py | 72 +++++++++++++++++ .../test_observability_deployment_assets.py | 13 +++ 6 files changed, 170 insertions(+), 5 deletions(-) diff --git a/gcp/observability/iam.template.yaml b/gcp/observability/iam.template.yaml index 062e90bbd..6e6528e56 100644 --- a/gcp/observability/iam.template.yaml +++ b/gcp/observability/iam.template.yaml @@ -33,4 +33,5 @@ modal_workload_identity: assertion.workspace_id == "${MODAL_WORKSPACE_ID}" && (assertion.environment_name == "main" || assertion.environment_name == "staging") && (assertion.app_name == "policyengine-simulation-gateway" || - assertion.app_name.matches("^policyengine-simulation-py[0-9]+-[0-9]+-[0-9]+$")) + assertion.app_name.matches("^policyengine-simulation-py[0-9]+-[0-9]+-[0-9]+$") || + assertion.app_name.matches("^policyengine-simulation-v2-py[0-9]+-[0-9]+-[0-9]+$")) diff --git a/gcp/observability/log-routing.template.yaml b/gcp/observability/log-routing.template.yaml index 8d84ed8a6..6a6ad20a6 100644 --- a/gcp/observability/log-routing.template.yaml +++ b/gcp/observability/log-routing.template.yaml @@ -34,4 +34,5 @@ default_bucket_exclusion: jsonPayload.schema_version="policyengine.observability.v2" AND jsonPayload."service.namespace"="policyengine.api-v1" AND (jsonPayload."service.name"="policyengine-simulation-gateway" OR - jsonPayload."service.name"=~"^policyengine-simulation-py[0-9]+-[0-9]+-[0-9]+$") + jsonPayload."service.name"=~"^policyengine-simulation-py[0-9]+-[0-9]+-[0-9]+$" OR + jsonPayload."service.name"=~"^policyengine-simulation-v2-py[0-9]+-[0-9]+-[0-9]+$") diff --git a/gcp/observability/workload-inventory.template.yaml b/gcp/observability/workload-inventory.template.yaml index f74db600e..46e97b554 100644 --- a/gcp/observability/workload-inventory.template.yaml +++ b/gcp/observability/workload-inventory.template.yaml @@ -52,6 +52,10 @@ modal: value: policyengine-simulation-py name_constraint: '^policyengine-simulation-py[0-9]+-[0-9]+-[0-9]+$' signals: [logs, traces, metrics] + - match: prefix + value: policyengine-simulation-v2-py + name_constraint: '^policyengine-simulation-v2-py[0-9]+-[0-9]+-[0-9]+$' + signals: [logs, traces, metrics] excluded: modal_exact: diff --git a/policyengine_api/asgi_factory.py b/policyengine_api/asgi_factory.py index cc1977ac7..80a1350a8 100644 --- a/policyengine_api/asgi_factory.py +++ b/policyengine_api/asgi_factory.py @@ -10,7 +10,6 @@ from fastapi import FastAPI, Request from fastapi.exception_handlers import request_validation_exception_handler from fastapi.exceptions import RequestValidationError -from fastapi.routing import APIRoute from policyengine_api.constants import VERSION from policyengine_api.fastapi_routes.dependencies import NativeRouteDependencies from policyengine_api.fastapi_routes.health import build_core_health_router @@ -29,6 +28,7 @@ RouteImplementationSettings, ) from policyengine_api.migration_logging import log_migration_request +from policyengine_api.observability import get_runtime from policyengine_api.request_context import ( REQUEST_ID_HEADER, _asgi_observability_id, @@ -36,12 +36,14 @@ generate_request_id, resolve_observability_id, ) +from policyengine_observability import ObservabilityRuntime from policyengine_api.observability.identifiers import OBSERVABILITY_ID_HEADER from policyengine_api.request_context import current_observability_id from starlette.datastructures import MutableHeaders from starlette.middleware.cors import CORSMiddleware from starlette.middleware.gzip import GZipMiddleware from starlette.responses import PlainTextResponse, Response +from starlette.routing import Match, Mount from starlette.types import ASGIApp @@ -59,12 +61,23 @@ def _apply_observability_id_header( response.headers[OBSERVABILITY_ID_HEADER] = observability_id +def _is_native_request(app: FastAPI, scope: dict) -> bool: + """Return whether FastAPI, rather than the mounted Flask app, handles it.""" + + for route in app.router.routes: + match, _ = route.matches(scope) + if match is Match.FULL: + return not isinstance(route, Mount) + return False + + def create_asgi_app( wsgi_app, *, route_settings: RouteImplementationSettings | None = None, dependencies: NativeRouteDependencies | None = None, shutdown_callback: Callable[[], None] | None = None, + observability_runtime: ObservabilityRuntime | None = None, ) -> ASGIApp: """Create the Stage 2 FastAPI shell around the existing Flask app.""" @@ -72,6 +85,7 @@ def create_asgi_app( route_settings = RouteImplementationSettings.from_environment() if dependencies is None: dependencies = NativeRouteDependencies.defaults() + request_runtime = observability_runtime or get_runtime() @asynccontextmanager async def lifespan(_app: FastAPI): @@ -145,9 +159,38 @@ async def add_request_context_and_migration_logging(request, call_next): request.state.policyengine_observability_id = observability_id context_token = _asgi_request_id.set(request_id) observability_context_token = _asgi_observability_id.set(observability_id) + native_request = _is_native_request(app, request.scope) + initial_route = request.url.path + + if native_request: + try: + runtime_request_id = request_runtime.begin_request( + headers={ + **dict(request.headers), + REQUEST_ID_HEADER: request_id, + OBSERVABILITY_ID_HEADER: observability_id, + }, + method=request.method, + route=initial_route, + ) + if isinstance(runtime_request_id, str) and runtime_request_id: + request_id = runtime_request_id + MutableHeaders(scope=request.scope)[REQUEST_ID_HEADER] = request_id + request.state.policyengine_request_id = request_id + _asgi_request_id.reset(context_token) + context_token = _asgi_request_id.set(request_id) + except Exception: + pass + try: + request_runtime.set_context( + request_id=request_id, + observability_id=observability_id, + ) + except Exception: + pass def log_native_route(status_code: int) -> None: - if not isinstance(request.scope.get("route"), APIRoute): + if not native_request: return try: log_migration_request( @@ -165,12 +208,42 @@ def log_native_route(status_code: int) -> None: except Exception: pass + def finish_native_route( + status_code: int, + error: BaseException | None = None, + ) -> None: + if not native_request: + return + resolved_route = getattr(request.scope.get("route"), "path", initial_route) + try: + request_runtime.update_request_route(resolved_route) + except Exception: + pass + try: + request_runtime.update_request_status(status_code) + except Exception: + pass + try: + request_runtime.end_request( + status_code=status_code, + error=error, + ) + except Exception: + pass + try: try: response = await call_next(request) - except Exception: + except Exception as error: log_native_route(500) + finish_native_route(500, error) raise + if native_request: + try: + for name, value in request_runtime.response_headers().items(): + response.headers[name] = value + except Exception: + pass _apply_request_id_header(response, request_id) response_observability_id = ( response.headers.get(OBSERVABILITY_ID_HEADER) @@ -179,6 +252,7 @@ def log_native_route(status_code: int) -> None: ) _apply_observability_id_header(response, response_observability_id) log_native_route(response.status_code) + finish_native_route(response.status_code) return response finally: _asgi_request_id.reset(context_token) diff --git a/tests/unit/test_asgi_factory.py b/tests/unit/test_asgi_factory.py index ec911df4d..6956bc655 100644 --- a/tests/unit/test_asgi_factory.py +++ b/tests/unit/test_asgi_factory.py @@ -312,6 +312,78 @@ def capture_request(**kwargs): generate_request_id.assert_called_once_with() +def test_native_route_uses_observability_request_lifecycle(): + runtime = Mock() + runtime.begin_request.return_value = "request-123" + runtime.response_headers.return_value = { + "traceparent": "00-00000000000000000000000000000001-0000000000000001-01" + } + observability_id = "00000000-0000-4000-8000-000000000001" + + response = TestClient( + create_asgi_app( + create_test_wsgi_app(), + observability_runtime=runtime, + ) + ).get( + "/health", + headers={ + REQUEST_ID_HEADER: "request-123", + OBSERVABILITY_ID_HEADER: observability_id, + }, + ) + + assert response.status_code == 200 + assert response.headers["traceparent"].startswith("00-") + runtime.begin_request.assert_called_once() + assert runtime.begin_request.call_args.kwargs["method"] == "GET" + assert runtime.begin_request.call_args.kwargs["route"] == "/health" + runtime.set_context.assert_called_once_with( + request_id="request-123", + observability_id=observability_id, + ) + runtime.update_request_route.assert_called_once_with("/health") + runtime.update_request_status.assert_called_once_with(200) + runtime.end_request.assert_called_once_with(status_code=200, error=None) + + +def test_flask_fallback_does_not_duplicate_observability_request_lifecycle(): + runtime = Mock() + + response = TestClient( + create_asgi_app( + create_test_wsgi_app(), + observability_runtime=runtime, + ) + ).get("/fallback") + + assert response.status_code == 202 + runtime.begin_request.assert_not_called() + runtime.end_request.assert_not_called() + + +def test_native_route_survives_observability_runtime_failures(): + runtime = Mock() + runtime.begin_request.side_effect = RuntimeError("begin unavailable") + runtime.set_context.side_effect = RuntimeError("context unavailable") + runtime.response_headers.side_effect = RuntimeError("headers unavailable") + runtime.update_request_route.side_effect = RuntimeError("route unavailable") + runtime.update_request_status.side_effect = RuntimeError("status unavailable") + runtime.end_request.side_effect = RuntimeError("finish unavailable") + + response = TestClient( + create_asgi_app( + create_test_wsgi_app(), + observability_runtime=runtime, + ) + ).get("/health") + + assert response.status_code == 200 + assert response.json() == {"status": "healthy"} + assert response.headers[REQUEST_ID_HEADER] + assert response.headers[OBSERVABILITY_ID_HEADER] + + def test_native_route_does_not_accept_x_request_id_as_an_alias(): with ( patch( diff --git a/tests/unit/test_observability_deployment_assets.py b/tests/unit/test_observability_deployment_assets.py index 5f048a608..98c1f53f5 100644 --- a/tests/unit/test_observability_deployment_assets.py +++ b/tests/unit/test_observability_deployment_assets.py @@ -46,9 +46,22 @@ def test_authorization_assets_exclude_unrelated_applications() -> None: assert excluded not in routing assert "policyengine-simulation-gateway" in iam assert "policyengine-simulation-py" in iam + assert "policyengine-simulation-v2-py" in iam + assert "policyengine-simulation-v2-py" in routing assert 'jsonPayload."service.namespace"' in routing +def test_stage12_modal_apps_are_in_the_workload_identity_allowlist() -> None: + iam = (DEPLOY / "iam.template.yaml").read_text() + inventory = (DEPLOY / "workload-inventory.template.yaml").read_text() + routing = (DEPLOY / "log-routing.template.yaml").read_text() + stage12_pattern = "^policyengine-simulation-v2-py[0-9]+-[0-9]+-[0-9]+$" + + assert stage12_pattern in iam + assert stage12_pattern in inventory + assert stage12_pattern in routing + + def test_deployment_templates_use_environment_placeholders() -> None: templates = [ DEPLOY / "iam.template.yaml", From 9d2a1869ee30bfef123a0ff86c1b3fd32097f7f9 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Fri, 25 Sep 2026 00:16:51 +0400 Subject: [PATCH 10/10] Keep budget-window observability IDs durable --- .../services/budget_window_cache.py | 404 ++++++++++-------- policyengine_api/services/economy_service.py | 158 +++++-- tests/fixtures/services/economy_service.py | 9 +- .../routes/test_spm_year_worker_polling.py | 24 +- .../unit/services/test_budget_window_cache.py | 233 +++++----- tests/unit/services/test_economy_service.py | 175 ++++++-- 6 files changed, 618 insertions(+), 385 deletions(-) diff --git a/policyengine_api/services/budget_window_cache.py b/policyengine_api/services/budget_window_cache.py index ef5261051..8d2eb08ad 100644 --- a/policyengine_api/services/budget_window_cache.py +++ b/policyengine_api/services/budget_window_cache.py @@ -1,7 +1,10 @@ -"""Shared, namespaced budget-window result cache and coordination claims.""" +"""Shared, namespaced budget-window state and coordination claims.""" + +from __future__ import annotations import time -from typing import Any +from dataclasses import dataclass +from typing import Any, Literal from policyengine_api.runtime_cache.claims import ExpiringClaimStore from policyengine_api.runtime_cache.core import ( @@ -17,15 +20,93 @@ BUDGET_WINDOW_CACHE_FAMILY = "budget-window" -BUDGET_WINDOW_CACHE_SCHEMA_VERSION = 1 -BUDGET_WINDOW_STARTING_PREFIX = "starting:" +BUDGET_WINDOW_CACHE_SCHEMA_VERSION = 2 BUDGET_WINDOW_STARTING_TTL_SECONDS = 300 BUDGET_WINDOW_BATCH_TTL_SECONDS = 86_400 BUDGET_WINDOW_RESULT_TTL_SECONDS = 2_592_000 +BudgetWindowStateStatus = Literal["starting", "submitted", "completed", "failed"] +BudgetWindowFailureType = Literal["spm_validation", "execution"] + + +@dataclass(frozen=True) +class BudgetWindowCacheState: + """One atomic cache document for a budget-window report.""" + + status: BudgetWindowStateStatus + observability_id: str | None = None + submission_claim_id: str | None = None + batch_job_id: str | None = None + result: dict[str, Any] | None = None + failure_type: BudgetWindowFailureType | None = None + error: dict[str, Any] | None = None + + def to_payload(self) -> dict[str, Any]: + return { + key: value + for key, value in { + "status": self.status, + "observability_id": self.observability_id, + "submission_claim_id": self.submission_claim_id, + "batch_job_id": self.batch_job_id, + "result": self.result, + "failure_type": self.failure_type, + "error": self.error, + }.items() + if value is not None + } + + @classmethod + def from_payload(cls, payload: object) -> BudgetWindowCacheState | None: + if not isinstance(payload, dict): + return None + status = payload.get("status") + if status not in {"starting", "submitted", "completed", "failed"}: + return None + observability_id = payload.get("observability_id") + if observability_id is not None and not isinstance(observability_id, str): + return None + submission_claim_id = payload.get("submission_claim_id") + if submission_claim_id is not None and not isinstance(submission_claim_id, str): + return None + batch_job_id = payload.get("batch_job_id") + if batch_job_id is not None and not isinstance(batch_job_id, str): + return None + result = payload.get("result") + if result is not None and not isinstance(result, dict): + return None + failure_type = payload.get("failure_type") + if failure_type is not None and failure_type not in { + "spm_validation", + "execution", + }: + return None + error = payload.get("error") + if error is not None and not isinstance(error, dict): + return None + + if status == "starting" and not submission_claim_id: + return None + if status == "submitted" and not batch_job_id: + return None + if status == "completed" and result is None: + return None + if status == "failed" and (failure_type is None or error is None): + return None + + return cls( + status=status, + observability_id=observability_id, + submission_claim_id=submission_claim_id, + batch_job_id=batch_job_id, + result=result, + failure_type=failure_type, + error=error, + ) + class BudgetWindowCache: - """Recoverable results plus fail-closed expensive-work coordination.""" + """Atomic report state plus fail-closed expensive-work coordination.""" def __init__( self, @@ -71,20 +152,16 @@ def build_key( ) @staticmethod - def _result_key(cache_key: str) -> str: - return f"{cache_key}:result" - - @staticmethod - def _error_key(cache_key: str) -> str: - return f"{cache_key}:terminal-error" + def _state_key(cache_key: str) -> str: + return f"{cache_key}:state" @staticmethod - def _batch_key(cache_key: str) -> str: - return f"{cache_key}:batch-job-id" - - @staticmethod - def _observability_key(cache_key: str) -> str: - return f"{cache_key}:observability-id" + def _encoded_state(state: BudgetWindowCacheState) -> str: + return encode_envelope( + BUDGET_WINDOW_CACHE_FAMILY, + BUDGET_WINDOW_CACHE_SCHEMA_VERSION, + state.to_payload(), + ) @staticmethod def _handle_cache_error( @@ -101,157 +178,116 @@ def _handle_cache_error( severity="WARNING", ) - def get_completed_result(self, cache_key: str) -> dict[str, Any] | None: - return self._get_payload(self._result_key(cache_key), "result") - - def get_terminal_error(self, cache_key: str) -> dict[str, str] | None: - """Replay a typed failure independently of completed success payloads.""" - error = self._get_payload(self._error_key(cache_key), "terminal-error") - if ( - error is not None - and set(error) == {"code", "message"} - and isinstance(error["code"], str) - and isinstance(error["message"], str) - ): - return error - return None - - def _get_payload(self, key: str, kind: str) -> dict[str, Any] | None: - started_at = time.perf_counter() - try: - payload = self.client.get(key) - except Exception: - self._handle_cache_error( - f"read-{kind}", - event="connection-failed", - started_at=started_at, - ) - return None - result = decode_envelope( - payload, - family=BUDGET_WINDOW_CACHE_FAMILY, - schema_version=BUDGET_WINDOW_CACHE_SCHEMA_VERSION, - ) - if payload is not None and result is None: - self._handle_cache_error( - f"decode-{kind}", - event="decode-failed", - started_at=started_at, - ) - else: - record_cache_event( - family=BUDGET_WINDOW_CACHE_FAMILY, - event="hit" if isinstance(result, dict) else "miss", - operation=f"read-{kind}", - started_at=started_at, - ) - return result if isinstance(result, dict) else None - - def set_completed_result( - self, - cache_key: str, - result: dict[str, Any], - ) -> bool: - return self._set_payload(self._result_key(cache_key), result, "result") - - def set_terminal_error(self, cache_key: str, error: dict[str, str]) -> bool: - """Retain deterministic typed failures for the existing result lifetime.""" - return self._set_payload(self._error_key(cache_key), error, "terminal-error") + def get_state(self, cache_key: str) -> BudgetWindowCacheState | None: + """Read the complete report state or fail closed on cache outage.""" - def _set_payload(self, key: str, result: dict[str, Any], kind: str) -> bool: started_at = time.perf_counter() + state_key = self._state_key(cache_key) try: - stored = self.client.set( - key, - encode_envelope( - BUDGET_WINDOW_CACHE_FAMILY, - BUDGET_WINDOW_CACHE_SCHEMA_VERSION, - result, - ), - ex=jittered_ttl(BUDGET_WINDOW_RESULT_TTL_SECONDS), - ) - except Exception: - self._handle_cache_error( - f"write-{kind}", - event="write-failed", - started_at=started_at, - ) - return False - record_cache_event( - family=BUDGET_WINDOW_CACHE_FAMILY, - event="write", - operation=f"write-{kind}", - started_at=started_at, - ) - return bool(stored) - - def get_batch_job_id(self, cache_key: str) -> str | None: - started_at = time.perf_counter() - try: - value = self.client.get(self._batch_key(cache_key)) + encoded = self.client.get(state_key) except Exception as error: self._handle_cache_error( - "read-batch-id", + "read-state", event="coordination-failed", started_at=started_at, ) raise CacheCoordinationError( "budget-window coordination state is unavailable" ) from error - if not isinstance(value, str) or not value: - record_cache_event( - family=BUDGET_WINDOW_CACHE_FAMILY, - event="coordination-miss", - operation="read-batch-id", - started_at=started_at, - ) - return None - if value.startswith(BUDGET_WINDOW_STARTING_PREFIX): - record_cache_event( - family=BUDGET_WINDOW_CACHE_FAMILY, - event="claim-contended", - operation="read-batch-id", + + payload = decode_envelope( + encoded, + family=BUDGET_WINDOW_CACHE_FAMILY, + schema_version=BUDGET_WINDOW_CACHE_SCHEMA_VERSION, + ) + state = BudgetWindowCacheState.from_payload(payload) + if encoded is not None and state is None: + self._handle_cache_error( + "decode-state", + event="decode-failed", started_at=started_at, ) + self._clear_invalid_state(state_key, encoded) return None + record_cache_event( family=BUDGET_WINDOW_CACHE_FAMILY, - event="coordination-hit", - operation="read-batch-id", + event="hit" if state is not None else "miss", + operation="read-state", started_at=started_at, ) - return value - - def get_observability_id(self, cache_key: str) -> str | None: - """Read optional diagnostic state without affecting calculation flow.""" + return state + def _clear_invalid_state(self, state_key: str, encoded: object) -> None: + if isinstance(encoded, bytes): + try: + encoded = encoded.decode("utf-8") + except UnicodeDecodeError: + return + if not isinstance(encoded, str): + return try: - value = self.client.get(self._observability_key(cache_key)) - except Exception: - return None - return value if isinstance(value, str) and value else None + self._claims.release(state_key, encoded) + except CacheCoordinationError: + return + + def claim_batch_start( + self, + cache_key: str, + claim_token: str, + observability_id: str | None, + ) -> bool: + state = BudgetWindowCacheState( + status="starting", + observability_id=observability_id, + submission_claim_id=claim_token, + ) + return self._claims.acquire( + self._state_key(cache_key), + self._encoded_state(state), + ttl_seconds=BUDGET_WINDOW_STARTING_TTL_SECONDS, + ) - def claim_batch_start(self, cache_key: str, claim_token: str) -> bool: + def clear_starting_claim( + self, + cache_key: str, + claim_token: str, + observability_id: str | None, + ) -> None: + state = BudgetWindowCacheState( + status="starting", + observability_id=observability_id, + submission_claim_id=claim_token, + ) try: - return self._claims.acquire( - self._batch_key(cache_key), - f"{BUDGET_WINDOW_STARTING_PREFIX}{claim_token}", - ttl_seconds=BUDGET_WINDOW_STARTING_TTL_SECONDS, + self._claims.release( + self._state_key(cache_key), + self._encoded_state(state), ) except CacheCoordinationError: - raise + return - def store_batch_job_id(self, cache_key: str, batch_job_id: str) -> None: + def store_submitted( + self, + cache_key: str, + batch_job_id: str, + observability_id: str | None, + ) -> None: + state = BudgetWindowCacheState( + status="submitted", + observability_id=observability_id, + batch_job_id=batch_job_id, + ) started_at = time.perf_counter() try: stored = self.client.set( - self._batch_key(cache_key), - batch_job_id, + self._state_key(cache_key), + self._encoded_state(state), ex=BUDGET_WINDOW_BATCH_TTL_SECONDS, ) except Exception as error: self._handle_cache_error( - "write-batch-id", + "write-submitted-state", event="coordination-failed", started_at=started_at, ) @@ -260,61 +296,95 @@ def store_batch_job_id(self, cache_key: str, batch_job_id: str) -> None: ) from error if not stored: self._handle_cache_error( - "write-batch-id", + "write-submitted-state", event="coordination-failed", started_at=started_at, ) raise CacheCoordinationError( - "budget-window batch identifier could not be stored" + "budget-window submitted state could not be stored" ) record_cache_event( family=BUDGET_WINDOW_CACHE_FAMILY, event="coordination-write", - operation="write-batch-id", + operation="write-submitted-state", started_at=started_at, ) - def store_observability_id( + def set_completed_result( self, cache_key: str, + result: dict[str, Any], observability_id: str | None, - ) -> None: - """Store diagnostic state best effort; exporter state cannot block work.""" + ) -> bool: + return self._set_recoverable_state( + cache_key, + BudgetWindowCacheState( + status="completed", + observability_id=observability_id, + result=result, + ), + operation="write-completed-state", + ) - if not observability_id: - return - try: - self.client.set( - self._observability_key(cache_key), - observability_id, - ex=BUDGET_WINDOW_BATCH_TTL_SECONDS, - ) - except Exception: - return + def set_terminal_error( + self, + cache_key: str, + error: dict[str, str], + observability_id: str | None, + ) -> bool: + return self._set_recoverable_state( + cache_key, + BudgetWindowCacheState( + status="failed", + observability_id=observability_id, + failure_type="spm_validation", + error=error, + ), + operation="write-spm-failure-state", + ) - def clear_starting_claim(self, cache_key: str, claim_token: str) -> None: - try: - self._claims.release( - self._batch_key(cache_key), - f"{BUDGET_WINDOW_STARTING_PREFIX}{claim_token}", - ) - except CacheCoordinationError: - return + def set_execution_failure( + self, + cache_key: str, + result: dict[str, Any], + observability_id: str | None, + ) -> bool: + return self._set_recoverable_state( + cache_key, + BudgetWindowCacheState( + status="failed", + observability_id=observability_id, + failure_type="execution", + error=result, + ), + operation="write-execution-failure-state", + ) - def clear_batch_job_id(self, cache_key: str) -> None: + def _set_recoverable_state( + self, + cache_key: str, + state: BudgetWindowCacheState, + *, + operation: str, + ) -> bool: started_at = time.perf_counter() try: - self.client.delete(self._batch_key(cache_key)) + stored = self.client.set( + self._state_key(cache_key), + self._encoded_state(state), + ex=jittered_ttl(BUDGET_WINDOW_RESULT_TTL_SECONDS), + ) except Exception: self._handle_cache_error( - "clear-batch-id", - event="coordination-failed", + operation, + event="write-failed", started_at=started_at, ) - return + return False record_cache_event( family=BUDGET_WINDOW_CACHE_FAMILY, - event="coordination-cleared", - operation="clear-batch-id", + event="write", + operation=operation, started_at=started_at, ) + return bool(stored) diff --git a/policyengine_api/services/economy_service.py b/policyengine_api/services/economy_service.py index 435441d67..e15782bea 100644 --- a/policyengine_api/services/economy_service.py +++ b/policyengine_api/services/economy_service.py @@ -414,12 +414,30 @@ def get_budget_window_economic_impact( ) cache_key = self._build_budget_window_cache_key(setup_options) - cached_error = self._budget_window_cache.get_terminal_error(cache_key) - if cached_error is not None: - raise SPMValidationError(**cached_error) + cached_state = self._budget_window_cache.get_state(cache_key) + if cached_state is not None: + self._adopt_budget_window_observability_id( + cached_state.observability_id, + setup_options=setup_options, + ) - cached_result = self._budget_window_cache.get_completed_result(cache_key) - if cached_result is not None: + if cached_state is not None and cached_state.status == "failed": + if cached_state.failure_type == "spm_validation": + cached_error = cached_state.error or {} + raise SPMValidationError( + code=str(cached_error.get("code", "SPM_VALIDATION_ERROR")), + message=str( + cached_error.get( + "message", "Stored budget-window validation failed" + ) + ), + ) + return BudgetWindowEconomicImpactResult.model_validate( + cached_state.error + ).model_copy(update={"cache_status": "failure-hit"}) + + if cached_state is not None and cached_state.status == "completed": + cached_result = cached_state.result or {} try: validate_worker_result( cached_result, @@ -434,7 +452,9 @@ def get_budget_window_economic_impact( # and the read above replays it, instead of re-deriving the # same failure from the same payload on every later poll. self._budget_window_cache.set_terminal_error( - cache_key, error.to_dict() + cache_key, + error.to_dict(), + setup_options.observability_id, ) raise return BudgetWindowEconomicImpactResult.completed( @@ -442,28 +462,34 @@ def get_budget_window_economic_impact( cache_status="result-hit", ) - batch_job_id = self._budget_window_cache.get_batch_job_id(cache_key) - if batch_job_id: - stored_observability_id = adopt_observability_id( - self._budget_window_cache.get_observability_id(cache_key) - ) - if stored_observability_id is not None: - setup_options.observability_id = stored_observability_id - observability_runtime.set_context( - observability_id=stored_observability_id - ) + if cached_state is not None and cached_state.status == "submitted": return self._get_budget_window_result_from_batch_job_id( - batch_job_id=batch_job_id, + batch_job_id=cached_state.batch_job_id or "", spm=setup_options.options.get("spm"), cache_key=cache_key, total_years=len(years), queued_years_on_submit=years, cache_status="batch-id-hit", + observability_id=setup_options.observability_id, + ) + + if cached_state is not None and cached_state.status == "starting": + return self._build_budget_window_computing_result( + total_years=len(years), + completed_years=[], + computing_years=[], + queued_years=years, + progress=0, + cache_status="starting-claim-hit", ) claim_token = setup_options.submission_claim_id cache_status = "starting-claim-hit" - if self._budget_window_cache.claim_batch_start(cache_key, claim_token): + if self._budget_window_cache.claim_batch_start( + cache_key, + claim_token, + setup_options.observability_id, + ): cache_status = "miss" try: batch_execution = self._start_budget_window_batch( @@ -472,33 +498,53 @@ def get_budget_window_economic_impact( window_size=window_size, max_parallel=max_active_years, ) - self._budget_window_cache.store_batch_job_id( - cache_key, batch_execution.batch_job_id + resolved_observability_id = ( + self._adopt_budget_window_observability_id( + batch_execution.observability_id, + setup_options=setup_options, + ) ) - self._budget_window_cache.store_observability_id( + self._budget_window_cache.store_submitted( cache_key, - batch_execution.observability_id - or setup_options.observability_id, + batch_execution.batch_job_id, + resolved_observability_id, ) except httpx.HTTPStatusError as error: - self._budget_window_cache.clear_starting_claim( - cache_key, claim_token - ) if ( error.response.status_code in BUDGET_WINDOW_SUBMISSION_VALIDATION_ERROR_STATUS_CODES ): - return BudgetWindowEconomicImpactResult.failed( + failed_result = BudgetWindowEconomicImpactResult.failed( self._build_budget_window_submission_error_message(error), queued_years=years, cache_status=cache_status, ) + self._budget_window_cache.set_execution_failure( + cache_key, + failed_result.model_dump(mode="json"), + setup_options.observability_id, + ) + return failed_result + self._budget_window_cache.clear_starting_claim( + cache_key, + claim_token, + setup_options.observability_id, + ) raise except Exception: self._budget_window_cache.clear_starting_claim( - cache_key, claim_token + cache_key, + claim_token, + setup_options.observability_id, ) raise + else: + claimed_state = self._budget_window_cache.get_state(cache_key) + if claimed_state is not None: + self._adopt_budget_window_observability_id( + claimed_state.observability_id, + setup_options=setup_options, + ) return self._build_budget_window_computing_result( total_years=len(years), @@ -527,6 +573,19 @@ def _build_budget_window_cache_key( api_version=setup_options.api_version, ) + @staticmethod + def _adopt_budget_window_observability_id( + value: str | None, + *, + setup_options: EconomicImpactSetupOptions, + ) -> str: + resolved = adopt_observability_id(value) + if resolved is None: + return setup_options.observability_id + setup_options.observability_id = resolved + observability_runtime.set_context(observability_id=resolved) + return resolved + def _build_budget_window_batch_payload( self, *, @@ -623,11 +682,21 @@ def _get_budget_window_result_from_batch_job_id( queued_years_on_submit: list[str], spm: dict | None = None, cache_status: Optional[str] = None, + observability_id: str, ) -> BudgetWindowEconomicImpactResult: + resolved_observability_id = observability_id try: batch_execution = self._simulation_gateway.get_budget_window_batch_by_id( batch_job_id ) + adopted_observability_id = adopt_observability_id( + batch_execution.observability_id + ) + if adopted_observability_id is not None: + resolved_observability_id = adopted_observability_id + observability_runtime.set_context( + observability_id=resolved_observability_id + ) if batch_execution.status in EXECUTION_STATUSES_SUCCESS: result = batch_execution.result if isinstance(result, dict) and result: @@ -635,26 +704,34 @@ def _get_budget_window_result_from_batch_job_id( result, spm, expected_years=queued_years_on_submit ) except SPMValidationError as error: - if self._budget_window_cache.set_terminal_error(cache_key, error.to_dict()): - self._budget_window_cache.clear_batch_job_id(cache_key) + self._budget_window_cache.set_terminal_error( + cache_key, + error.to_dict(), + resolved_observability_id, + ) raise if batch_execution.status in EXECUTION_STATUSES_SUCCESS: result = batch_execution.result if not isinstance(result, dict) or not result: - self._budget_window_cache.clear_batch_job_id(cache_key) - return BudgetWindowEconomicImpactResult.failed( + failed_result = BudgetWindowEconomicImpactResult.failed( "Budget-window batch completed without a result", completed_years=batch_execution.completed_years, computing_years=batch_execution.running_years, queued_years=batch_execution.queued_years or queued_years_on_submit, cache_status=cache_status, ) - result_stored = self._budget_window_cache.set_completed_result( - cache_key, result + self._budget_window_cache.set_execution_failure( + cache_key, + failed_result.model_dump(mode="json"), + resolved_observability_id, + ) + return failed_result + self._budget_window_cache.set_completed_result( + cache_key, + result, + resolved_observability_id, ) - if result_stored: - self._budget_window_cache.clear_batch_job_id(cache_key) return BudgetWindowEconomicImpactResult.completed( result, cache_status=cache_status, @@ -662,14 +739,19 @@ def _get_budget_window_result_from_batch_job_id( if batch_execution.status in EXECUTION_STATUSES_FAILURE: error_message = batch_execution.error or "Budget-window batch failed" - self._budget_window_cache.clear_batch_job_id(cache_key) - return BudgetWindowEconomicImpactResult.failed( + failed_result = BudgetWindowEconomicImpactResult.failed( error_message, completed_years=batch_execution.completed_years, computing_years=batch_execution.running_years, queued_years=batch_execution.queued_years or queued_years_on_submit, cache_status=cache_status, ) + self._budget_window_cache.set_execution_failure( + cache_key, + failed_result.model_dump(mode="json"), + resolved_observability_id, + ) + return failed_result if batch_execution.status in EXECUTION_STATUSES_PENDING: return self._build_budget_window_computing_result( diff --git a/tests/fixtures/services/economy_service.py b/tests/fixtures/services/economy_service.py index 35cb466ec..f2a4c4b51 100644 --- a/tests/fixtures/services/economy_service.py +++ b/tests/fixtures/services/economy_service.py @@ -160,14 +160,13 @@ def mock_budget_window_cache(): """Mock Redis-backed budget-window cache.""" mock_cache = MagicMock() mock_cache.build_key.return_value = "budget-window-cache-key" - mock_cache.get_terminal_error.return_value = None - mock_cache.get_completed_result.return_value = None - mock_cache.get_batch_job_id.return_value = None + mock_cache.get_state.return_value = None mock_cache.claim_batch_start.return_value = True - mock_cache.store_batch_job_id.return_value = None + mock_cache.store_submitted.return_value = None mock_cache.clear_starting_claim.return_value = None mock_cache.set_completed_result.return_value = True - mock_cache.clear_batch_job_id.return_value = None + mock_cache.set_terminal_error.return_value = True + mock_cache.set_execution_failure.return_value = True with patch( "policyengine_api.services.economy_service.budget_window_cache", diff --git a/tests/unit/routes/test_spm_year_worker_polling.py b/tests/unit/routes/test_spm_year_worker_polling.py index 797a57097..c79f023f4 100644 --- a/tests/unit/routes/test_spm_year_worker_polling.py +++ b/tests/unit/routes/test_spm_year_worker_polling.py @@ -156,7 +156,11 @@ def transport(request): ) if budget_window: cache_key = service._build_budget_window_cache_key(setup) - window_cache.store_batch_job_id(cache_key, job_id) + window_cache.store_submitted( + cache_key, + job_id, + setup.observability_id, + ) url = "/us/economy/1/over/2/budget-window" query = {"start_year": "2035", "window_size": "2"} else: @@ -193,9 +197,12 @@ def transport(request): "errors": [typed_error], } if budget_window: - assert window_cache.get_completed_result(cache_key) is None - assert window_cache.get_batch_job_id(cache_key) is None - assert window_cache.get_terminal_error(cache_key) == typed_error + state = window_cache.get_state(cache_key) + assert state is not None + assert state.status == "failed" + assert state.failure_type == "spm_validation" + assert state.error == typed_error + assert state.observability_id == setup.observability_id else: stored = annual_cache.get_by_execution_id(job_id) assert stored.status == "error" @@ -318,6 +325,7 @@ def transport(request): segmented_result("2036", "reform"), ], }, + setup.observability_id, ) url = "/us/economy/1/over/2/budget-window" query = {"start_year": "2035", "window_size": "2"} @@ -354,8 +362,12 @@ def transport(request): assert response.json == first if budget_window: - terminal = window_cache.get_terminal_error(cache_key) - assert terminal["code"] == "SPM_CONFIGURATION_UNAVAILABLE" + state = window_cache.get_state(cache_key) + assert state is not None + assert state.status == "failed" + assert state.failure_type == "spm_validation" + assert state.error is not None + assert state.error["code"] == "SPM_CONFIGURATION_UNAVAILABLE" else: stored = annual_cache.get_by_execution_id(job_id) assert stored.status == "error" diff --git a/tests/unit/services/test_budget_window_cache.py b/tests/unit/services/test_budget_window_cache.py index 0bb51ee9d..823205959 100644 --- a/tests/unit/services/test_budget_window_cache.py +++ b/tests/unit/services/test_budget_window_cache.py @@ -64,38 +64,51 @@ def test_build_key_is_stable_for_request_identity(): ) assert first == second - assert first.startswith("policyengine:test:api:budget-window:v1:") + assert first.startswith("policyengine:test:api:budget-window:v2:") -def test_claim_batch_start_allows_one_starter(): +def test_claim_batch_start_allows_one_starter_and_preserves_identity(): cache = BudgetWindowCache(client=FakeRedis()) + cache_key = "budget_window:v2:us:key" - assert cache.claim_batch_start("budget_window:v1:us:key", "process-1") is True - assert cache.claim_batch_start("budget_window:v1:us:key", "process-2") is False - assert cache.get_batch_job_id("budget_window:v1:us:key") is None + assert cache.claim_batch_start(cache_key, "claim-1", "obs-1") is True + assert cache.claim_batch_start(cache_key, "claim-2", "obs-2") is False + state = cache.get_state(cache_key) + assert state is not None + assert state.status == "starting" + assert state.submission_claim_id == "claim-1" + assert state.observability_id == "obs-1" -def test_store_batch_job_id_replaces_starting_claim(): + +def test_store_submitted_replaces_starting_state(): cache = BudgetWindowCache(client=FakeRedis()) - cache.claim_batch_start("budget_window:v1:us:key", "process-1") + cache_key = "budget_window:v2:us:key" + cache.claim_batch_start(cache_key, "claim-1", "obs-1") - cache.store_batch_job_id("budget_window:v1:us:key", "fc-parent") + cache.store_submitted(cache_key, "fc-parent", "obs-1") - assert cache.get_batch_job_id("budget_window:v1:us:key") == "fc-parent" + state = cache.get_state(cache_key) + assert state is not None + assert state.status == "submitted" + assert state.batch_job_id == "fc-parent" + assert state.observability_id == "obs-1" -def test_completed_result_round_trips(): +def test_completed_result_round_trips_with_identity(): cache = BudgetWindowCache(client=FakeRedis()) result = {"kind": "budgetWindow", "totals": {"budgetaryImpact": 10}} - cache.set_completed_result("budget_window:v1:us:key", result) + cache.set_completed_result("budget_window:v2:us:key", result, "obs-1") - assert cache.get_completed_result("budget_window:v1:us:key") == result + state = cache.get_state("budget_window:v2:us:key") + assert state is not None + assert state.status == "completed" + assert state.result == result + assert state.observability_id == "obs-1" -def test_terminal_error_round_trips_separately_from_success_and_other_selections( - monkeypatch, -): +def test_spm_validation_failure_round_trips_for_only_its_selection(monkeypatch): import policyengine_api.services.budget_window_cache as module monkeypatch.setattr(module, "jittered_ttl", lambda _ttl: 123) @@ -113,17 +126,36 @@ def test_terminal_error_round_trips_separately_from_success_and_other_selections } failed_key = cache.build_key(**identity, options_hash="canonical-selection-a") other_key = cache.build_key(**identity, options_hash="canonical-selection-b") - assert cache.set_terminal_error(failed_key, error) - cache = BudgetWindowCache(client=backend) - assert cache.get_terminal_error(failed_key) == error - assert cache.get_terminal_error(other_key) is None - assert cache.get_completed_result(failed_key) is None + + assert cache.set_terminal_error(failed_key, error, "obs-1") + + stored = BudgetWindowCache(client=backend).get_state(failed_key) + assert stored is not None + assert stored.status == "failed" + assert stored.failure_type == "spm_validation" + assert stored.error == error + assert stored.observability_id == "obs-1" + assert cache.get_state(other_key) is None assert set(backend._expires.values()) == {123} backend.advance(123) - assert cache.get_terminal_error(failed_key) is None + assert cache.get_state(failed_key) is None + + +def test_execution_failure_round_trips_with_identity(): + cache = BudgetWindowCache(client=FakeRedis()) + result = {"status": "error", "error": "simulation failed"} + + assert cache.set_execution_failure("budget_window:v2:us:key", result, "obs-1") + + state = cache.get_state("budget_window:v2:us:key") + assert state is not None + assert state.status == "failed" + assert state.failure_type == "execution" + assert state.error == result + assert state.observability_id == "obs-1" -def test_completed_result_ttl_is_jittered_but_coordination_ttls_are_exact( +def test_recoverable_state_ttl_is_jittered_but_coordination_ttls_are_exact( monkeypatch, ): import policyengine_api.services.budget_window_cache as module @@ -131,166 +163,103 @@ def test_completed_result_ttl_is_jittered_but_coordination_ttls_are_exact( monkeypatch.setattr(module, "jittered_ttl", lambda _ttl: 123) redis_client = FakeRedis() cache = BudgetWindowCache(client=redis_client) - cache_key = "budget_window:v1:us:key" + cache_key = "budget_window:v2:us:key" + state_key = f"{cache_key}:state" - assert cache.set_completed_result(cache_key, {"ok": True}) - assert redis_client._expires[f"{cache_key}:result"] == 123 + assert cache.set_completed_result(cache_key, {"ok": True}, "obs-1") + assert redis_client._expires[state_key] == 123 - assert cache.claim_batch_start(cache_key, "process-1") - assert ( - redis_client._expires[f"{cache_key}:batch-job-id"] - == BUDGET_WINDOW_STARTING_TTL_SECONDS - ) + redis_client.delete(state_key) + assert cache.claim_batch_start(cache_key, "claim-1", "obs-1") + assert redis_client._expires[state_key] == BUDGET_WINDOW_STARTING_TTL_SECONDS - cache.store_batch_job_id(cache_key, "batch-1") - assert ( - redis_client._expires[f"{cache_key}:batch-job-id"] - == BUDGET_WINDOW_BATCH_TTL_SECONDS - ) + cache.store_submitted(cache_key, "batch-1", "obs-1") + assert redis_client._expires[state_key] == BUDGET_WINDOW_BATCH_TTL_SECONDS -def test_get_completed_result_returns_none_for_empty_payload(): - redis_client = FakeRedis() - redis_client.values["budget_window:v1:us:key:result"] = "" - cache = BudgetWindowCache(client=redis_client) - - assert cache.get_completed_result("budget_window:v1:us:key") is None - - -def test_get_completed_result_returns_none_for_invalid_json(monkeypatch): +@pytest.mark.parametrize("invalid_value", ["", "{not-json", "123"]) +def test_get_state_removes_invalid_payload(invalid_value, monkeypatch): mock_logger = MagicMock() - monkeypatch.setattr( - "policyengine_api.runtime_cache.core.logger", - mock_logger, - ) + monkeypatch.setattr("policyengine_api.runtime_cache.core.logger", mock_logger) redis_client = FakeRedis() - redis_client.values["budget_window:v1:us:key:result"] = "{not-json" + state_key = "budget_window:v2:us:key:state" + redis_client.values[state_key] = invalid_value cache = BudgetWindowCache(client=redis_client) - assert cache.get_completed_result("budget_window:v1:us:key") is None - assert mock_logger.log_struct.call_args.kwargs["severity"] == "WARNING" + assert cache.get_state("budget_window:v2:us:key") is None + assert state_key not in redis_client.values + assert any( + call.kwargs.get("severity") == "WARNING" + for call in mock_logger.log_struct.call_args_list + ) -def test_get_completed_result_treats_read_errors_as_misses(monkeypatch): +def test_get_state_reraises_read_errors(monkeypatch): mock_logger = MagicMock() - monkeypatch.setattr( - "policyengine_api.runtime_cache.core.logger", - mock_logger, - ) + monkeypatch.setattr("policyengine_api.runtime_cache.core.logger", mock_logger) cache = BudgetWindowCache(client=RaisingRedis(method="get")) - assert cache.get_completed_result("budget_window:v1:us:key") is None + with pytest.raises(CacheCoordinationError): + cache.get_state("budget_window:v2:us:key") assert mock_logger.log_struct.call_args.kwargs["severity"] == "WARNING" -def test_set_completed_result_does_not_invalidate_compute_on_write_error(monkeypatch): +def test_completed_result_write_error_does_not_change_returned_result(monkeypatch): mock_logger = MagicMock() - monkeypatch.setattr( - "policyengine_api.runtime_cache.core.logger", - mock_logger, - ) + monkeypatch.setattr("policyengine_api.runtime_cache.core.logger", mock_logger) cache = BudgetWindowCache(client=RaisingRedis(method="set")) - assert not cache.set_completed_result("budget_window:v1:us:key", {"ok": True}) - - assert mock_logger.log_struct.call_args.kwargs["severity"] == "WARNING" - - -def test_get_batch_job_id_ignores_empty_non_string_and_starting_values(): - redis_client = FakeRedis() - cache = BudgetWindowCache(client=redis_client) - - redis_client.values["budget_window:v1:us:key:batch-job-id"] = "" - assert cache.get_batch_job_id("budget_window:v1:us:key") is None - - redis_client.values["budget_window:v1:us:key:batch-job-id"] = 123 - assert cache.get_batch_job_id("budget_window:v1:us:key") is None - - redis_client.values["budget_window:v1:us:key:batch-job-id"] = "starting:process-1" - assert cache.get_batch_job_id("budget_window:v1:us:key") is None - - -def test_get_batch_job_id_reraises_read_errors(monkeypatch): - mock_logger = MagicMock() - monkeypatch.setattr( - "policyengine_api.runtime_cache.core.logger", - mock_logger, + assert not cache.set_completed_result( + "budget_window:v2:us:key", {"ok": True}, "obs-1" ) - cache = BudgetWindowCache(client=RaisingRedis(method="get")) - - with pytest.raises(CacheCoordinationError): - cache.get_batch_job_id("budget_window:v1:us:key") - assert mock_logger.log_struct.call_args.kwargs["severity"] == "WARNING" def test_claim_batch_start_reraises_claim_errors(monkeypatch): mock_logger = MagicMock() - monkeypatch.setattr( - "policyengine_api.runtime_cache.core.logger", - mock_logger, - ) + monkeypatch.setattr("policyengine_api.runtime_cache.core.logger", mock_logger) cache = BudgetWindowCache(client=RaisingRedis(method="set")) with pytest.raises(CacheCoordinationError): - cache.claim_batch_start("budget_window:v1:us:key", "process-1") + cache.claim_batch_start("budget_window:v2:us:key", "claim-1", "obs-1") assert mock_logger.log_struct.call_args.kwargs["severity"] == "WARNING" -def test_store_batch_job_id_reraises_write_errors(monkeypatch): +def test_store_submitted_reraises_write_errors(monkeypatch): mock_logger = MagicMock() - monkeypatch.setattr( - "policyengine_api.runtime_cache.core.logger", - mock_logger, - ) + monkeypatch.setattr("policyengine_api.runtime_cache.core.logger", mock_logger) cache = BudgetWindowCache(client=RaisingRedis(method="set")) with pytest.raises(CacheCoordinationError): - cache.store_batch_job_id("budget_window:v1:us:key", "fc-parent") + cache.store_submitted("budget_window:v2:us:key", "fc-parent", "obs-1") assert mock_logger.log_struct.call_args.kwargs["severity"] == "WARNING" -def test_clear_starting_claim_deletes_only_matching_token(): +def test_clear_starting_claim_deletes_only_matching_document(): redis_client = FakeRedis() cache = BudgetWindowCache(client=redis_client) - cache.claim_batch_start("budget_window:v1:us:key", "process-1") + cache_key = "budget_window:v2:us:key" + state_key = f"{cache_key}:state" + cache.claim_batch_start(cache_key, "claim-1", "obs-1") - cache.clear_starting_claim("budget_window:v1:us:key", "process-2") - - assert ( - redis_client.values["budget_window:v1:us:key:batch-job-id"] - == "starting:process-1" - ) + cache.clear_starting_claim(cache_key, "claim-2", "obs-1") + assert state_key in redis_client.values - cache.clear_starting_claim("budget_window:v1:us:key", "process-1") + cache.clear_starting_claim(cache_key, "claim-1", "different-observability-id") + assert state_key in redis_client.values - assert "budget_window:v1:us:key:batch-job-id" not in redis_client.values + cache.clear_starting_claim(cache_key, "claim-1", "obs-1") + assert state_key not in redis_client.values -def test_clear_starting_claim_logs_and_swallows_errors(monkeypatch): +def test_clear_starting_claim_swallows_coordination_errors(monkeypatch): mock_logger = MagicMock() - monkeypatch.setattr( - "policyengine_api.runtime_cache.core.logger", - mock_logger, - ) - cache = BudgetWindowCache(client=RaisingRedis(method="get")) - - cache.clear_starting_claim("budget_window:v1:us:key", "process-1") - - assert mock_logger.log_struct.call_args.kwargs["severity"] == "WARNING" - - -def test_clear_batch_job_id_logs_and_swallows_errors(monkeypatch): - mock_logger = MagicMock() - monkeypatch.setattr( - "policyengine_api.runtime_cache.core.logger", - mock_logger, - ) - cache = BudgetWindowCache(client=RaisingRedis(method="delete")) + monkeypatch.setattr("policyengine_api.runtime_cache.core.logger", mock_logger) + cache = BudgetWindowCache(client=RaisingRedis(method="eval")) - cache.clear_batch_job_id("budget_window:v1:us:key") + cache.clear_starting_claim("budget_window:v2:us:key", "claim-1", "obs-1") assert mock_logger.log_struct.call_args.kwargs["severity"] == "WARNING" diff --git a/tests/unit/services/test_economy_service.py b/tests/unit/services/test_economy_service.py index 0e65d9ee1..4b735b72c 100644 --- a/tests/unit/services/test_economy_service.py +++ b/tests/unit/services/test_economy_service.py @@ -5,6 +5,7 @@ import httpx import pytest from policyengine_api.runtime_cache.core import CacheCoordinationError +from policyengine_api.services.budget_window_cache import BudgetWindowCacheState from policyengine_api.services.reform_impacts_service import ( ReformImpactHandoffError, ) @@ -990,10 +991,14 @@ def test__given_no_cached_batch__submits_parent_batch_and_returns_queued_result( assert submitted_payload["target"] == "general" assert "time_period" not in submitted_payload mock_budget_window_cache.claim_batch_start.assert_called_once_with( - "budget-window-cache-key", MOCK_SUBMISSION_CLAIM_ID + "budget-window-cache-key", + MOCK_SUBMISSION_CLAIM_ID, + MOCK_OBSERVABILITY_ID, ) - mock_budget_window_cache.store_batch_job_id.assert_called_once_with( - "budget-window-cache-key", "fc-budget-123" + mock_budget_window_cache.store_submitted.assert_called_once_with( + "budget-window-cache-key", + "fc-budget-123", + MOCK_OBSERVABILITY_ID, ) mock_reform_impacts_service.set_reform_impact.assert_not_called() @@ -1028,8 +1033,10 @@ def test__given_completed_cached_result__returns_completed_batch_result( "budgetaryImpact": 90, }, } - mock_budget_window_cache.get_completed_result.return_value = ( - completed_result + mock_budget_window_cache.get_state.return_value = BudgetWindowCacheState( + status="completed", + result=completed_result, + observability_id=MOCK_OBSERVABILITY_ID, ) result = economy_service.get_budget_window_economic_impact(**base_params) @@ -1048,7 +1055,11 @@ def test__given_cached_batch_id__returns_running_batch_progress( mock_simulation_entrypoint, mock_budget_window_cache, ): - mock_budget_window_cache.get_batch_job_id.return_value = "fc-budget-123" + mock_budget_window_cache.get_state.return_value = BudgetWindowCacheState( + status="submitted", + batch_job_id="fc-budget-123", + observability_id=MOCK_OBSERVABILITY_ID, + ) mock_simulation_entrypoint.get_budget_window_batch_by_id.return_value = ( create_mock_budget_window_batch_execution( batch_job_id="fc-budget-123", @@ -1088,7 +1099,11 @@ def test__given_completed_batch_poll__caches_result_and_returns_completed( "annualImpacts": [], "totals": {}, } - mock_budget_window_cache.get_batch_job_id.return_value = "fc-budget-123" + mock_budget_window_cache.get_state.return_value = BudgetWindowCacheState( + status="submitted", + batch_job_id="fc-budget-123", + observability_id=MOCK_OBSERVABILITY_ID, + ) mock_simulation_entrypoint.get_budget_window_batch_by_id.return_value = ( create_mock_budget_window_batch_execution( batch_job_id="fc-budget-123", @@ -1105,10 +1120,9 @@ def test__given_completed_batch_poll__caches_result_and_returns_completed( assert result.data == completed_result assert result.cache_status == "batch-id-hit" mock_budget_window_cache.set_completed_result.assert_called_once_with( - "budget-window-cache-key", completed_result - ) - mock_budget_window_cache.clear_batch_job_id.assert_called_once_with( - "budget-window-cache-key" + "budget-window-cache-key", + completed_result, + MOCK_OBSERVABILITY_ID, ) @pytest.mark.parametrize("malformed_result", [None, {}, []]) @@ -1120,7 +1134,11 @@ def test__given_completed_batch_without_result__returns_error_without_caching( mock_budget_window_cache, malformed_result, ): - mock_budget_window_cache.get_batch_job_id.return_value = "fc-budget-123" + mock_budget_window_cache.get_state.return_value = BudgetWindowCacheState( + status="submitted", + batch_job_id="fc-budget-123", + observability_id=MOCK_OBSERVABILITY_ID, + ) mock_simulation_entrypoint.get_budget_window_batch_by_id.return_value = ( create_mock_budget_window_batch_execution( batch_job_id="fc-budget-123", @@ -1142,9 +1160,7 @@ def test__given_completed_batch_without_result__returns_error_without_caching( assert result.queued_years == ["2028"] assert result.cache_status == "batch-id-hit" mock_budget_window_cache.set_completed_result.assert_not_called() - mock_budget_window_cache.clear_batch_job_id.assert_called_once_with( - "budget-window-cache-key" - ) + mock_budget_window_cache.set_execution_failure.assert_called_once() def test__given_completed_batch_cache_write_fails__does_not_clear_batch_id( self, @@ -1161,7 +1177,11 @@ def test__given_completed_batch_cache_write_fails__does_not_clear_batch_id( "annualImpacts": [], "totals": {}, } - mock_budget_window_cache.get_batch_job_id.return_value = "fc-budget-123" + mock_budget_window_cache.get_state.return_value = BudgetWindowCacheState( + status="submitted", + batch_job_id="fc-budget-123", + observability_id=MOCK_OBSERVABILITY_ID, + ) mock_budget_window_cache.set_completed_result.return_value = False mock_simulation_entrypoint.get_budget_window_batch_by_id.return_value = ( create_mock_budget_window_batch_execution( @@ -1177,7 +1197,11 @@ def test__given_completed_batch_cache_write_fails__does_not_clear_batch_id( assert result.status == ImpactStatus.OK assert result.data == completed_result - mock_budget_window_cache.clear_batch_job_id.assert_not_called() + mock_budget_window_cache.set_completed_result.assert_called_once_with( + "budget-window-cache-key", + completed_result, + MOCK_OBSERVABILITY_ID, + ) def test__given_failed_batch_poll__returns_failed( self, @@ -1186,7 +1210,11 @@ def test__given_failed_batch_poll__returns_failed( mock_simulation_entrypoint, mock_budget_window_cache, ): - mock_budget_window_cache.get_batch_job_id.return_value = "fc-budget-123" + mock_budget_window_cache.get_state.return_value = BudgetWindowCacheState( + status="submitted", + batch_job_id="fc-budget-123", + observability_id=MOCK_OBSERVABILITY_ID, + ) mock_simulation_entrypoint.get_budget_window_batch_by_id.return_value = ( create_mock_budget_window_batch_execution( batch_job_id="fc-budget-123", @@ -1208,9 +1236,7 @@ def test__given_failed_batch_poll__returns_failed( assert result.queued_years == ["2028"] assert result.cache_status == "batch-id-hit" mock_budget_window_cache.set_completed_result.assert_not_called() - mock_budget_window_cache.clear_batch_job_id.assert_called_once_with( - "budget-window-cache-key" - ) + mock_budget_window_cache.set_execution_failure.assert_called_once() def test_typed_error_write_failure_retains_batch_identity( self, @@ -1220,7 +1246,11 @@ def test_typed_error_write_failure_retains_batch_identity( mock_budget_window_cache, ): error = SPMValidationError("SPM_YEAR_UNAVAILABLE", "No forecast for 2036") - mock_budget_window_cache.get_batch_job_id.return_value = "expired-job" + mock_budget_window_cache.get_state.return_value = BudgetWindowCacheState( + status="submitted", + batch_job_id="expired-job", + observability_id=MOCK_OBSERVABILITY_ID, + ) mock_budget_window_cache.set_terminal_error.return_value = False mock_simulation_entrypoint.get_budget_window_batch_by_id.side_effect = error @@ -1228,7 +1258,6 @@ def test_typed_error_write_failure_retains_batch_identity( economy_service.get_budget_window_economic_impact(**base_params) assert raised.value is error - mock_budget_window_cache.clear_batch_job_id.assert_not_called() mock_budget_window_cache.set_completed_result.assert_not_called() mock_simulation_entrypoint.run_budget_window_batch.assert_not_called() @@ -1240,7 +1269,11 @@ def test_untyped_poll_error_keeps_existing_retry_behavior( mock_budget_window_cache, ): error = make_http_status_error(422, payload={"detail": "Unknown error"}) - mock_budget_window_cache.get_batch_job_id.return_value = "existing-job" + mock_budget_window_cache.get_state.return_value = BudgetWindowCacheState( + status="submitted", + batch_job_id="existing-job", + observability_id=MOCK_OBSERVABILITY_ID, + ) mock_simulation_entrypoint.get_budget_window_batch_by_id.side_effect = error with pytest.raises(httpx.HTTPStatusError) as raised: @@ -1248,7 +1281,6 @@ def test_untyped_poll_error_keeps_existing_retry_behavior( assert raised.value is error mock_budget_window_cache.set_terminal_error.assert_not_called() - mock_budget_window_cache.clear_batch_job_id.assert_not_called() mock_simulation_entrypoint.run_budget_window_batch.assert_not_called() def test__given_existing_start_claim__does_not_submit_duplicate_batch( @@ -1258,14 +1290,75 @@ def test__given_existing_start_claim__does_not_submit_duplicate_batch( mock_simulation_entrypoint, mock_budget_window_cache, ): + winning_observability_id = "00000000-0000-4000-8000-000000000099" mock_budget_window_cache.claim_batch_start.return_value = False + mock_budget_window_cache.get_state.side_effect = [ + None, + BudgetWindowCacheState( + status="starting", + submission_claim_id="winning-claim", + observability_id=winning_observability_id, + ), + ] - result = economy_service.get_budget_window_economic_impact(**base_params) + with ( + patch( + "policyengine_api.services.economy_service.observability_runtime.set_context" + ) as set_context, + patch( + "policyengine_api.services.economy_service.adopt_observability_id", + return_value=winning_observability_id, + ), + ): + result = economy_service.get_budget_window_economic_impact( + **base_params + ) assert result.status == ImpactStatus.COMPUTING assert result.progress == 0 assert result.queued_years == ["2026", "2027", "2028"] assert result.cache_status == "starting-claim-hit" + set_context.assert_any_call(observability_id=winning_observability_id) + mock_simulation_entrypoint.run_budget_window_batch.assert_not_called() + + def test__given_cached_execution_failure__replays_failure_and_identity( + self, + economy_service, + base_params, + mock_simulation_entrypoint, + mock_budget_window_cache, + ): + stored_observability_id = "00000000-0000-4000-8000-000000000099" + mock_budget_window_cache.get_state.return_value = BudgetWindowCacheState( + status="failed", + observability_id=stored_observability_id, + failure_type="execution", + error={ + "status": "error", + "error": "Budget window failed for 2027", + "completed_years": ["2026"], + "queued_years": ["2028"], + }, + ) + + with ( + patch( + "policyengine_api.services.economy_service.observability_runtime.set_context" + ) as set_context, + patch( + "policyengine_api.services.economy_service.adopt_observability_id", + return_value=stored_observability_id, + ), + ): + result = economy_service.get_budget_window_economic_impact( + **base_params + ) + + assert result.status == ImpactStatus.ERROR + assert result.error == "Budget window failed for 2027" + assert result.cache_status == "failure-hit" + set_context.assert_any_call(observability_id=stored_observability_id) + mock_simulation_entrypoint.get_budget_window_batch_by_id.assert_not_called() mock_simulation_entrypoint.run_budget_window_batch.assert_not_called() def test__given_gateway_raises_before_returning_batch__clears_start_claim( @@ -1283,7 +1376,9 @@ def test__given_gateway_raises_before_returning_batch__clears_start_claim( economy_service.get_budget_window_economic_impact(**base_params) mock_budget_window_cache.clear_starting_claim.assert_called_once_with( - "budget-window-cache-key", MOCK_SUBMISSION_CLAIM_ID + "budget-window-cache-key", + MOCK_SUBMISSION_CLAIM_ID, + MOCK_OBSERVABILITY_ID, ) @pytest.mark.parametrize("status_code", [400, 422]) @@ -1317,10 +1412,8 @@ def test__given_modal_rejects_batch_submission_for_validation__returns_failed_re assert result.computing_years == [] assert result.queued_years == ["2026", "2027", "2028"] assert result.cache_status == "miss" - mock_budget_window_cache.clear_starting_claim.assert_called_once_with( - "budget-window-cache-key", MOCK_SUBMISSION_CLAIM_ID - ) - mock_budget_window_cache.store_batch_job_id.assert_not_called() + mock_budget_window_cache.clear_starting_claim.assert_not_called() + mock_budget_window_cache.set_execution_failure.assert_called_once() @pytest.mark.parametrize("status_code", [401, 403, 429, 500]) def test__given_modal_non_validation_error_on_batch_submission__raises( @@ -1339,9 +1432,11 @@ def test__given_modal_non_validation_error_on_batch_submission__raises( economy_service.get_budget_window_economic_impact(**base_params) mock_budget_window_cache.clear_starting_claim.assert_called_once_with( - "budget-window-cache-key", MOCK_SUBMISSION_CLAIM_ID + "budget-window-cache-key", + MOCK_SUBMISSION_CLAIM_ID, + MOCK_OBSERVABILITY_ID, ) - mock_budget_window_cache.store_batch_job_id.assert_not_called() + mock_budget_window_cache.store_submitted.assert_not_called() @pytest.mark.parametrize( ("payload", "expected_message"), @@ -1451,9 +1546,11 @@ def test__given_reordered_options__uses_same_budget_window_cache_identity( mock_simulation_entrypoint, mock_budget_window_cache, ): - mock_budget_window_cache.get_completed_result.return_value = { - "kind": "budgetWindow" - } + mock_budget_window_cache.get_state.return_value = BudgetWindowCacheState( + status="completed", + result={"kind": "budgetWindow"}, + observability_id=MOCK_OBSERVABILITY_ID, + ) economy_service.get_budget_window_economic_impact( **{ @@ -1504,7 +1601,11 @@ def test__given_unexpected_batch_status__raises_value_error( mock_simulation_entrypoint, mock_budget_window_cache, ): - mock_budget_window_cache.get_batch_job_id.return_value = "fc-budget-123" + mock_budget_window_cache.get_state.return_value = BudgetWindowCacheState( + status="submitted", + batch_job_id="fc-budget-123", + observability_id=MOCK_OBSERVABILITY_ID, + ) mock_simulation_entrypoint.get_budget_window_batch_by_id.return_value = ( create_mock_budget_window_batch_execution( batch_job_id="fc-budget-123",