fix: package simulation runtime for integration

Author: anth-volk

.github/scripts/modal-deploy-app.sh

@@ -1,24 +1,24 @@
 #!/bin/bash
 # Deploy simulation API to Modal
 # Usage: ./modal-deploy-app.sh <modal-environment>
-# Required env vars: POLICYENGINE_US_VERSION, POLICYENGINE_UK_VERSION
+# Required env vars: POLICYENGINE_VERSION, POLICYENGINE_US_VERSION, POLICYENGINE_UK_VERSION
 # These should come from the bundled policyengine.py release manifest.
 #
 # Deploys two apps:
 # 1. policyengine-simulation-gateway - Stable gateway with fixed URL
-# 2. policyengine-simulation-us{X}-uk{Y} - Versioned simulation app
+# 2. policyengine-simulation-py{X} - Versioned simulation app
 
 set -euo pipefail
 
 MODAL_ENV="${1:?Modal environment required}"
 
 # Generate versioned simulation app name (dots replaced with dashes for URL safety)
-US_VERSION_SAFE="${POLICYENGINE_US_VERSION//./-}"
-UK_VERSION_SAFE="${POLICYENGINE_UK_VERSION//./-}"
-SIMULATION_APP_NAME="policyengine-simulation-us${US_VERSION_SAFE}-uk${UK_VERSION_SAFE}"
+POLICYENGINE_VERSION_SAFE="${POLICYENGINE_VERSION//./-}"
+SIMULATION_APP_NAME="policyengine-simulation-py${POLICYENGINE_VERSION_SAFE}"
 
 echo "========================================"
 echo "Deploying to Modal environment: $MODAL_ENV"
+echo "  policyengine.py version: ${POLICYENGINE_VERSION}"
 echo "  US version: ${POLICYENGINE_US_VERSION}"
 echo "  UK version: ${POLICYENGINE_UK_VERSION}"
 echo "========================================"
@@ -41,6 +41,7 @@ echo ""
 echo "Step 3: Updating version registries..."
 uv run python -m src.modal.utils.update_version_registry \
     --app-name "$SIMULATION_APP_NAME" \
+    --policyengine-version "${POLICYENGINE_VERSION}" \
     --us-version "${POLICYENGINE_US_VERSION}" \
     --uk-version "${POLICYENGINE_UK_VERSION}" \
     --environment "$MODAL_ENV"

.github/workflows/modal-deploy.reusable.yml

@@ -15,6 +15,9 @@ on:
       simulation_api_url:
         description: 'The deployed simulation API URL'
         value: ${{ jobs.deploy.outputs.simulation_api_url }}
+      policyengine_version:
+        description: 'The deployed policyengine.py package version'
+        value: ${{ jobs.deploy.outputs.policyengine_version }}
       us_version:
         description: 'The deployed policyengine-us package version'
         value: ${{ jobs.deploy.outputs.us_version }}
@@ -35,6 +38,7 @@ jobs:
     environment: ${{ inputs.environment }}
     outputs:
       simulation_api_url: ${{ steps.get-url.outputs.simulation_api_url }}
+      policyengine_version: ${{ steps.versions.outputs.policyengine_version }}
       us_version: ${{ steps.versions.outputs.us_version }}
       us_data_version: ${{ steps.versions.outputs.us_data_version }}
       uk_version: ${{ steps.versions.outputs.uk_version }}
@@ -82,9 +86,10 @@ jobs:
       env:
         MODAL_TOKEN_ID: ${{ secrets.MODAL_TOKEN_ID }}
         MODAL_TOKEN_SECRET: ${{ secrets.MODAL_TOKEN_SECRET }}
+        POLICYENGINE_VERSION: ${{ steps.versions.outputs.policyengine_version }}
         POLICYENGINE_US_VERSION: ${{ steps.versions.outputs.us_version }}
         POLICYENGINE_UK_VERSION: ${{ steps.versions.outputs.uk_version }}
-      run: ../../.github/scripts/modal-deploy-app.sh "${{ inputs.modal_environment }}" src/modal/app.py
+      run: ../../.github/scripts/modal-deploy-app.sh "${{ inputs.modal_environment }}"
 
     - name: Get deployed URL
       id: get-url

projects/policyengine-api-simulation/fixtures/gateway/test_endpoints.py

@@ -2,6 +2,69 @@
 
 import pytest
 
+TEST_APP_RELEASE_BUNDLE = {
+    "app_name": "policyengine-simulation-py4-10-0",
+    "policyengine_version": "4.10.0",
+    "us": {
+        "model_version": "1.500.0",
+        "data_version": "1.110.12",
+        "default_dataset": "enhanced_cps_2024",
+        "default_dataset_uri": "hf://policyengine/policyengine-us-data/enhanced_cps_2024.h5@1.110.12",
+        "dataset_uris": {
+            "enhanced_cps_2024": "hf://policyengine/policyengine-us-data/enhanced_cps_2024.h5@1.110.12",
+            "cps_2023": "hf://policyengine/policyengine-us-data/cps_2023.h5@1.110.12",
+            "pooled_3_year_cps_2023": "hf://policyengine/policyengine-us-data/pooled_3_year_cps_2023.h5@1.110.12",
+        },
+        "dataset_aliases": {
+            "enhanced_cps": "enhanced_cps_2024",
+            "enhanced_cps_2024": "enhanced_cps_2024",
+            "cps": "cps_2023",
+            "cps_2023": "cps_2023",
+            "pooled_cps": "pooled_3_year_cps_2023",
+            "pooled_3_year_cps_2023": "pooled_3_year_cps_2023",
+        },
+    },
+    "uk": {
+        "model_version": "2.66.0",
+        "data_version": "1.40.3",
+        "default_dataset": "enhanced_frs_2023_24",
+        "default_dataset_uri": "hf://policyengine/policyengine-uk-data-private/enhanced_frs_2023_24.h5@1.40.3",
+        "dataset_uris": {
+            "enhanced_frs_2023_24": "hf://policyengine/policyengine-uk-data-private/enhanced_frs_2023_24.h5@1.40.3",
+            "frs_2023_24": "hf://policyengine/policyengine-uk-data-private/frs_2023_24.h5@1.40.3",
+        },
+        "dataset_aliases": {
+            "enhanced_frs": "enhanced_frs_2023_24",
+            "enhanced_frs_2023_24": "enhanced_frs_2023_24",
+            "frs": "frs_2023_24",
+            "frs_2023_24": "frs_2023_24",
+        },
+    },
+}
+
+TEST_APP_NAMES = (
+    "policyengine-simulation-py4-10-0",
+    "policyengine-simulation-py3-9-0",
+)
+
+
+def resolve_test_dataset_uri(country: str, dataset: str | None) -> str | None:
+    if dataset is None:
+        return None
+    if "://" in dataset:
+        return dataset
+    country_bundle = TEST_APP_RELEASE_BUNDLE[country]
+    dataset_name, revision = (
+        dataset.rsplit("@", maxsplit=1) if "@" in dataset else (dataset, None)
+    )
+    dataset_name = country_bundle["dataset_aliases"].get(dataset_name, dataset_name)
+    dataset_uri = country_bundle["dataset_uris"].get(dataset_name, dataset_name)
+    if revision is not None and dataset_uri == dataset_name:
+        return dataset
+    if revision is not None and dataset_uri.startswith("hf://"):
+        dataset_uri = f"{dataset_uri.rsplit('@', maxsplit=1)[0]}@{revision}"
+    return dataset_uri
+
 
 class MockDict:
     """Mock for Modal.Dict to simulate version registry."""
@@ -107,7 +170,11 @@ def mock_modal(monkeypatch):
     from src.modal.gateway import endpoints
 
     mock_func = MockFunction()
-    mock_dicts = {}
+    mock_dicts = {
+        "simulation-api-app-release-bundles": {
+            app_name: TEST_APP_RELEASE_BUNDLE for app_name in TEST_APP_NAMES
+        }
+    }
     MockFunctionCall.registry = {}
     MockFunctionCall.from_id_errors = {}
 
@@ -134,6 +201,20 @@ class MockModal:
 
     monkeypatch.setattr(endpoints, "modal", MockModal)
     monkeypatch.setattr(budget_window_state, "modal", MockModal)
