Da item pagination (#263)

* Unit test and large scale test pass

* specific tests pass

* remove unwanted prints

* missing type annotation

* Update nucleus/constants.py

Co-authored-by: Drake Wong <40375132+drakejwong@users.noreply.github.com>

* fix slice tests

Co-authored-by: Ubuntu <diego.ardila@scale.com>
Co-authored-by: Drake Wong <40375132+drakejwong@users.noreply.github.com>

Author: 3 people

conftest.py

@@ -1,10 +1,14 @@
 import os
+from typing import TYPE_CHECKING
 
 import pytest
 
 import nucleus
 from tests.helpers import TEST_DATASET_ITEMS, TEST_DATASET_NAME
 
+if TYPE_CHECKING:
+    from nucleus import NucleusClient
+
 assert "NUCLEUS_PYTEST_API_KEY" in os.environ, (
     "You must set the 'NUCLEUS_PYTEST_API_KEY' environment variable to a valid "
     "Nucleus API key to run the test suite"
@@ -20,12 +24,12 @@ def CLIENT():
 
 
 @pytest.fixture()
-def dataset(CLIENT):
-    ds = CLIENT.create_dataset(TEST_DATASET_NAME)
-    ds.append(TEST_DATASET_ITEMS)
-    yield ds
+def dataset(CLIENT: "NucleusClient"):
+    test_dataset = CLIENT.create_dataset(TEST_DATASET_NAME, is_scene=False)
+    test_dataset.append(TEST_DATASET_ITEMS)
+    yield test_dataset
 
-    CLIENT.delete_dataset(ds.id)
+    CLIENT.delete_dataset(test_dataset.id)
 
 
 @pytest.fixture()

nucleus/constants.py

@@ -75,6 +75,7 @@
 JOB_LAST_KNOWN_STATUS_KEY = "job_last_known_status"
 JOB_TYPE_KEY = "job_type"
 JOB_CREATION_TIME_KEY = "job_creation_time"
+LAST_PAGE = "lastPage"
 LABEL_KEY = "label"
 LABELS_KEY = "labels"
 MASK_URL_KEY = "mask_url"
@@ -87,6 +88,8 @@
 NUCLEUS_ENDPOINT = "https://api.scale.com/v1/nucleus"
 NUM_SENSORS_KEY = "num_sensors"
 ORIGINAL_IMAGE_URL_KEY = "original_image_url"
+PAGE_SIZE = "pageSize"
+PAGE_TOKEN = "pageToken"
 P1_KEY = "p1"
 P2_KEY = "p2"
 POINTCLOUD_KEY = "pointcloud"
@@ -97,6 +100,7 @@
 PREDICTIONS_PROCESSED_KEY = "predictions_processed"
 REFERENCE_IDS_KEY = "reference_ids"
 REFERENCE_ID_KEY = "reference_id"
+BACKEND_REFERENCE_ID_KEY = "ref_id"  # TODO(355762): Our backend returns this instead of the "proper" key sometimes.
 REQUEST_ID_KEY = "requestId"
 SCENES_KEY = "scenes"
 SERIALIZED_REQUEST_KEY = "serialized_request"

nucleus/dataset.py

@@ -1,5 +1,5 @@
 import os
-from typing import Any, Dict, List, Optional, Sequence, Union
+from typing import Any, Dict, Iterable, List, Optional, Sequence, Union
 
 import requests
 
@@ -18,6 +18,7 @@
     convert_export_payload,
     format_dataset_item_response,
     format_prediction_response,
+    paginate_generator,
     serialize_and_write_to_presigned_url,
 )
 
@@ -32,6 +33,7 @@
     EMBEDDING_DIMENSION_KEY,
     EMBEDDINGS_URL_KEY,
     EXPORTED_ROWS,
+    ITEMS_KEY,
     KEEP_HISTORY_KEY,
     MESSAGE_KEY,
     NAME_KEY,
@@ -51,7 +53,7 @@
 )
 from .dataset_item_uploader import DatasetItemUploader
 from .deprecation_warning import deprecated
-from .errors import DatasetItemRetrievalError
+from .errors import NucleusAPIError
 from .metadata_manager import ExportMetadataType, MetadataManager
 from .payload_constructor import (
     construct_append_scenes_payload,
@@ -160,25 +162,51 @@ def size(self) -> int:
         dataset_size = DatasetSize.parse_obj(response)
         return dataset_size.count
 
+    def items_generator(self, page_size=100000) -> Iterable[DatasetItem]:
+        """Generator yielding all dataset items in the dataset.
+
+
+        ::
+            sum_example_field = 0
+            for item in dataset.items_generator():
+                sum += item.metadata["example_field"]
+
+        Args:
+            page_size (int, optional): Number of items to return per page. If you are
+                experiencing timeouts while using this generator, you can try lowering
+                the page size.
+
+        Yields:
+            an iterable of DatasetItem objects.
+        """
+        json_generator = paginate_generator(
+            client=self._client,
+            endpoint=f"dataset/{self.id}/itemsPage",
+            result_key=ITEMS_KEY,
+            page_size=page_size,
+        )
+        for item_json in json_generator:
+            yield DatasetItem.from_json(item_json)
+
     @property
     def items(self) -> List[DatasetItem]:
-        """List of all DatasetItem objects in the Dataset."""
-        response = self._client.make_request(
-            {}, f"dataset/{self.id}/datasetItems", requests.get
-        )
-        dataset_items = response.get("dataset_items", None)
-        error = response.get("error", None)
-        constructed_dataset_items = []
-        if dataset_items:
-            for item in dataset_items:
-                image_url = item.get("original_image_url")
-                metadata = item.get("metadata", None)
-                ref_id = item.get("ref_id", None)
-                dataset_item = DatasetItem(image_url, ref_id, metadata)
-                constructed_dataset_items.append(dataset_item)
-        elif error:
-            raise DatasetItemRetrievalError(message=error)
-        return constructed_dataset_items
+        """List of all DatasetItem objects in the Dataset.
+
+        For fetching more than 200k items see :meth:`NucleusDataset.items_generator`.
+        """
+        try:
+            response = self._client.make_request(
+                {}, f"dataset/{self.id}/datasetItems", requests.get
+            )
+        except NucleusAPIError as e:
+            if e.status_code == 503:
+                e.message += "\nThe server timed out while trying to load your items. Please try iterating over dataset.items_generator() instead."
+            raise e
+        dataset_item_jsons = response.get("dataset_items", None)
+        return [
+            DatasetItem.from_json(item_json)
+            for item_json in dataset_item_jsons
+        ]
 
     @property
     def scenes(self) -> List[ScenesListEntry]:

nucleus/dataset_item.py

@@ -7,6 +7,7 @@
 
 from .annotation import Point3D, is_local_path
 from .constants import (
+    BACKEND_REFERENCE_ID_KEY,
     CAMERA_MODEL_KEY,
     CAMERA_PARAMS_KEY,
     CX_KEY,
@@ -290,6 +291,8 @@ def from_json(cls, payload: dict):
         image_url = payload.get(IMAGE_URL_KEY, None) or payload.get(
             ORIGINAL_IMAGE_URL_KEY, None
         )
+        if BACKEND_REFERENCE_ID_KEY in payload:
+            payload[REFERENCE_ID_KEY] = payload[BACKEND_REFERENCE_ID_KEY]
         return cls(
             image_location=image_url,
             pointcloud_location=payload.get(POINTCLOUD_URL_KEY, None),

nucleus/errors.py

@@ -40,28 +40,29 @@ class NucleusAPIError(Exception):
     def __init__(
         self, endpoint, command, requests_response=None, aiohttp_response=None
     ):
-        message = f"Your client is on version {nucleus_client_version}. If you have not recently done so, please make sure you have updated to the latest version of the client by running pip install --upgrade scale-nucleus\n"
+        self.message = f"Your client is on version {nucleus_client_version}. If you have not recently done so, please make sure you have updated to the latest version of the client by running pip install --upgrade scale-nucleus\n"
         if requests_response is not None:
-            message += f"Tried to {command.__name__} {endpoint}, but received {requests_response.status_code}: {requests_response.reason}."
+            self.status_code = requests_response.status_code
+            self.message += f"Tried to {command.__name__} {endpoint}, but received {requests_response.status_code}: {requests_response.reason}."
             if hasattr(requests_response, "text"):
                 if requests_response.text:
-                    message += (
+                    self.message += (
                         f"\nThe detailed error is:\n{requests_response.text}"
                     )
 
         if aiohttp_response is not None:
             status, reason, data = aiohttp_response
-            message += f"Tried to {command.__name__} {endpoint}, but received {status}: {reason}."
+            self.status_code = status
+            self.message += f"Tried to {command.__name__} {endpoint}, but received {status}: {reason}."
             if data:
-                message += f"\nThe detailed error is:\n{data}"
+                self.message += f"\nThe detailed error is:\n{data}"
 
         if any(
-            infra_flake_message in message
+            infra_flake_message in self.message
             for infra_flake_message in INFRA_FLAKE_MESSAGES
         ):
-            message += "\n This likely indicates temporary downtime of the API, please try again in a minute or two"
-
-        super().__init__(message)
+            self.message += "\n This likely indicates temporary downtime of the API, please try again in a minute or two"
+        super().__init__(self.message)
 
 
 class NoAPIKey(Exception):

nucleus/slice.py

@@ -4,13 +4,15 @@
 import requests
 
 from nucleus.annotation import Annotation
-from nucleus.constants import EXPORTED_ROWS
+from nucleus.constants import EXPORTED_ROWS, ITEMS_KEY
 from nucleus.dataset_item import DatasetItem
+from nucleus.errors import NucleusAPIError
 from nucleus.job import AsyncJob
 from nucleus.utils import (
     KeyErrorDict,
     convert_export_payload,
     format_dataset_item_response,
+    paginate_generator,
 )
 
 
@@ -57,30 +59,6 @@ def __eq__(self, other):
                 return True
         return False
 
-    def _fetch_all(self) -> dict:
-        """Retrieves info and all items of the Slice.
-
-        Returns:
-            A dict mapping keys to the corresponding info retrieved.
-            ::
-
-                {
-                    "name": Union[str, int],
-                    "slice_id": str,
-                    "dataset_id": str,
-                    "dataset_items": List[{
-                        "id": str,
-                        "metadata": Dict[str, Union[str, int, float]],
-                        "ref_id": str,
-                        "original_image_url": str
-                    }]
-                }
-        """
-        response = self._client.make_request(
-            {}, f"slice/{self.id}", requests_command=requests.get
-        )
-        return response
-
     @property
     def slice_id(self):
         warnings.warn(
@@ -103,10 +81,52 @@ def dataset_id(self):
             self._dataset_id = self.info()["dataset_id"]
         return self._dataset_id
 
+    def items_generator(self, page_size=100000):
+        """Generator yielding all dataset items in the dataset.
+
+        ::
+            sum_example_field = 0
+            for item in slice.items_generator():
+                sum += item.metadata["example_field"]
+
+        Args:
+            page_size (int, optional): Number of items to return per page. If you are
+                experiencing timeouts while using this generator, you can try lowering
+                the page size.
+
+        Yields:
+            an iterable of DatasetItem objects.
+        """
+        json_generator = paginate_generator(
+            client=self._client,
+            endpoint=f"slice/{self.id}/itemsPage",
+            result_key=ITEMS_KEY,
+            page_size=page_size,
+        )
+        for item_json in json_generator:
+            yield DatasetItem.from_json(item_json)
+
     @property
     def items(self):
-        """All DatasetItems contained in the Slice."""
-        return self._fetch_all()["dataset_items"]
+        """All DatasetItems contained in the Slice.
+
+        For fetching more than 200k items see :meth:`Slice.items_generator`.
+
+        """
+        try:
+            dataset_item_jsons = self._client.make_request(
+                {}, f"slice/{self.id}", requests_command=requests.get
+            )[
+                "dataset_items"
+            ]  # Unfortunately, we didn't use a standard value here, so not using a constant for the key
+            return [
+                DatasetItem.from_json(dataset_item_json)
+                for dataset_item_json in dataset_item_jsons
+            ]
+        except NucleusAPIError as e:
+            if e.status_code == 503:
+                e.message += "/n Your request timed out while trying to get all the items in the slice. Please try slice.items_generator() instead."
+            raise e
 
     def info(self) -> dict:
         """Retrieves the name, slice_id, and dataset_id of the Slice.

nucleus/utils.py

@@ -4,7 +4,7 @@
 import json
 import uuid
 from collections import defaultdict
-from typing import IO, Dict, List, Sequence, Type, Union
+from typing import IO, TYPE_CHECKING, Dict, List, Sequence, Type, Union
 
 import requests
 from requests.models import HTTPError
@@ -19,6 +19,7 @@
     PolygonAnnotation,
     SegmentationAnnotation,
 )
+from nucleus.errors import NucleusAPIError
 
 from .constants import (
     ANNOTATION_TYPES,
@@ -27,8 +28,11 @@
     CATEGORY_TYPE,
     CUBOID_TYPE,
     ITEM_KEY,
+    LAST_PAGE,
     LINE_TYPE,
     MULTICATEGORY_TYPE,
+    PAGE_SIZE,
+    PAGE_TOKEN,
     POLYGON_TYPE,
     REFERENCE_ID_KEY,
     SEGMENTATION_TYPE,
@@ -50,6 +54,9 @@
     '\\\\"': '"',
 }
 
+if TYPE_CHECKING:
+    from . import NucleusClient
+
 
 class KeyErrorDict(dict):
     """Wrapper for response dicts with deprecated keys.
@@ -292,3 +299,27 @@ def replace_double_slashes(s: str) -> str:
     for key, val in STRING_REPLACEMENTS.items():
         s = s.replace(key, val)
     return s
+
+
+def paginate_generator(
+    client: "NucleusClient",
+    endpoint: str,
+    result_key: str,
+    page_size: int = 100000,
+):
+    last_page = False
+    page_token = None
+    while not last_page:
+        try:
+            response = client.make_request(
+                {PAGE_TOKEN: page_token, PAGE_SIZE: page_size},
+                endpoint,
+                requests.post,
+            )
+        except NucleusAPIError as e:
+            if e.status_code == 503:
+                e.message += f"/n Your request timed out while trying to get a page size of {page_size}. Try lowering the page_size."
+            raise e
+        page_token, last_page = response[PAGE_TOKEN], response[LAST_PAGE]
+        for json_value in response[result_key]:
+            yield json_value