From 9196aa806c5834a24b0701ca72d043586e1e6140 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 5 Nov 2026 17:55:51 +0400 Subject: [PATCH] Move the executor to policyengine.py 6.14.0 and keep US 2022 and 2023 policyengine.py 6.13.0 made create_datasets and ensure_datasets raise UnsupportedDatasetYearError for a year the source dataset doesn't support, which for the US means 2022 and 2023. The Public API still offers those years for society-wide runs, so _load_dataset now passes allow_unsupported_years=True until the Public API refuses them. The precompute calls to ensure_datasets stay strict. Fixes #1012 Co-Authored-By: Claude Opus 5.5 --- changelog.d/1012.changed.md | 1 + .../pyproject.toml | 4 +- .../simulation_runtime.py | 11 ++- ...est_simulation_output_dataset_selection.py | 68 ++++++++++++++++++- .../policyengine-simulation-executor/uv.lock | 10 +-- 5 files changed, 85 insertions(+), 9 deletions(-) create mode 100644 changelog.d/1012.changed.md diff --git a/changelog.d/1012.changed.md b/changelog.d/1012.changed.md new file mode 100644 index 000000000..a7d52580b --- /dev/null +++ b/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. diff --git a/projects/policyengine-simulation-executor/pyproject.toml b/projects/policyengine-simulation-executor/pyproject.toml index db24a0567..f44caddf9 100644 --- a/projects/policyengine-simulation-executor/pyproject.toml +++ b/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", diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py index f6a81a6d1..5a9e88595 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py +++ b/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) diff --git a/projects/policyengine-simulation-executor/tests/test_simulation_output_dataset_selection.py b/projects/policyengine-simulation-executor/tests/test_simulation_output_dataset_selection.py index 70d66cf96..d02138e2c 100644 --- a/projects/policyengine-simulation-executor/tests/test_simulation_output_dataset_selection.py +++ b/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, + } ] diff --git a/projects/policyengine-simulation-executor/uv.lock b/projects/policyengine-simulation-executor/uv.lock index 9163cc649..895ffd8bd 100644 --- a/projects/policyengine-simulation-executor/uv.lock +++ b/projects/policyengine-simulation-executor/uv.lock @@ -1747,7 +1747,7 @@ wheels = [ [[package]] name = "policyengine" -version = "6.12.0" +version = "6.14.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "h5py" }, @@ -1759,9 +1759,9 @@ dependencies = [ { name = "pydantic" }, { name = "requests" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/c4/ed/6f1b85109092d740a8121c98933c4848af6e6355241012688b13347a7d97/policyengine-6.12.0.tar.gz", hash = "sha256:b8d4a6346206015c508f416f2e48e150a37df3d3103966e0b294d4f15e269cf6", size = 1590873, upload-time = "2026-11-02T14:23:39.916Z" } +sdist = { url = "https://files.pythonhosted.org/packages/67/ec/5e8c266599d461c355c7bfdb84a5fa5a5e15a86e04995874e82eab07b2bc/policyengine-6.14.0.tar.gz", hash = "sha256:aebc3668d72c6f151803446928530f104e36e5b83b7f73c272aeb08397d300d9", size = 1607946, upload-time = "2026-11-05T05:24:53.677Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/3a/8e/cdae9cd07eea17b33ee682249c758326fdfb4f9c837bec6832484c7428c2/policyengine-6.12.0-py3-none-any.whl", hash = "sha256:de08e9eb4e2de9e12010d1bff0a224e10658cce0c7dafee4783eec46978edb59", size = 557145, upload-time = "2026-11-02T14:23:37.915Z" }, + { url = "https://files.pythonhosted.org/packages/88/fb/5a9f0dd968032ad6c89e0495b4627a24daff45d8459d0f38d4a6cdd1e5ca/policyengine-6.14.0-py3-none-any.whl", hash = "sha256:61b14267bbe772c35b348b0731a9cdb21de863511d52264a264a5cf74129b282", size = 558163, upload-time = "2026-11-05T05:24:51.946Z" }, ] [package.optional-dependencies] @@ -1951,7 +1951,7 @@ requires-dist = [ { name = "openapi-python-client", marker = "extra == 'build'", specifier = ">=0.21.6" }, { name = "opentelemetry-instrumentation-fastapi", specifier = ">=0.65b0,<0.66" }, { name = "opentelemetry-instrumentation-sqlalchemy", specifier = ">=0.65b0,<0.66" }, - { name = "policyengine", extras = ["models"], specifier = "==6.12.0" }, + { name = "policyengine", extras = ["models"], specifier = "==6.14.0" }, { name = "policyengine-fastapi", editable = "../../libs/policyengine-fastapi" }, { name = "policyengine-observability", extras = ["fastapi", "google", "otlp-grpc"], specifier = ">=3.0.4,<4" }, { name = "policyengine-simulation-contract", extras = ["modal"], editable = "../../libs/policyengine-simulation-contract" }, @@ -1973,7 +1973,7 @@ modal-simulation-image = [ { name = "fastapi", specifier = ">=0.115.0" }, { name = "google-cloud-storage", specifier = ">=2" }, { name = "huggingface-hub", specifier = ">=1,<2" }, - { name = "policyengine", extras = ["models"], specifier = "==6.12.0" }, + { name = "policyengine", extras = ["models"], specifier = "==6.14.0" }, { name = "policyengine-observability", extras = ["fastapi", "google", "otlp-grpc"], specifier = ">=3.0.4,<4" }, { name = "psycopg", extras = ["binary"], specifier = ">=3.2,<4" }, { name = "pyarrow", specifier = ">=20,<24" },