diff --git a/.github/scripts/deploy-cloud-run-simulation-entry.sh b/.github/scripts/deploy-cloud-run-simulation-entry.sh index 0956ae9a5..45784c03a 100755 --- a/.github/scripts/deploy-cloud-run-simulation-entry.sh +++ b/.github/scripts/deploy-cloud-run-simulation-entry.sh @@ -100,6 +100,9 @@ required_variables=( OLD_GATEWAY_AUTH_CLIENT_ID_VALUE OLD_GATEWAY_AUTH_CLIENT_SECRET_SECRET_NAME SIMULATION_DATABASE_URL_SECRET_NAME + MODAL_TOKEN_ID_SECRET_NAME + MODAL_TOKEN_SECRET_SECRET_NAME + MODAL_WORKSPACE_ID_VALUE OBSERVABILITY_SERVICE_NAMESPACE OBSERVABILITY_TRACE_PROJECT_ID OTEL_EXPORTER_OTLP_ENDPOINT @@ -126,6 +129,10 @@ esac || fail "TAG is invalid" [[ "${ENTRYPOINT_RUNTIME_SERVICE_ACCOUNT}" =~ ^[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.gserviceaccount\.com$ ]] \ || fail "ENTRYPOINT_RUNTIME_SERVICE_ACCOUNT is invalid" +# The Modal workspace whose call-report function may report that calls +# ended (the entrypoint also checks the environment, MODAL_ENVIRONMENT). +[[ "${MODAL_WORKSPACE_ID_VALUE}" =~ ^ac-[A-Za-z0-9]+$ ]] \ + || fail "MODAL_WORKSPACE_ID_VALUE must be a Modal workspace ID" [[ "${ENTRYPOINT_MIN_INSTANCES}" =~ ^[0-9]+$ ]] \ || fail "ENTRYPOINT_MIN_INSTANCES must be a non-negative integer" [[ "${ENTRYPOINT_MAX_INSTANCES}" =~ ^[0-9]+$ ]] \ @@ -143,13 +150,16 @@ add_runtime_secret OLD_GATEWAY_AUTH_CLIENT_SECRET \ # The runtime role's URL for the entrypoint's own simulation schema. The # migration job, which uses a separate migrator URL, runs before this deploy. add_runtime_secret SIMULATION_DATABASE_URL "${SIMULATION_DATABASE_URL_SECRET_NAME}" +# sim's own Modal token, on every deploy: the entrypoint starts runs' Modal +# calls and asks Modal about a call when Modal reports that it ended. Until +# the Stage 12 comparison path is removed, its Modal calls use the same token. +add_runtime_secret MODAL_TOKEN_ID "${MODAL_TOKEN_ID_SECRET_NAME}" +add_runtime_secret MODAL_TOKEN_SECRET "${MODAL_TOKEN_SECRET_SECRET_NAME}" if [[ "${DEPLOY_STAGE12_V2}" == "true" ]]; then stage12_variables=( STAGE12_ENABLED_VALUE STAGE12_ARTIFACT_BUCKET_VALUE - MODAL_TOKEN_ID_SECRET_NAME - MODAL_TOKEN_SECRET_SECRET_NAME STAGE12_DATABASE_URL_SECRET_NAME ) for variable_name in "${stage12_variables[@]}"; do @@ -161,8 +171,6 @@ if [[ "${DEPLOY_STAGE12_V2}" == "true" ]]; then || fail "STAGE12_ARTIFACT_BUCKET_VALUE is invalid" validate_secret_reference "${STAGE12_DATABASE_URL_SECRET_NAME}" - add_stage12_runtime_secret MODAL_TOKEN_ID "${MODAL_TOKEN_ID_SECRET_NAME}" - add_stage12_runtime_secret MODAL_TOKEN_SECRET "${MODAL_TOKEN_SECRET_SECRET_NAME}" add_stage12_runtime_secret STAGE12_DATABASE_URL "${STAGE12_DATABASE_URL_SECRET_NAME}" fi @@ -188,6 +196,8 @@ jq -n ' OTEL_METRICS_EXPORTER: "otlp", OTEL_TRACES_SAMPLER_ARG: "1.0", POLICYENGINE_OTEL_GOOGLE_AUDIENCE: env.POLICYENGINE_OTEL_GOOGLE_AUDIENCE, + SIMULATION_CALL_REPORT_MODAL_WORKSPACE_ID: env.MODAL_WORKSPACE_ID_VALUE, + SIMULATION_CALL_REPORT_MODAL_ENVIRONMENT: env.MODAL_ENVIRONMENT, STAGE12_ENABLED: ( if env.DEPLOY_STAGE12_V2 == "true" then env.STAGE12_ENABLED_VALUE else "0" end ) diff --git a/.github/scripts/deploy-cloud-run-simulation-safety-check.sh b/.github/scripts/deploy-cloud-run-simulation-safety-check.sh new file mode 100755 index 000000000..54d830beb --- /dev/null +++ b/.github/scripts/deploy-cloud-run-simulation-safety-check.sh @@ -0,0 +1,132 @@ +#!/usr/bin/env bash +# Point the safety-check Cloud Run job at the entrypoint revision that serves +# traffic, and keep its Cloud Scheduler entry. +# +# The job runs `python -m policyengine_simulation_entry.safety_check`: one +# recording pass per execution. It is updated only from a revision that was +# promoted (or restored by a rollback), so a rejected candidate never records +# anything. It copies the revision's image digest, service account, Modal +# token and database URL (at the secret versions the revision pins) and its +# plain telemetry settings; nothing else. A revision deployed before the +# safety check existed has no job code, so its schedule is paused instead. +# +# With the argument `pause`, only pause the schedule: a workflow does this +# when pointing the job at a revision failed part way, so the job never keeps +# running the image traffic just moved away from. +# +# Environment: PROJECT_ID, REGION, SERVICE, and unless pausing, +# SAFETY_CHECK_REVISION, SAFETY_CHECK_SCHEDULER_SERVICE_ACCOUNT, RUNNER_TEMP. +set -euo pipefail + +fail() { + printf '%s\n' "$*" >&2 + exit 1 +} + +pause_schedule() { + if gcloud scheduler jobs describe "${job}" "${scheduler_location[@]}" >/dev/null 2>&1; then + gcloud scheduler jobs pause "${job}" "${scheduler_location[@]}" --quiet + fi +} + +# Every five minutes: SAFETY_CHECK_INTERVAL in +# policyengine_simulation_entry.run_recording. +SCHEDULE="*/5 * * * *" +SECRETS=(SIMULATION_DATABASE_URL MODAL_TOKEN_ID MODAL_TOKEN_SECRET) +PLAIN_VARIABLES_FILTER='^(APP_ENVIRONMENT|CLOUD_RUN_REGION|OBSERVABILITY_[A-Z_]+|OTEL_[A-Z_]+|POLICYENGINE_OTEL_GOOGLE_AUDIENCE)$' + +for variable_name in PROJECT_ID REGION SERVICE; do + [[ -n "${!variable_name:-}" ]] || fail "${variable_name} is required" +done +[[ "${PROJECT_ID}" =~ ^[a-z][a-z0-9-]{4,28}[a-z0-9]$ ]] || fail "PROJECT_ID is invalid" +[[ "${REGION}" =~ ^[a-z]+-[a-z]+[0-9]+$ ]] || fail "REGION is invalid" +[[ "${SERVICE}" =~ ^[a-z]([a-z0-9-]{0,47}[a-z0-9])?$ ]] || fail "SERVICE is invalid" + +job="${SERVICE}-safety-check" +location=(--project "${PROJECT_ID}" --region "${REGION}") +scheduler_location=(--project "${PROJECT_ID}" --location "${REGION}") + +if [[ "${1:-}" == "pause" ]]; then + pause_schedule + printf 'The safety check schedule for %s is paused.\n' "${SERVICE}" + exit 0 +fi + +for variable_name in SAFETY_CHECK_REVISION SAFETY_CHECK_SCHEDULER_SERVICE_ACCOUNT RUNNER_TEMP; do + [[ -n "${!variable_name:-}" ]] || fail "${variable_name} is required" +done +[[ "${SAFETY_CHECK_REVISION}" =~ ^[a-z]([a-z0-9-]{0,61}[a-z0-9])?$ ]] \ + || fail "SAFETY_CHECK_REVISION is invalid" +[[ "${SAFETY_CHECK_SCHEDULER_SERVICE_ACCOUNT}" =~ ^[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.gserviceaccount\.com$ ]] \ + || fail "SAFETY_CHECK_SCHEDULER_SERVICE_ACCOUNT is invalid" + +revision_json="$(gcloud run revisions describe "${SAFETY_CHECK_REVISION}" "${location[@]}" --format=json)" +revision_service="$(jq -r '.metadata.labels["serving.knative.dev/service"] // ""' <<<"${revision_json}")" +[[ "${revision_service}" == "${SERVICE}" ]] \ + || fail "${SAFETY_CHECK_REVISION} belongs to ${revision_service:-no service}, not ${SERVICE}" + +# Every revision deployed with the safety check has the simulation database +# and the expected Modal caller (SIMULATION_CALL_REPORT_MODAL_WORKSPACE_ID, +# which the deploy has required since then); an older one has no job code. +if ! jq -e ' + .spec.containers[0].env // [] + | any(.name == "SIMULATION_DATABASE_URL") + and any(.name == "SIMULATION_CALL_REPORT_MODAL_WORKSPACE_ID") +' >/dev/null <<<"${revision_json}"; then + pause_schedule + printf '%s predates the safety check, so its schedule is paused.\n' "${SAFETY_CHECK_REVISION}" + exit 0 +fi + +image="$(jq -er '.status.imageDigest // .spec.containers[0].image' <<<"${revision_json}")" +service_account="$(jq -er '.spec.serviceAccountName' <<<"${revision_json}")" + +job_secrets=() +for secret_variable in "${SECRETS[@]}"; do + reference="$(jq -er --arg name "${secret_variable}" ' + [.spec.containers[0].env[] | select(.name == $name) | .valueFrom.secretKeyRef + | select(.name != null and .key != null) | "\(.name):\(.key)"] + | if length == 1 then .[0] else error("not one secret reference") end + ' <<<"${revision_json}")" || fail "${SAFETY_CHECK_REVISION} has no pinned secret for ${secret_variable}" + job_secrets+=("${secret_variable}=${reference}") +done +job_secrets_csv="$(IFS=,; printf '%s' "${job_secrets[*]}")" + +environment_file="${RUNNER_TEMP}/simulation-safety-check-environment.yaml" +trap 'rm -f "${environment_file}"' EXIT +jq --arg filter "${PLAIN_VARIABLES_FILTER}" ' + [.spec.containers[0].env[] | select(.value != null and (.name | test($filter))) + | {(.name): .value}] | add // {} +' <<<"${revision_json}" > "${environment_file}" + +gcloud run jobs deploy "${job}" "${location[@]}" \ + --image "${image}" \ + --service-account "${service_account}" \ + --set-secrets "${job_secrets_csv}" \ + --env-vars-file "${environment_file}" \ + --command /opt/venv/bin/python \ + --args=-m,policyengine_simulation_entry.safety_check \ + --tasks 1 \ + --max-retries 0 \ + --task-timeout 600s \ + --cpu 1 \ + --memory 512Mi \ + --labels "serving-revision=${SAFETY_CHECK_REVISION}" \ + --quiet + +schedule_action=create +if gcloud scheduler jobs describe "${job}" "${scheduler_location[@]}" >/dev/null 2>&1; then + schedule_action=update +fi +gcloud scheduler jobs "${schedule_action}" http "${job}" "${scheduler_location[@]}" \ + --schedule "${SCHEDULE}" \ + --time-zone Etc/UTC \ + --uri "https://${REGION}-run.googleapis.com/apis/run.googleapis.com/v1/namespaces/${PROJECT_ID}/jobs/${job}:run" \ + --http-method POST \ + --oauth-service-account-email "${SAFETY_CHECK_SCHEDULER_SERVICE_ACCOUNT}" \ + --attempt-deadline 60s \ + --quiet +if [[ "${schedule_action}" == "update" ]]; then + gcloud scheduler jobs resume "${job}" "${scheduler_location[@]}" --quiet +fi +printf 'The safety check runs %s from %s.\n' "${image}" "${SAFETY_CHECK_REVISION}" diff --git a/.github/scripts/simulation-entry-main-url.sh b/.github/scripts/simulation-entry-main-url.sh new file mode 100755 index 000000000..8e86373d8 --- /dev/null +++ b/.github/scripts/simulation-entry-main-url.sh @@ -0,0 +1,24 @@ +#!/bin/bash +# Write the Simulation Entrypoint's main URL (its stable Cloud Run URL, which +# reaches the revision that serves traffic) to $GITHUB_OUTPUT as `url`, for +# the v2 Modal app's report_call_ended function. +# +# Environment: PROJECT_ID, REGION, SERVICE, GITHUB_OUTPUT. +set -euo pipefail + +fail() { + printf '%s\n' "$*" >&2 + exit 1 +} + +for variable_name in PROJECT_ID REGION SERVICE GITHUB_OUTPUT; do + [[ -n "${!variable_name:-}" ]] || fail "${variable_name} is required" +done + +url="$(gcloud run services describe "${SERVICE}" \ + --project "${PROJECT_ID}" \ + --region "${REGION}" \ + --format='value(status.url)')" +[[ "${url}" =~ ^https://[a-z0-9.-]+$ ]] \ + || fail "${SERVICE} has no https main URL: ${url:-none}" +printf 'url=%s\n' "${url}" >> "${GITHUB_OUTPUT}" diff --git a/.github/workflows/simulation-deploy.reusable.yml b/.github/workflows/simulation-deploy.reusable.yml index ea4f48585..7324468fe 100644 --- a/.github/workflows/simulation-deploy.reusable.yml +++ b/.github/workflows/simulation-deploy.reusable.yml @@ -279,6 +279,7 @@ jobs: STAGE12_DATABASE_URL_SECRET_NAME: ${{ vars.STAGE12_DATABASE_URL_SECRET_NAME }} MODAL_TOKEN_ID_SECRET_NAME: ${{ vars.MODAL_TOKEN_ID_SECRET_NAME }} MODAL_TOKEN_SECRET_SECRET_NAME: ${{ vars.MODAL_TOKEN_SECRET_SECRET_NAME }} + MODAL_WORKSPACE_ID_VALUE: ${{ vars.MODAL_WORKSPACE_ID }} run: .github/scripts/deploy-cloud-run-simulation-entry.sh - name: Resolve exact candidate metadata @@ -670,12 +671,22 @@ jobs: STAGE12_CACHE_BUCKET: ${{ vars.STAGE12_CACHE_BUCKET }} run: ../../.github/scripts/stage12-cache-precompute.sh "${{ inputs.modal_environment }}" + # The v2 functions report to sim's main URL that a call ended. + - name: Find the Simulation Entrypoint's main URL + id: entry_url + env: + PROJECT_ID: ${{ vars.SIMULATION_ENTRYPOINT_GCP_PROJECT_ID }} + REGION: ${{ vars.SIMULATION_ENTRYPOINT_GCP_REGION || 'us-central1' }} + SERVICE: ${{ inputs.entrypoint_service }} + run: .github/scripts/simulation-entry-main-url.sh + - name: Deploy the additional versioned v2 Modal application working-directory: projects/policyengine-simulation-executor env: MODAL_TOKEN_ID: ${{ secrets.MODAL_TOKEN_ID }} MODAL_TOKEN_SECRET: ${{ secrets.MODAL_TOKEN_SECRET }} MODAL_APP_NAME: ${{ steps.application.outputs.name }} + SIMULATION_ENTRY_URL: ${{ steps.entry_url.outputs.url }} STAGE12_EXPECT_POLICYENGINE_VERSION: ${{ needs.prepare.outputs.policyengine_version }} STAGE12_EXPECT_US_VERSION: ${{ needs.prepare.outputs.us_version }} STAGE12_EXPECT_UK_VERSION: ${{ needs.prepare.outputs.uk_version }} @@ -1011,6 +1022,15 @@ jobs: SIMULATION_ENTRYPOINT_TEST_ENFORCE_PERMISSIONS: ${{ vars.SIMULATION_ENTRYPOINT_ENFORCE_PERMISSIONS }} run: uv run --frozen pytest authenticated_tests/ -v + # Only a verified, promoted revision records outcomes in the safety + # check, so a rejected candidate never records anything. + - name: Point the safety check at the promoted revision + id: safety_check + env: + SAFETY_CHECK_REVISION: ${{ env.TARGET_REVISION }} + SAFETY_CHECK_SCHEDULER_SERVICE_ACCOUNT: ${{ vars.SIMULATION_SAFETY_CHECK_SCHEDULER_SERVICE_ACCOUNT }} + run: .github/scripts/deploy-cloud-run-simulation-safety-check.sh + - name: Restore previous revision after failed verification if: ${{ failure() && steps.promote.outcome == 'success' }} env: @@ -1021,6 +1041,18 @@ jobs: SIMULATION_ENTRYPOINT_EXPECTED_CURRENT_REVISION: ${{ env.TARGET_REVISION }} run: .github/scripts/set-cloud-run-simulation-entry-revision.sh + - name: Point the safety check back at the previous revision + id: safety_check_restore + if: ${{ failure() && steps.safety_check.outcome == 'failure' && env.PREVIOUS_REVISION != '' }} + env: + SAFETY_CHECK_REVISION: ${{ env.PREVIOUS_REVISION }} + SAFETY_CHECK_SCHEDULER_SERVICE_ACCOUNT: ${{ vars.SIMULATION_SAFETY_CHECK_SCHEDULER_SERVICE_ACCOUNT }} + run: .github/scripts/deploy-cloud-run-simulation-safety-check.sh + + - name: Pause the safety check after a failed restore + if: ${{ failure() && steps.safety_check_restore.outcome == 'failure' }} + run: .github/scripts/deploy-cloud-run-simulation-safety-check.sh pause + finalize_routing: name: Finalize the Modal routing publication # Finish the durable journal after every gate, including failed uploads diff --git a/.github/workflows/simulation-rollback.yml b/.github/workflows/simulation-rollback.yml index 10054ea5d..c51ef0f8e 100644 --- a/.github/workflows/simulation-rollback.yml +++ b/.github/workflows/simulation-rollback.yml @@ -218,6 +218,19 @@ jobs: .github/scripts/set-cloud-run-simulation-entry-revision.sh fi + - name: Point the safety check at the restored revision + id: safety_check + if: ${{ inputs.entrypoint_revision != '' }} + env: + SAFETY_CHECK_REVISION: ${{ inputs.entrypoint_revision }} + SAFETY_CHECK_SCHEDULER_SERVICE_ACCOUNT: ${{ vars.SIMULATION_SAFETY_CHECK_SCHEDULER_SERVICE_ACCOUNT }} + run: .github/scripts/deploy-cloud-run-simulation-safety-check.sh + + # The job must not keep running the image traffic just moved away from. + - name: Pause the safety check after a failed update + if: ${{ failure() && steps.safety_check.outcome == 'failure' }} + run: .github/scripts/deploy-cloud-run-simulation-safety-check.sh pause + - name: Verify the stable service serves the restored release if: ${{ !inputs.resolve_pending_publication }} env: diff --git a/changelog.d/947.added.md b/changelog.d/947.added.md new file mode 100644 index 000000000..20ff473cc --- /dev/null +++ b/changelog.d/947.added.md @@ -0,0 +1 @@ +Added simulation runs and report runs, in migration revision `0003` of the `simulation` schema. Each run has its own status and Modal call, every status change is one conditional update, and only the instance that claims a run starts its Modal call. When a v2 Modal call ends, the new `report_call_ended` Modal function reports it on `POST /internal/modal-calls/ended`, which accepts only a Modal identity token from that function; sim then reads the call's outcome from Modal and records it on the runs. A new Cloud Run job, `-safety-check`, runs every five minutes on Cloud Scheduler and records calls that ended without a report. No route creates runs yet. diff --git a/changelog.d/947.changed.md b/changelog.d/947.changed.md new file mode 100644 index 000000000..a3b61fcdc --- /dev/null +++ b/changelog.d/947.changed.md @@ -0,0 +1,3 @@ +Every Simulation Entrypoint deploy now mounts sim's own Modal token as `MODAL_TOKEN_ID` and `MODAL_TOKEN_SECRET`, no longer only when `DEPLOY_STAGE12_V2` is `true`. + +Every deploy now also needs, in both `beta` and `prod`, the `MODAL_TOKEN_ID_SECRET_NAME` and `MODAL_TOKEN_SECRET_SECRET_NAME` environment variables naming sim's own Modal token secrets, the `MODAL_WORKSPACE_ID` environment variable, and the `SIMULATION_SAFETY_CHECK_SCHEDULER_SERVICE_ACCOUNT` environment variable with a scheduler service account allowed to run the safety-check job. Without the first three the entrypoint deploy stops; without the last, the promotion's safety-check step fails and traffic goes back to the previous revision. diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/call_reports.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/call_reports.py new file mode 100644 index 000000000..4b9a32cf9 --- /dev/null +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/call_reports.py @@ -0,0 +1,120 @@ +"""The report a Modal call sends the Simulation Entrypoint when it ends. + +The last step of each top-level v2 Modal function (the economy coordinator +and the household functions) reports that its call ended: it starts the small +Modal function ``CALL_REPORT_FUNCTION_NAME``, which posts a +``CallEndedReport`` to the entrypoint's ``CALL_ENDED_PATH`` with the Modal +identity token of its container as a bearer token. The report carries the +call ID and the run IDs only, never a result: the entrypoint checks that the +request comes from that Modal function, asks Modal for the call's outcome +itself, and records it. + +The entrypoint answers with a ``CallReportAnswer``. While it cannot record +the outcome yet (the call has not finished, its ID is not stored yet, Modal +or the database cannot be reached), it answers 503 with a status in +``RETRY_LATER_STATUSES``, and the sender reports again after a wait. A +repeated report changes nothing. +""" + +from __future__ import annotations + +from enum import StrEnum +from typing import Self +from uuid import UUID + +from pydantic import BaseModel, ConfigDict, Field, model_validator + +# The entrypoint's route for call-ended reports. It is not in the public +# OpenAPI document. +CALL_ENDED_PATH = "/internal/modal-calls/ended" +# The Modal function that sends reports; the entrypoint accepts identity +# tokens of this function only. +CALL_REPORT_FUNCTION_NAME = "report_call_ended" +# Modal sets this variable to the container's identity token in every +# function container. +MODAL_IDENTITY_ENVIRONMENT_VARIABLE = "MODAL_IDENTITY_TOKEN" +# A Modal function call ID, as Modal prints it. +MODAL_CALL_ID_PATTERN = r"^fc-[A-Za-z0-9]{1,64}$" +# A call computes one simulation run, or a household report's baseline and +# reform runs together. +MOST_SIMULATION_RUNS_PER_CALL = 2 + + +class CallEndedReport(BaseModel): + """A Modal call's report that it ended: its ID and the runs it computes. + + Exactly one of ``simulation_run_ids`` (a simulation call: one run, or a + household report's two) and ``report_run_id`` (an economy comparison + step's call) is set. + """ + + model_config = ConfigDict(frozen=True, extra="forbid") + + modal_call_id: str = Field(pattern=MODAL_CALL_ID_PATTERN) + simulation_run_ids: tuple[UUID, ...] | None = Field( + default=None, min_length=1, max_length=MOST_SIMULATION_RUNS_PER_CALL + ) + report_run_id: UUID | None = None + + @model_validator(mode="after") + def _names_one_kind_of_run(self) -> Self: + if (self.simulation_run_ids is None) == (self.report_run_id is None): + msg = "name either simulation_run_ids or report_run_id" + raise ValueError(msg) + if self.simulation_run_ids is not None and len( + set(self.simulation_run_ids) + ) != len(self.simulation_run_ids): + msg = "simulation_run_ids names a run twice" + raise ValueError(msg) + return self + + +class CallReportStatus(StrEnum): + """What the entrypoint did with a report.""" + + # The call's outcome is now recorded on its runs. + RECORDED = "recorded" + # The runs already hold the call's outcome; nothing changed. + ALREADY_RECORDED = "already_recorded" + # The runs are still being dispatched and do not hold the call ID yet. + NOT_DISPATCHED_YET = "not_dispatched_yet" + # Modal holds no output for the call yet. + STILL_RUNNING = "still_running" + # Modal could not be asked about the call. + MODAL_UNREACHABLE = "modal_unreachable" + # The serving revision cannot record this kind of result yet. + NO_RESULT_RECORDER = "no_result_recorder" + # The database could not be read or written. + NOT_WRITTEN = "not_written" + # The entrypoint is not configured to accept reports. + NOT_CONFIGURED = "not_configured" + # No run, or not every run, the report names exists. + UNKNOWN_RUNS = "unknown_runs" + # The runs hold another Modal call, or finished without holding one (for + # example failed as abandoned before the call's ID was stored). + OTHER_CALL = "other_call" + + +RETRY_LATER_STATUSES = frozenset( + { + CallReportStatus.NOT_DISPATCHED_YET, + CallReportStatus.STILL_RUNNING, + CallReportStatus.MODAL_UNREACHABLE, + CallReportStatus.NO_RESULT_RECORDER, + CallReportStatus.NOT_WRITTEN, + CallReportStatus.NOT_CONFIGURED, + } +) + + +class CallReportAnswer(BaseModel): + """The entrypoint's answer to a report.""" + + model_config = ConfigDict(frozen=True) + + status: CallReportStatus + + @property + def retry_later(self) -> bool: + """Whether the sender should report again after a wait.""" + return self.status in RETRY_LATER_STATUSES diff --git a/libs/policyengine-simulation-contract/tests/test_call_reports.py b/libs/policyengine-simulation-contract/tests/test_call_reports.py new file mode 100644 index 000000000..8ccce497b --- /dev/null +++ b/libs/policyengine-simulation-contract/tests/test_call_reports.py @@ -0,0 +1,99 @@ +"""The report a Modal function's call sends sim when it ends, and sim's answer.""" + +from __future__ import annotations + +from uuid import uuid4 + +import pytest +from policyengine_simulation_contract.call_reports import ( + RETRY_LATER_STATUSES, + CallEndedReport, + CallReportAnswer, + CallReportStatus, +) +from pydantic import ValidationError + +CALL_ID = "fc-01JB3W7X9Z8Q2R5T6Y4U1I0O9P" + + +def test_a_report_names_the_simulation_runs_its_call_computes() -> None: + baseline, reform = uuid4(), uuid4() + + report = CallEndedReport( + modal_call_id=CALL_ID, simulation_run_ids=(baseline, reform) + ) + + assert report.simulation_run_ids == (baseline, reform) + assert report.report_run_id is None + + +def test_a_report_names_the_report_run_whose_comparison_it_computes() -> None: + report_run_id = uuid4() + + report = CallEndedReport(modal_call_id=CALL_ID, report_run_id=report_run_id) + + assert report.report_run_id == report_run_id + assert report.simulation_run_ids is None + + +def test_a_report_round_trips_as_json() -> None: + report = CallEndedReport(modal_call_id=CALL_ID, simulation_run_ids=(uuid4(),)) + + assert CallEndedReport.model_validate_json(report.model_dump_json()) == report + + +@pytest.mark.parametrize( + "fields", + [ + {}, + {"simulation_run_ids": (uuid4(),), "report_run_id": uuid4()}, + {"simulation_run_ids": ()}, + {"simulation_run_ids": (uuid4(), uuid4(), uuid4())}, + ], + ids=["neither", "both", "no-runs", "three-runs"], +) +def test_a_report_names_either_one_or_two_simulation_runs_or_one_report_run( + fields, +) -> None: + with pytest.raises(ValidationError): + CallEndedReport(modal_call_id=CALL_ID, **fields) + + +def test_a_report_names_each_simulation_run_once() -> None: + run_id = uuid4() + + with pytest.raises(ValidationError, match="twice"): + CallEndedReport(modal_call_id=CALL_ID, simulation_run_ids=(run_id, run_id)) + + +@pytest.mark.parametrize( + "call_id", ["", "fc-", "in-01JB3W7X9Z", "fc-01JB/../x", "fc-" + "A" * 80] +) +def test_a_report_names_a_modal_function_call(call_id) -> None: + with pytest.raises(ValidationError): + CallEndedReport(modal_call_id=call_id, report_run_id=uuid4()) + + +def test_a_report_carries_nothing_else() -> None: + with pytest.raises(ValidationError): + CallEndedReport.model_validate( + {"modal_call_id": CALL_ID, "report_run_id": str(uuid4()), "result": {}} + ) + + +def test_the_sender_retries_only_while_sim_cannot_record_yet() -> None: + assert { + CallReportStatus.NOT_DISPATCHED_YET, + CallReportStatus.STILL_RUNNING, + CallReportStatus.MODAL_UNREACHABLE, + CallReportStatus.NO_RESULT_RECORDER, + CallReportStatus.NOT_WRITTEN, + CallReportStatus.NOT_CONFIGURED, + } == RETRY_LATER_STATUSES + assert CallReportStatus.RECORDED not in RETRY_LATER_STATUSES + assert CallReportStatus.UNKNOWN_RUNS not in RETRY_LATER_STATUSES + + +def test_an_answer_says_whether_to_report_again() -> None: + assert CallReportAnswer(status=CallReportStatus.STILL_RUNNING).retry_later + assert not CallReportAnswer(status=CallReportStatus.ALREADY_RECORDED).retry_later diff --git a/libs/policyengine-simulation-persistence/README.md b/libs/policyengine-simulation-persistence/README.md index 1436930e8..faeff7e29 100644 --- a/libs/policyengine-simulation-persistence/README.md +++ b/libs/policyengine-simulation-persistence/README.md @@ -14,17 +14,54 @@ writes them with DML only. | --- | --- | | `simulations` | What to run, independent of model versions: kind (`economy` or `household`), country, policy (null for current law), dynamic, and either a household or a region code (`state/ca`) with an optional dataset name (null means the region's default dataset in each catalog version) | | `reports` | Which two simulations a report compares (an economy report needs a reform; a household report may have only a baseline), both of the report's own country and kind, with its year and options | +| `simulation_runs` | One execution of a simulation for a catalog version and year: the region code and dataset row it resolved in that version (economy runs only), its status, Modal call ID, timings, error, trigger (`initial`, `manual` or `system`) and attempt number | +| `report_runs` | One calculation of a report: a baseline and a reform simulation run of the same catalog version and year (no reform run for a single-simulation household report), the trigger, and, for an economy report, the comparison step's own status and Modal call | -Both are definitions: the runtime role may read and insert them but not update -or delete them. Rows migrated from the Public API's v1 database keep their v1 +`simulations` and `reports` are definitions: the runtime role may read and +insert them but not update or delete them. Rows migrated from the Public API's v1 database keep their v1 IDs in the unique `legacy_simulation_id` and `legacy_report_id` columns, and the stores' `create_migrated` methods return the existing row when the same v1 ID is migrated again. Foreign keys into the Public API's tables in the same database (`policies` and -`households`, each together with its country, and `dynamics`) are -`ON DELETE RESTRICT`. Creating them needs `REFERENCES` on those tables for the -migrator role; the runtime role needs no privilege on them. +`households`, each together with its country, `dynamics`, +`tax_benefit_model_versions`, and `regions` and `datasets` each together with +its catalog version) are `ON DELETE RESTRICT`. Creating them needs +`REFERENCES` on those tables for the migrator role; the runtime role needs no +privilege on them. + +### Runs + +A simulation run's setup is its catalog version (one policyengine.py bundle +version for one country, which pins the country package, data release and +datasets) and its year. Its region and dataset keys include its catalog +version, through the Public API's unique constraints +`uq_regions_model_version_code` and `uq_datasets_id_model_version`, so a run +can use only its own version's rows; the caller resolves them in that version +(the simulation's region code, and the dataset it names or the region's +default dataset). A report run's simulation run keys include its catalog +version and year, so both of its runs share them. + +`RunStore.start_simulation_run` reuses the newest attempt for the same +simulation, catalog version and year while it is unfinished, or when it +succeeded within the last 28 days (output files are kept 30); a failed run is +never reused. Otherwise, and always when forced, it adds the next attempt. +Concurrent starts settle on one row through the unique attempt number. +`RunStore.create_report_run` joins one run of each of the report's simulations +and takes the catalog version and year from them; a report's status is its +newest report run's. The runtime role may read, insert and update runs but +not delete them. + +`RunTransitions` makes every status change one conditional update on the +expected status, so a second writer changes nothing: only the caller that +moves runs from `pending` to `dispatching` starts their Modal call, the runs +one call computes change together, a safety-check pass claims a due call by +moving its next check time (a call's own report needs no claim, since its +outcome is written conditionally too), results are written in the +transaction that marks a run succeeded, and dispatches left `dispatching` (and comparison +steps whose simulation run failed) are failed so the work can be submitted +again. A comparison step can start only once both of its simulation runs +have succeeded; starting it is left to the stage 2 report path. ## Ownership rules @@ -65,10 +102,9 @@ privileges on tables the migrator creates later; revision 0002 revokes `UPDATE` on the definition tables, which never change. Both roles must exist before the first migration runs; an operator creates them. Before revision 0002 is deployed, the owner of the Public API's tables must also grant the -migrator `REFERENCES` on `policies`, `households` and `dynamics` (the -operator step grants it on `tax_benefit_model_versions`, `datasets` and -`regions` too, for later revisions); without it the migration fails and -applies nothing. The migrator needs either +migrator `REFERENCES` on `policies`, `households` and `dynamics`, and +before revision 0003 on `tax_benefit_model_versions`, `datasets` and +`regions`; without it the migration fails and applies nothing. The migrator needs either `CREATE` on the database or a `simulation` schema created in advance and owned by it: the environment creates the schema only when it is missing, because PostgreSQL checks `CREATE` on the database even for diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/__init__.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/__init__.py index 5e87ac336..be715f3e5 100644 --- a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/__init__.py +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/__init__.py @@ -15,6 +15,26 @@ ReportRecord, ReportStore, ) +from policyengine_simulation_persistence.run_store import ( + SUCCEEDED_RUN_REUSE_WINDOW, + RunStore, +) +from policyengine_simulation_persistence.run_transitions import ( + AbandonedDispatches, + DueSimulationCall, + RunError, + RunErrorCode, + RunTransitions, +) +from policyengine_simulation_persistence.runs import ( + NewReportRun, + NewSimulationRun, + ReportRunRecord, + RunStatus, + RunTrigger, + SimulationRunRecord, + StartedSimulationRun, +) from policyengine_simulation_persistence.schema import ( ALEMBIC_VERSION_TABLE, HEAD_REVISION, @@ -34,14 +54,28 @@ "HEAD_REVISION", "RUNTIME_ROLE", "SIMULATION_SCHEMA", + "SUCCEEDED_RUN_REUSE_WINDOW", + "AbandonedDispatches", + "DueSimulationCall", "NewReport", + "NewReportRun", "NewSimulation", + "NewSimulationRun", "ReportKind", "ReportRecord", + "ReportRunRecord", "ReportStore", + "RunError", + "RunErrorCode", + "RunStatus", + "RunStore", + "RunTransitions", + "RunTrigger", "SimulationKind", "SimulationRecord", + "SimulationRunRecord", "SimulationStore", + "StartedSimulationRun", "create_simulation_engine", "normalized_database_url", "read_schema_revision", diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/public_api.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/public_api.py index 65f709a1f..5dfd974fd 100644 --- a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/public_api.py +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/public_api.py @@ -20,7 +20,7 @@ from __future__ import annotations -from sqlalchemy import Column, MetaData, String, Table, Uuid +from sqlalchemy import Column, MetaData, String, Table, UniqueConstraint, Uuid public_api_metadata = MetaData() @@ -41,3 +41,30 @@ public_api_metadata, Column("id", Uuid(as_uuid=True), primary_key=True), ) +# A catalog version: one published set of model metadata for one country and +# one policyengine.py version. Each publication creates new region and dataset +# rows, unique by code or name within their version. +tax_benefit_model_versions = Table( + "tax_benefit_model_versions", + public_api_metadata, + Column("id", Uuid(as_uuid=True), primary_key=True), +) +regions = Table( + "regions", + public_api_metadata, + Column("id", Uuid(as_uuid=True), primary_key=True), + Column("tax_benefit_model_version_id", Uuid(as_uuid=True), nullable=False), + Column("code", String(255), nullable=False), + UniqueConstraint( + "tax_benefit_model_version_id", "code", name="uq_regions_model_version_code" + ), +) +datasets = Table( + "datasets", + public_api_metadata, + Column("id", Uuid(as_uuid=True), primary_key=True), + Column("tax_benefit_model_version_id", Uuid(as_uuid=True), nullable=False), + UniqueConstraint( + "id", "tax_benefit_model_version_id", name="uq_datasets_id_model_version" + ), +) diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/reports.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/reports.py index 99faccf40..ec911e10f 100644 --- a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/reports.py +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/reports.py @@ -29,6 +29,7 @@ Integer, Table, Text, + UniqueConstraint, Uuid, func, select, @@ -106,6 +107,9 @@ class ReportKind(StrEnum): name="fk_reports_reform_simulation", ondelete="RESTRICT", ), + # Report runs refer to a report together with its kind, so only an economy + # report run has a comparison step. + UniqueConstraint("id", "kind"), ) Index("ix_reports_baseline_simulation_id", reports.c.baseline_simulation_id) Index("ix_reports_reform_simulation_id", reports.c.reform_simulation_id) diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/run_store.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/run_store.py new file mode 100644 index 000000000..e04547fc4 --- /dev/null +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/run_store.py @@ -0,0 +1,356 @@ +"""Find, create and read simulation runs and report runs. + +Starting a simulation run reuses the newest attempt for the same simulation, +catalog version and year while it is unfinished, or succeeded within the +reuse window; otherwise, and always when forced, it adds the next attempt. +Concurrent starts settle on one row through the unique attempt number: the +loser of the insert returns the winner's run. +""" + +from __future__ import annotations + +from datetime import UTC, datetime, timedelta +from typing import TYPE_CHECKING +from uuid import UUID, uuid4 + +from sqlalchemy import select +from sqlalchemy.dialects.postgresql import insert as postgres_insert + +from policyengine_simulation_persistence.reports import ReportKind, reports +from policyengine_simulation_persistence.runs import ( + UNFINISHED_STATUSES, + NewReportRun, + NewSimulationRun, + ReportRunRecord, + RunStatus, + SimulationRunRecord, + StartedSimulationRun, + report_runs, + simulation_runs, +) +from policyengine_simulation_persistence.simulations import ( + SimulationKind, + simulations, +) + +if TYPE_CHECKING: + from collections.abc import Callable + + from sqlalchemy import ColumnElement, Connection, Engine, RowMapping + +# Output files are kept 30 days. A succeeded run is reused only if it finished +# within the last 28, so its files outlive the comparison step that reads them. +SUCCEEDED_RUN_REUSE_WINDOW = timedelta(days=28) + + +def utc_now() -> datetime: + """Return the current time in UTC.""" + return datetime.now(UTC) + + +def _check_region_code(simulation: RowMapping, new_run: NewSimulationRun) -> None: + """Reject an economy run whose region code is not its simulation's. + + A missing region or dataset, or one on a household run, is left to the + database's population check. + """ + if ( + simulation["kind"] == SimulationKind.ECONOMY.value + and new_run.region_code is not None + and new_run.region_code != simulation["region_code"] + ): + msg = "an economy run's region_code must be its simulation's region code" + raise ValueError(msg) + + +def _check_report_run_simulations( + report: RowMapping, + baseline: RowMapping | None, + reform: RowMapping | None, + new_report_run: NewReportRun, +) -> RowMapping: + """Require one run of each of the report's simulations, and no other. + + Return the baseline run's key, from which the report run takes its catalog + version and year. + """ + expected_reform_run = report["reform_simulation_id"] is not None + if ( + baseline is None + or baseline["simulation_id"] != report["baseline_simulation_id"] + or expected_reform_run != (new_report_run.reform_simulation_run_id is not None) + or ( + expected_reform_run + and ( + reform is None + or reform["simulation_id"] != report["reform_simulation_id"] + ) + ) + ): + msg = ( + "a report run needs one existing run of each of the report's " + "simulations: its baseline and, when it has one, its reform" + ) + raise ValueError(msg) + return baseline + + +class RunStore: + """Typed SQLAlchemy Core reads and inserts for runs.""" + + def __init__( + self, + engine: Engine, + *, + clock: Callable[[], datetime] = utc_now, + ) -> None: + """Use ``engine`` (the runtime role) and ``clock`` for every timestamp.""" + self.engine = engine + self._clock = clock + + # --- Simulation runs --------------------------------------------------- + + def start_simulation_run( + self, new_run: NewSimulationRun, *, force: bool = False + ) -> StartedSimulationRun: + """Return the reusable run for this setup, or create the next attempt. + + A run is reused while it is unfinished, or when it succeeded within + ``SUCCEEDED_RUN_REUSE_WINDOW``; a failed run never is. ``force`` always + adds the next attempt, unless another caller added it at the same + time, in which case that run is returned. Raise LookupError when the + simulation does not exist. + """ + now = self._clock() + with self.engine.begin() as connection: + simulation = ( + connection.execute( + select(simulations.c.kind, simulations.c.region_code).where( + simulations.c.id == new_run.simulation_id + ) + ) + .mappings() + .first() + ) + if simulation is None: + msg = f"no simulation {new_run.simulation_id}" + raise LookupError(msg) + _check_region_code(simulation, new_run) + latest = self._latest_simulation_run( + connection, + new_run.simulation_id, + new_run.tax_benefit_model_version_id, + new_run.year, + ) + if latest is not None and not force and self._reusable(latest, now): + return StartedSimulationRun(run=latest, created=False) + attempt = 1 if latest is None else latest.attempt + 1 + created = ( + connection.execute( + postgres_insert(simulation_runs) + .values( + id=uuid4(), + simulation_id=new_run.simulation_id, + simulation_kind=simulation["kind"], + tax_benefit_model_version_id=( + new_run.tax_benefit_model_version_id + ), + year=new_run.year, + attempt=attempt, + region_code=new_run.region_code, + dataset_id=new_run.dataset_id, + trigger=new_run.trigger.value, + status=RunStatus.PENDING.value, + created_at=now, + updated_at=now, + ) + .on_conflict_do_nothing( + index_elements=[ + simulation_runs.c.simulation_id, + simulation_runs.c.tax_benefit_model_version_id, + simulation_runs.c.year, + simulation_runs.c.attempt, + ] + ) + .returning(*simulation_runs.c) + ) + .mappings() + .first() + ) + if created is not None: + return StartedSimulationRun( + run=SimulationRunRecord.from_row(created), created=True + ) + # Another caller added this attempt after the read above; the + # insert waited for it to commit, so it is visible now. + concurrent = self._latest_simulation_run( + connection, + new_run.simulation_id, + new_run.tax_benefit_model_version_id, + new_run.year, + ) + if concurrent is None: # pragma: no cover - the conflict proves a row + msg = "the conflicting simulation run disappeared" + raise RuntimeError(msg) + return StartedSimulationRun(run=concurrent, created=False) + + @staticmethod + def _reusable(run: SimulationRunRecord, now: datetime) -> bool: + if run.status in UNFINISHED_STATUSES: + return True + return ( + run.status is RunStatus.SUCCEEDED + and run.completed_at is not None + and run.completed_at >= now - SUCCEEDED_RUN_REUSE_WINDOW + ) + + @staticmethod + def _latest_simulation_run( + connection: Connection, + simulation_id: UUID, + catalog_version_id: UUID, + year: int, + ) -> SimulationRunRecord | None: + row = ( + connection.execute( + select(simulation_runs) + .where( + simulation_runs.c.simulation_id == simulation_id, + simulation_runs.c.tax_benefit_model_version_id + == catalog_version_id, + simulation_runs.c.year == year, + ) + .order_by(simulation_runs.c.attempt.desc()) + .limit(1) + ) + .mappings() + .first() + ) + return None if row is None else SimulationRunRecord.from_row(row) + + def get_simulation_run(self, run_id: UUID) -> SimulationRunRecord | None: + """Return the simulation run with this ID, or None.""" + with self.engine.connect() as connection: + row = ( + connection.execute( + select(simulation_runs).where(simulation_runs.c.id == run_id) + ) + .mappings() + .first() + ) + return None if row is None else SimulationRunRecord.from_row(row) + + def latest_simulation_run( + self, simulation_id: UUID, catalog_version_id: UUID, year: int + ) -> SimulationRunRecord | None: + """Return the newest attempt for this simulation, catalog version and year.""" + with self.engine.connect() as connection: + return self._latest_simulation_run( + connection, simulation_id, catalog_version_id, year + ) + + # --- Report runs ------------------------------------------------------- + + def create_report_run(self, new_report_run: NewReportRun) -> ReportRunRecord: + """Add a report run joining two existing simulation runs. + + The report run takes its catalog version and year from its baseline + run; the database refuses a reform run of another version or year. + An economy report run starts with a pending comparison step. Raise + LookupError when the report does not exist, and ValueError when the + runs are not runs of the report's simulations. + """ + now = self._clock() + with self.engine.begin() as connection: + report = ( + connection.execute( + select( + reports.c.kind, + reports.c.baseline_simulation_id, + reports.c.reform_simulation_id, + ).where(reports.c.id == new_report_run.report_id) + ) + .mappings() + .first() + ) + if report is None: + msg = f"no report {new_report_run.report_id}" + raise LookupError(msg) + baseline = self._run_key( + connection, new_report_run.baseline_simulation_run_id + ) + reform = ( + None + if new_report_run.reform_simulation_run_id is None + else self._run_key(connection, new_report_run.reform_simulation_run_id) + ) + baseline = _check_report_run_simulations( + report, baseline, reform, new_report_run + ) + economy = report["kind"] == ReportKind.ECONOMY.value + row = ( + connection.execute( + postgres_insert(report_runs) + .values( + id=uuid4(), + report_id=new_report_run.report_id, + report_kind=report["kind"], + tax_benefit_model_version_id=( + baseline["tax_benefit_model_version_id"] + ), + year=baseline["year"], + baseline_simulation_run_id=( + new_report_run.baseline_simulation_run_id + ), + reform_simulation_run_id=( + new_report_run.reform_simulation_run_id + ), + trigger=new_report_run.trigger.value, + comparison_status=( + RunStatus.PENDING.value if economy else None + ), + created_at=now, + updated_at=now, + ) + .returning(*report_runs.c) + ) + .mappings() + .one() + ) + return ReportRunRecord.from_row(row) + + @staticmethod + def _run_key(connection: Connection, run_id: UUID) -> RowMapping | None: + return ( + connection.execute( + select( + simulation_runs.c.simulation_id, + simulation_runs.c.tax_benefit_model_version_id, + simulation_runs.c.year, + ).where(simulation_runs.c.id == run_id) + ) + .mappings() + .first() + ) + + def get_report_run(self, report_run_id: UUID) -> ReportRunRecord | None: + """Return the report run with this ID, or None.""" + return self._one_report_run(report_runs.c.id == report_run_id) + + def latest_report_run(self, report_id: UUID) -> ReportRunRecord | None: + """Return the report's newest report run, whose status is the report's.""" + return self._one_report_run(report_runs.c.report_id == report_id) + + def _one_report_run(self, condition: ColumnElement[bool]) -> ReportRunRecord | None: + with self.engine.connect() as connection: + row = ( + connection.execute( + select(report_runs) + .where(condition) + .order_by(report_runs.c.created_at.desc(), report_runs.c.id.desc()) + .limit(1) + ) + .mappings() + .first() + ) + return None if row is None else ReportRunRecord.from_row(row) diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/run_transitions.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/run_transitions.py new file mode 100644 index 000000000..9d83327f2 --- /dev/null +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/run_transitions.py @@ -0,0 +1,615 @@ +"""Conditional status changes for simulation runs and comparison steps. + +Every change is one ``UPDATE ... WHERE = ``. When two +Simulation Entrypoint instances act on the same run at once, PostgreSQL makes +the second wait for the first's row lock, then re-checks the condition and +changes nothing. So no lease table or advisory lock is needed: + +- only the instance that moves a run from ``pending`` to ``dispatching`` + starts its Modal call, and it then records the call ID (``running``); +- the runs that share one Modal call, such as a household report's baseline + and reform, change status together, all or none; +- a recording pass claims a running call's next check by moving its + ``next_check_at`` forward, so two passes do not check the same call at + once, and only one of them records its outcome; +- a call's results are written in the same transaction that marks its runs + ``succeeded``, so a run never succeeds without its results; +- a run left in ``dispatching`` (the instance stopped between claiming it and + recording its call) is failed after a while, so it can be submitted again. + +An economy report run's comparison step follows the same statuses on its own +``comparison_*`` columns, and may start only once both of its simulation runs +have succeeded. +""" + +from __future__ import annotations + +from collections import defaultdict +from dataclasses import dataclass +from enum import StrEnum +from typing import TYPE_CHECKING + +from sqlalchemy import and_, exists, func, or_, select, update + +from policyengine_simulation_persistence.run_store import utc_now +from policyengine_simulation_persistence.runs import ( + ReportRunRecord, + RunStatus, + SimulationRunRecord, + bounded_error_message, + report_runs, + simulation_runs, +) + +if TYPE_CHECKING: + from collections.abc import Callable, Collection + from datetime import datetime, timedelta + from uuid import UUID + + from sqlalchemy import ColumnElement, Connection, Engine + + type SimulationResultWriter = Callable[ + [Connection, tuple[SimulationRunRecord, ...]], None + ] + type ComparisonResultWriter = Callable[[Connection, ReportRunRecord], None] + + +class RunErrorCode(StrEnum): + """Why a run or comparison step failed.""" + + # Modal refused the call or never acknowledged it, or the acknowledged + # call could not be recorded (the entrypoint then tries to cancel it). + DISPATCH_FAILED = "dispatch_failed" + # No Modal call was recorded in time after the dispatch was claimed. + DISPATCH_ABANDONED = "dispatch_abandoned" + # The Modal function raised. + WORKER_ERROR = "worker_error" + WORKER_TIMEOUT = "worker_timeout" + # Modal no longer holds the output, or does not know the call. + RESULT_EXPIRED = "result_expired" + CALL_NOT_FOUND = "call_not_found" + # The output failed validation against its result model. + INVALID_RESULT = "invalid_result" + # A comparison step whose baseline or reform simulation run failed. + SIMULATION_RUN_FAILED = "simulation_run_failed" + + +@dataclass(frozen=True) +class RunError: + """A failure to record on a run; the message must carry no user data.""" + + code: RunErrorCode + message: str + + +@dataclass(frozen=True) +class DueSimulationCall: + """A running Modal call whose runs this instance claimed to check now.""" + + modal_call_id: str + runs: tuple[SimulationRunRecord, ...] + + +@dataclass(frozen=True) +class AbandonedDispatches: + """Runs and comparison steps failed because their dispatch never finished.""" + + simulation_run_ids: tuple[UUID, ...] + report_run_ids: tuple[UUID, ...] + + +class RunTransitions: + """Conditional status writes; each returns whether it changed anything.""" + + def __init__( + self, + engine: Engine, + *, + clock: Callable[[], datetime] = utc_now, + ) -> None: + """Use ``engine`` (the runtime role) and ``clock`` for every timestamp.""" + self.engine = engine + self._clock = clock + + # --- Simulation runs --------------------------------------------------- + + def _change_simulation_runs( + self, + run_ids: Collection[UUID], + *, + expected: RunStatus, + values: dict[str, object], + ) -> bool: + """Change every run in ``run_ids`` from ``expected``, or none of them.""" + wanted = set(run_ids) + if not wanted: + msg = "a status change needs at least one run" + raise ValueError(msg) + with self.engine.connect() as connection: + transaction = connection.begin() + changed = connection.execute( + update(simulation_runs) + .where( + simulation_runs.c.id.in_(wanted), + simulation_runs.c.status == expected.value, + ) + .values(updated_at=self._clock(), **values) + .returning(simulation_runs.c.id) + ).scalars() + if set(changed) != wanted: + transaction.rollback() + return False + transaction.commit() + return True + + def claim_simulation_dispatch(self, run_ids: Collection[UUID]) -> bool: + """Move pending runs to ``dispatching``; only the winner starts the call. + + ``run_ids`` are the runs one Modal call computes: one run, or a + household report's baseline and reform. All must be pending. + """ + return self._change_simulation_runs( + run_ids, + expected=RunStatus.PENDING, + values={"status": RunStatus.DISPATCHING.value}, + ) + + def record_simulation_dispatched( + self, + run_ids: Collection[UUID], + *, + modal_call_id: str, + first_check_at: datetime, + ) -> bool: + """Record the acknowledged Modal call and move the runs to ``running``.""" + return self._change_simulation_runs( + run_ids, + expected=RunStatus.DISPATCHING, + values={ + "status": RunStatus.RUNNING.value, + "modal_call_id": modal_call_id, + "dispatched_at": self._clock(), + "next_check_at": first_check_at, + }, + ) + + def record_simulation_dispatch_failure( + self, run_ids: Collection[UUID], error: RunError + ) -> bool: + """Fail dispatching runs whose Modal call was refused.""" + return self._change_simulation_runs( + run_ids, + expected=RunStatus.DISPATCHING, + values={ + "status": RunStatus.FAILED.value, + "error_code": error.code.value, + "error_message": bounded_error_message(error.message), + "completed_at": self._clock(), + }, + ) + + def simulation_runs_hold_call( + self, run_ids: Collection[UUID], modal_call_id: str + ) -> bool: + """Return whether every one of the runs recorded ``modal_call_id``. + + A dispatcher whose record raised reads this to learn whether the + record committed all the same, its reply having been lost. + """ + wanted = set(run_ids) + with self.engine.connect() as connection: + holding = connection.execute( + select(func.count()) + .select_from(simulation_runs) + .where( + simulation_runs.c.id.in_(wanted), + simulation_runs.c.modal_call_id == modal_call_id, + ) + ).scalar_one() + return bool(wanted) and holding == len(wanted) + + def claim_due_simulation_calls( + self, *, limit: int, lease: timedelta + ) -> list[DueSimulationCall]: + """Claim up to ``limit`` running runs that are due a check, by call. + + Each claimed run's next check moves ``lease`` ahead, so no other + instance checks it until the lease ends; if this instance stops, the + run becomes due again then. Rows another instance holds are skipped. + Two runs of one call are normally claimed together; if the limit + splits them, two instances may ask Modal about the same call, and the + conditional writes still record its outcome once. + """ + now = self._clock() + due = ( + select(simulation_runs.c.id) + .where( + simulation_runs.c.status == RunStatus.RUNNING.value, + simulation_runs.c.next_check_at <= now, + ) + .order_by(simulation_runs.c.next_check_at) + .limit(limit) + .with_for_update(skip_locked=True) + .scalar_subquery() + ) + with self.engine.begin() as connection: + rows = ( + connection.execute( + update(simulation_runs) + .where( + simulation_runs.c.id.in_(due), + simulation_runs.c.status == RunStatus.RUNNING.value, + simulation_runs.c.next_check_at <= now, + ) + .values( + next_check_at=now + lease, + check_count=simulation_runs.c.check_count + 1, + updated_at=now, + ) + .returning(*simulation_runs.c) + ) + .mappings() + .all() + ) + runs_by_call: dict[str, list[SimulationRunRecord]] = defaultdict(list) + for row in rows: + run = SimulationRunRecord.from_row(row) + if run.modal_call_id is not None: + runs_by_call[run.modal_call_id].append(run) + return [ + DueSimulationCall(modal_call_id=call_id, runs=tuple(runs)) + for call_id, runs in sorted(runs_by_call.items()) + ] + + def _change_simulation_call( + self, + connection: Connection, + modal_call_id: str, + values: dict[str, object], + ) -> tuple[SimulationRunRecord, ...]: + rows = ( + connection.execute( + update(simulation_runs) + .where( + simulation_runs.c.modal_call_id == modal_call_id, + simulation_runs.c.status == RunStatus.RUNNING.value, + ) + .values(updated_at=self._clock(), **values) + .returning(*simulation_runs.c) + ) + .mappings() + .all() + ) + return tuple(SimulationRunRecord.from_row(row) for row in rows) + + def complete_simulation_call( + self, modal_call_id: str, write_results: SimulationResultWriter + ) -> bool: + """Mark a call's running runs succeeded and write their results. + + ``write_results`` receives the transaction and the succeeded runs, and + writes their result rows. If it raises, nothing is recorded and the + exception propagates. Return False, writing nothing, when no run of + the call is still running. + """ + with self.engine.begin() as connection: + runs = self._change_simulation_call( + connection, + modal_call_id, + { + "status": RunStatus.SUCCEEDED.value, + "completed_at": self._clock(), + "next_check_at": None, + }, + ) + if not runs: + return False + write_results(connection, runs) + return True + + def fail_simulation_call(self, modal_call_id: str, error: RunError) -> bool: + """Mark a call's running runs failed; False if none is still running.""" + with self.engine.begin() as connection: + return bool( + self._change_simulation_call( + connection, + modal_call_id, + { + "status": RunStatus.FAILED.value, + "error_code": error.code.value, + "error_message": bounded_error_message(error.message), + "completed_at": self._clock(), + "next_check_at": None, + }, + ) + ) + + # --- Comparison steps of economy report runs ---------------------------- + + def _change_comparison( + self, + connection: Connection, + report_run_id: UUID, + *, + expected: RunStatus, + values: dict[str, object], + conditions: tuple[ColumnElement[bool], ...] = (), + ) -> ReportRunRecord | None: + row = ( + connection.execute( + update(report_runs) + .where( + report_runs.c.id == report_run_id, + report_runs.c.comparison_status == expected.value, + *conditions, + ) + .values(updated_at=self._clock(), **values) + .returning(*report_runs.c) + ) + .mappings() + .first() + ) + return None if row is None else ReportRunRecord.from_row(row) + + def claim_comparison_dispatch(self, report_run_id: UUID) -> bool: + """Move a pending comparison step to ``dispatching``. + + Only the winner starts the comparison's Modal call. The step can start + only once both of its simulation runs have succeeded. + """ + baseline = simulation_runs.alias("baseline_run") + reform = simulation_runs.alias("reform_run") + both_runs_succeeded = ( + exists().where( + and_( + baseline.c.id == report_runs.c.baseline_simulation_run_id, + baseline.c.status == RunStatus.SUCCEEDED.value, + ) + ), + exists().where( + and_( + reform.c.id == report_runs.c.reform_simulation_run_id, + reform.c.status == RunStatus.SUCCEEDED.value, + ) + ), + ) + with self.engine.begin() as connection: + return ( + self._change_comparison( + connection, + report_run_id, + expected=RunStatus.PENDING, + values={"comparison_status": RunStatus.DISPATCHING.value}, + conditions=both_runs_succeeded, + ) + is not None + ) + + def record_comparison_dispatched( + self, + report_run_id: UUID, + *, + modal_call_id: str, + first_check_at: datetime, + ) -> bool: + """Record the comparison's acknowledged Modal call; it is now running.""" + with self.engine.begin() as connection: + return ( + self._change_comparison( + connection, + report_run_id, + expected=RunStatus.DISPATCHING, + values={ + "comparison_status": RunStatus.RUNNING.value, + "comparison_modal_call_id": modal_call_id, + "comparison_dispatched_at": self._clock(), + "comparison_next_check_at": first_check_at, + }, + ) + is not None + ) + + def record_comparison_dispatch_failure( + self, report_run_id: UUID, error: RunError + ) -> bool: + """Fail a dispatching comparison step whose Modal call was refused.""" + with self.engine.begin() as connection: + return ( + self._change_comparison( + connection, + report_run_id, + expected=RunStatus.DISPATCHING, + values=self._comparison_failure(error), + ) + is not None + ) + + def comparison_holds_call(self, report_run_id: UUID, modal_call_id: str) -> bool: + """Return whether the comparison step recorded ``modal_call_id``. + + As ``simulation_runs_hold_call``, for a comparison step. + """ + with self.engine.connect() as connection: + return connection.execute( + select( + exists().where( + report_runs.c.id == report_run_id, + report_runs.c.comparison_modal_call_id == modal_call_id, + ) + ) + ).scalar_one() + + def claim_due_comparisons( + self, *, limit: int, lease: timedelta + ) -> list[ReportRunRecord]: + """Claim up to ``limit`` running comparison steps that are due a check.""" + now = self._clock() + due = ( + select(report_runs.c.id) + .where( + report_runs.c.comparison_status == RunStatus.RUNNING.value, + report_runs.c.comparison_next_check_at <= now, + ) + .order_by(report_runs.c.comparison_next_check_at) + .limit(limit) + .with_for_update(skip_locked=True) + .scalar_subquery() + ) + with self.engine.begin() as connection: + rows = ( + connection.execute( + update(report_runs) + .where( + report_runs.c.id.in_(due), + report_runs.c.comparison_status == RunStatus.RUNNING.value, + report_runs.c.comparison_next_check_at <= now, + ) + .values( + comparison_next_check_at=now + lease, + comparison_check_count=report_runs.c.comparison_check_count + 1, + updated_at=now, + ) + .returning(*report_runs.c) + ) + .mappings() + .all() + ) + return [ReportRunRecord.from_row(row) for row in rows] + + def complete_comparison( + self, report_run_id: UUID, write_results: ComparisonResultWriter + ) -> bool: + """Mark a running comparison step succeeded and write its results. + + ``write_results`` runs in the same transaction; if it raises, nothing + is recorded and the exception propagates. + """ + with self.engine.begin() as connection: + report_run = self._change_comparison( + connection, + report_run_id, + expected=RunStatus.RUNNING, + values={ + "comparison_status": RunStatus.SUCCEEDED.value, + "comparison_completed_at": self._clock(), + "comparison_next_check_at": None, + }, + ) + if report_run is None: + return False + write_results(connection, report_run) + return True + + def fail_comparison(self, report_run_id: UUID, error: RunError) -> bool: + """Mark a running comparison step failed; False if it is not running.""" + with self.engine.begin() as connection: + return ( + self._change_comparison( + connection, + report_run_id, + expected=RunStatus.RUNNING, + values=self._comparison_failure(error), + ) + is not None + ) + + def _comparison_failure(self, error: RunError) -> dict[str, object]: + return { + "comparison_status": RunStatus.FAILED.value, + "comparison_error_code": error.code.value, + "comparison_error_message": bounded_error_message(error.message), + "comparison_completed_at": self._clock(), + "comparison_next_check_at": None, + } + + def fail_comparisons_of_failed_runs(self) -> tuple[UUID, ...]: + """Fail pending comparison steps whose baseline or reform run failed. + + Such a step can never start, so it fails with ``simulation_run_failed`` + and its report run's status says so. Return the report runs changed. + """ + baseline = simulation_runs.alias("failed_baseline_run") + reform = simulation_runs.alias("failed_reform_run") + either_run_failed = or_( + exists().where( + and_( + baseline.c.id == report_runs.c.baseline_simulation_run_id, + baseline.c.status == RunStatus.FAILED.value, + ) + ), + exists().where( + and_( + reform.c.id == report_runs.c.reform_simulation_run_id, + reform.c.status == RunStatus.FAILED.value, + ) + ), + ) + failure = self._comparison_failure( + RunError( + RunErrorCode.SIMULATION_RUN_FAILED, + "A simulation run this comparison needs failed.", + ) + ) + with self.engine.begin() as connection: + return tuple( + connection.execute( + update(report_runs) + .where( + report_runs.c.comparison_status == RunStatus.PENDING.value, + either_run_failed, + ) + .values(updated_at=self._clock(), **failure) + .returning(report_runs.c.id) + ).scalars() + ) + + # --- Dispatches that never finished ------------------------------------ + + def fail_abandoned_dispatches( + self, *, claimed_before: datetime + ) -> AbandonedDispatches: + """Fail runs and comparison steps left in ``dispatching``. + + A run stays ``dispatching`` when no Modal call was recorded for it: + the instance that claimed it stopped, or could not tell whether its + record committed and left the call running. Failing it, with + ``dispatch_abandoned``, lets the same simulation be run again. + """ + now = self._clock() + message = "No Modal call was recorded for the dispatch in time." + with self.engine.begin() as connection: + failed_runs = connection.execute( + update(simulation_runs) + .where( + simulation_runs.c.status == RunStatus.DISPATCHING.value, + simulation_runs.c.updated_at < claimed_before, + ) + .values( + status=RunStatus.FAILED.value, + error_code=RunErrorCode.DISPATCH_ABANDONED.value, + error_message=message, + completed_at=now, + updated_at=now, + ) + .returning(simulation_runs.c.id) + ).scalars() + simulation_run_ids = tuple(failed_runs) + report_run_ids = tuple( + connection.execute( + update(report_runs) + .where( + report_runs.c.comparison_status == RunStatus.DISPATCHING.value, + report_runs.c.updated_at < claimed_before, + ) + .values( + updated_at=now, + **self._comparison_failure( + RunError(RunErrorCode.DISPATCH_ABANDONED, message) + ), + ) + .returning(report_runs.c.id) + ).scalars() + ) + return AbandonedDispatches( + simulation_run_ids=simulation_run_ids, report_run_ids=report_run_ids + ) diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/runs.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/runs.py new file mode 100644 index 000000000..3ad05ea99 --- /dev/null +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/runs.py @@ -0,0 +1,487 @@ +"""``simulation_runs`` and ``report_runs``: executions of simulations and reports. + +A simulation run is one execution of a simulation for a catalog version and a +year. It records the region and dataset rows it resolved in that catalog +version, its own status and Modal call, its timings and error, what started it +(its trigger) and its attempt number. Runs are unique per simulation, catalog +version, year and attempt, so a run that already exists is reused, and a forced +re-run adds the next attempt. + +A report run is one calculation of a report: a baseline and a reform +simulation run of the same catalog version and year (the reform is null for a +single-simulation household report), plus, for an economy report, the +comparison step's own status and Modal call. A household report has no +comparison step: one Modal call computes both households, so its simulation +runs share that call. + +Each status moves ``pending`` -> ``dispatching`` -> ``running`` -> ``succeeded`` +or ``failed``. The region and dataset keys include the catalog version, through +the Public API's unique constraints ``uq_regions_model_version_code`` and +``uq_datasets_id_model_version``, so a run can use only its own version's +rows. A report run's simulation run keys include its catalog version and year, +so both runs share them. +""" + +from __future__ import annotations + +from dataclasses import dataclass, fields +from enum import StrEnum +from typing import TYPE_CHECKING, Any, Self + +from sqlalchemy import ( + CheckConstraint, + Column, + DateTime, + ForeignKeyConstraint, + Index, + Integer, + String, + Table, + Text, + UniqueConstraint, + Uuid, + text, +) + +from policyengine_simulation_persistence import public_api +from policyengine_simulation_persistence.reports import ( + FIRST_YEAR, + LAST_YEAR, + ReportKind, + reports, +) +from policyengine_simulation_persistence.schema import simulation_metadata +from policyengine_simulation_persistence.simulations import ( + CATALOG_KEY_LENGTH, + SimulationKind, + simulations, +) + +if TYPE_CHECKING: + from datetime import datetime + from uuid import UUID + + from sqlalchemy import RowMapping + +MAX_ERROR_MESSAGE_LENGTH = 2000 + + +class RunStatus(StrEnum): + """Where a simulation run, or a report run's comparison step, stands.""" + + PENDING = "pending" + DISPATCHING = "dispatching" + RUNNING = "running" + SUCCEEDED = "succeeded" + FAILED = "failed" + + +UNFINISHED_STATUSES = (RunStatus.PENDING, RunStatus.DISPATCHING, RunStatus.RUNNING) +FINISHED_STATUSES = (RunStatus.SUCCEEDED, RunStatus.FAILED) + + +class RunTrigger(StrEnum): + """What started a run. + + ``initial`` is a report's or calculation's first submission, ``manual`` a + user's re-run, and ``system`` a run sim or a migration started itself, for + example the first opening of a migrated report. + """ + + INITIAL = "initial" + MANUAL = "manual" + SYSTEM = "system" + + +def _values(enum: type[StrEnum]) -> str: + return ", ".join(f"'{member.value}'" for member in enum) + + +def _status_checks(prefix: str, *, optional: bool) -> tuple[CheckConstraint, ...]: + """Return the checks that keep one status and its columns consistent. + + ``prefix`` is empty for a simulation run and ``comparison_`` for a report + run's comparison step, whose status is null (``optional``) when the report + run has no comparison step. + """ + status = f"{prefix}status" + unless_absent = f"{status} IS NULL OR " if optional else "" + return ( + CheckConstraint( + f"{unless_absent}{status} IN ({_values(RunStatus)})", + name=f"{prefix}status", + ), + # A Modal call exists once dispatching succeeded, never before; a + # failed run may have one or not. + CheckConstraint( + f"{unless_absent}" + f"({status} IN ('pending', 'dispatching') " + f"AND {prefix}modal_call_id IS NULL) " + f"OR ({status} IN ('running', 'succeeded') " + f"AND {prefix}modal_call_id IS NOT NULL) " + f"OR {status} = 'failed'", + name=f"{prefix}modal_call", + ), + CheckConstraint( + f"{unless_absent}" + f"({status} IN ('succeeded', 'failed')) " + f"= ({prefix}completed_at IS NOT NULL)", + name=f"{prefix}completed_at", + ), + CheckConstraint( + f"{unless_absent}({status} = 'failed') = ({prefix}error_code IS NOT NULL)", + name=f"{prefix}error", + ), + # Only a running call is checked, and every running call is due a check. + CheckConstraint( + f"{unless_absent}" + f"({status} = 'running') = ({prefix}next_check_at IS NOT NULL)", + name=f"{prefix}next_check_at", + ), + CheckConstraint(f"{prefix}check_count >= 0", name=f"{prefix}check_count"), + ) + + +def _execution_columns(prefix: str, *, required: bool) -> list[Column[Any]]: + """Return the status, Modal call, error and timing columns of one execution.""" + return [ + Column(f"{prefix}status", Text, nullable=not required), + Column(f"{prefix}modal_call_id", Text, nullable=True), + Column(f"{prefix}error_code", Text, nullable=True), + Column(f"{prefix}error_message", Text, nullable=True), + Column( + f"{prefix}check_count", + Integer, + nullable=False, + server_default=text("0"), + ), + Column(f"{prefix}dispatched_at", DateTime(timezone=True), nullable=True), + Column(f"{prefix}next_check_at", DateTime(timezone=True), nullable=True), + Column(f"{prefix}completed_at", DateTime(timezone=True), nullable=True), + ] + + +simulation_runs = Table( + "simulation_runs", + simulation_metadata, + Column("id", Uuid(as_uuid=True), primary_key=True), + Column("simulation_id", Uuid(as_uuid=True), nullable=False), + Column("simulation_kind", Text, nullable=False), + Column("tax_benefit_model_version_id", Uuid(as_uuid=True), nullable=False), + Column("year", Integer, nullable=False), + Column("attempt", Integer, nullable=False), + Column("region_code", String(CATALOG_KEY_LENGTH), nullable=True), + Column("dataset_id", Uuid(as_uuid=True), nullable=True), + Column("trigger", Text, nullable=False), + *_execution_columns("", required=True), + Column("created_at", DateTime(timezone=True), nullable=False), + Column("updated_at", DateTime(timezone=True), nullable=False), + CheckConstraint(f"year BETWEEN {FIRST_YEAR} AND {LAST_YEAR}", name="year"), + CheckConstraint("attempt >= 1", name="attempt"), + CheckConstraint(f"trigger IN ({_values(RunTrigger)})", name="trigger"), + CheckConstraint( + "(simulation_kind = 'household' AND region_code IS NULL " + "AND dataset_id IS NULL) " + "OR (simulation_kind = 'economy' AND region_code IS NOT NULL " + "AND dataset_id IS NOT NULL)", + name="population_matches_kind", + ), + *_status_checks("", optional=False), + UniqueConstraint( + "simulation_id", + "tax_benefit_model_version_id", + "year", + "attempt", + name="uq_simulation_runs_attempt", + ), + # Report runs refer to a simulation run with its catalog version and year, + # so both runs of a report run share them. + UniqueConstraint( + "id", + "tax_benefit_model_version_id", + "year", + name="uq_simulation_runs_id_catalog_version_year", + ), + ForeignKeyConstraint( + ["simulation_id", "simulation_kind"], + [simulations.c.id, simulations.c.kind], + name="fk_simulation_runs_simulation", + ondelete="RESTRICT", + ), + ForeignKeyConstraint( + ["tax_benefit_model_version_id"], + [public_api.tax_benefit_model_versions.c.id], + name="fk_simulation_runs_catalog_version", + ondelete="RESTRICT", + ), + ForeignKeyConstraint( + ["tax_benefit_model_version_id", "region_code"], + [ + public_api.regions.c.tax_benefit_model_version_id, + public_api.regions.c.code, + ], + name="fk_simulation_runs_region", + ondelete="RESTRICT", + ), + ForeignKeyConstraint( + ["dataset_id", "tax_benefit_model_version_id"], + [ + public_api.datasets.c.id, + public_api.datasets.c.tax_benefit_model_version_id, + ], + name="fk_simulation_runs_dataset", + ondelete="RESTRICT", + ), +) +# A Public API delete of a catalog version, region or dataset checks these for +# references. +Index( + "ix_simulation_runs_catalog_version_region", + simulation_runs.c.tax_benefit_model_version_id, + simulation_runs.c.region_code, +) +Index("ix_simulation_runs_dataset_id", simulation_runs.c.dataset_id) +# Every run of one Modal call (a household report's two runs share one). +Index( + "ix_simulation_runs_modal_call_id", + simulation_runs.c.modal_call_id, + postgresql_where=simulation_runs.c.modal_call_id.is_not(None), +) +# A recording pass reads only running runs that are due a check, and +# dispatches that stopped. +Index( + "ix_simulation_runs_due_checks", + simulation_runs.c.next_check_at, + postgresql_where=simulation_runs.c.status == RunStatus.RUNNING.value, +) +Index( + "ix_simulation_runs_dispatching", + simulation_runs.c.updated_at, + postgresql_where=simulation_runs.c.status == RunStatus.DISPATCHING.value, +) + +_SIMULATION_RUN_KEY = [ + simulation_runs.c.id, + simulation_runs.c.tax_benefit_model_version_id, + simulation_runs.c.year, +] + +report_runs = Table( + "report_runs", + simulation_metadata, + Column("id", Uuid(as_uuid=True), primary_key=True), + Column("report_id", Uuid(as_uuid=True), nullable=False), + Column("report_kind", Text, nullable=False), + Column("tax_benefit_model_version_id", Uuid(as_uuid=True), nullable=False), + Column("year", Integer, nullable=False), + Column("baseline_simulation_run_id", Uuid(as_uuid=True), nullable=False), + Column("reform_simulation_run_id", Uuid(as_uuid=True), nullable=True), + Column("trigger", Text, nullable=False), + *_execution_columns("comparison_", required=False), + Column("created_at", DateTime(timezone=True), nullable=False), + Column("updated_at", DateTime(timezone=True), nullable=False), + CheckConstraint(f"year BETWEEN {FIRST_YEAR} AND {LAST_YEAR}", name="year"), + CheckConstraint(f"trigger IN ({_values(RunTrigger)})", name="trigger"), + # An economy report run compares two simulation runs in its own comparison + # step; a household report run has no comparison step. + CheckConstraint( + "(report_kind = 'economy' AND reform_simulation_run_id IS NOT NULL " + "AND comparison_status IS NOT NULL) " + "OR (report_kind = 'household' AND comparison_status IS NULL " + "AND comparison_modal_call_id IS NULL " + "AND comparison_error_code IS NULL " + "AND comparison_dispatched_at IS NULL " + "AND comparison_next_check_at IS NULL " + "AND comparison_completed_at IS NULL " + "AND comparison_check_count = 0)", + name="comparison_matches_kind", + ), + *_status_checks("comparison_", optional=True), + ForeignKeyConstraint( + ["report_id", "report_kind"], + [reports.c.id, reports.c.kind], + name="fk_report_runs_report", + ondelete="RESTRICT", + ), + ForeignKeyConstraint( + [ + "baseline_simulation_run_id", + "tax_benefit_model_version_id", + "year", + ], + _SIMULATION_RUN_KEY, + name="fk_report_runs_baseline_simulation_run", + ondelete="RESTRICT", + ), + ForeignKeyConstraint( + ["reform_simulation_run_id", "tax_benefit_model_version_id", "year"], + _SIMULATION_RUN_KEY, + name="fk_report_runs_reform_simulation_run", + ondelete="RESTRICT", + ), +) +# A report's status is its newest report run's. +Index( + "ix_report_runs_report_id_created_at", + report_runs.c.report_id, + report_runs.c.created_at, +) +Index( + "ix_report_runs_baseline_simulation_run_id", + report_runs.c.baseline_simulation_run_id, +) +Index( + "ix_report_runs_reform_simulation_run_id", + report_runs.c.reform_simulation_run_id, +) +Index( + "ix_report_runs_due_comparison_checks", + report_runs.c.comparison_next_check_at, + postgresql_where=report_runs.c.comparison_status == RunStatus.RUNNING.value, +) +# Each recording pass looks for pending comparison steps whose run failed. +Index( + "ix_report_runs_pending_comparisons", + report_runs.c.id, + postgresql_where=report_runs.c.comparison_status == RunStatus.PENDING.value, +) +Index( + "ix_report_runs_comparisons_dispatching", + report_runs.c.updated_at, + postgresql_where=report_runs.c.comparison_status == RunStatus.DISPATCHING.value, +) + + +def bounded_error_message(message: str) -> str: + """Return ``message`` cut to the length the run tables store.""" + return message[:MAX_ERROR_MESSAGE_LENGTH] + + +@dataclass(frozen=True) +class NewSimulationRun: + """A simulation run to find or create. + + An economy run names the region and dataset rows its caller resolved in + the run's catalog version: the simulation's region code, and the dataset + the simulation names or, when it names none, the region's default dataset + in that version. The database checks that both belong to that version. A + household run names neither. + """ + + simulation_id: UUID + tax_benefit_model_version_id: UUID + year: int + trigger: RunTrigger + region_code: str | None = None + dataset_id: UUID | None = None + + def __post_init__(self) -> None: + """Reject a year outside the supported range without touching the database.""" + if not FIRST_YEAR <= self.year <= LAST_YEAR: + msg = f"year must be between {FIRST_YEAR} and {LAST_YEAR}" + raise ValueError(msg) + + +@dataclass(frozen=True) +class NewReportRun: + """A report run to create from two simulation runs that already exist. + + The baseline run must be a run of the report's baseline simulation, and + the reform run one of its reform simulation; an economy report run needs + both. They must share a catalog version and year, which the report run + takes from them. + """ + + report_id: UUID + baseline_simulation_run_id: UUID + trigger: RunTrigger + reform_simulation_run_id: UUID | None = None + + +def _record_values(cls: type, row: RowMapping) -> dict[str, Any]: + return {field.name: row[field.name] for field in fields(cls)} + + +@dataclass(frozen=True) +class SimulationRunRecord: + """One ``simulation.simulation_runs`` row.""" + + id: UUID + simulation_id: UUID + simulation_kind: SimulationKind + tax_benefit_model_version_id: UUID + year: int + attempt: int + region_code: str | None + dataset_id: UUID | None + trigger: RunTrigger + status: RunStatus + modal_call_id: str | None + error_code: str | None + error_message: str | None + check_count: int + dispatched_at: datetime | None + next_check_at: datetime | None + completed_at: datetime | None + created_at: datetime + updated_at: datetime + + @property + def finished(self) -> bool: + """Whether the run has succeeded or failed.""" + return self.status in FINISHED_STATUSES + + @classmethod + def from_row(cls, row: RowMapping) -> Self: + """Build the record from a selected ``simulation_runs`` row.""" + values = _record_values(cls, row) + values["simulation_kind"] = SimulationKind(values["simulation_kind"]) + values["trigger"] = RunTrigger(values["trigger"]) + values["status"] = RunStatus(values["status"]) + return cls(**values) + + +@dataclass(frozen=True) +class ReportRunRecord: + """One ``simulation.report_runs`` row. + + The comparison fields are None for a household report run, which has no + comparison step. + """ + + id: UUID + report_id: UUID + report_kind: ReportKind + tax_benefit_model_version_id: UUID + year: int + baseline_simulation_run_id: UUID + reform_simulation_run_id: UUID | None + trigger: RunTrigger + comparison_status: RunStatus | None + comparison_modal_call_id: str | None + comparison_error_code: str | None + comparison_error_message: str | None + comparison_check_count: int + comparison_dispatched_at: datetime | None + comparison_next_check_at: datetime | None + comparison_completed_at: datetime | None + created_at: datetime + updated_at: datetime + + @classmethod + def from_row(cls, row: RowMapping) -> Self: + """Build the record from a selected ``report_runs`` row.""" + values = _record_values(cls, row) + values["report_kind"] = ReportKind(values["report_kind"]) + values["trigger"] = RunTrigger(values["trigger"]) + if values["comparison_status"] is not None: + values["comparison_status"] = RunStatus(values["comparison_status"]) + return cls(**values) + + +@dataclass(frozen=True) +class StartedSimulationRun: + """The simulation run a start resolved to, and whether this call created it.""" + + run: SimulationRunRecord + created: bool diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/schema.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/schema.py index 7e4080638..7ef7081c6 100644 --- a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/schema.py +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/schema.py @@ -17,7 +17,7 @@ # The newest revision in migrations/versions. Each migration PR updates it, and # tests/test_schema.py fails when it differs from the chain's head. The # readiness check reports it next to the database's revision. -HEAD_REVISION = "0002" +HEAD_REVISION = "0003" NAMING_CONVENTION = { "ix": "ix_%(table_name)s_%(column_0_N_name)s", diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/simulations.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/simulations.py index 0ad0e7629..76f134c26 100644 --- a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/simulations.py +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/simulations.py @@ -106,6 +106,9 @@ class SimulationKind(StrEnum): # Reports refer to a simulation together with its country and kind, so a # report can only compare simulations of its own country and kind. UniqueConstraint("id", "country", "kind"), + # Simulation runs refer to a simulation together with its kind, so a run + # names a region and dataset exactly when its simulation is an economy. + UniqueConstraint("id", "kind"), ForeignKeyConstraint( ["policy_id", "country"], [public_api.policies.c.id, public_api.policies.c.country_id], diff --git a/libs/policyengine-simulation-persistence/tests/public_api_support.py b/libs/policyengine-simulation-persistence/tests/public_api_support.py index cf144c68d..2ace7b50b 100644 --- a/libs/policyengine-simulation-persistence/tests/public_api_support.py +++ b/libs/policyengine-simulation-persistence/tests/public_api_support.py @@ -21,6 +21,24 @@ from sqlalchemy import Engine PUBLIC_API_TABLE_STATEMENTS = ( + "CREATE TABLE public.tax_benefit_model_versions (id uuid PRIMARY KEY)", + "CREATE TABLE public.datasets (" + " id uuid PRIMARY KEY," + " name varchar(255) NOT NULL," + " tax_benefit_model_version_id uuid NOT NULL" + " REFERENCES public.tax_benefit_model_versions (id)," + " CONSTRAINT uq_datasets_model_version_name" + " UNIQUE (tax_benefit_model_version_id, name)," + " CONSTRAINT uq_datasets_id_model_version" + " UNIQUE (id, tax_benefit_model_version_id))", + "CREATE TABLE public.regions (" + " id uuid PRIMARY KEY," + " code varchar(255) NOT NULL," + " tax_benefit_model_version_id uuid NOT NULL" + " REFERENCES public.tax_benefit_model_versions (id)," + " default_dataset_id uuid NOT NULL," + " CONSTRAINT uq_regions_model_version_code" + " UNIQUE (tax_benefit_model_version_id, code))", "CREATE TABLE public.policies (" " id uuid PRIMARY KEY," " country_id varchar(2) NOT NULL," @@ -31,7 +49,14 @@ " CONSTRAINT uq_households_id_country UNIQUE (id, country_id))", "CREATE TABLE public.dynamics (id uuid PRIMARY KEY)", ) -REFERENCED_TABLES = ("policies", "households", "dynamics") +REFERENCED_TABLES = ( + "tax_benefit_model_versions", + "datasets", + "regions", + "policies", + "households", + "dynamics", +) def superuser_engine(superuser_url: str, **options: object) -> Engine: @@ -65,13 +90,57 @@ def create_public_api_tables(superuser_url: str, *, grant_references: bool) -> N @dataclass(frozen=True) class PublicApiRows: - """Public API rows a simulation can refer to.""" + """Public API rows a simulation or run can refer to. + + Two catalog versions each have California (``state/ca``) with its own + default dataset, as each publication creates new rows. + """ us_policy_id: UUID uk_policy_id: UUID us_household_id: UUID uk_household_id: UUID dynamic_id: UUID + catalog_version_id: UUID + newer_catalog_version_id: UUID + dataset_id: UUID + newer_dataset_id: UUID + + +def _catalog_statements(rows: PublicApiRows) -> tuple[tuple[str, dict], ...]: + return ( + ( + "INSERT INTO public.tax_benefit_model_versions (id) VALUES " + "(:version), (:newer_version)", + { + "version": rows.catalog_version_id, + "newer_version": rows.newer_catalog_version_id, + }, + ), + ( + "INSERT INTO public.datasets (id, name, tax_benefit_model_version_id) " + "VALUES (:dataset, 'populace_us_2024', :version), " + "(:newer_dataset, 'populace_us_2024_acs_local', :newer_version)", + { + "dataset": rows.dataset_id, + "newer_dataset": rows.newer_dataset_id, + "version": rows.catalog_version_id, + "newer_version": rows.newer_catalog_version_id, + }, + ), + ( + "INSERT INTO public.regions " + "(id, code, tax_benefit_model_version_id, default_dataset_id) VALUES " + "(gen_random_uuid(), 'state/ca', :version, :dataset), " + "(gen_random_uuid(), 'state/ca', :newer_version, :newer_dataset)", + { + "version": rows.catalog_version_id, + "newer_version": rows.newer_catalog_version_id, + "dataset": rows.dataset_id, + "newer_dataset": rows.newer_dataset_id, + }, + ), + ) def insert_public_api_rows(superuser_url: str) -> PublicApiRows: @@ -82,8 +151,13 @@ def insert_public_api_rows(superuser_url: str) -> PublicApiRows: us_household_id=uuid4(), uk_household_id=uuid4(), dynamic_id=uuid4(), + catalog_version_id=uuid4(), + newer_catalog_version_id=uuid4(), + dataset_id=uuid4(), + newer_dataset_id=uuid4(), ) statements = ( + *_catalog_statements(rows), ( "INSERT INTO public.policies (id, country_id) VALUES " "(:us_policy, 'us'), (:uk_policy, 'uk')", diff --git a/libs/policyengine-simulation-persistence/tests/run_support.py b/libs/policyengine-simulation-persistence/tests/run_support.py new file mode 100644 index 000000000..99a2427bb --- /dev/null +++ b/libs/policyengine-simulation-persistence/tests/run_support.py @@ -0,0 +1,85 @@ +"""Builders and direct status changes for the run tests.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from policyengine_simulation_persistence import ( + NewSimulationRun, + RunStore, + RunTrigger, + create_simulation_engine, +) +from public_api_support import superuser_engine +from sqlalchemy import text + +if TYPE_CHECKING: + from uuid import UUID + + from policyengine_simulation_persistence import SimulationRecord + from public_api_support import PublicApiRows + + +def run_store(database) -> RunStore: + """Return a store that connects as the runtime role.""" + return RunStore(create_simulation_engine(database.runtime_url)) + + +def economy_run( + simulation: SimulationRecord, rows: PublicApiRows, **overrides: object +) -> NewSimulationRun: + """Return a 2026 run of an economy simulation on the first catalog version.""" + values: dict[str, object] = { + "simulation_id": simulation.id, + "tax_benefit_model_version_id": rows.catalog_version_id, + "year": 2026, + "trigger": RunTrigger.INITIAL, + "region_code": "state/ca", + "dataset_id": rows.dataset_id, + } + values.update(overrides) + return NewSimulationRun(**values) # type: ignore[arg-type] + + +def household_run( + simulation: SimulationRecord, rows: PublicApiRows, **overrides: object +) -> NewSimulationRun: + """Return a 2026 run of a household simulation on the first catalog version.""" + values: dict[str, object] = { + "simulation_id": simulation.id, + "tax_benefit_model_version_id": rows.catalog_version_id, + "year": 2026, + "trigger": RunTrigger.INITIAL, + } + values.update(overrides) + return NewSimulationRun(**values) # type: ignore[arg-type] + + +def set_simulation_run_status( + database, run_id: UUID, status: str, *, completed_days_ago: int = 0 +) -> None: + """Move a run to a status directly, as the superuser, for read tests.""" + finished = status in {"succeeded", "failed"} + engine = superuser_engine(database.superuser_url) + try: + with engine.begin() as connection: + connection.execute( + text( + "UPDATE simulation.simulation_runs SET status = :status, " + "modal_call_id = CASE WHEN :status IN ('running', 'succeeded') " + "THEN 'fc-test' ELSE modal_call_id END, " + "error_code = CASE WHEN :status = 'failed' " + "THEN 'test_failure' ELSE NULL END, " + "completed_at = CASE WHEN :finished " + "THEN now() - make_interval(days => :days) ELSE NULL END " + "WHERE id = :id" + ), + { + "status": status, + "finished": finished, + "days": completed_days_ago, + "id": run_id, + }, + ) + finally: + engine.dispose() diff --git a/libs/policyengine-simulation-persistence/tests/test_recording_load.py b/libs/policyengine-simulation-persistence/tests/test_recording_load.py new file mode 100644 index 000000000..28b096c68 --- /dev/null +++ b/libs/policyengine-simulation-persistence/tests/test_recording_load.py @@ -0,0 +1,223 @@ +"""A load check of the conditional writes that record outcomes, on PostgreSQL. + +Several simulated writers, each on its own engine, repeatedly fail stalled +dispatches, claim due calls and record an outcome for each, as safety-check +passes do, until every run has finished. Some outcomes are written late, +after another writer could claim the same call and record it, as a report +from Modal can arrive while a pass checks the same call. A third of the calls +are household pairs (two runs under one call). Every call must be recorded +once, its results written once and all of its runs finished once, and every +stalled dispatch failed once, whatever the interleaving. The elapsed time and +the number of passes are printed, as a measurement; nothing depends on them. +The same check of the reporting route against the staging database is an +operator step. +""" + +from __future__ import annotations + +import random +import time +from concurrent.futures import ThreadPoolExecutor +from datetime import timedelta +from threading import Barrier, Lock + +import pytest +from policyengine_simulation_persistence import ( + RunError, + RunErrorCode, + RunStatus, + RunTransitions, + create_simulation_engine, +) +from policyengine_simulation_persistence.runs import simulation_runs +from run_support import household_run +from simulation_support import household_simulation, simulation_store +from sqlalchemy import func, select +from transition_support import ManualClock, clocked_stores + +pytestmark = pytest.mark.postgres + +CALLS = 480 +# Every third call is a household report's pair: two runs under one call. +PAIR_EVERY = 3 +# Runs claimed by an instance that stopped before recording their call. +STALLED_DISPATCHES = 30 +INSTANCES = 6 +CLAIM_LIMIT = 20 +LEASE = timedelta(minutes=1) +ABANDONED_AFTER = timedelta(minutes=10) +MAXIMUM_PASSES = 400 + + +def _household_run(stores, simulations, rows, legacy_id: int): + simulation = simulations.create( + household_simulation(rows, legacy_simulation_id=legacy_id) + ) + return stores.runs.start_simulation_run(household_run(simulation, rows)).run + + +def _start_running_calls(database, rows, clock) -> dict[str, int]: + """Create running calls, a third of them household pairs, and stalled runs. + + Return each call's number of runs. + """ + stores = clocked_stores(database, clock) + simulations = simulation_store(database) + legacy_ids = iter(range(1, 10_000)) + runs_per_call = {} + for index in range(CALLS): + count = 2 if index % PAIR_EVERY == 0 else 1 + run_ids = [ + _household_run(stores, simulations, rows, next(legacy_ids)).id + for _ in range(count) + ] + call_id = f"fc-load-{index}" + stores.transitions.claim_simulation_dispatch(run_ids) + stores.transitions.record_simulation_dispatched( + run_ids, modal_call_id=call_id, first_check_at=clock() + ) + runs_per_call[call_id] = count + for _ in range(STALLED_DISPATCHES): + run = _household_run(stores, simulations, rows, next(legacy_ids)) + stores.transitions.claim_simulation_dispatch([run.id]) + stores.transitions.engine.dispose() + return runs_per_call + + +class _Outcomes: + """Remember every recorded outcome, from every instance.""" + + def __init__(self) -> None: + self.lock = Lock() + self.written_calls: list[str] = [] + self.failed_calls: list[str] = [] + self.abandoned_runs: list[object] = [] + + def write(self, call_id: str): + def write_results(_connection, _runs) -> None: + with self.lock: + self.written_calls.append(call_id) + + return write_results + + +def _record(transitions, outcomes, call_id: str, *, succeed: bool) -> None: + if succeed: + transitions.complete_simulation_call(call_id, outcomes.write(call_id)) + elif transitions.fail_simulation_call( + call_id, RunError(RunErrorCode.WORKER_ERROR, "The Modal call raised.") + ): + with outcomes.lock: + outcomes.failed_calls.append(call_id) + + +def _instance(database, clock, outcomes, barrier, seed: int) -> int: + """Claim and record due calls until none is left; return the pass count. + + Some outcomes are recorded only in a later pass, after an instance with + nothing to claim has moved the clock past every lease, so another + instance may claim the same call and record it first: the late write + must then change nothing. + """ + chooser = random.Random(seed) # noqa: S311 # test interleaving, not security + transitions = RunTransitions( + create_simulation_engine(database.runtime_url), clock=clock + ) + late: list[tuple[str, bool]] = [] + barrier.wait() + passes = 0 + try: + while passes < MAXIMUM_PASSES: + passes += 1 + for call_id, succeed in late: + _record(transitions, outcomes, call_id, succeed=succeed) + late.clear() + abandoned = transitions.fail_abandoned_dispatches( + claimed_before=clock() - ABANDONED_AFTER + ) + with outcomes.lock: + outcomes.abandoned_runs.extend(abandoned.simulation_run_ids) + calls = transitions.claim_due_simulation_calls( + limit=CLAIM_LIMIT, lease=LEASE + ) + if not calls: + if ( + not _unfinished(transitions) + and len(outcomes.abandoned_runs) == STALLED_DISPATCHES + ): + return passes + clock.advance(LEASE) + continue + for call in calls: + outcome = chooser.random() + if outcome < 0.3: + # Still running: checked again once the claim's lease ends. + continue + if outcome < 0.5: + late.append((call.modal_call_id, outcome < 0.45)) + else: + _record( + transitions, outcomes, call.modal_call_id, succeed=outcome < 0.9 + ) + return passes + finally: + transitions.engine.dispose() + + +def _unfinished(transitions: RunTransitions) -> int: + with transitions.engine.connect() as connection: + return connection.execute( + select(func.count()).where( + simulation_runs.c.status == RunStatus.RUNNING.value + ) + ).scalar_one() + + +def test_concurrent_instances_record_every_call_once( + migrated_database, public_api_rows, capsys +) -> None: + clock = ManualClock() + runs_per_call = _start_running_calls(migrated_database, public_api_rows, clock) + outcomes = _Outcomes() + barrier = Barrier(INSTANCES) + started = time.monotonic() + + with ThreadPoolExecutor(max_workers=INSTANCES) as pool: + passes = list( + pool.map( + lambda seed: _instance( + migrated_database, clock, outcomes, barrier, seed + ), + range(INSTANCES), + ) + ) + + elapsed = time.monotonic() - started + recorded = outcomes.written_calls + outcomes.failed_calls + assert sorted(recorded) == sorted(runs_per_call) + assert len(outcomes.abandoned_runs) == len(set(outcomes.abandoned_runs)) + assert len(outcomes.abandoned_runs) == STALLED_DISPATCHES + stores = clocked_stores(migrated_database, clock) + with stores.transitions.engine.connect() as connection: + statuses = dict( + connection.execute( + select(simulation_runs.c.status, func.count()).group_by( + simulation_runs.c.status + ) + ).all() + ) + stores.transitions.engine.dispose() + assert statuses == { + RunStatus.SUCCEEDED.value: sum( + runs_per_call[call_id] for call_id in outcomes.written_calls + ), + RunStatus.FAILED.value: STALLED_DISPATCHES + + sum(runs_per_call[call_id] for call_id in outcomes.failed_calls), + } + with capsys.disabled(): + print( # noqa: T201 # the measurement this check records + f"\nrecording load check: {len(runs_per_call)} calls " + f"({sum(runs_per_call.values())} runs), {STALLED_DISPATCHES} stalled " + f"dispatches, {INSTANCES} instances, {sum(passes)} passes, " + f"{elapsed:.2f} s" + ) diff --git a/libs/policyengine-simulation-persistence/tests/test_run_transitions.py b/libs/policyengine-simulation-persistence/tests/test_run_transitions.py new file mode 100644 index 000000000..e54fce729 --- /dev/null +++ b/libs/policyengine-simulation-persistence/tests/test_run_transitions.py @@ -0,0 +1,546 @@ +"""Conditional status changes: a second writer never changes anything.""" + +from __future__ import annotations + +from concurrent.futures import ThreadPoolExecutor +from datetime import timedelta +from threading import Barrier + +import pytest +from policyengine_simulation_persistence import ( + NewReportRun, + ReportKind, + RunError, + RunErrorCode, + RunStatus, + RunTransitions, + RunTrigger, + create_simulation_engine, +) +from report_support import economy_pair, economy_report, household_baseline +from report_support import report_store as definition_report_store +from run_support import economy_run, household_run +from simulation_support import household_simulation, simulation_store +from transition_support import ManualClock, clocked_stores + +pytestmark = pytest.mark.postgres + +LEASE = timedelta(minutes=1) +FIRST_CHECK = timedelta(seconds=2) +WORKER_ERROR = RunError(RunErrorCode.WORKER_ERROR, "The Modal call raised KeyError.") + + +def _pending_run(stores, database, rows): + baseline, _ = economy_pair(database, rows) + return stores.runs.start_simulation_run(economy_run(baseline, rows)).run + + +def _running_run(stores, database, rows, call_id="fc-economy"): + run = _pending_run(stores, database, rows) + assert stores.transitions.claim_simulation_dispatch([run.id]) + assert stores.transitions.record_simulation_dispatched( + [run.id], modal_call_id=call_id, first_check_at=stores.clock() + FIRST_CHECK + ) + return stores.runs.get_simulation_run(run.id) + + +def _household_pair(stores, database, rows): + baseline = household_baseline(database, rows) + reform = simulation_store(database).create( + household_simulation(rows, policy_id=rows.us_policy_id) + ) + return ( + stores.runs.start_simulation_run(household_run(baseline, rows)).run, + stores.runs.start_simulation_run(household_run(reform, rows)).run, + ) + + +def _no_results(_connection, _records) -> None: + return None + + +def test_only_one_claim_moves_a_run_to_dispatching( + migrated_database, public_api_rows +) -> None: + stores = clocked_stores(migrated_database) + run = _pending_run(stores, migrated_database, public_api_rows) + + assert stores.transitions.claim_simulation_dispatch([run.id]) is True + assert stores.transitions.claim_simulation_dispatch([run.id]) is False + + claimed = stores.runs.get_simulation_run(run.id) + assert claimed.status is RunStatus.DISPATCHING + assert claimed.modal_call_id is None + + +def test_concurrent_claims_start_one_modal_call( + migrated_database, public_api_rows +) -> None: + stores = clocked_stores(migrated_database) + run = _pending_run(stores, migrated_database, public_api_rows) + instances = 8 + barrier = Barrier(instances) + + def claim(_instance: int) -> bool: + transitions = RunTransitions( + create_simulation_engine(migrated_database.runtime_url) + ) + try: + barrier.wait() + return transitions.claim_simulation_dispatch([run.id]) + finally: + transitions.engine.dispose() + + with ThreadPoolExecutor(max_workers=instances) as pool: + outcomes = list(pool.map(claim, range(instances))) + + assert outcomes.count(True) == 1 + + +def test_a_dispatched_run_records_its_call_and_first_check( + migrated_database, public_api_rows +) -> None: + stores = clocked_stores(migrated_database) + run = _running_run(stores, migrated_database, public_api_rows) + + assert run.status is RunStatus.RUNNING + assert run.modal_call_id == "fc-economy" + assert run.dispatched_at == stores.clock() + assert run.next_check_at == stores.clock() + FIRST_CHECK + assert not stores.transitions.record_simulation_dispatched( + [run.id], modal_call_id="fc-other", first_check_at=stores.clock() + ) + + +def test_runs_say_whether_they_recorded_a_call( + migrated_database, public_api_rows +) -> None: + """A dispatcher whose record raised reads whether it committed all the same.""" + stores = clocked_stores(migrated_database) + running = _running_run(stores, migrated_database, public_api_rows) + dispatching = _pending_run(stores, migrated_database, public_api_rows) + stores.transitions.claim_simulation_dispatch([dispatching.id]) + transitions = stores.transitions + + assert transitions.simulation_runs_hold_call([running.id], "fc-economy") + assert not transitions.simulation_runs_hold_call([running.id], "fc-other") + assert not transitions.simulation_runs_hold_call([dispatching.id], "fc-economy") + assert not transitions.simulation_runs_hold_call( + [running.id, dispatching.id], "fc-economy" + ) + assert not transitions.simulation_runs_hold_call([], "fc-economy") + + +def test_a_refused_dispatch_fails_the_run_and_frees_the_setup( + migrated_database, public_api_rows +) -> None: + stores = clocked_stores(migrated_database) + run = _pending_run(stores, migrated_database, public_api_rows) + stores.transitions.claim_simulation_dispatch([run.id]) + + assert stores.transitions.record_simulation_dispatch_failure( + [run.id], RunError(RunErrorCode.DISPATCH_FAILED, "Modal refused the call.") + ) + + failed = stores.runs.get_simulation_run(run.id) + assert failed.status is RunStatus.FAILED + assert failed.error_code == "dispatch_failed" + assert failed.completed_at == stores.clock() + simulation = simulation_store(migrated_database).get(run.simulation_id) + retry = stores.runs.start_simulation_run(economy_run(simulation, public_api_rows)) + assert retry.created is True + assert retry.run.attempt == 2 + + +def test_a_household_pair_shares_one_call_and_changes_together( + migrated_database, public_api_rows +) -> None: + stores = clocked_stores(migrated_database) + baseline_run, reform_run = _household_pair( + stores, migrated_database, public_api_rows + ) + pair = [baseline_run.id, reform_run.id] + written = [] + + assert stores.transitions.claim_simulation_dispatch(pair) + assert stores.transitions.record_simulation_dispatched( + pair, modal_call_id="fc-household", first_check_at=stores.clock() + ) + assert stores.transitions.complete_simulation_call( + "fc-household", lambda _connection, runs: written.extend(runs) + ) + + assert {run.id for run in written} == set(pair) + for run_id in pair: + run = stores.runs.get_simulation_run(run_id) + assert run.status is RunStatus.SUCCEEDED + assert run.modal_call_id == "fc-household" + assert run.next_check_at is None + + +def test_a_group_claim_changes_nothing_unless_every_run_is_pending( + migrated_database, public_api_rows +) -> None: + stores = clocked_stores(migrated_database) + baseline_run, reform_run = _household_pair( + stores, migrated_database, public_api_rows + ) + assert stores.transitions.claim_simulation_dispatch([baseline_run.id]) + + assert not stores.transitions.claim_simulation_dispatch( + [baseline_run.id, reform_run.id] + ) + + assert stores.runs.get_simulation_run(reform_run.id).status is RunStatus.PENDING + with pytest.raises(ValueError, match="at least one run"): + stores.transitions.claim_simulation_dispatch([]) + + +def test_due_calls_are_claimed_for_one_lease_at_a_time( + migrated_database, public_api_rows +) -> None: + stores = clocked_stores(migrated_database) + run = _running_run(stores, migrated_database, public_api_rows) + transitions = stores.transitions + + assert transitions.claim_due_simulation_calls(limit=10, lease=LEASE) == [] + stores.clock.advance(FIRST_CHECK) + claimed = transitions.claim_due_simulation_calls(limit=10, lease=LEASE) + assert [call.modal_call_id for call in claimed] == ["fc-economy"] + assert claimed[0].runs[0].id == run.id + assert claimed[0].runs[0].check_count == 1 + assert claimed[0].runs[0].next_check_at == stores.clock() + LEASE + assert transitions.claim_due_simulation_calls(limit=10, lease=LEASE) == [] + + # The lease ends when the claiming instance stops before recording anything. + stores.clock.advance(LEASE) + reclaimed = transitions.claim_due_simulation_calls(limit=10, lease=LEASE) + assert reclaimed[0].runs[0].check_count == 2 + + +def test_concurrent_instances_claim_disjoint_calls( + migrated_database, public_api_rows +) -> None: + clock = ManualClock() + stores = clocked_stores(migrated_database, clock) + simulations = simulation_store(migrated_database) + run_ids = set() + for index in range(12): + simulation = simulations.create( + household_simulation(public_api_rows, legacy_simulation_id=index + 1) + ) + run = stores.runs.start_simulation_run( + household_run(simulation, public_api_rows) + ).run + stores.transitions.claim_simulation_dispatch([run.id]) + stores.transitions.record_simulation_dispatched( + [run.id], modal_call_id=f"fc-{index}", first_check_at=clock() + ) + run_ids.add(run.id) + instances = 4 + barrier = Barrier(instances) + + def claim(_instance: int) -> list[str]: + transitions = RunTransitions( + create_simulation_engine(migrated_database.runtime_url), clock=clock + ) + try: + barrier.wait() + return [ + call.modal_call_id + for call in transitions.claim_due_simulation_calls(limit=5, lease=LEASE) + ] + finally: + transitions.engine.dispose() + + with ThreadPoolExecutor(max_workers=instances) as pool: + claims = list(pool.map(claim, range(instances))) + + claimed_calls = [call for instance_claims in claims for call in instance_claims] + assert len(claimed_calls) == len(set(claimed_calls)) == len(run_ids) + + +def test_results_are_written_in_the_transaction_that_marks_success( + migrated_database, public_api_rows +) -> None: + stores = clocked_stores(migrated_database) + run = _running_run(stores, migrated_database, public_api_rows) + + def failing_writer(_connection, _runs) -> None: + msg = "the result rows could not be written" + raise RuntimeError(msg) + + with pytest.raises(RuntimeError, match="result rows"): + stores.transitions.complete_simulation_call("fc-economy", failing_writer) + + assert stores.runs.get_simulation_run(run.id).status is RunStatus.RUNNING + assert stores.transitions.complete_simulation_call("fc-economy", _no_results) + succeeded = stores.runs.get_simulation_run(run.id) + assert succeeded.status is RunStatus.SUCCEEDED + assert succeeded.completed_at == stores.clock() + + +def test_a_finished_call_is_recorded_once(migrated_database, public_api_rows) -> None: + stores = clocked_stores(migrated_database) + _running_run(stores, migrated_database, public_api_rows) + calls = [] + + assert stores.transitions.complete_simulation_call("fc-economy", _no_results) + + assert not stores.transitions.complete_simulation_call( + "fc-economy", lambda _connection, runs: calls.append(runs) + ) + assert calls == [] + assert not stores.transitions.fail_simulation_call("fc-economy", WORKER_ERROR) + + +def test_a_failed_call_records_a_bounded_error( + migrated_database, public_api_rows +) -> None: + stores = clocked_stores(migrated_database) + run = _running_run(stores, migrated_database, public_api_rows) + + assert stores.transitions.fail_simulation_call( + "fc-economy", RunError(RunErrorCode.WORKER_ERROR, "x" * 5000) + ) + + failed = stores.runs.get_simulation_run(run.id) + assert failed.status is RunStatus.FAILED + assert failed.error_code == "worker_error" + assert len(failed.error_message) == 2000 + assert failed.next_check_at is None + + +def _economy_report_run(stores, database, rows): + baseline, reform = economy_pair(database, rows) + report = definition_report_store(database).create(economy_report(baseline, reform)) + runs = [ + stores.runs.start_simulation_run(economy_run(simulation, rows)).run + for simulation in (baseline, reform) + ] + report_run = stores.runs.create_report_run( + NewReportRun( + report_id=report.id, + baseline_simulation_run_id=runs[0].id, + reform_simulation_run_id=runs[1].id, + trigger=RunTrigger.INITIAL, + ) + ) + return report_run, runs + + +def _succeed(stores, run, call_id) -> None: + stores.transitions.claim_simulation_dispatch([run.id]) + stores.transitions.record_simulation_dispatched( + [run.id], modal_call_id=call_id, first_check_at=stores.clock() + ) + assert stores.transitions.complete_simulation_call(call_id, _no_results) + + +def test_a_comparison_starts_only_after_both_runs_succeed( + migrated_database, public_api_rows +) -> None: + stores = clocked_stores(migrated_database) + report_run, (baseline_run, reform_run) = _economy_report_run( + stores, migrated_database, public_api_rows + ) + + assert not stores.transitions.claim_comparison_dispatch(report_run.id) + _succeed(stores, baseline_run, "fc-baseline") + assert not stores.transitions.claim_comparison_dispatch(report_run.id) + _succeed(stores, reform_run, "fc-reform") + + assert stores.transitions.claim_comparison_dispatch(report_run.id) + assert not stores.transitions.claim_comparison_dispatch(report_run.id) + assert ( + stores.runs.get_report_run(report_run.id).comparison_status + is RunStatus.DISPATCHING + ) + + +def test_a_comparison_runs_and_records_its_results_once( + migrated_database, public_api_rows +) -> None: + stores = clocked_stores(migrated_database) + report_run, (baseline_run, reform_run) = _economy_report_run( + stores, migrated_database, public_api_rows + ) + _succeed(stores, baseline_run, "fc-baseline") + _succeed(stores, reform_run, "fc-reform") + transitions = stores.transitions + transitions.claim_comparison_dispatch(report_run.id) + written = [] + + assert transitions.record_comparison_dispatched( + report_run.id, modal_call_id="fc-comparison", first_check_at=stores.clock() + ) + claimed = transitions.claim_due_comparisons(limit=10, lease=LEASE) + assert [due.id for due in claimed] == [report_run.id] + assert claimed[0].comparison_check_count == 1 + assert transitions.claim_due_comparisons(limit=10, lease=LEASE) == [] + assert transitions.complete_comparison( + report_run.id, lambda _connection, record: written.append(record) + ) + + assert [record.id for record in written] == [report_run.id] + finished = stores.runs.get_report_run(report_run.id) + assert finished.comparison_status is RunStatus.SUCCEEDED + assert finished.comparison_modal_call_id == "fc-comparison" + assert finished.comparison_next_check_at is None + assert not transitions.complete_comparison(report_run.id, _no_results) + assert not transitions.fail_comparison(report_run.id, WORKER_ERROR) + + +def test_a_comparison_failure_is_recorded(migrated_database, public_api_rows) -> None: + stores = clocked_stores(migrated_database) + report_run, (baseline_run, reform_run) = _economy_report_run( + stores, migrated_database, public_api_rows + ) + _succeed(stores, baseline_run, "fc-baseline") + _succeed(stores, reform_run, "fc-reform") + transitions = stores.transitions + transitions.claim_comparison_dispatch(report_run.id) + transitions.record_comparison_dispatched( + report_run.id, modal_call_id="fc-comparison", first_check_at=stores.clock() + ) + + assert transitions.fail_comparison(report_run.id, WORKER_ERROR) + + failed = stores.runs.get_report_run(report_run.id) + assert failed.comparison_status is RunStatus.FAILED + assert failed.comparison_error_code == "worker_error" + assert failed.comparison_completed_at == stores.clock() + + +def test_a_comparison_says_whether_it_recorded_a_call( + migrated_database, public_api_rows +) -> None: + stores = clocked_stores(migrated_database) + report_run, (baseline_run, reform_run) = _economy_report_run( + stores, migrated_database, public_api_rows + ) + _succeed(stores, baseline_run, "fc-baseline") + _succeed(stores, reform_run, "fc-reform") + transitions = stores.transitions + transitions.claim_comparison_dispatch(report_run.id) + + assert not transitions.comparison_holds_call(report_run.id, "fc-comparison") + + transitions.record_comparison_dispatched( + report_run.id, modal_call_id="fc-comparison", first_check_at=stores.clock() + ) + + assert transitions.comparison_holds_call(report_run.id, "fc-comparison") + assert not transitions.comparison_holds_call(report_run.id, "fc-other") + + +def test_a_refused_comparison_dispatch_fails_the_step( + migrated_database, public_api_rows +) -> None: + stores = clocked_stores(migrated_database) + report_run, (baseline_run, reform_run) = _economy_report_run( + stores, migrated_database, public_api_rows + ) + _succeed(stores, baseline_run, "fc-baseline") + _succeed(stores, reform_run, "fc-reform") + stores.transitions.claim_comparison_dispatch(report_run.id) + + assert stores.transitions.record_comparison_dispatch_failure( + report_run.id, RunError(RunErrorCode.DISPATCH_FAILED, "Modal refused it.") + ) + + failed = stores.runs.get_report_run(report_run.id) + assert failed.comparison_status is RunStatus.FAILED + assert failed.comparison_error_code == "dispatch_failed" + assert failed.comparison_modal_call_id is None + + +def test_household_report_runs_have_no_comparison_to_claim( + migrated_database, public_api_rows +) -> None: + stores = clocked_stores(migrated_database) + baseline_run, reform_run = _household_pair( + stores, migrated_database, public_api_rows + ) + baseline = simulation_store(migrated_database).get(baseline_run.simulation_id) + reform = simulation_store(migrated_database).get(reform_run.simulation_id) + report = definition_report_store(migrated_database).create( + economy_report(baseline, reform, kind=ReportKind.HOUSEHOLD) + ) + report_run = stores.runs.create_report_run( + NewReportRun( + report_id=report.id, + baseline_simulation_run_id=baseline_run.id, + reform_simulation_run_id=reform_run.id, + trigger=RunTrigger.INITIAL, + ) + ) + + assert not stores.transitions.claim_comparison_dispatch(report_run.id) + + +def test_dispatches_left_for_ten_minutes_are_failed( + migrated_database, public_api_rows +) -> None: + stores = clocked_stores(migrated_database) + report_run, (baseline_run, reform_run) = _economy_report_run( + stores, migrated_database, public_api_rows + ) + _succeed(stores, baseline_run, "fc-baseline") + _succeed(stores, reform_run, "fc-reform") + stores.transitions.claim_comparison_dispatch(report_run.id) + stalled = _pending_run(stores, migrated_database, public_api_rows) + stores.transitions.claim_simulation_dispatch([stalled.id]) + stores.clock.advance(timedelta(minutes=10, seconds=1)) + recent = household_run( + household_baseline(migrated_database, public_api_rows), public_api_rows + ) + recent_run = stores.runs.start_simulation_run(recent).run + stores.transitions.claim_simulation_dispatch([recent_run.id]) + + abandoned = stores.transitions.fail_abandoned_dispatches( + claimed_before=stores.clock() - timedelta(minutes=10) + ) + + assert abandoned.simulation_run_ids == (stalled.id,) + assert abandoned.report_run_ids == (report_run.id,) + failed_run = stores.runs.get_simulation_run(stalled.id) + assert failed_run.status is RunStatus.FAILED + assert failed_run.error_code == "dispatch_abandoned" + # The claimer may have stopped, or may not know whether its record + # committed; the message says only what is known. + assert failed_run.error_message == ( + "No Modal call was recorded for the dispatch in time." + ) + assert ( + stores.runs.get_report_run(report_run.id).comparison_error_code + == "dispatch_abandoned" + ) + assert stores.runs.get_simulation_run(recent_run.id).status is RunStatus.DISPATCHING + + +def test_a_comparison_of_a_failed_run_fails_too( + migrated_database, public_api_rows +) -> None: + stores = clocked_stores(migrated_database) + report_run, (baseline_run, reform_run) = _economy_report_run( + stores, migrated_database, public_api_rows + ) + waiting, _ = _economy_report_run(stores, migrated_database, public_api_rows) + _succeed(stores, baseline_run, "fc-baseline") + stores.transitions.claim_simulation_dispatch([reform_run.id]) + stores.transitions.record_simulation_dispatched( + [reform_run.id], modal_call_id="fc-reform", first_check_at=stores.clock() + ) + stores.transitions.fail_simulation_call("fc-reform", WORKER_ERROR) + + failed = stores.transitions.fail_comparisons_of_failed_runs() + + assert failed == (report_run.id,) + record = stores.runs.get_report_run(report_run.id) + assert record.comparison_status is RunStatus.FAILED + assert record.comparison_error_code == "simulation_run_failed" + assert stores.runs.get_report_run(waiting.id).comparison_status is ( + RunStatus.PENDING + ) + assert stores.transitions.fail_comparisons_of_failed_runs() == () diff --git a/libs/policyengine-simulation-persistence/tests/test_runs.py b/libs/policyengine-simulation-persistence/tests/test_runs.py new file mode 100644 index 000000000..6ad236bf5 --- /dev/null +++ b/libs/policyengine-simulation-persistence/tests/test_runs.py @@ -0,0 +1,495 @@ +"""simulation_runs and report_runs: executions of simulations and reports.""" + +from __future__ import annotations + +from datetime import timedelta +from uuid import uuid4 + +import pytest +from policyengine_simulation_persistence import ( + SUCCEEDED_RUN_REUSE_WINDOW, + NewReportRun, + ReportKind, + ReportRunRecord, + RunStatus, + RunTrigger, + SimulationRunRecord, + create_simulation_engine, +) +from public_api_support import delete_public_api_row +from report_support import ( + economy_pair, + economy_report, + household_baseline, + report_store, +) +from run_support import ( + economy_run, + household_run, + run_store, + set_simulation_run_status, +) +from simulation_support import ( + economy_simulation, + household_simulation, + simulation_store, +) +from sqlalchemy import text +from sqlalchemy.exc import IntegrityError, ProgrammingError + +pytestmark = pytest.mark.postgres + + +def test_an_economy_run_records_the_region_and_dataset_it_resolved( + migrated_database, public_api_rows +) -> None: + _, reform = economy_pair(migrated_database, public_api_rows) + store = run_store(migrated_database) + + started = store.start_simulation_run(economy_run(reform, public_api_rows)) + + run = started.run + assert started.created is True + assert isinstance(run, SimulationRunRecord) + assert run.simulation_id == reform.id + assert run.simulation_kind == "economy" + assert run.tax_benefit_model_version_id == public_api_rows.catalog_version_id + assert run.year == 2026 + assert run.attempt == 1 + assert run.region_code == "state/ca" + assert run.dataset_id == public_api_rows.dataset_id + assert run.status is RunStatus.PENDING + assert run.trigger is RunTrigger.INITIAL + assert run.modal_call_id is None + assert run.completed_at is None + assert store.get_simulation_run(run.id) == run + + +def test_a_run_that_exists_is_reused(migrated_database, public_api_rows) -> None: + baseline, _ = economy_pair(migrated_database, public_api_rows) + store = run_store(migrated_database) + first = store.start_simulation_run(economy_run(baseline, public_api_rows)) + + again = store.start_simulation_run( + economy_run(baseline, public_api_rows, trigger=RunTrigger.MANUAL) + ) + + assert again.created is False + assert again.run == first.run + + +def test_a_forced_rerun_adds_the_next_attempt( + migrated_database, public_api_rows +) -> None: + baseline, _ = economy_pair(migrated_database, public_api_rows) + store = run_store(migrated_database) + first = store.start_simulation_run(economy_run(baseline, public_api_rows)) + + forced = store.start_simulation_run( + economy_run(baseline, public_api_rows, trigger=RunTrigger.MANUAL), force=True + ) + + assert forced.created is True + assert forced.run.attempt == first.run.attempt + 1 + assert forced.run.trigger is RunTrigger.MANUAL + latest = store.latest_simulation_run( + baseline.id, public_api_rows.catalog_version_id, 2026 + ) + assert latest == forced.run + + +def test_a_failed_run_is_not_reused(migrated_database, public_api_rows) -> None: + baseline, _ = economy_pair(migrated_database, public_api_rows) + store = run_store(migrated_database) + first = store.start_simulation_run(economy_run(baseline, public_api_rows)) + set_simulation_run_status(migrated_database, first.run.id, "failed") + + again = store.start_simulation_run(economy_run(baseline, public_api_rows)) + + assert again.created is True + assert again.run.attempt == 2 + + +@pytest.mark.parametrize( + ("completed_days_ago", "reused"), + [(27, True), (29, False)], + ids=["within-28-days", "older-than-28-days"], +) +def test_a_succeeded_run_is_reused_only_within_28_days( + migrated_database, public_api_rows, completed_days_ago, reused +) -> None: + baseline, _ = economy_pair(migrated_database, public_api_rows) + store = run_store(migrated_database) + first = store.start_simulation_run(economy_run(baseline, public_api_rows)) + set_simulation_run_status( + migrated_database, + first.run.id, + "succeeded", + completed_days_ago=completed_days_ago, + ) + + again = store.start_simulation_run(economy_run(baseline, public_api_rows)) + + assert again.created is not reused + assert again.run.attempt == (1 if reused else 2) + assert timedelta(days=28) == SUCCEEDED_RUN_REUSE_WINDOW + + +def test_each_catalog_version_and_year_has_its_own_runs( + migrated_database, public_api_rows +) -> None: + baseline, _ = economy_pair(migrated_database, public_api_rows) + store = run_store(migrated_database) + first = store.start_simulation_run(economy_run(baseline, public_api_rows)) + + newer = store.start_simulation_run( + economy_run( + baseline, + public_api_rows, + tax_benefit_model_version_id=public_api_rows.newer_catalog_version_id, + dataset_id=public_api_rows.newer_dataset_id, + ) + ) + next_year = store.start_simulation_run( + economy_run(baseline, public_api_rows, year=2027) + ) + + assert len({first.run.id, newer.run.id, next_year.run.id}) == 3 + assert newer.created is True + assert next_year.created is True + assert newer.run.attempt == next_year.run.attempt == 1 + + +@pytest.mark.parametrize( + ("overrides", "constraint"), + [ + ({"dataset_attribute": "newer_dataset_id"}, "fk_simulation_runs_dataset"), + ({"region_code": None}, "population_matches_kind"), + ({"dataset_attribute": None}, "population_matches_kind"), + ], + ids=[ + "dataset-of-another-version", + "economy-without-region", + "economy-without-dataset", + ], +) +def test_an_economy_run_uses_only_its_own_versions_dataset( + migrated_database, public_api_rows, overrides, constraint +) -> None: + baseline, _ = economy_pair(migrated_database, public_api_rows) + values = dict(overrides) + if "dataset_attribute" in values: + attribute = values.pop("dataset_attribute") + values["dataset_id"] = ( + None if attribute is None else getattr(public_api_rows, attribute) + ) + + with pytest.raises(IntegrityError, match=constraint): + run_store(migrated_database).start_simulation_run( + economy_run(baseline, public_api_rows, **values) + ) + + +def test_an_economy_run_needs_its_region_in_its_catalog_version( + migrated_database, public_api_rows +) -> None: + new_york = simulation_store(migrated_database).create( + economy_simulation(public_api_rows, region_code="state/ny") + ) + + with pytest.raises(IntegrityError, match="fk_simulation_runs_region"): + run_store(migrated_database).start_simulation_run( + economy_run(new_york, public_api_rows, region_code="state/ny") + ) + + +def test_an_economy_run_uses_its_simulations_region( + migrated_database, public_api_rows +) -> None: + baseline, _ = economy_pair(migrated_database, public_api_rows) + + with pytest.raises(ValueError, match="region_code"): + run_store(migrated_database).start_simulation_run( + economy_run(baseline, public_api_rows, region_code="state/ny") + ) + + +def test_a_household_run_has_no_region_or_dataset( + migrated_database, public_api_rows +) -> None: + simulation = household_baseline(migrated_database, public_api_rows) + store = run_store(migrated_database) + + started = store.start_simulation_run(household_run(simulation, public_api_rows)) + assert started.run.simulation_kind == "household" + assert started.run.region_code is None + + with pytest.raises(IntegrityError, match="population_matches_kind"): + store.start_simulation_run( + household_run( + simulation, + public_api_rows, + year=2027, + region_code="state/ca", + dataset_id=public_api_rows.dataset_id, + ) + ) + + +def test_a_run_of_an_unknown_simulation_is_refused( + migrated_database, public_api_rows +) -> None: + baseline, _ = economy_pair(migrated_database, public_api_rows) + + with pytest.raises(LookupError, match="simulation"): + run_store(migrated_database).start_simulation_run( + economy_run(baseline, public_api_rows, simulation_id=uuid4()) + ) + + +def test_the_public_api_cannot_delete_a_dataset_a_run_used( + migrated_database, public_api_rows +) -> None: + baseline, _ = economy_pair(migrated_database, public_api_rows) + run_store(migrated_database).start_simulation_run( + economy_run(baseline, public_api_rows) + ) + + with pytest.raises(IntegrityError, match="violates foreign key constraint"): + delete_public_api_row( + migrated_database.superuser_url, "datasets", public_api_rows.dataset_id + ) + + +def _economy_report_with_runs(database, rows): + baseline, reform = economy_pair(database, rows) + report = report_store(database).create(economy_report(baseline, reform)) + store = run_store(database) + baseline_run = store.start_simulation_run(economy_run(baseline, rows)).run + reform_run = store.start_simulation_run(economy_run(reform, rows)).run + return report, baseline_run, reform_run + + +def test_an_economy_report_run_joins_two_runs_and_waits_for_its_comparison( + migrated_database, public_api_rows +) -> None: + report, baseline_run, reform_run = _economy_report_with_runs( + migrated_database, public_api_rows + ) + store = run_store(migrated_database) + + report_run = store.create_report_run( + NewReportRun( + report_id=report.id, + baseline_simulation_run_id=baseline_run.id, + reform_simulation_run_id=reform_run.id, + trigger=RunTrigger.INITIAL, + ) + ) + + assert isinstance(report_run, ReportRunRecord) + assert report_run.report_kind == "economy" + assert report_run.tax_benefit_model_version_id == ( + public_api_rows.catalog_version_id + ) + assert report_run.year == 2026 + assert report_run.comparison_status is RunStatus.PENDING + assert report_run.comparison_modal_call_id is None + assert store.get_report_run(report_run.id) == report_run + assert store.latest_report_run(report.id) == report_run + + +def test_both_runs_of_a_report_run_share_its_version_and_year( + migrated_database, public_api_rows +) -> None: + report, baseline_run, _ = _economy_report_with_runs( + migrated_database, public_api_rows + ) + store = run_store(migrated_database) + reform = report_store(migrated_database).get(report.id) + assert reform is not None + later_reform_run = store.start_simulation_run( + economy_run( + simulation_store(migrated_database).get(reform.reform_simulation_id), + public_api_rows, + year=2027, + ) + ).run + + with pytest.raises(IntegrityError, match="fk_report_runs_reform_simulation_run"): + store.create_report_run( + NewReportRun( + report_id=report.id, + baseline_simulation_run_id=baseline_run.id, + reform_simulation_run_id=later_reform_run.id, + trigger=RunTrigger.INITIAL, + ) + ) + + +def test_a_report_run_uses_only_runs_of_the_reports_simulations( + migrated_database, public_api_rows +) -> None: + report, baseline_run, reform_run = _economy_report_with_runs( + migrated_database, public_api_rows + ) + + with pytest.raises(ValueError, match="report's simulations"): + run_store(migrated_database).create_report_run( + NewReportRun( + report_id=report.id, + baseline_simulation_run_id=reform_run.id, + reform_simulation_run_id=baseline_run.id, + trigger=RunTrigger.INITIAL, + ) + ) + + +def test_an_economy_report_run_needs_its_reform_run( + migrated_database, public_api_rows +) -> None: + report, baseline_run, _ = _economy_report_with_runs( + migrated_database, public_api_rows + ) + + with pytest.raises(ValueError, match="report's simulations"): + run_store(migrated_database).create_report_run( + NewReportRun( + report_id=report.id, + baseline_simulation_run_id=baseline_run.id, + trigger=RunTrigger.INITIAL, + ) + ) + + +def test_a_household_report_run_has_no_comparison_step( + migrated_database, public_api_rows +) -> None: + store = simulation_store(migrated_database) + baseline = household_baseline(migrated_database, public_api_rows) + reform = store.create( + household_simulation(public_api_rows, policy_id=public_api_rows.us_policy_id) + ) + report = report_store(migrated_database).create( + economy_report(baseline, reform, kind=ReportKind.HOUSEHOLD) + ) + runs = run_store(migrated_database) + baseline_run = runs.start_simulation_run( + household_run(baseline, public_api_rows) + ).run + reform_run = runs.start_simulation_run(household_run(reform, public_api_rows)).run + + report_run = runs.create_report_run( + NewReportRun( + report_id=report.id, + baseline_simulation_run_id=baseline_run.id, + reform_simulation_run_id=reform_run.id, + trigger=RunTrigger.SYSTEM, + ) + ) + + assert report_run.report_kind == "household" + assert report_run.comparison_status is None + assert report_run.trigger is RunTrigger.SYSTEM + + +def test_a_single_simulation_household_report_run_has_no_reform_run( + migrated_database, public_api_rows +) -> None: + baseline = household_baseline(migrated_database, public_api_rows) + report = report_store(migrated_database).create( + economy_report( + baseline, baseline, kind=ReportKind.HOUSEHOLD, reform_simulation_id=None + ) + ) + runs = run_store(migrated_database) + baseline_run = runs.start_simulation_run( + household_run(baseline, public_api_rows) + ).run + + report_run = runs.create_report_run( + NewReportRun( + report_id=report.id, + baseline_simulation_run_id=baseline_run.id, + trigger=RunTrigger.INITIAL, + ) + ) + + assert report_run.reform_simulation_run_id is None + + +def test_the_newest_report_run_is_the_reports_current_run( + migrated_database, public_api_rows +) -> None: + report, baseline_run, reform_run = _economy_report_with_runs( + migrated_database, public_api_rows + ) + store = run_store(migrated_database) + new_report_run = NewReportRun( + report_id=report.id, + baseline_simulation_run_id=baseline_run.id, + reform_simulation_run_id=reform_run.id, + trigger=RunTrigger.INITIAL, + ) + store.create_report_run(new_report_run) + + rerun = store.create_report_run( + NewReportRun( + report_id=report.id, + baseline_simulation_run_id=baseline_run.id, + reform_simulation_run_id=reform_run.id, + trigger=RunTrigger.MANUAL, + ) + ) + + assert store.latest_report_run(report.id) == rerun + assert store.latest_report_run(uuid4()) is None + + +def test_a_report_run_of_an_unknown_report_is_refused( + migrated_database, public_api_rows +) -> None: + _, baseline_run, reform_run = _economy_report_with_runs( + migrated_database, public_api_rows + ) + + with pytest.raises(LookupError, match="report"): + run_store(migrated_database).create_report_run( + NewReportRun( + report_id=uuid4(), + baseline_simulation_run_id=baseline_run.id, + reform_simulation_run_id=reform_run.id, + trigger=RunTrigger.INITIAL, + ) + ) + + +@pytest.mark.parametrize("table_name", ["simulation_runs", "report_runs"]) +def test_the_runtime_role_updates_runs_but_cannot_delete_them( + migrated_database, public_api_rows, table_name +) -> None: + report, baseline_run, reform_run = _economy_report_with_runs( + migrated_database, public_api_rows + ) + run_store(migrated_database).create_report_run( + NewReportRun( + report_id=report.id, + baseline_simulation_run_id=baseline_run.id, + reform_simulation_run_id=reform_run.id, + trigger=RunTrigger.INITIAL, + ) + ) + engine = create_simulation_engine(migrated_database.runtime_url) + try: + with engine.begin() as connection: + updated = connection.execute( + text(f"UPDATE simulation.{table_name} SET updated_at = now()") # noqa: S608 # fixed table names + ) + assert updated.rowcount >= 1 + with ( + pytest.raises(ProgrammingError, match="permission denied"), + engine.begin() as connection, + ): + connection.execute(text(f"DELETE FROM simulation.{table_name}")) # noqa: S608 # fixed table names + finally: + engine.dispose() diff --git a/libs/policyengine-simulation-persistence/tests/test_schema.py b/libs/policyengine-simulation-persistence/tests/test_schema.py index 8ac5b7b05..4a6f42383 100644 --- a/libs/policyengine-simulation-persistence/tests/test_schema.py +++ b/libs/policyengine-simulation-persistence/tests/test_schema.py @@ -116,4 +116,9 @@ def test_importing_the_package_registers_every_table() -> None: importlib.import_module(f"policyengine_simulation_persistence.{module.name}") assert set(simulation_metadata.tables) == package_tables - assert {"simulation.simulations", "simulation.reports"} <= package_tables + assert { + "simulation.simulations", + "simulation.reports", + "simulation.simulation_runs", + "simulation.report_runs", + } <= package_tables diff --git a/libs/policyengine-simulation-persistence/tests/transition_support.py b/libs/policyengine-simulation-persistence/tests/transition_support.py new file mode 100644 index 000000000..426007b8f --- /dev/null +++ b/libs/policyengine-simulation-persistence/tests/transition_support.py @@ -0,0 +1,49 @@ +"""A clock the tests move by hand, and stores that share it.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from datetime import UTC, datetime, timedelta + +from policyengine_simulation_persistence import ( + RunStore, + RunTransitions, + create_simulation_engine, +) + + +@dataclass +class ManualClock: + """Return a fixed time until the test advances it.""" + + now: datetime = field( + default_factory=lambda: datetime(2026, 10, 22, 12, 0, tzinfo=UTC) + ) + + def __call__(self) -> datetime: + """Return the current fake time.""" + return self.now + + def advance(self, interval: timedelta) -> None: + """Move the fake time forward by ``interval``.""" + self.now += interval + + +@dataclass(frozen=True) +class ClockedStores: + """A run store and transitions on one runtime engine and one clock.""" + + runs: RunStore + transitions: RunTransitions + clock: ManualClock + + +def clocked_stores(database, clock: ManualClock | None = None) -> ClockedStores: + """Return stores that connect as the runtime role and share ``clock``.""" + clock = clock or ManualClock() + engine = create_simulation_engine(database.runtime_url) + return ClockedStores( + runs=RunStore(engine, clock=clock), + transitions=RunTransitions(engine, clock=clock), + clock=clock, + ) diff --git a/migrations/versions/0003_simulation_runs_and_report_runs.py b/migrations/versions/0003_simulation_runs_and_report_runs.py new file mode 100644 index 000000000..f7a2b5d43 --- /dev/null +++ b/migrations/versions/0003_simulation_runs_and_report_runs.py @@ -0,0 +1,336 @@ +"""Add simulation runs and report runs. + +``simulation_runs`` records each execution of a simulation for a catalog +version and year: the region and dataset rows it resolved in that version, its +status, Modal call, timings, error, trigger and attempt number, unique per +simulation, catalog version, year and attempt. Composite foreign keys through +the Public API's ``uq_regions_model_version_code`` and +``uq_datasets_id_model_version`` keep a run on its own version's rows, and need +``REFERENCES`` on ``tax_benefit_model_versions``, ``regions`` and ``datasets`` +for the migrator role. ``report_runs`` records each calculation of a report: +two simulation runs of one catalog version and year and, for an economy +report, the comparison step's status and Modal call. The runtime role may read, +insert and update both tables. Nothing writes them yet. + +Revision ID: 0003 +Revises: 0002 +Create Date: 2026-10-22 +""" + +from __future__ import annotations + +import sqlalchemy as sa +from alembic import op + +revision = "0003" +down_revision = "0002" +branch_labels = None +depends_on = None + +SCHEMA = "simulation" +STATUSES = "'pending', 'dispatching', 'running', 'succeeded', 'failed'" +TRIGGERS = "'initial', 'manual', 'system'" + + +def _public_api(table_and_column: str) -> str: + return f"public.{table_and_column}" + + +def _execution_columns(prefix: str, *, required: bool) -> list[sa.Column]: + return [ + sa.Column(f"{prefix}status", sa.Text(), nullable=not required), + sa.Column(f"{prefix}modal_call_id", sa.Text(), nullable=True), + sa.Column(f"{prefix}error_code", sa.Text(), nullable=True), + sa.Column(f"{prefix}error_message", sa.Text(), nullable=True), + sa.Column( + f"{prefix}check_count", + sa.Integer(), + server_default=sa.text("0"), + nullable=False, + ), + sa.Column(f"{prefix}dispatched_at", sa.DateTime(timezone=True), nullable=True), + sa.Column(f"{prefix}next_check_at", sa.DateTime(timezone=True), nullable=True), + sa.Column(f"{prefix}completed_at", sa.DateTime(timezone=True), nullable=True), + ] + + +def _status_checks( + table: str, prefix: str, *, optional: bool +) -> list[sa.CheckConstraint]: + status = f"{prefix}status" + unless_absent = f"{status} IS NULL OR " if optional else "" + return [ + sa.CheckConstraint( + f"{unless_absent}{status} IN ({STATUSES})", + name=op.f(f"ck_{table}_{prefix}status"), + ), + sa.CheckConstraint( + f"{unless_absent}" + f"({status} IN ('pending', 'dispatching') " + f"AND {prefix}modal_call_id IS NULL) " + f"OR ({status} IN ('running', 'succeeded') " + f"AND {prefix}modal_call_id IS NOT NULL) " + f"OR {status} = 'failed'", + name=op.f(f"ck_{table}_{prefix}modal_call"), + ), + sa.CheckConstraint( + f"{unless_absent}" + f"({status} IN ('succeeded', 'failed')) " + f"= ({prefix}completed_at IS NOT NULL)", + name=op.f(f"ck_{table}_{prefix}completed_at"), + ), + sa.CheckConstraint( + f"{unless_absent}({status} = 'failed') = ({prefix}error_code IS NOT NULL)", + name=op.f(f"ck_{table}_{prefix}error"), + ), + sa.CheckConstraint( + f"{unless_absent}" + f"({status} = 'running') = ({prefix}next_check_at IS NOT NULL)", + name=op.f(f"ck_{table}_{prefix}next_check_at"), + ), + sa.CheckConstraint( + f"{prefix}check_count >= 0", name=op.f(f"ck_{table}_{prefix}check_count") + ), + ] + + +def _create_simulation_runs() -> None: + op.create_unique_constraint( + "uq_simulations_id_kind", "simulations", ["id", "kind"], schema=SCHEMA + ) + op.create_table( + "simulation_runs", + sa.Column("id", sa.Uuid(), nullable=False), + sa.Column("simulation_id", sa.Uuid(), nullable=False), + sa.Column("simulation_kind", sa.Text(), nullable=False), + sa.Column("tax_benefit_model_version_id", sa.Uuid(), nullable=False), + sa.Column("year", sa.Integer(), nullable=False), + sa.Column("attempt", sa.Integer(), nullable=False), + sa.Column("region_code", sa.String(length=255), nullable=True), + sa.Column("dataset_id", sa.Uuid(), nullable=True), + sa.Column("trigger", sa.Text(), nullable=False), + *_execution_columns("", required=True), + sa.Column("created_at", sa.DateTime(timezone=True), nullable=False), + sa.Column("updated_at", sa.DateTime(timezone=True), nullable=False), + sa.CheckConstraint( + "year BETWEEN 1900 AND 2200", name=op.f("ck_simulation_runs_year") + ), + sa.CheckConstraint("attempt >= 1", name=op.f("ck_simulation_runs_attempt")), + sa.CheckConstraint( + f"trigger IN ({TRIGGERS})", name=op.f("ck_simulation_runs_trigger") + ), + sa.CheckConstraint( + "(simulation_kind = 'household' AND region_code IS NULL " + "AND dataset_id IS NULL) " + "OR (simulation_kind = 'economy' AND region_code IS NOT NULL " + "AND dataset_id IS NOT NULL)", + name=op.f("ck_simulation_runs_population_matches_kind"), + ), + *_status_checks("simulation_runs", "", optional=False), + sa.ForeignKeyConstraint( + ["simulation_id", "simulation_kind"], + [f"{SCHEMA}.simulations.id", f"{SCHEMA}.simulations.kind"], + name="fk_simulation_runs_simulation", + ondelete="RESTRICT", + ), + sa.ForeignKeyConstraint( + ["tax_benefit_model_version_id"], + [_public_api("tax_benefit_model_versions.id")], + name="fk_simulation_runs_catalog_version", + ondelete="RESTRICT", + ), + sa.ForeignKeyConstraint( + ["tax_benefit_model_version_id", "region_code"], + [ + _public_api("regions.tax_benefit_model_version_id"), + _public_api("regions.code"), + ], + name="fk_simulation_runs_region", + ondelete="RESTRICT", + ), + sa.ForeignKeyConstraint( + ["dataset_id", "tax_benefit_model_version_id"], + [ + _public_api("datasets.id"), + _public_api("datasets.tax_benefit_model_version_id"), + ], + name="fk_simulation_runs_dataset", + ondelete="RESTRICT", + ), + sa.PrimaryKeyConstraint("id", name="pk_simulation_runs"), + sa.UniqueConstraint( + "simulation_id", + "tax_benefit_model_version_id", + "year", + "attempt", + name="uq_simulation_runs_attempt", + ), + sa.UniqueConstraint( + "id", + "tax_benefit_model_version_id", + "year", + name="uq_simulation_runs_id_catalog_version_year", + ), + schema=SCHEMA, + ) + op.create_index( + "ix_simulation_runs_catalog_version_region", + "simulation_runs", + ["tax_benefit_model_version_id", "region_code"], + schema=SCHEMA, + ) + op.create_index( + "ix_simulation_runs_dataset_id", + "simulation_runs", + ["dataset_id"], + schema=SCHEMA, + ) + op.create_index( + "ix_simulation_runs_modal_call_id", + "simulation_runs", + ["modal_call_id"], + schema=SCHEMA, + postgresql_where=sa.text("modal_call_id IS NOT NULL"), + ) + op.create_index( + "ix_simulation_runs_due_checks", + "simulation_runs", + ["next_check_at"], + schema=SCHEMA, + postgresql_where=sa.text("status = 'running'"), + ) + op.create_index( + "ix_simulation_runs_dispatching", + "simulation_runs", + ["updated_at"], + schema=SCHEMA, + postgresql_where=sa.text("status = 'dispatching'"), + ) + + +def _create_report_runs() -> None: + op.create_unique_constraint( + "uq_reports_id_kind", "reports", ["id", "kind"], schema=SCHEMA + ) + simulation_run_key = [ + f"{SCHEMA}.simulation_runs.id", + f"{SCHEMA}.simulation_runs.tax_benefit_model_version_id", + f"{SCHEMA}.simulation_runs.year", + ] + op.create_table( + "report_runs", + sa.Column("id", sa.Uuid(), nullable=False), + sa.Column("report_id", sa.Uuid(), nullable=False), + sa.Column("report_kind", sa.Text(), nullable=False), + sa.Column("tax_benefit_model_version_id", sa.Uuid(), nullable=False), + sa.Column("year", sa.Integer(), nullable=False), + sa.Column("baseline_simulation_run_id", sa.Uuid(), nullable=False), + sa.Column("reform_simulation_run_id", sa.Uuid(), nullable=True), + sa.Column("trigger", sa.Text(), nullable=False), + *_execution_columns("comparison_", required=False), + sa.Column("created_at", sa.DateTime(timezone=True), nullable=False), + sa.Column("updated_at", sa.DateTime(timezone=True), nullable=False), + sa.CheckConstraint( + "year BETWEEN 1900 AND 2200", name=op.f("ck_report_runs_year") + ), + sa.CheckConstraint( + f"trigger IN ({TRIGGERS})", name=op.f("ck_report_runs_trigger") + ), + sa.CheckConstraint( + "(report_kind = 'economy' AND reform_simulation_run_id IS NOT NULL " + "AND comparison_status IS NOT NULL) " + "OR (report_kind = 'household' AND comparison_status IS NULL " + "AND comparison_modal_call_id IS NULL " + "AND comparison_error_code IS NULL " + "AND comparison_dispatched_at IS NULL " + "AND comparison_next_check_at IS NULL " + "AND comparison_completed_at IS NULL " + "AND comparison_check_count = 0)", + name=op.f("ck_report_runs_comparison_matches_kind"), + ), + *_status_checks("report_runs", "comparison_", optional=True), + sa.ForeignKeyConstraint( + ["report_id", "report_kind"], + [f"{SCHEMA}.reports.id", f"{SCHEMA}.reports.kind"], + name="fk_report_runs_report", + ondelete="RESTRICT", + ), + sa.ForeignKeyConstraint( + ["baseline_simulation_run_id", "tax_benefit_model_version_id", "year"], + simulation_run_key, + name="fk_report_runs_baseline_simulation_run", + ondelete="RESTRICT", + ), + sa.ForeignKeyConstraint( + ["reform_simulation_run_id", "tax_benefit_model_version_id", "year"], + simulation_run_key, + name="fk_report_runs_reform_simulation_run", + ondelete="RESTRICT", + ), + sa.PrimaryKeyConstraint("id", name="pk_report_runs"), + schema=SCHEMA, + ) + op.create_index( + "ix_report_runs_report_id_created_at", + "report_runs", + ["report_id", "created_at"], + schema=SCHEMA, + ) + for column in ("baseline_simulation_run_id", "reform_simulation_run_id"): + op.create_index( + f"ix_report_runs_{column}", "report_runs", [column], schema=SCHEMA + ) + op.create_index( + "ix_report_runs_due_comparison_checks", + "report_runs", + ["comparison_next_check_at"], + schema=SCHEMA, + postgresql_where=sa.text("comparison_status = 'running'"), + ) + op.create_index( + "ix_report_runs_pending_comparisons", + "report_runs", + ["id"], + schema=SCHEMA, + postgresql_where=sa.text("comparison_status = 'pending'"), + ) + op.create_index( + "ix_report_runs_comparisons_dispatching", + "report_runs", + ["updated_at"], + schema=SCHEMA, + postgresql_where=sa.text("comparison_status = 'dispatching'"), + ) + + +def upgrade() -> None: + """Create both tables; the default privileges let the runtime role use them.""" + _create_simulation_runs() + _create_report_runs() + + +def downgrade() -> None: + """Drop the tables ``upgrade`` created, with their indexes and keys.""" + for index in ( + "ix_report_runs_pending_comparisons", + "ix_report_runs_comparisons_dispatching", + "ix_report_runs_due_comparison_checks", + "ix_report_runs_reform_simulation_run_id", + "ix_report_runs_baseline_simulation_run_id", + "ix_report_runs_report_id_created_at", + ): + op.drop_index(index, table_name="report_runs", schema=SCHEMA) + op.drop_table("report_runs", schema=SCHEMA) + op.drop_constraint("uq_reports_id_kind", "reports", schema=SCHEMA, type_="unique") + for index in ( + "ix_simulation_runs_dispatching", + "ix_simulation_runs_due_checks", + "ix_simulation_runs_modal_call_id", + "ix_simulation_runs_dataset_id", + "ix_simulation_runs_catalog_version_region", + ): + op.drop_index(index, table_name="simulation_runs", schema=SCHEMA) + op.drop_table("simulation_runs", schema=SCHEMA) + op.drop_constraint( + "uq_simulations_id_kind", "simulations", schema=SCHEMA, type_="unique" + ) diff --git a/projects/policyengine-simulation-entry/DEPLOYMENT.md b/projects/policyengine-simulation-entry/DEPLOYMENT.md index 87ddd1f5e..4f445fd98 100644 --- a/projects/policyengine-simulation-entry/DEPLOYMENT.md +++ b/projects/policyengine-simulation-entry/DEPLOYMENT.md @@ -56,6 +56,83 @@ rollback checks do not, because an older image may run against a newer schema. The ownership rules and roles are described in `libs/policyengine-simulation-persistence/README.md`. +## Recording run outcomes + +The Simulation Entrypoint records a run's outcome in the `simulation` schema +within a request; nothing runs between requests, so Cloud Run keeps no CPU +allocated outside them. Every deploy mounts sim's own Modal token as +`MODAL_TOKEN_ID` and `MODAL_TOKEN_SECRET`, whatever the Stage 12 setting: the +deploy step resolves the Secret Manager secrets named by the `beta` and `prod` +environment variables `MODAL_TOKEN_ID_SECRET_NAME` and +`MODAL_TOKEN_SECRET_SECRET_NAME` to their newest enabled versions. The runtime +service account must be able to read both secrets. The Stage 12 comparison +path's Modal calls use the same token until it is removed. + +A Modal call reports that it ended on `POST /internal/modal-calls/ended` +(`policyengine_simulation_entry.call_report_route`), with the Modal identity +token of the `report_call_ended` function as a bearer token. The entrypoint +checks the token against Modal's published keys and accepts only that +function, in a v2 app, of the workspace in the `beta`/`prod` environment +variable `MODAL_WORKSPACE_ID` and of the deployment's own Modal environment; +the deploy passes both to the runtime as +`SIMULATION_CALL_REPORT_MODAL_WORKSPACE_ID` and +`SIMULATION_CALL_REPORT_MODAL_ENVIRONMENT`, and stops when the workspace ID is +missing. The service already allows unauthenticated requests through Cloud +Run (`--allow-unauthenticated`): callers prove who they are to the +entrypoint, so Modal needs no Google credentials. + +The reports come from the v2 Modal app's `report_call_ended` function, which +the v2 functions start as their last step (`src/modal/call_reports.py` in the +executor). It posts to the entrypoint's main URL: the v2 deploy job reads the +service's stable URL (`status.url`) with +`.github/scripts/simulation-entry-main-url.sh` and passes it to `modal deploy` +as `SIMULATION_ENTRY_URL`. Only the revision that serves traffic receives +these requests. `beta` does not promote its candidate yet, so in staging the +main URL reaches an older revision until it does (plan item 1.39), and +staging calls are recorded by the safety check only once it exists there. + +## Safety check + +A call whose report never arrives is recorded by the safety check, the Cloud +Run job `-safety-check`. Each execution runs one recording +pass (`python -m policyengine_simulation_entry.safety_check`) and exits: it +fails runs left `dispatching` for ten minutes and records what Modal reports +for calls dispatched, or last checked, at least five minutes ago. Cloud +Scheduler starts it every five minutes (`*/5 * * * *`, UTC): a lost report is +recorded within about ten minutes, and each execution is a single short +container start, so the job costs little when every call reports. + +`.github/scripts/deploy-cloud-run-simulation-safety-check.sh` points the job +at one entrypoint revision. It copies that revision's image digest, service +account, the secret versions it pins for `SIMULATION_DATABASE_URL`, +`MODAL_TOKEN_ID` and `MODAL_TOKEN_SECRET`, and its plain telemetry settings, +then creates or updates the schedule. Only two places call it, so the job +only ever runs a revision that serves traffic and a rejected candidate never +records anything: + +- the promotion job, after the promoted revision passes every check; if this + step fails, the previous revision is restored and the job is pointed back + at it, and if that fails too, the schedule is paused; +- the rollback workflow, after it restores entrypoint traffic; if this step + fails, the schedule is paused, so the job never keeps running the image + traffic just moved away from. + +A revision deployed before the safety check existed (without +`SIMULATION_CALL_REPORT_MODAL_WORKSPACE_ID`) has no job code, so pointing the +job at it pauses the schedule instead. + +`beta` does not promote its candidate yet, so staging has no safety check +until it does (plan item 1.39). + +One-time operator steps, in each environment: enable the Cloud Scheduler +API; create the scheduler's service account, give it `roles/run.invoker` +on the job (or the project), and set its email in the GitHub environment +variable `SIMULATION_SAFETY_CHECK_SCHEDULER_SERVICE_ACCOUNT`; give the deploy +service account permission to deploy Cloud Run jobs and manage Cloud +Scheduler jobs, and `iam.serviceAccounts.actAs` on the scheduler's service +account. The job uses the entrypoint's runtime service account, which can +already read its secrets. + ## Candidate qualification For each eligible revision, the workflow first deploys `beta` diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/app.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/app.py index 894394d50..82d360a01 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/app.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/app.py @@ -10,6 +10,7 @@ import logging import uuid from contextlib import asynccontextmanager +from dataclasses import dataclass from typing import TYPE_CHECKING from uuid import UUID @@ -51,6 +52,11 @@ CallerAuthorizer, ) from policyengine_simulation_entry.backend import OldGatewayBackend, SimulationBackend +from policyengine_simulation_entry.call_report_route import ( + CallReportRecorder, + ModalCallerVerifier, + register_call_report_route, +) from policyengine_simulation_entry.config import Settings from policyengine_simulation_entry.entry_runtime import ( BACKEND_RESPONSE_HEADER, @@ -65,6 +71,7 @@ route_template, set_observability_header, ) +from policyengine_simulation_entry.modal_identity import ModalIdentityVerifier from policyengine_simulation_entry.schemas import ( SimulationDatabaseReadiness, SimulationEntrypointReadinessResponse, @@ -88,6 +95,7 @@ "STAGE12_MODAL_SUBMISSION_TIMEOUT_SECONDS", "ComparisonBackend", "ResponseIdentifier", + "RunServices", "SimulationDatabaseConnection", "app", "create_app", @@ -97,6 +105,20 @@ SIMULATION_DATABASE_READINESS_TIMEOUT_SECONDS = 3.0 +@dataclass(frozen=True) +class RunServices: + """The simulation database and the call-ended route's seams, for tests. + + ``create_app`` builds each one it is not given: the database from + ``SIMULATION_DATABASE_URL``, the call recorder from that database, and the + Modal identity verifier from the expected Modal caller. + """ + + database: SimulationDatabaseConnection | None = None + call_recorder: CallReportRecorder | None = None + modal_identity_verifier: ModalCallerVerifier | None = None + + def _temporary_submission_payload( report: ComparisonReportRecord, ) -> TemporaryStage12SubmissionResponse: @@ -564,15 +586,23 @@ def create_app( backend: SimulationBackend | None = None, auth_dependency: AuthenticationDependency | None = None, comparison_backend: ComparisonBackend | None = None, - simulation_database: SimulationDatabaseConnection | None = None, + run_services: RunServices | None = None, ) -> FastAPI: """Build the app with injectable auth/backend seams for hermetic tests.""" runtime_settings = settings or Settings.from_env() - runtime_database = simulation_database + services = run_services or RunServices() + runtime_database = services.database if runtime_database is None and runtime_settings.simulation_database_url: runtime_database = SimulationDatabase.from_url( runtime_settings.simulation_database_url ) + call_recorder = services.call_recorder + if call_recorder is None and isinstance(runtime_database, SimulationDatabase): + call_recorder = runtime_database.call_recorder() + modal_identity_verifier = services.modal_identity_verifier + expected_modal_caller = runtime_settings.expected_modal_caller + if modal_identity_verifier is None and expected_modal_caller is not None: + modal_identity_verifier = ModalIdentityVerifier(expected_modal_caller) @asynccontextmanager async def lifespan(_: FastAPI) -> AsyncIterator[None]: @@ -635,6 +665,9 @@ async def lifespan(_: FastAPI) -> AsyncIterator[None]: reader=simulation_reader, ) _register_service_routes(app, entry) + register_call_report_route( + app, verifier=modal_identity_verifier, recorder=call_recorder + ) return app diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/call_report_route.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/call_report_route.py new file mode 100644 index 000000000..f21e56325 --- /dev/null +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/call_report_route.py @@ -0,0 +1,161 @@ +"""The route on which sim's call-report Modal function reports that a call ended. + +``POST /internal/modal-calls/ended`` (``CALL_ENDED_PATH``), not in the public +OpenAPI document. The request carries the Modal identity token of the +``report_call_ended`` function as a bearer token, and a ``CallEndedReport`` +naming the call and its runs, never a result. The token is checked first, so +a caller without it learns nothing about the body. Then the reported call's +outcome is recorded, if Modal has one (``RunRecorder.record_reported_call``). + +Answers, each a ``CallReportAnswer``: + +- 200: ``recorded`` or ``already_recorded``; a repeated report changes + nothing. +- 503 with ``Retry-After``: a status in ``RETRY_LATER_STATUSES``; the sender + reports again after a wait, and the safety check covers a call whose + sender gives up. +- 404 ``unknown_runs`` and 409 ``other_call``: the report cannot be recorded; + the sender stops. +- 401: the token is missing or refused; 422: the body is not a report. + +Only the revision that serves traffic receives these requests, since the +sender posts to sim's main URL. +""" + +from __future__ import annotations + +import asyncio +import logging +from typing import TYPE_CHECKING, Protocol + +from fastapi import Request, status +from fastapi.responses import JSONResponse +from policyengine_simulation_contract.call_reports import ( + CALL_ENDED_PATH, + CallEndedReport, + CallReportAnswer, + CallReportStatus, +) +from pydantic import ValidationError + +from policyengine_simulation_entry.modal_identity import ( + ModalIdentityError, + ModalKeysUnavailableError, +) + +if TYPE_CHECKING: + from fastapi import FastAPI + + from policyengine_simulation_entry.modal_identity import ModalIdentity + +logger = logging.getLogger(__name__) + +# How long the sender waits before reporting again, in seconds. +RETRY_AFTER_SECONDS = 5 +_REFUSED_STATUSES = { + CallReportStatus.UNKNOWN_RUNS: status.HTTP_404_NOT_FOUND, + CallReportStatus.OTHER_CALL: status.HTTP_409_CONFLICT, +} +_BEARER_PREFIX = "bearer " + + +class CallReportRecorder(Protocol): + """Record a reported call's outcome (``RunRecorder``).""" + + async def record_reported_call(self, report: CallEndedReport) -> CallReportStatus: + """Record the call's outcome if Modal has one, and say what was done.""" + ... + + +class ModalCallerVerifier(Protocol): + """Check a Modal identity token (``ModalIdentityVerifier``).""" + + def verify(self, token: str) -> ModalIdentity: + """Return the token's identity, or raise ModalIdentityError.""" + ... + + +def _answer(report_status: CallReportStatus) -> JSONResponse: + answer = CallReportAnswer(status=report_status) + if answer.retry_later: + return JSONResponse( + answer.model_dump(mode="json"), + status_code=status.HTTP_503_SERVICE_UNAVAILABLE, + headers={"Retry-After": str(RETRY_AFTER_SECONDS)}, + ) + return JSONResponse( + answer.model_dump(mode="json"), + status_code=_REFUSED_STATUSES.get(report_status, status.HTTP_200_OK), + ) + + +def _unauthorized() -> JSONResponse: + return JSONResponse( + {"detail": "A Modal identity token of the call-report function is required."}, + status_code=status.HTTP_401_UNAUTHORIZED, + headers={"WWW-Authenticate": "Bearer"}, + ) + + +def _bearer_token(request: Request) -> str | None: + header = request.headers.get("authorization", "") + if not header.lower().startswith(_BEARER_PREFIX): + return None + return header[len(_BEARER_PREFIX) :].strip() or None + + +async def _caller( + request: Request, verifier: ModalCallerVerifier +) -> ModalIdentity | JSONResponse: + """Return the verified Modal caller, or the answer that refuses or defers it.""" + token = _bearer_token(request) + if token is None: + return _unauthorized() + try: + return await asyncio.to_thread(verifier.verify, token) + except ModalKeysUnavailableError: + logger.warning("call_report_modal_keys_unavailable") + return _answer(CallReportStatus.MODAL_UNREACHABLE) + except ModalIdentityError as error: + logger.warning("call_report_refused", extra={"reason": str(error)}) + return _unauthorized() + # Anything else is this instance's own problem, not a refusal: the sender + # reports again after the wait. + except Exception: + logger.exception("call_report_token_not_checked") + return _answer(CallReportStatus.MODAL_UNREACHABLE) + + +def register_call_report_route( + app: FastAPI, + *, + verifier: ModalCallerVerifier | None, + recorder: CallReportRecorder | None, +) -> None: + """Add the call-ended route; without both seams it answers 503.""" + + @app.post(CALL_ENDED_PATH, include_in_schema=False) + async def report_call_ended(request: Request) -> JSONResponse: + if verifier is None or recorder is None: + return _answer(CallReportStatus.NOT_CONFIGURED) + identity = await _caller(request, verifier) + if isinstance(identity, JSONResponse): + return identity + try: + report = CallEndedReport.model_validate_json(await request.body()) + except ValidationError as error: + return JSONResponse( + {"detail": error.errors(include_url=False, include_input=False)}, + status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, + ) + report_status = await recorder.record_reported_call(report) + logger.info( + "call_report_answered", + extra={ + "modal_call_id": report.modal_call_id, + "call_report_status": report_status.value, + "modal_app_name": identity.app_name, + "modal_container_id": identity.container_id, + }, + ) + return _answer(report_status) diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/config.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/config.py index 5d23e8c54..6ac11436f 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/config.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/config.py @@ -3,6 +3,7 @@ from __future__ import annotations import os +import re from dataclasses import dataclass from urllib.parse import urlparse @@ -12,10 +13,15 @@ assert_separate_manifest_names, ) +from policyengine_simulation_entry.modal_identity import ExpectedModalCaller + PRODUCTION_ENVIRONMENTS = frozenset({"main", "prod", "production"}) MODAL_GATEWAY_HOST_SUFFIX = ".modal.run" # Stage 12 comparison artifacts and rows are kept for exactly this many days. REQUIRED_STAGE12_RETENTION_DAYS = 30 +# Modal workspace IDs and environment names, as Modal prints them. +MODAL_WORKSPACE_ID_PATTERN = re.compile(r"ac-[A-Za-z0-9]+") +MODAL_ENVIRONMENT_NAME_PATTERN = re.compile(r"[a-z0-9][a-z0-9-]{0,63}") class ConfigurationError(RuntimeError): @@ -96,6 +102,10 @@ class Settings: # outside deployed environments; nothing on the old-gateway or Stage 12 # paths reads it. simulation_database_url: str = "" + # The Modal workspace and environment whose call-report function may + # report that calls ended. Both empty: the reporting route answers 503. + call_report_modal_workspace_id: str = "" + call_report_modal_environment: str = "" @classmethod def from_env(cls) -> Settings: @@ -147,6 +157,12 @@ def from_env(cls) -> Settings: stage12_artifact_bucket=os.getenv("STAGE12_ARTIFACT_BUCKET", ""), stage12_retention_days=int(os.getenv("STAGE12_RETENTION_DAYS", "30")), simulation_database_url=os.getenv("SIMULATION_DATABASE_URL", ""), + call_report_modal_workspace_id=os.getenv( + "SIMULATION_CALL_REPORT_MODAL_WORKSPACE_ID", "" + ), + call_report_modal_environment=os.getenv( + "SIMULATION_CALL_REPORT_MODAL_ENVIRONMENT", "" + ), ) @property @@ -154,6 +170,18 @@ def production(self) -> bool: """Whether this is a production environment, where auth must stay on.""" return self.environment in PRODUCTION_ENVIRONMENTS + @property + def expected_modal_caller(self) -> ExpectedModalCaller | None: + """The Modal caller whose call-ended reports are accepted, if configured.""" + if not ( + self.call_report_modal_workspace_id and self.call_report_modal_environment + ): + return None + return ExpectedModalCaller( + workspace_id=self.call_report_modal_workspace_id, + environment_name=self.call_report_modal_environment, + ) + @property def stage12_resources_configured(self) -> bool: """Whether the Stage 12 manifest environment, database and bucket are set.""" @@ -176,6 +204,7 @@ def validate(self) -> None: self._validate_old_gateway() self._validate_stage12() self._validate_simulation_database() + self._validate_call_reports() def _validate_caller_auth(self) -> None: """Require caller authentication in production, with an HTTPS issuer.""" @@ -281,3 +310,22 @@ def _validate_simulation_database(self) -> None: ): msg = "SIMULATION_DATABASE_URL must use the canonical postgresql:// scheme." raise ConfigurationError(msg) + + def _validate_call_reports(self) -> None: + """Require one Modal workspace ID and environment name, or neither.""" + workspace = self.call_report_modal_workspace_id + environment = self.call_report_modal_environment + if not workspace and not environment: + return + if not MODAL_WORKSPACE_ID_PATTERN.fullmatch(workspace): + msg = ( + "SIMULATION_CALL_REPORT_MODAL_WORKSPACE_ID must be a Modal " + "workspace ID." + ) + raise ConfigurationError(msg) + if not MODAL_ENVIRONMENT_NAME_PATTERN.fullmatch(environment): + msg = ( + "SIMULATION_CALL_REPORT_MODAL_ENVIRONMENT must be a Modal " + "environment name." + ) + raise ConfigurationError(msg) diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/modal_calls.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/modal_calls.py new file mode 100644 index 000000000..96ec899bf --- /dev/null +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/modal_calls.py @@ -0,0 +1,223 @@ +"""Start Modal calls and ask about them without waiting, for simulation runs. + +The Simulation Entrypoint starts each run's Modal call with ``spawn`` and later +asks Modal about it with ``get(timeout=0)``, which returns at once. Modal stops +a pre-empted worker with KeyboardInterrupt, records no result for it and +re-runs the same input under the same call ID, so the call reads as still +running and finishes later: pre-emption is never a failure. + +These adapters belong to the main path. The Stage 12 comparison path keeps its +own Modal code; the two share nothing. +""" + +from __future__ import annotations + +import asyncio +import traceback +from dataclasses import dataclass +from typing import Any, Protocol + +import modal +from policyengine_simulation_persistence import RunError, RunErrorCode + + +@dataclass(frozen=True) +class ModalTarget: + """A deployed Modal function: its application and function names.""" + + application: str + function: str + + +@dataclass(frozen=True) +class ModalRequest: + """One call to start: the target, its positional arguments and telemetry. + + ``arguments`` hold calculation inputs and functional metadata only. + ``observability_context`` travels as Modal's separate keyword argument. + """ + + target: ModalTarget + arguments: tuple[object, ...] + observability_context: dict[str, Any] | None = None + + +@dataclass(frozen=True) +class CallRunning: + """The call has produced no output yet, or Modal is re-running it.""" + + +@dataclass(frozen=True) +class CallSucceeded: + """The call returned ``output``.""" + + output: object + + +@dataclass(frozen=True) +class CallFailed: + """The call failed; the error carries no calculation data.""" + + error: RunError + + +type CallOutcome = CallRunning | CallSucceeded | CallFailed + + +class CallInspectionError(RuntimeError): + """Modal could not be asked about a call; ask again later.""" + + +class CallDispatcher(Protocol): + """Start a Modal call without waiting for it, or cancel one.""" + + async def spawn(self, request: ModalRequest) -> str: + """Return the call ID once Modal acknowledges the call.""" + ... + + async def cancel(self, modal_call_id: str) -> None: + """Ask Modal to stop a call whose runs could not record it.""" + ... + + +class CallInspector(Protocol): + """Ask Modal about a call without waiting for it.""" + + async def inspect(self, modal_call_id: str) -> CallOutcome: + """Return the call's outcome; raise CallInspectionError if unknown.""" + ... + + +class ModalCallDispatcher: + """Spawn deployed Modal functions of one Modal environment.""" + + def __init__(self, environment: str) -> None: + """Look functions up in ``environment``.""" + self._environment = environment + + async def spawn(self, request: ModalRequest) -> str: + """Spawn the request's function and return Modal's call ID. + + Raise RuntimeError when Modal acknowledges the call without an ID. + """ + function = modal.Function.from_name( + request.target.application, + request.target.function, + environment_name=self._environment, + ) + call = await function.spawn.aio( + *request.arguments, + observability_context=request.observability_context, + ) + modal_call_id = getattr(call, "object_id", None) + if not isinstance(modal_call_id, str) or not modal_call_id: + msg = "Modal acknowledged the call without a call ID" + raise RuntimeError(msg) + return modal_call_id + + async def cancel(self, modal_call_id: str) -> None: + """Ask Modal to stop the call; its output will never be collected.""" + await modal.FunctionCall.from_id(modal_call_id).cancel.aio() + + +# Where Modal raises a finished call's own outcome in this process: the +# worker's exception (whatever its type, followed in the traceback by the +# worker's own frames), the call's timeout, a failure Modal records for the +# call, or an output that cannot be decoded. Anything raised elsewhere came +# from this process's own request to Modal. +_CALL_RESULT_MODULE = "modal._utils.function_utils" +_CALL_RESULT_FUNCTION = "_process_result" +# Before raising anything, that function fetches an output Modal stored as a +# blob (any spawned call's output over 8 KiB) with this process's own +# requests; an error raised in them is not the call's outcome. +_OUTPUT_DOWNLOAD_MODULE = "modal._utils.blob_utils" +# Signals to this process itself, which must never be taken for an outcome. +OWN_PROCESS_SIGNALS: tuple[type[BaseException], ...] = ( + asyncio.CancelledError, + KeyboardInterrupt, + SystemExit, + GeneratorExit, +) + + +def raised_from_call_result(error: BaseException) -> bool: + """Return whether Modal raised ``error`` as a finished call's outcome. + + That is so when Modal's result processing raised it, or re-raised the + worker's exception, but not when the error came from the output download + that processing makes first. + """ + frames = [frame for frame, _ in traceback.walk_tb(error.__traceback__)] + for position, frame in enumerate(frames): + if ( + frame.f_globals.get("__name__") == _CALL_RESULT_MODULE + and frame.f_code.co_name == _CALL_RESULT_FUNCTION + ): + # The frame it called, if any: a download, or the worker's own. + callee = frames[position + 1] if position + 1 < len(frames) else None + return ( + callee is None + or callee.f_globals.get("__name__") != _OUTPUT_DOWNLOAD_MODULE + ) + return False + + +def _failed(code: RunErrorCode, message: str) -> CallFailed: + return CallFailed(RunError(code, message)) + + +class ModalCallInspector: + """Read a Modal call's output if it has one, without waiting.""" + + async def inspect(self, modal_call_id: str) -> CallOutcome: + """Return the call's outcome, or raise CallInspectionError. + + A worker's exception is re-raised here by Modal and can be of any + type, including one Modal's client also raises (a built-in + TimeoutError, a Modal or grpclib error) or a BaseException such as + SystemExit. So the outcome is decided by where the exception was + raised: from the finished call's result, it is the call's outcome; + otherwise it came from this process's own request to Modal, and a + signal to this process propagates. + """ + try: + call = modal.FunctionCall.from_id(modal_call_id) + output = await call.get.aio(timeout=0) + except BaseException as error: + if raised_from_call_result(error): + return _call_failure(error) + if isinstance(error, OWN_PROCESS_SIGNALS) or not isinstance( + error, Exception + ): + raise + return _request_outcome(error) + return CallSucceeded(output) + + +def _call_failure(error: BaseException) -> CallFailed: + """Turn a finished call's own outcome into a run failure.""" + if isinstance(error, modal.exception.FunctionTimeoutError): + return _failed(RunErrorCode.WORKER_TIMEOUT, "The Modal call timed out.") + # Record only the type: the message may contain calculation inputs. + return _failed( + RunErrorCode.WORKER_ERROR, f"The Modal call raised {type(error).__name__}." + ) + + +def _request_outcome(error: Exception) -> CallOutcome: + """Read an error from this process's own request to Modal. + + Raise CallInspectionError for a problem reaching Modal, so the call is + asked about again later. + """ + # Modal raises the built-in TimeoutError while the call has no output, + # including while it re-runs a pre-empted call. + if type(error) is TimeoutError: + return CallRunning() + if isinstance(error, modal.exception.OutputExpiredError): + return _failed( + RunErrorCode.RESULT_EXPIRED, "Modal no longer holds the call's output." + ) + if isinstance(error, modal.exception.NotFoundError): + return _failed(RunErrorCode.CALL_NOT_FOUND, "Modal does not know the call.") + raise CallInspectionError(type(error).__name__) from error diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/modal_identity.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/modal_identity.py new file mode 100644 index 000000000..cecb7190c --- /dev/null +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/modal_identity.py @@ -0,0 +1,155 @@ +"""Check that a request comes from sim's call-report Modal function. + +Modal gives every function container an OIDC identity token +(``MODAL_IDENTITY_TOKEN``), signed with RS256 by ``MODAL_OIDC_ISSUER``, whose +claims name the container's workspace, environment, app and function. A +call-ended report carries it as a bearer token. The entrypoint accepts it +only when Modal's published key verifies its signature, its issuer, audience +and expiry are valid, and it names the expected workspace and environment, a +v2 app (``V2_APPLICATION_PREFIX``) and the call-report function +(``CALL_REPORT_FUNCTION_NAME``). A token of any other Modal function, such as +a worker running a user's policy, is refused. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Protocol + +import jwt +from policyengine_simulation_contract.call_reports import CALL_REPORT_FUNCTION_NAME +from policyengine_simulation_contract.stage12_manifest import V2_APPLICATION_PREFIX + +if TYPE_CHECKING: + from collections.abc import Mapping + +MODAL_OIDC_ISSUER = "https://oidc.modal.com" +MODAL_OIDC_AUDIENCE = "oidc.modal.com" +MODAL_JWKS_URL = "https://oidc.modal.com/.well-known/jwks.json" +# Modal lists RS256 as its only signing algorithm. +MODAL_SIGNING_ALGORITHMS = ("RS256",) +# Tolerated clock difference between Modal and this instance. +CLOCK_LEEWAY_SECONDS = 30 + + +class ModalIdentityError(Exception): + """The token does not prove that the request comes from the expected caller.""" + + +class ModalKeysUnavailableError(Exception): + """Modal's published keys could not be fetched, so no token can be checked.""" + + +class SigningKeySource(Protocol): + """Where the key that signed a token is found (``jwt.PyJWKClient``).""" + + def get_signing_key_from_jwt(self, token: str) -> jwt.PyJWK: + """Return the published key that signed ``token``.""" + ... + + +@dataclass(frozen=True) +class ExpectedModalCaller: + """The only Modal caller whose reports are accepted.""" + + workspace_id: str + environment_name: str + app_name_prefix: str = V2_APPLICATION_PREFIX + function_name: str = CALL_REPORT_FUNCTION_NAME + + def __post_init__(self) -> None: + """Refuse an expectation that any workspace or environment would meet.""" + if not self.workspace_id: + msg = "the expected Modal workspace is not set" + raise ValueError(msg) + if not self.environment_name: + msg = "the expected Modal environment is not set" + raise ValueError(msg) + + +@dataclass(frozen=True) +class ModalIdentity: + """The verified claims of a Modal container's identity token.""" + + app_name: str + function_name: str + container_id: str + + +def _claim(claims: Mapping[str, object], name: str) -> str: + value = claims.get(name) + return value if isinstance(value, str) else "" + + +class ModalIdentityVerifier: + """Verify Modal identity tokens against Modal's published keys.""" + + def __init__( + self, + expected: ExpectedModalCaller, + *, + signing_keys: SigningKeySource | None = None, + ) -> None: + """Expect ``expected``; fetch keys from ``MODAL_JWKS_URL`` unless given.""" + self.expected = expected + self.signing_keys: SigningKeySource = signing_keys or jwt.PyJWKClient( + MODAL_JWKS_URL + ) + + def verify(self, token: str) -> ModalIdentity: + """Return the token's identity, or raise ModalIdentityError. + + Raise ModalKeysUnavailableError when Modal's keys cannot be fetched. + Fetching them makes a network request the first time and whenever + the key cache expires, so call this in a worker thread. + """ + try: + signing_key = self.signing_keys.get_signing_key_from_jwt(token).key + except jwt.PyJWKClientConnectionError as error: + msg = "Modal's signing keys could not be fetched" + raise ModalKeysUnavailableError(msg) from error + except jwt.PyJWTError as error: + msg = f"the token names no key Modal publishes: {type(error).__name__}" + raise ModalIdentityError(msg) from error + # PyJWKClient wraps only URLError and TimeoutError; a key set that + # arrives truncated, reset or not as JSON raises something else, and + # is no more the caller's fault. + except Exception as error: + msg = f"Modal's signing keys could not be read: {type(error).__name__}" + raise ModalKeysUnavailableError(msg) from error + try: + claims = jwt.decode( + token, + signing_key, + algorithms=list(MODAL_SIGNING_ALGORITHMS), + audience=MODAL_OIDC_AUDIENCE, + issuer=MODAL_OIDC_ISSUER, + leeway=CLOCK_LEEWAY_SECONDS, + options={"require": ["exp", "iat", "iss", "aud"]}, + ) + except jwt.PyJWTError as error: + msg = f"the token is not valid: {type(error).__name__}" + raise ModalIdentityError(msg) from error + return self._caller(claims) + + def _caller(self, claims: Mapping[str, object]) -> ModalIdentity: + expected = self.expected + if _claim(claims, "workspace_id") != expected.workspace_id: + msg = "the token names another Modal workspace" + raise ModalIdentityError(msg) + if _claim(claims, "environment_name") != expected.environment_name: + msg = "the token names another Modal environment" + raise ModalIdentityError(msg) + app_name = _claim(claims, "app_name") + if not app_name.startswith(expected.app_name_prefix): + msg = "the token names another Modal app" + raise ModalIdentityError(msg) + function_name = _claim(claims, "function_name") + if function_name != expected.function_name: + msg = "the token names another Modal function" + raise ModalIdentityError(msg) + return ModalIdentity( + app_name=app_name, + function_name=function_name, + container_id=_claim(claims, "container_id"), + ) diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_dispatch.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_dispatch.py new file mode 100644 index 000000000..acdc426a6 --- /dev/null +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_dispatch.py @@ -0,0 +1,312 @@ +"""Start each simulation run's, or comparison step's, Modal call exactly once. + +Starting a call takes three conditional writes around Modal's ``spawn``: + +1. claim: move the runs from ``pending`` to ``dispatching``. Only the instance + whose claim changes them goes on; every other caller stops here. +2. spawn the Modal call and wait only for Modal's acknowledgement. +3. record the call ID, moving the runs to ``running``, with the safety + check's first look due one interval later (the call itself reports when + it ends); or, when starting the call raised, fail them. + +If the instance stops between 1 and 3, the runs stay ``dispatching`` until the +safety check fails them as abandoned. No route starts calls yet. +""" + +from __future__ import annotations + +import asyncio +import logging +from dataclasses import dataclass +from enum import StrEnum +from typing import TYPE_CHECKING, Protocol + +from policyengine_simulation_persistence import RunError, RunErrorCode +from policyengine_simulation_persistence.run_store import utc_now + +from policyengine_simulation_entry.run_recording import SAFETY_CHECK_INTERVAL + +if TYPE_CHECKING: + from collections.abc import Callable, Collection + from datetime import datetime + from uuid import UUID + + from policyengine_simulation_entry.modal_calls import CallDispatcher, ModalRequest + +logger = logging.getLogger(__name__) + + +def _event_name(prefix: str, outcome: str) -> str: + """Name a log event for a simulation run or a comparison step.""" + return f"{prefix}_{outcome}" + + +# A call reports when it ends; the safety check first looks at one that has +# not reported after this long. +FIRST_CHECK_DELAY = SAFETY_CHECK_INTERVAL + + +class DispatchOutcome(StrEnum): + """What one attempt to start a call did.""" + + # This caller claimed the runs and Modal acknowledged the call. + STARTED = "started" + # Another caller claimed the runs first, or they were no longer pending. + NOT_CLAIMED = "not_claimed" + # Modal refused the call; the runs are failed. + DISPATCH_FAILED = "dispatch_failed" + # Modal acknowledged the call, but the runs had stopped dispatching before + # the call ID could be recorded, or recording it failed; the call is + # cancelled. + NOT_RECORDED = "not_recorded" + # Recording the call raised and the database could not say whether the + # record committed; the call is left running. If it did, the call's report + # or the safety check records its outcome; if not, the safety check fails + # the runs as abandoned. + RECORD_UNKNOWN = "record_unknown" + + +class DispatchTransitions(Protocol): + """The ``RunTransitions`` writes that starting a call needs.""" + + def claim_simulation_dispatch(self, run_ids: Collection[UUID]) -> bool: + """Move pending runs to dispatching.""" + ... + + def record_simulation_dispatched( + self, + run_ids: Collection[UUID], + *, + modal_call_id: str, + first_check_at: datetime, + ) -> bool: + """Record the call and move the runs to running.""" + ... + + def record_simulation_dispatch_failure( + self, run_ids: Collection[UUID], error: RunError + ) -> bool: + """Fail dispatching runs.""" + ... + + def simulation_runs_hold_call( + self, run_ids: Collection[UUID], modal_call_id: str + ) -> bool: + """Return whether every one of the runs recorded the call.""" + ... + + def claim_comparison_dispatch(self, report_run_id: UUID) -> bool: + """Move a pending comparison step to dispatching.""" + ... + + def record_comparison_dispatched( + self, + report_run_id: UUID, + *, + modal_call_id: str, + first_check_at: datetime, + ) -> bool: + """Record the comparison's call and move it to running.""" + ... + + def record_comparison_dispatch_failure( + self, report_run_id: UUID, error: RunError + ) -> bool: + """Fail a dispatching comparison step.""" + ... + + def comparison_holds_call(self, report_run_id: UUID, modal_call_id: str) -> bool: + """Return whether the comparison step recorded the call.""" + ... + + +@dataclass(frozen=True) +class _DispatchSteps: + """The writes around one spawn, and a read, bound to the runs they change.""" + + claim: Callable[[], bool] + record: Callable[[str, datetime], bool] + fail: Callable[[RunError], bool] + holds_call: Callable[[str], bool] + # "simulation_run" or "report_run_comparison": the log event names' start. + event_prefix: str + + +@dataclass +class RunDispatchService: + """Start Modal calls for simulation runs and comparison steps.""" + + transitions: DispatchTransitions + dispatcher: CallDispatcher + clock: Callable[[], datetime] = utc_now + + async def start_simulation_runs( + self, run_ids: Collection[UUID], request: ModalRequest + ) -> DispatchOutcome: + """Start the one Modal call that computes ``run_ids``. + + Pass a household report's baseline and reform runs together: one call + computes both, and they change status together. + """ + runs = frozenset(run_ids) + return await self._start( + _DispatchSteps( + claim=lambda: self.transitions.claim_simulation_dispatch(runs), + record=lambda call_id, first_check_at: ( + self.transitions.record_simulation_dispatched( + runs, modal_call_id=call_id, first_check_at=first_check_at + ) + ), + fail=lambda error: self.transitions.record_simulation_dispatch_failure( + runs, error + ), + holds_call=lambda call_id: self.transitions.simulation_runs_hold_call( + runs, call_id + ), + event_prefix="simulation_run", + ), + request, + subject={"simulation_run_ids": sorted(str(run_id) for run_id in runs)}, + ) + + async def start_comparison( + self, report_run_id: UUID, request: ModalRequest + ) -> DispatchOutcome: + """Start an economy report run's comparison step. + + The claim succeeds only once both of its simulation runs have + succeeded. + """ + return await self._start( + _DispatchSteps( + claim=lambda: self.transitions.claim_comparison_dispatch(report_run_id), + record=lambda call_id, first_check_at: ( + self.transitions.record_comparison_dispatched( + report_run_id, + modal_call_id=call_id, + first_check_at=first_check_at, + ) + ), + fail=lambda error: self.transitions.record_comparison_dispatch_failure( + report_run_id, error + ), + holds_call=lambda call_id: self.transitions.comparison_holds_call( + report_run_id, call_id + ), + event_prefix="report_run_comparison", + ), + request, + subject={"report_run_id": str(report_run_id)}, + ) + + async def _start( + self, + steps: _DispatchSteps, + request: ModalRequest, + *, + subject: dict[str, object], + ) -> DispatchOutcome: + if not await asyncio.to_thread(steps.claim): + return DispatchOutcome.NOT_CLAIMED + try: + modal_call_id = await self.dispatcher.spawn(request) + # Any failure to start the call fails the runs; they can be run again. + except Exception as error: # noqa: BLE001 # recorded on the runs + logger.warning( + _event_name(steps.event_prefix, "dispatch_failed"), + extra={**subject, "error_type": type(error).__name__}, + ) + await asyncio.to_thread( + steps.fail, + RunError( + RunErrorCode.DISPATCH_FAILED, + f"Starting the Modal call raised {type(error).__name__}.", + ), + ) + return DispatchOutcome.DISPATCH_FAILED + return await self._record(steps, modal_call_id, subject) + + async def _record( + self, steps: _DispatchSteps, modal_call_id: str, subject: dict[str, object] + ) -> DispatchOutcome: + """Record the acknowledged call, or cancel it if it cannot be recorded. + + A call whose runs never record it would run to its end with nobody + collecting its output, and its runs would later be failed as + abandoned; a resubmission would then start a second call. So such a + call is cancelled. + """ + context = {**subject, "modal_call_id": modal_call_id} + try: + recorded = await asyncio.to_thread( + steps.record, modal_call_id, self.clock() + FIRST_CHECK_DELAY + ) + except Exception as error: # settled below + logger.exception( + _event_name(steps.event_prefix, "call_record_failed"), + extra={**context, "error_type": type(error).__name__}, + ) + return await self._settle_failed_record(steps, modal_call_id, context) + if not recorded: + # The runs stopped dispatching while Modal acknowledged the call, + # for example because a recording pass failed them as abandoned. + logger.warning( + _event_name(steps.event_prefix, "call_not_recorded"), extra=context + ) + await self._cancel(modal_call_id, context, steps.event_prefix) + return DispatchOutcome.NOT_RECORDED + logger.info(_event_name(steps.event_prefix, "dispatched"), extra=context) + return DispatchOutcome.STARTED + + async def _cancel( + self, modal_call_id: str, context: dict[str, object], event_prefix: str + ) -> None: + try: + await self.dispatcher.cancel(modal_call_id) + except Exception as error: # logged; nothing else can be done + logger.exception( + _event_name(event_prefix, "call_not_cancelled"), + extra={**context, "error_type": type(error).__name__}, + ) + + async def _settle_failed_record( + self, + steps: _DispatchSteps, + modal_call_id: str, + context: dict[str, object], + ) -> DispatchOutcome: + """Settle a call whose record raised: it may have committed all the same. + + A record that committed but whose reply was lost leaves the runs + running under the call, which must keep running. Only an uncommitted + record leaves them dispatching, so they are failed if they still are, + and then the call is cancelled. Otherwise the runs are read: holding + the call, it started; not holding it (failed meanwhile as abandoned), + the call is cancelled. If the database cannot answer, the call is left + running, since cancelling it could stop a recorded call. + """ + try: + failed = await asyncio.to_thread( + steps.fail, + RunError( + RunErrorCode.DISPATCH_FAILED, + "The Modal call could not be recorded.", + ), + ) + holds_call = not failed and await asyncio.to_thread( + steps.holds_call, modal_call_id + ) + except Exception as error: # the call is left running; see above + logger.exception( + _event_name(steps.event_prefix, "call_record_unknown"), + extra={**context, "error_type": type(error).__name__}, + ) + return DispatchOutcome.RECORD_UNKNOWN + if holds_call: + logger.info(_event_name(steps.event_prefix, "dispatched"), extra=context) + return DispatchOutcome.STARTED + logger.warning( + _event_name(steps.event_prefix, "call_not_recorded"), extra=context + ) + await self._cancel(modal_call_id, context, steps.event_prefix) + return DispatchOutcome.NOT_RECORDED diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_recording.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_recording.py new file mode 100644 index 000000000..395f23153 --- /dev/null +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_recording.py @@ -0,0 +1,596 @@ +"""Record each Modal call's outcome on its runs, one pass of due checks at a time. + +A pass: + +1. fails runs and comparison steps left ``dispatching`` for ten minutes, so + they can be submitted again, and pending comparison steps whose baseline + or reform run failed, since they can never start; +2. claims the running calls that are due a check, by moving their next check + one lease ahead, so two passes do not check the same call at once; +3. asks Modal about each claimed call without waiting, and records what it + found with conditional writes: a finished call's result is validated and + its result rows are written in the transaction that marks its runs + ``succeeded``; a failed or timed-out call is marked ``failed`` with its + error; a call still running, or finished but of a kind this revision has + no recorder for, is checked again once its claim's lease ends. + +Pre-emption is never a failure: Modal re-runs a pre-empted call under the +same call ID, so it reads as still running. No caller polls, and Modal holds +no database credentials: sim writes every outcome. +""" + +from __future__ import annotations + +import asyncio +import logging +from dataclasses import dataclass, field +from datetime import timedelta +from enum import StrEnum +from typing import TYPE_CHECKING, Protocol + +from policyengine_simulation_contract.call_reports import CallReportStatus +from policyengine_simulation_persistence import ( + DueSimulationCall, + RunError, + RunErrorCode, + RunStatus, +) +from policyengine_simulation_persistence.run_store import utc_now + +from policyengine_simulation_entry.modal_calls import ( + OWN_PROCESS_SIGNALS, + CallFailed, + CallRunning, +) + +if TYPE_CHECKING: + from collections.abc import Callable, Mapping + from datetime import datetime + from uuid import UUID + + from policyengine_simulation_contract.call_reports import CallEndedReport + from policyengine_simulation_persistence import ( + AbandonedDispatches, + ReportKind, + ReportRunRecord, + SimulationKind, + SimulationRunRecord, + ) + from sqlalchemy import Connection + + from policyengine_simulation_entry.modal_calls import CallInspector, CallOutcome + +logger = logging.getLogger(__name__) + + +def _event_name(prefix: str, outcome: str) -> str: + """Name a log event for a simulation run or a comparison step.""" + return f"{prefix}_{outcome}" + + +# Cloud Scheduler starts the safety check this often (the deploy's +# schedule): a call that has not reported is checked about this long after +# its dispatch, then once per interval while it runs. +SAFETY_CHECK_INTERVAL = timedelta(minutes=5) +# A claimed call is not checked by another pass for this long, so each +# execution of the safety check checks a running call once; if the pass +# stops, the call becomes due again afterwards. +CHECK_LEASE = SAFETY_CHECK_INTERVAL +# A run left dispatching this long was claimed by an instance that stopped. +DISPATCH_ABANDONED_AFTER = timedelta(minutes=10) +CLAIM_BATCH_SIZE = 50 +# Each check writes through the runtime engine's small pool, which requests +# share (policyengine_simulation_persistence.engine: 2 connections plus 3 +# overflow), so a pass uses at most three at once. +INSPECTION_CONCURRENCY = 3 + + +class SimulationResultRecorder(Protocol): + """Validate one kind of simulation call's output and write its result rows.""" + + def validate(self, output: object) -> object: + """Return the validated result; raise ValueError when it is invalid.""" + ... + + def write( + self, + connection: Connection, + runs: tuple[SimulationRunRecord, ...], + result: object, + ) -> None: + """Write the result rows of ``runs`` in the transaction that succeeds them.""" + ... + + +class ComparisonResultRecorder(Protocol): + """Validate one kind of comparison's output and write its result rows.""" + + def validate(self, output: object) -> object: + """Return the validated result; raise ValueError when it is invalid.""" + ... + + def write( + self, connection: Connection, report_run: ReportRunRecord, result: object + ) -> None: + """Write the comparison's result rows in the transaction that succeeds it.""" + ... + + +@dataclass(frozen=True) +class ResultRecorders: + """The result recorders by kind of run. + + A finished call of a kind with no recorder here is checked again later, + so that a revision which has one records it. + """ + + simulations: Mapping[SimulationKind, SimulationResultRecorder] = field( + default_factory=dict + ) + comparisons: Mapping[ReportKind, ComparisonResultRecorder] = field( + default_factory=dict + ) + + +class RecordingTransitions(Protocol): + """The ``RunTransitions`` writes a recording pass needs.""" + + def claim_due_simulation_calls( + self, *, limit: int, lease: timedelta + ) -> list[DueSimulationCall]: + """Claim running simulation calls that are due a check.""" + ... + + def complete_simulation_call( + self, + modal_call_id: str, + write_results: Callable[[Connection, tuple[SimulationRunRecord, ...]], None], + ) -> bool: + """Mark a call's runs succeeded and write their results.""" + ... + + def fail_simulation_call(self, modal_call_id: str, error: RunError) -> bool: + """Mark a call's runs failed.""" + ... + + def claim_due_comparisons( + self, *, limit: int, lease: timedelta + ) -> list[ReportRunRecord]: + """Claim running comparison steps that are due a check.""" + ... + + def complete_comparison( + self, + report_run_id: UUID, + write_results: Callable[[Connection, ReportRunRecord], None], + ) -> bool: + """Mark a comparison step succeeded and write its results.""" + ... + + def fail_comparison(self, report_run_id: UUID, error: RunError) -> bool: + """Mark a comparison step failed.""" + ... + + def fail_abandoned_dispatches( + self, *, claimed_before: datetime + ) -> AbandonedDispatches: + """Fail runs and comparison steps left dispatching.""" + ... + + def fail_comparisons_of_failed_runs(self) -> tuple[UUID, ...]: + """Fail pending comparison steps whose simulation runs failed.""" + ... + + +class RunReader(Protocol): + """The ``RunStore`` reads a reported call needs.""" + + def get_simulation_run(self, run_id: UUID) -> SimulationRunRecord | None: + """Return the simulation run with this ID, or None.""" + ... + + def get_report_run(self, report_run_id: UUID) -> ReportRunRecord | None: + """Return the report run with this ID, or None.""" + ... + + +class _Recorded(StrEnum): + """What recording one call's outcome did.""" + + SUCCEEDED = "succeeded" + FAILED = "failed" + STILL_RUNNING = "still_running" + WITHOUT_RECORDER = "without_recorder" + NOT_WRITTEN = "not_written" + # Another writer recorded the outcome first. + UNCHANGED = "unchanged" + + +_REPORT_STATUSES = { + _Recorded.SUCCEEDED: CallReportStatus.RECORDED, + _Recorded.FAILED: CallReportStatus.RECORDED, + _Recorded.STILL_RUNNING: CallReportStatus.STILL_RUNNING, + _Recorded.WITHOUT_RECORDER: CallReportStatus.NO_RESULT_RECORDER, + _Recorded.NOT_WRITTEN: CallReportStatus.NOT_WRITTEN, + _Recorded.UNCHANGED: CallReportStatus.ALREADY_RECORDED, +} +# A run or comparison step in these statuses may not hold its call ID yet. +_NOT_YET_DISPATCHED = frozenset({RunStatus.PENDING, RunStatus.DISPATCHING}) + + +@dataclass +class RecordingPass: + """What one pass found and recorded.""" + + checked: int = 0 + succeeded: int = 0 + failed: int = 0 + still_running: int = 0 + not_reached: int = 0 + abandoned: int = 0 + comparisons_without_runs: int = 0 + without_recorder: int = 0 + + @property + def found_work(self) -> bool: + """Whether the pass checked or failed anything.""" + return bool(self.checked or self.abandoned or self.comparisons_without_runs) + + def count(self, recorded: _Recorded) -> None: + """Count what recording one checked call did.""" + counters = { + _Recorded.SUCCEEDED: "succeeded", + _Recorded.FAILED: "failed", + _Recorded.STILL_RUNNING: "still_running", + _Recorded.WITHOUT_RECORDER: "without_recorder", + _Recorded.NOT_WRITTEN: "not_reached", + } + counter = counters.get(recorded) + if counter is not None: + setattr(self, counter, getattr(self, counter) + 1) + + +@dataclass(frozen=True) +class _ResultSteps: + """Validate a finished call's output, then record it with its result rows.""" + + validate: Callable[[object], object] + complete: Callable[[object], bool] + + +@dataclass(frozen=True) +class _CallHandlers: + """How to record one claimed call's outcome, bound to its runs. + + ``result_steps`` is None when no recorder handles the runs' kind. + """ + + subject: dict[str, object] + # "simulation_run" or "report_run_comparison": the log event names' start. + event_prefix: str + result_steps: _ResultSteps | None + fail: Callable[[RunError], bool] + + +def _invalid_result(error: ValueError) -> RunError: + count = getattr(error, "error_count", None) + detail = f" with {count()} errors" if callable(count) else "" + return RunError( + RunErrorCode.INVALID_RESULT, + f"The Modal call's output failed validation{detail} ({type(error).__name__}).", + ) + + +@dataclass +class RunRecorder: + """Record Modal call outcomes: one reported call, or a pass of due checks.""" + + transitions: RecordingTransitions + runs: RunReader + inspector: CallInspector + recorders: ResultRecorders = field(default_factory=ResultRecorders) + clock: Callable[[], datetime] = utc_now + + async def record_due_outcomes(self) -> RecordingPass: + """Run one pass: fail abandoned dispatches, then check every due call.""" + recording = RecordingPass() + await self._fail_stopped_work(recording) + calls = await asyncio.to_thread( + lambda: self.transitions.claim_due_simulation_calls( + limit=CLAIM_BATCH_SIZE, lease=CHECK_LEASE + ) + ) + comparisons = await asyncio.to_thread( + lambda: self.transitions.claim_due_comparisons( + limit=CLAIM_BATCH_SIZE, lease=CHECK_LEASE + ) + ) + handlers = [ + *(self._simulation_handlers(call) for call in calls), + *(self._comparison_handlers(report_run) for report_run in comparisons), + ] + limit = asyncio.Semaphore(INSPECTION_CONCURRENCY) + + async def check(modal_call_id: str, call_handlers: _CallHandlers) -> None: + async with limit: + await self._check(modal_call_id, call_handlers, recording) + + await asyncio.gather( + *( + check(modal_call_id, call_handlers) + for modal_call_id, call_handlers in handlers + ) + ) + return recording + + async def _fail_stopped_work(self, recording: RecordingPass) -> None: + """Fail dispatches that stopped, and comparisons that can never start.""" + abandoned = await asyncio.to_thread( + lambda: self.transitions.fail_abandoned_dispatches( + claimed_before=self.clock() - DISPATCH_ABANDONED_AFTER + ) + ) + recording.abandoned = len(abandoned.simulation_run_ids) + len( + abandoned.report_run_ids + ) + for event_prefix, field_name, ids in ( + ("simulation_run", "simulation_run_ids", abandoned.simulation_run_ids), + ("report_run_comparison", "report_run_ids", abandoned.report_run_ids), + ): + if ids: + logger.warning( + _event_name(event_prefix, "dispatches_abandoned"), + extra={field_name: [str(identifier) for identifier in ids]}, + ) + without_runs = await asyncio.to_thread( + self.transitions.fail_comparisons_of_failed_runs + ) + recording.comparisons_without_runs = len(without_runs) + if without_runs: + logger.warning( + "report_run_comparisons_failed_with_their_runs", + extra={ + "report_run_ids": [ + str(report_run_id) for report_run_id in without_runs + ] + }, + ) + + def _simulation_handlers( + self, call: DueSimulationCall + ) -> tuple[str, _CallHandlers]: + kinds = {run.simulation_kind for run in call.runs} + recorder = ( + self.recorders.simulations.get(next(iter(kinds))) + if len(kinds) == 1 + else None + ) + call_id = call.modal_call_id + result_steps = None + if recorder is not None: + writer = recorder + result_steps = _ResultSteps( + validate=writer.validate, + complete=lambda result: self.transitions.complete_simulation_call( + call_id, + lambda connection, runs: writer.write(connection, runs, result), + ), + ) + return call_id, _CallHandlers( + subject={ + "modal_call_id": call_id, + "simulation_run_ids": [str(run.id) for run in call.runs], + }, + event_prefix="simulation_run", + result_steps=result_steps, + fail=lambda error: self.transitions.fail_simulation_call(call_id, error), + ) + + def _comparison_handlers( + self, report_run: ReportRunRecord + ) -> tuple[str, _CallHandlers]: + recorder = self.recorders.comparisons.get(report_run.report_kind) + report_run_id = report_run.id + call_id = report_run.comparison_modal_call_id or "" + result_steps = None + if recorder is not None: + writer = recorder + result_steps = _ResultSteps( + validate=writer.validate, + complete=lambda result: self.transitions.complete_comparison( + report_run_id, + lambda connection, record: writer.write(connection, record, result), + ), + ) + return call_id, _CallHandlers( + subject={"modal_call_id": call_id, "report_run_id": str(report_run_id)}, + event_prefix="report_run_comparison", + result_steps=result_steps, + fail=lambda error: self.transitions.fail_comparison(report_run_id, error), + ) + + async def _check( + self, modal_call_id: str, handlers: _CallHandlers, recording: RecordingPass + ) -> None: + recording.checked += 1 + try: + outcome = await self.inspector.inspect(modal_call_id) + recorded = await self._record(outcome, handlers) + # Modal could not be reached, or a database write failed (for example + # a dropped connection). Each check ends here, whatever was raised + # (only a signal to this process passes), so a pass returns only once + # all of its checks have, and its concurrency limit holds. The claim's + # lease makes the call due again; nothing was recorded. + except BaseException as error: + if isinstance(error, OWN_PROCESS_SIGNALS): + raise + recording.not_reached += 1 + logger.warning( + _event_name(handlers.event_prefix, "check_failed"), + extra={**handlers.subject, **_error_fields(error)}, + ) + return + recording.count(recorded) + + async def record_reported_call(self, report: CallEndedReport) -> CallReportStatus: + """Record the outcome of a call whose Modal function reported its end. + + The report is only a hint: the runs are read, Modal is asked about the + call without waiting, and the outcome is recorded with the same + conditional writes as a recording pass, so a repeated report changes + nothing. While the outcome cannot be recorded yet, the answer is one + of ``RETRY_LATER_STATUSES`` and nothing is written; the sender + reports again, and the safety check finds the call otherwise. Only a + signal to this process propagates. + """ + try: + found = await asyncio.to_thread(self._reported_call, report) + except Exception as error: # the database could not be read + logger.exception( + "call_report_not_read", + extra={"modal_call_id": report.modal_call_id, **_error_fields(error)}, + ) + return CallReportStatus.NOT_WRITTEN + if isinstance(found, CallReportStatus): + return found + try: + outcome = await self.inspector.inspect(report.modal_call_id) + except BaseException as error: + if isinstance(error, OWN_PROCESS_SIGNALS): + raise + logger.warning( + _event_name(found.event_prefix, "check_failed"), + extra={**found.subject, **_error_fields(error)}, + ) + return CallReportStatus.MODAL_UNREACHABLE + try: + recorded = await self._record(outcome, found) + except BaseException as error: + if isinstance(error, OWN_PROCESS_SIGNALS): + raise + logger.warning( + _event_name(found.event_prefix, "outcome_not_written"), + extra={**found.subject, **_error_fields(error)}, + ) + return CallReportStatus.NOT_WRITTEN + return _REPORT_STATUSES[recorded] + + def _reported_call( + self, report: CallEndedReport + ) -> CallReportStatus | _CallHandlers: + """Return how to record the reported call, or why it cannot be yet.""" + call_id = report.modal_call_id + if report.simulation_run_ids is not None: + runs = [ + self.runs.get_simulation_run(run_id) + for run_id in report.simulation_run_ids + ] + known = [run for run in runs if run is not None] + if len(known) != len(runs): + return CallReportStatus.UNKNOWN_RUNS + state = _reported_state( + call_id, [(run.modal_call_id, run.status) for run in known] + ) + if state is not None: + return state + due = DueSimulationCall(modal_call_id=call_id, runs=tuple(known)) + return self._simulation_handlers(due)[1] + report_run = ( + None + if report.report_run_id is None + else self.runs.get_report_run(report.report_run_id) + ) + if report_run is None or report_run.comparison_status is None: + return CallReportStatus.UNKNOWN_RUNS + state = _reported_state( + call_id, + [(report_run.comparison_modal_call_id, report_run.comparison_status)], + ) + if state is not None: + return state + return self._comparison_handlers(report_run)[1] + + async def _record(self, outcome: CallOutcome, handlers: _CallHandlers) -> _Recorded: + # A call still running, or of a kind with no recorder here, is checked + # again once its claim's lease ends; nothing is written. + if isinstance(outcome, CallRunning): + return _Recorded.STILL_RUNNING + if isinstance(outcome, CallFailed): + return await self._fail(handlers, outcome.error) + steps = handlers.result_steps + if steps is None: + # A revision without a recorder for this kind leaves the call for + # one that has it; Modal's output retention bounds the wait. + logger.warning( + _event_name(handlers.event_prefix, "without_result_recorder"), + extra=handlers.subject, + ) + return _Recorded.WITHOUT_RECORDER + return await self._record_result(outcome.output, steps, handlers) + + async def _record_result( + self, output: object, steps: _ResultSteps, handlers: _CallHandlers + ) -> _Recorded: + """Validate a finished call's output and write it with its success.""" + try: + result = await asyncio.to_thread(steps.validate, output) + except ValueError as error: + return await self._fail(handlers, _invalid_result(error)) + try: + completed = await asyncio.to_thread(steps.complete, result) + # Writing the result rows failed, for example because the database + # was unreachable: nothing was recorded, and the call is checked again. + except Exception as error: + logger.exception( + _event_name(handlers.event_prefix, "result_not_written"), + extra={**handlers.subject, **_error_fields(error)}, + ) + return _Recorded.NOT_WRITTEN + if not completed: + return _Recorded.UNCHANGED + logger.info( + _event_name(handlers.event_prefix, "succeeded"), extra=handlers.subject + ) + return _Recorded.SUCCEEDED + + async def _fail(self, handlers: _CallHandlers, error: RunError) -> _Recorded: + if not await asyncio.to_thread(handlers.fail, error): + return _Recorded.UNCHANGED + logger.warning( + _event_name(handlers.event_prefix, "failed"), + extra={**handlers.subject, "error_code": error.code.value}, + ) + return _Recorded.FAILED + + +def _reported_state( + modal_call_id: str, executions: list[tuple[str | None, RunStatus]] +) -> CallReportStatus | None: + """Return why a reported call cannot be recorded, or None if it can. + + ``executions`` holds each reported run's (or the comparison step's) call + ID and status. + """ + held_call_ids = {held for held, _ in executions} + if held_call_ids == {modal_call_id}: + if any(status is RunStatus.RUNNING for _, status in executions): + return None + return CallReportStatus.ALREADY_RECORDED + if held_call_ids - {modal_call_id, None}: + return CallReportStatus.OTHER_CALL + # Some runs hold no call ID yet. While they are being dispatched, the + # dispatcher is about to store it; a finished run without it (failed as + # abandoned, say) never will. + if all( + status in _NOT_YET_DISPATCHED for held, status in executions if held is None + ): + return CallReportStatus.NOT_DISPATCHED_YET + return CallReportStatus.OTHER_CALL + + +def _error_fields(error: BaseException) -> dict[str, str]: + """Name an error, and the error it was raised from, for a log record.""" + fields = {"error_type": type(error).__name__} + if error.__cause__ is not None: + fields["error_cause_type"] = type(error.__cause__).__name__ + return fields diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/safety_check.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/safety_check.py new file mode 100644 index 000000000..d630ee44e --- /dev/null +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/safety_check.py @@ -0,0 +1,126 @@ +"""The safety check: one recording pass, run as a Cloud Run job, then exit. + +A Modal call reports when it ends, and sim records its outcome then +(``call_report_route``). The safety check covers a call whose report never +arrived: Cloud Scheduler starts this job every ``SAFETY_CHECK_INTERVAL``, and +each execution runs one ``RunRecorder`` pass, which + +- fails runs and comparison steps left ``dispatching`` for ten minutes (sim + stopped between claiming them and recording their Modal call), so they can + be submitted again, and pending comparison steps whose run failed; +- asks Modal about running calls first dispatched, or last checked, at least + one interval ago, and records what it reports, with the same conditional + writes as a report. + +The job runs only the image of the revision that serves traffic: the deploy +updates it when that revision is promoted, and a rollback moves it back, so +a rejected candidate never records anything. Run it with +``python -m policyengine_simulation_entry.safety_check``; it reads +``SIMULATION_DATABASE_URL`` and sim's Modal token (``MODAL_TOKEN_ID``, +``MODAL_TOKEN_SECRET``). +""" + +from __future__ import annotations + +import asyncio +import logging +import os +import sys +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 + +if TYPE_CHECKING: + from collections.abc import Callable, Mapping + + from policyengine_simulation_entry.run_recording import RecordingPass + +logger = logging.getLogger(__name__) + +EXIT_PASSED = 0 +EXIT_PASS_FAILED = 1 +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``).""" + + async def record_due_outcomes(self) -> RecordingPass: + """Fail abandoned dispatches, then check every due call.""" + ... + + +def _database_recorder(database_url: str) -> DueOutcomeRecorder: + return SimulationDatabase.from_url(database_url).call_recorder() + + +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. + + 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") + print("SIMULATION_DATABASE_URL is not set.", file=sys.stderr) # noqa: T201 # the job's own error output + return EXIT_NOT_CONFIGURED + try: + recording = asyncio.run(build_recorder(database_url).record_due_outcomes()) + except BaseException as error: + if isinstance(error, OWN_PROCESS_SIGNALS): + raise + logger.exception( + "safety_check_pass_failed", extra={"error_type": type(error).__name__} + ) + return EXIT_PASS_FAILED + # 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.getLogger().setLevel(logging.INFO) + sys.exit(main()) diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/simulation_database.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/simulation_database.py index e0ca4f1bd..4401f8b7b 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/simulation_database.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/simulation_database.py @@ -6,10 +6,15 @@ from typing import TYPE_CHECKING from policyengine_simulation_persistence import ( + RunStore, + RunTransitions, create_simulation_engine, read_schema_revision, ) +from policyengine_simulation_entry.modal_calls import ModalCallInspector +from policyengine_simulation_entry.run_recording import ResultRecorders, RunRecorder + if TYPE_CHECKING: from sqlalchemy import Engine @@ -37,6 +42,20 @@ async def read_revision(self) -> str | None: """Return the applied migration revision, or None before the first one.""" return await asyncio.to_thread(self._read_revision) + def call_recorder(self) -> RunRecorder: + """Return what records a reported call's outcome on its runs. + + No kind of run has a result recorder yet (stage 2 registers them), so + a finished call is answered no_result_recorder and left for the + safety check; no route creates runs yet. + """ + return RunRecorder( + transitions=RunTransitions(self.engine), + runs=RunStore(self.engine), + inspector=ModalCallInspector(), + recorders=ResultRecorders(), + ) + def dispose(self) -> None: """Close every pooled connection.""" self.engine.dispose() diff --git a/projects/policyengine-simulation-entry/tests/run_fakes.py b/projects/policyengine-simulation-entry/tests/run_fakes.py new file mode 100644 index 000000000..acee770ed --- /dev/null +++ b/projects/policyengine-simulation-entry/tests/run_fakes.py @@ -0,0 +1,504 @@ +"""In-memory stand-ins for the run transitions and Modal, for the run tests. + +``FakeRunTransitions`` keeps each simulation run's and comparison step's status +and applies the same conditional rules as ``RunTransitions``: a change applies +only from the expected status, and the runs of one call change together. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field, replace +from datetime import UTC, datetime, timedelta +from threading import Lock +from typing import TYPE_CHECKING, Any +from uuid import UUID, uuid4 + +from policyengine_simulation_entry.modal_calls import ( + CallInspectionError, + CallRunning, + ModalRequest, + ModalTarget, +) +from policyengine_simulation_persistence import ( + AbandonedDispatches, + DueSimulationCall, + RunError, + RunErrorCode, + RunStatus, + SimulationKind, +) + +if TYPE_CHECKING: + from collections.abc import Collection + + from policyengine_simulation_entry.modal_calls import CallOutcome + + +@dataclass +class ManualClock: + """Return a fixed time until a test advances it.""" + + now: datetime = field( + default_factory=lambda: datetime(2026, 10, 22, 12, 0, tzinfo=UTC) + ) + + def __call__(self) -> datetime: + """Return the current fake time.""" + return self.now + + def advance(self, interval: timedelta) -> None: + """Move the fake time forward.""" + self.now += interval + + +@dataclass +class FakeExecution: + """One simulation run, or one comparison step, as the fake keeps it.""" + + id: UUID + kind: str + status: RunStatus = RunStatus.PENDING + modal_call_id: str | None = None + next_check_at: datetime | None = None + check_count: int = 0 + error: RunError | None = None + updated_at: datetime | None = None + written_results: list[Any] = field(default_factory=list) + + +@dataclass(frozen=True) +class FakeRunRecord: + """The run fields the recorder reads from a claimed run.""" + + id: UUID + simulation_kind: SimulationKind + modal_call_id: str | None + check_count: int + status: RunStatus = RunStatus.RUNNING + + +@dataclass(frozen=True) +class FakeReportRunRecord: + """The report run fields the recorder reads from a claimed comparison.""" + + id: UUID + report_kind: str + comparison_modal_call_id: str | None + comparison_check_count: int + comparison_status: RunStatus | None = RunStatus.RUNNING + + +class FakeRunTransitions: + """Conditional status changes over in-memory runs and comparison steps.""" + + def __init__(self, clock: ManualClock) -> None: + """Start with no runs.""" + self.clock = clock + # Callers run these methods in worker threads, as with the database. + self.lock = Lock() + self.runs: dict[UUID, FakeExecution] = {} + self.comparisons: dict[UUID, FakeExecution] = {} + self.comparison_runs_succeeded: dict[UUID, bool] = {} + self.comparisons_with_failed_runs: set[UUID] = set() + + # --- Building state -------------------------------------------------- + + def add_run(self, kind: str = "economy", **values: Any) -> UUID: + """Add a simulation run and return its ID.""" + run = FakeExecution(id=uuid4(), kind=kind, **values) + self.runs[run.id] = run + return run.id + + def add_comparison( + self, *, runs_succeeded: bool = True, runs_failed: bool = False, **values: Any + ) -> UUID: + """Add an economy report run's comparison step and return its ID.""" + comparison = FakeExecution(id=uuid4(), kind="economy", **values) + self.comparisons[comparison.id] = comparison + self.comparison_runs_succeeded[comparison.id] = ( + runs_succeeded and not runs_failed + ) + if runs_failed: + self.comparisons_with_failed_runs.add(comparison.id) + return comparison.id + + def running_call(self, call_id: str, *kinds: str) -> list[UUID]: + """Add runs already running under ``call_id``, due a check now.""" + return [ + self.add_run( + kind, + status=RunStatus.RUNNING, + modal_call_id=call_id, + next_check_at=self.clock(), + ) + for kind in kinds or ("economy",) + ] + + # --- Dispatch ---------------------------------------------------------- + + def _change_runs( + self, run_ids: Collection[UUID], expected: RunStatus, **values: Any + ) -> bool: + with self.lock: + runs = [self.runs[run_id] for run_id in run_ids] + if not runs or any(run.status is not expected for run in runs): + return False + for run in runs: + for name, value in values.items(): + setattr(run, name, value) + run.updated_at = self.clock() + return True + + def claim_simulation_dispatch(self, run_ids: Collection[UUID]) -> bool: + """Move pending runs to dispatching.""" + return self._change_runs( + run_ids, RunStatus.PENDING, status=RunStatus.DISPATCHING + ) + + def record_simulation_dispatched( + self, + run_ids: Collection[UUID], + *, + modal_call_id: str, + first_check_at: datetime, + ) -> bool: + """Record the call and move the runs to running.""" + return self._change_runs( + run_ids, + RunStatus.DISPATCHING, + status=RunStatus.RUNNING, + modal_call_id=modal_call_id, + next_check_at=first_check_at, + ) + + def record_simulation_dispatch_failure( + self, run_ids: Collection[UUID], error: RunError + ) -> bool: + """Fail dispatching runs.""" + return self._change_runs( + run_ids, RunStatus.DISPATCHING, status=RunStatus.FAILED, error=error + ) + + def _change_comparison( + self, report_run_id: UUID, expected: RunStatus, **values: Any + ) -> FakeExecution | None: + with self.lock: + comparison = self.comparisons[report_run_id] + if comparison.status is not expected: + return None + for name, value in values.items(): + setattr(comparison, name, value) + comparison.updated_at = self.clock() + return comparison + + def claim_comparison_dispatch(self, report_run_id: UUID) -> bool: + """Move a pending comparison step whose runs succeeded to dispatching.""" + if not self.comparison_runs_succeeded[report_run_id]: + return False + return ( + self._change_comparison( + report_run_id, RunStatus.PENDING, status=RunStatus.DISPATCHING + ) + is not None + ) + + def record_comparison_dispatched( + self, + report_run_id: UUID, + *, + modal_call_id: str, + first_check_at: datetime, + ) -> bool: + """Record the comparison's call.""" + return ( + self._change_comparison( + report_run_id, + RunStatus.DISPATCHING, + status=RunStatus.RUNNING, + modal_call_id=modal_call_id, + next_check_at=first_check_at, + ) + is not None + ) + + def record_comparison_dispatch_failure( + self, report_run_id: UUID, error: RunError + ) -> bool: + """Fail a dispatching comparison step.""" + return ( + self._change_comparison( + report_run_id, + RunStatus.DISPATCHING, + status=RunStatus.FAILED, + error=error, + ) + is not None + ) + + def simulation_runs_hold_call( + self, run_ids: Collection[UUID], modal_call_id: str + ) -> bool: + """Return whether every run recorded ``modal_call_id``.""" + with self.lock: + return all( + self.runs[run_id].modal_call_id == modal_call_id for run_id in run_ids + ) + + def comparison_holds_call(self, report_run_id: UUID, modal_call_id: str) -> bool: + """Return whether the comparison step recorded ``modal_call_id``.""" + with self.lock: + return self.comparisons[report_run_id].modal_call_id == modal_call_id + + # --- Reading ------------------------------------------------------- + + def get_simulation_run(self, run_id: UUID) -> FakeRunRecord | None: + """Return the run as a record, or None when there is no such run.""" + run = self.runs.get(run_id) + if run is None: + return None + return FakeRunRecord( + id=run.id, + simulation_kind=SimulationKind(run.kind), + modal_call_id=run.modal_call_id, + check_count=run.check_count, + status=run.status, + ) + + def get_report_run(self, report_run_id: UUID) -> FakeReportRunRecord | None: + """Return the comparison step's report run, or None when there is none.""" + comparison = self.comparisons.get(report_run_id) + if comparison is None: + return None + return FakeReportRunRecord( + id=comparison.id, + report_kind=comparison.kind, + comparison_modal_call_id=comparison.modal_call_id, + comparison_check_count=comparison.check_count, + comparison_status=comparison.status, + ) + + # --- Recording ------------------------------------------------------- + + def claim_due_simulation_calls( + self, *, limit: int, lease: timedelta + ) -> list[DueSimulationCall]: + """Claim running runs that are due, grouped by call.""" + now = self.clock() + due = [ + run + for run in self.runs.values() + if run.status is RunStatus.RUNNING + and run.next_check_at is not None + and run.next_check_at <= now + ][:limit] + calls: dict[str, list[FakeRunRecord]] = {} + for run in due: + run.next_check_at = now + lease + run.check_count += 1 + assert run.modal_call_id is not None + calls.setdefault(run.modal_call_id, []).append( + FakeRunRecord( + id=run.id, + simulation_kind=SimulationKind(run.kind), + modal_call_id=run.modal_call_id, + check_count=run.check_count, + ) + ) + return [ + DueSimulationCall(modal_call_id=call_id, runs=tuple(runs)) # type: ignore[arg-type] + for call_id, runs in sorted(calls.items()) + ] + + def _call_runs(self, modal_call_id: str) -> list[FakeExecution]: + return [ + run + for run in self.runs.values() + if run.modal_call_id == modal_call_id and run.status is RunStatus.RUNNING + ] + + def complete_simulation_call(self, modal_call_id: str, write_results) -> bool: + """Mark a call's runs succeeded after the writer succeeds.""" + runs = self._call_runs(modal_call_id) + if not runs: + return False + records = tuple( + FakeRunRecord( + id=run.id, + simulation_kind=SimulationKind(run.kind), + modal_call_id=run.modal_call_id, + check_count=run.check_count, + ) + for run in runs + ) + connection = object() + write_results(connection, records) + for run in runs: + run.status = RunStatus.SUCCEEDED + run.next_check_at = None + run.written_results.append(connection) + return True + + def fail_simulation_call(self, modal_call_id: str, error: RunError) -> bool: + """Fail a call's running runs.""" + runs = self._call_runs(modal_call_id) + for run in runs: + run.status = RunStatus.FAILED + run.error = error + run.next_check_at = None + return bool(runs) + + def claim_due_comparisons( + self, *, limit: int, lease: timedelta + ) -> list[FakeReportRunRecord]: + """Claim running comparison steps that are due.""" + now = self.clock() + due = [ + comparison + for comparison in self.comparisons.values() + if comparison.status is RunStatus.RUNNING + and comparison.next_check_at is not None + and comparison.next_check_at <= now + ][:limit] + for comparison in due: + comparison.next_check_at = now + lease + comparison.check_count += 1 + return [ + FakeReportRunRecord( + id=comparison.id, + report_kind="economy", + comparison_modal_call_id=comparison.modal_call_id, + comparison_check_count=comparison.check_count, + ) + for comparison in due + ] + + def complete_comparison(self, report_run_id: UUID, write_results) -> bool: + """Mark a running comparison step succeeded after the writer succeeds.""" + comparison = self.comparisons[report_run_id] + if comparison.status is not RunStatus.RUNNING: + return False + connection = object() + write_results( + connection, + FakeReportRunRecord( + id=comparison.id, + report_kind="economy", + comparison_modal_call_id=comparison.modal_call_id, + comparison_check_count=comparison.check_count, + ), + ) + comparison.status = RunStatus.SUCCEEDED + comparison.next_check_at = None + comparison.written_results.append(connection) + return True + + def fail_comparison(self, report_run_id: UUID, error: RunError) -> bool: + """Fail a running comparison step.""" + return ( + self._change_comparison( + report_run_id, + RunStatus.RUNNING, + status=RunStatus.FAILED, + error=error, + next_check_at=None, + ) + is not None + ) + + def fail_comparisons_of_failed_runs(self) -> tuple[UUID, ...]: + """Fail pending comparison steps whose simulation runs failed.""" + failed = [] + for comparison_id in sorted(self.comparisons_with_failed_runs): + changed = self._change_comparison( + comparison_id, + RunStatus.PENDING, + status=RunStatus.FAILED, + error=RunError(RunErrorCode.SIMULATION_RUN_FAILED, "A run failed."), + ) + if changed is not None: + failed.append(comparison_id) + return tuple(failed) + + def fail_abandoned_dispatches( + self, *, claimed_before: datetime + ) -> AbandonedDispatches: + """Fail runs and comparison steps dispatching since before the cutoff.""" + failed_runs = [] + for run in self.runs.values(): + if ( + run.status is RunStatus.DISPATCHING + and run.updated_at is not None + and run.updated_at < claimed_before + ): + run.status = RunStatus.FAILED + failed_runs.append(run.id) + failed_comparisons = [] + for comparison in self.comparisons.values(): + if ( + comparison.status is RunStatus.DISPATCHING + and comparison.updated_at is not None + and comparison.updated_at < claimed_before + ): + comparison.status = RunStatus.FAILED + failed_comparisons.append(comparison.id) + return AbandonedDispatches( + simulation_run_ids=tuple(failed_runs), + report_run_ids=tuple(failed_comparisons), + ) + + +class FakeDispatcher: + """Record spawned requests and return call IDs, or fail like Modal.""" + + def __init__( + self, + failure: Exception | None = None, + cancel_failure: Exception | None = None, + ) -> None: + """Fail spawns with ``failure`` and cancels with ``cancel_failure``.""" + self.failure = failure + self.cancel_failure = cancel_failure + self.requests: list[ModalRequest] = [] + self.cancelled: list[str] = [] + + async def cancel(self, modal_call_id: str) -> None: + """Record the cancelled call, or raise the configured failure.""" + self.cancelled.append(modal_call_id) + if self.cancel_failure is not None: + raise self.cancel_failure + + async def spawn(self, request: ModalRequest) -> str: + """Return a new call ID, or raise the configured failure.""" + self.requests.append(request) + if self.failure is not None: + raise self.failure + return f"fc-{len(self.requests)}" + + +class FakeInspector: + """Return a scripted outcome per call ID, or running when none is set.""" + + def __init__(self) -> None: + """Start with every call running.""" + self.outcomes: dict[str, CallOutcome | CallInspectionError] = {} + self.inspected: list[str] = [] + + async def inspect(self, modal_call_id: str) -> CallOutcome: + """Return the scripted outcome, or raise the scripted error.""" + self.inspected.append(modal_call_id) + outcome = self.outcomes.get(modal_call_id, CallRunning()) + if isinstance(outcome, CallInspectionError): + raise outcome + return outcome + + +SAMPLE_REQUEST = ModalRequest( + target=ModalTarget(application="policyengine-simulation-v2", function="run"), + arguments=({"country": "us"},), + observability_context={"traceparent": "00-test"}, +) + + +def execution(transitions: FakeRunTransitions, run_id: UUID) -> FakeExecution: + """Return a copy of a run as the fake keeps it.""" + return replace(transitions.runs[run_id]) diff --git a/projects/policyengine-simulation-entry/tests/test_call_report_deployment.py b/projects/policyengine-simulation-entry/tests/test_call_report_deployment.py new file mode 100644 index 000000000..cb09bfad3 --- /dev/null +++ b/projects/policyengine-simulation-entry/tests/test_call_report_deployment.py @@ -0,0 +1,139 @@ +"""Deploy settings for the route on which Modal reports that a call ended. + +The entrypoint accepts reports only from the call-report function of one +Modal workspace and environment. The workspace ID comes from the `beta` and +`prod` GitHub environment variable `MODAL_WORKSPACE_ID`; the environment is +the deployment's own Modal environment. +""" + +from __future__ import annotations + +import json +import subprocess +from pathlib import Path + +import pytest +from test_deployment_assets_cloud_run import cloud_run_deploy_environment +from test_deployment_assets_support import ( + CLOUD_RUN_DEPLOY_SCRIPT, + REUSABLE_DEPLOY_WORKFLOW, + workflow_jobs, +) +from workflow_script_support import FakeCommands, github_output, run_script + + +def _deploy(tmp_path, **overrides): + environment = cloud_run_deploy_environment(tmp_path) | overrides + result = subprocess.run( + [CLOUD_RUN_DEPLOY_SCRIPT], + check=False, + capture_output=True, + text=True, + env=environment, + ) + capture = Path(environment["GCLOUD_ENV_FILE_CAPTURE"]) + runtime = ( + json.loads(capture.read_text(encoding="utf-8")) if capture.exists() else {} + ) + return result, runtime + + +@pytest.mark.parametrize( + ("app_environment", "modal_environment"), + [("staging", "staging"), ("production", "main")], +) +def test_the_entrypoint_expects_reports_from_its_own_modal_environment( + tmp_path, app_environment, modal_environment +) -> None: + result, runtime = _deploy( + tmp_path, + APP_ENVIRONMENT=app_environment, + MODAL_ENVIRONMENT=modal_environment, + MODAL_WORKSPACE_ID_VALUE="ac-workspace0001", + ) + + assert result.returncode == 0, result.stderr + assert runtime["SIMULATION_CALL_REPORT_MODAL_WORKSPACE_ID"] == "ac-workspace0001" + assert runtime["SIMULATION_CALL_REPORT_MODAL_ENVIRONMENT"] == modal_environment + + +@pytest.mark.parametrize("workspace_id", ["", "workspace0001", "ac-", "ac-a b"]) +def test_a_deploy_without_a_modal_workspace_id_stops(tmp_path, workspace_id) -> None: + result, runtime = _deploy(tmp_path, MODAL_WORKSPACE_ID_VALUE=workspace_id) + + assert result.returncode != 0 + assert "MODAL_WORKSPACE_ID_VALUE" in result.stderr + assert runtime == {} + + +def test_the_deploy_step_passes_the_modal_workspace_id() -> None: + deploy_entrypoint = workflow_jobs(REUSABLE_DEPLOY_WORKFLOW)["deploy_entrypoint"] + environment = next( + step["env"] + for step in deploy_entrypoint["steps"] + if "deploy-cloud-run-simulation-entry.sh" in step.get("run", "") + ) + + assert environment["MODAL_WORKSPACE_ID_VALUE"] == "${{ vars.MODAL_WORKSPACE_ID }}" + + +URL_SCRIPT = "simulation-entry-main-url.sh" + + +def _find_url(tmp_path, service_url, **overrides): + fakes = FakeCommands.create(tmp_path) + fakes.add("gcloud", f"printf '%s\\n' '{service_url}'\n") + output = tmp_path / "github-output" + settings = { + "PROJECT_ID": "policyengine", + "REGION": "us-central1", + "SERVICE": "policyengine-simulation-entry", + "GITHUB_OUTPUT": str(output), + } + result = run_script(URL_SCRIPT, env=fakes.env(**(settings | overrides))) + return result, fakes, github_output(output) + + +def test_the_modal_functions_report_to_the_entrypoints_main_url(tmp_path) -> None: + url = "https://policyengine-simulation-entry-abc-uc.a.run.app" + + result, fakes, outputs = _find_url(tmp_path, url) + + assert result.returncode == 0, result.stderr + assert outputs == {"url": url} + (describe,) = fakes.calls("gcloud") + assert describe[1:5] == [ + "run", + "services", + "describe", + "policyengine-simulation-entry", + ] + assert "--format=value(status.url)" in describe + + +@pytest.mark.parametrize("service_url", ["", "http://insecure.example", "not a url"]) +def test_a_main_url_that_is_not_https_stops_the_deploy(tmp_path, service_url) -> None: + result, _, outputs = _find_url(tmp_path, service_url) + + assert result.returncode != 0 + assert outputs == {} + + +def test_the_v2_deploy_passes_the_main_url_to_the_modal_app() -> None: + job = workflow_jobs(REUSABLE_DEPLOY_WORKFLOW)["deploy_stage12_v2"] + steps = job["steps"] + find = next( + index for index, step in enumerate(steps) if URL_SCRIPT in step.get("run", "") + ) + deploy = next( + index + for index, step in enumerate(steps) + if "src/modal/v2_app.py" in step.get("run", "") + ) + + assert find < deploy + assert steps[find]["id"] == "entry_url" + assert steps[find]["env"]["SERVICE"] == "${{ inputs.entrypoint_service }}" + assert steps[deploy]["env"]["SIMULATION_ENTRY_URL"] == ( + "${{ steps.entry_url.outputs.url }}" + ) diff --git a/projects/policyengine-simulation-entry/tests/test_call_report_recording.py b/projects/policyengine-simulation-entry/tests/test_call_report_recording.py new file mode 100644 index 000000000..76ef79bd0 --- /dev/null +++ b/projects/policyengine-simulation-entry/tests/test_call_report_recording.py @@ -0,0 +1,280 @@ +"""Recording a call's outcome when its Modal function reports that it ended.""" + +from __future__ import annotations + +import asyncio +from uuid import uuid4 + +import pytest +from policyengine_simulation_contract.call_reports import ( + CallEndedReport, + CallReportStatus, +) +from policyengine_simulation_entry.modal_calls import ( + CallFailed, + CallInspectionError, + CallSucceeded, +) +from policyengine_simulation_entry.run_recording import ResultRecorders, RunRecorder +from policyengine_simulation_persistence import ( + ReportKind, + RunError, + RunErrorCode, + RunStatus, + SimulationKind, +) +from run_fakes import FakeInspector, FakeRunTransitions, ManualClock +from test_run_recording import RecordingRecorder + +WORKER_TIMEOUT = RunError(RunErrorCode.WORKER_TIMEOUT, "The Modal call timed out.") + + +def _recorder(recorders=None): + clock = ManualClock() + transitions = FakeRunTransitions(clock) + inspector = FakeInspector() + recorder = RunRecorder( + transitions=transitions, + runs=transitions, + inspector=inspector, + recorders=recorders or ResultRecorders(), + clock=clock, + ) + return recorder, transitions, inspector + + +def _economy_recorder(): + budget = RecordingRecorder() + recorder, transitions, inspector = _recorder( + ResultRecorders(simulations={SimulationKind.ECONOMY: budget}) + ) + return recorder, transitions, inspector, budget + + +def _report(recorder, **fields) -> CallReportStatus: + return asyncio.run(recorder.record_reported_call(CallEndedReport(**fields))) + + +def test_a_finished_call_is_recorded_when_it_reports() -> None: + recorder, transitions, inspector, budget = _economy_recorder() + (run_id,) = transitions.running_call("fc-done") + inspector.outcomes["fc-done"] = CallSucceeded({"budgetary_impact": 2.5}) + + status = _report(recorder, modal_call_id="fc-done", simulation_run_ids=(run_id,)) + + assert status is CallReportStatus.RECORDED + assert transitions.runs[run_id].status is RunStatus.SUCCEEDED + assert len(budget.writes) == 1 + + +def test_a_repeated_report_changes_nothing() -> None: + recorder, transitions, inspector, budget = _economy_recorder() + (run_id,) = transitions.running_call("fc-done") + inspector.outcomes["fc-done"] = CallSucceeded({"budgetary_impact": 2.5}) + _report(recorder, modal_call_id="fc-done", simulation_run_ids=(run_id,)) + + status = _report(recorder, modal_call_id="fc-done", simulation_run_ids=(run_id,)) + + assert status is CallReportStatus.ALREADY_RECORDED + assert len(budget.writes) == 1 + assert inspector.inspected == ["fc-done"] + + +def test_a_household_pair_is_recorded_from_one_report() -> None: + household = RecordingRecorder() + recorder, transitions, inspector = _recorder( + ResultRecorders(simulations={SimulationKind.HOUSEHOLD: household}) + ) + pair = transitions.running_call("fc-pair", "household", "household") + inspector.outcomes["fc-pair"] = CallSucceeded({"budgetary_impact": 0}) + + status = _report(recorder, modal_call_id="fc-pair", simulation_run_ids=tuple(pair)) + + assert status is CallReportStatus.RECORDED + assert {transitions.runs[run_id].status for run_id in pair} == {RunStatus.SUCCEEDED} + assert len(household.writes) == 1 + + +def test_a_failed_call_is_recorded_as_failed() -> None: + recorder, transitions, inspector = _recorder() + (run_id,) = transitions.running_call("fc-failed") + inspector.outcomes["fc-failed"] = CallFailed(WORKER_TIMEOUT) + + status = _report(recorder, modal_call_id="fc-failed", simulation_run_ids=(run_id,)) + + assert status is CallReportStatus.RECORDED + assert transitions.runs[run_id].status is RunStatus.FAILED + assert transitions.runs[run_id].error.code is RunErrorCode.WORKER_TIMEOUT + + +def test_a_call_modal_still_runs_is_left_for_another_report() -> None: + recorder, transitions, _, _ = _economy_recorder() + (run_id,) = transitions.running_call("fc-running") + next_check_at = transitions.runs[run_id].next_check_at + + status = _report(recorder, modal_call_id="fc-running", simulation_run_ids=(run_id,)) + + assert status is CallReportStatus.STILL_RUNNING + assert transitions.runs[run_id].status is RunStatus.RUNNING + assert transitions.runs[run_id].next_check_at == next_check_at + + +def test_a_call_whose_id_is_not_stored_yet_is_left_for_another_report() -> None: + """A fast household call can report before its dispatcher stores its ID.""" + recorder, transitions, inspector, _ = _economy_recorder() + run_id = transitions.add_run(status=RunStatus.DISPATCHING) + + status = _report(recorder, modal_call_id="fc-fast", simulation_run_ids=(run_id,)) + + assert status is CallReportStatus.NOT_DISPATCHED_YET + assert inspector.inspected == [] + assert transitions.runs[run_id].status is RunStatus.DISPATCHING + + +def test_a_report_naming_an_unknown_run_is_refused() -> None: + recorder, transitions, inspector, _ = _economy_recorder() + (run_id,) = transitions.running_call("fc-done") + + status = _report( + recorder, modal_call_id="fc-done", simulation_run_ids=(run_id, uuid4()) + ) + + assert status is CallReportStatus.UNKNOWN_RUNS + assert inspector.inspected == [] + + +@pytest.mark.parametrize( + "run_state", + [ + {"status": RunStatus.RUNNING, "modal_call_id": "fc-other"}, + {"status": RunStatus.SUCCEEDED, "modal_call_id": "fc-other"}, + # Failed as abandoned before the call's ID was stored. + {"status": RunStatus.FAILED, "modal_call_id": None}, + ], + ids=["running-another-call", "finished-another-call", "failed-without-call"], +) +def test_a_report_for_runs_that_do_not_hold_its_call_is_refused(run_state) -> None: + recorder, transitions, inspector, _ = _economy_recorder() + run_id = transitions.add_run(**run_state) + + status = _report(recorder, modal_call_id="fc-late", simulation_run_ids=(run_id,)) + + assert status is CallReportStatus.OTHER_CALL + assert inspector.inspected == [] + assert transitions.runs[run_id].status is run_state["status"] + + +def test_an_unreachable_modal_is_left_for_another_report() -> None: + recorder, transitions, inspector, _ = _economy_recorder() + (run_id,) = transitions.running_call("fc-done") + inspector.outcomes["fc-done"] = CallInspectionError("ConnectionError") + + status = _report(recorder, modal_call_id="fc-done", simulation_run_ids=(run_id,)) + + assert status is CallReportStatus.MODAL_UNREACHABLE + assert transitions.runs[run_id].status is RunStatus.RUNNING + + +def test_a_kind_without_a_recorder_is_left_for_a_revision_with_one() -> None: + recorder, transitions, inspector = _recorder() + (run_id,) = transitions.running_call("fc-done") + inspector.outcomes["fc-done"] = CallSucceeded({"budgetary_impact": 1}) + + status = _report(recorder, modal_call_id="fc-done", simulation_run_ids=(run_id,)) + + assert status is CallReportStatus.NO_RESULT_RECORDER + assert transitions.runs[run_id].status is RunStatus.RUNNING + + +def test_a_failed_read_is_left_for_another_report() -> None: + recorder, transitions, inspector, _ = _economy_recorder() + (run_id,) = transitions.running_call("fc-done") + + def unreachable_database(_run_id): + msg = "connection lost" + raise ConnectionError(msg) + + transitions.get_simulation_run = unreachable_database # type: ignore[method-assign] + + status = _report(recorder, modal_call_id="fc-done", simulation_run_ids=(run_id,)) + + assert status is CallReportStatus.NOT_WRITTEN + assert inspector.inspected == [] + + +@pytest.mark.parametrize( + ("outcome", "write"), + [ + (CallSucceeded({"budgetary_impact": 2.5}), "complete_simulation_call"), + (CallFailed(WORKER_TIMEOUT), "fail_simulation_call"), + ], + ids=["result", "failure"], +) +def test_a_failed_write_is_left_for_another_report(outcome, write) -> None: + recorder, transitions, inspector, _ = _economy_recorder() + (run_id,) = transitions.running_call("fc-done") + inspector.outcomes["fc-done"] = outcome + + def unreachable_database(*_arguments): + msg = "connection lost" + raise ConnectionError(msg) + + setattr(transitions, write, unreachable_database) + + status = _report(recorder, modal_call_id="fc-done", simulation_run_ids=(run_id,)) + + assert status is CallReportStatus.NOT_WRITTEN + assert transitions.runs[run_id].status is RunStatus.RUNNING + + +def test_a_signal_to_this_process_is_not_swallowed() -> None: + recorder, transitions, inspector, _ = _economy_recorder() + (run_id,) = transitions.running_call("fc-done") + inspector.outcomes["fc-done"] = CallFailed(WORKER_TIMEOUT) + + def interrupted(*_arguments): + raise KeyboardInterrupt + + transitions.fail_simulation_call = interrupted # type: ignore[method-assign] + + with pytest.raises(KeyboardInterrupt): + _report(recorder, modal_call_id="fc-done", simulation_run_ids=(run_id,)) + + +def test_a_comparison_call_is_recorded_when_it_reports() -> None: + comparisons = RecordingRecorder() + recorder, transitions, inspector = _recorder( + ResultRecorders(comparisons={ReportKind.ECONOMY: comparisons}) + ) + report_run_id = transitions.add_comparison( + status=RunStatus.RUNNING, modal_call_id="fc-comparison" + ) + inspector.outcomes["fc-comparison"] = CallSucceeded({"budgetary_impact": 3}) + + status = _report( + recorder, modal_call_id="fc-comparison", report_run_id=report_run_id + ) + + assert status is CallReportStatus.RECORDED + assert transitions.comparisons[report_run_id].status is RunStatus.SUCCEEDED + assert [record.id for _, record, _ in comparisons.writes] == [report_run_id] + + +def test_a_comparison_report_for_an_unknown_report_run_is_refused() -> None: + recorder, _, inspector, _ = _economy_recorder() + + status = _report(recorder, modal_call_id="fc-comparison", report_run_id=uuid4()) + + assert status is CallReportStatus.UNKNOWN_RUNS + assert inspector.inspected == [] + + +def test_a_comparison_whose_call_id_is_not_stored_yet_waits() -> None: + recorder, transitions, _, _ = _economy_recorder() + report_run_id = transitions.add_comparison(status=RunStatus.DISPATCHING) + + status = _report( + recorder, modal_call_id="fc-comparison", report_run_id=report_run_id + ) + + assert status is CallReportStatus.NOT_DISPATCHED_YET diff --git a/projects/policyengine-simulation-entry/tests/test_call_report_route.py b/projects/policyengine-simulation-entry/tests/test_call_report_route.py new file mode 100644 index 000000000..ebdeb996b --- /dev/null +++ b/projects/policyengine-simulation-entry/tests/test_call_report_route.py @@ -0,0 +1,198 @@ +"""The route on which sim's call-report Modal function reports that a call ended.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from uuid import uuid4 + +import pytest +from conftest import FakeBackend, make_settings +from fastapi import HTTPException +from fastapi.testclient import TestClient +from policyengine_simulation_contract.call_reports import ( + CALL_ENDED_PATH, + RETRY_LATER_STATUSES, + CallEndedReport, + CallReportStatus, +) +from policyengine_simulation_entry.app import RunServices, create_app +from policyengine_simulation_entry.modal_identity import ( + ModalIdentity, + ModalIdentityError, + ModalKeysUnavailableError, +) + +GOOD_TOKEN = "token-of-the-call-report-function" +RUN_ID = uuid4() +REPORT = {"modal_call_id": "fc-ended", "simulation_run_ids": [str(RUN_ID)]} + + +@dataclass +class _Verifier: + """Accept one token; refuse or fail on the others.""" + + keys_unavailable: bool = False + verified: list[str] = field(default_factory=list) + + def verify(self, token: str) -> ModalIdentity: + self.verified.append(token) + if self.keys_unavailable: + raise ModalKeysUnavailableError + if token != GOOD_TOKEN: + msg = "the token names another Modal function" + raise ModalIdentityError(msg) + return ModalIdentity( + app_name="policyengine-simulation-v2-py6-10-0", + function_name="report_call_ended", + container_id="ta-container0001", + ) + + +@dataclass +class _Recorder: + """Answer every report with one status, and keep the reports.""" + + status: CallReportStatus = CallReportStatus.RECORDED + reports: list[CallEndedReport] = field(default_factory=list) + + async def record_reported_call(self, report: CallEndedReport) -> CallReportStatus: + self.reports.append(report) + return self.status + + +def _reject_every_caller(): + raise HTTPException(status_code=403) + + +def _client(verifier=None, recorder=None) -> TestClient: + return TestClient( + create_app( + settings=make_settings(), + backend=FakeBackend(), + auth_dependency=_reject_every_caller, + run_services=RunServices( + call_recorder=recorder, modal_identity_verifier=verifier + ), + ) + ) + + +def _post(client, body=REPORT, token=GOOD_TOKEN): + headers = {} if token is None else {"Authorization": f"Bearer {token}"} + return client.post(CALL_ENDED_PATH, json=body, headers=headers) + + +def test_a_report_from_the_call_report_function_is_recorded() -> None: + verifier, recorder = _Verifier(), _Recorder() + + response = _post(_client(verifier, recorder)) + + assert response.status_code == 200 + assert response.json() == {"status": "recorded"} + assert verifier.verified == [GOOD_TOKEN] + assert recorder.reports == [CallEndedReport.model_validate(REPORT)] + + +def test_a_repeated_report_is_answered_as_already_recorded() -> None: + recorder = _Recorder(status=CallReportStatus.ALREADY_RECORDED) + + response = _post(_client(_Verifier(), recorder)) + + assert response.status_code == 200 + assert response.json() == {"status": "already_recorded"} + + +@pytest.mark.parametrize("status", sorted(RETRY_LATER_STATUSES - {"not_configured"})) +def test_a_report_sim_cannot_record_yet_is_answered_with_a_wait(status) -> None: + response = _post(_client(_Verifier(), _Recorder(status=CallReportStatus(status)))) + + assert response.status_code == 503 + assert response.json() == {"status": status} + assert int(response.headers["Retry-After"]) > 0 + + +@pytest.mark.parametrize( + ("status", "http_status"), + [(CallReportStatus.UNKNOWN_RUNS, 404), (CallReportStatus.OTHER_CALL, 409)], +) +def test_a_report_that_cannot_be_recorded_is_refused(status, http_status) -> None: + response = _post(_client(_Verifier(), _Recorder(status=status))) + + assert response.status_code == http_status + assert response.json() == {"status": status.value} + assert "Retry-After" not in response.headers + + +@pytest.mark.parametrize("token", [None, "token-of-a-policy-worker"]) +def test_a_request_without_the_functions_token_is_refused(token) -> None: + recorder = _Recorder() + + response = _post(_client(_Verifier(), recorder), token=token) + + assert response.status_code == 401 + assert recorder.reports == [] + + +def test_an_unauthenticated_request_learns_nothing_about_its_body() -> None: + response = _post(_client(_Verifier(), _Recorder()), body={"x": 1}, token=None) + + assert response.status_code == 401 + + +def test_a_malformed_report_is_refused_after_authentication() -> None: + recorder = _Recorder() + + response = _post( + _client(_Verifier(), recorder), + body={"modal_call_id": "fc-ended", "simulation_run_ids": [], "result": {}}, + ) + + assert response.status_code == 422 + assert recorder.reports == [] + + +def test_unreachable_modal_keys_are_answered_with_a_wait() -> None: + recorder = _Recorder() + + response = _post(_client(_Verifier(keys_unavailable=True), recorder)) + + assert response.status_code == 503 + assert response.json() == {"status": "modal_unreachable"} + assert recorder.reports == [] + + +@pytest.mark.parametrize( + ("verifier", "recorder"), + [(None, _Recorder()), (_Verifier(), None)], + ids=["no-modal-caller", "no-database"], +) +def test_an_instance_not_configured_for_reports_answers_with_a_wait( + verifier, recorder +) -> None: + response = _post(_client(verifier, recorder)) + + assert response.status_code == 503 + assert response.json() == {"status": "not_configured"} + + +def test_the_route_is_not_in_the_public_openapi_document() -> None: + document = _client(_Verifier(), _Recorder()).get("/openapi.json").json() + + assert CALL_ENDED_PATH not in document["paths"] + + +class _FailingVerifier: + def verify(self, token: str) -> ModalIdentity: + msg = "an error the verifier did not expect" + raise RuntimeError(msg) + + +def test_an_unexpected_verifier_failure_is_answered_with_a_wait() -> None: + """The sender reports again, instead of taking a 500 as a refusal.""" + recorder = _Recorder() + + response = _post(_client(_FailingVerifier(), recorder)) + + assert response.status_code == 503 + assert response.json() == {"status": "modal_unreachable"} + assert recorder.reports == [] diff --git a/projects/policyengine-simulation-entry/tests/test_config.py b/projects/policyengine-simulation-entry/tests/test_config.py index 23ed45320..eb8a6b72e 100644 --- a/projects/policyengine-simulation-entry/tests/test_config.py +++ b/projects/policyengine-simulation-entry/tests/test_config.py @@ -2,7 +2,7 @@ import pytest from conftest import make_settings -from policyengine_simulation_entry.config import ConfigurationError +from policyengine_simulation_entry.config import ConfigurationError, Settings def test_production_refuses_disabled_auth(): @@ -335,3 +335,43 @@ def test_the_simulation_database_url_must_be_canonical(database_url): make_settings(simulation_database_url=database_url).validate() assert "secret" not in str(error.value) + + +def test_call_reports_are_off_without_a_modal_workspace(monkeypatch): + monkeypatch.delenv("SIMULATION_CALL_REPORT_MODAL_WORKSPACE_ID", raising=False) + monkeypatch.delenv("SIMULATION_CALL_REPORT_MODAL_ENVIRONMENT", raising=False) + + settings = Settings.from_env() + + assert settings.expected_modal_caller is None + + +def test_call_reports_expect_the_configured_workspace_and_environment(monkeypatch): + monkeypatch.setenv("SIMULATION_CALL_REPORT_MODAL_WORKSPACE_ID", "ac-workspace0001") + monkeypatch.setenv("SIMULATION_CALL_REPORT_MODAL_ENVIRONMENT", "staging") + + expected = Settings.from_env().expected_modal_caller + + assert expected is not None + assert expected.workspace_id == "ac-workspace0001" + assert expected.environment_name == "staging" + + +@pytest.mark.parametrize( + ("workspace", "environment"), + [ + ("ac-workspace0001", ""), + ("", "staging"), + ("workspace0001", "staging"), + ("ac-workspace0001", "Not An Environment"), + ], + ids=["no-environment", "no-workspace", "not-a-workspace-id", "bad-environment"], +) +def test_call_report_settings_must_name_one_modal_workspace_and_environment( + workspace, environment +): + with pytest.raises(ConfigurationError, match="SIMULATION_CALL_REPORT"): + make_settings( + call_report_modal_workspace_id=workspace, + call_report_modal_environment=environment, + ).validate() diff --git a/projects/policyengine-simulation-entry/tests/test_deployment_assets_cloud_run.py b/projects/policyengine-simulation-entry/tests/test_deployment_assets_cloud_run.py index a6d870841..bc251481f 100644 --- a/projects/policyengine-simulation-entry/tests/test_deployment_assets_cloud_run.py +++ b/projects/policyengine-simulation-entry/tests/test_deployment_assets_cloud_run.py @@ -116,6 +116,7 @@ def cloud_run_deploy_environment(tmp_path: Path) -> dict[str, str]: "STAGE12_DATABASE_URL_SECRET_NAME": "stage12-staging-database-url", "MODAL_TOKEN_ID_SECRET_NAME": "modal-token-id", "MODAL_TOKEN_SECRET_SECRET_NAME": "modal-token-secret", + "MODAL_WORKSPACE_ID_VALUE": "ac-workspace0001", "SIMULATION_DATABASE_URL_SECRET_NAME": ( "simulation-staging-runtime-database-url" ), diff --git a/projects/policyengine-simulation-entry/tests/test_modal_identity.py b/projects/policyengine-simulation-entry/tests/test_modal_identity.py new file mode 100644 index 000000000..7f4c542ae --- /dev/null +++ b/projects/policyengine-simulation-entry/tests/test_modal_identity.py @@ -0,0 +1,185 @@ +"""Checking that a request comes from sim's call-report Modal function.""" + +from __future__ import annotations + +import time +from dataclasses import dataclass + +import jwt +import pytest +from cryptography.hazmat.primitives.asymmetric import rsa +from policyengine_simulation_entry.modal_identity import ( + MODAL_OIDC_AUDIENCE, + MODAL_OIDC_ISSUER, + ExpectedModalCaller, + ModalIdentityError, + ModalIdentityVerifier, + ModalKeysUnavailableError, +) + +WORKSPACE_ID = "ac-workspace0001" +EXPECTED = ExpectedModalCaller(workspace_id=WORKSPACE_ID, environment_name="staging") +SIGNING_KEY = rsa.generate_private_key(public_exponent=65537, key_size=2048) +OTHER_KEY = rsa.generate_private_key(public_exponent=65537, key_size=2048) + + +@dataclass +class _KeyFromJwks: + """Stands in for jwt.PyJWKClient: the public key Modal publishes.""" + + public_key: object + fetches: int = 0 + + def get_signing_key_from_jwt(self, _token: str): + self.fetches += 1 + return jwt.PyJWK.from_dict( + { + **jwt.algorithms.RSAAlgorithm.to_jwk(self.public_key, as_dict=True), + "kid": "modal-key", + "alg": "RS256", + } + ) + + +def _claims(**overrides): + now = int(time.time()) + claims = { + "iss": MODAL_OIDC_ISSUER, + "aud": MODAL_OIDC_AUDIENCE, + "sub": "modal:workspace_id:ac-workspace0001:...", + "iat": now, + "exp": now + 3600, + "workspace_id": WORKSPACE_ID, + "environment_name": "staging", + "app_name": "policyengine-simulation-v2-py6-10-0", + "function_name": "report_call_ended", + "container_id": "ta-container0001", + } + claims.update(overrides) + return {name: value for name, value in claims.items() if value is not None} + + +def _token(key=SIGNING_KEY, **overrides) -> str: + return jwt.encode( + _claims(**overrides), key, algorithm="RS256", headers={"kid": "modal-key"} + ) + + +def _verifier() -> ModalIdentityVerifier: + return ModalIdentityVerifier( + EXPECTED, signing_keys=_KeyFromJwks(SIGNING_KEY.public_key()) + ) + + +def test_a_token_of_the_call_report_function_is_accepted() -> None: + identity = _verifier().verify(_token()) + + assert identity.app_name == "policyengine-simulation-v2-py6-10-0" + assert identity.function_name == "report_call_ended" + assert identity.container_id == "ta-container0001" + + +@pytest.mark.parametrize( + ("overrides", "reason"), + [ + ({"workspace_id": "ac-someoneelse"}, "workspace"), + ({"environment_name": "main"}, "environment"), + ({"app_name": "policyengine-simulation-py6-10-0"}, "app"), + ({"function_name": "coordinate_report"}, "function"), + ({"workspace_id": None}, "workspace"), + ({"function_name": None}, "function"), + ], + ids=[ + "other-workspace", + "other-environment", + "main-path-app", + "other-function", + "no-workspace", + "no-function", + ], +) +def test_a_token_of_any_other_modal_caller_is_refused(overrides, reason) -> None: + with pytest.raises(ModalIdentityError, match=reason): + _verifier().verify(_token(**overrides)) + + +@pytest.mark.parametrize( + "token", + [ + _token(key=OTHER_KEY), + _token(iss="https://oidc.example.com"), + _token(aud="sim"), + _token(exp=int(time.time()) - 3600), + "not-a-token", + ], + ids=["wrong-key", "wrong-issuer", "wrong-audience", "expired", "malformed"], +) +def test_a_token_modal_did_not_sign_for_us_is_refused(token) -> None: + with pytest.raises(ModalIdentityError): + _verifier().verify(token) + + +def test_a_token_signed_with_another_algorithm_is_refused() -> None: + token = jwt.encode( + _claims(), "shared-secret-of-at-least-32-bytes!", algorithm="HS256" + ) + + with pytest.raises(ModalIdentityError): + _verifier().verify(token) + + +def test_the_expected_caller_needs_a_workspace_and_an_environment() -> None: + with pytest.raises(ValueError, match="workspace"): + ExpectedModalCaller(workspace_id="", environment_name="staging") + with pytest.raises(ValueError, match="environment"): + ExpectedModalCaller(workspace_id=WORKSPACE_ID, environment_name="") + + +class _UnreachableJwks: + def get_signing_key_from_jwt(self, _token: str): + msg = "Fail to fetch data from the url" + raise jwt.PyJWKClientConnectionError(msg) + + +def test_unreachable_modal_keys_are_not_a_refusal() -> None: + """The sender should report again later, not give up on its report.""" + verifier = ModalIdentityVerifier(EXPECTED, signing_keys=_UnreachableJwks()) + + with pytest.raises(ModalKeysUnavailableError): + verifier.verify(_token()) + + +class _BrokenJwks: + """Modal's key set answered, but not with a key set PyJWT could read.""" + + def __init__(self, failure: Exception) -> None: + self.failure = failure + + def get_signing_key_from_jwt(self, _token: str): + raise self.failure + + +@pytest.mark.parametrize( + "failure", + [ + ValueError("Expecting value: line 1 column 1 (char 0)"), + ConnectionResetError("reset while reading the key set"), + OSError("truncated response"), + ], + ids=["not-json", "reset", "truncated"], +) +def test_any_failure_to_fetch_modal_keys_is_not_a_refusal(failure) -> None: + verifier = ModalIdentityVerifier(EXPECTED, signing_keys=_BrokenJwks(failure)) + + with pytest.raises(ModalKeysUnavailableError): + verifier.verify(_token()) + + +def test_a_token_naming_no_published_key_is_still_refused() -> None: + verifier = ModalIdentityVerifier( + EXPECTED, + signing_keys=_BrokenJwks(jwt.PyJWKClientError("Unable to find a signing key")), + ) + + with pytest.raises(ModalIdentityError): + verifier.verify(_token()) diff --git a/projects/policyengine-simulation-entry/tests/test_modal_token_deployment.py b/projects/policyengine-simulation-entry/tests/test_modal_token_deployment.py new file mode 100644 index 000000000..24563081a --- /dev/null +++ b/projects/policyengine-simulation-entry/tests/test_modal_token_deployment.py @@ -0,0 +1,88 @@ +"""The entrypoint's Modal token, and no CPU kept outside requests. + +The entrypoint starts runs' Modal calls and asks Modal about a call when +Modal reports that it ended, so every deploy mounts sim's own Modal token, +whether or not it also deploys the Stage 12 path. It records outcomes only +within requests, so Cloud Run keeps no CPU allocated between them. +""" + +from __future__ import annotations + +import subprocess +from pathlib import Path + +import pytest +from test_deployment_assets_cloud_run import cloud_run_deploy_environment +from test_deployment_assets_support import ( + CLOUD_RUN_DEPLOY_SCRIPT, + REUSABLE_DEPLOY_WORKFLOW, + workflow_jobs, +) + + +def _run_deploy(tmp_path, **overrides): + environment = cloud_run_deploy_environment(tmp_path) | overrides + result = subprocess.run( + [CLOUD_RUN_DEPLOY_SCRIPT], + check=False, + capture_output=True, + text=True, + env=environment, + ) + arguments_capture = Path(environment["GCLOUD_ARGUMENTS_CAPTURE"]) + arguments = ( + arguments_capture.read_text(encoding="utf-8").splitlines() + if arguments_capture.exists() + else [] + ) + return result, arguments + + +@pytest.mark.parametrize("deploy_stage12", ["true", "false"]) +def test_no_deploy_keeps_cpu_allocated_outside_requests( + tmp_path, deploy_stage12 +) -> None: + result, arguments = _run_deploy(tmp_path, DEPLOY_STAGE12_V2=deploy_stage12) + + assert result.returncode == 0, result.stderr + assert "--no-cpu-throttling" not in arguments + + +@pytest.mark.parametrize("deploy_stage12", ["true", "false"]) +def test_every_deploy_mounts_sims_own_modal_token(tmp_path, deploy_stage12) -> None: + result, arguments = _run_deploy(tmp_path, DEPLOY_STAGE12_V2=deploy_stage12) + + assert result.returncode == 0, result.stderr + secrets = arguments[arguments.index("--set-secrets") + 1].split(",") + assert "MODAL_TOKEN_ID=modal-token-id:3" in secrets + assert "MODAL_TOKEN_SECRET=modal-token-secret:3" in secrets + assert len(secrets) == len(set(secrets)) + + +@pytest.mark.parametrize( + "variable", ["MODAL_TOKEN_ID_SECRET_NAME", "MODAL_TOKEN_SECRET_SECRET_NAME"] +) +def test_a_deploy_without_the_modal_token_stops(tmp_path, variable) -> None: + result, arguments = _run_deploy( + tmp_path, DEPLOY_STAGE12_V2="false", **{variable: ""} + ) + + assert result.returncode != 0 + assert variable in result.stderr + assert arguments == [] + + +def test_the_deploy_step_always_passes_the_modal_token_secret_names() -> None: + deploy_entrypoint = workflow_jobs(REUSABLE_DEPLOY_WORKFLOW)["deploy_entrypoint"] + environment = next( + step["env"] + for step in deploy_entrypoint["steps"] + if "deploy-cloud-run-simulation-entry.sh" in step.get("run", "") + ) + + assert environment["MODAL_TOKEN_ID_SECRET_NAME"] == ( + "${{ vars.MODAL_TOKEN_ID_SECRET_NAME }}" + ) + assert environment["MODAL_TOKEN_SECRET_SECRET_NAME"] == ( + "${{ vars.MODAL_TOKEN_SECRET_SECRET_NAME }}" + ) diff --git a/projects/policyengine-simulation-entry/tests/test_run_dispatch.py b/projects/policyengine-simulation-entry/tests/test_run_dispatch.py new file mode 100644 index 000000000..a90f420bc --- /dev/null +++ b/projects/policyengine-simulation-entry/tests/test_run_dispatch.py @@ -0,0 +1,629 @@ +"""Starting a run's Modal call exactly once, and reading Modal's answers.""" + +from __future__ import annotations + +import asyncio +import builtins + +import grpclib.const +import grpclib.exceptions +import modal +import pytest +from modal._serialization import pickle_traceback, serialize +from modal._utils import blob_utils +from modal._utils.function_utils import _process_result +from modal_proto import api_pb2 +from policyengine_simulation_entry import modal_calls +from policyengine_simulation_entry.modal_calls import ( + CallFailed, + CallInspectionError, + CallRunning, + CallSucceeded, + ModalCallDispatcher, + ModalCallInspector, +) +from policyengine_simulation_entry.run_dispatch import ( + FIRST_CHECK_DELAY, + DispatchOutcome, + RunDispatchService, +) +from policyengine_simulation_persistence import RunError, RunErrorCode, RunStatus +from run_fakes import ( + SAMPLE_REQUEST, + FakeDispatcher, + FakeRunTransitions, + ManualClock, + execution, +) + + +def _service(dispatcher=None): + clock = ManualClock() + transitions = FakeRunTransitions(clock) + service = RunDispatchService( + transitions=transitions, + dispatcher=dispatcher or FakeDispatcher(), + clock=clock, + ) + return service, transitions, clock + + +@pytest.mark.asyncio +async def test_a_claimed_run_starts_its_call_and_records_it() -> None: + service, transitions, clock = _service() + run_id = transitions.add_run() + + outcome = await service.start_simulation_runs([run_id], SAMPLE_REQUEST) + + assert outcome is DispatchOutcome.STARTED + run = execution(transitions, run_id) + assert run.status is RunStatus.RUNNING + assert run.modal_call_id == "fc-1" + assert run.next_check_at == clock() + FIRST_CHECK_DELAY + assert service.dispatcher.requests == [SAMPLE_REQUEST] + + +@pytest.mark.asyncio +async def test_concurrent_starts_spawn_one_call() -> None: + service, transitions, _ = _service() + run_id = transitions.add_run() + + outcomes = await asyncio.gather( + *(service.start_simulation_runs([run_id], SAMPLE_REQUEST) for _ in range(5)) + ) + + assert outcomes.count(DispatchOutcome.STARTED) == 1 + assert outcomes.count(DispatchOutcome.NOT_CLAIMED) == 4 + assert len(service.dispatcher.requests) == 1 + + +@pytest.mark.asyncio +async def test_a_household_pair_starts_one_shared_call() -> None: + service, transitions, _ = _service() + pair = [transitions.add_run("household"), transitions.add_run("household")] + + outcome = await service.start_simulation_runs(pair, SAMPLE_REQUEST) + + assert outcome is DispatchOutcome.STARTED + assert {execution(transitions, run_id).modal_call_id for run_id in pair} == {"fc-1"} + + +@pytest.mark.asyncio +async def test_a_refused_call_fails_the_runs_without_raising() -> None: + service, transitions, _ = _service( + FakeDispatcher(failure=modal.exception.NotFoundError("no such function")) + ) + run_id = transitions.add_run() + + outcome = await service.start_simulation_runs([run_id], SAMPLE_REQUEST) + + assert outcome is DispatchOutcome.DISPATCH_FAILED + run = execution(transitions, run_id) + assert run.status is RunStatus.FAILED + assert run.error.code is RunErrorCode.DISPATCH_FAILED + assert "NotFoundError" in run.error.message + assert "no such function" not in run.error.message + # Modal may have created the call before the error, so the message says + # only that starting it raised. + assert run.error.message == "Starting the Modal call raised NotFoundError." + + +@pytest.mark.asyncio +async def test_a_call_whose_runs_stopped_dispatching_is_cancelled(caplog) -> None: + service, transitions, _ = _service() + run_id = transitions.add_run() + + class SlowDispatcher(FakeDispatcher): + async def spawn(self, request): + transitions.runs[run_id].status = RunStatus.FAILED + return await super().spawn(request) + + service.dispatcher = SlowDispatcher() + + outcome = await service.start_simulation_runs([run_id], SAMPLE_REQUEST) + + assert outcome is DispatchOutcome.NOT_RECORDED + assert service.dispatcher.cancelled == ["fc-1"] + assert "simulation_run_call_not_recorded" in caplog.text + + +@pytest.mark.asyncio +async def test_a_call_that_cannot_be_recorded_is_cancelled_and_its_runs_failed( + caplog, +) -> None: + service, transitions, _ = _service() + run_id = transitions.add_run() + + def unreachable_database(*_arguments, **_keywords): + msg = "connection lost" + raise ConnectionError(msg) + + transitions.record_simulation_dispatched = unreachable_database # type: ignore[method-assign] + + outcome = await service.start_simulation_runs([run_id], SAMPLE_REQUEST) + + assert outcome is DispatchOutcome.NOT_RECORDED + assert service.dispatcher.cancelled == ["fc-1"] + run = execution(transitions, run_id) + assert run.status is RunStatus.FAILED + assert run.error.code is RunErrorCode.DISPATCH_FAILED + assert run.error.message == "The Modal call could not be recorded." + assert any( + getattr(record, "modal_call_id", None) == "fc-1" for record in caplog.records + ) + + +@pytest.mark.asyncio +async def test_a_failed_cancel_is_logged_without_raising(caplog) -> None: + service, transitions, _ = _service( + FakeDispatcher(cancel_failure=RuntimeError("modal down")) + ) + run_id = transitions.add_run() + + def unreachable_database(*_arguments, **_keywords): + msg = "connection lost" + raise ConnectionError(msg) + + transitions.record_simulation_dispatched = unreachable_database # type: ignore[method-assign] + + outcome = await service.start_simulation_runs([run_id], SAMPLE_REQUEST) + + assert outcome is DispatchOutcome.NOT_RECORDED + assert "simulation_run_call_not_cancelled" in caplog.text + + +def _commit_then_lose_acknowledgement(record): + """Return ``record`` changed to commit and then raise, as a lost reply does.""" + + def committed_without_acknowledgement(*arguments, **keywords): + assert record(*arguments, **keywords) + msg = "connection lost after the commit" + raise ConnectionError(msg) + + return committed_without_acknowledgement + + +@pytest.mark.asyncio +async def test_an_unacknowledged_record_keeps_its_call() -> None: + service, transitions, _ = _service() + run_id = transitions.add_run() + transitions.record_simulation_dispatched = _commit_then_lose_acknowledgement( # type: ignore[method-assign] + transitions.record_simulation_dispatched + ) + + outcome = await service.start_simulation_runs([run_id], SAMPLE_REQUEST) + + assert outcome is DispatchOutcome.STARTED + assert service.dispatcher.cancelled == [] + run = execution(transitions, run_id) + assert run.status is RunStatus.RUNNING + assert run.modal_call_id == "fc-1" + + +@pytest.mark.asyncio +async def test_an_unacknowledged_comparison_record_keeps_its_call() -> None: + service, transitions, _ = _service() + comparison = transitions.add_comparison() + transitions.record_comparison_dispatched = _commit_then_lose_acknowledgement( # type: ignore[method-assign] + transitions.record_comparison_dispatched + ) + + outcome = await service.start_comparison(comparison, SAMPLE_REQUEST) + + assert outcome is DispatchOutcome.STARTED + assert service.dispatcher.cancelled == [] + assert transitions.comparisons[comparison].status is RunStatus.RUNNING + + +@pytest.mark.asyncio +async def test_a_failed_record_of_runs_failed_meanwhile_cancels_the_call() -> None: + service, transitions, _ = _service() + run_id = transitions.add_run() + + def abandoned_meanwhile(*_arguments, **_keywords): + transitions.runs[run_id].status = RunStatus.FAILED + msg = "connection lost" + raise ConnectionError(msg) + + transitions.record_simulation_dispatched = abandoned_meanwhile # type: ignore[method-assign] + + outcome = await service.start_simulation_runs([run_id], SAMPLE_REQUEST) + + assert outcome is DispatchOutcome.NOT_RECORDED + assert service.dispatcher.cancelled == ["fc-1"] + + +@pytest.mark.asyncio +async def test_a_call_whose_record_cannot_be_read_is_left_running(caplog) -> None: + """With the database unreachable, the record may have committed all the same.""" + service, transitions, _ = _service() + run_id = transitions.add_run() + + def unreachable_database(*_arguments, **_keywords): + msg = "connection lost" + raise ConnectionError(msg) + + transitions.record_simulation_dispatched = unreachable_database # type: ignore[method-assign] + transitions.record_simulation_dispatch_failure = unreachable_database # type: ignore[method-assign] + + outcome = await service.start_simulation_runs([run_id], SAMPLE_REQUEST) + + assert outcome is DispatchOutcome.RECORD_UNKNOWN + assert service.dispatcher.cancelled == [] + assert any( + record.message == "simulation_run_call_record_unknown" + and getattr(record, "modal_call_id", None) == "fc-1" + for record in caplog.records + ) + + +def test_the_dispatcher_cancels_a_call_by_its_id(monkeypatch) -> None: + cancelled = [] + + class _FakeCancel: + async def aio(self): + cancelled.append("cancelled") + + class _FakeCall: + cancel = _FakeCancel() + + monkeypatch.setattr( + modal_calls.modal.FunctionCall, "from_id", lambda call_id: _FakeCall() + ) + + asyncio.run(ModalCallDispatcher("staging").cancel("fc-orphan")) + + assert cancelled == ["cancelled"] + + +@pytest.mark.asyncio +async def test_a_comparison_starts_only_once_its_runs_succeeded() -> None: + service, transitions, _ = _service() + waiting = transitions.add_comparison(runs_succeeded=False) + ready = transitions.add_comparison() + + assert ( + await service.start_comparison(waiting, SAMPLE_REQUEST) + is DispatchOutcome.NOT_CLAIMED + ) + assert ( + await service.start_comparison(ready, SAMPLE_REQUEST) is DispatchOutcome.STARTED + ) + assert transitions.comparisons[ready].modal_call_id == "fc-1" + assert len(service.dispatcher.requests) == 1 + + +@pytest.mark.asyncio +async def test_a_refused_comparison_call_fails_the_step(caplog) -> None: + service, transitions, _ = _service(FakeDispatcher(failure=RuntimeError("down"))) + comparison = transitions.add_comparison() + + outcome = await service.start_comparison(comparison, SAMPLE_REQUEST) + + assert outcome is DispatchOutcome.DISPATCH_FAILED + assert transitions.comparisons[comparison].status is RunStatus.FAILED + assert [record.message for record in caplog.records] == [ + "report_run_comparison_dispatch_failed" + ] + + +class _FakeCall: + def __init__(self, behaviour): + self.behaviour = behaviour + self.get = self + + async def aio(self, *, timeout): + assert timeout == 0 + if isinstance(self.behaviour, BaseException): + raise self.behaviour + return self.behaviour + + +def _inspect(monkeypatch, behaviour): + monkeypatch.setattr( + modal_calls.modal.FunctionCall, + "from_id", + lambda call_id: _FakeCall(behaviour), + ) + return asyncio.run(ModalCallInspector().inspect("fc-test")) + + +def test_a_call_without_output_is_running(monkeypatch) -> None: + assert _inspect(monkeypatch, TimeoutError()) == CallRunning() + + +def test_a_finished_call_returns_its_output(monkeypatch) -> None: + assert _inspect(monkeypatch, {"budget": 1.0}) == CallSucceeded({"budget": 1.0}) + + +@pytest.mark.parametrize( + ("raised", "code"), + [ + (modal.exception.OutputExpiredError(), RunErrorCode.RESULT_EXPIRED), + (modal.exception.NotFoundError("gone"), RunErrorCode.CALL_NOT_FOUND), + ], + ids=["expired", "not-found"], +) +def test_modal_answers_about_the_call_itself_fail_the_run( + monkeypatch, raised, code +) -> None: + outcome = _inspect(monkeypatch, raised) + + assert isinstance(outcome, CallFailed) + assert outcome.error.code is code + + +def test_a_failed_call_records_only_the_failure_type(monkeypatch) -> None: + outcome = _inspect_finished(monkeypatch, KeyError("household secret")) + + assert outcome == CallFailed( + RunError(RunErrorCode.WORKER_ERROR, "The Modal call raised KeyError.") + ) + + +@pytest.mark.parametrize( + "raised", + [ + modal.exception.ConnectionError("unreachable"), + modal.exception.ServiceError("unavailable"), + modal.exception.PermissionDeniedError("token revoked"), + modal.exception.VersionError("client too old"), + modal.exception.ClientClosed("closed"), + grpclib.exceptions.StreamTerminatedError("stream reset"), + grpclib.exceptions.GRPCError(grpclib.const.Status.UNAVAILABLE, "down"), + RuntimeError("an unexpected client problem"), + ], + ids=[ + "connection", + "service", + "permission", + "version", + "client-closed", + "stream-terminated", + "grpc", + "other-client-error", + ], +) +def test_a_modal_problem_is_retried_later(monkeypatch, raised) -> None: + with pytest.raises(CallInspectionError): + _inspect(monkeypatch, raised) + + +class _FinishedCall: + """A call whose output Modal holds: the worker's exception, or a timeout. + + Reading it goes through Modal's own result processing, which raises the + worker's exception in this process, as a real finished call does. + """ + + def __init__(self, worker_exception, status=None): + self.result = api_pb2.GenericResult( + status=status or api_pb2.GenericResult.GENERIC_STATUS_FAILURE, + data=serialize(worker_exception), + exception=repr(worker_exception), + ) + self.get = self + + async def aio(self, *, timeout): + assert timeout == 0 + return await _process_result( + self.result, api_pb2.DATA_FORMAT_PICKLE, None, None + ) + + +def _inspect_finished(monkeypatch, worker_exception, status=None): + monkeypatch.setattr( + modal_calls.modal.FunctionCall, + "from_id", + lambda call_id: _FinishedCall(worker_exception, status), + ) + return asyncio.run(ModalCallInspector().inspect("fc-test")) + + +@pytest.mark.parametrize( + "worker_exception", + [ + SystemExit(2), + builtins.BaseExceptionGroup("worker", [ValueError("one")]), + KeyboardInterrupt(), + TimeoutError("a child call timed out"), + modal.exception.ConnectionError("the worker's own Modal call failed"), + grpclib.exceptions.GRPCError(grpclib.const.Status.UNAVAILABLE, "down"), + ], + ids=[ + "system-exit", + "exception-group", + "keyboard-interrupt", + "timeout", + "modal-connection", + "grpc", + ], +) +def test_whatever_the_worker_raised_fails_the_run( + monkeypatch, worker_exception +) -> None: + """The worker's own exception, even of a type the client also raises, is final.""" + outcome = _inspect_finished(monkeypatch, worker_exception) + + assert isinstance(outcome, CallFailed) + assert outcome.error.code is RunErrorCode.WORKER_ERROR + assert type(worker_exception).__name__ in outcome.error.message + + +def test_a_call_modal_failed_internally_fails_the_run(monkeypatch) -> None: + outcome = _inspect_finished( + monkeypatch, + RuntimeError("ignored"), + status=api_pb2.GenericResult.GENERIC_STATUS_INTERNAL_FAILURE, + ) + + assert isinstance(outcome, CallFailed) + assert outcome.error.code is RunErrorCode.WORKER_ERROR + + +def test_a_call_modal_timed_out_fails_as_a_timeout(monkeypatch) -> None: + outcome = _inspect_finished( + monkeypatch, + RuntimeError("ignored"), + status=api_pb2.GenericResult.GENERIC_STATUS_TIMEOUT, + ) + + assert outcome == CallFailed( + RunError(RunErrorCode.WORKER_TIMEOUT, "The Modal call timed out.") + ) + + +BLOB_DOWNLOAD_URL = "https://blob.test/output" + + +class _BlobStub: + """Modal's API as its result processing uses it to fetch a stored output.""" + + def __init__(self, failure=None): + self.failure = failure + + async def BlobGet(self, request): # noqa: N802 # Modal's method name + assert request.blob_id == "bl-test" + if self.failure is not None: + raise self.failure + return api_pb2.BlobGetResponse(download_url=BLOB_DOWNLOAD_URL) + + +class _BlobStoredCall: + """A finished call whose output Modal stores as a blob. + + Modal does so for any spawned call's output over 8 KiB, so nearly every + real result is read this way: Modal's result processing first makes this + process's own requests for the blob, then decodes it. + """ + + def __init__(self, stub): + self.result = api_pb2.GenericResult( + status=api_pb2.GenericResult.GENERIC_STATUS_SUCCESS, + data_blob_id="bl-test", + ) + self.stub = stub + self.get = self + + async def aio(self, *, timeout): + assert timeout == 0 + return await _process_result( + self.result, api_pb2.DATA_FORMAT_PICKLE, self.stub, None + ) + + +def _inspect_blob_stored(monkeypatch, stub, download): + async def fake_download(download_url): + assert download_url == BLOB_DOWNLOAD_URL + if isinstance(download, BaseException): + raise download + return download + + monkeypatch.setattr(blob_utils, "_download_from_url", fake_download) + monkeypatch.setattr( + modal_calls.modal.FunctionCall, + "from_id", + lambda call_id: _BlobStoredCall(stub), + ) + return asyncio.run(ModalCallInspector().inspect("fc-test")) + + +def test_an_output_stored_as_a_blob_is_read(monkeypatch) -> None: + outcome = _inspect_blob_stored(monkeypatch, _BlobStub(), serialize({"budget": 1.0})) + + assert outcome == CallSucceeded({"budget": 1.0}) + + +@pytest.mark.parametrize( + ("blob_request_failure", "download_failure"), + [ + (modal.exception.ExecutionError("blob request failed"), None), + (grpclib.exceptions.GRPCError(grpclib.const.Status.UNAVAILABLE, "down"), None), + (modal.exception.ConnectionError("unreachable"), None), + (None, modal.exception.ExecutionError("Get from url failed with status 503")), + ], + ids=["blob-request-execution-error", "blob-request-grpc", "connection", "download"], +) +def test_a_failed_output_download_is_retried_later( + monkeypatch, blob_request_failure, download_failure +) -> None: + """Requests for a stored output are the entrypoint's own, not the worker's.""" + with pytest.raises(CallInspectionError): + _inspect_blob_stored( + monkeypatch, _BlobStub(blob_request_failure), download_failure + ) + + +def test_an_output_download_that_times_out_is_checked_again(monkeypatch) -> None: + outcome = _inspect_blob_stored(monkeypatch, _BlobStub(), TimeoutError()) + + assert outcome == CallRunning() + + +def _worker_exception_from_its_own_download(): + """Return what a worker raises when its own Modal blob download fails.""" + + async def worker(): + await blob_utils.blob_download( + "bl-test", _BlobStub(modal.exception.ConnectionError("worker's download")) + ) + + try: + asyncio.run(worker()) + except modal.exception.ConnectionError as error: + return error + raise AssertionError + + +def test_a_worker_exception_with_its_traceback_fails_the_run(monkeypatch) -> None: + """The worker's frames follow Modal's result processing, whatever they hold.""" + worker_exception = _worker_exception_from_its_own_download() + serialized_tb, tb_line_cache = pickle_traceback(worker_exception, "ta-test") + call = _FinishedCall(worker_exception) + call.result.serialized_tb = serialized_tb + call.result.tb_line_cache = tb_line_cache + monkeypatch.setattr(modal_calls.modal.FunctionCall, "from_id", lambda call_id: call) + + outcome = asyncio.run(ModalCallInspector().inspect("fc-test")) + + assert outcome == CallFailed( + RunError(RunErrorCode.WORKER_ERROR, "The Modal call raised ConnectionError.") + ) + + +@pytest.mark.parametrize( + "signal", + [asyncio.CancelledError(), KeyboardInterrupt(), SystemExit(0)], + ids=["cancelled", "keyboard-interrupt", "system-exit"], +) +def test_a_signal_to_this_process_is_not_swallowed(monkeypatch, signal) -> None: + """Only an exception that came from the call's result is the worker's.""" + with pytest.raises(type(signal)): + _inspect(monkeypatch, signal) + + +def test_the_dispatcher_passes_observability_as_its_own_argument(monkeypatch) -> None: + spawned = {} + + class _FakeSpawn: + async def aio(self, *arguments, observability_context): + spawned.update(arguments=arguments, context=observability_context) + return type("Call", (), {"object_id": "fc-spawned"})() + + class _FakeFunction: + spawn = _FakeSpawn() + + def from_name(application, function, *, environment_name): + spawned.update(target=(application, function, environment_name)) + return _FakeFunction() + + monkeypatch.setattr(modal_calls.modal.Function, "from_name", from_name) + + call_id = asyncio.run(ModalCallDispatcher("staging").spawn(SAMPLE_REQUEST)) + + assert call_id == "fc-spawned" + assert spawned == { + "target": ("policyengine-simulation-v2", "run", "staging"), + "arguments": ({"country": "us"},), + "context": {"traceparent": "00-test"}, + } diff --git a/projects/policyengine-simulation-entry/tests/test_run_recording.py b/projects/policyengine-simulation-entry/tests/test_run_recording.py new file mode 100644 index 000000000..a2408ea8a --- /dev/null +++ b/projects/policyengine-simulation-entry/tests/test_run_recording.py @@ -0,0 +1,403 @@ +"""Recording each Modal call's outcome on its runs, one pass at a time.""" + +from __future__ import annotations + +import asyncio +import builtins +from datetime import timedelta + +import pytest +from conftest import FakeBackend, make_settings +from fastapi.testclient import TestClient +from policyengine_simulation_entry import run_recording +from policyengine_simulation_entry.app import RunServices, create_app +from policyengine_simulation_entry.modal_calls import ( + CallFailed, + CallInspectionError, + CallSucceeded, +) +from policyengine_simulation_entry.run_recording import ( + CHECK_LEASE, + RecordingPass, + ResultRecorders, + RunRecorder, +) +from policyengine_simulation_persistence import ( + HEAD_REVISION, + ReportKind, + RunError, + RunErrorCode, + RunStatus, + SimulationKind, +) +from policyengine_simulation_persistence.engine import MAX_OVERFLOW, POOL_SIZE +from pydantic import BaseModel +from run_fakes import FakeInspector, FakeRunTransitions, ManualClock + + +class BudgetResult(BaseModel): + budgetary_impact: float + + +WORKER_ERROR = RunError(RunErrorCode.WORKER_ERROR, "The Modal call raised KeyError.") + + +class RecordingRecorder: + """Validate outputs as ``BudgetResult`` and remember what it wrote.""" + + def __init__(self, failure: Exception | None = None) -> None: + self.failure = failure + self.writes = [] + + def validate(self, output): + return BudgetResult.model_validate(output) + + def write(self, connection, records, result) -> None: + if self.failure is not None: + raise self.failure + self.writes.append((connection, records, result)) + + +def _recorder(recorders=None): + clock = ManualClock() + transitions = FakeRunTransitions(clock) + inspector = FakeInspector() + recorder = RunRecorder( + transitions=transitions, + runs=transitions, + inspector=inspector, + recorders=recorders or ResultRecorders(), + clock=clock, + ) + return recorder, transitions, inspector, clock + + +def _run(recorder) -> RecordingPass: + return asyncio.run(recorder.record_due_outcomes()) + + +def test_a_running_call_is_checked_again_after_its_claims_lease() -> None: + recorder, transitions, inspector, clock = _recorder() + (run_id,) = transitions.running_call("fc-1") + + recording = _run(recorder) + + assert inspector.inspected == ["fc-1"] + assert recording.still_running == 1 + run = transitions.runs[run_id] + assert run.status is RunStatus.RUNNING + assert run.check_count == 1 + assert run.next_check_at == clock() + CHECK_LEASE + assert _run(recorder).checked == 0 + clock.advance(CHECK_LEASE) + assert _run(recorder).checked == 1 + + +def test_a_failed_call_fails_its_runs() -> None: + recorder, transitions, inspector, _ = _recorder() + (run_id,) = transitions.running_call("fc-1") + inspector.outcomes["fc-1"] = CallFailed( + RunError(RunErrorCode.WORKER_TIMEOUT, "The Modal call timed out.") + ) + + recording = _run(recorder) + + assert recording.failed == 1 + run = transitions.runs[run_id] + assert run.status is RunStatus.FAILED + assert run.error.code is RunErrorCode.WORKER_TIMEOUT + + +def test_a_finished_call_is_validated_and_written_with_its_success() -> None: + budget = RecordingRecorder() + recorder, transitions, inspector, _ = _recorder( + ResultRecorders(simulations={SimulationKind.ECONOMY: budget}) + ) + (run_id,) = transitions.running_call("fc-1") + inspector.outcomes["fc-1"] = CallSucceeded({"budgetary_impact": -1.5}) + + recording = _run(recorder) + + assert recording.succeeded == 1 + run = transitions.runs[run_id] + assert run.status is RunStatus.SUCCEEDED + ((connection, records, result),) = budget.writes + assert [record.id for record in records] == [run_id] + assert result == BudgetResult(budgetary_impact=-1.5) + assert run.written_results == [connection] + + +def test_a_household_pair_is_written_once_for_both_runs() -> None: + household = RecordingRecorder() + recorder, transitions, inspector, _ = _recorder( + ResultRecorders(simulations={SimulationKind.HOUSEHOLD: household}) + ) + pair = transitions.running_call("fc-pair", "household", "household") + inspector.outcomes["fc-pair"] = CallSucceeded({"budgetary_impact": 0}) + + recording = _run(recorder) + + assert recording.checked == 1 + assert inspector.inspected == ["fc-pair"] + ((_, records, _),) = household.writes + assert {record.id for record in records} == set(pair) + assert {transitions.runs[run_id].status for run_id in pair} == {RunStatus.SUCCEEDED} + + +def test_a_kind_without_a_recorder_waits_for_an_instance_with_one() -> None: + """An older revision still running beside a new one must not fail the run.""" + recorder, transitions, inspector, clock = _recorder() + (run_id,) = transitions.running_call("fc-1") + inspector.outcomes["fc-1"] = CallSucceeded({"budgetary_impact": 1}) + + recording = _run(recorder) + + assert recording.without_recorder == 1 + run = transitions.runs[run_id] + assert run.status is RunStatus.RUNNING + assert run.error is None + assert run.next_check_at == clock() + CHECK_LEASE + + +def test_an_invalid_output_fails_its_runs_without_repeating_it() -> None: + recorder, transitions, inspector, _ = _recorder( + ResultRecorders(simulations={SimulationKind.ECONOMY: RecordingRecorder()}) + ) + (run_id,) = transitions.running_call("fc-1") + inspector.outcomes["fc-1"] = CallSucceeded({"budgetary_impact": "household 42"}) + + recording = _run(recorder) + + assert recording.failed == 1 + error = transitions.runs[run_id].error + assert error.code is RunErrorCode.INVALID_RESULT + assert "1 errors" in error.message + assert "household 42" not in error.message + + +def test_an_unreachable_modal_leaves_the_call_for_after_its_lease() -> None: + recorder, transitions, inspector, clock = _recorder() + (run_id,) = transitions.running_call("fc-1") + inspector.outcomes["fc-1"] = CallInspectionError("ConnectionError") + + recording = _run(recorder) + + assert recording.not_reached == 1 + run = transitions.runs[run_id] + assert run.status is RunStatus.RUNNING + assert run.next_check_at == clock() + CHECK_LEASE + clock.advance(CHECK_LEASE) + assert _run(recorder).checked == 1 + + +def test_a_failed_result_write_records_nothing_and_is_retried() -> None: + recorder, transitions, inspector, clock = _recorder( + ResultRecorders( + simulations={ + SimulationKind.ECONOMY: RecordingRecorder( + failure=ConnectionError("database unreachable") + ) + } + ) + ) + (run_id,) = transitions.running_call("fc-1") + inspector.outcomes["fc-1"] = CallSucceeded({"budgetary_impact": 1}) + + recording = _run(recorder) + + assert recording.not_reached == 1 + assert recording.succeeded == 0 + run = transitions.runs[run_id] + assert run.status is RunStatus.RUNNING + assert run.next_check_at == clock() + CHECK_LEASE + + +def test_an_outcome_another_instance_recorded_first_is_not_counted() -> None: + recorder, transitions, _, _ = _recorder( + ResultRecorders(simulations={SimulationKind.ECONOMY: RecordingRecorder()}) + ) + (run_id,) = transitions.running_call("fc-1") + + class RacingInspector(FakeInspector): + async def inspect(self, modal_call_id): + transitions.runs[run_id].status = RunStatus.SUCCEEDED + return CallSucceeded({"budgetary_impact": 1}) + + recorder.inspector = RacingInspector() + + recording = _run(recorder) + + assert recording.checked == 1 + assert recording.succeeded == 0 + + +def test_comparison_steps_are_checked_and_recorded() -> None: + comparisons = RecordingRecorder() + recorder, transitions, inspector, clock = _recorder( + ResultRecorders(comparisons={ReportKind.ECONOMY: comparisons}) + ) + finished = transitions.add_comparison( + status=RunStatus.RUNNING, modal_call_id="fc-done", next_check_at=clock() + ) + waiting = transitions.add_comparison( + status=RunStatus.RUNNING, modal_call_id="fc-wait", next_check_at=clock() + ) + inspector.outcomes["fc-done"] = CallSucceeded({"budgetary_impact": 3}) + + recording = _run(recorder) + + assert recording.succeeded == 1 + assert recording.still_running == 1 + assert transitions.comparisons[finished].status is RunStatus.SUCCEEDED + assert [record.id for _, record, _ in comparisons.writes] == [finished] + assert transitions.comparisons[waiting].next_check_at == clock() + CHECK_LEASE + + +def test_dispatches_left_for_ten_minutes_are_failed_by_a_pass() -> None: + recorder, transitions, _, clock = _recorder() + stalled = transitions.add_run(status=RunStatus.DISPATCHING, updated_at=clock()) + clock.advance(timedelta(minutes=10, seconds=1)) + recent = transitions.add_run(status=RunStatus.DISPATCHING, updated_at=clock()) + + recording = _run(recorder) + + assert recording.abandoned == 1 + assert transitions.runs[stalled].status is RunStatus.FAILED + assert transitions.runs[recent].status is RunStatus.DISPATCHING + + +def test_a_pass_fails_comparisons_whose_simulation_run_failed() -> None: + recorder, transitions, _, _ = _recorder() + stuck = transitions.add_comparison(runs_failed=True) + waiting = transitions.add_comparison(runs_succeeded=False) + + recording = _run(recorder) + + assert recording.comparisons_without_runs == 1 + assert recording.found_work + assert transitions.comparisons[stuck].status is RunStatus.FAILED + assert transitions.comparisons[stuck].error.code is ( + RunErrorCode.SIMULATION_RUN_FAILED + ) + assert transitions.comparisons[waiting].status is RunStatus.PENDING + + +def test_comparison_outcomes_are_logged_as_comparison_events(caplog) -> None: + recorder, transitions, inspector, clock = _recorder() + transitions.add_comparison( + status=RunStatus.RUNNING, modal_call_id="fc-failed", next_check_at=clock() + ) + inspector.outcomes["fc-failed"] = CallFailed( + RunError(RunErrorCode.WORKER_ERROR, "The Modal call raised KeyError.") + ) + + _run(recorder) + + assert [record.message for record in caplog.records] == [ + "report_run_comparison_failed" + ] + + +def test_the_loop_leaves_database_connections_for_requests() -> None: + assert run_recording.INSPECTION_CONCURRENCY < POOL_SIZE + MAX_OVERFLOW + + +def test_a_database_error_in_one_check_ends_with_its_pass(caplog) -> None: + """Every check finishes before the pass returns, so the limit holds.""" + recorder, transitions, inspector, clock = _recorder() + (broken_run,) = transitions.running_call("fc-broken") + (failed_run,) = transitions.running_call("fc-failed") + for call_id in ("fc-broken", "fc-failed"): + inspector.outcomes[call_id] = CallFailed(WORKER_ERROR) + fail_call = transitions.fail_simulation_call + + def drop_one_connection(modal_call_id, error): + if modal_call_id == "fc-broken": + msg = "connection lost" + raise ConnectionError(msg) + return fail_call(modal_call_id, error) + + transitions.fail_simulation_call = drop_one_connection # type: ignore[method-assign] + + recording = _run(recorder) + + assert recording.not_reached == 1 + assert recording.failed == 1 + assert transitions.runs[failed_run].status is RunStatus.FAILED + assert transitions.runs[broken_run].next_check_at == clock() + CHECK_LEASE + assert "simulation_run_check_failed" in [ + record.message for record in caplog.records + ] + + +def test_abandoned_comparison_steps_are_logged_as_comparison_events(caplog) -> None: + recorder, transitions, _, clock = _recorder() + transitions.add_comparison(status=RunStatus.DISPATCHING, updated_at=clock()) + clock.advance(timedelta(minutes=10, seconds=1)) + + _run(recorder) + + assert [record.message for record in caplog.records] == [ + "report_run_comparison_dispatches_abandoned" + ] + + +class _NotAnException(BaseException): + """An error that derives from BaseException only, like a worker's exit.""" + + +def test_a_non_exception_error_in_one_check_ends_with_its_pass(caplog) -> None: + recorder, transitions, inspector, clock = _recorder() + (run_id,) = transitions.running_call("fc-broken") + inspector.outcomes["fc-broken"] = CallFailed(WORKER_ERROR) + + def raise_a_group(modal_call_id, error): + message = "writes" + raise builtins.BaseExceptionGroup(message, [_NotAnException()]) + + transitions.fail_simulation_call = raise_a_group # type: ignore[method-assign] + + recording = _run(recorder) + + assert recording.not_reached == 1 + assert transitions.runs[run_id].next_check_at == clock() + CHECK_LEASE + assert "simulation_run_check_failed" in [ + record.message for record in caplog.records + ] + + +def test_a_signal_to_this_process_stops_a_check() -> None: + recorder, transitions, inspector, _ = _recorder() + transitions.running_call("fc-1") + inspector.outcomes["fc-1"] = CallFailed(WORKER_ERROR) + + def interrupted(modal_call_id, error): + raise KeyboardInterrupt + + transitions.fail_simulation_call = interrupted # type: ignore[method-assign] + + with pytest.raises(KeyboardInterrupt): + _run(recorder) + + +class _DatabaseThatOnlyReadsItsRevision: + async def read_revision(self): + return HEAD_REVISION + + def dispose(self) -> None: + return None + + +def test_an_instance_runs_nothing_in_the_background() -> None: + """Outcomes are recorded when Modal reports them, within a request.""" + app = create_app( + settings=make_settings( + simulation_database_url="postgresql://runtime:secret@db.example/postgres" + ), + backend=FakeBackend(), + auth_dependency=lambda: None, + run_services=RunServices(database=_DatabaseThatOnlyReadsItsRevision()), + ) + + with TestClient(app) as client: + assert client.get("/health").status_code == 200 diff --git a/projects/policyengine-simulation-entry/tests/test_safety_check.py b/projects/policyengine-simulation-entry/tests/test_safety_check.py new file mode 100644 index 000000000..18b49b0ef --- /dev/null +++ b/projects/policyengine-simulation-entry/tests/test_safety_check.py @@ -0,0 +1,144 @@ +"""The safety check: one recording pass per execution of a Cloud Run job.""" + +from __future__ import annotations + +import subprocess +import sys +from datetime import timedelta + +import pytest +from policyengine_simulation_entry import safety_check +from policyengine_simulation_entry.run_dispatch import FIRST_CHECK_DELAY +from policyengine_simulation_entry.run_recording import ( + CHECK_LEASE, + SAFETY_CHECK_INTERVAL, + RecordingPass, +) + +DATABASE_URL = "postgresql://runtime:secret@db.example/postgres" + + +class _OnePass: + """A recorder whose pass returns a scripted result, or raises.""" + + def __init__(self, result) -> None: + self.result = result + self.passes = 0 + + async def record_due_outcomes(self) -> RecordingPass: + self.passes += 1 + if isinstance(self.result, BaseException): + raise self.result + return self.result + + +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 + + telemetry = telemetry or _Telemetry() + status = safety_check.main( + environment=environment, + build_recorder=build, + start_observability=telemetry.start, + ) + return status, built_for + + +def test_each_execution_runs_one_pass_and_exits() -> None: + recorder = _OnePass(RecordingPass(checked=4, succeeded=1, failed=1)) + + status, built_for = _run({"SIMULATION_DATABASE_URL": DATABASE_URL}, recorder) + + assert status == 0 + assert recorder.passes == 1 + assert built_for == [DATABASE_URL] + + +def test_the_pass_is_logged_with_what_it_found(caplog) -> None: + caplog.set_level("INFO") + recorder = _OnePass(RecordingPass(checked=4, succeeded=1, abandoned=2)) + + _run({"SIMULATION_DATABASE_URL": DATABASE_URL}, recorder) + + (record,) = [ + record for record in caplog.records if record.message == "safety_check_pass" + ] + # 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: + recorder = _OnePass(RecordingPass()) + + status, built_for = _run({}, recorder) + + assert status == safety_check.EXIT_NOT_CONFIGURED + assert built_for == [] + + +def test_a_pass_that_raises_fails_the_execution() -> None: + status, _ = _run( + {"SIMULATION_DATABASE_URL": DATABASE_URL}, _OnePass(RuntimeError("down")) + ) + + assert status == safety_check.EXIT_PASS_FAILED + + +def test_a_signal_to_the_job_is_not_swallowed() -> None: + with pytest.raises(KeyboardInterrupt): + _run({"SIMULATION_DATABASE_URL": DATABASE_URL}, _OnePass(KeyboardInterrupt())) + + +def test_the_job_runs_the_module_as_its_command() -> None: + """The deploy runs ``python -m policyengine_simulation_entry.safety_check``.""" + result = subprocess.run( + [sys.executable, "-m", "policyengine_simulation_entry.safety_check"], + check=False, + capture_output=True, + text=True, + env={"PATH": "/usr/bin:/bin"}, + ) + + assert result.returncode == safety_check.EXIT_NOT_CONFIGURED + assert "SIMULATION_DATABASE_URL" in result.stderr + + +def test_a_call_is_first_checked_one_interval_after_dispatch() -> None: + """A call reports when it ends; the safety check looks only after a while.""" + assert timedelta(minutes=5) == SAFETY_CHECK_INTERVAL + assert FIRST_CHECK_DELAY == SAFETY_CHECK_INTERVAL + assert CHECK_LEASE == SAFETY_CHECK_INTERVAL diff --git a/projects/policyengine-simulation-entry/tests/test_safety_check_deployment.py b/projects/policyengine-simulation-entry/tests/test_safety_check_deployment.py new file mode 100644 index 000000000..01b009b51 --- /dev/null +++ b/projects/policyengine-simulation-entry/tests/test_safety_check_deployment.py @@ -0,0 +1,357 @@ +"""Deploying the safety-check job from the revision that serves traffic. + +The job runs only the image of the promoted (or restored) entrypoint +revision, so a rejected candidate never records anything. The deploy copies +that revision's image, service account, database URL and Modal token, and +keeps a Cloud Scheduler entry that starts the job every five minutes. +""" + +from __future__ import annotations + +import json +from datetime import timedelta + +import pytest +from policyengine_simulation_entry.run_recording import SAFETY_CHECK_INTERVAL +from test_deployment_assets_support import ( + REUSABLE_DEPLOY_WORKFLOW, + ROLLBACK_WORKFLOW, + workflow_jobs, +) +from workflow_script_support import FakeCommands, run_script + +SCRIPT = "deploy-cloud-run-simulation-safety-check.sh" +SERVICE = "policyengine-simulation-entry" +REVISION = "policyengine-simulation-entry-p-123" +IMAGE = ( + "us-central1-docker.pkg.dev/policyengine/policyengine-simulation-entry/" + "policyengine-simulation-entry@sha256:" + "a" * 64 +) +RUNTIME_ACCOUNT = "sim-entry-prod-runtime@policyengine.iam.gserviceaccount.com" +SCHEDULER_ACCOUNT = "sim-safety-check-prod@policyengine.iam.gserviceaccount.com" +JOB = f"{SERVICE}-safety-check" + + +def _secret(name, secret, version): + return { + "name": name, + "valueFrom": {"secretKeyRef": {"name": secret, "key": version}}, + } + + +def _revision(*, service=SERVICE, environment=None): + return { + "apiVersion": "serving.knative.dev/v1", + "kind": "Revision", + "metadata": { + "name": REVISION, + "labels": {"serving.knative.dev/service": service}, + }, + "spec": { + "serviceAccountName": RUNTIME_ACCOUNT, + "containers": [ + { + "image": "us-central1-docker.pkg.dev/policyengine/x/y:sha", + "env": environment + if environment is not None + else [ + {"name": "APP_ENVIRONMENT", "value": "production"}, + { + "name": "SIMULATION_CALL_REPORT_MODAL_WORKSPACE_ID", + "value": "ac-workspace0001", + }, + {"name": "OLD_GATEWAY_URL", "value": "https://legacy.example"}, + {"name": "OBSERVABILITY_SERVICE_NAMESPACE", "value": "sim"}, + {"name": "CLOUD_RUN_REGION", "value": "us-central1"}, + _secret( + "SIMULATION_DATABASE_URL", "simulation-runtime-url", "4" + ), + _secret("MODAL_TOKEN_ID", "modal-token-id", "3"), + _secret("MODAL_TOKEN_SECRET", "modal-token-secret", "3"), + _secret( + "OLD_GATEWAY_AUTH_CLIENT_SECRET", "old-client-secret", "2" + ), + ], + } + ], + }, + "status": {"imageDigest": IMAGE}, + } + + +def _fakes(tmp_path, revision, *, schedule_exists): + fakes = FakeCommands.create(tmp_path) + revision_file = tmp_path / "revision.json" + revision_file.write_text(json.dumps(revision), encoding="utf-8") + environment_capture = tmp_path / "job-environment.json" + fakes.add( + "gcloud", + f""" + case "$1 $2 $3" in + "run revisions describe") cat "{revision_file}" ;; + "run jobs deploy") + while (($#)); do + if [[ "$1" == "--env-vars-file" ]]; then + cp "$2" "{environment_capture}" + fi + shift + done + ;; + "scheduler jobs describe") {"exit 0" if schedule_exists else "exit 1"} ;; + *) ;; + esac + """, + ) + return fakes, environment_capture + + +def _deploy(tmp_path, revision=None, *, schedule_exists=False, **overrides): + fakes, environment_capture = _fakes( + tmp_path, revision or _revision(), schedule_exists=schedule_exists + ) + settings = { + "PROJECT_ID": "policyengine", + "REGION": "us-central1", + "SERVICE": SERVICE, + "SAFETY_CHECK_REVISION": REVISION, + "SAFETY_CHECK_SCHEDULER_SERVICE_ACCOUNT": SCHEDULER_ACCOUNT, + "RUNNER_TEMP": str(tmp_path), + } + result = run_script(SCRIPT, env=fakes.env(**(settings | overrides))) + return result, fakes, environment_capture + + +def _call(fakes, *prefix): + calls = [ + call + for call in fakes.calls("gcloud") + if call[1 : 1 + len(prefix)] == list(prefix) + ] + assert len(calls) == 1, calls + return calls[0] + + +def _flag(call, name): + return call[call.index(name) + 1] + + +def test_the_job_runs_the_serving_revisions_exact_image_and_identity(tmp_path) -> None: + result, fakes, _ = _deploy(tmp_path) + + assert result.returncode == 0, result.stderr + deploy = _call(fakes, "run", "jobs", "deploy") + assert deploy[4] == JOB + assert _flag(deploy, "--image") == IMAGE + assert _flag(deploy, "--service-account") == RUNTIME_ACCOUNT + assert _flag(deploy, "--command") == "/opt/venv/bin/python" + assert "--args=-m,policyengine_simulation_entry.safety_check" in deploy + assert _flag(deploy, "--tasks") == "1" + assert _flag(deploy, "--max-retries") == "0" + assert _flag(deploy, "--project") == "policyengine" + assert _flag(deploy, "--region") == "us-central1" + + +def test_the_job_gets_only_the_secrets_it_reads_at_the_revisions_versions( + tmp_path, +) -> None: + result, fakes, environment_capture = _deploy(tmp_path) + + assert result.returncode == 0, result.stderr + secrets = _flag(_call(fakes, "run", "jobs", "deploy"), "--set-secrets").split(",") + assert sorted(secrets) == [ + "MODAL_TOKEN_ID=modal-token-id:3", + "MODAL_TOKEN_SECRET=modal-token-secret:3", + "SIMULATION_DATABASE_URL=simulation-runtime-url:4", + ] + environment = json.loads(environment_capture.read_text(encoding="utf-8")) + assert environment == { + "APP_ENVIRONMENT": "production", + "OBSERVABILITY_SERVICE_NAMESPACE": "sim", + "CLOUD_RUN_REGION": "us-central1", + } + + +@pytest.mark.parametrize("schedule_exists", [False, True]) +def test_the_schedule_starts_the_job_every_interval(tmp_path, schedule_exists) -> None: + result, fakes, _ = _deploy(tmp_path, schedule_exists=schedule_exists) + + assert result.returncode == 0, result.stderr + action = "update" if schedule_exists else "create" + schedule = _call(fakes, "scheduler", "jobs", action, "http") + assert schedule[5] == JOB + assert _flag(schedule, "--location") == "us-central1" + assert _flag(schedule, "--schedule") == "*/5 * * * *" + assert _flag(schedule, "--uri") == ( + "https://us-central1-run.googleapis.com/apis/run.googleapis.com/v1/" + f"namespaces/policyengine/jobs/{JOB}:run" + ) + assert _flag(schedule, "--http-method") == "POST" + assert _flag(schedule, "--oauth-service-account-email") == SCHEDULER_ACCOUNT + + +def test_the_schedule_matches_the_safety_check_interval() -> None: + assert timedelta(minutes=5) == SAFETY_CHECK_INTERVAL + + +def test_a_revision_of_another_service_is_refused(tmp_path) -> None: + result, fakes, _ = _deploy(tmp_path, _revision(service="another-service")) + + assert result.returncode != 0 + assert "another-service" in result.stderr + assert [call[1:3] for call in fakes.calls("gcloud")] == [["run", "revisions"]] + + +def test_a_revision_without_the_simulation_database_pauses_the_schedule( + tmp_path, +) -> None: + """An older revision restored by a rollback has no safety check.""" + revision = _revision( + environment=[{"name": "APP_ENVIRONMENT", "value": "production"}] + ) + + result, fakes, _ = _deploy(tmp_path, revision, schedule_exists=True) + + assert result.returncode == 0, result.stderr + assert [call[1:4] for call in fakes.calls("gcloud")] == [ + ["run", "revisions", "describe"], + ["scheduler", "jobs", "describe"], + ["scheduler", "jobs", "pause"], + ] + + +@pytest.mark.parametrize( + ("variable", "value"), + [ + ("SAFETY_CHECK_SCHEDULER_SERVICE_ACCOUNT", ""), + ("SAFETY_CHECK_SCHEDULER_SERVICE_ACCOUNT", "not-an-account"), + ("SAFETY_CHECK_REVISION", ""), + ("SAFETY_CHECK_REVISION", "Bad Revision"), + ("SERVICE", ""), + ], +) +def test_incomplete_settings_stop_before_any_change(tmp_path, variable, value) -> None: + result, fakes, _ = _deploy(tmp_path, **{variable: value}) + + assert result.returncode != 0 + assert variable in result.stderr + assert fakes.calls("gcloud") == [] + + +def _steps(job): + return job["steps"] + + +def _step_named(steps, name): + return next(index for index, step in enumerate(steps) if step.get("name") == name) + + +def test_promotion_points_the_job_at_the_revision_after_it_is_verified() -> None: + steps = _steps(workflow_jobs(REUSABLE_DEPLOY_WORKFLOW)["promote_entrypoint"]) + promote = _step_named(steps, "Promote exact tested revision") + safety = _step_named(steps, "Point the safety check at the promoted revision") + restore = _step_named(steps, "Restore previous revision after failed verification") + restore_safety = _step_named( + steps, "Point the safety check back at the previous revision" + ) + + verification = [ + index + for index, step in enumerate(steps) + if index > promote and step.get("name", "").startswith("Verify") + ] + assert verification + assert max(verification) < safety < restore < restore_safety + assert SCRIPT in steps[safety]["run"] + assert steps[safety]["env"]["SAFETY_CHECK_REVISION"] == "${{ env.TARGET_REVISION }}" + assert steps[safety]["env"]["SAFETY_CHECK_SCHEDULER_SERVICE_ACCOUNT"] == ( + "${{ vars.SIMULATION_SAFETY_CHECK_SCHEDULER_SERVICE_ACCOUNT }}" + ) + assert "steps.safety_check.outcome == 'failure'" in steps[restore_safety]["if"] + assert steps[restore_safety]["env"]["SAFETY_CHECK_REVISION"] == ( + "${{ env.PREVIOUS_REVISION }}" + ) + + +def test_no_candidate_deploy_touches_the_job() -> None: + jobs = workflow_jobs(REUSABLE_DEPLOY_WORKFLOW) + for name, job in jobs.items(): + runs = " ".join(step.get("run", "") for step in job.get("steps", [])) + if SCRIPT in runs: + assert name == "promote_entrypoint" + + +def test_a_rollback_points_the_job_at_the_restored_revision() -> None: + steps = _steps(workflow_jobs(ROLLBACK_WORKFLOW)["rollback"]) + restore = _step_named(steps, "Restore entrypoint traffic") + safety = _step_named(steps, "Point the safety check at the restored revision") + + assert restore < safety + assert steps[safety]["if"] == "${{ inputs.entrypoint_revision != '' }}" + assert steps[safety]["env"]["SAFETY_CHECK_REVISION"] == ( + "${{ inputs.entrypoint_revision }}" + ) + + +def test_a_revision_from_before_the_safety_check_pauses_the_schedule(tmp_path) -> None: + """A 1.04 or 1.05 revision mounts the database and token but has no job code.""" + environment = [ + entry + for entry in _revision()["spec"]["containers"][0]["env"] + if entry["name"] != "SIMULATION_CALL_REPORT_MODAL_WORKSPACE_ID" + ] + + result, fakes, _ = _deploy( + tmp_path, _revision(environment=environment), schedule_exists=True + ) + + assert result.returncode == 0, result.stderr + assert [call[1:4] for call in fakes.calls("gcloud")] == [ + ["run", "revisions", "describe"], + ["scheduler", "jobs", "describe"], + ["scheduler", "jobs", "pause"], + ] + + +@pytest.mark.parametrize("schedule_exists", [True, False]) +def test_pausing_stops_the_schedule_without_a_revision( + tmp_path, schedule_exists +) -> None: + fakes, _ = _fakes(tmp_path, _revision(), schedule_exists=schedule_exists) + settings = { + "PROJECT_ID": "policyengine", + "REGION": "us-central1", + "SERVICE": SERVICE, + } + + result = run_script(SCRIPT, "pause", env=fakes.env(**settings)) + + assert result.returncode == 0, result.stderr + expected = [["scheduler", "jobs", "describe"]] + if schedule_exists: + expected.append(["scheduler", "jobs", "pause"]) + assert [call[1:4] for call in fakes.calls("gcloud")] == expected + + +def test_a_failed_update_in_a_rollback_pauses_the_safety_check() -> None: + """The job must not keep running the image traffic just moved away from.""" + steps = _steps(workflow_jobs(ROLLBACK_WORKFLOW)["rollback"]) + safety = _step_named(steps, "Point the safety check at the restored revision") + pause = _step_named(steps, "Pause the safety check after a failed update") + + assert safety < pause + assert steps[safety]["id"] == "safety_check" + assert steps[pause]["if"] == ( + "${{ failure() && steps.safety_check.outcome == 'failure' }}" + ) + assert steps[pause]["run"] == f".github/scripts/{SCRIPT} pause" + + +def test_a_failed_restore_in_a_promotion_pauses_the_safety_check() -> None: + steps = _steps(workflow_jobs(REUSABLE_DEPLOY_WORKFLOW)["promote_entrypoint"]) + restore = _step_named(steps, "Point the safety check back at the previous revision") + pause = _step_named(steps, "Pause the safety check after a failed restore") + + assert restore < pause + assert steps[restore]["id"] == "safety_check_restore" + assert "steps.safety_check_restore.outcome == 'failure'" in steps[pause]["if"] + assert steps[pause]["run"] == f".github/scripts/{SCRIPT} pause" diff --git a/projects/policyengine-simulation-entry/tests/test_simulation_database.py b/projects/policyengine-simulation-entry/tests/test_simulation_database.py index 95b1defae..c5fd00f91 100644 --- a/projects/policyengine-simulation-entry/tests/test_simulation_database.py +++ b/projects/policyengine-simulation-entry/tests/test_simulation_database.py @@ -9,7 +9,7 @@ from fastapi.testclient import TestClient from policyengine_simulation_entry import app as app_module from policyengine_simulation_entry import simulation_database as database_module -from policyengine_simulation_entry.app import create_app +from policyengine_simulation_entry.app import RunServices, create_app from policyengine_simulation_entry.simulation_database import SimulationDatabase from policyengine_simulation_persistence import HEAD_REVISION from sqlalchemy.exc import OperationalError @@ -43,7 +43,7 @@ def _client(backend, database): settings=make_settings(simulation_database_url=DATABASE_URL), backend=backend, auth_dependency=lambda: None, - simulation_database=database, + run_services=RunServices(database=database), ) ) diff --git a/projects/policyengine-simulation-executor/src/modal/call_reports.py b/projects/policyengine-simulation-executor/src/modal/call_reports.py new file mode 100644 index 000000000..81ef33db2 --- /dev/null +++ b/projects/policyengine-simulation-executor/src/modal/call_reports.py @@ -0,0 +1,268 @@ +"""Report to the Simulation Entrypoint that a v2 Modal call ended. + +Interface for the v2 functions' last step (the economy coordinator, item +2.01, and the household functions, item 2.03). Wrap the work of a top-level +function whose call sim started:: + + with reporting_call_end( + modal_call_id=modal.current_function_call_id(), + simulation_run_ids=run_ids, # or report_run_id=report_run_id + start_report=report_call_ended.spawn, + ): + result = calculate(...) + return result + +When the block ends, normally or by raising an ``Exception``, the report is +started as a separate call of the ``report_call_ended`` Modal function and +the block's own result or exception goes on unchanged. Starting it never +raises. A signal sends nothing: Modal delivers pre-emption as +``KeyboardInterrupt``, re-runs the call under the same ID, and the re-run +reports. Segment calls report to their coordinator, not to sim. + +Why a separate call: this function's own output exists only once it has +returned, and sim asks Modal for the outcome itself, without waiting, so a +report sent from inside the call would always find it still running. +``report_call_ended`` runs ``send_call_ended_report_from_container``: it +posts the ``CallEndedReport`` (the call ID and run IDs, never a result) to +``SIMULATION_ENTRY_URL`` + ``CALL_ENDED_PATH`` with the container's Modal +identity token, and reports again after a wait while sim answers 503 or +cannot be reached, up to ``CALL_REPORT_ATTEMPTS`` times. A report that never +arrives is recorded by sim's safety check within about ten minutes. + +This module is v2 code only: the main path's ``run_simulation`` does not use +it. +""" + +from __future__ import annotations + +import json +import logging +import os +import time +import urllib.error +import urllib.request +from contextlib import contextmanager +from enum import StrEnum +from typing import TYPE_CHECKING, Protocol + +from policyengine_simulation_contract.call_reports import ( + CALL_ENDED_PATH, + MODAL_IDENTITY_ENVIRONMENT_VARIABLE, + CallEndedReport, +) +from pydantic import ValidationError + +if TYPE_CHECKING: + from collections.abc import Callable, Iterator, Mapping, Sequence + from uuid import UUID + +logger = logging.getLogger(__name__) + +# The Simulation Entrypoint's main URL, which the deploy bakes into the +# report_call_ended function's environment. +ENTRY_URL_VARIABLE = "SIMULATION_ENTRY_URL" +CALL_REPORT_ATTEMPTS = 6 +FIRST_WAIT_SECONDS = 1.0 +LONGEST_WAIT_SECONDS = 30.0 +REQUEST_TIMEOUT_SECONDS = 15.0 +# sim answers these when it cannot record the outcome yet. +_RETRY_LATER_HTTP_STATUSES = frozenset({429, 502, 503, 504}) + + +class ReportDelivery(StrEnum): + """What sending one report came to.""" + + # sim recorded the outcome, or already had it. + DELIVERED = "delivered" + # sim refused the report; sending it again would not help. + REFUSED = "refused" + # sim could not record it after every attempt; the safety check will. + GAVE_UP = "gave_up" + # The container has no URL for sim or no identity token. + NOT_CONFIGURED = "not_configured" + + +class Poster(Protocol): + """Send one HTTP POST; return the status, headers and body.""" + + def __call__( + self, + url: str, + body: bytes, + headers: Mapping[str, str], + timeout: float, + ) -> tuple[int, Mapping[str, str], bytes]: + """Post ``body`` to ``url``; raise OSError when nothing came back.""" + ... + + +def post_with_urllib( + url: str, body: bytes, headers: Mapping[str, str], timeout: float +) -> tuple[int, Mapping[str, str], bytes]: + """Post with the standard library; an HTTP error status is an answer.""" + request = urllib.request.Request( # noqa: S310 # an https URL the deploy sets + url, data=body, headers=dict(headers), method="POST" + ) + try: + with urllib.request.urlopen(request, timeout=timeout) as response: # noqa: S310 # as above + return response.status, dict(response.headers), response.read() + except urllib.error.HTTPError as error: + return error.code, dict(error.headers), error.read() + + +def _wait_seconds(attempt: int, headers: Mapping[str, str]) -> float: + retry_after = {name.lower(): value for name, value in headers.items()}.get( + "retry-after", "" + ) + backoff = min(FIRST_WAIT_SECONDS * 2**attempt, LONGEST_WAIT_SECONDS) + if retry_after.isdigit(): + return min(max(float(retry_after), backoff), LONGEST_WAIT_SECONDS) + return backoff + + +def _answer_status(body: bytes) -> str: + try: + answer = json.loads(body) + except ValueError: + return "" + return str(answer.get("status", "")) if isinstance(answer, dict) else "" + + +def send_call_ended_report( + report: CallEndedReport, + *, + entry_url: str, + identity_token: str, + post: Poster = post_with_urllib, + sleep: Callable[[float], None] = time.sleep, +) -> ReportDelivery: + """Post ``report`` to sim, reporting again while sim cannot record it yet. + + Never raises an ``Exception``: a report that cannot be delivered is + logged, and sim's safety check records the call instead. + """ + if not entry_url or not identity_token: + logger.warning( + "call_report_not_configured", + extra={"modal_call_id": report.modal_call_id}, + ) + return ReportDelivery.NOT_CONFIGURED + url = entry_url.rstrip("/") + CALL_ENDED_PATH + body = report.model_dump_json().encode() + headers = { + "Authorization": f"Bearer {identity_token}", + "Content-Type": "application/json", + } + answer_headers: Mapping[str, str] = {} + for attempt in range(CALL_REPORT_ATTEMPTS): + if attempt: + sleep(_wait_seconds(attempt - 1, answer_headers)) + try: + status, answer_headers, answer_body = post( + url, body, headers, REQUEST_TIMEOUT_SECONDS + ) + except OSError as error: # no answer: unreachable, reset or timed out + answer_headers = {} + logger.warning( + "call_report_not_answered", + extra={ + "modal_call_id": report.modal_call_id, + "attempt": attempt + 1, + "error_type": type(error).__name__, + }, + ) + continue + fields = { + "modal_call_id": report.modal_call_id, + "attempt": attempt + 1, + "http_status": status, + "call_report_status": _answer_status(answer_body), + } + if status < 300: # noqa: PLR2004 # any 2xx answer + logger.info("call_report_delivered", extra=fields) + return ReportDelivery.DELIVERED + if status not in _RETRY_LATER_HTTP_STATUSES: + logger.warning("call_report_refused", extra=fields) + return ReportDelivery.REFUSED + logger.info("call_report_not_recorded_yet", extra=fields) + logger.warning( + "call_report_given_up", + extra={"modal_call_id": report.modal_call_id, "attempts": CALL_REPORT_ATTEMPTS}, + ) + return ReportDelivery.GAVE_UP + + +def send_call_ended_report_from_container( + report: Mapping[str, object], + *, + post: Poster = post_with_urllib, + sleep: Callable[[float], None] = time.sleep, +) -> str: + """Run in ``report_call_ended``: send ``report`` with this container's settings. + + Return the ``ReportDelivery`` value. A malformed report is refused. + """ + try: + call_report = CallEndedReport.model_validate(report) + except ValidationError: + logger.warning("call_report_malformed") + return ReportDelivery.REFUSED.value + return send_call_ended_report( + call_report, + entry_url=os.environ.get(ENTRY_URL_VARIABLE, ""), + identity_token=os.environ.get(MODAL_IDENTITY_ENVIRONMENT_VARIABLE, ""), + post=post, + sleep=sleep, + ).value + + +def _start_report( + report: CallEndedReport, start_report: Callable[[dict[str, object]], object] +) -> None: + try: + start_report(report.model_dump(mode="json")) + except Exception as error: # a failed report must never fail the calculation + logger.exception( + "call_report_not_started", + extra={ + "modal_call_id": report.modal_call_id, + "error_type": type(error).__name__, + }, + ) + + +@contextmanager +def reporting_call_end( + *, + modal_call_id: str | None, + start_report: Callable[[dict[str, object]], object], + simulation_run_ids: Sequence[UUID] | None = None, + report_run_id: UUID | None = None, +) -> Iterator[None]: + """Report the call's end when the block returns or raises an Exception. + + ``start_report`` starts the ``report_call_ended`` function (its + ``spawn``). Without a call ID (outside Modal), or with run IDs the report + cannot carry, nothing is reported. + """ + report: CallEndedReport | None = None + try: + if modal_call_id is not None: + report = CallEndedReport( + modal_call_id=modal_call_id, + simulation_run_ids=( + None if simulation_run_ids is None else tuple(simulation_run_ids) + ), + report_run_id=report_run_id, + ) + except ValidationError: + logger.warning("call_report_malformed", extra={"modal_call_id": modal_call_id}) + if report is None: + yield + return + try: + yield + except Exception: + _start_report(report, start_report) + raise + _start_report(report, start_report) diff --git a/projects/policyengine-simulation-executor/src/modal/v2_app.py b/projects/policyengine-simulation-executor/src/modal/v2_app.py index 69eedc493..ceb8ce9d2 100644 --- a/projects/policyengine-simulation-executor/src/modal/v2_app.py +++ b/projects/policyengine-simulation-executor/src/modal/v2_app.py @@ -94,6 +94,23 @@ def _deploy_time_cache_inputs() -> tuple[str, dict | None]: STAGE12_CACHE_BUCKET, STAGE12_CACHE_MANIFEST = _deploy_time_cache_inputs() +def _deploy_time_entry_url() -> str: + """Return sim's main URL, which the deploy passes, for report_call_ended.""" + if not modal.is_local(): + return "" + url = os.environ.get("SIMULATION_ENTRY_URL", "") + if url and (url != url.strip() or not url.startswith("https://")): + msg = "SIMULATION_ENTRY_URL must be an https URL" + raise RuntimeError(msg) + return url + + +SIMULATION_ENTRY_URL = _deploy_time_entry_url() +# Every attempt to report, each with its request timeout, and every wait +# between them (src/modal/call_reports.py), with a margin. +CALL_REPORT_TIMEOUT_SECONDS = 300 + + def _external_assertions(environment: dict[str, str]) -> dict[str, str]: mapping = { "STAGE12_EXPECT_POLICYENGINE_VERSION": "policyengine_version", @@ -463,3 +480,29 @@ def coordinate_report( ) finally: runtime.shutdown() + + +@app.function( + image=coordinator_image, + cpu=0.25, + memory=512, + timeout=CALL_REPORT_TIMEOUT_SECONDS, + retries=0, + max_containers=20, + # No secrets: sim accepts this function's own Modal identity token. + env={"SIMULATION_ENTRY_URL": SIMULATION_ENTRY_URL}, +) +def report_call_ended(report: dict) -> str: + """Tell sim that the call in ``report`` ended; return how that went. + + The v2 functions start this as their last step + (``src.modal.call_reports.reporting_call_end``) and then return, so the + reported call has usually returned by the time sim asks Modal for its + outcome; when it has not, sim answers still_running and this reports + again after a wait. + """ + from src.modal.call_reports import ( # noqa: PLC0415 # Modal function body: loaded in the worker container, not at deploy time + send_call_ended_report_from_container, + ) + + return send_call_ended_report_from_container(report) diff --git a/projects/policyengine-simulation-executor/tests/test_call_report_function.py b/projects/policyengine-simulation-executor/tests/test_call_report_function.py new file mode 100644 index 000000000..9ce5cd33d --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_call_report_function.py @@ -0,0 +1,84 @@ +"""The v2 app's report_call_ended function, which tells sim that a call ended.""" + +from __future__ import annotations + +import os +import subprocess +import sys +from pathlib import Path + +import pytest +from policyengine_simulation_contract.call_reports import CALL_REPORT_FUNCTION_NAME +from src.modal import call_reports +from test_stage12_modal_app import _load + +ENTRY_URL = "https://policyengine-simulation-entry-abc.a.run.app" +PROJECT_ROOT = Path(__file__).resolve().parents[1] + + +def test_the_v2_app_declares_the_function_sim_accepts_reports_from( + monkeypatch, +) -> None: + monkeypatch.setenv("SIMULATION_ENTRY_URL", ENTRY_URL) + module = _load(monkeypatch) + + declaration = dict(module.app.function_calls)[CALL_REPORT_FUNCTION_NAME] + + assert declaration["image"] is module.coordinator_image + assert declaration["env"] == {"SIMULATION_ENTRY_URL": ENTRY_URL} + assert declaration["retries"] == 0 + # Every attempt, each with its request timeout, and every wait between them. + assert declaration["timeout"] >= call_reports.CALL_REPORT_ATTEMPTS * ( + call_reports.REQUEST_TIMEOUT_SECONDS + call_reports.LONGEST_WAIT_SECONDS + ) + # It needs no secret: its own identity token proves who it is. + assert "secrets" not in declaration + + +def test_the_function_sends_the_report_with_its_containers_settings( + monkeypatch, +) -> None: + module = _load(monkeypatch) + sent = [] + monkeypatch.setattr( + call_reports, + "send_call_ended_report_from_container", + lambda report: sent.append(report) or "delivered", + ) + report = {"modal_call_id": "fc-ended", "report_run_id": "x"} + + assert module.report_call_ended(report) == "delivered" + assert sent == [report] + + +@pytest.mark.parametrize("url", ["http://entry.example", "entry.example", " https://x"]) +def test_a_deploy_with_a_url_that_is_not_https_stops(monkeypatch, url) -> None: + monkeypatch.setenv("SIMULATION_ENTRY_URL", url) + + with pytest.raises(RuntimeError, match="SIMULATION_ENTRY_URL"): + _load(monkeypatch) + + +def test_the_v2_app_imports_with_the_real_modal_package() -> None: + """`modal deploy` imports the module with Modal's own decorators. + + There a declared function is a ``modal.Function``, not the decorated + Python function, so nothing at import may rely on its Python attributes. + """ + environment = { + name: value + for name, value in os.environ.items() + if not name.startswith(("MODAL_", "STAGE12_", "SIMULATION_ENTRY_")) + } + + result = subprocess.run( + [sys.executable, "-c", "import src.modal.v2_app"], + cwd=PROJECT_ROOT, + env=environment, + capture_output=True, + text=True, + check=False, + timeout=300, + ) + + assert result.returncode == 0, result.stderr[-2000:] diff --git a/projects/policyengine-simulation-executor/tests/test_call_reports.py b/projects/policyengine-simulation-executor/tests/test_call_reports.py new file mode 100644 index 000000000..4d33d5219 --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_call_reports.py @@ -0,0 +1,237 @@ +"""Reporting to the Simulation Entrypoint that a v2 Modal call ended.""" + +from __future__ import annotations + +import json +import urllib.error +from dataclasses import dataclass, field +from uuid import uuid4 + +import pytest +from policyengine_simulation_contract.call_reports import ( + CALL_ENDED_PATH, + CallEndedReport, +) +from src.modal.call_reports import ( + CALL_REPORT_ATTEMPTS, + ReportDelivery, + reporting_call_end, + send_call_ended_report, + send_call_ended_report_from_container, +) + +ENTRY_URL = "https://policyengine-simulation-entry-abc.a.run.app" +RUN_ID = uuid4() +REPORT = CallEndedReport(modal_call_id="fc-ended", simulation_run_ids=(RUN_ID,)) + + +@dataclass +class _Entrypoint: + """Answer each request with the next scripted (status, headers, body) or error.""" + + answers: list + requests: list = field(default_factory=list) + + def __call__(self, url, body, headers, timeout): + self.requests.append({"url": url, "body": body, "headers": headers}) + answer = self.answers.pop(0) + if isinstance(answer, BaseException): + raise answer + status, answer_headers, answer_body = answer + return status, answer_headers, json.dumps(answer_body).encode() + + +@dataclass +class _Waits: + seconds: list = field(default_factory=list) + + def __call__(self, seconds: float) -> None: + self.seconds.append(seconds) + + +def _send(answers, **settings): + entrypoint, waits = _Entrypoint(list(answers)), _Waits() + delivery = send_call_ended_report( + REPORT, + entry_url=settings.get("entry_url", ENTRY_URL), + identity_token=settings.get("identity_token", "modal-identity-token"), + post=entrypoint, + sleep=waits, + ) + return delivery, entrypoint, waits + + +def test_a_report_is_posted_with_the_containers_identity_token() -> None: + delivery, entrypoint, waits = _send([(200, {}, {"status": "recorded"})]) + + assert delivery is ReportDelivery.DELIVERED + (request,) = entrypoint.requests + assert request["url"] == ENTRY_URL + CALL_ENDED_PATH + assert request["headers"]["Authorization"] == "Bearer modal-identity-token" + assert request["headers"]["Content-Type"] == "application/json" + assert CallEndedReport.model_validate_json(request["body"]) == REPORT + assert waits.seconds == [] + + +def test_a_report_sim_cannot_record_yet_is_sent_again_after_its_wait() -> None: + delivery, entrypoint, waits = _send( + [ + (503, {"Retry-After": "3"}, {"status": "still_running"}), + (503, {}, {"status": "not_dispatched_yet"}), + (200, {}, {"status": "already_recorded"}), + ] + ) + + assert delivery is ReportDelivery.DELIVERED + assert len(entrypoint.requests) == 3 + assert waits.seconds[0] == 3 + assert waits.seconds[1] > 0 + + +@pytest.mark.parametrize( + "failure", + [urllib.error.URLError("unreachable"), TimeoutError(), ConnectionResetError()], + ids=["unreachable", "timeout", "reset"], +) +def test_a_report_that_does_not_arrive_is_sent_again(failure) -> None: + delivery, entrypoint, _ = _send([failure, (200, {}, {"status": "recorded"})]) + + assert delivery is ReportDelivery.DELIVERED + assert len(entrypoint.requests) == 2 + + +@pytest.mark.parametrize("status", [401, 404, 409, 422]) +def test_a_refused_report_is_not_sent_again(status) -> None: + delivery, entrypoint, waits = _send([(status, {}, {"status": "other_call"})]) + + assert delivery is ReportDelivery.REFUSED + assert len(entrypoint.requests) == 1 + assert waits.seconds == [] + + +def test_the_sender_gives_up_after_a_few_attempts_with_growing_waits() -> None: + answers = [(503, {}, {"status": "still_running"})] * CALL_REPORT_ATTEMPTS + + delivery, entrypoint, waits = _send(answers) + + assert delivery is ReportDelivery.GAVE_UP + assert len(entrypoint.requests) == CALL_REPORT_ATTEMPTS + assert len(waits.seconds) == CALL_REPORT_ATTEMPTS - 1 + assert waits.seconds == sorted(waits.seconds) + assert sum(waits.seconds) < 120 + + +def test_an_unexpected_answer_is_not_a_crash() -> None: + delivery, _, _ = _send([(200, {}, "not an answer")]) + + assert delivery is ReportDelivery.DELIVERED + + +def test_a_report_without_sims_url_or_a_token_is_not_sent() -> None: + delivery, entrypoint, _ = _send([], entry_url="") + assert delivery is ReportDelivery.NOT_CONFIGURED + delivery, entrypoint, _ = _send([], identity_token="") + assert delivery is ReportDelivery.NOT_CONFIGURED + assert entrypoint.requests == [] + + +def test_the_report_function_reads_its_settings_from_the_container( + monkeypatch, +) -> None: + monkeypatch.setenv("SIMULATION_ENTRY_URL", ENTRY_URL) + monkeypatch.setenv("MODAL_IDENTITY_TOKEN", "modal-identity-token") + entrypoint = _Entrypoint([(200, {}, {"status": "recorded"})]) + + delivery = send_call_ended_report_from_container( + REPORT.model_dump(mode="json"), post=entrypoint, sleep=_Waits() + ) + + assert delivery == "delivered" + assert entrypoint.requests[0]["url"] == ENTRY_URL + CALL_ENDED_PATH + + +@dataclass +class _ReportStarter: + """Stands in for ``report_call_ended.spawn``.""" + + failure: BaseException | None = None + started: list = field(default_factory=list) + + def __call__(self, report: dict) -> None: + if self.failure is not None: + raise self.failure + self.started.append(report) + + +def test_a_call_that_returns_reports_its_end() -> None: + starter = _ReportStarter() + + with reporting_call_end( + modal_call_id="fc-ended", simulation_run_ids=(RUN_ID,), start_report=starter + ): + pass + + assert [CallEndedReport.model_validate(report) for report in starter.started] == [ + REPORT + ] + + +def _failing_comparison(report_run_id, starter) -> None: + with reporting_call_end( + modal_call_id="fc-failed", report_run_id=report_run_id, start_report=starter + ): + msg = "calculation failed" + raise ValueError(msg) + + +def test_a_call_that_fails_reports_its_end_and_still_fails() -> None: + starter = _ReportStarter() + report_run_id = uuid4() + + with pytest.raises(ValueError, match="calculation failed"): + _failing_comparison(report_run_id, starter) + + assert starter.started == [ + { + "modal_call_id": "fc-failed", + "simulation_run_ids": None, + "report_run_id": str(report_run_id), + } + ] + + +@pytest.mark.parametrize("signal", [KeyboardInterrupt(), SystemExit(1)]) +def test_a_pre_empted_call_reports_nothing(signal) -> None: + """Modal re-runs a pre-empted call under its ID, and the re-run reports.""" + starter = _ReportStarter() + + with ( + pytest.raises(type(signal)), + reporting_call_end( + modal_call_id="fc-ended", simulation_run_ids=(RUN_ID,), start_report=starter + ), + ): + raise signal + + assert starter.started == [] + + +def test_a_report_that_cannot_start_never_fails_the_calculation() -> None: + starter = _ReportStarter(failure=RuntimeError("Modal unreachable")) + + with reporting_call_end( + modal_call_id="fc-ended", simulation_run_ids=(RUN_ID,), start_report=starter + ): + pass + + +def test_a_call_without_a_modal_call_id_reports_nothing() -> None: + """Run outside Modal (a local test), the call has no ID to report.""" + starter = _ReportStarter() + + with reporting_call_end( + modal_call_id=None, simulation_run_ids=(RUN_ID,), start_report=starter + ): + pass + + assert starter.started == [] diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py b/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py index 671aa4d12..1850f4bda 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py @@ -326,6 +326,7 @@ def test_v2_app_declares_validation_workers_and_non_http_coordinator( "run_single_simulation_us", "run_single_simulation_uk", "coordinate_report", + "report_call_ended", } assert functions["run_single_simulation_us"]["image"] is module.us_worker_image assert functions["run_single_simulation_uk"]["image"] is module.uk_worker_image