diff --git a/src/temporal/t.stac/libstac/staclib.py b/src/temporal/t.stac/libstac/staclib.py index 1121e02cde..bbbc0636ce 100644 --- a/src/temporal/t.stac/libstac/staclib.py +++ b/src/temporal/t.stac/libstac/staclib.py @@ -25,7 +25,7 @@ from grass.exceptions import CalledModuleError from grass.pygrass.vector import VectorTopo from grass.pygrass.vector.geometry import Point, Centroid, Boundary -from concurrent.futures import ThreadPoolExecutor +from concurrent.futures import ProcessPoolExecutor from grass.jupyter.reprojection_renderer import ReprojectionRenderer @@ -848,6 +848,12 @@ def create_metadata_vector(vector, metadata): def import_grass_raster(params): assets, resample_method, extent, resolution, resolution_value, memory = params + # Limit GDAL to 2 concurrent HTTP connections per process to avoid + # overwhelming the network when multiple processes run in parallel. + os.environ.setdefault("GDAL_HTTP_MAX_RETRY", "3") + os.environ.setdefault("GDAL_HTTP_RETRY_DELAY", "2") + os.environ.setdefault("CPL_VSIL_CURL_CACHE_SIZE", "0") + os.environ.setdefault("GDAL_MAX_CONNECTIONS", "2") sys.stdout.write(f"Downloading Asset: {assets}\n") input_url = check_url_type(assets["href"]) sys.stdout.write(f"Import Url: {input_url}\n") @@ -870,6 +876,9 @@ def import_grass_raster(params): gs.fatal(_("Error importing raster: {}").format(e.stderr)) +MAX_CONCURRENT_DOWNLOADS = 8 + + def download_assets( assets, resample_method, @@ -881,11 +890,6 @@ def download_assets( ): """Downloads a list of images from the given URLs to the given filenames.""" number_of_assets = len(assets) - resample_extent_list = [resample_extent] * number_of_assets - resolution_list = [resolution] * number_of_assets - resolution_value_list = [resolution_value] * number_of_assets - resample_method_list = [resample_method] * number_of_assets - memory_list = [memory] * number_of_assets max_cpus = os.cpu_count() - 1 if nprocs > max_cpus: gs.warning( @@ -895,29 +899,36 @@ def download_assets( ) nprocs = max_cpus - def execute_import_grass_raster(pbar=None): - with ThreadPoolExecutor(max_workers=nprocs) as executor: - try: - for _a in executor.map( - import_grass_raster, - zip( - assets, - resample_method_list, - resample_extent_list, - resolution_list, - resolution_value_list, - memory_list, - ), - ): - if pbar: - pbar.update(1) - except Exception as e: - gs.fatal(_("Error importing raster: {}").format(str(e))) + if nprocs > MAX_CONCURRENT_DOWNLOADS: + gs.warning( + _( + "Capping concurrent downloads at {max} to avoid overwhelming the network " + "(requested {nprocs})." + ).format(max=MAX_CONCURRENT_DOWNLOADS, nprocs=nprocs) + ) + nprocs = MAX_CONCURRENT_DOWNLOADS - tqdm = _import_tqdm(False) - if tqdm is None: + params = [ + (asset, resample_method, resample_extent, resolution, resolution_value, memory) + for asset in assets + ] + + tqdm_cls = _import_tqdm(False) + if tqdm_cls is None: gs.warning(_("tqdm module not found. Progress bar will not be displayed.")) - execute_import_grass_raster() + + def _run(pbar=None): + # ProcessPoolExecutor (not threads) is required: grass.script relies on + # process-global state (GISRC, WIND_OVERRIDE, GIS_LOCK, computational + # region) and the per-worker GDAL_* env vars set in import_grass_raster. + # Concurrent r.import calls in the same process race on that state. + with ProcessPoolExecutor(max_workers=nprocs) as executor: + for _ in executor.map(import_grass_raster, params): + if pbar: + pbar.update(1) + + if tqdm_cls is not None: + with tqdm_cls(total=number_of_assets, desc="Downloading assets") as pbar: + _run(pbar) else: - with tqdm(total=number_of_assets, desc="Downloading assets") as pbar: - execute_import_grass_raster(pbar) + _run() diff --git a/src/temporal/t.stac/libstac/testsuite/test_staclib.py b/src/temporal/t.stac/libstac/testsuite/test_staclib.py index 3e05fb0d2e..6f0b59a901 100644 --- a/src/temporal/t.stac/libstac/testsuite/test_staclib.py +++ b/src/temporal/t.stac/libstac/testsuite/test_staclib.py @@ -1,6 +1,7 @@ #!/usr/bin/env python3 import io +import os import sys import json import grass.script as gs @@ -8,7 +9,8 @@ from grass.gunittest.main import test from grass.pygrass.utils import get_lib_path from grass.pygrass.vector.geometry import Point -from unittest.mock import patch, MagicMock +from grass.exceptions import CalledModuleError +from unittest.mock import patch, MagicMock, call import base64 import tempfile @@ -480,5 +482,259 @@ def test_wgs84_geojson_from_vector_raises_exception( mock_fatal.assert_called_once() +class TestImportGrassRaster(TestCase): + """Tests for import_grass_raster worker function.""" + + def _make_params(self, href="http://example.com/asset0.tif", file_name="asset_0"): + return ( + {"href": href, "file_name": file_name}, + "bilinear", + "region", + "value", + 30.0, + 300, + ) + + @patch("staclib.gs.parse_command") + @patch("staclib.check_url_type", return_value="http://example.com/asset0.tif") + def test_calls_r_import_with_correct_args(self, mock_check_url, mock_parse): + libstac.import_grass_raster(self._make_params()) + mock_parse.assert_called_once_with( + "r.import", + input="http://example.com/asset0.tif", + output="asset_0", + resample="bilinear", + extent="region", + resolution="value", + resolution_value=30.0, + title="asset_0", + memory=300, + quiet=True, + ) + + @patch("staclib.gs.parse_command") + @patch("staclib.check_url_type", return_value="http://example.com/asset0.tif") + def test_check_url_type_called_with_href(self, mock_check_url, mock_parse): + params = self._make_params() + libstac.import_grass_raster(params) + mock_check_url.assert_called_once_with("http://example.com/asset0.tif") + + @patch("staclib.gs.parse_command") + @patch("staclib.check_url_type", return_value="http://example.com/asset0.tif") + def test_sets_gdal_env_vars(self, mock_check_url, mock_parse): + for key in ( + "GDAL_HTTP_MAX_RETRY", + "GDAL_HTTP_RETRY_DELAY", + "CPL_VSIL_CURL_CACHE_SIZE", + "GDAL_MAX_CONNECTIONS", + ): + os.environ.pop(key, None) + + libstac.import_grass_raster(self._make_params()) + + self.assertEqual(os.environ.get("GDAL_HTTP_MAX_RETRY"), "3") + self.assertEqual(os.environ.get("GDAL_HTTP_RETRY_DELAY"), "2") + self.assertEqual(os.environ.get("CPL_VSIL_CURL_CACHE_SIZE"), "0") + self.assertEqual(os.environ.get("GDAL_MAX_CONNECTIONS"), "2") + + @patch("staclib.gs.parse_command") + @patch("staclib.check_url_type", return_value="http://example.com/asset0.tif") + def test_gdal_env_vars_not_overwritten_if_already_set( + self, mock_check_url, mock_parse + ): + os.environ["GDAL_MAX_CONNECTIONS"] = "10" + try: + libstac.import_grass_raster(self._make_params()) + self.assertEqual(os.environ.get("GDAL_MAX_CONNECTIONS"), "10") + finally: + os.environ.pop("GDAL_MAX_CONNECTIONS", None) + + @patch("staclib.gs.fatal") + @patch("staclib.check_url_type", return_value="http://example.com/asset0.tif") + @patch("staclib.gs.parse_command") + def test_called_module_error_calls_fatal( + self, mock_parse, mock_check_url, mock_fatal + ): + err = CalledModuleError("r.import", [], 1, "r.import failed") + mock_parse.side_effect = err + mock_fatal.side_effect = SystemExit + + with self.assertRaises(SystemExit): + libstac.import_grass_raster(self._make_params()) + + mock_fatal.assert_called_once_with("Error importing raster: r.import failed") + + +class TestDownloadAssets(TestCase): + """Tests for download_assets.""" + + def _make_assets(self, n=2): + return [ + {"href": f"http://example.com/asset{i}.tif", "file_name": f"asset_{i}"} + for i in range(n) + ] + + def _setup_pool_mock(self, mock_pool_cls, n_assets, map_results=None): + """Wire mock_pool_cls to return a mock executor whose map() yields + one result per asset.""" + if map_results is None: + map_results = [None] * n_assets + + mock_executor = MagicMock() + mock_executor.map.return_value = iter(map_results) + mock_pool_cls.return_value.__enter__ = MagicMock(return_value=mock_executor) + mock_pool_cls.return_value.__exit__ = MagicMock(return_value=False) + return mock_executor + + # ------------------------------------------------------------------ + # nprocs capping + # ------------------------------------------------------------------ + + @patch("staclib.gs.warning") + @patch("staclib.ProcessPoolExecutor") + @patch("os.cpu_count", return_value=4) + def test_nprocs_capped_at_cpu_count_minus_one( + self, mock_cpu, mock_pool_cls, mock_warning + ): + assets = self._make_assets(2) + self._setup_pool_mock(mock_pool_cls, 2) + + # nprocs=10 exceeds cpu_count-1=3 + libstac.download_assets(assets, "bilinear", "region", "value", 30.0, nprocs=10) + + mock_warning.assert_any_call( + "Number of processes {nprocs} is greater than the number of CPUs {max_cpus}." + ) + mock_pool_cls.assert_called_once_with(max_workers=3) + + @patch("staclib.gs.warning") + @patch("staclib.ProcessPoolExecutor") + @patch("os.cpu_count", return_value=64) + def test_nprocs_capped_at_max_concurrent_downloads( + self, mock_cpu, mock_pool_cls, mock_warning + ): + assets = self._make_assets(2) + self._setup_pool_mock(mock_pool_cls, 2) + + # nprocs=32 exceeds MAX_CONCURRENT_DOWNLOADS=8 but not cpu_count-1=63 + libstac.download_assets(assets, "bilinear", "region", "value", 30.0, nprocs=32) + + mock_warning.assert_any_call( + "Capping concurrent downloads at {max} to avoid overwhelming the network " + "(requested {nprocs}).".format( + max=libstac.MAX_CONCURRENT_DOWNLOADS, nprocs=32 + ) + ) + mock_pool_cls.assert_called_once_with( + max_workers=libstac.MAX_CONCURRENT_DOWNLOADS + ) + + # ------------------------------------------------------------------ + # executor.map behaviour + # ------------------------------------------------------------------ + + @patch("staclib.ProcessPoolExecutor") + def test_executor_created_with_nprocs(self, mock_pool_cls): + assets = self._make_assets(2) + self._setup_pool_mock(mock_pool_cls, 2) + + libstac.download_assets(assets, "bilinear", "region", "value", 30.0, nprocs=3) + + mock_pool_cls.assert_called_once_with(max_workers=3) + + @patch("staclib.ProcessPoolExecutor") + def test_map_called_once_with_all_params(self, mock_pool_cls): + n = 4 + assets = self._make_assets(n) + mock_executor = self._setup_pool_mock(mock_pool_cls, n) + + libstac.download_assets(assets, "bilinear", "region", "value", 30.0, nprocs=2) + + self.assertEqual(mock_executor.map.call_count, 1) + fn_arg, params_arg = mock_executor.map.call_args.args + self.assertIs(fn_arg, libstac.import_grass_raster) + self.assertEqual(len(params_arg), n) + + @patch("staclib.ProcessPoolExecutor") + def test_correct_params_passed_to_map(self, mock_pool_cls): + assets = self._make_assets(2) + mock_executor = self._setup_pool_mock(mock_pool_cls, 2) + + libstac.download_assets( + assets, "bilinear", "region", "value", 30.0, memory=512, nprocs=2 + ) + + _fn, params_arg = mock_executor.map.call_args.args + for i, asset in enumerate(assets): + self.assertEqual( + params_arg[i], + (asset, "bilinear", "region", "value", 30.0, 512), + ) + + @patch("staclib.ProcessPoolExecutor") + def test_empty_assets_map_with_empty_list(self, mock_pool_cls): + mock_executor = self._setup_pool_mock(mock_pool_cls, 0) + + libstac.download_assets([], "bilinear", "region", "value", 30.0, nprocs=2) + + _fn, params_arg = mock_executor.map.call_args.args + self.assertEqual(params_arg, []) + + @patch("staclib.ProcessPoolExecutor") + def test_exception_from_worker_propagates(self, mock_pool_cls): + assets = self._make_assets(2) + + def raising_iter(): + yield None + raise RuntimeError("Download failed") + + mock_executor = MagicMock() + mock_executor.map.return_value = raising_iter() + mock_pool_cls.return_value.__enter__ = MagicMock(return_value=mock_executor) + mock_pool_cls.return_value.__exit__ = MagicMock(return_value=False) + + with self.assertRaises(RuntimeError, msg="Download failed"): + libstac.download_assets( + assets, "bilinear", "region", "value", 30.0, nprocs=1 + ) + + # ------------------------------------------------------------------ + # tqdm integration + # ------------------------------------------------------------------ + + @patch("staclib._import_tqdm") + @patch("staclib.ProcessPoolExecutor") + def test_pbar_updated_once_per_asset(self, mock_pool_cls, mock_import_tqdm): + n = 3 + assets = self._make_assets(n) + self._setup_pool_mock(mock_pool_cls, n) + + mock_pbar = MagicMock() + mock_tqdm_cls = MagicMock() + mock_tqdm_cls.return_value.__enter__ = MagicMock(return_value=mock_pbar) + mock_tqdm_cls.return_value.__exit__ = MagicMock(return_value=False) + mock_import_tqdm.return_value = mock_tqdm_cls + + libstac.download_assets(assets, "bilinear", "region", "value", 30.0, nprocs=2) + + mock_tqdm_cls.assert_called_once_with(total=n, desc="Downloading assets") + self.assertEqual(mock_pbar.update.call_count, n) + + @patch("staclib.gs.warning") + @patch("staclib._import_tqdm", return_value=None) + @patch("staclib.ProcessPoolExecutor") + def test_warning_and_runs_without_tqdm( + self, mock_pool_cls, mock_import_tqdm, mock_warning + ): + assets = self._make_assets(2) + self._setup_pool_mock(mock_pool_cls, 2) + + libstac.download_assets(assets, "bilinear", "region", "value", 30.0, nprocs=1) + + mock_warning.assert_any_call( + "tqdm module not found. Progress bar will not be displayed." + ) + + if __name__ == "__main__": test()