Repository: voxel51/fiftyone

Issue:
[BUG] MongoDB timeout when iterating through large datasets
Iterating through a detection dataset with ~40K images resulted in a pymongo error after roughly 1 hour.

``` shell
---------------------------------------------------------------------------
CursorNotFound                            Traceback (most recent call last)
<ipython-input-13-40dc943154af> in <module>
      5 header = "ImageID,Source,LabelName,Confidence,XMin,XMax,YMin,YMax,IsOccluded,IsTruncated,IsGroupOf,IsDepiction,IsInside"
      6 lines = [header]
----> 7 for sample in eval_view:
      8     img_id = sample.open_images_id
      9     for det in sample[gt_det_field].detections:
​
~/work/fiftyone/fiftyone/fiftyone/core/view.py in iter_samples(self)
    254         filtered_fields = self._get_filtered_fields()
    255 
--> 256         for d in self._aggregate(detach_frames=True):
    257             try:
    258                 doc = self._dataset._sample_dict_to_doc(d)
​
~/venvs/oi-test/lib/python3.6/site-packages/pymongo/command_cursor.py in next(self)
    257         # Block until a document is returnable.
    258         while self.alive:
--> 259             doc = self._try_next(True)
    260             if doc is not None:
    261                 return doc
​
~/venvs/oi-test/lib/python3.6/site-packages/pymongo/command_cursor.py in _try_next(self, get_more_allowed)
    268         """Advance the cursor blocking for at most one getMore command."""
    269         if not len(self.__data) and not self.__killed and get_more_allowed:
--> 270             self._refresh()
    271         if len(self.__data):
    272             coll = self.__collection
​
~/venvs/oi-test/lib/python3.6/site-packages/pymongo/command_cursor.py in _refresh(self)
    204                                     self.__collection.database.client,
    205                                     self.__max_await_time_ms,
--> 206                                     False))
    207         else:  # Cursor id is zero nothing else to return
    208             self.__killed = True
​
~/venvs/oi-test/lib/python3.6/site-packages/pymongo/command_cursor.py in __send_message(self, operation)
    138         try:
    139             response = client._run_operation_with_response(
--> 140                 operation, self._unpack_response, address=self.__address)
    141         except OperationFailure:
    142             kill()
​
~/venvs/oi-test/lib/python3.6/site-packages/pymongo/mongo_client.py in _run_operation_with_response(self, operation, unpack_res, exhaust, address)
   1370             address=address,
   1371             retryable=isinstance(operation, message._Query),
-> 1372             exhaust=exhaust)
   1373 
   1374     def _retry_with_session(self, retryable, func, session, bulk):
​
~/venvs/oi-test/lib/python3.6/site-packages/pymongo/mongo_client.py in _retryable_read(self, func, read_pref, session, address, retryable, exhaust)
   1469                         # not support retryable reads, raise the last error.
   1470                         raise last_error
-> 1471                     return func(session, server, sock_info, slave_ok)
   1472             except ServerSelectionTimeoutError:
   1473                 if retrying:
​
~/venvs/oi-test/lib/python3.6/site-packages/pymongo/mongo_client.py in _cmd(session, server, sock_info, slave_ok)
   1364                 self._event_listeners,
   1365                 exhaust,
-> 1366                 unpack_res)
   1367 
   1368         return self._retryable_read(
​
~/venvs/oi-test/lib/python3.6/site-packages/pymongo/server.py in run_operation_with_response(self, sock_info, operation, set_slave_okay, listeners, exhaust, unpack_res)
    135                     first, operation.session)
    136                 _check_command_response(
--> 137                     first, sock_info.max_wire_version)
    138         except Exception as exc:
    139             if publish:
​
~/venvs/oi-test/lib/python3.6/site-packages/pymongo/helpers.py in _check_command_response(response, max_wire_version, msg, allowable_errors, parse_write_concern_error)
    162             elif code == 43:
    163                 raise CursorNotFound(errmsg, code, response,
--> 164                                      max_wire_version)
    165 
    166             msg = msg or "%s"
​
CursorNotFound: cursor id 8666854562855358436 not found, full error: {'ok': 0.0, 'errmsg': 'cursor id 8666854562855358436 not found', 'code': 43, 'codeName': 'CursorNotFound'}
```


