Move the simulation executor to policyengine.py 6.14.0 and keep accepting US 2022 and 2023

changelog.d/1012.changed.md

@@ -0,0 +1 @@
+Updated policyengine.py from 6.12.0 to 6.14.0 in the simulation runtime. Simulations for US 2022 and 2023 still run: policyengine.py now refuses dataset years before the US data year by default, and the simulation service opts out of that check until the Public API stops offering those years.

projects/policyengine-simulation-executor/pyproject.toml

@@ -21,7 +21,7 @@ dependencies = [
     "policyengine-stage12-persistence",
     # PolicyEngine.py's models extra is the only selector for core, country
     # models, and the SPM calculator. Do not duplicate those versions here.
-    "policyengine[models]==6.12.0",
+    "policyengine[models]==6.14.0",
     "tables>=3.10.2",
     "modal>=0.73.0",
     "policyengine-observability[fastapi,google,otlp-grpc]>=3.0.4,<4",
@@ -46,7 +46,7 @@ packages = ["src/policyengine_simulation_executor"]
 # certified datasets and writes their receipt.
 [dependency-groups]
 modal-simulation-image = [
-    "policyengine[models]==6.12.0",
+    "policyengine[models]==6.14.0",
     "fastapi>=0.115.0",
     "tables>=3.10.2",
     "policyengine-observability[fastapi,google,otlp-grpc]>=3.0.4,<4",

projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py

@@ -244,7 +244,16 @@ def _load_dataset(
         )
 
     start = time.monotonic()
-    load_options = {"years": [year], "data_folder": data_folder}
+    # Temporary: .py refuses a year its source dataset doesn't support, such
+    # as US 2022 and 2023, but the Public API still offers those years for
+    # society-wide runs. The simulation service accepts every year the Public
+    # API offers until the Public API stops offering US 2022 and 2023 and
+    # refuses them with a 400; then this repository removes the opt-out.
+    load_options: dict[str, Any] = {
+        "years": [year],
+        "data_folder": data_folder,
+        "allow_unsupported_years": True,
+    }
     if not selection.is_default:
         load_options["datasets"] = [selection.name]
     datasets = country_module.ensure_datasets(**load_options)

projects/policyengine-simulation-executor/tests/test_simulation_output_dataset_selection.py

@@ -4,6 +4,7 @@
 from types import SimpleNamespace
 
 import pytest
+from policyengine.provenance import UnsupportedDatasetYearError
 from policyengine_simulation_executor.release_bundle import (
     BUNDLE_RECEIPT_FILENAME,
     get_country_release_bundle,
@@ -113,6 +114,64 @@ def ensure_datasets(**kwargs):
         {
             "years": [2026],
             "data_folder": str(tmp_path),
+            "allow_unsupported_years": True,
+        }
+    ]
+
+
+@pytest.mark.parametrize("year", [2022, 2023])
+def test_load_dataset_lets_a_request_load_a_us_year_before_the_data_year(
+    year, tmp_path, monkeypatch
+):
+    """A request for US 2022 or 2023 loads, although .py refuses those years.
+
+    The fake loader has .py's keyword-only ``ensure_datasets`` arguments and,
+    like .py since 6.13.0, raises ``UnsupportedDatasetYearError`` for a year
+    before the 2024 data year unless the caller passes
+    ``allow_unsupported_years=True``.
+    """
+    monkeypatch.setenv("POLICYENGINE_DATA_FOLDER", str(tmp_path))
+    country = "us"
+    bundle = get_country_release_bundle(country)
+    supported_years = list(range(2024, 2036))
+    loaded_dataset = SimpleNamespace()
+    ensure_calls = []
+
+    def ensure_datasets(
+        *,
+        datasets=None,
+        years,
+        data_folder,
+        allow_unmanaged=False,
+        allow_unsupported_years=False,
+    ):
+        ensure_calls.append(
+            {
+                "years": years,
+                "data_folder": data_folder,
+                "allow_unsupported_years": allow_unsupported_years,
+            }
+        )
+        unsupported = [
+            requested for requested in years if requested not in supported_years
+        ]
+        if unsupported and not allow_unsupported_years:
+            raise UnsupportedDatasetYearError(
+                country, bundle.default_dataset, unsupported, supported_years
+            )
+        return {bundle.default_dataset: loaded_dataset}
+
+    loaded = _load_dataset(
+        {"country": country, "time_period": str(year)},
+        country_module=SimpleNamespace(ensure_datasets=ensure_datasets),
+    )
+
+    assert loaded is loaded_dataset
+    assert ensure_calls == [
+        {
+            "years": [year],
+            "data_folder": str(tmp_path),
+            "allow_unsupported_years": True,
         }
     ]
 
@@ -193,6 +252,7 @@ def ensure_datasets(**kwargs):
             "data_folder": _nondefault_data_folder(
                 "us", nondefault, bundle.dataset_uris[nondefault]
             ),
+            "allow_unsupported_years": True,
         }
     ]
 
@@ -231,6 +291,7 @@ def ensure_datasets(**kwargs):
                 if use_default
                 else _nondefault_data_folder("us", name, bundle.dataset_uris[name])
             ),
+            "allow_unsupported_years": True,
         }
     ]
 
@@ -300,7 +361,12 @@ def ensure_datasets(**kwargs):
     )
 
     assert ensure_calls == [
-        {"datasets": [nondefault], "years": [2026], "data_folder": state_folder}
+        {
+            "datasets": [nondefault],
+            "years": [2026],
+            "data_folder": state_folder,
+            "allow_unsupported_years": True,
+        }
     ]