+    monkeypatch.setattr(
+        endpoints,
+        "with_hf_revision",
+        lambda dataset_uri, revision: (
+            f"{dataset_uri.rsplit('@', maxsplit=1)[0]}@{revision}"
+            if dataset_uri.startswith("hf://")
+            else dataset_uri
+        ),
+    )
+    monkeypatch.setattr(
+        endpoints,
+        "validate_hf_dataset_uri",
+        lambda dataset_uri: dataset_uri,
+    )
 
     return {
         "func": mock_func,

projects/policyengine-api-simulation/pyproject.toml

@@ -9,7 +9,7 @@ readme = "README.md"
 authors = [
     {name = "PolicyEngine", email = "hello@policyengine.org"},
 ]
-license = {file = "../../LICENSE"}
+license = "AGPL-3.0-only"
 requires-python = ">=3.13,<3.14"
 dependencies = [
     "opentelemetry-instrumentation-sqlalchemy (>=0.51b0,<0.52)",

projects/policyengine-api-simulation/src/modal/app.py

@@ -2,7 +2,7 @@
 PolicyEngine Simulation - Versioned Modal App
 
 This app contains the heavy simulation workload with snapshotted models.
-Each deployment creates a versioned app (e.g., policyengine-simulation-us1-459-0-uk2-65-9).
+Each deployment creates a versioned app (e.g., policyengine-simulation-py4-10-0).
 
 The gateway app (policyengine-simulation-gateway) routes requests to these versioned apps.
 """
@@ -13,7 +13,7 @@
 from src.modal._image_setup import snapshot_models
 from src.modal.dependency_pins import project_dependency_pin
 from src.modal.logging_redaction import redact_params_for_logging
-from src.modal.release_bundle import get_bundled_country_model_version
+from policyengine_api_simulation.release_bundle import get_bundled_country_model_version
 
 POLICYENGINE_VERSION = os.environ.get("POLICYENGINE_VERSION") or project_dependency_pin(
     "policyengine"
@@ -31,20 +31,19 @@
 ) or get_bundled_country_model_version("uk")
 
 
-def get_app_name(us_version: str, uk_version: str) -> str:
+def get_app_name(policyengine_version: str) -> str:
     """
-    Generate versioned app name from package versions.
+    Generate versioned app name from the policyengine.py package version.
 
     Replaces dots with dashes for URL safety.
-    Example: us1.459.0, uk2.65.9 -> policyengine-simulation-us1-459-0-uk2-65-9
+    Example: 4.10.0 -> policyengine-simulation-py4-10-0
     """
-    us_safe = us_version.replace(".", "-")
-    uk_safe = uk_version.replace(".", "-")
-    return f"policyengine-simulation-us{us_safe}-uk{uk_safe}"
+    policyengine_safe = policyengine_version.replace(".", "-")
+    return f"policyengine-simulation-py{policyengine_safe}"
 
 
 # App name can be overridden via environment variable, otherwise generated from versions
-APP_NAME = os.environ.get("MODAL_APP_NAME", get_app_name(US_VERSION, UK_VERSION))
+APP_NAME = os.environ.get("MODAL_APP_NAME", get_app_name(POLICYENGINE_VERSION))
 
 # App definition with versioned name
 app = modal.App(APP_NAME)
@@ -67,7 +66,11 @@ def get_app_name(us_version: str, uk_version: str) -> str:
         "tables>=3.10.2",
         "logfire",
     )
-    .add_local_python_source("src.modal", copy=True)
+    .add_local_python_source(
+        "src.modal",
+        "policyengine_api_simulation",
+        copy=True,
+    )
     .run_function(snapshot_models)
 )
 
@@ -106,7 +109,7 @@ def run_simulation(params: dict) -> dict:
     """
     import logfire
 
-    from src.modal.simulation import run_simulation_impl
+    from policyengine_api_simulation.simulation_runtime import run_simulation_impl
 
     configure_logfire()
 

projects/policyengine-api-simulation/src/modal/gateway/app.py

@@ -26,7 +26,11 @@
         # the auth module at runtime here.
         "cryptography>=41.0.0",
     )
-    .add_local_python_source("src.modal", copy=True)
+    .add_local_python_source(
+        "src.modal",
+        "policyengine_api_simulation",
+        copy=True,
+    )
     .add_local_python_source("policyengine_fastapi", copy=True)
 )