Prevent oversized payloads (#280)

* Prevent oversized payloads

* upped version

Author: cmpajot

CHANGELOG.md

@@ -5,6 +5,11 @@ All notable changes to the [Nucleus Python Client](https://github.com/scaleapi/n
 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.10.1](https://github.com/scaleapi/nucleus-python-client/releases/tag/v0.10.1) - 2022-04-21
+
+### Added
+- Added check for payload size
+
 ## [0.10.0](https://github.com/scaleapi/nucleus-python-client/releases/tag/v0.10.0)) - 2022-04-21
 
 ### Added

nucleus/constants.py

@@ -85,6 +85,7 @@
 LABEL_KEY = "label"
 LABELS_KEY = "labels"
 MASK_URL_KEY = "mask_url"
+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"
 MODEL_ID_KEY = "model_id"

nucleus/utils.py

@@ -32,6 +32,7 @@
     KEYPOINTS_TYPE,
     LAST_PAGE,
     LINE_TYPE,
+    MAX_PAYLOAD_SIZE,
     MULTICATEGORY_TYPE,
     PAGE_SIZE,
     PAGE_TOKEN,
@@ -247,6 +248,13 @@ def serialize_and_write(
     ],
     file_pointer,
 ):
+    """Helper function serialize and write payload to file
+
+    Args:
+        upload_units: Sequence of items, annotations or scenes
+        file_pointer: Pointer of the file to write to
+    """
+    bytes_written = 0
     if len(upload_units) == 0:
         raise ValueError(
             "Expecting at least one object when serializing objects to upload, but got zero.  Please try again."
@@ -256,9 +264,9 @@ def serialize_and_write(
             if isinstance(
                 unit, (DatasetItem, Annotation, LidarScene, VideoScene)
             ):
-                file_pointer.write(unit.to_json() + "\n")
+                bytes_written += file_pointer.write(unit.to_json() + "\n")
             else:
-                file_pointer.write(json.dumps(unit) + "\n")
+                bytes_written += file_pointer.write(json.dumps(unit) + "\n")
         except TypeError as e:
             type_name = type(unit).__name__
             message = (
@@ -273,6 +281,10 @@ def serialize_and_write(
             )
             message += f"The specific error was {e}"
             raise ValueError(message) from e
+    if bytes_written > MAX_PAYLOAD_SIZE:
+        raise ValueError(
+            f"Payload of {bytes_written} bytes exceed maximum payload size of {MAX_PAYLOAD_SIZE} bytes. Please reduce payload size and try again."
+        )
 
 
 def upload_to_presigned_url(presigned_url: str, file_pointer: IO):

pyproject.toml

@@ -21,7 +21,7 @@ exclude = '''
 
 [tool.poetry]
 name = "scale-nucleus"
-version = "0.10.0"
+version = "0.10.1"
 description = "The official Python client library for Nucleus, the Data Platform for AI"
 license =  "MIT"
 authors = ["Scale AI Nucleus Team <nucleusapi@scaleapi.com>"]