feat: allow uploading manually timestamped videos

nominal/_utils.py

@@ -41,13 +41,14 @@ def from_path_dataset(cls, path: Path | str) -> FileType:
 
 
 class FileTypes:
+    BINARY: FileType = FileType("", "application/octet-stream")
     CSV: FileType = FileType(".csv", "text/csv")
     CSV_GZ: FileType = FileType(".csv.gz", "text/csv")
-    # https://issues.apache.org/jira/browse/PARQUET-1889
-    PARQUET: FileType = FileType(".parquet", "application/vnd.apache.parquet")
+    JSON: FileType = FileType(".json", "application/json")
     MP4: FileType = FileType(".mp4", "video/mp4")
-    BINARY: FileType = FileType("", "application/octet-stream")
     MCAP: FileType = FileType(".mcap", "application/octet-stream")
+    # https://issues.apache.org/jira/browse/PARQUET-1889
+    PARQUET: FileType = FileType(".parquet", "application/vnd.apache.parquet")
 
 
 @contextmanager

nominal/core/_multipart.py

@@ -2,17 +2,23 @@
 
 import concurrent.futures
 import logging
+import pathlib
+import urllib.parse
 from functools import partial
 from queue import Queue
 from typing import BinaryIO, Iterable
 
 import requests
 
 from nominal._api.scout_service_api import ingest_api, upload_api
+from nominal._utils import FileType
 from nominal.exceptions import NominalMultipartUploadFailed
 
 logger = logging.getLogger(__name__)
 
