Add support for zarr multipart upload

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

@@ -203,10 +203,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 @@
         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

@@ -1043,7 +1043,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 +1052,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 +1065,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 +1080,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 +1160,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}