Add support for zarr multipart upload

DEVELOPMENT.md

@@ -337,6 +337,11 @@ New markers must be registered in `pytest_configure()` in
 - `DANDI_TESTS_NO_VCR` — When set, the use of vcrpy to playback captured HTTP
   requests during testing will be disabled
 
+- `DANDI_TESTS_OVERSIZED_ZARR` -- When set to a non-empty value, enables the
+  tests that upload a Zarr entry larger than `S3_MAX_SINGLE_PART_UPLOAD`
+  (5 GiB).  These write several GiB to disk, take several minutes, and require
+  an archive that supports multipart Zarr upload, so they are opt-in.
+
 - `DANDI_TESTS_INSTANCE_NAME` -- Sets the instance name for the dandi-archive instance used for
   testing. Defaults to `"DANDI"`. Useful for testing dandi-cli against a dandi-archive instance with
   a particular vendor information.

dandi/cli/cmd_digest.py

@@ -11,7 +11,15 @@
     "--digest",
     "digest_alg",
     type=click.Choice(
-        ["dandi-etag", "md5", "sha1", "sha256", "sha512", "zarr-checksum"],
+        [
+            "dandi-etag",
+            "md5",
+            "sha1",
+            "sha256",
+            "sha512",
+            "zarr-checksum",
+            "zarr-checksum-multipart",
+        ],
         case_sensitive=False,
     ),
     default="dandi-etag",
@@ -21,7 +29,15 @@
 @click.argument("paths", nargs=-1, type=click.Path(exists=True))
 @map_to_click_exceptions
 def digest(paths: tuple[str, ...], digest_alg: str) -> None:
-    """Calculate file digests"""
+    """Calculate file digests
+
+    A Zarr's checksum depends on the scheme its entries were uploaded with, so
+    the two schemes are named apart: use "zarr-checksum" for a Zarr uploaded
+    via single-part PUTs and "zarr-checksum-multipart" for one uploaded via S3
+    multipart upload, which is the scheme `dandi upload` uses for new Zarrs.
+
+    Example: dandi digest --digest zarr-checksum-multipart sample.zarr
+    """
     # Avoid heavy import by importing within function:
     from ..support.digests import get_digest
 

dandi/cli/cmd_ls.py

@@ -371,7 +371,10 @@ def fn():
                             digest = "0" * 32 + "-0--0"
                         else:
                             lgr.info("Calculating digest for %s", path)
-                            digest = get_digest(path, digest="zarr-checksum")
+                            # `dandi upload` creates new Zarrs with multipart
+                            # upload, so report the checksum this Zarr would
+                            # have in the archive.
+                            digest = get_digest(path, digest="zarr-checksum-multipart")
                         rec = get_metadata(path, Digest.dandi_zarr(digest))
                     else:
                         if use_fake_digest:

dandi/cli/tests/test_digest.py

@@ -106,3 +106,49 @@ def test_digest_zarr_with_excluded_dotfiles(
     r = runner.invoke(digest, ["--digest", "zarr-checksum", "sample.zarr"])
     assert r.exit_code == 0
     assert r.output == f"sample.zarr: {expected}\n"
+
+
+@pytest.mark.ai_generated
+def test_digest_zarr_multipart_differs_from_singlepart(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """
+    The two Zarr checksums are distinct digests of the same content, since S3
+    stores an entry under a plain MD5 ETag for a single-part upload and under a
+    multipart ETag for a multipart one.
+    """
+    runner = CliRunner()
+    monkeypatch.chdir(tmp_path)
+    dt = np.dtype("<i8")
+    zarr.save(
+        "sample.zarr", np.arange(1000, dtype=dt), np.arange(1000, 0, -1, dtype=dt)
+    )
+    singlepart = runner.invoke(digest, ["--digest", "zarr-checksum", "sample.zarr"])
+    assert singlepart.exit_code == 0
+    multipart = runner.invoke(
+        digest, ["--digest", "zarr-checksum-multipart", "sample.zarr"]
+    )
+    assert multipart.exit_code == 0
+    assert multipart.output != singlepart.output
+    # The single-part checksum is the one keyed by serialisation format in the
+    # tests above; both are Zarr checksums of the same shape.
+    expected = _EXPECTED_SAMPLE_ZARR_DIGEST_BY_FORMAT[
+        zarr_format_of(Path("sample.zarr"))
+    ]
+    assert singlepart.output == f"sample.zarr: {expected}\n"
+
+
+@pytest.mark.ai_generated
+def test_digest_zarr_multipart_of_file(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """
+    Applied to a single file, the multipart Zarr checksum is that entry's own
+    multipart ETag, as `dandi upload` would store it.
+    """
+    runner = CliRunner()
+    monkeypatch.chdir(tmp_path)
+    Path("file.txt").write_bytes(b"123")
+    r = runner.invoke(digest, ["--digest", "zarr-checksum-multipart", "file.txt"])
+    assert r.exit_code == 0
+    assert r.output == "file.txt: d022646351048ac0ba397d12dfafa304-1\n"

dandi/consts.py

@@ -204,10 +204,22 @@ def urls(self) -> Iterator[str]:
 ZARR_MIME_TYPE = "application/x-zarr"
 
 #: Maximum file size for a single S3 PUT upload (5 GiB).
-#: S3 rejects single-part PUTs larger than this; such files would need
-#: multipart upload which is not yet supported for zarr chunks.
+#: S3 rejects single-part PUTs larger than this.  A Zarr that contains any
+#: entry above this size must therefore be uploaded via S3 multipart upload;
+#: the archive records this per Zarr in an immutable ``upload_type`` field, set
+#: when the Zarr is created.  All entries of a ``multipart`` Zarr are uploaded
+#: via multipart upload (and digested with their S3 multipart ETag), while all
+#: entries of a ``singlepart`` Zarr are uploaded via single-part PUT (and
+#: digested with plain MD5); the two schemes cannot be mixed within one Zarr,
+#: since its checksum is an aggregate over per-entry S3 ETags.
 S3_MAX_SINGLE_PART_UPLOAD = 5 * 1024**3
 
+#: Values of a Zarr's ``upload_type`` field in the archive API, recording
+#: whether its entries are uploaded via single-part PUT or S3 multipart upload.
+#: The scheme is fixed when the Zarr is created (see `S3_MAX_SINGLE_PART_UPLOAD`).
+ZARR_UPLOAD_TYPE_SINGLEPART = "singlepart"
+ZARR_UPLOAD_TYPE_MULTIPART = "multipart"
+
 #: Maximum number of Zarr directory entries to upload at once
 ZARR_UPLOAD_BATCH_SIZE = 255
 

dandi/dandiapi.py

@@ -2239,12 +2239,24 @@ def from_server_data(
         cls, asset: BaseRemoteZarrAsset, data: ZarrEntryServerData
     ) -> RemoteZarrEntry:
         """:meta private:"""
+        # Avoid heavy import by importing within function:
+        from dandi.support.digests import is_multipart_etag
+
+        # An entry's digest is the ETag S3 stores it under, which is a plain MD5
+        # for a single-part upload and a multipart ETag (a.k.a. the DANDI etag)
+        # for a multipart one; the algorithm has to follow suit, or the entry
+        # cannot be verified on download.
+        algorithm = (
+            models.DigestType.dandi_etag
+            if is_multipart_etag(data.etag)
+            else models.DigestType.md5
+        )
         return cls(
             client=asset.client,
             zarr_id=asset.zarr,
             parts=tuple(data.key.split("/")),
             modified=data.last_modified,
-            digest=Digest(algorithm=models.DigestType.md5, value=data.etag),
+            digest=Digest(algorithm=algorithm, value=data.etag),
             size=data.size,
         )
 

dandi/download.py

@@ -788,10 +788,13 @@ def _download_file(
         # TODO: reuse that sorting based on speed
         for algo, digest in digests.items():
             if algo == "dandi-etag" and size is not None:
-                # Instantiate outside the lambda so that mypy is assured that
-                # `size` is not None:
-                hasher = ETagHashlike(size)
-                digester = lambda: hasher  # noqa: E731
+                # Bind `size` to a local so that mypy is assured it is not
+                # None.  The lambda must construct a fresh `ETagHashlike` on
+                # each call, as the hashlib branch does: a download attempt
+                # that is retried starts the hashing over, and an
+                # `ETagHashlike` fed past `size` bytes raises `ValueError`.
+                etag_size = size
+                digester = lambda: ETagHashlike(etag_size)  # noqa: E731
             else:
                 digester = getattr(hashlib, algo, None)
             if digester is not None:
@@ -1043,7 +1046,7 @@ def _download_zarr(
     zarr_entry_filter: Callable[[str], bool] | None = None,
 ) -> Iterator[dict]:
     # Avoid heavy import by importing within function:
-    from .support.digests import get_zarr_checksum
+    from .support.digests import get_zarr_checksum, get_zarr_multipart_checksum
 
     # we will collect them all while starting the download
     # with the first page of entries received from the server.
@@ -1052,7 +1055,10 @@ def _download_zarr(
     pc = ProgressCombiner(zarr_size=asset.size)
 
     def digest_callback(path: str, algoname: str, d: str) -> None:
-        if algoname == "md5":
+        # A Zarr's entries are all digested the same way -- with plain MD5 for a
+        # single-part Zarr and with the multipart ETag for a multipart one --
+        # and either is what the Zarr's checksum is computed from.
+        if algoname in ("md5", "dandi-etag"):
             digests[path] = d
 
     def downloads_gen():
@@ -1062,7 +1068,12 @@ def downloads_gen():
                 continue
             entries.append(entry)
             etag = entry.digest
-            assert etag.algorithm is DigestType.md5
+            # An entry of a multipart Zarr is stored under an S3 multipart ETag
+            # rather than a plain MD5 (see `RemoteZarrEntry.from_server_data`).
+            assert etag.algorithm in (DigestType.md5, DigestType.dandi_etag)
+            etag_algo = (
+                "dandi-etag" if etag.algorithm is DigestType.dandi_etag else "md5"
+            )
             yield pairing(
                 entry_path,
                 _download_file(
@@ -1072,7 +1083,7 @@ def downloads_gen():
                     size=entry.size,
                     mtime=entry.modified,
                     existing=existing,
-                    digests={"md5": etag.value},
+                    digests={etag_algo: etag.value},
                     lock=lock,
                     digest_callback=partial(digest_callback, entry_path),
                 ),
@@ -1152,7 +1163,16 @@ def downloads_gen():
 
         if "skipped" not in final_out["message"]:
             zarr_checksum = asset.get_digest().value
-            local_checksum = get_zarr_checksum(zarr_basepath, known=digests)
+            # A multipart Zarr's checksum aggregates its entries' multipart
+            # ETags, so it has to be recomputed the same way; every entry of a
+            # Zarr uses the one scheme the Zarr was created with.
+            multipart = any(
+                e.digest.algorithm is DigestType.dandi_etag for e in entries
+            )
+            checksummer = (
+                get_zarr_multipart_checksum if multipart else get_zarr_checksum
+            )
+            local_checksum = checksummer(zarr_basepath, known=digests)
             if zarr_checksum != local_checksum:
                 msg = f"Zarr checksum: downloaded {local_checksum} != {zarr_checksum}"
                 yield {"checksum": "differs", "status": "error", "message": msg}

dandi/exceptions.py

@@ -93,6 +93,21 @@ class UploadError(Exception):
     pass
 
 
+class BlobExistsError(UploadError):
+    """
+    Raised when the archive reports, via an HTTP 409 from the ``initialize``
+    endpoint of a multipart upload, that the blob being uploaded is already
+    present.  Only that endpoint identifies the existing blob; a 409 from any
+    other point of an upload is an ordinary error and propagates as such.
+    """
+
+    def __init__(self, blob_id: str) -> None:
+        super().__init__(f"Blob already exists on server with ID {blob_id}")
+        #: The ID of the pre-existing blob, from the response's ``Location``
+        #: header
+        self.blob_id = blob_id
+
+
 class UploadValidationError(UploadError):
     """An upload could not proceed because an asset failed validation."""