Repository: microsoft/torchgeo

Issue:
Parallel data loading doesn't work on macOS/Windows
When trying to to benchmark the data loader, if I run with `num_workers > 0`, it crashes with the following error message:
```
Traceback (most recent call last):
  File "benchmark.py", line 299, in <module>
    main(args)
  File "benchmark.py", line 181, in main
    for i, batch in enumerate(dataloader):
  File "/Users/Adam/.spack/.spack-env/view/lib/python3.8/site-packages/torch/utils/data/dataloader.py", line 521, in __next__
    data = self._next_data()
  File "/Users/Adam/.spack/.spack-env/view/lib/python3.8/site-packages/torch/utils/data/dataloader.py", line 1203, in _next_data
    return self._process_data(data)
  File "/Users/Adam/.spack/.spack-env/view/lib/python3.8/site-packages/torch/utils/data/dataloader.py", line 1229, in _process_data
    data.reraise()
  File "/Users/Adam/.spack/.spack-env/view/lib/python3.8/site-packages/torch/_utils.py", line 425, in reraise
    raise self.exc_type(msg)
ValueError: Caught ValueError in DataLoader worker process 0.
Original Traceback (most recent call last):
  File "/Users/Adam/.spack/.spack-env/view/lib/python3.8/site-packages/torch/utils/data/_utils/worker.py", line 287, in _worker_loop
    data = fetcher.fetch(index)
  File "/Users/Adam/.spack/.spack-env/view/lib/python3.8/site-packages/torch/utils/data/_utils/fetch.py", line 44, in fetch
    data = [self.dataset[idx] for idx in possibly_batched_index]
  File "/Users/Adam/.spack/.spack-env/view/lib/python3.8/site-packages/torch/utils/data/_utils/fetch.py", line 44, in <listcomp>
    data = [self.dataset[idx] for idx in possibly_batched_index]
  File "/Users/Adam/torchgeo/torchgeo/datasets/geo.py", line 719, in __getitem__
    if not query.intersects(self.bounds):
  File "/Users/Adam/torchgeo/torchgeo/datasets/geo.py", line 760, in bounds
    minx = max([ds.bounds[0] for ds in self.datasets])
  File "/Users/Adam/torchgeo/torchgeo/datasets/geo.py", line 760, in <listcomp>
    minx = max([ds.bounds[0] for ds in self.datasets])
  File "/Users/Adam/torchgeo/torchgeo/datasets/geo.py", line 135, in bounds
    return BoundingBox(*self.index.bounds)
  File "/Users/Adam/torchgeo/torchgeo/datasets/utils.py", line 213, in __new__
    raise ValueError(f"Bounding box is invalid: 'minx={minx}' > 'maxx={maxx}'")
ValueError: Bounding box is invalid: 'minx=1.7976931348623157e+308' > 'maxx=-1.7976931348623157e+308'
```
Since this doesn't occur in serial or on Linux, I'm guessing this has something to do with the fact that Python's multiprocessing module switched from `fork` to `spawn` as the default start method on macOS for Python 3.8+.


Discussion:
Also, the fact that this wasn't caught by our unit tests means we need better integration tests. We do test our samplers in parallel, but not with a real `GeoDataset`.
This issue is caused by https://github.com/Toblerity/rtree/issues/87, rtree indices do not pickle correctly:
```console
$ python
Python 3.8.12 (default, Oct 25 2021, 13:41:41) 
[Clang 12.0.0 (clang-1200.0.32.29)] on darwin
Type "help", "copyright", "credits" or "license" for more information.
>>> import pickle
>>> from rtree import index
>>> idx = index.Index()
>>> idx.insert(0, (0, 1, 2, 3))
>>> print(idx)
rtree.index.Index(bounds=[0.0, 1.0, 2.0, 3.0], size=1)
>>> x = pickle.dumps(idx)
>>> y = pickle.loads(x)
>>> print(y)
rtree.index.Index(bounds=[1.7976931348623157e+308, 1.7976931348623157e+308, -1.7976931348623157e+308, -1.7976931348623157e+308], size=0)
```
This also happens on Windows and is due to the whole fork vs spawn issue.

Resolving patch:
diff --git a/torchgeo/datasets/geo.py b/torchgeo/datasets/geo.py
--- a/torchgeo/datasets/geo.py
+++ b/torchgeo/datasets/geo.py
@@ -172,6 +172,41 @@ def __str__(self) -> str:
     bbox: {self.bounds}
     size: {len(self)}"""
 
+    # NOTE: This hack should be removed once the following issue is fixed:
+    # https://github.com/Toblerity/rtree/issues/87
+
+    def __getstate__(
+        self,
+    ) -> Tuple[
+        Dict[Any, Any],
+        List[Tuple[int, Tuple[float, float, float, float, float, float], str]],
+    ]:
+        """Define how instances are pickled.
+
+        Returns:
+            the state necessary to unpickle the instance
+        """
+        objects = self.index.intersection(self.index.bounds, objects=True)
+        tuples = [(item.id, item.bounds, item.object) for item in objects]
+        return self.__dict__, tuples
+
+    def __setstate__(
+        self,
+        state: Tuple[
+            Dict[Any, Any],
+            List[Tuple[int, Tuple[float, float, float, float, float, float], str]],
+        ],
+    ) -> None:
+        """Define how to unpickle an instance.
+
+        Args:
+            state: the state of the instance when it was pickled
+        """
+        attrs, tuples = state
+        self.__dict__.update(attrs)
+        for item in tuples:
+            self.index.insert(*item)
+
     @property
     def bounds(self) -> BoundingBox:
         """Bounds of the index.


Test patch:
diff --git a/tests/datasets/test_geo.py b/tests/datasets/test_geo.py
--- a/tests/datasets/test_geo.py
+++ b/tests/datasets/test_geo.py
@@ -2,6 +2,7 @@
 # Licensed under the MIT License.
 
 import os
+import pickle
 from pathlib import Path
 from typing import Dict
 
@@ -121,6 +122,14 @@ def test_str(self, dataset: GeoDataset) -> None:
         assert "bbox: BoundingBox" in out
         assert "size: 1" in out
 
+    def test_picklable(self, dataset: GeoDataset) -> None:
+        x = pickle.dumps(dataset)
+        y = pickle.loads(x)
+        assert dataset.crs == y.crs
+        assert dataset.res == y.res
+        assert len(dataset) == len(y)
+        assert dataset.bounds == y.bounds
+
     def test_abstract(self) -> None:
         with pytest.raises(TypeError, match="Can't instantiate abstract class"):
             GeoDataset()  # type: ignore[abstract]