Merge pull request #437 from PolicyEngine/pr/observability-contracts-and-telemetry-v2

Add observability scaffolding and telemetry propagation to the simulation gateway

Author: anth-volk

libs/policyengine-fastapi/fixtures/__init__.py

@@ -0,0 +1 @@
+"""Fixture modules for policyengine-fastapi tests."""

libs/policyengine-fastapi/fixtures/ping/__init__.py

@@ -0,0 +1 @@
+"""Fixture modules for ping tests."""

…icyengine-fastapi/tests/ping/conftest.py …cyengine-fastapi/fixtures/ping/shared.pylibs/policyengine-fastapi/tests/ping/conftest.py renamed to libs/policyengine-fastapi/fixtures/ping/shared.py libs/policyengine-fastapi/tests/ping/conftest.py renamed to libs/policyengine-fastapi/fixtures/ping/shared.py

@@ -1,7 +1,9 @@
+"""Shared fixtures for ping endpoint tests."""
+
+from fastapi import FastAPI
 from fastapi.testclient import TestClient
 import pytest
 
-from fastapi import FastAPI
 from policyengine_fastapi import ping
 from policyengine_fastapi.health import HealthRegistry
 

libs/policyengine-fastapi/src/policyengine_fastapi/__init__.py

@@ -1,3 +1,11 @@
 """
 Re-usable elements that support FastAPI and SQLModel indipendent of this specific app.
 """
+
+from .observability import (
+    ObservabilityConfig as ObservabilityConfig,
+    SimulationStage as SimulationStage,
+    TracerCaptureMode as TracerCaptureMode,
+    build_observability as build_observability,
+    get_observability as get_observability,
+)

libs/policyengine-fastapi/src/policyengine_fastapi/observability/__init__.py

@@ -0,0 +1,39 @@
+"""
+Shared observability contracts and utilities for API v2 services.
+
+This stays in `policyengine-fastapi` for now because only one active service
+needs it and the main value today is shared contract stability, not a separate
+distribution boundary. Extract it to a dedicated package only when a second
+service has materially different runtime or release needs.
+"""
+
+from .config import (
+    ObservabilityConfig as ObservabilityConfig,
+    parse_header_value_pairs as parse_header_value_pairs,
+)
+from .contracts import (
+    SimulationCompositeTraceResponse as SimulationCompositeTraceResponse,
+    SimulationLifecycleEvent as SimulationLifecycleEvent,
+    SimulationRunSummary as SimulationRunSummary,
+    SimulationTelemetryEnvelope as SimulationTelemetryEnvelope,
+    SimulationTimelineEntry as SimulationTimelineEntry,
+    TracerArtifactManifest as TracerArtifactManifest,
+    VersionStageMetricResponse as VersionStageMetricResponse,
+)
+from .correlation import (
+    generate_run_id as generate_run_id,
+    stable_config_hash as stable_config_hash,
+)
+from .emitters import (
+    Observability as Observability,
+    NoOpObservability as NoOpObservability,
+    NoOpSpan as NoOpSpan,
+)
+from .provider import (
+    build_observability as build_observability,
+    get_observability as get_observability,
+)
+from .stages import (
+    SimulationStage as SimulationStage,
+    TracerCaptureMode as TracerCaptureMode,
+)

libs/policyengine-fastapi/src/policyengine_fastapi/observability/config.py

@@ -0,0 +1,48 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+from .stages import TracerCaptureMode
+
+
+def parse_header_value_pairs(raw: str | None) -> dict[str, str]:
+    """Parse OTLP headers from a comma or newline separated key=value string."""
+
+    if raw is None:
+        return {}
+
+    stripped = raw.strip()
+    if not stripped:
+        return {}
+
+    headers: dict[str, str] = {}
+    for pair in stripped.replace("\n", ",").split(","):
+        candidate = pair.strip()
+        if not candidate:
+            continue
+        key, separator, value = candidate.partition("=")
+        if not separator:
+            raise ValueError(
+                "Expected OTLP headers in key=value format separated by commas"
+            )
+        headers[key.strip()] = value.strip()
+
+    return headers
+
+
+@dataclass(frozen=True)
+class ObservabilityConfig:
+    enabled: bool = False
+    shadow_mode: bool = True
+    service_name: str = "policyengine-observability"
+    environment: str = "production"
+    otlp_endpoint: str | None = None
+    otlp_headers: dict[str, str] = field(default_factory=dict)
+    artifact_bucket: str | None = None
+    artifact_prefix: str = "simulation-observability"
+    tracer_capture_mode: TracerCaptureMode = TracerCaptureMode.DISABLED
+    slow_run_threshold_seconds: float = 30.0
+
+    @classmethod
+    def disabled(cls, service_name: str = "policyengine-observability"):
+        return cls(enabled=False, service_name=service_name)

libs/policyengine-fastapi/src/policyengine_fastapi/observability/contracts.py

@@ -0,0 +1,110 @@
+from __future__ import annotations
+
+from datetime import datetime
+from typing import Any
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from .stages import SimulationStage, TracerCaptureMode
+
+
+class ObservabilityModel(BaseModel):
+    model_config = ConfigDict(extra="forbid", use_enum_values=True)
+
+
+class CorrelatedRunFields(ObservabilityModel):
+    run_id: str
+    process_id: str | None = None
+    job_id: str | None = None
+    trace_id: str | None = None
+    request_id: str | None = None
+    country: str | None = None
+    simulation_kind: str | None = None
+    geography_code: str | None = None
+    geography_type: str | None = None
+    country_package_name: str | None = None
+    country_package_version: str | None = None
+    policyengine_version: str | None = None
+    data_version: str | None = None
+    modal_app_name: str | None = None
+    config_hash: str | None = None
+
+
+class SimulationLifecycleEvent(CorrelatedRunFields):
+    event_name: str
+    stage: SimulationStage
+    status: str
+    timestamp: datetime
+    service: str
+    duration_seconds: float | None = None
+    details: dict[str, Any] = Field(default_factory=dict)
+
+
+class TracerArtifactManifest(CorrelatedRunFields):
+    scenario: str
+    capture_mode: TracerCaptureMode
+    artifact_format: str
+    storage_uri: str
+    summary_uri: str | None = None
+    node_count: int = 0
+    root_count: int = 0
+    max_depth: int = 0
+    total_calculation_time_seconds: float = 0.0
+    total_formula_time_seconds: float = 0.0
+    generated_at: datetime
+
+
+class SimulationTelemetryEnvelope(ObservabilityModel):
+    run_id: str
+    process_id: str | None = None
+    request_id: str | None = None
+    traceparent: str | None = None
+    requested_at: datetime | None = None
+    simulation_kind: str | None = None
+    geography_code: str | None = None
+    geography_type: str | None = None
+    config_hash: str | None = None
+    capture_mode: TracerCaptureMode = TracerCaptureMode.DISABLED
+
+
+class SimulationRunSummary(CorrelatedRunFields):
+    status: str
+    requested_at: datetime | None = None
+    returned_at: datetime | None = None
+    total_duration_seconds: float | None = None
+
+
+class SimulationTimelineEntry(ObservabilityModel):
+    stage: SimulationStage
+    started_at: datetime
+    ended_at: datetime
+    duration_seconds: float
+    service: str
+
+
+class SimulationCompositeTraceResponse(ObservabilityModel):
+    run: SimulationRunSummary
+    timeline: list[SimulationTimelineEntry] = Field(default_factory=list)
+    spans: dict[str, Any] = Field(default_factory=dict)
+    logs: dict[str, Any] = Field(default_factory=dict)
+    tracer: dict[str, Any] = Field(default_factory=dict)
+
+
+class StageRuntimeSummary(ObservabilityModel):
+    mean_seconds: float | None = None
+    p50_seconds: float | None = None
+    p95_seconds: float | None = None
+
+
+class VersionStageMetrics(ObservabilityModel):
+    country_package_version: str
+    launch_source: str = "modal_version_registry"
+    launched: bool = True
+    observed_run_count: int = 0
+    stages: dict[str, StageRuntimeSummary] = Field(default_factory=dict)
+
+
+class VersionStageMetricResponse(ObservabilityModel):
+    country: str
+    window: dict[str, datetime]
+    versions: list[VersionStageMetrics] = Field(default_factory=list)

libs/policyengine-fastapi/src/policyengine_fastapi/observability/correlation.py

@@ -0,0 +1,20 @@
+from __future__ import annotations
+
+from hashlib import sha256
+import json
+from typing import Any
+from uuid import uuid4
+
+
+def generate_run_id() -> str:
+    return str(uuid4())
+
+
+def stable_config_hash(payload: Any) -> str:
+    serialised = json.dumps(
+        payload,
+        sort_keys=True,
+        separators=(",", ":"),
+        default=str,
+    )
+    return f"sha256:{sha256(serialised.encode()).hexdigest()}"

libs/policyengine-fastapi/src/policyengine_fastapi/observability/emitters.py

@@ -0,0 +1,84 @@
+from __future__ import annotations
+
+from contextlib import AbstractContextManager
+from dataclasses import dataclass, field
+from typing import Protocol
+from typing import Any, Mapping
+
+from .config import ObservabilityConfig
+from .contracts import SimulationLifecycleEvent, TracerArtifactManifest
+
+
+class NoOpSpan(AbstractContextManager["NoOpSpan"]):
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        return None
+
+    def set_attribute(self, key: str, value: Any) -> None:
+        return None
+
+    def add_event(self, name: str, attributes: Mapping[str, Any] | None = None) -> None:
+        return None
+
+
+class Observability(Protocol):
+    config: ObservabilityConfig
+
+    def emit_lifecycle_event(self, event: SimulationLifecycleEvent) -> None: ...
+
+    def emit_counter(
+        self,
+        name: str,
+        value: int = 1,
+        attributes: Mapping[str, str] | None = None,
+    ) -> None: ...
+
+    def emit_histogram(
+        self,
+        name: str,
+        value: float,
+        attributes: Mapping[str, str] | None = None,
+    ) -> None: ...
+
+    def record_artifact_manifest(self, manifest: TracerArtifactManifest) -> None: ...
+
+    def span(
+        self, name: str, attributes: Mapping[str, Any] | None = None
+    ) -> NoOpSpan: ...
+
+    def flush(self) -> None: ...
+
+
+@dataclass
+class NoOpObservability:
+    config: ObservabilityConfig = field(default_factory=ObservabilityConfig.disabled)
+
+    def emit_lifecycle_event(self, event: SimulationLifecycleEvent) -> None:
+        return None
+
+    def emit_counter(
+        self,
+        name: str,
+        value: int = 1,
+        attributes: Mapping[str, str] | None = None,
+    ) -> None:
+        return None
+
+    def emit_histogram(
+        self,
+        name: str,
+        value: float,
+        attributes: Mapping[str, str] | None = None,
+    ) -> None:
+        return None
+
+    def record_artifact_manifest(self, manifest: TracerArtifactManifest) -> None:
+        return None
+
+    def span(self, name: str, attributes: Mapping[str, Any] | None = None) -> NoOpSpan:
+        return NoOpSpan()
+
+    def flush(self) -> None:
+        return None

libs/policyengine-fastapi/src/policyengine_fastapi/observability/provider.py

@@ -0,0 +1,25 @@
+from __future__ import annotations
+
+from .config import ObservabilityConfig
+from .emitters import NoOpObservability, Observability
+
+
+def build_observability(
+    config: ObservabilityConfig | None = None,
+) -> Observability:
+    """Central construction point for observability emitters.
+
+    Commit 1 intentionally returns a no-op implementation even when enabled.
+    Later commits can swap in a real backend here without changing callers.
+    """
+
+    if config is None:
+        config = ObservabilityConfig.disabled()
+
+    return NoOpObservability(config=config)
+
+
+def get_observability(
+    config: ObservabilityConfig | None = None,
+) -> Observability:
+    return build_observability(config)