allow for concurrent task fetches (#411)

jean-lucas · web-flow · commit f2dbd31d63a8 · 2023-11-22T16:08:31.000+01:00
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -6,6 +6,23 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 
+## [0.16.11](https://github.com/scaleapi/nucleus-python-client/releases/tag/v0.16.11) - 2023-11-22
+
+### Added
+
+- Method to allow for concurrent task fetches for pointcloud data
+
+Example:
+```python
+>>> task_ids = ['task_1', 'task_2']
+>>> resp = client.download_pointcloud_tasks(task_ids=task_ids, frame_num=1)
+>>> resp
+{
+  'task_1': [Point3D(x=5, y=10.7, z=-2.3), ...],
+  'task_2': [Point3D(x=1.3 y=11.1, z=1.5), ...],
+}
+```
+
 ## [0.16.10](https://github.com/scaleapi/nucleus-python-client/releases/tag/v0.16.10) - 2023-11-22
 
 Allow creating a dataset by crawling all images in a directory, recursively. Also supports privacy mode datasets.
@@ -53,7 +70,7 @@ This would create a dataset `my-dataset`, and when opened in Nucleus, the images
 
 ### Fixes
 
-- Minor fixes to video scene upload on privacy moce
+- Minor fixes to video scene upload on privacy mode
 
 ## [0.16.8](https://github.com/scaleapi/nucleus-python-client/releases/tag/v0.16.8) - 2023-11-16
 
diff --git a/nucleus/__init__.py b/nucleus/__init__.py
@@ -70,6 +70,7 @@
     SegmentationAnnotation,
 )
 from .async_job import AsyncJob, EmbeddingsExportJob
