From ed14c42feb7637bc0217375c1274da8ecb650b22 Mon Sep 17 00:00:00 2001 From: Laurent Courty Date: Wed, 28 Oct 2026 10:34:36 -0600 Subject: [PATCH 1/3] Use RegionManager --- src/itzi/grass/interface.py | 139 +++++++++++++++++++---------- src/itzi/grass/utils.py | 43 ++++++--- tests/grass/test_ensemble.py | 6 +- tests/grass/test_mask.py | 49 +++++++++-- tests/grass/test_region.py | 166 +++++++++++++++++++++++++++++++++++ 5 files changed, 338 insertions(+), 65 deletions(-) create mode 100644 tests/grass/test_region.py diff --git a/src/itzi/grass/interface.py b/src/itzi/grass/interface.py index ae5b0f9f..e4e47a0f 100644 --- a/src/itzi/grass/interface.py +++ b/src/itzi/grass/interface.py @@ -12,19 +12,23 @@ GNU General Public License for more details. """ +from __future__ import annotations + import copy import os -from collections.abc import Mapping +from collections.abc import Iterator, Mapping +from contextlib import ExitStack, contextmanager +from ctypes import byref from datetime import datetime, timedelta from pathlib import Path -from types import MappingProxyType -from typing import ClassVar, NamedTuple +from types import MappingProxyType, TracebackType +from typing import ClassVar, NamedTuple, Self import grass.pygrass.utils as gutils import grass.script as gscript import grass.temporal as tgis import numpy as np -from grass.lib import raster as libraster +from grass.lib import gis as libgis, raster as libraster from grass.pygrass import raster from grass.pygrass.gis.region import Region from grass.pygrass.vector import VectorTopo @@ -51,6 +55,7 @@ resolve_effective_mask, set_null, split_identifier, + temporary_region, ) # color rules @@ -128,7 +133,6 @@ def __init__( assert isinstance(end_time, datetime), "end_time not a datetime object!" assert start_time <= end_time, "start_time > end_time!" - self.region_id = region_id self.start_time = start_time self.end_time = end_time self.dtype = dtype @@ -136,46 +140,88 @@ def __init__( # LatLon is not supported if gscript.locn_is_latlong(): msgr.fatal("latlong location is not supported. Please use a projected location") - # Set region - if self.region_id: - gscript.use_temp_region() - gscript.run_command("g.region", region=region_id) - self.region = Region() - # Reload the region without remapping GRASS's retained mask descriptor. + self._contexts = ExitStack() + try: + self._contexts.enter_context(temporary_region(region_id)) + self._contexts.callback(self._refresh_mask) + # Initialize caller raster windows before Region reads the simulation region. + libraster.Rast__init_window() + self.region = Region() + self.xr = self.region.cols + self.yr = self.region.rows + if self.xr < 3 or self.yr < 3: + msgr.fatal("GRASS Region should be at least 3 cells by 3 cells") + self.mask_source = resolve_effective_mask(raster_mask_id).source + self.overwrite = gscript.overwrite() + self._domain_data = read_domain(None) + self._contexts.enter_context( + gscript.MaskManager(mask_name=self.mask_source, remove=False) + ) + except BaseException: + self.cleanup() + raise + + @contextmanager + def _raster_window(self) -> Iterator[None]: + """Use the simulation window only while opening and closing Itzi-owned rasters.""" + with temporary_region(None), ExitStack() as contexts: + contexts.callback(self._refresh_mask) + previous_input = libgis.Cell_head() + previous_output = libgis.Cell_head() + libraster.Rast_get_input_window(byref(previous_input)) + libraster.Rast_get_output_window(byref(previous_output)) + contexts.callback(self._restore_raster_windows, previous_input, previous_output) + libgis.G_set_window(self.region.byref()) + libraster.Rast_unset_window() + libraster.Rast__init_window() + yield + + @staticmethod + def _refresh_mask() -> None: + # Mask reopening checks the GIS projection but reads in the raster input window. + with temporary_region(None): + raster_input = libgis.Cell_head() + libraster.Rast_get_input_window(byref(raster_input)) + libgis.G_set_window(byref(raster_input)) + libraster.Rast__check_for_auto_masking() + + @staticmethod + def _restore_raster_windows( + previous_input: libgis.Cell_head, previous_output: libgis.Cell_head + ) -> None: + current_input = libgis.Cell_head() + current_output = libgis.Cell_head() + libraster.Rast_get_input_window(byref(current_input)) + libraster.Rast_get_output_window(byref(current_output)) + if bytes(current_input) == bytes(previous_input) and bytes(current_output) == bytes( + previous_output + ): + return + # Caller rasters retain their mappings and are never read in the temporary window. + libgis.G_set_window(byref(previous_input)) libraster.Rast_unset_window() - libraster.Rast_get_window(self.region.byref()) - self.xr = self.region.cols - self.yr = self.region.rows - # Check if region is at least 3x3 - if self.xr < 3 or self.yr < 3: - msgr.fatal("GRASS Region should be at least 3 cells by 3 cells") - self.mask_source = resolve_effective_mask(raster_mask_id).source - self.overwrite = gscript.overwrite() - self._domain_data = read_domain(None) - mask_manager = gscript.MaskManager(mask_name=self.mask_source, remove=False) - mask_manager.activate() - self._mask_manager: gscript.MaskManager | None = mask_manager - - def __enter__(self): + libraster.Rast_get_input_window(byref(current_input)) + if bytes(previous_input) != bytes(previous_output): + libraster.Rast_set_output_window(byref(previous_output)) + + def __enter__(self) -> Self: return self - def __exit__(self, exc_type, exc_value, traceback): + def __exit__( + self, + exc_type: type[BaseException] | None, + exc_value: BaseException | None, + traceback: TracebackType | None, + ) -> None: self.cleanup() - def __del__(self): + def __del__(self) -> None: self.cleanup() def cleanup(self) -> None: - """Restore the previous mask and remove the temporary region.""" - if manager := getattr(self, "_mask_manager", None): - manager.deactivate() - self._mask_manager = None - # Refresh only the mask, preserving unrelated open raster handles. - libraster.Rast__check_for_auto_masking() - if getattr(self, "region_id", None): - msgr.debug("Remove temp region...") - gscript.del_temp_region() - self.region_id = None + """Restore the caller's mask, temporary region, and GIS working window.""" + if contexts := getattr(self, "_contexts", None): + contexts.close() def grass_dtype(self, dtype: str) -> RasterMapType: if dtype in self.dtype_conv["DCELL"]: @@ -344,7 +390,7 @@ def read_raster_map(self, raster_id: str) -> np.ndarray: if not mapset: msgr.fatal(f"Raster <{raster_id}> not found") - with raster.RasterRow(name, mapset=mapset, mode="r") as source: + with self._raster_window(), raster.RasterRow(name, mapset=mapset, mode="r") as source: array = np.empty((self.yr, self.xr), dtype=self.dtype) buffer = raster.Buffer((self.xr,), mtype=source.mtype) for row_index in range(self.yr): @@ -354,13 +400,16 @@ def read_raster_map(self, raster_id: str) -> np.ndarray: def write_raster_map(self, arr: np.ndarray, raster_id: str, mkey: str, hmin: float) -> None: """Take a numpy array and write it to GRASS DB""" map_type = self.grass_dtype(str(arr.dtype)) - with raster.RasterRow( - raster_id, - mapset=get_current_mapset(), - mode="w", - mtype=map_type, - overwrite=self.overwrite, - ) as newraster: + with ( + self._raster_window(), + raster.RasterRow( + raster_id, + mapset=get_current_mapset(), + mode="w", + mtype=map_type, + overwrite=self.overwrite, + ) as newraster, + ): newrow = raster.Buffer((arr.shape[1],), mtype=map_type) for row in arr: newrow[:] = row[:] diff --git a/src/itzi/grass/utils.py b/src/itzi/grass/utils.py index 82c5b81f..42d1bccb 100644 --- a/src/itzi/grass/utils.py +++ b/src/itzi/grass/utils.py @@ -14,8 +14,10 @@ from __future__ import annotations -from collections.abc import Mapping -from ctypes import create_string_buffer +import os +from collections.abc import Iterator, Mapping +from contextlib import contextmanager +from ctypes import byref, create_string_buffer from datetime import datetime from pathlib import Path from typing import TYPE_CHECKING, Literal, NamedTuple @@ -152,15 +154,39 @@ def active_grass_params(requested: GrassParams) -> GrassParams: ) +@contextmanager +def temporary_region(region_id: str | None) -> Iterator[None]: + """Select a temporary region and restore the caller's GIS working window.""" + previous_window = libgis.Cell_head() + libgis.G_get_set_window(byref(previous_window)) + previous_grass_region = os.environ.get("GRASS_REGION") + manager = gscript.RegionManager(region=region_id) if region_id is not None else None + try: + if manager is not None: + # GRASS_REGION takes precedence over the manager's WIND_OVERRIDE. + os.environ.pop("GRASS_REGION", None) + manager.activate() + yield + finally: + try: + if manager is not None: + try: + manager.deactivate() + finally: + # GRASS 8.5 leaves its saved region behind for nested overrides. + libgis.G_remove("windows", manager.region_name) + finally: + if manager is not None and previous_grass_region is not None: + os.environ["GRASS_REGION"] = previous_grass_region + libgis.G_unset_window() + libgis.G_set_window(byref(previous_window)) + + def read_domain(region_id: str | None) -> DomainData: """Read the selected computational region without leaving it changed.""" if gscript.locn_is_latlong(): raise EnsembleError("latlong locations are not supported") - using_temp_region = region_id is not None - if using_temp_region: - gscript.use_temp_region() - gscript.run_command("g.region", region=region_id) - try: + with temporary_region(region_id): region = Region() if region.cols < 3 or region.rows < 3: raise EnsembleError("GRASS Region should be at least 3 cells by 3 cells") @@ -173,9 +199,6 @@ def read_domain(region_id: str | None) -> DomainData: cols=region.cols, crs_wkt=get_crs_wkt(), ) - finally: - if using_temp_region: - gscript.del_temp_region() def split_identifier(identifier: str) -> tuple[str, str]: diff --git a/tests/grass/test_ensemble.py b/tests/grass/test_ensemble.py index faa702b9..8ebb4d6b 100644 --- a/tests/grass/test_ensemble.py +++ b/tests/grass/test_ensemble.py @@ -260,12 +260,8 @@ def test_worker_rejects_changed_mask_region_and_inputs_after_resolution(tmp_path def test_dry_run_rejects_raster_outside_domain_without_writing(tmp_path: Path) -> None: from grass.pygrass import utils as gutils - gscript.use_temp_region() - try: - gscript.run_command("g.region", n=150, s=100, e=150, w=100, res=10) + with gscript.RegionManager(n=150, s=100, e=150, w=100, res=10): gscript.mapcalc("stage_outside=1") - finally: - gscript.del_temp_region() path = write_study( tmp_path, input={"friction": "stage_outside"}, diff --git a/tests/grass/test_mask.py b/tests/grass/test_mask.py index a001a7ab..45ab196a 100644 --- a/tests/grass/test_mask.py +++ b/tests/grass/test_mask.py @@ -161,12 +161,16 @@ def test_runner_restores_mask_when_lifecycle_fails( test_data_path: str, phase: str, ) -> None: + from grass.lib import raster as libraster + from itzi.configreader import ConfigReader from itzi.grass.session import GrassParams from itzi.simulation_runner import SimulationRunner gscript.mapcalc("mask_values=if(col()==1, 0, 1)") monkeypatch.setenv("GRASS_MASK", "previous_mask") + previous_region = gscript.region() + previous_shape = (libraster.Rast_window_rows(), libraster.Rast_window_cols()) reader = ConfigReader(str(Path(test_data_path) / "5by5" / "5by5.ini")) config = reader.sim_config.model_copy(update={"output_map_names": {}}) if phase == "initialize": @@ -176,7 +180,7 @@ def test_runner_restores_mask_when_lifecycle_fails( "input_map_names": dict(config.input_map_names) | {"ground_elevation": "bad_dem"} } ) - runner = SimulationRunner(config, GrassParams(mask="mask_values")) + runner = SimulationRunner(config, GrassParams(region="3by5", mask="mask_values")) if phase == "initialize": with pytest.raises(RuntimeError, match="only NULL/NaN cells"): runner.initialize() @@ -191,19 +195,26 @@ def fail_finalize(self: Simulation) -> None: with pytest.raises(RuntimeError, match="finalization failed"): runner.finalize() assert os.environ["GRASS_MASK"] == "previous_mask" + assert gscript.region() == previous_region + assert "WIND_OVERRIDE" not in os.environ + assert (libraster.Rast_window_rows(), libraster.Rast_window_cols()) == previous_shape @pytest.mark.usefixtures("grass_5by5") @pytest.mark.parametrize("phase", ["finalize", "delete"]) @pytest.mark.parametrize("explicit_mask", [False, True]) @pytest.mark.parametrize("previous_mask", [False, True]) +@pytest.mark.parametrize("region", [None, "3by3"]) def test_runner_cleanup_preserves_open_raster_and_restores_reads( monkeypatch: pytest.MonkeyPatch, test_data_path: str, phase: str, explicit_mask: bool, previous_mask: bool, + region: str | None, ) -> None: + gscript.mapcalc("gradient=float(row()*10+col())") + gscript.run_command("g.region", n=40, s=10, e=40, w=10, res=10, save="3by3") gscript.mapcalc("previous_mask=if(col()==1, 0, 1)") if previous_mask: monkeypatch.setenv("GRASS_MASK", "previous_mask@5by5") @@ -221,6 +232,8 @@ def test_runner_cleanup_preserves_open_raster_and_restores_reads( from pathlib import Path import numpy as np +import grass.script as gscript +from grass.lib import raster as libraster from grass.pygrass.raster import RasterRow from itzi.configreader import ConfigReader @@ -232,11 +245,34 @@ def test_runner_cleanup_preserves_open_raster_and_restores_reads( config = reader.sim_config.model_copy(update={"output_map_names": {}}) original_mask = os.environ.get("GRASS_MASK") mask = "5by3" if sys.argv[3] == "True" else None - runner = SimulationRunner(config, GrassParams(mask=mask)).initialize() - with RasterRow("n", mapset="5by5", mode="r") as source: + region = sys.argv[4] or None + original_region = gscript.region() + runner = SimulationRunner(config, GrassParams(region=region, mask=mask)).initialize() + with RasterRow("gradient", mapset="5by5", mode="r") as source: + original_row = np.asarray(source.get_row(0)).copy() + assert original_row.shape == (5,) + np.testing.assert_allclose( + original_row[np.isfinite(original_row)], + np.arange(11, 16)[np.isfinite(original_row)], + ) np.testing.assert_array_equal( - np.isnan(np.asarray(source.get_row(0))), runner.g_interface.get_npmask()[0] + np.isnan(original_row)[1:4] if region else np.isnan(original_row), + runner.g_interface.get_npmask()[0], ) + if region: + expected = np.arange(20, 50, 10)[:, None] + np.arange(2, 5) + np.testing.assert_allclose( + runner.g_interface.read_raster_map("gradient@5by5"), expected + ) + runner.g_interface.write_raster_map( + expected.astype(np.float32), "roundtrip@5by5", "friction", 0 + ) + info = gscript.raster_info("roundtrip@5by5") + assert (int(info["rows"]), int(info["cols"])) == (3, 3) + np.testing.assert_allclose( + runner.g_interface.read_raster_map("roundtrip@5by5"), expected + ) + np.testing.assert_array_equal(np.asarray(source.get_row(0)), original_row) if sys.argv[2] == "finalize": runner.finalize() else: @@ -245,7 +281,9 @@ def test_runner_cleanup_preserves_open_raster_and_restores_reads( expected = [original_mask is not None, False, False, False, False] restored_row = np.asarray(source.get_row(0)) np.testing.assert_array_equal(np.isnan(restored_row), expected) - np.testing.assert_allclose(restored_row[1:], 0.05) + np.testing.assert_allclose(restored_row[1:], np.arange(12, 16)) + assert gscript.region() == original_region + assert (libraster.Rast_window_rows(), libraster.Rast_window_cols()) == (5, 5) next_runner = SimulationRunner(config, GrassParams(region="3by5", mask=mask)).initialize() assert next_runner.g_interface.read_raster_map("n@5by5").shape == (3, 5) @@ -255,6 +293,7 @@ def test_runner_cleanup_preserves_open_raster_and_restores_reads( test_data_path, phase, str(explicit_mask), + region or "", ], env=os.environ | {"PYTHONPATH": os.pathsep.join(sys.path)}, capture_output=True, diff --git a/tests/grass/test_region.py b/tests/grass/test_region.py new file mode 100644 index 00000000..29b1e1a6 --- /dev/null +++ b/tests/grass/test_region.py @@ -0,0 +1,166 @@ +"""Temporary regions must agree across GRASS commands and in-process raster access.""" + +from __future__ import annotations + +import copy +import os +from ctypes import byref +from datetime import datetime, timedelta + +import grass.script as gscript +import numpy as np +import pytest + + +def working_windows() -> tuple[bytes, bytes, bytes]: + from grass.lib import gis as libgis, raster as libraster + + gis = libgis.Cell_head() + raster_input = libgis.Cell_head() + raster_output = libgis.Cell_head() + libgis.G_get_set_window(byref(gis)) + libraster.Rast_get_input_window(byref(raster_input)) + libraster.Rast_get_output_window(byref(raster_output)) + return bytes(gis), bytes(raster_input), bytes(raster_output) + + +@pytest.mark.usefixtures("grass_5by5") +@pytest.mark.parametrize("inherited", ["none", "WIND_OVERRIDE", "GRASS_REGION", "both"]) +def test_successive_regions_restore_overrides_windows_and_saved_regions( + monkeypatch: pytest.MonkeyPatch, inherited: str +) -> None: + from grass.lib import raster as libraster + from grass.pygrass.gis.region import Region + + from itzi.grass.interface import GrassInterface + + gscript.run_command("g.region", n=40, s=0, e=50, w=0, res=10, save="outer") + if inherited in ("WIND_OVERRIDE", "both"): + monkeypatch.setenv("WIND_OVERRIDE", "outer") + if inherited in ("GRASS_REGION", "both"): + monkeypatch.setenv("GRASS_REGION", gscript.region_env(region="outer")) + Region() + libraster.Rast_unset_window() + previous_windows = working_windows() + previous_region = gscript.region() + previous_overrides = (os.environ.get("WIND_OVERRIDE"), os.environ.get("GRASS_REGION")) + previous_saved_regions = gscript.list_strings("region") + + start = datetime(2020, 1, 1) + for index, region_id in enumerate(("3by5", None, "3by5")): + rows = 3 if region_id else int(previous_region["rows"]) + with GrassInterface( + start, start + timedelta(seconds=1), np.float32, region_id, None + ) as interface: + assert (interface.yr, interface.xr) == (rows, 5) + assert (interface.get_domain_data().rows, interface.get_domain_data().cols) == ( + rows, + 5, + ) + assert (gscript.region()["rows"], gscript.region()["cols"]) == (rows, 5) + values = interface.read_raster_map("n@5by5") + assert values.shape == (rows, 5) + np.testing.assert_allclose(values, 0.05) + output = f"region_output_{index}@5by5" + interface.write_raster_map(values, output, "friction", 0) + info = gscript.raster_info(output) + assert (int(info["rows"]), int(info["cols"])) == (rows, 5) + interface.cleanup() + assert gscript.region() == previous_region + assert working_windows() == previous_windows + assert ( + os.environ.get("WIND_OVERRIDE"), + os.environ.get("GRASS_REGION"), + ) == previous_overrides + assert gscript.list_strings("region") == previous_saved_regions + + +@pytest.mark.usefixtures("grass_5by5") +@pytest.mark.parametrize("failure", ["missing_region", "small_region", "missing_mask"]) +def test_constructor_failure_restores_region_immediately( + monkeypatch: pytest.MonkeyPatch, failure: str +) -> None: + from grass.exceptions import CalledModuleError + from grass.lib import raster as libraster + from grass.pygrass.gis.region import Region + + from itzi.ensemble.models import EnsembleError + from itzi.grass.interface import GrassInterface + from itzi.messenger import FatalError + + gscript.run_command("g.region", n=20, s=0, e=20, w=0, res=10, save="small_region") + monkeypatch.setenv("WIND_OVERRIDE", "3by5") + Region() + libraster.Rast_unset_window() + previous_windows = working_windows() + previous_saved_regions = gscript.list_strings("region") + start = datetime(2020, 1, 1) + region = failure if failure != "missing_mask" else "3by5" + mask = "missing_mask" if failure == "missing_mask" else None + with pytest.raises((CalledModuleError, FatalError, EnsembleError)): + GrassInterface(start, start + timedelta(seconds=1), np.float32, region, mask) + assert os.environ["WIND_OVERRIDE"] == "3by5" + assert working_windows() == previous_windows + assert gscript.list_strings("region") == previous_saved_regions + + +@pytest.mark.usefixtures("grass_5by5") +def test_custom_gis_and_split_raster_windows_are_restored() -> None: + from grass.lib import raster as libraster + from grass.pygrass.gis.region import Region + + from itzi.grass.interface import GrassInterface + + configured = Region() + libraster.Rast_unset_window() + gis = copy.deepcopy(configured) + gis.north = 40 + gis.adjust() + raster_input = copy.deepcopy(configured) + raster_input.north = 30 + raster_input.adjust() + raster_output = copy.deepcopy(configured) + raster_output.east = 40 + raster_output.adjust() + try: + raster_input.set_current() + libraster.Rast_unset_window() + libraster.Rast_get_input_window(raster_input.byref()) + libraster.Rast_set_output_window(raster_output.byref()) + gis.set_current() + previous_windows = working_windows() + start = datetime(2020, 1, 1) + with GrassInterface( + start, start + timedelta(seconds=1), np.float32, "3by5", None + ) as interface: + assert interface.read_raster_map("n@5by5").shape == (3, 5) + assert working_windows() == previous_windows + finally: + configured.set_current() + libraster.Rast_unset_window() + + +@pytest.mark.usefixtures("grass_5by5") +def test_domain_lookup_preserves_working_windows_and_open_raster() -> None: + from grass.lib import raster as libraster + from grass.pygrass.gis.region import Region + from grass.pygrass.raster import RasterRow + + from itzi.grass.utils import read_domain + + configured = Region() + libraster.Rast_unset_window() + libraster.Rast_get_input_window(configured.byref()) + gis = copy.deepcopy(configured) + gis.north = 40 + gis.adjust() + gis.set_current() + try: + with RasterRow("n", mapset="5by5", mode="r") as source: + previous_windows = working_windows() + domain = read_domain("3by5") + assert (domain.rows, domain.cols) == (3, 5) + assert working_windows() == previous_windows + np.testing.assert_allclose(np.asarray(source.get_row(0)), 0.05) + finally: + configured.set_current() From 448eeac6326600dafc88c853164219f73283f5b4 Mon Sep 17 00:00:00 2001 From: Laurent Courty Date: Wed, 28 Oct 2026 11:54:20 -0600 Subject: [PATCH 2/3] GrassSessionManager owns the whole grass context: mapset, region, and mask --- docs/conf_file.rst | 13 +- docs/yaml_file.rst | 10 +- src/itzi/bmi_itzi.py | 37 +-- src/itzi/ensemble/resolution.py | 19 +- src/itzi/ensemble/run_validation.py | 27 +- src/itzi/grass/interface.py | 123 ++------- src/itzi/grass/session.py | 346 ++++++++++++++++++------- src/itzi/grass/utils.py | 76 ++---- src/itzi/itzi.py | 27 +- src/itzi/simulation_runner.py | 97 +++---- tests/cli/test_cli.py | 19 +- tests/grass/conftest.py | 25 +- tests/grass/test_bmi.py | 48 ++++ tests/grass/test_ensemble.py | 73 +++++- tests/grass/test_grass_input.py | 26 +- tests/grass/test_itzi.py | 64 +++-- tests/grass/test_mask.py | 127 ++++----- tests/grass/test_region.py | 108 ++++---- tests/grass/test_session_context.py | 163 ++++++++++++ tests/grass/test_temporal_overwrite.py | 19 +- tests/grass/test_tutorial.py | 51 ++-- tests/test_grass_session.py | 88 ++++++- 22 files changed, 971 insertions(+), 615 deletions(-) create mode 100644 tests/grass/test_session_context.py diff --git a/docs/conf_file.rst b/docs/conf_file.rst index 4454e241..4a3ec55e 100644 --- a/docs/conf_file.rst +++ b/docs/conf_file.rst @@ -643,5 +643,14 @@ If Itzï is run from within the GRASS shell, this section is not necessary. With GNU/Linux, *grass_bin* could simply be ``grass``. -The *region* and *mask* parameters are optionals and are applied only during the simulation. -After the simulation, those parameters are returned to the previous *region* and *mask* setting. +The *region* and *mask* parameters are optional and apply for the whole simulation, +including initialization and output finalization. +The selected region is applied privately, so another process changing the mapset's default region does not +change a running job. +Raster and mask contents are not protected against external edits. +After itzi ends, whether with success or failure, the original region, mask, and session settings are restored +without rewriting saved regions, the shared default region, or mask rasters. + +An explicit *grassdata*, *location*, and *mapset* must be supplied together and select an existing, accessible mapset. +Omitted region and mask settings inherit the active settings when the target is unchanged, +and use the target's defaults when it differs. diff --git a/docs/yaml_file.rst b/docs/yaml_file.rst index c668a0be..54f4f106 100644 --- a/docs/yaml_file.rst +++ b/docs/yaml_file.rst @@ -103,7 +103,15 @@ An explicit session requires all three of the following options: session (defaults to ``grass``); it cannot be used without those three options. ``region`` optionally names a saved region; ``mask`` optionally names a raster mask. Both can also be used to override the active region or mask for the run. -When a session is already active, an explicit database, project, and mapset must match it. +When a session is already active, a different explicit context is selected temporarily using a private session configuration. +The context must already exist and be accessible. +Omitted region and mask settings inherit the active settings for an unchanged context; +for a different context, they use that context's defaults. +The selected region and mask remain active through initialization, execution, and finalization. +A private region snapshot isolates the job from changes to the shared default region, +but raster and mask contents can still be edited externally. +Context exit restores the original settings on success or failure without rewriting the shared default region, +saved regions, or user mask rasters. All members of an ensemble share this GRASS context and domain. Only a projected location / project is accepted, not latitude/longitude. The region must have at least 3 rows and 3 columns, and a mask cannot exclude the entire simulation domain. diff --git a/src/itzi/bmi_itzi.py b/src/itzi/bmi_itzi.py index 2059972c..9e55cefe 100644 --- a/src/itzi/bmi_itzi.py +++ b/src/itzi/bmi_itzi.py @@ -64,14 +64,14 @@ class BmiItzi(Bmi): {arr_def.csdms_name: arr_def.var_loc for arr_def in input_output_array_definition} ) - def __init__(self): + def __init__(self) -> None: """Create a BmiItzi model that is ready for initialization.""" self.itzi: SimulationRunner - self.grass_session_manager: GrassSessionManager + self.grass_session_manager: GrassSessionManager | None = None # Model control functions # - def initialize(self, config_file=None) -> None: + def initialize(self, config_file: str | None = None) -> None: """Initialize the Itzï model. Parameters @@ -83,13 +83,16 @@ def initialize(self, config_file=None) -> None: sim_params = conf_data.sim_config grass_params = conf_data.grass_params self.grass_session_manager = GrassSessionManager(grass_params) - self.grass_session_manager.open() - - self.itzi = SimulationRunner( - sim_params, - grass_params, - stats_file=conf_data.stats_file, - ).initialize() + try: + self.grass_session_manager.open() + self.itzi = SimulationRunner( + sim_params, + stats_file=conf_data.stats_file, + ).initialize() + except BaseException: + self.grass_session_manager.close() + self.grass_session_manager = None + raise def update(self): """Advance model by one time step.""" @@ -106,10 +109,14 @@ def update_until(self, then): then = timedelta(seconds=float(then)) self.itzi.sim.update_until(then) - def finalize(self): + def finalize(self) -> None: """Finalize model.""" - self.itzi.finalize() - self.grass_session_manager.close() + manager, self.grass_session_manager = self.grass_session_manager, None + if manager is not None: + try: + self.itzi.finalize() + finally: + manager.close() # Model information functions # @@ -381,7 +388,3 @@ def get_grid_y(self, grid, y): def get_grid_z(self, grid, z): raise NotImplementedError("get_grid_z") - - def __del__(self): - """Close the GRASS session even if finalize() is not called""" - self.grass_session_manager.close() diff --git a/src/itzi/ensemble/resolution.py b/src/itzi/ensemble/resolution.py index e48d5f5b..8cd63786 100644 --- a/src/itzi/ensemble/resolution.py +++ b/src/itzi/ensemble/resolution.py @@ -46,18 +46,15 @@ def resolve_ensemble( return () requested_params = _requested_grass_params(simulations[0]) - with GrassSessionManager(requested_params): + with GrassSessionManager(requested_params) as context: from itzi.grass.utils import ( active_grass_params, - ensure_min_version, read_domain, - resolve_effective_mask, ) - ensure_min_version() actual_params = active_grass_params(requested_params) - domain = read_domain(simulations[0].grass.region) - effective_mask = resolve_effective_mask(simulations[0].grass.mask) + domain = read_domain() + effective_mask = context.effective_mask input_cache: dict[str, tuple[str, Literal["raster", "strds"]]] = {} swmm_cache: dict[Path, str] = {} results: list[ResolvedSimulation | ValidationFailure] = [] @@ -431,17 +428,19 @@ def _reject_output_source_alias(map_id: str, sources: set[str]) -> None: raise EnsembleError(f"output {map_id!r} aliases an ensemble input or mask") -def verify_resolved_simulation(simulation: ResolvedSimulation) -> None: +def verify_resolved_simulation( + simulation: ResolvedSimulation, context: GrassSessionManager +) -> None: """Re-resolve worker facts and reject a changed GRASS environment.""" - from itzi.grass.utils import active_grass_params, read_domain, resolve_effective_mask + from itzi.grass.utils import active_grass_params, read_domain actual = active_grass_params(simulation.grass_params) if actual != simulation.grass_params: raise EnsembleError("GRASS context changed after resolution") - domain = read_domain(simulation.grass_params.region) + domain = read_domain() if domain != simulation.domain_data: raise EnsembleError("GRASS domain changed after resolution") - mask = resolve_effective_mask(simulation.grass_params.mask) + mask = context.effective_mask if mask != simulation.effective_mask: raise EnsembleError("effective GRASS mask changed after resolution") _verify_input_kinds(simulation.simulation_config.input_map_names, dict(simulation.input_kinds)) diff --git a/src/itzi/ensemble/run_validation.py b/src/itzi/ensemble/run_validation.py index 154c0f34..b29d2394 100644 --- a/src/itzi/ensemble/run_validation.py +++ b/src/itzi/ensemble/run_validation.py @@ -21,6 +21,7 @@ from itzi.ensemble.models import EnsembleError from itzi.ensemble.resolution import _last_record_index, verify_resolved_simulation from itzi.grass.names import derived_record_name +from itzi.grass.session import GrassSessionManager from itzi.manifest import validate_file_destination if TYPE_CHECKING: @@ -30,20 +31,16 @@ def validate_run_requirements(simulation: ResolvedSimulation) -> None: """Validate one resolved simulation without creating user artifacts.""" - from itzi.grass.interface import GrassInterface - from itzi.grass.utils import ensure_min_version - - ensure_min_version() - verify_resolved_simulation(simulation) - config = simulation.simulation_config - interface = GrassInterface( - start_time=config.start_time, - end_time=config.end_time, - dtype=np.float32, - region_id=simulation.grass_params.region, - raster_mask_id=simulation.grass_params.mask, - ) - try: + with GrassSessionManager(simulation.grass_params) as context: + from itzi.grass.interface import GrassInterface + + verify_resolved_simulation(simulation, context) + config = simulation.simulation_config + interface = GrassInterface( + start_time=config.start_time, + end_time=config.end_time, + dtype=np.float32, + ) _validate_inputs(simulation, interface) _validate_grass_outputs(simulation, interface) if simulation.artifacts.statistics_file is not None: @@ -52,8 +49,6 @@ def validate_run_requirements(simulation: ResolvedSimulation) -> None: overwrite=interface.overwrite, description="statistics file", ) - finally: - interface.cleanup() def _validate_inputs(simulation: ResolvedSimulation, interface: GrassInterface) -> None: diff --git a/src/itzi/grass/interface.py b/src/itzi/grass/interface.py index e4e47a0f..30012a2b 100644 --- a/src/itzi/grass/interface.py +++ b/src/itzi/grass/interface.py @@ -16,21 +16,17 @@ import copy import os -from collections.abc import Iterator, Mapping -from contextlib import ExitStack, contextmanager -from ctypes import byref +from collections.abc import Mapping from datetime import datetime, timedelta from pathlib import Path -from types import MappingProxyType, TracebackType -from typing import ClassVar, NamedTuple, Self +from types import MappingProxyType +from typing import ClassVar, NamedTuple import grass.pygrass.utils as gutils import grass.script as gscript import grass.temporal as tgis import numpy as np -from grass.lib import gis as libgis, raster as libraster from grass.pygrass import raster -from grass.pygrass.gis.region import Region from grass.pygrass.vector import VectorTopo from grass.pygrass.vector.geometry import Line, Point from grass.pygrass.vector.table import Link, Table @@ -55,7 +51,6 @@ resolve_effective_mask, set_null, split_identifier, - temporary_region, ) # color rules @@ -126,8 +121,6 @@ def __init__( start_time: datetime, end_time: datetime, dtype: type[np.generic] | np.dtype | str, - region_id: str | None, - raster_mask_id: str | None, ) -> None: assert isinstance(start_time, datetime), "start_time not a datetime object!" assert isinstance(end_time, datetime), "end_time not a datetime object!" @@ -137,91 +130,11 @@ def __init__( self.end_time = end_time self.dtype = dtype - # LatLon is not supported - if gscript.locn_is_latlong(): - msgr.fatal("latlong location is not supported. Please use a projected location") - self._contexts = ExitStack() - try: - self._contexts.enter_context(temporary_region(region_id)) - self._contexts.callback(self._refresh_mask) - # Initialize caller raster windows before Region reads the simulation region. - libraster.Rast__init_window() - self.region = Region() - self.xr = self.region.cols - self.yr = self.region.rows - if self.xr < 3 or self.yr < 3: - msgr.fatal("GRASS Region should be at least 3 cells by 3 cells") - self.mask_source = resolve_effective_mask(raster_mask_id).source - self.overwrite = gscript.overwrite() - self._domain_data = read_domain(None) - self._contexts.enter_context( - gscript.MaskManager(mask_name=self.mask_source, remove=False) - ) - except BaseException: - self.cleanup() - raise - - @contextmanager - def _raster_window(self) -> Iterator[None]: - """Use the simulation window only while opening and closing Itzi-owned rasters.""" - with temporary_region(None), ExitStack() as contexts: - contexts.callback(self._refresh_mask) - previous_input = libgis.Cell_head() - previous_output = libgis.Cell_head() - libraster.Rast_get_input_window(byref(previous_input)) - libraster.Rast_get_output_window(byref(previous_output)) - contexts.callback(self._restore_raster_windows, previous_input, previous_output) - libgis.G_set_window(self.region.byref()) - libraster.Rast_unset_window() - libraster.Rast__init_window() - yield - - @staticmethod - def _refresh_mask() -> None: - # Mask reopening checks the GIS projection but reads in the raster input window. - with temporary_region(None): - raster_input = libgis.Cell_head() - libraster.Rast_get_input_window(byref(raster_input)) - libgis.G_set_window(byref(raster_input)) - libraster.Rast__check_for_auto_masking() - - @staticmethod - def _restore_raster_windows( - previous_input: libgis.Cell_head, previous_output: libgis.Cell_head - ) -> None: - current_input = libgis.Cell_head() - current_output = libgis.Cell_head() - libraster.Rast_get_input_window(byref(current_input)) - libraster.Rast_get_output_window(byref(current_output)) - if bytes(current_input) == bytes(previous_input) and bytes(current_output) == bytes( - previous_output - ): - return - # Caller rasters retain their mappings and are never read in the temporary window. - libgis.G_set_window(byref(previous_input)) - libraster.Rast_unset_window() - libraster.Rast_get_input_window(byref(current_input)) - if bytes(previous_input) != bytes(previous_output): - libraster.Rast_set_output_window(byref(previous_output)) - - def __enter__(self) -> Self: - return self - - def __exit__( - self, - exc_type: type[BaseException] | None, - exc_value: BaseException | None, - traceback: TracebackType | None, - ) -> None: - self.cleanup() - - def __del__(self) -> None: - self.cleanup() - - def cleanup(self) -> None: - """Restore the caller's mask, temporary region, and GIS working window.""" - if contexts := getattr(self, "_contexts", None): - contexts.close() + self._domain_data = read_domain() + self.xr = self._domain_data.cols + self.yr = self._domain_data.rows + self.mask_source = resolve_effective_mask(None).source + self.overwrite = gscript.overwrite() def grass_dtype(self, dtype: str) -> RasterMapType: if dtype in self.dtype_conv["DCELL"]: @@ -385,12 +298,11 @@ def validate_output_stds_temporal_type( ) def read_raster_map(self, raster_id: str) -> np.ndarray: - """Read a raster using the mask selected for the simulation.""" name, mapset = split_identifier(raster_id) if not mapset: msgr.fatal(f"Raster <{raster_id}> not found") - with self._raster_window(), raster.RasterRow(name, mapset=mapset, mode="r") as source: + with raster.RasterRow(name, mapset=mapset, mode="r") as source: array = np.empty((self.yr, self.xr), dtype=self.dtype) buffer = raster.Buffer((self.xr,), mtype=source.mtype) for row_index in range(self.yr): @@ -400,16 +312,13 @@ def read_raster_map(self, raster_id: str) -> np.ndarray: def write_raster_map(self, arr: np.ndarray, raster_id: str, mkey: str, hmin: float) -> None: """Take a numpy array and write it to GRASS DB""" map_type = self.grass_dtype(str(arr.dtype)) - with ( - self._raster_window(), - raster.RasterRow( - raster_id, - mapset=get_current_mapset(), - mode="w", - mtype=map_type, - overwrite=self.overwrite, - ) as newraster, - ): + with raster.RasterRow( + raster_id, + mapset=get_current_mapset(), + mode="w", + mtype=map_type, + overwrite=self.overwrite, + ) as newraster: newrow = raster.Buffer((arr.shape[1],), mtype=map_type) for row in arr: newrow[:] = row[:] diff --git a/src/itzi/grass/session.py b/src/itzi/grass/session.py index 46bae5a6..190509f0 100644 --- a/src/itzi/grass/session.py +++ b/src/itzi/grass/session.py @@ -17,12 +17,26 @@ import os import subprocess import sys +from contextlib import ExitStack +from ctypes import byref from dataclasses import dataclass +from multiprocessing import get_start_method, set_start_method from pathlib import Path +from types import TracebackType +from typing import TYPE_CHECKING, Self import itzi.messenger as msgr _initialized_temporal_session: tuple[int, str, tuple[str, str, str]] | None = None +_synchronized_session: tuple[int, str, tuple[str, str, str]] | None = None + +if TYPE_CHECKING: + from grass.lib.gis import Cell_head + from grass.script.setup import SessionHandle + + from itzi.ensemble.models import EffectiveMask + +MIN_GRASS_VERSION = (8, 5) @dataclass(frozen=True) @@ -44,55 +58,158 @@ def __post_init__(self) -> None: raise ValueError("GRASS database, location, and mapset must be supplied together") +def ensure_min_version() -> None: + """Check compatibility before importing session-dependent GRASS libraries.""" + import grass.script as gscript + + full = gscript.parse_command("g.version", flags="g")["version"] + major, minor = (int(part) for part in full.split(".")[:2]) + if (major, minor) < MIN_GRASS_VERSION: + required = f"{MIN_GRASS_VERSION[0]}.{MIN_GRASS_VERSION[1]}" + raise RuntimeError( + f"itzi requires at least GRASS {required}, found version {major}.{minor}" + ) + + +def _restore_environment(previous: dict[str, str | None]) -> None: + for name, value in previous.items(): + if value is None: + os.environ.pop(name, None) + else: + os.environ[name] = value + + +def _validate_target(params: GrassParams) -> None: + assert params.grassdata and params.location and params.mapset + database = Path(params.grassdata).expanduser().resolve() + for path, permission in ( + (database, os.R_OK), + (database / params.location, os.R_OK), + (database / params.location / params.mapset, os.W_OK), + ): + if not path.is_dir() or not os.access(path, permission): + msgr.fatal(f"'{path}' does not exist or does not have adequate permissions") + + +def _finish_session(session: SessionHandle) -> None: + global _initialized_temporal_session, _synchronized_session + + try: + if _initialized_temporal_session is not None: + _stop_temporal() + finally: + _initialized_temporal_session = None + _synchronized_session = None + session.finish() + + +def _stop_temporal() -> None: + global _initialized_temporal_session + import grass.temporal as tgis + + try: + tgis.stop_subprocesses() + finally: + # Stopped RPC objects retain closed pipes and can race their checker threads. + tgis.core.c_library_interface = None + tgis.core.message_interface = None + _initialized_temporal_session = None + + +def _acquire_session(params: GrassParams, contexts: ExitStack) -> SessionHandle | None: + """Acquire a bare runtime, also used by multi-member validation workers.""" + try: + import grass.script as gscript + except ImportError: + python_path = subprocess.check_output( + [params.grass_bin or "grass", "--config", "python_path"], text=True + ).strip() + sys.path.append(python_path) + import grass.script as gscript + + if GrassSessionManager.current_context() is not None: + return None + if not params.grassdata: + msgr.fatal("No GRASS parameters to create a session.") + _validate_target(params) + previous = {name: os.environ.get(name) for name in ("GISRC", "GIS_LOCK")} + contexts.callback(_restore_environment, previous) + session = gscript.setup.init( + path=str(Path(params.grassdata).expanduser().resolve()), + location=params.location, + mapset=params.mapset, + grass_path=params.grass_bin or "grass", + ) + contexts.callback(_finish_session, session) + return session + + +def _refresh_session() -> None: + """Refresh C-library session paths and caches without rewriting the GISRC.""" + global _synchronized_session + from grass.lib import gis as libgis, raster as libraster + + libgis.G__read_gisrc_path() + names = [] + index = 0 + while name := libgis.G_get_env_name(index): + names.append(name) + index += 1 + for source in (libgis.G_VAR_MAPSET, libgis.G_VAR_GISRC): + for name in names: + libgis.G_setenv_nogisrc2(name, None, source) + libgis.G__read_gisrc_env() + libgis.G__read_mapset_env() + libgis.G_reset_mapsets() + libgis.G_unset_window() + libraster.Rast_unset_window() + context = GrassSessionManager.current_context() + assert context is not None + _synchronized_session = (os.getpid(), os.environ["GISRC"], context) + + +def _restore_windows(gis: Cell_head, window: Cell_head) -> None: + from grass.lib import gis as libgis, raster as libraster + + # The normal window API closes the retained automatic-mask descriptor first. + libgis.G_set_window(byref(window)) + try: + libraster.Rast_set_window(byref(window)) + finally: + libgis.G_set_window(byref(gis)) + + class GrassSessionManager: - """Manages GRASS session lifecycle.""" + """Own one fixed work region/mask and restore its caller's GRASS context. + + All Itzi raster handles must be closed before closing the manager. An attached + session remains available; only sessions acquired by this manager are finished. + """ - def __init__(self, grass_params: GrassParams): + def __init__(self, grass_params: GrassParams) -> None: self.grass_params = grass_params - self.grass_session = None + self.grass_session: SessionHandle | None = None + self.effective_mask: EffectiveMask + self._contexts: ExitStack | None = None @staticmethod def current_context() -> tuple[str, str, str] | None: """Return the active GRASS context, or ``None`` outside a session.""" - # Importing PyGRASS without a GISRC can terminate through GRASS' C API. - # The environment guard keeps configuration parsing and external launches - # independent from GRASS imports. if not os.environ.get("GISRC"): return None try: - from grass.pygrass.utils import getenv + import grass.script as gscript - values = (getenv("GISDBASE"), getenv("LOCATION_NAME"), getenv("MAPSET")) - except (ImportError, RuntimeError, SystemExit): + environment = gscript.gisenv() + values = tuple( + environment.get(name) for name in ("GISDBASE", "LOCATION_NAME", "MAPSET") + ) + except (ImportError, RuntimeError, SystemExit, OSError): return None if not all(values): return None return values - def _validate_requested_context(self) -> None: - """Ensure an explicit request matches the session selected for this process.""" - active = self.current_context() - if active is None: - msgr.fatal("Unable to determine the active GRASS context") - gisdb, location, mapset = ( - self.grass_params.grassdata, - self.grass_params.location, - self.grass_params.mapset, - ) - if gisdb is None: - return - requested_database = str(Path(gisdb).expanduser().resolve()) - active_database = str(Path(active[0]).expanduser().resolve()) - if (requested_database, location, mapset) != ( - active_database, - active[1], - active[2], - ): - msgr.fatal( - "Requested GRASS context does not match the active session " - f"({active_database}/{active[1]}/{active[2]})" - ) - @staticmethod def ensure_temporal_initialized() -> None: """Initialize the temporal framework once for the active GRASS session.""" @@ -107,77 +224,124 @@ def ensure_temporal_initialized() -> None: gscript.set_raise_on_error(True) session = (os.getpid(), os.environ["GISRC"], context) if _initialized_temporal_session != session: - tgis.init(raise_fatal_error=True) + _stop_temporal() + # Python's forkserver retains the first session's environment for later helpers. + method = get_start_method() + if method == "forkserver": + set_start_method("spawn", force=True) + try: + tgis.init(raise_fatal_error=True) + except BaseException: + _stop_temporal() + raise + finally: + if method == "forkserver": + set_start_method(method, force=True) _initialized_temporal_session = session tgis.set_raise_on_error(True) def open(self) -> None: - """Open a GRASS session if needed.""" - if self.current_context() is not None: - self._validate_requested_context() - self.ensure_temporal_initialized() + """Acquire a session, select the job context, and synchronize its windows.""" + if self._contexts is not None: return - - gisdb = self.grass_params.grassdata - location = self.grass_params.location - mapset = self.grass_params.mapset - # Check if mandatory GRASS parameters are present - if not gisdb or not location or not mapset: - msgr.fatal("No GRASS parameters to create a session.") - - # Check if the given parameters exist and can be accessed - error_msg = "'{}' does not exist or does not have adequate permissions" - if not os.access(gisdb, os.R_OK): - msgr.fatal(error_msg.format(gisdb)) - elif not os.access(os.path.join(gisdb, location), os.R_OK): - msgr.fatal(error_msg.format(location)) - elif not os.access(os.path.join(gisdb, location, mapset), os.W_OK): - msgr.fatal(error_msg.format(mapset)) - - # Set GRASS python path - if self.grass_params.grass_bin: - grassbin = self.grass_params.grass_bin - else: - grassbin = "grass" - grass_cmd = [grassbin, "--config", "python_path"] - grass_python_path = subprocess.check_output(grass_cmd, text=True).strip() - sys.path.append(grass_python_path) - # Now we can import grass modules - import grass.script as gscript - - # set up session - self.grass_session = gscript.setup.init( - path=gisdb, location=location, mapset=mapset, grass_path=grassbin - ) - if not os.environ.get("GISRC"): - session_env = getattr(self.grass_session, "env", None) - if session_env and session_env.get("GISRC"): - os.environ.update(session_env) + self._contexts = contexts = ExitStack() try: - self._validate_requested_context() + self.grass_session = _acquire_session(self.grass_params, contexts) + ensure_min_version() + self._select_context(contexts) self.ensure_temporal_initialized() - except Exception: + except BaseException: self.close() raise + def _select_context(self, contexts: ExitStack) -> None: + import grass.script as gscript + from grass.lib import gis as libgis, raster as libraster + + from itzi.grass.utils import read_domain, resolve_effective_mask + + active = self.current_context() + if active is None: + msgr.fatal("Unable to determine the active GRASS context") + identity = (os.getpid(), os.environ["GISRC"], active) + if _synchronized_session != identity: + _refresh_session() + gis = libgis.Cell_head() + window = libgis.Cell_head() + libgis.G_get_set_window(byref(gis)) + libraster.Rast_get_window(byref(window)) + params = self.grass_params + requested = ( + (str(Path(params.grassdata).expanduser().resolve()), params.location, params.mapset) + if params.grassdata + else active + ) + changed = requested != (str(Path(active[0]).resolve()), active[1], active[2]) + previous = { + name: os.environ.get(name) + for name in ("GISRC", "GRASS_REGION", "WIND_OVERRIDE", "GRASS_MASK") + } + if changed: + _validate_target(params) + target = gscript.setup.init( + path=requested[0], + location=params.location, + mapset=params.mapset, + grass_path=params.grass_bin or "grass", + env=os.environ.copy(), + ) + contexts.callback(target.finish) + contexts.callback(self.ensure_temporal_initialized) + contexts.callback(_restore_windows, gis, window) + if changed: + contexts.callback(_refresh_session) + contexts.callback(_restore_environment, previous) + if changed: + os.environ["GISRC"] = target.env["GISRC"] + for name in ("GRASS_REGION", "WIND_OVERRIDE", "GRASS_MASK"): + os.environ.pop(name, None) + _refresh_session() + self.effective_mask = resolve_effective_mask(params.mask) + # Explicit regions must not consult caller-only saved-region overrides. + if params.region is not None: + os.environ.pop("GRASS_REGION", None) + os.environ.pop("WIND_OVERRIDE", None) + contexts.enter_context(gscript.RegionManagerEnv(region=params.region)) + else: + # RegionManagerEnv() without arguments reads WIND, ignoring inherited overrides. + snapshot = gscript.region() + contexts.enter_context( + gscript.RegionManagerEnv( + **{key: snapshot[key] for key in ("n", "s", "e", "w", "rows", "cols")} + ) + ) + contexts.enter_context( + gscript.MaskManager(mask_name=self.effective_mask.source, remove=False) + ) + libgis.G_unset_window() + selected = libgis.Cell_head() + libgis.G_get_window(byref(selected)) + libgis.G_set_window(byref(selected)) + libraster.Rast_set_window(byref(selected)) + read_domain() + def close(self) -> None: - """Stop GRASS session.""" - global _initialized_temporal_session - if self.grass_session is not None: - try: - self.grass_session.finish() - except Exception as e: - print(f"Warning: Error cleaning up GRASS session: {e}") - _initialized_temporal_session = None - self.grass_session = None + """Restore process state and release owned resources, once.""" + contexts, self._contexts = self._contexts, None + try: + if contexts is not None: + contexts.close() + finally: + self.grass_session = None - def __enter__(self): + def __enter__(self) -> Self: self.open() return self - def __exit__(self, exc_type, exc_val, exc_tb): - self.close() - return False - - def __del__(self): + def __exit__( + self, + exc_type: type[BaseException] | None, + exc_val: BaseException | None, + exc_tb: TracebackType | None, + ) -> None: self.close() diff --git a/src/itzi/grass/utils.py b/src/itzi/grass/utils.py index 42d1bccb..6cdcd494 100644 --- a/src/itzi/grass/utils.py +++ b/src/itzi/grass/utils.py @@ -14,9 +14,7 @@ from __future__ import annotations -import os -from collections.abc import Iterator, Mapping -from contextlib import contextmanager +from collections.abc import Mapping from ctypes import byref, create_string_buffer from datetime import datetime from pathlib import Path @@ -28,7 +26,6 @@ import numpy as np from grass.lib import gis as libgis, raster as libraster from grass.pygrass.gis import Mapset -from grass.pygrass.gis.region import Region from itzi_core import INPUT_ARRAY_KEYS, DomainData import itzi.messenger as msgr @@ -38,9 +35,6 @@ if TYPE_CHECKING: from itzi.grass.interface import GrassInterface -MIN_GRASS_VERSION = (8, 5) - - type STDSType = Literal["strds", "stvds"] type RasterMapType = Literal["CELL", "FCELL", "DCELL"] @@ -132,16 +126,6 @@ def name_is_map(map_id: str) -> bool: return raster_exists(name, mapset) -def ensure_min_version() -> None: - full = gscript.parse_command("g.version", flags="g")["version"] - major, minor = (int(part) for part in full.split(".")[:2]) - if (major, minor) < MIN_GRASS_VERSION: - required = f"{MIN_GRASS_VERSION[0]}.{MIN_GRASS_VERSION[1]}" - raise RuntimeError( - f"itzi requires at least GRASS {required}, found version {major}.{minor}" - ) - - def active_grass_params(requested: GrassParams) -> GrassParams: """Return the active session context, retaining requested region/mask/bin.""" return GrassParams( @@ -154,51 +138,23 @@ def active_grass_params(requested: GrassParams) -> GrassParams: ) -@contextmanager -def temporary_region(region_id: str | None) -> Iterator[None]: - """Select a temporary region and restore the caller's GIS working window.""" - previous_window = libgis.Cell_head() - libgis.G_get_set_window(byref(previous_window)) - previous_grass_region = os.environ.get("GRASS_REGION") - manager = gscript.RegionManager(region=region_id) if region_id is not None else None - try: - if manager is not None: - # GRASS_REGION takes precedence over the manager's WIND_OVERRIDE. - os.environ.pop("GRASS_REGION", None) - manager.activate() - yield - finally: - try: - if manager is not None: - try: - manager.deactivate() - finally: - # GRASS 8.5 leaves its saved region behind for nested overrides. - libgis.G_remove("windows", manager.region_name) - finally: - if manager is not None and previous_grass_region is not None: - os.environ["GRASS_REGION"] = previous_grass_region - libgis.G_unset_window() - libgis.G_set_window(byref(previous_window)) - - -def read_domain(region_id: str | None) -> DomainData: - """Read the selected computational region without leaving it changed.""" +def read_domain() -> DomainData: + """Read the active work window without reloading the mapset's default.""" if gscript.locn_is_latlong(): raise EnsembleError("latlong locations are not supported") - with temporary_region(region_id): - region = Region() - if region.cols < 3 or region.rows < 3: - raise EnsembleError("GRASS Region should be at least 3 cells by 3 cells") - return DomainData( - north=region.north, - south=region.south, - east=region.east, - west=region.west, - rows=region.rows, - cols=region.cols, - crs_wkt=get_crs_wkt(), - ) + region = libgis.Cell_head() + libgis.G_get_set_window(byref(region)) + if region.cols < 3 or region.rows < 3: + raise EnsembleError("GRASS Region should be at least 3 cells by 3 cells") + return DomainData( + north=region.north, + south=region.south, + east=region.east, + west=region.west, + rows=region.rows, + cols=region.cols, + crs_wkt=get_crs_wkt(), + ) def split_identifier(identifier: str) -> tuple[str, str]: diff --git a/src/itzi/itzi.py b/src/itzi/itzi.py index 7f2228f3..c54fbbbe 100644 --- a/src/itzi/itzi.py +++ b/src/itzi/itzi.py @@ -31,6 +31,7 @@ from argparse import Namespace from collections.abc import Callable from concurrent.futures import ProcessPoolExecutor +from contextlib import ExitStack from datetime import timedelta from importlib.metadata import version from multiprocessing import Process, get_context @@ -51,7 +52,7 @@ validate_resolved_ensemble, verify_resolved_simulation, ) -from itzi.grass.session import GrassSessionManager +from itzi.grass.session import GrassSessionManager, _acquire_session from itzi.manifest import ( _create_manifest, _initial_member_states, @@ -94,11 +95,19 @@ def sim_runner_worker(conf_file: str, hotstart_file: str | None) -> None: with GrassSessionManager(grass_params): sim_runner = SimulationRunner( sim_params, - grass_params, hotstart_path=hotstart_file, stats_file=conf_data.stats_file, ) - sim_runner.initialize().run().finalize() + sim_runner.initialize() + try: + sim_runner.run() + except BaseException: + try: + sim_runner.finalize() + except Exception as finalize_error: + msgr.warning(f"Could not finalize partial outputs: {finalize_error}") + raise + sim_runner.finalize() except msgr.FatalError: raise SystemExit(1) from None except SystemExit as error: @@ -115,11 +124,10 @@ def resolved_sim_runner_worker(simulation: ResolvedSimulation) -> tuple[str, str msgr.raise_on_error = True msgr._itzi_logger.set_verbosity(msgr.verbosity()) try: - with GrassSessionManager(simulation.grass_params): - verify_resolved_simulation(simulation) + with GrassSessionManager(simulation.grass_params) as context: + verify_resolved_simulation(simulation, context) runner = SimulationRunner( simulation.simulation_config, - simulation.grass_params, stats_file=str(simulation.artifacts.statistics_file) if simulation.artifacts.statistics_file is not None else None, @@ -128,7 +136,7 @@ def resolved_sim_runner_worker(simulation: ResolvedSimulation) -> tuple[str, str runner.initialize() try: runner.run() - except Exception: + except BaseException: try: runner.finalize() except Exception as finalize_error: @@ -159,7 +167,10 @@ def run_validation_worker(simulations: tuple[ResolvedSimulation, ...]) -> dict[s msgr.raise_on_error = True msgr._itzi_logger.set_verbosity(msgr.verbosity()) failures: dict[str, str] = {} - with GrassSessionManager(simulations[0].grass_params): + if not simulations: + return failures + with ExitStack() as resources: + _acquire_session(simulations[0].grass_params, resources) from itzi.ensemble.run_validation import validate_run_requirements for simulation in simulations: diff --git a/src/itzi/simulation_runner.py b/src/itzi/simulation_runner.py index be87a991..0cc794a7 100644 --- a/src/itzi/simulation_runner.py +++ b/src/itzi/simulation_runner.py @@ -21,7 +21,7 @@ from itzi_core import SimulationBuilder, SimulationConfig import itzi.messenger as msgr -from itzi.grass.session import GrassParams +from itzi.grass.session import GrassSessionManager from itzi.providers.csv_output import ExclusiveCSVMassBalanceOutputProvider from itzi.providers.grass_input import GrassRasterInputProvider from itzi.providers.grass_output import GrassRasterOutputProvider, GrassVectorOutputProvider @@ -34,28 +34,26 @@ class SimulationRunner: """Provide the necessary tools to run one simulation. - Must be instantiated within an active GRASS session.""" + Construct, initialize, run, and finalize inside a configured GrassSessionManager.""" def __init__( self, sim_config: SimulationConfig, - grass_params: GrassParams, hotstart_path: str | None = None, stats_file: str | None = None, input_kinds: Mapping[str, Literal["raster", "strds"]] | None = None, ) -> None: + if GrassSessionManager.current_context() is None: + raise RuntimeError("SimulationRunner requires an active GRASS work context") from itzi.grass.utils import ( check_output_files, - ensure_min_version, get_current_mapset, qualify_output_id, resolve_input_identifier, ) - self.grass_params = grass_params self.hotstart_path = hotstart_path self.stats_file = stats_file - ensure_min_version() if input_kinds is None: resolved_inputs = { key: resolve_input_identifier(name) @@ -98,52 +96,46 @@ def initialize(self) -> Self: start_time=self.sim_config.start_time, end_time=self.sim_config.end_time, dtype=data_type, - region_id=self.grass_params.region, - raster_mask_id=self.grass_params.mask, ) - try: - msgr.verbose("Setting up GRASS simulation...") - - raster_input_provider = GrassRasterInputProvider( - grass_interface=self.g_interface, - input_map_names=self.sim_config.input_map_names, - default_start_time=self.sim_config.start_time, - default_end_time=self.sim_config.end_time, - input_kinds=self.input_kinds, - ) - raster_output_provider = GrassRasterOutputProvider( + msgr.verbose("Setting up GRASS simulation...") + + raster_input_provider = GrassRasterInputProvider( + grass_interface=self.g_interface, + input_map_names=self.sim_config.input_map_names, + default_start_time=self.sim_config.start_time, + default_end_time=self.sim_config.end_time, + input_kinds=self.input_kinds, + ) + raster_output_provider = GrassRasterOutputProvider( + grass_interface=self.g_interface, + out_map_names=self.sim_config.output_map_names, + hmin=self.sim_config.surface_flow_parameters.hmin, + temporal_type=self.sim_config.temporal_type, + ) + sim_builder = ( + SimulationBuilder(self.sim_config, self.g_interface.get_npmask(), data_type) + .with_input_provider(raster_input_provider) + .with_raster_output_provider(raster_output_provider) + ) + if self.sim_config.drainage_output is not None: + vector_output_provider = GrassVectorOutputProvider( grass_interface=self.g_interface, - out_map_names=self.sim_config.output_map_names, - hmin=self.sim_config.surface_flow_parameters.hmin, temporal_type=self.sim_config.temporal_type, + drainage_map_id=self.sim_config.drainage_output, ) - sim_builder = ( - SimulationBuilder(self.sim_config, self.g_interface.get_npmask(), data_type) - .with_input_provider(raster_input_provider) - .with_raster_output_provider(raster_output_provider) + sim_builder.with_vector_output_provider(vector_output_provider) + if self.stats_file: + stats_provider = ExclusiveCSVMassBalanceOutputProvider( + self.stats_file, overwrite=self.g_interface.overwrite ) - if self.sim_config.drainage_output is not None: - vector_output_provider = GrassVectorOutputProvider( - grass_interface=self.g_interface, - temporal_type=self.sim_config.temporal_type, - drainage_map_id=self.sim_config.drainage_output, - ) - sim_builder.with_vector_output_provider(vector_output_provider) - if self.stats_file: - stats_provider = ExclusiveCSVMassBalanceOutputProvider( - self.stats_file, overwrite=self.g_interface.overwrite - ) - sim_builder.with_mass_balance_output_provider(stats_provider) - if self.hotstart_path: - sim_builder.with_hotstart(self.hotstart_path) - self.sim: Simulation = sim_builder.build() - self.sim.initialize() - except BaseException: - self.g_interface.cleanup() - raise + sim_builder.with_mass_balance_output_provider(stats_provider) + if self.hotstart_path: + sim_builder.with_hotstart(self.hotstart_path) + self.sim: Simulation = sim_builder.build() + self.sim.initialize() return self - def run(self): + def run(self) -> Self: """Run a full simulation""" sim_start_time = datetime.now() last_progress_update = monotonic() - 0.5 @@ -163,15 +155,11 @@ def run(self): return self def finalize(self) -> Self: - """Tear down the simulation and return to previous state.""" - try: - self.sim.finalize() - finally: - if hasattr(self, "g_interface"): - self.g_interface.cleanup() + """Finalize model outputs inside the active work context.""" + self.sim.finalize() return self - def step(self): + def step(self) -> Self: """Do one simulation step.""" self.sim.update() return self @@ -179,8 +167,3 @@ def step(self): @property def origin(self): return (self.sim.domain_data.north, self.sim.domain_data.west) - - def __del__(self): - # Cleanup the grass interface object - if hasattr(self, "g_interface"): - self.g_interface.cleanup() diff --git a/tests/cli/test_cli.py b/tests/cli/test_cli.py index ca180db2..1dbef76a 100644 --- a/tests/cli/test_cli.py +++ b/tests/cli/test_cli.py @@ -298,22 +298,17 @@ def test_one_ensemble_runs_only_planned_members_and_counts_failures(tmp_path, mo def test_run_validation_worker_keeps_member_failures(monkeypatch): calls = [] - class FakeGrassSessionManager: - def __init__(self, params): - assert params == "shared" - - def __enter__(self): - calls.append("open") - - def __exit__(self, *_): - calls.append("close") + def acquire(params, stack): + assert params == "shared" + calls.append("open") + stack.callback(calls.append, "close") def validate_run_requirements(simulation): calls.append(simulation.simulation_id) if simulation.simulation_id == "bad": raise ValueError("invalid input") - monkeypatch.setattr("itzi.itzi.GrassSessionManager", FakeGrassSessionManager) + monkeypatch.setattr("itzi.itzi._acquire_session", acquire) fake_run_validation = ModuleType("itzi.ensemble.run_validation") setattr(fake_run_validation, "validate_run_requirements", validate_run_requirements) monkeypatch.setitem(sys.modules, "itzi.ensemble.run_validation", fake_run_validation) @@ -362,7 +357,7 @@ def finalize(self): monkeypatch.setattr("itzi.itzi.GrassSessionManager", FakeGrassSessionManager) monkeypatch.setattr("itzi.itzi.SimulationRunner", FakeSimulationRunner) - monkeypatch.setattr("itzi.itzi.verify_resolved_simulation", lambda _: calls.append("verify")) + monkeypatch.setattr("itzi.itzi.verify_resolved_simulation", lambda *_: calls.append("verify")) simulation = cast( ResolvedSimulation, SimpleNamespace( @@ -555,7 +550,7 @@ def finalize(self): sim_runner_worker("a.ini", "hotstart.zip") assert runner_arguments == { - "args": (sim_params, grass_params), + "args": (sim_params,), "kwargs": { "hotstart_path": "hotstart.zip", "stats_file": "statistics.csv", diff --git a/tests/grass/conftest.py b/tests/grass/conftest.py index c52676e2..ae70658d 100644 --- a/tests/grass/conftest.py +++ b/tests/grass/conftest.py @@ -182,10 +182,10 @@ def grass_5by5_template( gscript.mapcalc("infiltration_rate=2") gscript.mapcalc("loss_rate=1.5") gscript.mapcalc("inflow_rate=0.1") - with session.GrassSessionManager(session.GrassParams()): - mapset_path = Path(gscript.gisenv()["GISDBASE"]) / "xy" / "5by5" - template_path = tmp_path_factory.mktemp("grass_template") / "5by5" - shutil.copytree(mapset_path, template_path) + session.GrassSessionManager.ensure_temporal_initialized() + mapset_path = Path(gscript.gisenv()["GISDBASE"]) / "xy" / "5by5" + template_path = tmp_path_factory.mktemp("grass_template") / "5by5" + shutil.copytree(mapset_path, template_path) return mapset_path, template_path @@ -204,18 +204,19 @@ def grass_5by5(grass_5by5_template: tuple[Path, Path], tmp_path: Path) -> Iterat libgis.G_unset_window() libraster.Rast_unset_window() os.chdir(tmp_path) - with session.GrassSessionManager(session.GrassParams()): - yield + session.GrassSessionManager.ensure_temporal_initialized() + yield @pytest.fixture def grass_5by5_sim(grass_5by5: None, test_data_path: str) -> SimulationRunner: config_file = os.path.join(test_data_path, "5by5", "5by5.ini") conf_data = ConfigReader(config_file) - sim_runner = SimulationRunner( - conf_data.sim_config, - conf_data.grass_params, - ).initialize() - assert isinstance(sim_runner, SimulationRunner) - sim_runner.run().finalize() + with session.GrassSessionManager(conf_data.grass_params): + sim_runner = SimulationRunner(conf_data.sim_config).initialize() + assert isinstance(sim_runner, SimulationRunner) + try: + sim_runner.run() + finally: + sim_runner.finalize() return sim_runner diff --git a/tests/grass/test_bmi.py b/tests/grass/test_bmi.py index 91c53d90..f320ed6e 100644 --- a/tests/grass/test_bmi.py +++ b/tests/grass/test_bmi.py @@ -2,11 +2,16 @@ import os from collections.abc import Iterator +from pathlib import Path +from types import SimpleNamespace import numpy as np import pytest from itzi import BmiItzi +from itzi.configreader import ConfigReader +from itzi.grass.session import GrassParams, GrassSessionManager +from itzi.simulation_runner import SimulationRunner @pytest.fixture @@ -197,3 +202,46 @@ def test_get_grid_type(self, bmi_object): grid_id = bmi_object.get_var_grid("land_surface__elevation") grid_type = bmi_object.get_grid_type(grid_id) assert grid_type == "uniform_rectilinear" + + +@pytest.mark.usefixtures("grass_5by5") +@pytest.mark.parametrize("phase", ["__init__", "initialize", "finalize"]) +def test_bmi_failure_restores_attached_context_and_allows_reinitialization( + monkeypatch: pytest.MonkeyPatch, test_data_path: str, phase: str +) -> None: + import grass.script as gscript + from grass.lib import raster as libraster + + reader = ConfigReader(str(Path(test_data_path) / "5by5" / "5by5_bmi.ini")) + reader_config = SimpleNamespace( + sim_config=reader.sim_config.model_copy(update={"output_map_names": {}}), + grass_params=GrassParams(region="3by5", mask="5by3"), + stats_file=None, + ) + monkeypatch.setattr("itzi.bmi_itzi.ConfigReader", lambda _: reader_config) + previous = { + key: os.environ.get(key) + for key in ("GISRC", "GRASS_REGION", "WIND_OVERRIDE", "GRASS_MASK") + } + previous_region = gscript.region() + original_shape = (libraster.Rast_window_rows(), libraster.Rast_window_cols()) + model = BmiItzi() + + def fail(self, *args, **kwargs): + raise RuntimeError("BMI lifecycle failed") + + with monkeypatch.context() as fault: + fault.setattr(SimulationRunner, phase, fail) + with pytest.raises(RuntimeError, match="BMI lifecycle failed"): + model.initialize("config.ini") + model.finalize() + assert model.grass_session_manager is None + assert GrassSessionManager.current_context() is not None + assert {key: os.environ.get(key) for key in previous} == previous + assert gscript.region() == previous_region + assert (libraster.Rast_window_rows(), libraster.Rast_window_cols()) == original_shape + model.initialize("config.ini") + assert model.get_grid_shape(0) == (3, 5) + model.finalize() + model.finalize() + assert {key: os.environ.get(key) for key in previous} == previous diff --git a/tests/grass/test_ensemble.py b/tests/grass/test_ensemble.py index 8ebb4d6b..7eea12bd 100644 --- a/tests/grass/test_ensemble.py +++ b/tests/grass/test_ensemble.py @@ -18,7 +18,7 @@ from itzi.cli_parser import build_parser from itzi.ensemble.models import ResolvedSimulation, ValidationFailure from itzi.grass.names import derived_record_name -from itzi.grass.session import GrassSessionManager +from itzi.grass.session import GrassParams, GrassSessionManager from itzi.itzi import ( _resolve_ensemble_in_subprocess, itzi_run, @@ -171,6 +171,60 @@ def test_temporal_run_validation_failure_does_not_stop_next_member(tmp_path: Pat assert gutils.get_mapset_raster(f"{output.partition('@')[0]}_0000", gutils.getenv("MAPSET")) +@pytest.mark.usefixtures("grass_5by5") +@pytest.mark.parametrize("external_change", [False, True]) +def test_failed_member_restores_settings_before_next_default_region_verification( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, external_change: bool +) -> None: + from itzi.ensemble import run_validation + from itzi.ensemble.resolution import resolve_ensemble + + gscript.mapcalc("empty_mask=0") + first_path = write_study( + tmp_path, + grass={"region": "3by5", "mask": "empty_mask"}, + outputs={"rasters": {"prefix": "failed_{simulation}", "variables": ["water_depth"]}}, + ) + first = resolve_ensemble(load_batch([str(first_path)])[0][0].simulations)[0] + second_path = write_study( + tmp_path, + outputs={"rasters": {"prefix": "valid_{simulation}", "variables": ["water_depth"]}}, + ) + second = resolve_ensemble(load_batch([str(second_path)])[0][0].simulations)[0] + assert isinstance(first, ResolvedSimulation) + assert isinstance(second, ResolvedSimulation) + previous = { + key: os.environ.get(key) + for key in ("GISRC", "GRASS_REGION", "WIND_OVERRIDE", "GRASS_MASK") + } + rasters = gscript.list_strings("raster") + saved_regions = gscript.list_strings("region") + validate_inputs = run_validation._validate_inputs + + def validate(simulation, interface): + try: + validate_inputs(simulation, interface) + finally: + if external_change and simulation.simulation_id == first.simulation_id: + environment = os.environ.copy() + environment.pop("GRASS_REGION", None) + environment.pop("WIND_OVERRIDE", None) + gscript.run_command("g.region", n=40, s=0, e=50, w=0, res=10, env=environment) + + monkeypatch.setattr(run_validation, "_validate_inputs", validate) + failures = run_validation_worker((first, second)) + assert "mask excludes the complete" in failures[first.simulation_id] + if external_change: + assert "domain changed after resolution" in failures[second.simulation_id] + gscript.run_command("g.region", n=50, s=0, e=50, w=0, res=10) + else: + assert second.simulation_id not in failures + assert {key: os.environ.get(key) for key in previous} == previous + assert gscript.list_strings("raster") == rasters + assert gscript.list_strings("region") == saved_regions + assert run_validation_worker((second,)) == {} + + @pytest.mark.usefixtures("grass_5by5") def test_worker_rejects_changed_mask_region_and_inputs_after_resolution(tmp_path: Path) -> None: from grass.pygrass import utils as gutils @@ -187,26 +241,17 @@ def test_worker_rejects_changed_mask_region_and_inputs_after_resolution(tmp_path assert resolved.effective_mask.mode == "explicit" config = resolved.simulation_config - interface = GrassInterface( - config.start_time, - config.end_time, - np.float32, - None, - "stage_mask_values", - ) - try: + with GrassSessionManager(GrassParams(mask="stage_mask_values")): + interface = GrassInterface(config.start_time, config.end_time, np.float32) explicit = interface.get_npmask() assert explicit[:, :2].all() assert not explicit[:, 2:].any() friction = interface.read_raster_map(f"n@{gutils.getenv('MAPSET')}") np.testing.assert_array_equal(np.isnan(friction), explicit) assert np.isfinite(friction[~explicit]).all() - finally: - interface.cleanup() - with GrassInterface( - config.start_time, config.end_time, np.float32, None, None - ) as interface: + with GrassSessionManager(GrassParams()): + interface = GrassInterface(config.start_time, config.end_time, np.float32) np.testing.assert_array_equal(interface.get_npmask(), explicit) assert run_validation_worker((resolved,)).get(resolved.simulation_id) is None diff --git a/tests/grass/test_grass_input.py b/tests/grass/test_grass_input.py index 87227bff..836670dc 100644 --- a/tests/grass/test_grass_input.py +++ b/tests/grass/test_grass_input.py @@ -6,6 +6,8 @@ import numpy as np import pytest +from itzi.grass.session import GrassParams, GrassSessionManager + if TYPE_CHECKING: from itzi.grass.utils import MapData @@ -70,21 +72,23 @@ def test_domain_metadata_is_cached_per_interface(monkeypatch: pytest.MonkeyPatch reads = 0 read_domain = grass_interface.read_domain - def counted_read_domain(region_id: str | None) -> DomainData: + def counted_read_domain() -> DomainData: nonlocal reads reads += 1 - return read_domain(region_id) + return read_domain() monkeypatch.setattr(grass_interface, "read_domain", counted_read_domain) start = datetime(2020, 1, 1) end = start + timedelta(seconds=1) - with grass_interface.GrassInterface(start, end, np.float32, "3by5", None) as interface: + with GrassSessionManager(GrassParams(region="3by5")): + interface = grass_interface.GrassInterface(start, end, np.float32) domain = interface.get_domain_data() assert (domain.rows, domain.cols) == (3, 5) assert interface.get_domain_data() is domain assert reads == 1 - with grass_interface.GrassInterface(start, end, np.float32, None, None) as interface: + with GrassSessionManager(GrassParams()): + interface = grass_interface.GrassInterface(start, end, np.float32) new_domain = interface.get_domain_data() assert (new_domain.rows, new_domain.cols) == (5, 5) assert new_domain is not domain @@ -113,7 +117,8 @@ def test_visible_input_and_current_output_with_same_name(grass_xy_session) -> No gscript.run_command("g.region", n=3, s=0, e=3, w=0, res=1) start = datetime(2020, 1, 1) - with GrassInterface(start, start + timedelta(seconds=1), np.int32, None, None) as interface: + with GrassSessionManager(GrassParams()): + interface = GrassInterface(start, start + timedelta(seconds=1), np.int32) interface.write_raster_map( np.ones((3, 3), dtype=np.int32), output_id, "ground_elevation", 0 ) @@ -136,10 +141,10 @@ def test_direct_runner_qualifies_ids_before_creating_providers(test_data_path: s config = reader.sim_config.model_copy( update={"output_map_names": {"water_depth": "direct_id_check"}} ) - runner = SimulationRunner(config, reader.grass_params) - - assert runner.sim_config.input_map_names["ground_elevation"] == "z@5by5" - assert runner.sim_config.output_map_names["water_depth"] == "direct_id_check@5by5" + with GrassSessionManager(reader.grass_params): + runner = SimulationRunner(config) + assert runner.sim_config.input_map_names["ground_elevation"] == "z@5by5" + assert runner.sim_config.output_map_names["water_depth"] == "direct_id_check@5by5" @pytest.mark.usefixtures("grass_5by5") @@ -150,7 +155,8 @@ def test_vector_writer_uses_bare_name_from_qualified_output_id() -> None: from itzi.grass.interface import GrassInterface start = datetime(2020, 1, 1) - with GrassInterface(start, start + timedelta(seconds=1), np.float32, None, None) as interface: + with GrassSessionManager(GrassParams()): + interface = GrassInterface(start, start + timedelta(seconds=1), np.float32) interface.write_vector_map( DrainageNetworkTopology(nodes=(), links=()), DrainageNetworkAttributes(nodes=(), links=()), diff --git a/tests/grass/test_itzi.py b/tests/grass/test_itzi.py index 52f9b3ed..3363b553 100644 --- a/tests/grass/test_itzi.py +++ b/tests/grass/test_itzi.py @@ -1,5 +1,7 @@ import os +from collections.abc import Iterator from configparser import ConfigParser +from contextlib import contextmanager from datetime import timedelta from uuid import uuid4 @@ -54,6 +56,7 @@ def _create_timed_rain_inputs(name_prefix: str) -> dict[str, str]: } +@contextmanager def _build_timed_rain_runner( test_data_temp_path: str, *, @@ -61,7 +64,7 @@ def _build_timed_rain_runner( duration: str, record_step: str, prefix: str, -) -> SimulationRunner: +) -> Iterator[SimulationRunner]: current_mapset = gscript.read_command("g.mapset", flags="p").rstrip() config_dict = { "input": { @@ -91,12 +94,14 @@ def _build_timed_rain_runner( parser.write(file_handle) conf_data = ConfigReader(config_file) - GrassSessionManager.ensure_temporal_initialized() - return SimulationRunner( - conf_data.sim_config, - conf_data.grass_params, - stats_file=conf_data.stats_file, - ).initialize() + with GrassSessionManager(conf_data.grass_params): + runner = SimulationRunner( + conf_data.sim_config, stats_file=conf_data.stats_file + ).initialize() + try: + yield runner + finally: + runner.finalize() @pytest.mark.usefixtures("grass_5by5_sim") @@ -158,12 +163,12 @@ def test_region_mask(test_data_path): conf_data = ConfigReader(config_file) sim_params = conf_data.sim_config grass_params = conf_data.grass_params - sim_runner = SimulationRunner( - sim_params, - grass_params, - ) - # Run simulation - sim_runner.initialize().run().finalize() + with GrassSessionManager(grass_params): + sim_runner = SimulationRunner(sim_params).initialize() + try: + sim_runner.run() + finally: + sim_runner.finalize() # Check temporary mask and region assert ( int(gscript.parse_command("r.univar", map="out_5by5_max_flow_speed_0002", flags="g")["n"]) @@ -214,11 +219,8 @@ def test_fails_when_region_has_no_dem_data(test_data_temp_path): with pytest.raises( RuntimeError, match=r"input map contains only NULL/NaN cells" ): - SimulationRunner( - conf_data.sim_config, - conf_data.grass_params, - stats_file=conf_data.stats_file, - ).initialize() + with GrassSessionManager(conf_data.grass_params): + SimulationRunner(conf_data.sim_config, stats_file=conf_data.stats_file).initialize() @pytest.mark.usefixtures("grass_5by5") @@ -246,13 +248,14 @@ def test_empty_strds_fails_with_clear_error(test_data_temp_path: str, temporal_t RuntimeError, match=rf"STRDS <{strds_name}@{current_mapset}> has no temporal extent", ): - _build_timed_rain_runner( + with _build_timed_rain_runner( test_data_temp_path, input_names=input_names, duration="00:00:20", record_step="00:00:20", prefix=f"out_empty_{temporal_type}_{uuid4().hex[:8]}", - ) + ): + pass @pytest.mark.usefixtures("grass_5by5") @@ -296,13 +299,14 @@ def test_relative_strds_with_unsupported_unit_fails_with_clear_error( r"" ), ): - _build_timed_rain_runner( + with _build_timed_rain_runner( test_data_temp_path, input_names=input_names, duration="00:00:20", record_step="00:00:20", prefix=f"out_monthly_relative_{uuid4().hex[:8]}", - ) + ): + pass @pytest.mark.usefixtures("grass_5by5") @@ -321,15 +325,13 @@ def test_timed_grass_rain_switches_cleanly_around_boundary( expected_input_deadline_seconds: int, ): input_names = _create_timed_rain_inputs("timed_boundary") - sim_runner = _build_timed_rain_runner( + with _build_timed_rain_runner( test_data_temp_path, input_names=input_names, duration="00:00:20", record_step="00:00:20", prefix=f"out_timed_boundary_{target_seconds}_{uuid4().hex[:8]}", - ) - - try: + ) as sim_runner: simulation = sim_runner.sim simulation.update_until(timedelta(seconds=target_seconds)) @@ -340,22 +342,18 @@ def test_timed_grass_rain_switches_cleanly_around_boundary( assert simulation.next_ts["input"] == simulation.start_time + timedelta( seconds=expected_input_deadline_seconds ) - finally: - sim_runner.g_interface.cleanup() @pytest.mark.usefixtures("grass_5by5") def test_timed_grass_rain_is_applied_before_a_step_crosses_its_boundary(test_data_temp_path): input_names = _create_timed_rain_inputs("timed_crossing") - sim_runner = _build_timed_rain_runner( + with _build_timed_rain_runner( test_data_temp_path, input_names=input_names, duration="00:00:20", record_step="00:00:15", prefix=f"out_timed_crossing_{uuid4().hex[:8]}", - ) - - try: + ) as sim_runner: simulation = sim_runner.sim simulation.update() @@ -369,5 +367,3 @@ def test_timed_grass_rain_is_applied_before_a_step_crosses_its_boundary(test_dat simulation.raster_domain.get_array("rainfall_rate"), 360.0 / (1000 * 3600), ) - finally: - sim_runner.g_interface.cleanup() diff --git a/tests/grass/test_mask.py b/tests/grass/test_mask.py index 45ab196a..4846d0c2 100644 --- a/tests/grass/test_mask.py +++ b/tests/grass/test_mask.py @@ -13,6 +13,8 @@ import numpy as np import pytest +from itzi.grass.session import GrassParams, GrassSessionManager + if TYPE_CHECKING: from itzi_core import Simulation @@ -75,13 +77,11 @@ def test_mask_selection_has_identical_rules_and_preserves_rasters( original_rasters = gscript.list_strings("raster") expected = np.tile(excluded, (5, 1)) start = datetime(2020, 1, 1) - with GrassInterface( - start, - start + timedelta(seconds=1), - np.float32, - None, - source if selection == "explicit" else None, - ) as interface: + with GrassSessionManager( + GrassParams(mask=source if selection == "explicit" else None) + ) as manager: + interface = GrassInterface(start, start + timedelta(seconds=1), np.float32) + assert manager.effective_mask.mode == ("explicit" if selection == "explicit" else "active") assert os.environ["GRASS_MASK"] == source assert resolve_effective_mask(None) == EffectiveMask("active", source) np.testing.assert_array_equal(interface.get_npmask(), expected) @@ -95,7 +95,7 @@ def test_mask_selection_has_identical_rules_and_preserves_rasters( assert int(stats["n"]) == np.count_nonzero(~expected) assert os.environ.get("GRASS_MASK") == original_mask - interface.cleanup() + manager.close() assert os.environ.get("GRASS_MASK") == original_mask assert gscript.list_strings("raster") == original_rasters @@ -112,7 +112,8 @@ def test_missing_grass_mask_does_not_fall_back_to_mapset_mask( monkeypatch.setenv("GRASS_MASK", "missing_mask") assert resolve_effective_mask(None) == EffectiveMask("none", None) start = datetime(2020, 1, 1) - with GrassInterface(start, start + timedelta(seconds=1), np.float32, None, None) as interface: + with GrassSessionManager(GrassParams()): + interface = GrassInterface(start, start + timedelta(seconds=1), np.float32) assert not interface.get_npmask().any() np.testing.assert_allclose(interface.read_raster_map("n@5by5"), 0.05) assert os.environ["GRASS_MASK"] == "missing_mask" @@ -138,9 +139,8 @@ def test_qualified_mask_in_another_mapset_and_exception_restoration( start = datetime(2020, 1, 1) with pytest.raises(OpenError): - with GrassInterface( - start, start + timedelta(seconds=1), np.float32, None, "5by3" - ) as interface: + with GrassSessionManager(GrassParams(mask="5by3")): + interface = GrassInterface(start, start + timedelta(seconds=1), np.float32) assert os.environ["GRASS_MASK"] == "5by3@5by5" with pytest.raises(OpenError): interface.read_raster_map("missing_input@5by5") @@ -148,7 +148,8 @@ def test_qualified_mask_in_another_mapset_and_exception_restoration( interface.read_raster_map("missing_input@5by5") assert os.environ["GRASS_MASK"] == "shared_mask@PERMANENT" - with GrassInterface(start, start + timedelta(seconds=1), np.float32, None, None) as interface: + with GrassSessionManager(GrassParams()): + interface = GrassInterface(start, start + timedelta(seconds=1), np.float32) mask = interface.get_npmask() assert mask[:, 0].all() assert not mask[:, 1:].any() @@ -180,19 +181,18 @@ def test_runner_restores_mask_when_lifecycle_fails( "input_map_names": dict(config.input_map_names) | {"ground_elevation": "bad_dem"} } ) - runner = SimulationRunner(config, GrassParams(region="3by5", mask="mask_values")) - if phase == "initialize": - with pytest.raises(RuntimeError, match="only NULL/NaN cells"): - runner.initialize() - else: - runner.initialize() - assert os.environ["GRASS_MASK"] == "mask_values@5by5" + with pytest.raises( + RuntimeError, + match="only NULL/NaN cells" if phase == "initialize" else "finalization failed", + ): + with GrassSessionManager(GrassParams(region="3by5", mask="mask_values")): + runner = SimulationRunner(config).initialize() + assert os.environ["GRASS_MASK"] == "mask_values@5by5" - def fail_finalize(self: Simulation) -> None: - raise RuntimeError("finalization failed") + def fail_finalize(self: Simulation) -> None: + raise RuntimeError("finalization failed") - monkeypatch.setattr(type(runner.sim), "finalize", fail_finalize) - with pytest.raises(RuntimeError, match="finalization failed"): + monkeypatch.setattr(type(runner.sim), "finalize", fail_finalize) runner.finalize() assert os.environ["GRASS_MASK"] == "previous_mask" assert gscript.region() == previous_region @@ -201,14 +201,12 @@ def fail_finalize(self: Simulation) -> None: @pytest.mark.usefixtures("grass_5by5") -@pytest.mark.parametrize("phase", ["finalize", "delete"]) @pytest.mark.parametrize("explicit_mask", [False, True]) @pytest.mark.parametrize("previous_mask", [False, True]) @pytest.mark.parametrize("region", [None, "3by3"]) -def test_runner_cleanup_preserves_open_raster_and_restores_reads( +def test_work_context_restores_fresh_raster_reads( monkeypatch: pytest.MonkeyPatch, test_data_path: str, - phase: str, explicit_mask: bool, previous_mask: bool, region: str | None, @@ -240,58 +238,37 @@ def test_runner_cleanup_preserves_open_raster_and_restores_reads( from itzi.grass.session import GrassParams, GrassSessionManager from itzi.simulation_runner import SimulationRunner -with GrassSessionManager(GrassParams()): - reader = ConfigReader(str(Path(sys.argv[1]) / "5by5" / "5by5.ini")) - config = reader.sim_config.model_copy(update={"output_map_names": {}}) - original_mask = os.environ.get("GRASS_MASK") - mask = "5by3" if sys.argv[3] == "True" else None - region = sys.argv[4] or None - original_region = gscript.region() - runner = SimulationRunner(config, GrassParams(region=region, mask=mask)).initialize() - with RasterRow("gradient", mapset="5by5", mode="r") as source: - original_row = np.asarray(source.get_row(0)).copy() - assert original_row.shape == (5,) - np.testing.assert_allclose( - original_row[np.isfinite(original_row)], - np.arange(11, 16)[np.isfinite(original_row)], - ) - np.testing.assert_array_equal( - np.isnan(original_row)[1:4] if region else np.isnan(original_row), - runner.g_interface.get_npmask()[0], - ) - if region: - expected = np.arange(20, 50, 10)[:, None] + np.arange(2, 5) - np.testing.assert_allclose( - runner.g_interface.read_raster_map("gradient@5by5"), expected - ) - runner.g_interface.write_raster_map( - expected.astype(np.float32), "roundtrip@5by5", "friction", 0 - ) - info = gscript.raster_info("roundtrip@5by5") - assert (int(info["rows"]), int(info["cols"])) == (3, 3) - np.testing.assert_allclose( - runner.g_interface.read_raster_map("roundtrip@5by5"), expected - ) - np.testing.assert_array_equal(np.asarray(source.get_row(0)), original_row) - if sys.argv[2] == "finalize": - runner.finalize() - else: - del runner - assert os.environ.get("GRASS_MASK") == original_mask - expected = [original_mask is not None, False, False, False, False] - restored_row = np.asarray(source.get_row(0)) - np.testing.assert_array_equal(np.isnan(restored_row), expected) - np.testing.assert_allclose(restored_row[1:], np.arange(12, 16)) - assert gscript.region() == original_region - assert (libraster.Rast_window_rows(), libraster.Rast_window_cols()) == (5, 5) - - next_runner = SimulationRunner(config, GrassParams(region="3by5", mask=mask)).initialize() +reader = ConfigReader(str(Path(sys.argv[1]) / "5by5" / "5by5.ini")) +config = reader.sim_config.model_copy(update={"output_map_names": {}}) +original_mask = os.environ.get("GRASS_MASK") +mask = "5by3" if sys.argv[2] == "True" else None +region = sys.argv[3] or None +original_region = gscript.region() +with GrassSessionManager(GrassParams(region=region, mask=mask)): + runner = SimulationRunner(config).initialize() + values = runner.g_interface.read_raster_map("gradient@5by5") + expected = (np.arange(20, 50, 10)[:, None] + np.arange(2, 5) if region + else np.arange(10, 60, 10)[:, None] + np.arange(1, 6)).astype(np.float32) + excluded = runner.g_interface.get_npmask() + expected[excluded] = np.nan + np.testing.assert_array_equal(values, expected) + runner.g_interface.write_raster_map(values, "roundtrip@5by5", "friction", 0) + np.testing.assert_array_equal(runner.g_interface.read_raster_map("roundtrip@5by5"), expected) + runner.finalize() +assert os.environ.get("GRASS_MASK") == original_mask +assert gscript.region() == original_region +assert (libraster.Rast_window_rows(), libraster.Rast_window_cols()) == (5, 5) +with RasterRow("gradient", mapset="5by5", mode="r") as source: + restored_row = np.asarray(source.get_row(0)) + np.testing.assert_array_equal(np.isnan(restored_row), [original_mask is not None, False, False, False, False]) + np.testing.assert_allclose(restored_row[1:], np.arange(12, 16)) +with GrassSessionManager(GrassParams(region="3by5", mask=mask)): + next_runner = SimulationRunner(config).initialize() assert next_runner.g_interface.read_raster_map("n@5by5").shape == (3, 5) next_runner.finalize() - assert os.environ.get("GRASS_MASK") == original_mask +assert os.environ.get("GRASS_MASK") == original_mask """, test_data_path, - phase, str(explicit_mask), region or "", ], diff --git a/tests/grass/test_region.py b/tests/grass/test_region.py index 29b1e1a6..48250c20 100644 --- a/tests/grass/test_region.py +++ b/tests/grass/test_region.py @@ -2,7 +2,6 @@ from __future__ import annotations -import copy import os from ctypes import byref from datetime import datetime, timedelta @@ -11,6 +10,8 @@ import numpy as np import pytest +from itzi.grass.session import GrassParams, GrassSessionManager + def working_windows() -> tuple[bytes, bytes, bytes]: from grass.lib import gis as libgis, raster as libraster @@ -45,27 +46,34 @@ def test_successive_regions_restore_overrides_windows_and_saved_regions( previous_region = gscript.region() previous_overrides = (os.environ.get("WIND_OVERRIDE"), os.environ.get("GRASS_REGION")) previous_saved_regions = gscript.list_strings("region") + gscript.mapcalc("gradient=float(row()*10+col())") start = datetime(2020, 1, 1) for index, region_id in enumerate(("3by5", None, "3by5")): rows = 3 if region_id else int(previous_region["rows"]) - with GrassInterface( - start, start + timedelta(seconds=1), np.float32, region_id, None - ) as interface: + with GrassSessionManager(GrassParams(region=region_id)) as manager: + interface = GrassInterface(start, start + timedelta(seconds=1), np.float32) assert (interface.yr, interface.xr) == (rows, 5) assert (interface.get_domain_data().rows, interface.get_domain_data().cols) == ( rows, 5, ) assert (gscript.region()["rows"], gscript.region()["cols"]) == (rows, 5) - values = interface.read_raster_map("n@5by5") + values = interface.read_raster_map("gradient@5by5") assert values.shape == (rows, 5) - np.testing.assert_allclose(values, 0.05) + expected = ( + np.arange(1, 6) + + np.arange(10, (int(previous_region["rows"]) + 1) * 10, 10)[:, None] + ) + if region_id: + expected = expected[1:4] if previous_region["rows"] == 5 else expected[:3] + np.testing.assert_allclose(values, expected) output = f"region_output_{index}@5by5" interface.write_raster_map(values, output, "friction", 0) info = gscript.raster_info(output) assert (int(info["rows"]), int(info["cols"])) == (rows, 5) - interface.cleanup() + np.testing.assert_array_equal(interface.read_raster_map(output), values) + manager.close() assert gscript.region() == previous_region assert working_windows() == previous_windows assert ( @@ -85,7 +93,6 @@ def test_constructor_failure_restores_region_immediately( from grass.pygrass.gis.region import Region from itzi.ensemble.models import EnsembleError - from itzi.grass.interface import GrassInterface from itzi.messenger import FatalError gscript.run_command("g.region", n=20, s=0, e=20, w=0, res=10, save="small_region") @@ -94,73 +101,72 @@ def test_constructor_failure_restores_region_immediately( libraster.Rast_unset_window() previous_windows = working_windows() previous_saved_regions = gscript.list_strings("region") - start = datetime(2020, 1, 1) region = failure if failure != "missing_mask" else "3by5" mask = "missing_mask" if failure == "missing_mask" else None with pytest.raises((CalledModuleError, FatalError, EnsembleError)): - GrassInterface(start, start + timedelta(seconds=1), np.float32, region, mask) + GrassSessionManager(GrassParams(region=region, mask=mask)).open() assert os.environ["WIND_OVERRIDE"] == "3by5" assert working_windows() == previous_windows assert gscript.list_strings("region") == previous_saved_regions + with GrassSessionManager(GrassParams()): + assert gscript.region()["rows"] == 3 @pytest.mark.usefixtures("grass_5by5") -def test_custom_gis_and_split_raster_windows_are_restored() -> None: +def test_nested_jobs_restore_the_outer_work_window() -> None: from grass.lib import raster as libraster - from grass.pygrass.gis.region import Region from itzi.grass.interface import GrassInterface - configured = Region() - libraster.Rast_unset_window() - gis = copy.deepcopy(configured) - gis.north = 40 - gis.adjust() - raster_input = copy.deepcopy(configured) - raster_input.north = 30 - raster_input.adjust() - raster_output = copy.deepcopy(configured) - raster_output.east = 40 - raster_output.adjust() - try: - raster_input.set_current() - libraster.Rast_unset_window() - libraster.Rast_get_input_window(raster_input.byref()) - libraster.Rast_set_output_window(raster_output.byref()) - gis.set_current() + start = datetime(2020, 1, 1) + gscript.run_command("g.region", n=40, s=10, e=40, w=10, res=10, flags="u", save="3by3") + with GrassSessionManager(GrassParams(region="3by5", mask="5by3")): previous_windows = working_windows() - start = datetime(2020, 1, 1) - with GrassInterface( - start, start + timedelta(seconds=1), np.float32, "3by5", None - ) as interface: - assert interface.read_raster_map("n@5by5").shape == (3, 5) + previous_region = gscript.region() + previous_mask = os.environ["GRASS_MASK"] + with GrassSessionManager(GrassParams(region="3by3", mask="n")): + interface = GrassInterface(start, start + timedelta(seconds=1), np.float32) + assert interface.read_raster_map("n@5by5").shape == (3, 3) + assert (libraster.Rast_window_rows(), libraster.Rast_window_cols()) == (3, 3) assert working_windows() == previous_windows - finally: - configured.set_current() - libraster.Rast_unset_window() + assert gscript.region() == previous_region + assert os.environ["GRASS_MASK"] == previous_mask + with GrassSessionManager(GrassParams()): + assert working_windows() == previous_windows + assert os.environ["GRASS_MASK"] == previous_mask @pytest.mark.usefixtures("grass_5by5") -def test_domain_lookup_preserves_working_windows_and_open_raster() -> None: - from grass.lib import raster as libraster - from grass.pygrass.gis.region import Region +def test_domain_lookup_reads_the_active_window_without_reloading() -> None: from grass.pygrass.raster import RasterRow from itzi.grass.utils import read_domain - configured = Region() - libraster.Rast_unset_window() - libraster.Rast_get_input_window(configured.byref()) - gis = copy.deepcopy(configured) - gis.north = 40 - gis.adjust() - gis.set_current() - try: + with GrassSessionManager(GrassParams(region="3by5")): with RasterRow("n", mapset="5by5", mode="r") as source: previous_windows = working_windows() - domain = read_domain("3by5") + domain = read_domain() assert (domain.rows, domain.cols) == (3, 5) assert working_windows() == previous_windows np.testing.assert_allclose(np.asarray(source.get_row(0)), 0.05) - finally: - configured.set_current() + + +@pytest.mark.usefixtures("grass_5by5") +def test_default_region_snapshot_preserves_external_changes() -> None: + from itzi.grass.interface import GrassInterface + + start = datetime(2020, 1, 1) + with GrassSessionManager(GrassParams()): + interface = GrassInterface(start, start + timedelta(seconds=1), np.float32) + environment = os.environ.copy() + environment.pop("GRASS_REGION", None) + environment.pop("WIND_OVERRIDE", None) + gscript.run_command("g.region", n=40, s=0, e=50, w=0, res=10, env=environment) + assert gscript.region()["rows"] == 5 + assert interface.read_raster_map("n@5by5").shape == (5, 5) + with GrassSessionManager(GrassParams()): + assert gscript.region()["rows"] == 5 + assert gscript.region()["rows"] == 4 + with GrassSessionManager(GrassParams()): + next_interface = GrassInterface(start, start + timedelta(seconds=1), np.float32) + assert next_interface.read_raster_map("n@5by5").shape == (4, 5) diff --git a/tests/grass/test_session_context.py b/tests/grass/test_session_context.py new file mode 100644 index 00000000..89775a2a --- /dev/null +++ b/tests/grass/test_session_context.py @@ -0,0 +1,163 @@ +"""Isolate session transitions which can otherwise exit through the GRASS C API.""" + +import os +import subprocess +import sys +from pathlib import Path + +import pytest + + +def test_bootstrap_outside_grass_discovers_python_path_and_cleans_session(tmp_path: Path) -> None: + import grass.script as gscript + + gscript.setup.setup_runtime_env() + gscript.create_project(tmp_path, name="bootstrap") + with gscript.setup.init( + tmp_path, + location="bootstrap", + mapset="PERMANENT", + grass_path="grass", + env=os.environ.copy(), + ) as runtime: + gscript.run_command("g.region", n=30, s=0, e=50, w=0, res=10, env=runtime.env) + environment = os.environ | {"PYTHONPATH": str(Path(__file__).parents[2] / "src")} + environment.pop("GISRC", None) + result = subprocess.run( + [ + sys.executable, + "-c", + """ +import os +import sys +from pathlib import Path +from itzi.grass.session import GrassParams, GrassSessionManager + +assert "grass.script" not in sys.modules +assert GrassSessionManager.current_context() is None +with GrassSessionManager(GrassParams(sys.argv[1], "bootstrap", "PERMANENT")) as manager: + from itzi.grass.utils import read_domain + assert (read_domain().rows, read_domain().cols) == (3, 5) + gisrc = Path(os.environ["GISRC"]) +assert "GISRC" not in os.environ +assert not gisrc.exists() +""", + str(tmp_path), + ], + env=environment, + capture_output=True, + text=True, + timeout=60, + ) + assert result.returncode == 0, result.stdout + result.stderr + + +@pytest.mark.parametrize("inherited", ["none", "WIND_OVERRIDE", "GRASS_REGION", "both"]) +def test_mapset_switches_restore_caller_and_owned_sessions(tmp_path: Path, inherited: str) -> None: + result = subprocess.run( + [ + sys.executable, + "-c", + """ +import os +import sys +from datetime import datetime, timedelta +from pathlib import Path + +import numpy as np + +from itzi.grass.session import GrassParams, GrassSessionManager + +assert "grass.pygrass" not in sys.modules +assert "itzi.grass.interface" not in sys.modules +assert "itzi.grass.utils" not in sys.modules +import grass.script as gs + +database = Path(sys.argv[1]) +gs.setup.setup_runtime_env() +gs.create_project(database, name="project") +gs.create_project(database, name="second") +with gs.setup.init(database, location="project", mapset="PERMANENT", grass_path="grass"): + gs.run_command("g.region", n=50, s=0, e=50, w=0, res=10) + gs.mapcalc("gradient=float(row()*10+col())") + gs.mapcalc("MASK=if(col()==1, 0, 1)") + gs.run_command("g.region", n=40, s=10, e=50, w=0, res=10, flags="u", save="outer") + gs.run_command("g.mapset", flags="c", mapset="target") + gs.run_command("g.region", n=30, s=0, e=40, w=0, res=10) + gs.mapcalc("gradient=float(100+row()*10+col())") + gs.mapcalc("MASK=if(col()==4, 0, 1)") + gs.run_command("g.mapset", mapset="PERMANENT") + if sys.argv[2] in ("WIND_OVERRIDE", "both"): + os.environ["WIND_OVERRIDE"] = "outer" + if sys.argv[2] in ("GRASS_REGION", "both"): + os.environ["GRASS_REGION"] = gs.region_env(region="outer") + os.environ["GRASS_MASK"] = "MASK@PERMANENT" + previous = {key: os.environ.get(key) for key in ("GISRC", "GRASS_REGION", "WIND_OVERRIDE", "GRASS_MASK")} + gisrc_contents = Path(previous["GISRC"]).read_bytes() + original_region = gs.region() + saved_regions = gs.list_strings("region") + params = GrassParams(str(database), "project", "target") + start = datetime(2020, 1, 1) + for index in range(2): + with GrassSessionManager(params) as manager: + from grass.lib import raster as lr + import grass.temporal as tgis + from itzi.grass.interface import GrassInterface + assert gs.gisenv()["MAPSET"] == "target" + assert manager.current_context()[2] == "target" + assert tgis.get_current_mapset() == "target" + assert manager.effective_mask.source == "MASK@target" + assert (lr.Rast_window_rows(), lr.Rast_window_cols()) == (3, 4) + interface = GrassInterface(start, start + timedelta(seconds=1), np.float32) + values = interface.read_raster_map("gradient@target") + expected = 100 + np.arange(10, 40, 10)[:, None] + np.arange(1, 5) + np.testing.assert_allclose(values[:, :3], expected[:, :3]) + assert np.isnan(values[:, 3]).all() + interface.write_raster_map(values, f"roundtrip_{index}@target", "friction", 0) + np.testing.assert_array_equal(interface.read_raster_map(f"roundtrip_{index}@target"), values) + with GrassSessionManager(GrassParams()): + assert gs.region()["rows"] == 3 + assert gs.gisenv()["MAPSET"] == "target" + assert {key: os.environ.get(key) for key in previous} == previous + assert gs.region() == original_region + assert manager.current_context()[2] == "PERMANENT" + assert tgis.get_current_mapset() == "PERMANENT" + assert gs.list_strings("region") == saved_regions + assert Path(previous["GISRC"]).read_bytes() == gisrc_contents + with GrassSessionManager(GrassParams()): + restored = GrassInterface(start, start + timedelta(seconds=1), np.float32) + assert restored.yr == original_region["rows"] + assert restored.get_npmask()[:, 0].all() + assert not gs.find_file(f"roundtrip_{index}", mapset="PERMANENT")["name"] + os.environ.pop("GRASS_REGION", None) + os.environ.pop("WIND_OVERRIDE", None) + os.environ.pop("GRASS_MASK", None) + +# Cached modules and C libraries must follow newly acquired underlying sessions. +for project, rows in (("second", 3), ("project", 5), ("second", 4)): + with gs.setup.init(database, location=project, mapset="PERMANENT", grass_path="grass"): + gs.run_command("g.region", n=rows*10, s=0, e=50, w=0, res=10) + gs.mapcalc("owned_gradient=float(row()*10+col())", overwrite=True) + with GrassSessionManager(GrassParams()): + assert (lr.Rast_window_rows(), lr.Rast_window_cols()) == (rows, 5) + assert (lr.Rast_window_rows(), lr.Rast_window_cols()) == (rows, 5) + with GrassSessionManager(GrassParams(str(database), project, "PERMANENT")) as manager: + job_gisrc = os.environ["GISRC"] + interface = GrassInterface(start, start + timedelta(seconds=1), np.float32) + assert interface.yr == rows + assert manager.current_context()[1] == project + assert tgis.get_current_location() == project + assert interface.read_raster_map("owned_gradient@PERMANENT").shape == (rows, 5) + manager.close() + assert "GISRC" not in os.environ + assert not Path(job_gisrc).exists() +""", + str(tmp_path), + inherited, + ], + env=os.environ | {"PYTHONPATH": os.pathsep.join(sys.path)}, + capture_output=True, + text=True, + timeout=120, + ) + assert result.returncode == 0, result.stdout + result.stderr diff --git a/tests/grass/test_temporal_overwrite.py b/tests/grass/test_temporal_overwrite.py index fe2340df..b983c6d1 100644 --- a/tests/grass/test_temporal_overwrite.py +++ b/tests/grass/test_temporal_overwrite.py @@ -1,7 +1,9 @@ from __future__ import annotations import os +from collections.abc import Iterator from configparser import ConfigParser +from contextlib import contextmanager from uuid import uuid4 import grass.script as gscript @@ -13,11 +15,12 @@ from itzi.grass.session import GrassSessionManager +@contextmanager def _build_runner( test_data_temp_path: str, prefix: str, temporal_type: TemporalType, -) -> SimulationRunner: +) -> Iterator[SimulationRunner]: current_mapset = gscript.read_command("g.mapset", flags="p").rstrip() config_dict = { "input": { @@ -55,11 +58,8 @@ def _build_runner( parser.write(file_handle) conf_data = ConfigReader(config_file) - return SimulationRunner( - conf_data.sim_config, - conf_data.grass_params, - stats_file=conf_data.stats_file, - ) + with GrassSessionManager(conf_data.grass_params): + yield SimulationRunner(conf_data.sim_config, stats_file=conf_data.stats_file) @pytest.mark.usefixtures("grass_5by5") @@ -96,8 +96,9 @@ def test_runner_creation_fails_when_overwriting_strds_with_different_temporal_ty monkeypatch.setenv("GRASS_OVERWRITE", "1") prefix = f"out_overwrite_{existing_temporal_type}_{simulation_temporal_type}_{uuid4().hex[:8]}" - initial_runner = _build_runner(test_data_temp_path, prefix, existing_temporal_type) - initial_runner.initialize().run().finalize() + with _build_runner(test_data_temp_path, prefix, existing_temporal_type) as initial_runner: + initial_runner.initialize().run().finalize() with pytest.raises(RuntimeError, match=r"temporal type"): - _build_runner(test_data_temp_path, prefix, simulation_temporal_type).initialize() + with _build_runner(test_data_temp_path, prefix, simulation_temporal_type) as runner: + runner.initialize() diff --git a/tests/grass/test_tutorial.py b/tests/grass/test_tutorial.py index 19e597b3..fae93554 100644 --- a/tests/grass/test_tutorial.py +++ b/tests/grass/test_tutorial.py @@ -125,12 +125,15 @@ def test_tutorial(itzi_tutorial, test_data_path): config_file = os.path.join(test_data_path, "tutorial_files", "tutorial.ini") conf_data = ConfigReader(config_file) - sim_runner = SimulationRunner( - conf_data.sim_config, - conf_data.grass_params, - stats_file=conf_data.stats_file, - ).initialize() - sim_runner.run().finalize() + with session.GrassSessionManager(conf_data.grass_params): + sim_runner = SimulationRunner( + conf_data.sim_config, + stats_file=conf_data.stats_file, + ).initialize() + try: + sim_runner.run() + finally: + sim_runner.finalize() # Check the results hmax_maps = gscript.list_grouped("raster", pattern="nc_itzi_tutorial_max_water_depth_*")[ "PERMANENT" @@ -187,12 +190,15 @@ def test_tutorial_drainage(itzi_tutorial, test_data_path, helpers): # Run the simulation conf_data = ConfigReader(config_file) - sim_runner = SimulationRunner( - conf_data.sim_config, - conf_data.grass_params, - stats_file=conf_data.stats_file, - ).initialize() - sim_runner.run().finalize() + with session.GrassSessionManager(conf_data.grass_params): + sim_runner = SimulationRunner( + conf_data.sim_config, + stats_file=conf_data.stats_file, + ).initialize() + try: + sim_runner.run() + finally: + sim_runner.finalize() # Test consistency of stats file stat_file_path = pathlib.Path("nc_itzi_tutorial_drainage.csv") @@ -296,14 +302,13 @@ def test_tutorial_drainage_without_vector_output(itzi_tutorial, test_data_path): """Run coupled drainage without creating vector output.""" config_file = os.path.join(test_data_path, "tutorial_files", "tutorial_drainage.ini") conf_data = ConfigReader(config_file) - sim_runner = SimulationRunner( - conf_data.sim_config, - conf_data.grass_params, - stats_file=conf_data.stats_file, - ).initialize() - - try: - assert sim_runner.sim.report.vector_provider is None - sim_runner.run() - finally: - sim_runner.finalize() + with session.GrassSessionManager(conf_data.grass_params): + sim_runner = SimulationRunner( + conf_data.sim_config, + stats_file=conf_data.stats_file, + ).initialize() + try: + assert sim_runner.sim.report.vector_provider is None + sim_runner.run() + finally: + sim_runner.finalize() diff --git a/tests/test_grass_session.py b/tests/test_grass_session.py index ca9bc1c8..e43e90d4 100644 --- a/tests/test_grass_session.py +++ b/tests/test_grass_session.py @@ -34,15 +34,35 @@ def finish(self) -> None: self.finished = True -def test_manager_activates_and_finishes_created_session(monkeypatch) -> None: +@pytest.mark.parametrize("setup_fails", [False, True]) +def test_manager_activates_and_finishes_created_session( + monkeypatch, tmp_path, setup_fails +) -> None: session = FakeGrassSession({"GISRC": "/tmp/worker-gisrc"}) grass_script = ModuleType("grass.script") - grass_script.__dict__["setup"] = SimpleNamespace(init=lambda **_: session) + + def init(**kwargs): + os.environ.update(session.env) + return session + + grass_script.__dict__["setup"] = SimpleNamespace(init=init) grass_package = ModuleType("grass") grass_package.__dict__["script"] = grass_script monkeypatch.setattr("itzi.grass.session.os.access", lambda *_: True) + monkeypatch.setattr("itzi.grass.session._validate_target", lambda _: None) + monkeypatch.setattr("itzi.grass.session.ensure_min_version", lambda: None) monkeypatch.setattr(GrassSessionManager, "ensure_temporal_initialized", lambda self: None) + monkeypatch.setattr("itzi.grass.session._initialized_temporal_session", None) + restored = [] + + def select_context(self, stack): + assert os.environ["GISRC"] == session.env["GISRC"] + stack.callback(restored.append, "restored") + if setup_fails: + raise ValueError("setup failed") + + monkeypatch.setattr(GrassSessionManager, "_select_context", select_context) monkeypatch.setattr( "itzi.grass.session.subprocess.check_output", lambda *_args, **_kwargs: "/tmp" ) @@ -69,11 +89,19 @@ def test_manager_activates_and_finishes_created_session(monkeypatch) -> None: ) manager = GrassSessionManager(grass_params) - manager.open() + if setup_fails: + with pytest.raises(ValueError, match="setup failed"): + manager.open() + assert session.finished + else: + manager.open() + manager.open() + manager.close() manager.close() - assert os.environ["GISRC"] == "/tmp/worker-gisrc" + assert "GISRC" not in os.environ assert session.finished + assert restored == ["restored"] def test_temporal_initialization_once_per_active_session(monkeypatch) -> None: @@ -83,6 +111,8 @@ def test_temporal_initialization_once_per_active_session(monkeypatch) -> None: temporal = ModuleType("grass.temporal") temporal.__dict__["init"] = lambda **_: calls.append("init") temporal.__dict__["set_raise_on_error"] = lambda _: calls.append("temporal") + temporal.__dict__["stop_subprocesses"] = lambda: calls.append("stop") + temporal.__dict__["core"] = SimpleNamespace() package = ModuleType("grass") package.__dict__["script"] = script package.__dict__["temporal"] = temporal @@ -97,10 +127,56 @@ def test_temporal_initialization_once_per_active_session(monkeypatch) -> None: ) monkeypatch.setenv("GISRC", "/tmp/first-gisrc") - GrassSessionManager(GrassParams()).open() + GrassSessionManager.ensure_temporal_initialized() GrassSessionManager.ensure_temporal_initialized() assert calls.count("init") == 1 monkeypatch.setenv("GISRC", "/tmp/second-gisrc") - GrassSessionManager(GrassParams()).open() + GrassSessionManager.ensure_temporal_initialized() assert calls.count("init") == 2 + assert calls.count("stop") == 2 + + +def test_attached_cleanup_runs_all_callbacks_after_restoration_error(monkeypatch) -> None: + calls = [] + monkeypatch.setattr("itzi.grass.session._acquire_session", lambda *_: None) + monkeypatch.setattr("itzi.grass.session.ensure_min_version", lambda: None) + monkeypatch.setattr(GrassSessionManager, "ensure_temporal_initialized", lambda self: None) + + def fail(): + calls.append("failure") + raise RuntimeError("restore failed") + + def select_context(self, stack): + stack.callback(calls.append, "last") + stack.callback(fail) + stack.callback(calls.append, "first") + + monkeypatch.setattr(GrassSessionManager, "_select_context", select_context) + manager = GrassSessionManager(GrassParams()) + manager.open() + with pytest.raises(RuntimeError, match="restore failed"): + manager.close() + manager.close() + assert calls == ["first", "failure", "last"] + + +def test_unsupported_version_unwinds_owned_session_before_context_setup(monkeypatch) -> None: + calls = [] + + def acquire(params, stack): + stack.callback(calls.append, "finished") + + def version(): + raise RuntimeError("requires at least GRASS 8.5") + + monkeypatch.setattr("itzi.grass.session._acquire_session", acquire) + monkeypatch.setattr("itzi.grass.session.ensure_min_version", version) + monkeypatch.setattr( + GrassSessionManager, "_select_context", lambda *_: pytest.fail("context setup") + ) + manager = GrassSessionManager(GrassParams(region="missing")) + with pytest.raises(RuntimeError, match="requires at least GRASS 8.5"): + manager.open() + manager.close() + assert calls == ["finished"] From ce4038b174371280a1082fb4757a34cb65dff860 Mon Sep 17 00:00:00 2001 From: Laurent Courty Date: Wed, 28 Oct 2026 13:26:06 -0600 Subject: [PATCH 3/3] explicit dataclass for context --- src/itzi/grass/session.py | 35 ++++++++++++++++++++++++++++------- 1 file changed, 28 insertions(+), 7 deletions(-) diff --git a/src/itzi/grass/session.py b/src/itzi/grass/session.py index 190509f0..931c513b 100644 --- a/src/itzi/grass/session.py +++ b/src/itzi/grass/session.py @@ -27,9 +27,6 @@ import itzi.messenger as msgr -_initialized_temporal_session: tuple[int, str, tuple[str, str, str]] | None = None -_synchronized_session: tuple[int, str, tuple[str, str, str]] | None = None - if TYPE_CHECKING: from grass.lib.gis import Cell_head from grass.script.setup import SessionHandle @@ -39,7 +36,25 @@ MIN_GRASS_VERSION = (8, 5) -@dataclass(frozen=True) +@dataclass(frozen=True, slots=True) +class _GrassContext: + grassdata: str + location: str + mapset: str + + +@dataclass(frozen=True, slots=True) +class _SessionIdentity: + pid: int + gisrc: str + context: _GrassContext + + +_initialized_temporal_session: _SessionIdentity | None = None +_synchronized_session: _SessionIdentity | None = None + + +@dataclass(frozen=True, slots=True) class GrassParams: """Parameters for GRASS GIS session.""" @@ -165,7 +180,9 @@ def _refresh_session() -> None: libraster.Rast_unset_window() context = GrassSessionManager.current_context() assert context is not None - _synchronized_session = (os.getpid(), os.environ["GISRC"], context) + _synchronized_session = _SessionIdentity( + pid=os.getpid(), gisrc=os.environ["GISRC"], context=_GrassContext(*context) + ) def _restore_windows(gis: Cell_head, window: Cell_head) -> None: @@ -222,7 +239,9 @@ def ensure_temporal_initialized() -> None: import grass.temporal as tgis gscript.set_raise_on_error(True) - session = (os.getpid(), os.environ["GISRC"], context) + session = _SessionIdentity( + pid=os.getpid(), gisrc=os.environ["GISRC"], context=_GrassContext(*context) + ) if _initialized_temporal_session != session: _stop_temporal() # Python's forkserver retains the first session's environment for later helpers. @@ -263,7 +282,9 @@ def _select_context(self, contexts: ExitStack) -> None: active = self.current_context() if active is None: msgr.fatal("Unable to determine the active GRASS context") - identity = (os.getpid(), os.environ["GISRC"], active) + identity = _SessionIdentity( + pid=os.getpid(), gisrc=os.environ["GISRC"], context=_GrassContext(*active) + ) if _synchronized_session != identity: _refresh_session() gis = libgis.Cell_head()