scale-nucleus 0.16.1__tar.gz → 0.16.3__tar.gz

{scale_nucleus-0.16.1 → scale_nucleus-0.16.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: scale-nucleus
-Version: 0.16.1
+Version: 0.16.3
 Summary: The official Python client library for Nucleus, the Data Platform for AI
 Home-page: https://scale.com/nucleus
 License: MIT

{scale_nucleus-0.16.1 → scale_nucleus-0.16.3}/nucleus/async_job.py

@@ -1,6 +1,7 @@
 import time
 from dataclasses import dataclass
-from typing import Dict, List
+from enum import Enum
+from typing import Dict, List, Set
 
 import requests
 
@@ -16,6 +17,29 @@ from nucleus.utils import replace_double_slashes
 JOB_POLLING_INTERVAL = 5
 
 
+class JobStatus(str, Enum):
+    QUEUED = "Queued"
+    RUNNING = "Running"
+    COMPLETED = "Completed"
+    ERRORED_DEPRECATED = "Errored"
+    ERRORED_SERVER = "Errored_Server"  # Server Error
+    ERRORED_USER = "Errored_User"  # User Error
+    ERRORED_PARTIAL = "Errored_Partial"  # Partially Completed
+    ERRORED_HANGING = "Errored_Hanging"  # Hanging
+    CANCELLED = "Cancelled"
+    RETRIED = "Retried"
+
+
+JOB_ERROR_PREFIX = JobStatus.ERRORED_DEPRECATED
+JOB_ERROR_STATES: Set[JobStatus] = {
+    JobStatus.ERRORED_DEPRECATED,
+    JobStatus.ERRORED_SERVER,
+    JobStatus.ERRORED_USER,
+    JobStatus.ERRORED_PARTIAL,
+    JobStatus.ERRORED_HANGING,
+}
+
+
 @dataclass
 class AsyncJob:
     """Object used to check the status or errors of a long running asynchronous operation.
@@ -116,7 +140,9 @@ class AsyncJob:
                 f"Finished at {time.perf_counter() - start_time} s: {status}"
             )
         final_status = status
-        if final_status["status"] == "Errored":
+        if final_status["status"] in JOB_ERROR_STATES or final_status[
+            "status"
+        ].startswith(JOB_ERROR_PREFIX):
             raise JobError(final_status, self)
 
     @classmethod

{scale_nucleus-0.16.1 → scale_nucleus-0.16.3}/nucleus/constants.py

@@ -89,6 +89,7 @@ KEYPOINTS_SKELETON_KEY = "skeleton"
 LABEL_KEY = "label"
 LABELS_KEY = "labels"
 MASK_URL_KEY = "mask_url"
+MAX_ES_PAGE_SIZE = 10000  # Max number of document fetches allowed per ES page
 MAX_PAYLOAD_SIZE = 0x1FFFFFE8  # Set to max string size since we currently convert payloads to strings for processing on the server-side
 MESSAGE_KEY = "message"
 METADATA_KEY = "metadata"

{scale_nucleus-0.16.1 → scale_nucleus-0.16.3}/nucleus/dataset.py

@@ -46,6 +46,7 @@ from .constants import (
     ITEMS_KEY,
     JOB_REQ_LIMIT,
     KEEP_HISTORY_KEY,
+    MAX_ES_PAGE_SIZE,
     MESSAGE_KEY,
     NAME_KEY,
     OBJECT_IDS_KEY,
@@ -1972,12 +1973,32 @@ class Dataset:
             client=self._client,
             endpoint=f"dataset/{self.id}/queryItemsPage",
             result_key=ITEMS_KEY,
-            page_size=10000,  # max ES page size
+            page_size=MAX_ES_PAGE_SIZE,
             query=query,
         )
         for item_json in json_generator:
             yield DatasetItem.from_json(item_json)
 
+    def query_scenes(self, query: str) -> Iterable[Scene]:
+        """
+        Fetches all Scenes that pertain to a given structured query.
+
+        Args:
+            query: Structured query compatible with the `Nucleus query language <https://nucleus.scale.com/docs/query-language-reference>`_.
+
+        Returns:
+            A list of Scene query results.
+        """
+        json_generator = paginate_generator(
+            client=self._client,
+            endpoint=f"dataset/{self.id}/queryScenesPage",
+            result_key=ITEMS_KEY,
+            page_size=MAX_ES_PAGE_SIZE,
+            query=query,
+        )
+        for item_json in json_generator:
+            yield Scene.from_json(item_json, None, True)
+
     @property
     def tracks(self) -> List[Track]:
         """Tracks unique to this dataset.

{scale_nucleus-0.16.1 → scale_nucleus-0.16.3}/pyproject.toml

@@ -25,7 +25,7 @@ ignore = ["E501", "E741", "E731", "F401"]  # Easy ignore for getting it running
 
 [tool.poetry]
 name = "scale-nucleus"
-version = "0.16.1"
+version = "0.16.3"
 description = "The official Python client library for Nucleus, the Data Platform for AI"
 license =  "MIT"
 authors = ["Scale AI Nucleus Team <nucleusapi@scaleapi.com>"]