+from .async_utils import make_multiple_requests_concurrently
 from .camera_params import CameraParams
 from .connection import Connection
 from .constants import (
@@ -1065,6 +1066,50 @@ def download_pointcloud_task(
 
         return [Point3D.from_json(pt) for pt in points]
 
+    def download_pointcloud_tasks(
+        self, task_ids: List[str], frame_num: int
+    ) -> Dict[str, List[Union[Point3D, LidarPoint]]]:
+        """
+        Download the lidar point cloud data for a given set of tasks and frame number.
+
+        Parameters:
+            task_ids: list of task ids to fetch data from
+            frame_num: download point cloud for this particular frame
+
+        Returns:
+            A dictionary from task_id to list of Point3D objects
+
+        """
+        endpoints = [
+            f"task/{task_id}/frame/{frame_num}" for task_id in task_ids
+        ]
+        progressbar = self.tqdm_bar(
+            total=len(endpoints),
+            desc="Downloading pointcloud tasks",
+        )
+        results = make_multiple_requests_concurrently(
+            client=self,
+            requests=endpoints,
+            route=None,
+            progressbar=progressbar,
+        )
+        resp = {}
+
+        for result in results:
+            req, data = result
+            task_id = req.split("/")[1]  # task/<task id>/frame/1 => task_id
+            points = data.get(POINTS_KEY, None)
+            if points is None or len(points) == 0:
+                raise Exception("Response has invalid payload")
+
+            sample_point = points[0]
+            if I_KEY in sample_point.keys():
+                resp[task_id] = [LidarPoint.from_json(pt) for pt in points]
+
+            resp[task_id] = [Point3D.from_json(pt) for pt in points]
+
+        return resp
+
     @deprecated("Prefer calling Dataset.create_custom_index instead.")
     def create_custom_index(
         self, dataset_id: str, embeddings_urls: list, embedding_dim: int
diff --git a/nucleus/annotation_uploader.py b/nucleus/annotation_uploader.py
@@ -6,7 +6,7 @@
 from nucleus.async_utils import (
     FileFormField,
     FormDataContextHandler,
-    make_many_form_data_requests_concurrently,
+    make_multiple_requests_concurrently,
 )
 from nucleus.constants import MASK_TYPE, SERIALIZED_REQUEST_KEY
 from nucleus.errors import DuplicateIDError
@@ -150,7 +150,7 @@ def make_batched_file_form_data_requests(
             desc="Local segmentation mask file batches",
         )
 
-        return make_many_form_data_requests_concurrently(
+        return make_multiple_requests_concurrently(
             client=self._client,
             requests=requests,
             route=self._route,
diff --git a/nucleus/async_utils.py b/nucleus/async_utils.py
@@ -1,7 +1,15 @@
 import asyncio
 import time
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, BinaryIO, Callable, Sequence, Tuple
+from typing import (
+    TYPE_CHECKING,
+    BinaryIO,
+    Callable,
+    Optional,
+    Sequence,
+    Tuple,
+    Union,
+)
 
 import aiohttp
 import nest_asyncio
@@ -84,56 +92,71 @@ def get_event_loop():
     return loop
 
 
-def make_many_form_data_requests_concurrently(
+def make_multiple_requests_concurrently(
     client: "NucleusClient",
-    requests: Sequence[FormDataContextHandler],
-    route: str,
+    requests: Sequence[Union[FormDataContextHandler, str]],
+    route: Optional[str],
     progressbar: tqdm,
 ):
     """
     Makes an async post request with form data to a Nucleus endpoint.
 
     Args:
         client: The client to use for the request.
-        requests: Each requst should be a FormDataContextHandler object which will
-            handle generating form data, and opening/closing files for each request.
-        route: route for the request.
+        requests: a list of requests to make. This list either comprises a string of endpoints to request,
+        or a list of FormDataContextHandler object which will handle generating form data, and opening/closing files for each request.
+        route: A route is required when requests are for Form Data Post requests
         progressbar: A tqdm progress bar to use for showing progress to the user.
     """
     loop = get_event_loop()
     return loop.run_until_complete(
-        form_data_request_helper(client, requests, route, progressbar)
+        _request_helper(client, requests, route, progressbar)
     )
 
 
-async def form_data_request_helper(
+async def _request_helper(
     client: "NucleusClient",
-    requests: Sequence[FormDataContextHandler],
-    route: str,
+    requests: Sequence[Union[FormDataContextHandler, str]],
+    route: Optional[str],
     progressbar: tqdm,
 ):
     """
-    Makes an async post request with files to a Nucleus endpoint.
+    Makes an async requests to a Nucleus endpoint.
 
     Args:
         client: The client to use for the request.
-        requests: Each request should be a FormDataContextHandler object which will
-            handle generating form data, and opening/closing files for each request.
+        requests: a list of requests to make. This list either comprises a string of endpoints to request,
+        or a list of FormDataContextHandler object which will handle generating form data, and opening/closing files for each request.
         route: route for the request.
     """
     async with aiohttp.ClientSession() as session:
-        tasks = [
-            asyncio.ensure_future(
-                _post_form_data(
-                    client=client,
-                    request=request,
-                    route=route,
-                    session=session,
-                    progressbar=progressbar,
+        tasks = []
+        for request in requests:
+            if isinstance(request, FormDataContextHandler):
+                assert (
+                    route
+                ), "A route must be specified for FormDataContextHandler requests"
+                req = asyncio.ensure_future(
+                    _post_form_data(
+                        client=client,
+                        request=request,
+                        route=route,
+                        session=session,
+                        progressbar=progressbar,
+                    )
                 )
-            )
-            for request in requests
-        ]
+                tasks.append(req)
+            else:
+                req = asyncio.ensure_future(
+                    _make_request(
+                        client=client,
+                        request=request,
+                        session=session,
+                        progressbar=progressbar,
+                    )
+                )
+                tasks.append(req)
+
         return await asyncio.gather(*tasks)
 
 
@@ -165,36 +188,80 @@ async def _post_form_data(
                     auth=aiohttp.BasicAuth(client.api_key, ""),
                     timeout=DEFAULT_NETWORK_TIMEOUT_SEC,
                 ) as response:
-                    logger.info(
-                        "API request has response code %s", response.status
+                    data = await _parse_async_response(
+                        endpoint, session, response, sleep_time
                     )
-
-                    try:
-                        data = await response.json()
-                    except aiohttp.client_exceptions.ContentTypeError:
-                        # In case of 404, the server returns text
-                        data = await response.text()
-                    if (
-                        response.status in RetryStrategy.statuses
-                        and sleep_time != -1
-                    ):
-                        time.sleep(sleep_time)
+                    if data is None:
                         continue
 
-                    if response.status == 503:
-                        raise TimeoutError(
-                            "The request to upload your max is timing out, please lower local_files_per_upload_request in your api call."
-                        )
-
-                    if not response.ok:
-                        raise NucleusAPIError(
-                            endpoint,
-                            session.post,
-                            aiohttp_response=(
-                                response.status,
-                                response.reason,
-                                data,
-                            ),
-                        )
                     progressbar.update(1)
                     return data
+
+
+async def _make_request(
+    client: "NucleusClient",
+    request: str,
+    session: aiohttp.ClientSession,
+    progressbar: tqdm,
+):
+    """
+    Makes an async post request with files to a Nucleus endpoint.
+
+    Args:
+        client: The client to use for the request.
+        request: The request to make (See FormDataContextHandler for more details.)
+        route: route for the request.
+        session: The session to use for the request.
+
+    Returns:
+        A tuple (endpoint request string, response from endpoint)
+    """
+    endpoint = f"{client.endpoint}/{request}"
+    logger.info("GET %s", endpoint)
+
+    async with UPLOAD_SEMAPHORE:
+        for sleep_time in RetryStrategy.sleep_times() + [-1]:
+            async with session.get(
+                endpoint,
+                auth=aiohttp.BasicAuth(client.api_key, ""),
+                timeout=DEFAULT_NETWORK_TIMEOUT_SEC,
+            ) as response:
+                data = await _parse_async_response(
+                    endpoint, session, response, sleep_time
+                )
+                if data is None:
+                    continue
+
+                progressbar.update(1)
+                return (request, data)
+
+
+async def _parse_async_response(endpoint, session, response, sleep_time):
+    logger.info("API request has response code %s", response.status)
+
+    try:
+        data = await response.json()
+    except aiohttp.client_exceptions.ContentTypeError:
+        # In case of 404, the server returns text
+        data = await response.text()
+    if response.status in RetryStrategy.statuses and sleep_time != -1:
+        time.sleep(sleep_time)
+        return None
+
+    if response.status == 503:
+        raise TimeoutError(
+            "The request to upload your max is timing out, please lower local_files_per_upload_request in your api call."
+        )
+
+    if not response.ok:
+        raise NucleusAPIError(
+            endpoint,
+            session.get,
+            aiohttp_response=(
+                response.status,
+                response.reason,
+                data,
+            ),
+        )
+
+    return data
diff --git a/nucleus/dataset_item_uploader.py b/nucleus/dataset_item_uploader.py
@@ -14,7 +14,7 @@
     FileFormData,
     FileFormField,
     FormDataContextHandler,
-    make_many_form_data_requests_concurrently,
+    make_multiple_requests_concurrently,
 )
 
 from .constants import DATASET_ID_KEY, IMAGE_KEY, ITEMS_KEY, UPDATE_KEY
@@ -125,7 +125,7 @@ def _process_append_requests_local(
             desc=f"Uploading {len(items)} items in {len(requests)} batches",
         )
 
-        return make_many_form_data_requests_concurrently(
+        return make_multiple_requests_concurrently(
             self._client,
             requests,
             f"dataset/{dataset_id}/append",
diff --git a/tests/test_video_scene.py b/tests/test_video_scene.py
@@ -158,7 +158,7 @@ def test_video_scene_upload_async(dataset_video_scene):
 
 
 @pytest.mark.integration
-@pytest.mark.xfail(reason="SFN doesn't throw on validation error - 05.10.2023")
+@pytest.mark.skip(reason="SFN doesn't throw on validation error - 05.10.2023")
 def test_repeat_refid_video_scene_upload_async(dataset_video_scene):
     payload = TEST_VIDEO_SCENES_REPEAT_REF_IDS
     scenes = [
