t.stac: fixed issue with download and import asset potentially overwhelming network.

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()

src/temporal/t.stac/libstac/testsuite/test_staclib.py

@@ -1,14 +1,16 @@
 #!/usr/bin/env python3
 
 import io
+import os
 import sys
 import json
 import grass.script as gs
 from grass.gunittest.case import TestCase
 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()