Deploy Modal worker for API v4 migration (#466)

Author: MaxGhenis

projects/policyengine-api-simulation/pyproject.toml

@@ -18,8 +18,8 @@ dependencies = [
     "policyengine-fastapi",
     "policyengine==0.13.0",
     "policyengine-core>=3.23.5",
-    "policyengine-uk>=2.22.8",
-    "policyengine-us>=1.370.2",
+    "policyengine-uk==2.88.0",
+    "policyengine-us==1.653.3",
     "tables>=3.10.2",
     "modal>=0.73.0",
     "logfire>=3.0.0",

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

@@ -73,6 +73,7 @@ def _build_policyengine_bundle(
         resolved_dataset = None
     return PolicyEngineBundle(
         model_version=resolved_version,
+        data_version=payload.get("data_version"),
         dataset=resolved_dataset,
     )
 

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

@@ -20,7 +20,7 @@
 MAX_GATEWAY_REQUEST_BYTES = 262_144
 
 
-INTERNAL_PASSTHROUGH_FIELDS = frozenset({"_metadata"})
+INTERNAL_PASSTHROUGH_FIELDS = frozenset({"_metadata", "_runtime_bundle"})
 
 
 def _move_internal_telemetry_alias(value):
@@ -35,12 +35,11 @@ def _move_internal_telemetry_alias(value):
 def _strip_internal_passthrough_fields(value):
     """Drop internal-only fields the gateway adds to payloads in flight.
 
-    When the parent batch entrypoint forwards a payload to the worker we
-    attach ``_metadata`` describing resolved routing. That enrichment is
-    consumed by :mod:`src.modal.budget_window_context`, not by the request
-    model itself. We strip it before strict validation so ``extra="forbid"``
-    keeps catching unknown fields from *callers* without breaking the
-    internal round-trip.
+    Parent batch entrypoints attach ``_metadata`` describing resolved routing,
+    and the v4 API attaches ``_runtime_bundle`` describing provenance that the
+    gateway returns separately. Strip those fields before strict validation so
+    ``extra="forbid"`` keeps catching unknown caller fields without breaking
+    internal round-trips.
     """
 
     if not isinstance(value, dict):

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

@@ -311,6 +311,40 @@ def test__given_submission_with_uk_alias_data__then_bundle_dataset_is_versioned_
             == "hf://policyengine/policyengine-uk-data-private/enhanced_frs_2023_24.h5@1.40.3"
         )
 
+    def test__given_submission_with_runtime_bundle__then_accepts_internal_provenance(
+        self, mock_modal, client: TestClient
+    ):
+        mock_modal["dicts"]["simulation-api-us-versions"] = {
+            "latest": "1.500.0",
+            "1.500.0": "policyengine-simulation-us1-500-0-uk2-66-0",
+        }
+
+        request_body = {
+            "country": "us",
+            "scope": "macro",
+            "reform": {},
+            "data": "enhanced_cps_2024",
+            "data_version": "1.78.2",
+            "_runtime_bundle": {
+                "model_version": "1.500.0",
+                "data_version": "1.78.2",
+            },
+            "_metadata": {"process_id": "process-123"},
+        }
+
+        response = client.post("/simulate/economy/comparison", json=request_body)
+
+        assert response.status_code == 200
+        data = response.json()
+        assert data["policyengine_bundle"] == {
+            "model_version": "1.500.0",
+            "data_version": "1.78.2",
+            "dataset": "hf://policyengine/policyengine-us-data/enhanced_cps_2024.h5@1.77.0",
+        }
+        assert mock_modal["func"].last_payload["data_version"] == "1.78.2"
+        assert "_runtime_bundle" not in mock_modal["func"].last_payload
+        assert "_metadata" not in mock_modal["func"].last_payload
+
     def test__given_submission_with_unknown_alias_data__then_bundle_dataset_is_preserved(
         self, mock_modal, client: TestClient
     ):