Allow passing device_id=None to determine CUDA device id automatically

Make `device_id` optional (default to None), so that it can be inferred automatically by `cupy.cuda.runtime.getDevice`. Also added a unit test using xr.open_mfdataset.

Author: weiji14

cupy_xarray/cog3pio.py

@@ -17,7 +17,7 @@ class Cog3pioBackendEntrypoint(BackendEntrypoint):
     """
     Xarray backend to read GeoTIFF files using 'cog3pio' engine.
 
-    When using :py:func:`xarray.open_dataarray` with ``engine="cog3pio"``, the
+    When using :py:func:`xarray.open_dataarray` with ``engine="cog3pio"``, the optional
     ``device_id`` parameter can be set to the CUDA GPU id to do the decoding on.
 
     Examples
@@ -47,7 +47,7 @@ def open_dataset(  # type: ignore[override]
         filename_or_obj: str,
         *,
         drop_variables: str | Iterable[str] | None = None,
-        device_id: int,
+        device_id: int | None = None,
         # other backend specific keyword arguments
         # `chunks` and `cache` DO NOT go here, they are handled by xarray
         mask_and_scale=None,
@@ -60,14 +60,18 @@ def open_dataset(  # type: ignore[override]
         filename_or_obj : str
             File path or url to a TIFF (.tif) image file that can be read by the
             nvTIFF or image-tiff backend library.
-        device_id : int
-            CUDA device ID on which to place the created cupy array.
+        device_id : int | None
+            CUDA device ID on which to place the created cupy array. Default is None,
+            which means device_id will be inferred via
+            :py:func:`cupy.cuda.runtime.getDevice`.
 
         Returns
         -------
         xarray.Dataset
 
         """
+        if device_id is None:
+            device_id: int = cp.cuda.runtime.getDevice()
 
         with cp.cuda.Stream(ptds=True):
             cog = CudaCogReader(path=filename_or_obj, device_id=device_id)

cupy_xarray/tests/test_cog3pio.py

@@ -11,6 +11,7 @@
 cog3pio = pytest.importorskip("cog3pio")
 
 
+# %%
 def test_entrypoint():
     assert "cog3pio" in xr.backends.list_engines()
 
@@ -32,3 +33,27 @@ def test_xarray_backend_open_dataarray():
         assert da.y.min() == 3490250.0
         assert da.y.max() == 3599950.0
         assert da.dtype == "uint8"
+
+
+def test_xarray_backend_open_mfdataset():
+    """
+    Ensure that passing engine='cog3pio' to xarray.open_mfdataset works to read multiple
+    Cloud-optimized GeoTIFF files from http urls. Also testing that `device_id=None`
+    works.
+    """
+    ds: xr.Dataset = xr.open_mfdataset(
+        paths=[
+            "https://github.com/developmentseed/titiler/raw/1.2.0/src/titiler/mosaic/tests/fixtures/B01.tif",
+            "https://github.com/developmentseed/titiler/raw/1.2.0/src/titiler/mosaic/tests/fixtures/B09.tif",
+        ],
+        engine=Cog3pioBackendEntrypoint,
+        concat_dim="band",
+        combine="nested",
+        device_id=None,
+    )
+    assert ds.sizes == {"band": 2, "y": 183, "x": 183}
+    assert ds.x.min() == 700260.0
+    assert ds.x.max() == 809460.0
+    assert ds.y.min() == 3490500.0
+    assert ds.y.max() == 3599700.0
+    assert ds.raster.dtype == "uint16"