mapillary-tools 0.14.0a2__py3-none-any.whl → 0.14.0b1__py3-none-any.whl

mapillary_tools/uploader.py

@@ -1,12 +1,16 @@
 from __future__ import annotations
 
+import concurrent.futures
+
 import dataclasses
 import io
 import json
 import logging
 import os
 import struct
+import sys
 import tempfile
+import threading
 import time
 import typing as T
 import uuid
@@ -14,9 +18,19 @@ import zipfile
 from contextlib import contextmanager
 from pathlib import Path
 
+if sys.version_info >= (3, 11):
+    from typing import Required
+else:
+    from typing_extensions import Required
+
 import requests
 
-from . import api_v4, constants, exif_write, types, upload_api_v4
+from . import api_v4, config, constants, exif_write, types, upload_api_v4, utils
+from .serializer.description import (
+    desc_file_to_exif,
+    DescriptionJSONSerializer,
+    validate_image_desc,
+)
 
 
 LOG = logging.getLogger(__name__)
@@ -56,17 +70,21 @@ class UploaderProgress(T.TypedDict, total=True):
 class SequenceProgress(T.TypedDict, total=False):
     """Progress data at sequence level"""
 
-    # md5sum of the zipfile/BlackVue/CAMM in uploading
-    md5sum: str
+    # Used to check if it is uploaded or not
+    sequence_md5sum: Required[str]
+
+    # Used to resume from the previous upload,
+    # so it has to an unique identifier (hash) of the upload content
+    upload_md5sum: str
 
     # File type
-    file_type: str
+    file_type: Required[str]
 
     # How many sequences in total. It's always 1 when uploading Zipfile/BlackVue/CAMM
-    total_sequence_count: int
+    total_sequence_count: Required[int]
 
     # 0-based nth sequence. It is always 0 when uploading Zipfile/BlackVue/CAMM
-    sequence_idx: int
+    sequence_idx: Required[int]
 
     # How many images in the sequence. It's available only when uploading directories/Zipfiles
     sequence_image_count: int
@@ -121,6 +139,7 @@ class EventEmitter:
     def on(self, event: EventName):
         def _wrap(callback):
             self.events.setdefault(event, []).append(callback)
+            return callback
 
         return _wrap
 
@@ -148,30 +167,50 @@ class ZipImageSequence:
 
         for sequence_uuid, sequence in sequences.items():
             _validate_metadatas(sequence)
-            upload_md5sum = types.update_sequence_md5sum(sequence)
-
             # For atomicity we write into a WIP file and then rename to the final file
             wip_zip_filename = zip_dir.joinpath(
                 f".mly_zip_{uuid.uuid4()}_{sequence_uuid}_{os.getpid()}_{int(time.time())}"
             )
-            filename = _session_key(upload_md5sum, upload_api_v4.ClusterFileType.ZIP)
-            zip_filename = zip_dir.joinpath(filename)
-            with wip_file_context(wip_zip_filename, zip_filename) as wip_path:
+            with cls._wip_file_context(wip_zip_filename) as wip_path:
                 with wip_path.open("wb") as wip_fp:
-                    actual_md5sum = cls.zip_sequence_deterministically(sequence, wip_fp)
-                    assert actual_md5sum == upload_md5sum, "md5sum mismatch"
+                    cls.zip_sequence_fp(sequence, wip_fp)
+
+    @classmethod
+    @contextmanager
+    def _wip_file_context(cls, wip_path: Path):
+        try:
+            os.remove(wip_path)
+        except FileNotFoundError:
+            pass
+        try:
+            yield wip_path
+
+            with wip_path.open("rb") as fp:
+                upload_md5sum = utils.md5sum_fp(fp).hexdigest()
+
+            done_path = wip_path.parent.joinpath(
+                _session_key(upload_md5sum, api_v4.ClusterFileType.ZIP)
+            )
+
+            try:
+                os.remove(done_path)
+            except FileNotFoundError:
+                pass
+            wip_path.rename(done_path)
+        finally:
+            try:
+                os.remove(wip_path)
+            except FileNotFoundError:
+                pass
 
     @classmethod
-    def zip_sequence_deterministically(
+    def zip_sequence_fp(
         cls,
         sequence: T.Sequence[types.ImageMetadata],
         zip_fp: T.IO[bytes],
     ) -> str:
         """
-        Write a sequence of ImageMetadata into the zipfile handle. It should guarantee
-        that the same sequence always produces the same zipfile, because the
-        sequence md5sum will be used to upload the zipfile or resume the upload.
-
+        Write a sequence of ImageMetadata into the zipfile handle.
         The sequence has to be one sequence and sorted.
         """
 
@@ -180,21 +219,23 @@ class ZipImageSequence:
             f"Only one sequence is allowed but got {len(sequence_groups)}: {list(sequence_groups.keys())}"
         )
 
-        upload_md5sum = types.update_sequence_md5sum(sequence)
+        sequence_md5sum = types.update_sequence_md5sum(sequence)
 
         with zipfile.ZipFile(zip_fp, "w", zipfile.ZIP_DEFLATED) as zipf:
             for idx, metadata in enumerate(sequence):
-                # Use {idx}.jpg (suffix does not matter) as the archive name to ensure the
-                # resulting zipfile is deterministic. This determinism is based on the upload_md5sum,
-                # which is derived from a list of image md5sums
-                cls._write_imagebytes_in_zip(zipf, metadata, arcname=f"{idx}.jpg")
+                # Arcname should be unique, the name does not matter
+                arcname = f"{idx}.jpg"
+                zipinfo = zipfile.ZipInfo(arcname, date_time=(1980, 1, 1, 0, 0, 0))
+                zipf.writestr(zipinfo, cls._dump_image_bytes(metadata))
             assert len(sequence) == len(set(zipf.namelist()))
-            zipf.comment = json.dumps({"upload_md5sum": upload_md5sum}).encode("utf-8")
+            zipf.comment = json.dumps({"sequence_md5sum": sequence_md5sum}).encode(
+                "utf-8"
+            )
 
-        return upload_md5sum
+        return sequence_md5sum
 
     @classmethod
-    def extract_upload_md5sum(cls, zip_fp: T.IO[bytes]) -> str:
+    def extract_sequence_md5sum(cls, zip_fp: T.IO[bytes]) -> str:
         with zipfile.ZipFile(zip_fp, "r", zipfile.ZIP_DEFLATED) as ziph:
             comment = ziph.comment
 
@@ -209,17 +250,15 @@ class ZipImageSequence:
         except json.JSONDecodeError as ex:
             raise InvalidMapillaryZipFileError(str(ex)) from ex
 
-        upload_md5sum = zip_metadata.get("upload_md5sum")
+        sequence_md5sum = zip_metadata.get("sequence_md5sum")
 
-        if not upload_md5sum and not isinstance(upload_md5sum, str):
-            raise InvalidMapillaryZipFileError("No upload_md5sum found")
+        if not sequence_md5sum and not isinstance(sequence_md5sum, str):
+            raise InvalidMapillaryZipFileError("No sequence_md5sum found")
 
-        return upload_md5sum
+        return sequence_md5sum
 
     @classmethod
-    def _write_imagebytes_in_zip(
-        cls, zipf: zipfile.ZipFile, metadata: types.ImageMetadata, arcname: str
-    ):
+    def _dump_image_bytes(cls, metadata: types.ImageMetadata) -> bytes:
         try:
             edit = exif_write.ExifEdit(metadata.filename)
         except struct.error as ex:
@@ -227,24 +266,24 @@ class ZipImageSequence:
 
         # The cast is to fix the type checker error
         edit.add_image_description(
-            T.cast(T.Dict, types.desc_file_to_exif(types.as_desc(metadata)))
+            T.cast(
+                T.Dict,
+                desc_file_to_exif(DescriptionJSONSerializer.as_desc(metadata)),
+            )
         )
 
         try:
-            image_bytes = edit.dump_image_bytes()
+            return edit.dump_image_bytes()
         except struct.error as ex:
             raise ExifError(
                 f"Failed to dump EXIF bytes: {ex}", metadata.filename
             ) from ex
 
-        zipinfo = zipfile.ZipInfo(arcname, date_time=(1980, 1, 1, 0, 0, 0))
-        zipf.writestr(zipinfo, image_bytes)
-
     @classmethod
-    def prepare_zipfile_and_upload(
+    def upload_zipfile(
         cls,
-        zip_path: Path,
         uploader: Uploader,
+        zip_path: Path,
         progress: dict[str, T.Any] | None = None,
     ) -> str:
         if progress is None:
@@ -256,30 +295,34 @@ class ZipImageSequence:
                 raise InvalidMapillaryZipFileError("Zipfile has no files")
 
         with zip_path.open("rb") as zip_fp:
-            upload_md5sum = cls.extract_upload_md5sum(zip_fp)
+            sequence_md5sum = cls.extract_sequence_md5sum(zip_fp)
 
-        sequence_progress: SequenceProgress = {
+        # Send the copy of the input progress to each upload session, to avoid modifying the original one
+        mutable_progress: SequenceProgress = {
+            **T.cast(SequenceProgress, progress),
             "sequence_image_count": len(namelist),
+            "sequence_md5sum": sequence_md5sum,
             "file_type": types.FileType.ZIP.value,
-            "md5sum": upload_md5sum,
         }
 
-        session_key = _session_key(upload_md5sum, upload_api_v4.ClusterFileType.ZIP)
-
         with zip_path.open("rb") as zip_fp:
-            return uploader.upload_stream(
-                zip_fp,
-                upload_api_v4.ClusterFileType.ZIP,
-                session_key,
-                # Send the copy of the input progress to each upload session, to avoid modifying the original one
-                progress=T.cast(T.Dict[str, T.Any], {**progress, **sequence_progress}),
+            file_handle = uploader.upload_stream(
+                zip_fp, progress=T.cast(T.Dict[str, T.Any], mutable_progress)
             )
 
+        cluster_id = uploader.finish_upload(
+            file_handle,
+            api_v4.ClusterFileType.ZIP,
+            progress=T.cast(T.Dict[str, T.Any], mutable_progress),
+        )
+
+        return cluster_id
+
     @classmethod
-    def prepare_images_and_upload(
+    def zip_images_and_upload(
         cls,
-        image_metadatas: T.Sequence[types.ImageMetadata],
         uploader: Uploader,
+        image_metadatas: T.Sequence[types.ImageMetadata],
         progress: dict[str, T.Any] | None = None,
     ) -> T.Generator[tuple[str, UploadResult], None, None]:
         if progress is None:
@@ -288,14 +331,6 @@ class ZipImageSequence:
         sequences = types.group_and_sort_images(image_metadatas)
 
         for sequence_idx, (sequence_uuid, sequence) in enumerate(sequences.items()):
-            sequence_progress: SequenceProgress = {
-                "sequence_idx": sequence_idx,
-                "total_sequence_count": len(sequences),
-                "sequence_image_count": len(sequence),
-                "sequence_uuid": sequence_uuid,
-                "file_type": types.FileType.IMAGE.value,
-            }
-
             try:
                 _validate_metadatas(sequence)
             except Exception as ex:
@@ -304,25 +339,28 @@ class ZipImageSequence:
 
             with tempfile.NamedTemporaryFile() as fp:
                 try:
-                    upload_md5sum = cls.zip_sequence_deterministically(sequence, fp)
+                    sequence_md5sum = cls.zip_sequence_fp(sequence, fp)
                 except Exception as ex:
                     yield sequence_uuid, UploadResult(error=ex)
                     continue
 
-                sequence_progress["md5sum"] = upload_md5sum
+                sequence_progress: SequenceProgress = {
+                    "sequence_idx": sequence_idx,
+                    "total_sequence_count": len(sequences),
+                    "sequence_image_count": len(sequence),
+                    "sequence_uuid": sequence_uuid,
+                    "file_type": types.FileType.ZIP.value,
+                    "sequence_md5sum": sequence_md5sum,
+                }
 
-                session_key = _session_key(
-                    upload_md5sum, upload_api_v4.ClusterFileType.ZIP
-                )
+                mutable_progress: dict[str, T.Any] = {**progress, **sequence_progress}
 
                 try:
-                    cluster_id = uploader.upload_stream(
-                        fp,
-                        upload_api_v4.ClusterFileType.ZIP,
-                        session_key,
-                        progress=T.cast(
-                            T.Dict[str, T.Any], {**progress, **sequence_progress}
-                        ),
+                    file_handle = uploader.upload_stream(fp, progress=mutable_progress)
+                    cluster_id = uploader.finish_upload(
+                        file_handle,
+                        api_v4.ClusterFileType.ZIP,
+                        progress=mutable_progress,
                     )
                 except Exception as ex:
                     yield sequence_uuid, UploadResult(error=ex)
@@ -330,11 +368,115 @@ class ZipImageSequence:
 
             yield sequence_uuid, UploadResult(result=cluster_id)
 
+    @classmethod
+    def _upload_sequence(
+        cls,
+        uploader: Uploader,
+        sequence: T.Sequence[types.ImageMetadata],
+        progress: dict[str, T.Any] | None = None,
+    ) -> str:
+        if progress is None:
+            progress = {}
+
+        # FIXME: This is a hack to disable the event emitter inside the uploader
+        uploader_without_emitter = uploader.copy_uploader_without_emitter()
+
+        lock = threading.Lock()
+
+        def _upload_image(image_metadata: types.ImageMetadata) -> str:
+            mutable_progress = {
+                **(progress or {}),
+                "filename": str(image_metadata.filename),
+            }
+
+            bytes = cls._dump_image_bytes(image_metadata)
+            file_handle = uploader_without_emitter.upload_stream(
+                io.BytesIO(bytes), progress=mutable_progress
+            )
+
+            mutable_progress["chunk_size"] = image_metadata.filesize
+
+            with lock:
+                uploader.emitter.emit("upload_progress", mutable_progress)
+
+            return file_handle
+
+        _validate_metadatas(sequence)
+
+        progress["entity_size"] = sum(m.filesize or 0 for m in sequence)
+
+        # TODO: assert sequence is sorted
+
+        uploader.emitter.emit("upload_start", progress)
+
+        with concurrent.futures.ThreadPoolExecutor(
+            max_workers=constants.MAX_IMAGE_UPLOAD_WORKERS
+        ) as executor:
+            image_file_handles = list(executor.map(_upload_image, sequence))
+
+        manifest = {
+            "version": "1",
+            "upload_type": "images",
+            "image_handles": image_file_handles,
+        }
+
+        with io.BytesIO() as manifest_fp:
+            manifest_fp.write(json.dumps(manifest).encode("utf-8"))
+            manifest_fp.seek(0, io.SEEK_SET)
+            manifest_file_handle = uploader_without_emitter.upload_stream(
+                manifest_fp, session_key=f"{uuid.uuid4().hex}.json"
+            )
+
+        uploader.emitter.emit("upload_end", progress)
+
+        cluster_id = uploader.finish_upload(
+            manifest_file_handle,
+            api_v4.ClusterFileType.MLY_BUNDLE_MANIFEST,
+            progress=progress,
+        )
+
+        return cluster_id
+
+    @classmethod
+    def upload_images(
+        cls,
+        uploader: Uploader,
+        image_metadatas: T.Sequence[types.ImageMetadata],
+        progress: dict[str, T.Any] | None = None,
+    ) -> T.Generator[tuple[str, UploadResult], None, None]:
+        if progress is None:
+            progress = {}
+
+        sequences = types.group_and_sort_images(image_metadatas)
+
+        for sequence_idx, (sequence_uuid, sequence) in enumerate(sequences.items()):
+            sequence_md5sum = types.update_sequence_md5sum(sequence)
+
+            sequence_progress: SequenceProgress = {
+                "sequence_idx": sequence_idx,
+                "total_sequence_count": len(sequences),
+                "sequence_image_count": len(sequence),
+                "sequence_uuid": sequence_uuid,
+                "file_type": types.FileType.IMAGE.value,
+                "sequence_md5sum": sequence_md5sum,
+            }
+
+            mutable_progress: dict[str, T.Any] = {**progress, **sequence_progress}
+
+            try:
+                cluster_id = cls._upload_sequence(
+                    uploader, sequence, progress=mutable_progress
+                )
+            except Exception as ex:
+                yield sequence_uuid, UploadResult(error=ex)
+            else:
+                yield sequence_uuid, UploadResult(result=cluster_id)
+
 
 class Uploader:
     def __init__(
         self,
-        user_items: types.UserItem,
+        user_items: config.UserItem,
         emitter: EventEmitter | None = None,
         chunk_size: int = int(constants.UPLOAD_CHUNK_SIZE_MB * 1024 * 1024),
         dry_run=False,
@@ -351,17 +493,25 @@ class Uploader:
     def upload_stream(
         self,
         fp: T.IO[bytes],
-        cluster_filetype: upload_api_v4.ClusterFileType,
-        session_key: str,
+        session_key: str | None = None,
         progress: dict[str, T.Any] | None = None,
     ) -> str:
         if progress is None:
             progress = {}
 
+        if session_key is None:
+            fp.seek(0, io.SEEK_SET)
+            md5sum = utils.md5sum_fp(fp).hexdigest()
+            filetype = progress.get("file_type")
+            if filetype is not None:
+                session_key = _session_key(md5sum, types.FileType(filetype))
+            else:
+                session_key = md5sum
+
         fp.seek(0, io.SEEK_END)
         entity_size = fp.tell()
 
-        upload_service = self._create_upload_service(session_key, cluster_filetype)
+        upload_service = self._create_upload_service(session_key)
 
         progress["entity_size"] = entity_size
         progress["chunk_size"] = self.chunk_size
@@ -384,30 +534,64 @@ class Uploader:
 
         self.emitter.emit("upload_end", progress)
 
-        # TODO: retry here
-        cluster_id = self._finish_upload_retryable(upload_service, file_handle)
-        progress["cluster_id"] = cluster_id
+        return file_handle
+
+    def finish_upload(
+        self,
+        file_handle: str,
+        cluster_filetype: api_v4.ClusterFileType,
+        progress: dict[str, T.Any] | None = None,
+    ) -> str:
+        """Finish upload with safe retries guraranteed"""
+        if progress is None:
+            progress = {}
 
+        if self.dry_run:
+            cluster_id = "0"
+        else:
+            resp = api_v4.finish_upload(
+                self.user_items["user_upload_token"],
+                file_handle,
+                cluster_filetype,
+                organization_id=self.user_items.get("MAPOrganizationKey"),
+            )
+
+            data = resp.json()
+            cluster_id = data.get("cluster_id")
+
+            # TODO: validate cluster_id
+
+        progress["cluster_id"] = cluster_id
         self.emitter.emit("upload_finished", progress)
 
         return cluster_id
 
-    def _create_upload_service(
-        self, session_key: str, cluster_filetype: upload_api_v4.ClusterFileType
-    ) -> upload_api_v4.UploadService:
+    def copy_uploader_without_emitter(self) -> Uploader:
+        return Uploader(
+            self.user_items,
+            emitter=None,
+            chunk_size=self.chunk_size,
+            dry_run=self.dry_run,
+        )
+
+    def _create_upload_service(self, session_key: str) -> upload_api_v4.UploadService:
         upload_service: upload_api_v4.UploadService
 
         if self.dry_run:
+            upload_path = os.getenv("MAPILLARY_UPLOAD_ENDPOINT")
             upload_service = upload_api_v4.FakeUploadService(
                 user_access_token=self.user_items["user_upload_token"],
                 session_key=session_key,
-                cluster_filetype=cluster_filetype,
+                upload_path=Path(upload_path) if upload_path is not None else None,
+            )
+            LOG.info(
+                "Dry run mode enabled. Data will be uploaded to %s",
+                upload_service.upload_path.joinpath(session_key),
             )
         else:
             upload_service = upload_api_v4.UploadService(
                 user_access_token=self.user_items["user_upload_token"],
                 session_key=session_key,
-                cluster_filetype=cluster_filetype,
             )
 
         return upload_service
@@ -484,57 +668,14 @@ class Uploader:
 
         return upload_service.upload_shifted_chunks(shifted_chunks, begin_offset)
 
-    def _finish_upload_retryable(
-        self, upload_service: upload_api_v4.UploadService, file_handle: str
-    ) -> str:
-        """Finish upload with safe retries guraranteed"""
-
-        if self.dry_run:
-            cluster_id = "0"
-        else:
-            resp = api_v4.finish_upload(
-                self.user_items["user_upload_token"],
-                file_handle,
-                upload_service.cluster_filetype,
-                organization_id=self.user_items.get("MAPOrganizationKey"),
-            )
-
-            data = resp.json()
-            cluster_id = data.get("cluster_id")
-
-            # TODO: validate cluster_id
-
-        return cluster_id
-
 
 def _validate_metadatas(metadatas: T.Sequence[types.ImageMetadata]):
     for metadata in metadatas:
-        types.validate_image_desc(types.as_desc(metadata))
+        validate_image_desc(DescriptionJSONSerializer.as_desc(metadata))
         if not metadata.filename.is_file():
             raise FileNotFoundError(f"No such file {metadata.filename}")
 
 
-@contextmanager
-def wip_file_context(wip_path: Path, done_path: Path):
-    assert wip_path != done_path, "should not be the same file"
-    try:
-        os.remove(wip_path)
-    except FileNotFoundError:
-        pass
-    try:
-        yield wip_path
-        try:
-            os.remove(done_path)
-        except FileNotFoundError:
-            pass
-        wip_path.rename(done_path)
-    finally:
-        try:
-            os.remove(wip_path)
-        except FileNotFoundError:
-            pass
-
-
 def _is_immediate_retry(ex: Exception):
     if (
         isinstance(ex, requests.HTTPError)
@@ -568,14 +709,19 @@ def _is_retriable_exception(ex: Exception):
     return False
 
 
-_SUFFIX_MAP: dict[upload_api_v4.ClusterFileType, str] = {
-    upload_api_v4.ClusterFileType.ZIP: ".zip",
-    upload_api_v4.ClusterFileType.CAMM: ".mp4",
-    upload_api_v4.ClusterFileType.BLACKVUE: ".mp4",
-}
-
-
 def _session_key(
-    upload_md5sum: str, cluster_filetype: upload_api_v4.ClusterFileType
+    upload_md5sum: str, filetype: api_v4.ClusterFileType | types.FileType
 ) -> str:
-    return f"mly_tools_{upload_md5sum}{_SUFFIX_MAP[cluster_filetype]}"
+    _SUFFIX_MAP: dict[api_v4.ClusterFileType | types.FileType, str] = {
+        api_v4.ClusterFileType.ZIP: ".zip",
+        api_v4.ClusterFileType.CAMM: ".mp4",
+        api_v4.ClusterFileType.BLACKVUE: ".mp4",
+        types.FileType.IMAGE: ".jpg",
+        types.FileType.ZIP: ".zip",
+        types.FileType.BLACKVUE: ".mp4",
+        types.FileType.CAMM: ".mp4",
+        types.FileType.GOPRO: ".mp4",
+        types.FileType.VIDEO: ".mp4",
+    }
+
+    return f"mly_tools_{upload_md5sum}{_SUFFIX_MAP[filetype]}"