Discussion:

Resolving patch:
diff --git a/fiftyone/core/dataset.py b/fiftyone/core/dataset.py
--- a/fiftyone/core/dataset.py
+++ b/fiftyone/core/dataset.py
@@ -20,7 +20,7 @@
 from deprecated import deprecated
 import mongoengine.errors as moe
 from pymongo import UpdateMany, UpdateOne
-from pymongo.errors import BulkWriteError
+from pymongo.errors import CursorNotFound, BulkWriteError
 
 import eta.core.serial as etas
 import eta.core.utils as etau
@@ -1254,12 +1254,29 @@ def iter_samples(self):
         Returns:
             an iterator over :class:`fiftyone.core.sample.Sample` instances
         """
-        for d in self._aggregate(detach_frames=True):
-            doc = self._sample_dict_to_doc(d)
-            sample = fos.Sample.from_doc(doc, dataset=self)
+        pipeline = self._pipeline(detach_frames=True)
 
+        for sample in self._iter_samples(pipeline):
             yield sample
 
+    def _iter_samples(self, pipeline):
+        index = 0
+
+        try:
+            for d in foo.aggregate(self._sample_collection, pipeline):
+                doc = self._sample_dict_to_doc(d)
+                sample = fos.Sample.from_doc(doc, dataset=self)
+                index += 1
+                yield sample
+
+        except CursorNotFound:
+            # The cursor has timed out so we yield from a new one after
+            # skipping to the last offset
+
+            pipeline.append({"$skip": index})
+            for sample in self._iter_samples(pipeline):
+                yield sample
+
     def add_sample(self, sample, expand_schema=True):
         """Adds the given sample to the dataset.
 
diff --git a/fiftyone/core/view.py b/fiftyone/core/view.py
--- a/fiftyone/core/view.py
+++ b/fiftyone/core/view.py
@@ -10,6 +10,7 @@
 import numbers
 
 from bson import ObjectId
+from pymongo.errors import CursorNotFound
 
 import eta.core.utils as etau
 
@@ -262,23 +263,35 @@ def iter_samples(self):
         selected_fields, excluded_fields = self._get_selected_excluded_fields()
         filtered_fields = self._get_filtered_fields()
 
-        for d in self._aggregate(detach_frames=True):
-            try:
-                doc = self._dataset._sample_dict_to_doc(d)
-                sample = sample_cls(
-                    doc,
-                    self,
-                    selected_fields=selected_fields,
-                    excluded_fields=excluded_fields,
-                    filtered_fields=filtered_fields,
-                )
+        index = 0
 
+        try:
+            for d in self._aggregate(detach_frames=True):
+                try:
+                    doc = self._dataset._sample_dict_to_doc(d)
+                    sample = sample_cls(
+                        doc,
+                        self,
+                        selected_fields=selected_fields,
+                        excluded_fields=excluded_fields,
+                        filtered_fields=filtered_fields,
+                    )
+                    index += 1
+                    yield sample
+
+                except Exception as e:
+                    raise ValueError(
+                        "Failed to load sample from the database. This is "
+                        "likely due to an invalid stage in the DatasetView"
+                    ) from e
+
+        except CursorNotFound:
+            # The cursor has timed out so we yield from a new one after
+            # skipping to the last offset
+
+            view = self.skip(index)
+            for sample in view.iter_samples():
                 yield sample
-            except Exception as e:
-                raise ValueError(
-                    "Failed to load sample from the database. This is likely "
-                    "due to an invalid stage in the DatasetView"
-                ) from e
 
     def get_field_schema(
         self, ftype=None, embedded_doc_type=None, include_private=False


Test patch: