Log the safety check's passes through the observability runtime

From the extra review of the reworked 1.06: the job configured logging
with logging.basicConfig, whose plain format drops every field of a log
record, so a pass's counts never reached the job's logs. The job now
starts the process observability runtime, as the Modal workers do, logs
JSON like the web service, and shuts the runtime down on exit. The pass
summary carries its counts as policyengine_attributes, the only record
fields the observability log handler keeps.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Author: anth-volk

projects/policyengine-simulation-entry/src/policyengine_simulation_entry/safety_check.py

@@ -29,6 +29,10 @@
 from dataclasses import asdict
 from typing import TYPE_CHECKING, Protocol
 
+from policyengine_simulation_observability.observability import (
+    init_process_observability,
+)
+
 from policyengine_simulation_entry.modal_calls import OWN_PROCESS_SIGNALS
 from policyengine_simulation_entry.simulation_database import SimulationDatabase
 
@@ -44,6 +48,23 @@
 EXIT_NOT_CONFIGURED = 2
 
 
+class Telemetry(Protocol):
+    """The observability runtime's part the job uses."""
+
+    def shutdown(self) -> None:
+        """Flush and stop telemetry."""
+        ...
+
+
+def _start_observability(environment: str) -> Telemetry:
+    return init_process_observability(
+        service_name="policyengine-simulation-safety-check",
+        service_role="simulation_safety_check",
+        platform="google_cloud_run",
+        environment=environment,
+    )
+
+
 class DueOutcomeRecorder(Protocol):
     """Run one recording pass (``RunRecorder``)."""
 
@@ -60,13 +81,25 @@ def main(
     *,
     environment: Mapping[str, str] | None = None,
     build_recorder: Callable[[str], DueOutcomeRecorder] = _database_recorder,
+    start_observability: Callable[[str], Telemetry] = _start_observability,
 ) -> int:
     """Run one pass and return the process exit status.
 
-    Only a signal to this process propagates; any other error fails the
+    Logs go through the observability runtime, as the web service's do. Only
+    a signal to this process propagates; any other error fails the
     execution, and the next scheduled execution tries again.
     """
     values = os.environ if environment is None else environment
+    telemetry = start_observability(values.get("APP_ENVIRONMENT", "local"))
+    try:
+        return _run_pass(values, build_recorder)
+    finally:
+        telemetry.shutdown()
+
+
+def _run_pass(
+    values: Mapping[str, str], build_recorder: Callable[[str], DueOutcomeRecorder]
+) -> int:
     database_url = values.get("SIMULATION_DATABASE_URL", "")
     if not database_url:
         logger.error("safety_check_not_configured")
@@ -81,10 +114,13 @@ def main(
             "safety_check_pass_failed", extra={"error_type": type(error).__name__}
         )
         return EXIT_PASS_FAILED
-    logger.info("safety_check_pass", extra=asdict(recording))
+    # The observability log handler keeps a record's policyengine_attributes.
+    logger.info(
+        "safety_check_pass", extra={"policyengine_attributes": asdict(recording)}
+    )
     return EXIT_PASSED
 
 
 if __name__ == "__main__":
-    logging.basicConfig(level=logging.INFO)
+    logging.getLogger().setLevel(logging.INFO)
     sys.exit(main())

projects/policyengine-simulation-entry/tests/test_safety_check.py

@@ -32,14 +32,34 @@ async def record_due_outcomes(self) -> RecordingPass:
         return self.result
 
 
-def _run(environment, recorder):
+class _Telemetry:
+    """Stands in for the observability runtime the job starts and shuts down."""
+
+    def __init__(self) -> None:
+        self.environments = []
+        self.shut_down = False
+
+    def start(self, environment):
+        self.environments.append(environment)
+        return self
+
+    def shutdown(self) -> None:
+        self.shut_down = True
+
+
+def _run(environment, recorder, telemetry=None):
     built_for = []
 
     def build(database_url):
         built_for.append(database_url)
         return recorder
 
-    status = safety_check.main(environment=environment, build_recorder=build)
+    telemetry = telemetry or _Telemetry()
+    status = safety_check.main(
+        environment=environment,
+        build_recorder=build,
+        start_observability=telemetry.start,
+    )
     return status, built_for
 
 
@@ -62,9 +82,23 @@ def test_the_pass_is_logged_with_what_it_found(caplog) -> None:
     (record,) = [
         record for record in caplog.records if record.message == "safety_check_pass"
     ]
-    assert record.checked == 4
-    assert record.succeeded == 1
-    assert record.abandoned == 2
+    # The observability log handler keeps only these attributes of a record.
+    assert record.policyengine_attributes["checked"] == 4
+    assert record.policyengine_attributes["succeeded"] == 1
+    assert record.policyengine_attributes["abandoned"] == 2
+
+
+def test_the_job_logs_through_the_observability_runtime_and_shuts_it_down() -> None:
+    telemetry = _Telemetry()
+
+    _run(
+        {"SIMULATION_DATABASE_URL": DATABASE_URL, "APP_ENVIRONMENT": "production"},
+        _OnePass(RecordingPass()),
+        telemetry,
+    )
+
+    assert telemetry.environments == ["production"]
+    assert telemetry.shut_down
 
 
 def test_an_execution_without_the_database_url_fails_without_a_pass() -> None: