Use multi-part upload on files larger than 5GB

Author: jjnesbitt

dandi/consts.py

@@ -199,6 +199,9 @@ def urls(self) -> Iterator[str]:
 #: Maximum number of Zarr directory entries to delete at once
 ZARR_DELETE_BATCH_SIZE = 100
 
+#: Zarr chunks above this size (bytes) are uploaded via multipart upload
+ZARR_LARGE_CHUNK_THRESHOLD = 5 * 1024 * 1024 * 1024  # 5 GB
+
 BIDS_DATASET_DESCRIPTION = "dataset_description.json"
 
 BIDS_IGNORE_FILE = ".bidsignore"

dandi/files/zarr.py

@@ -10,10 +10,13 @@
 import json
 import os
 import os.path
+import re
 from pathlib import Path
+from threading import Lock
 from time import sleep
 from typing import Any, Optional
 import urllib.parse
+from xml.etree.ElementTree import fromstring
 
 from dandischema.models import BareAsset, DigestType
 from pydantic import BaseModel, ConfigDict, ValidationError
@@ -25,6 +28,7 @@
 from dandi.consts import (
     MAX_ZARR_DEPTH,
     ZARR_DELETE_BATCH_SIZE,
+    ZARR_LARGE_CHUNK_THRESHOLD,
     ZARR_MIME_TYPE,
     ZARR_UPLOAD_BATCH_SIZE,
 )
@@ -46,7 +50,7 @@
     pre_upload_size_check,
 )
 
-from .bases import LocalDirectoryAsset
+from .bases import LocalDirectoryAsset, _upload_blob_part
 from ..validate._types import (
     ORIGIN_VALIDATION_DANDI_ZARR,
     Origin,
@@ -742,13 +746,40 @@ def mkzarr() -> str:
                 ):
                     # Items to upload in this batch (may be retried e.g. due to
                     # 403 errors because of timed-out upload URLs)
-                    items_to_upload = list(items)
+                    all_items = list(items)
+                    large_items = [
+                        it for it in all_items if it.size > ZARR_LARGE_CHUNK_THRESHOLD
+                    ]
+                    items_to_upload = [
+                        it for it in all_items if it.size <= ZARR_LARGE_CHUNK_THRESHOLD
+                    ]
                     max_retries = 5
                     retry_count = 0
                     # Add all items to checksum tree (only done once)
-                    for it in items_to_upload:
+                    for it in all_items:
                         zcc.add_leaf(Path(it.entry_path), it.size, it.digest)
 
+                    # Upload chunks above 5GB individually via multipart upload
+                    for it in large_items:
+                        for status in upload_zarr_file_multipart(
+                            item=it,
+                            zarr_id=zarr_id,
+                            dandiset=dandiset,
+                            jobs=jobs,
+                        ):
+                            if status.get("status") == "done":
+                                changed = True
+                                bytes_uploaded += it.size
+                                yield {
+                                    "status": "uploading",
+                                    "progress": 100
+                                    * bytes_uploaded
+                                    / to_upload.total_size,
+                                    "current": bytes_uploaded,
+                                }
+                            else:
+                                yield status
+
                     while items_to_upload and retry_count <= max_retries:
                         # Prepare upload requests for current items
                         uploading = [it.upload_request() for it in items_to_upload]
@@ -903,6 +934,115 @@ def _handle_failed_items_and_raise(
     raise failed_items[0][1]
 
 
+def upload_zarr_file_multipart(
+    item: UploadItem,
+    zarr_id: str,
+    dandiset: RemoteDandiset,
+    jobs: int | None = None,
+):
+    # Avoid heavy import by importing within function:
+    from dandi.support.digests import get_dandietag
+
+    client = dandiset.client
+
+    yield {"status": "calculating etag"}
+    etagger = get_dandietag(item.filepath)
+    filetag = etagger.as_str()
+
+    yield {"status": "initiating upload"}
+    lgr.debug("%s: Beginning upload", item.filepath)
+    total_size = pre_upload_size_check(item.filepath)
+
+    resp = client.post(
+        "/uploads/zarr/initialize/",
+        json={
+            "contentSize": total_size,
+            "digest": {
+                "algorithm": "dandi:dandi-etag",
+                "value": filetag,
+            },
+            "zarr": {
+                "chunk_key": item.entry_path,
+                "zarr_id": zarr_id,
+            },
+        },
+    )
+
+    try:
+        upload_id = resp["upload_id"]
+        parts = resp["parts"]
+        if len(parts) != etagger.part_qty:
+            raise RuntimeError(
+                f"Server and client disagree on number of parts for upload;"
+                f" server says {len(parts)}, client says {etagger.part_qty}"
+            )
+        parts_out = []
+        bytes_uploaded = 0
+        lgr.debug("Uploading %s in %d parts", item.filepath, len(parts))
+        with RESTFullAPIClient("http://nil.nil") as storage:
+            with item.filepath.open("rb") as fp:
+                with ThreadPoolExecutor(max_workers=jobs or 5) as executor:
+                    lock = Lock()
+                    futures = [
+                        executor.submit(
+                            _upload_blob_part,
+                            storage_session=storage,
+                            fp=fp,
+                            lock=lock,
+                            etagger=etagger,
+                            asset_path=item.entry_path,
+                            part=part,
+                        )
+                        for part in parts
+                    ]
+                    for fut in as_completed(futures):
+                        out_part = fut.result()
+                        bytes_uploaded += out_part["size"]
+                        yield {
+                            "status": "uploading",
+                            "progress": 100 * bytes_uploaded / total_size,
+                            "current": bytes_uploaded,
+                        }
+                        parts_out.append(out_part)
+
+            lgr.debug("%s: Completing upload", item.entry_path)
+            resp = client.post(
+                f"/uploads/zarr/{upload_id}/complete/",
+                json={"parts": parts_out},
+            )
+            lgr.debug(
+                "%s: Announcing completion to %s",
+                item.entry_path,
+                resp["complete_url"],
+            )
+            r = storage.post(resp["complete_url"], data=resp["body"], json_resp=False)
+            lgr.debug(
+                "%s: Upload completed. Response content: %s",
+                item.entry_path,
+                r.content,
+            )
+            rxml = fromstring(r.text)
+            m = re.match(r"\{.+?\}", rxml.tag)
+            ns = m.group(0) if m else ""
+            final_etag = rxml.findtext(f"{ns}ETag")
+            if final_etag is not None:
+                final_etag = final_etag.strip('"')
+                if final_etag != filetag:
+                    raise RuntimeError(
+                        "Server and client disagree on final ETag of"
+                        f" uploaded file; server says {final_etag},"
+                        f" client says {filetag}"
+                    )
+            # else: Error? Warning?
+            resp = client.post(f"/uploads/zarr/{upload_id}/validate/")
+            yield {"status": "done"}
+    except Exception:
+        post_upload_size_check(item.filepath, total_size, True)
+        raise
+    else:
+        post_upload_size_check(item.filepath, total_size, False)
+
+
 def _upload_zarr_file(
     storage_session: RESTFullAPIClient,
     dandiset: RemoteDandiset,

dandi/tests/test_files.py

@@ -1,17 +1,16 @@
 from __future__ import annotations
 
-from operator import attrgetter
 import os
-from pathlib import Path
 import subprocess
-from unittest.mock import ANY
+from operator import attrgetter
+from pathlib import Path
+from unittest.mock import ANY, patch
 
-from dandischema.models import get_schema_version
 import numpy as np
 import pytest
 import zarr
+from dandischema.models import get_schema_version
 
-from .fixtures import SampleDandiset
 from .. import get_logger
 from ..consts import ZARR_MIME_TYPE, dandiset_metadata_file
 from ..dandiapi import AssetType, RemoteZarrAsset
@@ -29,6 +28,7 @@
     dandi_file,
     find_dandi_files,
 )
+from .fixtures import SampleDandiset
 
 lgr = get_logger()
 
@@ -536,6 +536,39 @@ def test_upload_zarr_entry_content_type(new_dandiset, tmp_path):
     assert r.headers["Content-Type"] == "application/json"
 
 
+@pytest.mark.ai_generated
+def test_upload_zarr_large_chunks(new_dandiset, tmp_path):
+    """Chunks above ZARR_LARGE_CHUNK_THRESHOLD are uploaded via upload_zarr_file_multipart."""
+    filepath = tmp_path / "example.zarr"
+    zarr.save(filepath, np.arange(1000), np.arange(1000, 0, -1))
+    zf = dandi_file(filepath)
+    assert isinstance(zf, ZarrAsset)
+
+    from ..files.zarr import upload_zarr_file_multipart
+
+    real_upload_zarr_file_multipart = upload_zarr_file_multipart
+    called_paths: list[str] = []
+
+    def spy_upload_zarr_file_multipart(item, *args, **kwargs):
+        called_paths.append(item.entry_path)
+        yield from real_upload_zarr_file_multipart(item, *args, **kwargs)
+
+    # Set threshold to 0 so every chunk is treated as "large"
+    with (
+        patch("dandi.files.zarr.ZARR_LARGE_CHUNK_THRESHOLD", 0),
+        patch(
+            "dandi.files.zarr.upload_zarr_file_multipart",
+            spy_upload_zarr_file_multipart,
+        ),
+    ):
+        asset = zf.upload(new_dandiset.dandiset, {})
+
+    assert isinstance(asset, RemoteZarrAsset)
+    # Every chunk file in the zarr should have been routed through upload_zarr_file_multipart
+    remote_entries = {str(e) for e in asset.iterfiles()}
+    assert remote_entries == set(called_paths)
+
+
 def test_validate_deep_zarr(tmp_path: Path) -> None:
     zarr_path = tmp_path / "foo.zarr"
     zarr.save(zarr_path, np.arange(1000), np.arange(1000, 0, -1))