+DEFAULT_CHUNK_SIZE = 64_000_000
+DEFAULT_NUM_WORKERS = 8
+
 
 def _sign_and_upload_part_job(
     upload_client: upload_api.UploadService,
@@ -54,8 +60,8 @@ def put_multipart_upload(
     filename: str,
     mimetype: str,
     upload_client: upload_api.UploadService,
-    chunk_size: int = 64_000_000,
-    max_workers: int = 8,
+    chunk_size: int = DEFAULT_CHUNK_SIZE,
+    max_workers: int = DEFAULT_NUM_WORKERS,
 ) -> str:
     """Execute a multipart upload to S3.
 
@@ -118,6 +124,52 @@ def put_multipart_upload(
         raise e
 
 
+def upload_multipart_io(
+    auth_header: str,
+    f: BinaryIO,
+    name: str,
+    file_type: FileType,
+    upload_client: upload_api.UploadService,
+    chunk_size: int = DEFAULT_CHUNK_SIZE,
+    max_workers: int = DEFAULT_NUM_WORKERS,
+) -> str:
+    urlsafe_name = urllib.parse.quote_plus(name)
+    safe_filename = f"{urlsafe_name}{file_type.extension}"
+    return put_multipart_upload(
+        auth_header,
+        f,
+        safe_filename,
+        file_type.mimetype,
+        upload_client,
+        chunk_size=chunk_size,
+        max_workers=max_workers,
+    )
+
+
+def upload_multipart_file(
+    auth_header: str,
+    file: pathlib.Path,
+    upload_client: upload_api.UploadService,
+    file_type: FileType | None = None,
+    chunk_size: int = DEFAULT_CHUNK_SIZE,
+    max_workers: int = DEFAULT_NUM_WORKERS,
+) -> str:
+    if file_type is None:
+        file_type = FileType.from_path(file)
+
+    file_name = file.stem.split(".")[0]
+    with file.open("rb") as file_handle:
+        return upload_multipart_io(
+            auth_header,
+            file_handle,
+            file_name,
+            file_type,
+            upload_client,
+            chunk_size=chunk_size,
+            max_workers=max_workers,
+        )
+
+
 def _abort(upload_client: upload_api.UploadService, auth_header: str, key: str, upload_id: str, e: Exception) -> None:
     logger.error(
         "aborting multipart upload due to an exception", exc_info=e, extra={"key": key, "upload_id": upload_id}

nominal/core/client.py

@@ -1,6 +1,8 @@
 from __future__ import annotations
 
-import urllib.parse
+import json
+import logging
+import tempfile
 from dataclasses import dataclass, field
 from datetime import datetime
 from io import TextIOBase
@@ -34,7 +36,7 @@
 )
 from nominal.core._clientsbunch import ClientsBunch
 from nominal.core._conjure_utils import _available_units, _build_unit_update
-from nominal.core._multipart import put_multipart_upload
+from nominal.core._multipart import upload_multipart_file, upload_multipart_io
 from nominal.core._utils import construct_user_agent_string, rid_from_instance_or_string
 from nominal.core.attachment import Attachment, _iter_get_attachments
 from nominal.core.channel import Channel
@@ -58,6 +60,8 @@
 
 from .asset import Asset
 
+logger = logging.getLogger(__name__)
+
 
 @dataclass(frozen=True)
 class NominalClient:
@@ -235,13 +239,12 @@ def create_mcap_dataset(
             )
 
         mcap_path = Path(path)
-        file_type = FileTypes.MCAP
-        urlsafe_name = urllib.parse.quote_plus(mcap_path.stem)
-        filename = f"{urlsafe_name}{file_type.extension}"
-        with open(mcap_path, "rb") as f:
-            s3_path = put_multipart_upload(
-                self._clients.auth_header, f, filename, file_type.mimetype, self._clients.upload
-            )
+        s3_path = upload_multipart_file(
+            self._clients.auth_header,
+            mcap_path,
+            self._clients.upload,
+            file_type=FileTypes.MCAP,
+        )
         source = ingest_api.IngestSource(s3=ingest_api.S3IngestSource(path=s3_path))
         request = ingest_api.IngestMcapRequest(
             channel_config=[],
@@ -288,12 +291,7 @@ def create_dataset_from_io(
             raise TypeError(f"dataset {dataset} must be open in binary mode, rather than text mode")
 
         file_type = FileType(*file_type)
-        urlsafe_name = urllib.parse.quote_plus(name)
-        filename = f"{urlsafe_name}{file_type.extension}"
-
-        s3_path = put_multipart_upload(
-            self._clients.auth_header, dataset, filename, file_type.mimetype, self._clients.upload
-        )
+        s3_path = upload_multipart_io(self._clients.auth_header, dataset, name, file_type, self._clients.upload)
         request = ingest_api.TriggerFileIngest(
             destination=ingest_api.IngestDestination(
                 new_dataset=ingest_api.NewDatasetIngestDestination(
@@ -319,36 +317,63 @@ def create_video_from_io(
         self,
         video: BinaryIO,
         name: str,
-        start: datetime | IntegralNanosecondsUTC,
+        start: datetime | IntegralNanosecondsUTC | None = None,
         description: str | None = None,
         file_type: tuple[str, str] | FileType = FileTypes.MP4,
         *,
         labels: Sequence[str] = (),
         properties: Mapping[str, str] | None = None,
+        frame_timestamps: Sequence[IntegralNanosecondsUTC] | None = None,
     ) -> Video:
         """Create a video from a file-like object.
 
         The video must be a file-like object in binary mode, e.g. open(path, "rb") or io.BytesIO.
         """
+        if (start is None and frame_timestamps is None) and (None not in (start, frame_timestamps)):
+            raise ValueError("One of 'start' or 'frame_timestamps' must be provided")
+
         if isinstance(video, TextIOBase):
             raise TypeError(f"video {video} must be open in binary mode, rather than text mode")
 
-        file_type = FileType(*file_type)
-        urlsafe_name = urllib.parse.quote_plus(name)
-        filename = f"{urlsafe_name}{file_type.extension}"
+        # if frame_timestamps is None:
+        if start is None:
+            with tempfile.NamedTemporaryFile(suffix=".json", delete_on_close=False) as tmp:
+                tmp.close()
+                timestamp_manifest_path = Path(tmp.name)
 
-        s3_path = put_multipart_upload(
-            self._clients.auth_header, video, filename, file_type.mimetype, self._clients.upload
-        )
+                logger.debug(f"Writing timestamp manifest to '{timestamp_manifest_path}'")
+                timestamp_manifest_path.write_text(json.dumps(frame_timestamps))
+
+                logger.debug("Uploading timestamp manifests to s3")
+                manifest_s3_path = upload_multipart_file(
+                    self._clients.auth_header, timestamp_manifest_path, self._clients.upload
+                )
+
+            timestamp_manifest = ingest_api.VideoTimestampManifest(
+                timestamp_manifests=ingest_api.TimestampManifest(
+                    sources=[
+                        ingest_api.IngestSource(
+                            s3=ingest_api.S3IngestSource(
+                                path=manifest_s3_path,
+                            )
+                        )
+                    ]
+                )
+            )
+        else:
+            timestamp_manifest = ingest_api.VideoTimestampManifest(
+                no_manifest=ingest_api.NoTimestampManifest(
+                    starting_timestamp=_SecondsNanos.from_flexible(start).to_ingest_api()
+                )
+            )
+
+        file_type = FileType(*file_type)
+        s3_path = upload_multipart_io(self._clients.auth_header, video, name, file_type, self._clients.upload)
         request = ingest_api.IngestVideoRequest(
             labels=list(labels),
             properties={} if properties is None else dict(properties),
             sources=[ingest_api.IngestSource(s3=ingest_api.S3IngestSource(path=s3_path))],
-            timestamps=ingest_api.VideoTimestampManifest(
-                no_manifest=ingest_api.NoTimestampManifest(
-                    starting_timestamp=_SecondsNanos.from_flexible(start).to_ingest_api()
-                )
-            ),
+            timestamps=timestamp_manifest,
             description=description,
             title=name,
         )
@@ -481,11 +506,12 @@ def create_attachment_from_io(
             raise TypeError(f"attachment {attachment} must be open in binary mode, rather than text mode")
 
         file_type = FileType(*file_type)
-        urlsafe_name = urllib.parse.quote_plus(name)
-        filename = f"{urlsafe_name}{file_type.extension}"
-
-        s3_path = put_multipart_upload(
-            self._clients.auth_header, attachment, filename, file_type.mimetype, self._clients.upload
+        s3_path = upload_multipart_io(
+            self._clients.auth_header,
+            attachment,
+            name,
+            file_type,
+            self._clients.upload,
         )
         request = attachments_api.CreateAttachmentRequest(
             description=description or "",
@@ -604,12 +630,7 @@ def create_video_from_mcap_io(
             raise TypeError(f"dataset {mcap} must be open in binary mode, rather than text mode")
 
         file_type = FileType(*file_type)
-        urlsafe_name = urllib.parse.quote_plus(name)
-        filename = f"{urlsafe_name}{file_type.extension}"
-
-        s3_path = put_multipart_upload(
-            self._clients.auth_header, mcap, filename, file_type.mimetype, self._clients.upload
-        )
+        s3_path = upload_multipart_io(self._clients.auth_header, mcap, name, file_type, self._clients.upload)
         request = ingest_api.IngestMcapRequest(
             channel_config=[
                 ingest_api.McapChannelConfig(

nominal/core/dataset.py

@@ -2,7 +2,6 @@
 
 import logging
 import time
-import urllib.parse
 from dataclasses import dataclass, field
 from datetime import timedelta
 from io import TextIOBase
@@ -27,7 +26,7 @@
 from nominal._utils import FileType, FileTypes
 from nominal.core._clientsbunch import HasAuthHeader
 from nominal.core._conjure_utils import _available_units, _build_unit_update
-from nominal.core._multipart import put_multipart_upload
+from nominal.core._multipart import upload_multipart_io
 from nominal.core._utils import HasRid, update_dataclass
 from nominal.core.channel import Channel, _get_series_values_csv
 from nominal.exceptions import NominalIngestError, NominalIngestFailed, NominalIngestMultiError
@@ -181,13 +180,15 @@ def add_to_dataset_from_io(
         if isinstance(dataset, TextIOBase):
             raise TypeError(f"dataset {dataset!r} must be open in binary mode, rather than text mode")
 
-        file_type = FileType(*file_type)
-
         self.poll_until_ingestion_completed()
-        urlsafe_name = urllib.parse.quote_plus(self.name)
-        filename = f"{urlsafe_name}{file_type.extension}"
-        s3_path = put_multipart_upload(
-            self._clients.auth_header, dataset, filename, file_type.mimetype, self._clients.upload
+
+        file_type = FileType(*file_type)
+        s3_path = upload_multipart_io(
+            self._clients.auth_header,
+            dataset,
+            self.name,
+            file_type,
+            self._clients.upload,
         )
         request = ingest_api.TriggerFileIngest(
             destination=ingest_api.IngestDestination(