diff --git a/docs/history.rst b/docs/history.rst index 38747916..ee7a6cb9 100644 --- a/docs/history.rst +++ b/docs/history.rst @@ -3,6 +3,7 @@ History Latest ------ +- BUG: Close cached raster files before interpreter shutdown (issue #929) 0.23.0 ------ diff --git a/rioxarray/_io.py b/rioxarray/_io.py index 2ac62f07..f653bcf2 100644 --- a/rioxarray/_io.py +++ b/rioxarray/_io.py @@ -6,12 +6,14 @@ Source file: https://github.com/pydata/xarray/blob/1d7bcbdc75b6d556c04e2c7d7a042e4379e15303/xarray/backends/rasterio_.py # noqa """ # pylint: disable=too-many-lines +import atexit import contextlib import importlib.metadata import os import re import threading import warnings +import weakref from collections import defaultdict from collections.abc import Hashable, Iterable from typing import Any, Optional, Union @@ -25,7 +27,7 @@ from xarray import Dataset, IndexVariable from xarray.backends.common import BackendArray from xarray.backends.file_manager import CachingFileManager, FileManager -from xarray.backends.locks import SerializableLock +from xarray.backends.locks import SerializableLock, acquire from xarray.coding import times, variables from xarray.core import indexing from xarray.core.dataarray import DataArray @@ -43,6 +45,38 @@ # TODO: should this be GDAL_LOCK instead? RASTERIO_LOCK = SerializableLock() NO_LOCK = contextlib.nullcontext() +# xarray disables CachingFileManager.__del__ during interpreter shutdown. +# Track the open files instead of their managers because concurrent manager +# finalizers can leave files in xarray's cache after the managers are gone. +_LIVE_RASTERIO_FILES: weakref.WeakKeyDictionary[Any, Any] = weakref.WeakKeyDictionary() +_LIVE_RASTERIO_FILES_LOCK = threading.Lock() + + +@atexit.register +def _close_live_rasterio_files(): + if not _LIVE_RASTERIO_FILES_LOCK.acquire( # pylint: disable=consider-using-with + blocking=False + ): + return + try: + live_files = list(_LIVE_RASTERIO_FILES.items()) + finally: + _LIVE_RASTERIO_FILES_LOCK.release() + + for riods, lock in live_files: + if acquire(lock, blocking=False): + try: + riods.close() + finally: + lock.release() + + +def _acquire_rasterio(manager, lock, needs_lock=True): + riods = manager.acquire(needs_lock=needs_lock) + if isinstance(manager, CachingFileManager): + with _LIVE_RASTERIO_FILES_LOCK: + _LIVE_RASTERIO_FILES[riods] = lock + return riods def _ensure_warped_vrt(riods, vrt_params): @@ -312,7 +346,7 @@ def __init__( self.mask_and_scale = mask_and_scale # cannot save riods as an attribute: this would break pickleability - riods = _ensure_warped_vrt(manager.acquire(), vrt_params) + riods = _ensure_warped_vrt(_acquire_rasterio(manager, lock), vrt_params) self.vrt_params = vrt_params self._shape = (riods.count, riods.height, riods.width) self._dtype = None @@ -430,7 +464,8 @@ def _getitem(self, key): else: with self.lock: riods = _ensure_warped_vrt( - self.manager.acquire(needs_lock=False), self.vrt_params + _acquire_rasterio(self.manager, self.lock, needs_lock=False), + self.vrt_params, ) out = riods.read(band_key, window=window, masked=self.masked) if self._unsigned_dtype is not None: @@ -1137,7 +1172,7 @@ def open_rasterio( ) else: manager = URIManager(file_opener, filename, mode="r", kwargs=open_kwargs) - riods = manager.acquire() + riods = _acquire_rasterio(manager, lock) captured_warnings = rio_warnings.copy() # raise the NotGeoreferencedWarning if applicable diff --git a/test/integration/test_integration__io.py b/test/integration/test_integration__io.py index 60fe9a44..bd5156d5 100644 --- a/test/integration/test_integration__io.py +++ b/test/integration/test_integration__io.py @@ -5,6 +5,7 @@ import os import pickle import shutil +import subprocess import sys import tempfile import warnings @@ -54,6 +55,93 @@ def _assert_tmmx_source(source): ) +def _assert_clean_rasterio_shutdown(script): + raster_path = os.path.abspath(os.path.join(TEST_INPUT_DATA_DIR, "2d_test.tif")) + result = subprocess.run( + [sys.executable, "-c", script, raster_path], + capture_output=True, + check=True, + text=True, + timeout=15, + ) + + assert result.stderr == "" + + +def test_open_rasterio_closes_files_at_interpreter_shutdown(): + _assert_clean_rasterio_shutdown( + """ +import rioxarray +import sys + +for _ in range(5): + rioxarray.open_rasterio(sys.argv[1]) +""" + ) + + +def test_failed_open_rasterio_closes_files_at_interpreter_shutdown(): + _assert_clean_rasterio_shutdown( + """ +import rioxarray +import rioxarray._io +import sys + +def raise_dtype_error(_): + raise RuntimeError("metadata failure") + +rioxarray._io._rasterio_to_numpy_dtype = raise_dtype_error +try: + rioxarray.open_rasterio(sys.argv[1]) +except RuntimeError: + saved_exception = sys.exc_info() + +assert saved_exception[0] is RuntimeError +""" + ) + + +def test_open_rasterio_shutdown_does_not_wait_for_locked_files(): + _assert_clean_rasterio_shutdown( + """ +import rioxarray +import rioxarray._io +import sys +import threading + +raster = rioxarray.open_rasterio(sys.argv[1]) +lock_acquired = threading.Event() + +def hold_file_lock(): + with rioxarray._io.RASTERIO_LOCK: + lock_acquired.set() + threading.Event().wait() + +threading.Thread(target=hold_file_lock, daemon=True).start() +assert lock_acquired.wait(timeout=5) +""" + ) + + +def test_open_rasterio_closes_orphaned_cached_files_at_shutdown(): + _assert_clean_rasterio_shutdown( + """ +from concurrent.futures import ThreadPoolExecutor +import gc +import rioxarray +import sys + +def open_and_read(_): + raster = rioxarray.open_rasterio(sys.argv[1]) + return raster.isel(x=0, y=0).values.tolist() + +with ThreadPoolExecutor(max_workers=16) as pool: + list(pool.map(open_and_read, range(100))) +gc.collect() +""" + ) + + @pytest.mark.parametrize( "subdataset, variable, group, match", [