diff --git a/dandi/cli/cmd_digest.py b/dandi/cli/cmd_digest.py index 639f1bda9..43572ee58 100644 --- a/dandi/cli/cmd_digest.py +++ b/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 diff --git a/dandi/cli/cmd_ls.py b/dandi/cli/cmd_ls.py index 1969236ad..f016e3ded 100644 --- a/dandi/cli/cmd_ls.py +++ b/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: diff --git a/dandi/cli/tests/test_digest.py b/dandi/cli/tests/test_digest.py index 9c8fcccbf..64d62ac86 100644 --- a/dandi/cli/tests/test_digest.py +++ b/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(" 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" diff --git a/dandi/consts.py b/dandi/consts.py index baca5465c..d42c08c5c 100644 --- a/dandi/consts.py +++ b/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 diff --git a/dandi/dandiapi.py b/dandi/dandiapi.py index ccf17239b..fe7dd18c6 100644 --- a/dandi/dandiapi.py +++ b/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, ) diff --git a/dandi/download.py b/dandi/download.py index a87d5380e..46e494e6f 100644 --- a/dandi/download.py +++ b/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} diff --git a/dandi/files/bases.py b/dandi/files/bases.py index db502dcc3..baca238fe 100644 --- a/dandi/files/bases.py +++ b/dandi/files/bases.py @@ -2,7 +2,7 @@ from abc import ABC, abstractmethod from collections import deque -from collections.abc import Iterator +from collections.abc import Generator, Iterator from concurrent.futures import ThreadPoolExecutor, as_completed from dataclasses import dataclass from datetime import datetime @@ -364,37 +364,17 @@ def iter_upload( ``"done"`` and an ``"asset"`` key containing the resulting `RemoteAsset`. """ - # Avoid heavy import by importing within function: - from dandi.support.digests import get_dandietag - asset_path = metadata.setdefault("path", self.path) set_asset_schema_key(metadata) client = dandiset.client - yield {"status": "calculating etag"} - etagger = get_dandietag(self.filepath) - filetag = etagger.as_str() - lgr.debug("Calculated dandi-etag of %s for %s", filetag, self.filepath) - digest = metadata.get("digest", {}) - if "dandi:dandi-etag" in digest: - if digest["dandi:dandi-etag"] != filetag: - raise RuntimeError( - f"{self.filepath}: File etag changed; was originally" - f" {digest['dandi:dandi-etag']} but is now {filetag}" - ) - yield {"status": "initiating upload"} - lgr.debug("%s: Beginning upload", asset_path) - total_size = pre_upload_size_check(self.filepath) try: - resp = client.post( - "/uploads/initialize/", - json={ - "contentSize": total_size, - "digest": { - "algorithm": "dandi:dandi-etag", - "value": filetag, - }, - "dandiset": dandiset.identifier, - }, + resp = yield from multipart_upload( + client=client, + filepath=self.filepath, + asset_path=asset_path, + init_fields={"dandiset": dandiset.identifier}, + expected_etag=metadata.get("digest", {}).get("dandi:dandi-etag"), + jobs=jobs, ) except requests.HTTPError as e: if e.response is not None and e.response.status_code == 409: @@ -403,80 +383,7 @@ def iter_upload( else: raise else: - 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", self.filepath, len(parts)) - with RESTFullAPIClient("http://nil.nil") as storage: - with self.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=asset_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", asset_path) - resp = client.post( - f"/uploads/{upload_id}/complete/", - json={"parts": parts_out}, - ) - lgr.debug( - "%s: Announcing completion to %s", - asset_path, - resp["complete_url"], - ) - r = storage.post( - resp["complete_url"], data=resp["body"], json_resp=False - ) - lgr.debug( - "%s: Upload completed. Response content: %s", - asset_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/{upload_id}/validate/") - blob_id = resp["blob_id"] - except Exception: - post_upload_size_check(self.filepath, total_size, True) - raise - else: - post_upload_size_check(self.filepath, total_size, False) + blob_id = resp["blob_id"] lgr.debug("%s: Assigning asset blob to dandiset & version", asset_path) yield {"status": "producing asset"} if replacing is not None: @@ -676,6 +583,30 @@ def size(self) -> int: return sum(p.size for p in self.iterfiles()) +def _retry_s3_upload(r: requests.Response) -> bool: + """ + Return whether a failed S3 upload request should be retried as-is. These + are transient conditions that a fresh attempt at the same presigned URL can + resolve; an expired URL (403), by contrast, has to be re-signed and so is + handled by the caller. + + :meta private: + """ + return ( + # Some sort of filesystem hiccup can cause requests to be unable to get + # the filesize, leading to it falling back to "chunked" transfer + # encoding, which S3 doesn't support. + r.status_code == 501 + and "header you provided implies functionality that is not implemented" + in r.text + ) or ( + # Network issue or rate limiting can cause a timeout, which results in a + # 400. Case: https://github.com/dandi/dandi-cli/issues/1662 + r.status_code == 400 + and "was not read from or written to within the timeout period" in r.text + ) + + def _upload_blob_part( storage_session: RESTFullAPIClient, fp: IO[bytes], @@ -711,7 +642,18 @@ def _upload_blob_part( data=chunk, json_resp=False, retry_statuses=[500], + retry_if=_retry_s3_upload, ) + if "ETag" not in r.headers: + # A part upload that reports no ETag cannot be completed, as the ETag is + # what the completion request identifies the part by. Raise it as an + # HTTP error carrying the response, so that the caller can decide from + # the response whether the condition is worth retrying. + raise requests.HTTPError( + f"{asset_path}: upload of part {part['part_number']} returned no" + f" ETag (status {r.status_code})", + response=r, + ) server_etag = r.headers["ETag"].strip('"') lgr.debug( "%s: Part upload finished ETag=%s Content-Length=%s", @@ -733,6 +675,131 @@ def _upload_blob_part( } +def multipart_upload( + client: RESTFullAPIClient, + filepath: Path, + asset_path: str, + init_fields: dict[str, Any], + expected_etag: str | None = None, + jobs: int | None = None, + upload_root: str = "/uploads", +) -> Generator[dict, None, dict]: + """ + Upload ``filepath`` to the archive via the S3 multipart upload API, + yielding progress `dict`\\s and returning the deserialized response of the + ``validate`` endpoint. + + Asset blobs and Zarr chunks have separate multipart upload endpoints, so + the caller selects one via ``upload_root`` — ``"/uploads"`` for an asset + blob or ``"/zarr/uploads"`` for a Zarr chunk — and supplies the matching + ``init_fields`` identifying what is being uploaded to the ``initialize`` + endpoint: ``{"dandiset": ...}`` for an asset blob or + ``{"zarr_id": ..., "chunk_key": ...}`` for a Zarr chunk. + + If ``expected_etag`` is non-`None` and does not match the etag computed for + ``filepath``, `RuntimeError` is raised. An HTTP 409 from ``initialize`` + (i.e., the blob already exists) propagates to the caller. + + :meta private: + """ + # Avoid heavy import by importing within function: + from dandi.support.digests import get_dandietag + + yield {"status": "calculating etag"} + etagger = get_dandietag(filepath) + filetag = etagger.as_str() + lgr.debug("Calculated dandi-etag of %s for %s", filetag, filepath) + if expected_etag is not None and expected_etag != filetag: + raise RuntimeError( + f"{filepath}: File etag changed; was originally" + f" {expected_etag} but is now {filetag}" + ) + yield {"status": "initiating upload"} + lgr.debug("%s: Beginning upload", asset_path) + total_size = pre_upload_size_check(filepath) + resp = client.post( + f"{upload_root}/initialize/", + json={ + "contentSize": total_size, + "digest": {"algorithm": "dandi:dandi-etag", "value": filetag}, + **init_fields, + }, + ) + 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", filepath, len(parts)) + with RESTFullAPIClient("http://nil.nil") as storage: + with 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=asset_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", asset_path) + resp = client.post( + f"{upload_root}/{upload_id}/complete/", + json={"parts": parts_out}, + ) + lgr.debug( + "%s: Announcing completion to %s", + asset_path, + resp["complete_url"], + ) + r = storage.post(resp["complete_url"], data=resp["body"], json_resp=False) + lgr.debug( + "%s: Upload completed. Response content: %s", + asset_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? + validated = client.post(f"{upload_root}/{upload_id}/validate/") + except Exception: + post_upload_size_check(filepath, total_size, True) + raise + else: + post_upload_size_check(filepath, total_size, False) + assert isinstance(validated, dict) + return validated + + def _check_required_fields( d: dict, required: list[str], file_path: str ) -> list[ValidationResult]: diff --git a/dandi/files/zarr.py b/dandi/files/zarr.py index 25394769e..1d5bf4847 100644 --- a/dandi/files/zarr.py +++ b/dandi/files/zarr.py @@ -2,12 +2,13 @@ from base64 import b64encode from collections import Counter -from collections.abc import Generator, Iterator +from collections.abc import Callable, Generator, Iterator from concurrent.futures import Future, ThreadPoolExecutor, as_completed from contextlib import closing from dataclasses import dataclass, field, replace from datetime import datetime from enum import Enum +from functools import partial import json import math import os @@ -21,7 +22,7 @@ if TYPE_CHECKING: from ..upload import ZarrMode -from dandischema.models import BareAsset, DigestType +from dandischema.models import BareAsset from pydantic import BaseModel, ConfigDict, ValidationError import requests from zarr_checksum.tree import ZarrChecksumTree @@ -34,6 +35,7 @@ ZARR_DELETE_BATCH_SIZE, ZARR_MIME_TYPE, ZARR_UPLOAD_BATCH_SIZE, + ZARR_UPLOAD_TYPE_MULTIPART, ) from dandi.dandiapi import ( RemoteAsset, @@ -55,7 +57,7 @@ pre_upload_size_check, ) -from .bases import LocalDirectoryAsset +from .bases import LocalDirectoryAsset, _retry_s3_upload, multipart_upload from ..validate._types import ( ORIGIN_VALIDATION_DANDI_ZARR, MissingFileContent, @@ -349,19 +351,18 @@ def iterdir(self) -> Iterator[LocalZarrEntry]: def get_digest(self) -> Digest: """ - Calculate the DANDI etag digest for the entry. If the entry is a - directory, the algorithm will be the DANDI Zarr checksum algorithm; if - it is a file, it will be MD5. + Calculate the digest of the entry as it would be stored in the archive. + If the entry is a directory, the algorithm is the DANDI Zarr checksum + algorithm; if it is a file, it is the S3 multipart ETag (DANDI etag), + matching how `dandi upload` stores an entry of a (multipart) Zarr. """ # Avoid heavy import by importing within function: - from dandi.support.digests import get_digest, get_zarr_checksum + from dandi.support.digests import get_digest, get_zarr_multipart_checksum if self.is_dir(): - return Digest.dandi_zarr(get_zarr_checksum(self.filepath)) + return Digest.dandi_zarr(get_zarr_multipart_checksum(self.filepath)) else: - return Digest( - algorithm=DigestType.md5, value=get_digest(self.filepath, "md5") - ) + return Digest.dandi_etag(get_digest(self.filepath, "dandi-etag")) @property def size(self) -> int: @@ -424,11 +425,12 @@ def filetree(self) -> LocalZarrEntry: def stat(self) -> ZarrStat: """Return various details about the Zarr asset""" + # Avoid heavy import by importing within function: + from dandi.support.digests import checksum_zarr_dir, dandietag_nocache + # Digest entries as they would be uploaded: `dandi upload` creates new + # Zarrs with multipart upload, so every entry gets its multipart ETag. def dirstat(dirpath: LocalZarrEntry) -> ZarrStat: - # Avoid heavy import by importing within function: - from dandi.support.digests import checksum_zarr_dir, md5file_nocache - size = 0 dir_info = {} file_info = {} @@ -441,7 +443,7 @@ def dirstat(dirpath: LocalZarrEntry) -> ZarrStat: files.extend(st.files) else: size += p.size - file_info[p.name] = (md5file_nocache(p.filepath), p.size) + file_info[p.name] = (dandietag_nocache(p.filepath), p.size) files.append(p) return ZarrStat( size=size, @@ -454,9 +456,11 @@ def dirstat(dirpath: LocalZarrEntry) -> ZarrStat: def get_digest(self) -> Digest: """Calculate a dandi-zarr-checksum digest for the asset""" # Avoid heavy import by importing within function: - from dandi.support.digests import get_zarr_checksum + from dandi.support.digests import get_zarr_multipart_checksum - return Digest.dandi_zarr(get_zarr_checksum(self.filepath)) + # `dandi upload` creates new Zarrs with multipart upload, so report the + # checksum the asset will have once uploaded. + return Digest.dandi_zarr(get_zarr_multipart_checksum(self.filepath)) def get_metadata( self, @@ -594,11 +598,29 @@ def iter_upload( lgr.debug("%s: Producing asset", asset_path) yield {"status": "producing asset"} - def mkzarr() -> str: + # Avoid heavy import by importing within function: + from dandi.support.digests import zarr_has_oversized_entry + + # New Zarrs are created with multipart upload enabled. The archive + # records the scheme per Zarr in an immutable ``upload_type`` field set + # at creation time; all of a Zarr's entries must use the same scheme, + # since its checksum is an aggregate over per-entry S3 ETags. A + # pre-existing Zarr keeps whatever scheme it was created with. An entry + # too large for a single-part S3 PUT *requires* multipart upload, so + # such content cannot be uploaded to a single-part Zarr (a pre-existing + # one, or any Zarr on an archive predating the field, which always + # creates single-part Zarrs). + needs_multipart = zarr_has_oversized_entry(self.filepath) + + def mkzarr() -> tuple[str, bool]: try: r = client.post( "/zarr/", - json={"name": asset_path, "dandiset": dandiset.identifier}, + json={ + "name": asset_path, + "dandiset": dandiset.identifier, + "upload_type": ZARR_UPLOAD_TYPE_MULTIPART, + }, ) except requests.HTTPError as e: if e.response is not None and ( @@ -618,12 +640,17 @@ def mkzarr() -> str: }, ) zarr_id = old_zarr["zarr_id"] + zarr_multipart = _zarr_is_multipart(old_zarr) else: raise else: zarr_id = r["zarr_id"] + # Honour the scheme the server actually recorded: an archive + # that predates the ``upload_type`` field ignores it and always + # creates a single-part Zarr. + zarr_multipart = _zarr_is_multipart(r) assert isinstance(zarr_id, str) - return zarr_id + return zarr_id, zarr_multipart if replacing is not None: lgr.debug("%s: Replacing pre-existing asset", asset_path) @@ -632,22 +659,33 @@ def mkzarr() -> str: "%s: Pre-existing asset is a Zarr; reusing & updating", asset_path ) zarr_id = replacing.zarr + # The reused Zarr's upload scheme was fixed when it was created. + multipart = _zarr_is_multipart(client.get(f"/zarr/{zarr_id}/")) else: lgr.debug( "%s: Pre-existing asset is not a Zarr; minting new Zarr", asset_path ) - zarr_id = mkzarr() + zarr_id, multipart = mkzarr() r = client.put( replacing.api_path, json={"metadata": metadata, "zarr_id": zarr_id}, ) else: lgr.debug("%s: Minting new Zarr", asset_path) - zarr_id = mkzarr() + zarr_id, multipart = mkzarr() r = client.post( f"{dandiset.version_api_path}assets/", json={"metadata": metadata, "zarr_id": zarr_id}, ) + + if needs_multipart and not multipart: + raise UploadError( + f"{asset_path}: this Zarr contains an entry larger than" + f" {S3_MAX_SINGLE_PART_UPLOAD / 1024**3:.0f} GiB and so requires" + f" multipart upload, but the target Zarr does not support it." + f" The archive may not support multipart Zarr upload, or the" + f" Zarr being replaced was created as single-part." + ) a = RemoteAsset.from_data(dandiset, r) assert isinstance(a, RemoteZarrAsset) mismatched = True @@ -659,7 +697,7 @@ def mkzarr() -> str: str(e): e for e in a.iterfiles() } total_size = 0 - to_upload = EntryUploadTracker() + to_upload = EntryUploadTracker(multipart=multipart) if old_zarr_entries: to_delete: list[RemoteZarrEntry] = [] digesting: list[Future[tuple[LocalZarrEntry, str, bool]]] = [] @@ -713,6 +751,7 @@ def mkzarr() -> str: asset_path, local_entry, remote_entry.digest.value, + multipart, ) ) for dgstfut in as_completed(digesting): @@ -780,111 +819,54 @@ def mkzarr() -> str: closing(to_upload.get_items()) as upload_items, ): bytes_uploaded = 0 + # A batch of entries is uploaded the same way under either + # scheme -- concurrently, with retries -- and the schemes differ + # only in how an attempt's items are handed to the worker pool. + submit_batch: _BatchSubmitter + if multipart: + submit_batch = partial( + _submit_zarr_entries, + client=client, + zarr_id=zarr_id, + jobs=jobs, + asset_path=asset_path, + ) + via = " via multipart upload" + retry_reason = "retryable errors" + else: + submit_batch = partial( + _submit_zarr_files, + client=client, + zarr_id=zarr_id, + storage=storage, + dandiset=dandiset, + ) + via = "" + retry_reason = "403 errors" for i, items in enumerate( chunked(upload_items, ZARR_UPLOAD_BATCH_SIZE), start=1 ): - # 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) - max_retries = 5 - retry_count = 0 - current_jobs = jobs or 5 + batch = list(items) # Add all items to checksum tree (only done once) - for it in items_to_upload: + for it in batch: zcc.add_leaf(Path(it.entry_path), it.size, it.digest) - 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] - - if retry_count == 0: - lgr.debug( - "%s: Uploading Zarr file batch #%d (%s)", - asset_path, - i, - pluralize(len(uploading), "file"), - ) - else: - lgr.debug( - "%s: Retrying %s from batch #%d (attempt %d/%d)", - asset_path, - pluralize(len(uploading), "file"), - i, - retry_count, - max_retries, - ) - - # Get signed URLs for items - r = client.post(f"/zarr/{zarr_id}/files/", json=uploading) - - # Upload files in parallel - with ThreadPoolExecutor(max_workers=current_jobs) as executor: - futures = [ - executor.submit( - _upload_zarr_file, - storage_session=storage, - dandiset=dandiset, - upload_url=signed_url, - item=it, - ) - for (signed_url, it) in zip(r, items_to_upload) - ] - - changed = True - retry_items = [] - failed_items = [] - - for fut in as_completed(futures): - result = fut.result() - - if result.status == UploadStatus.SUCCESS: - bytes_uploaded += result.size - yield { - "status": "uploading", - "progress": 100 - * bytes_uploaded - / to_upload.total_size, - "current": bytes_uploaded, - } - elif result.status == UploadStatus.RETRY_NEEDED: - retry_items.append(result.item) - else: - assert result.status == UploadStatus.FAILED - failed_items.append((result.item, result.error)) - - # Handle failed items (non-403 errors) - if failed_items: - _handle_failed_items_and_raise( - executor, failed_items, futures - ) - - # Prepare for next iteration with retry items - if items_to_upload := retry_items: - retry_count += 1 - current_jobs = max(1, math.ceil(current_jobs / 2)) - if retry_count <= max_retries: - lgr.info( - "%s: %s got 403 errors, requesting new URLs" - " (attempt %d/%d, workers: %d)", - asset_path, - pluralize(len(items_to_upload), "file"), - retry_count, - max_retries, - current_jobs, - ) - # Exponential backoff with jitter before retry - sleep( - min(2**retry_count * 5, 120) - + random.uniform(0, 5) - ) - - # Check if we exhausted retries - if items_to_upload: - nfiles_str = pluralize(len(items_to_upload), "file") - raise UploadError( - f"{asset_path}: failed to upload {nfiles_str} " - f"after {max_retries} retries due to repeated 403 errors" - ) + for size in _upload_zarr_batch( + asset_path=asset_path, + batch_number=i, + items=batch, + submit=submit_batch, + jobs=jobs, + via=via, + retry_reason=retry_reason, + ): + bytes_uploaded += size + changed = True + yield { + "status": "uploading", + "progress": 100 * bytes_uploaded / to_upload.total_size, + "current": bytes_uploaded, + } lgr.debug("%s: Completing upload of batch #%d", asset_path, i) lgr.debug("%s: All files uploaded", asset_path) if zarr_mode == "full": @@ -969,6 +951,232 @@ def mkzarr() -> str: yield {"status": "done", "asset": a} +#: Hands an attempt's worth of a batch's items to a worker pool, returning one +#: future per item +_BatchSubmitter = Callable[ + [ThreadPoolExecutor, list["UploadItem"]], list[Future[UploadResult]] +] + +#: Number of times a batch of Zarr entries is retried before giving up +_MAX_BATCH_RETRIES = 5 + + +def _upload_zarr_batch( + *, + asset_path: str, + batch_number: int, + items: list[UploadItem], + submit: _BatchSubmitter, + jobs: int | None, + retry_reason: str, + via: str = "", +) -> Iterator[int]: + """ + Upload one batch of Zarr entries, yielding the size of each entry as it + finishes so that the caller can report progress. + + ``submit`` hands the batch's items to a worker pool. It is called afresh + on every attempt, since a retry is generally needed because the presigned + URLs of the previous attempt timed out and so have to be re-requested. + Entries that failed for a retryable reason are re-submitted with the worker + count halved, after an exponential backoff, for up to + `_MAX_BATCH_RETRIES` attempts; any other failure aborts the batch at once. + + ``retry_reason`` names the retryable condition for log messages, and + ``via`` describes the upload scheme. + + :meta private: + """ + items_to_upload = list(items) + retry_count = 0 + current_jobs = jobs or 5 + while items_to_upload and retry_count <= _MAX_BATCH_RETRIES: + if retry_count == 0: + lgr.debug( + "%s: Uploading Zarr batch #%d (%s)%s", + asset_path, + batch_number, + pluralize(len(items_to_upload), "file"), + via, + ) + else: + lgr.debug( + "%s: Retrying %s from batch #%d (attempt %d/%d)", + asset_path, + pluralize(len(items_to_upload), "file"), + batch_number, + retry_count, + _MAX_BATCH_RETRIES, + ) + with ThreadPoolExecutor(max_workers=current_jobs) as executor: + futures = submit(executor, items_to_upload) + retry_items: list[UploadItem] = [] + failed_items: list[tuple[UploadItem, Exception | None]] = [] + try: + for fut in as_completed(futures): + result = fut.result() + if result.status == UploadStatus.SUCCESS: + yield result.size + elif result.status == UploadStatus.RETRY_NEEDED: + retry_items.append(result.item) + else: + assert result.status == UploadStatus.FAILED + failed_items.append((result.item, result.error)) + except BaseException: + for f in futures: + f.cancel() + raise + + if failed_items: + _handle_failed_items_and_raise(executor, failed_items, futures) + + if items_to_upload := retry_items: + retry_count += 1 + current_jobs = max(1, math.ceil(current_jobs / 2)) + if retry_count <= _MAX_BATCH_RETRIES: + lgr.info( + "%s: %s got %s, requesting new URLs" + " (attempt %d/%d, workers: %d)", + asset_path, + pluralize(len(items_to_upload), "file"), + retry_reason, + retry_count, + _MAX_BATCH_RETRIES, + current_jobs, + ) + # Exponential backoff with jitter before retry + sleep(min(2**retry_count * 5, 120) + random.uniform(0, 5)) + + if items_to_upload: + nfiles_str = pluralize(len(items_to_upload), "file") + raise UploadError( + f"{asset_path}: failed to upload {nfiles_str} after" + f" {_MAX_BATCH_RETRIES} retries due to repeated {retry_reason}" + ) + + +def _submit_zarr_files( + executor: ThreadPoolExecutor, + items: list[UploadItem], + *, + client: RESTFullAPIClient, + zarr_id: str, + storage: RESTFullAPIClient, + dandiset: RemoteDandiset, +) -> list[Future[UploadResult]]: + """ + Submit the entries of a single-part Zarr for upload, each via a presigned + PUT. The URLs are requested anew on every attempt, as a retry is usually + needed because they timed out. + + :meta private: + """ + signed_urls = client.post( + f"/zarr/{zarr_id}/files/", json=[it.upload_request() for it in items] + ) + return [ + executor.submit( + _upload_zarr_file, + storage_session=storage, + dandiset=dandiset, + upload_url=signed_url, + item=it, + ) + for (signed_url, it) in zip(signed_urls, items) + ] + + +def _submit_zarr_entries( + executor: ThreadPoolExecutor, + items: list[UploadItem], + *, + client: RESTFullAPIClient, + zarr_id: str, + jobs: int | None, + asset_path: str, +) -> list[Future[UploadResult]]: + """ + Submit the entries of a multipart Zarr for upload, each via its own S3 + multipart upload. Each entry's upload is driven to completion in a worker + thread, since a generator cannot yield from within one, while the entry + still parallelizes its own parts across ``jobs`` threads, which matters for + the occasional very large entry. + + :meta private: + """ + return [ + executor.submit( + _upload_zarr_entry, + client=client, + zarr_id=zarr_id, + item=it, + jobs=jobs, + asset_path=asset_path, + ) + for it in items + ] + + +def _upload_zarr_entry( + *, + client: RESTFullAPIClient, + zarr_id: str, + item: UploadItem, + jobs: int | None, + asset_path: str, +) -> UploadResult: + """ + Upload one entry of a multipart Zarr via S3 multipart upload, reporting the + outcome as an `UploadResult` rather than raising, so that the batch can + retry or abort as the error warrants. + + :meta private: + """ + lgr.debug( + "%s: Uploading Zarr entry %s (%.2f GiB) via multipart upload", + asset_path, + item.entry_path, + item.size / 1024**3, + ) + try: + for _status in _upload_zarr_entry_multipart( + client=client, zarr_id=zarr_id, item=item, jobs=jobs + ): + pass + except requests.HTTPError as e: + # A 403 means the presigned part URLs timed out; the other conditions + # are transient S3 hiccups (see `_retry_s3_upload`). Either way the + # entry is retried from a fresh multipart upload, since its part URLs + # cannot be re-signed in place. + if e.response is not None and ( + e.response.status_code == 403 or _retry_s3_upload(e.response) + ): + lgr.debug( + "Got %d error uploading Zarr entry %s (%d bytes), will retry" + " with a new multipart upload: %s", + e.response.status_code, + item.filepath, + item.size, + str(e), + ) + return UploadResult(status=UploadStatus.RETRY_NEEDED, item=item) + return UploadResult(status=UploadStatus.FAILED, item=item, error=e) + except Exception as e: + return UploadResult(status=UploadStatus.FAILED, item=item, error=e) + return UploadResult(status=UploadStatus.SUCCESS, item=item, size=item.size) + + +def _zarr_is_multipart(zarr: dict[str, Any]) -> bool: + """ + Return whether the archive's serialization of a Zarr indicates multipart + upload. An archive predating the ``upload_type`` field omits it, which + means single-part upload. + + :meta private: + """ + return zarr.get("upload_type") == ZARR_UPLOAD_TYPE_MULTIPART + + def _handle_failed_items_and_raise( executor: ThreadPoolExecutor, failed_items: list, futures: list ) -> None: @@ -996,6 +1204,65 @@ def _handle_failed_items_and_raise( raise failed_items[0][1] +def _upload_zarr_entry_multipart( + client: RESTFullAPIClient, + zarr_id: str, + item: UploadItem, + jobs: int | None = None, +) -> Generator[dict, None, None]: + """ + Upload an entry of a multipart Zarr via S3 multipart upload, yielding the + status `dict`\\s of the underlying multipart upload. + + :meta private: + """ + init_fields: dict[str, Any] = {"zarr_id": zarr_id, "chunk_key": item.entry_path} + if item.content_type is not None: + # S3 fixes an object's Content-Type when the multipart upload is + # created, which happens server-side, so -- unlike the single-part + # path, which sets it as a header on its own PUT -- it has to be sent + # to the initialize endpoint. + init_fields["content_type"] = item.content_type + try: + resp = yield from multipart_upload( + client=client, + filepath=item.filepath, + asset_path=item.entry_path, + init_fields=init_fields, + expected_etag=item.digest, + jobs=jobs, + upload_root="/zarr/uploads", + ) + except requests.HTTPError as e: + # Only the ``initialize`` request tells us whether the archive supports + # this upload at all; a failure of a part upload or of completion is an + # ordinary upload error, which the caller may retry. + if ( + e.response is not None + and e.response.status_code in (400, 404) + and "/zarr/uploads/initialize" in str(getattr(e.response, "url", "")) + ): + # A 404 means the archive lacks the Zarr multipart upload endpoint + # entirely; a 400 means it rejected this multipart upload (e.g. the + # Zarr is not marked for multipart upload). Either way a multipart + # upload of a Zarr entry cannot succeed against this archive. + raise UploadError( + f"{item.entry_path}: server rejected the multipart upload of" + f" this Zarr entry; the archive may not support multipart upload" + f" of Zarr chunks. Server response: {e.response.text}" + ) from e + raise + # The server reports back the key it stored the entry under. If that isn't + # the key we asked for, the entry landed elsewhere in the Zarr, which would + # otherwise surface only as an unexplained Zarr checksum mismatch. + chunk_key = resp.get("chunk_key") + if chunk_key is not None and chunk_key != item.entry_path: + raise UploadError( + f"{item.entry_path}: server stored this Zarr entry under the" + f" unexpected key {chunk_key!r}" + ) + + def _upload_zarr_file( storage_session: RESTFullAPIClient, dandiset: RemoteDandiset, @@ -1036,7 +1303,7 @@ def _upload_zarr_file( upload_url, data=fp, json_resp=False, - retry_if=_retry_zarr_file, + retry_if=_retry_s3_upload, headers=headers, timeout=(60, 7200), ) @@ -1074,22 +1341,6 @@ def _upload_zarr_file( return UploadResult(item=item, status=UploadStatus.SUCCESS, size=item.size) -def _retry_zarr_file(r: requests.Response) -> bool: - return ( - # Some sort of filesystem hiccup can cause requests to be unable to get the - # filesize, leading to it falling back to "chunked" transfer encoding, - # which S3 doesn't support. - r.status_code == 501 - and "header you provided implies functionality that is not implemented" - in r.text - ) or ( - # Network issue or rate limiting can cause a timeout, which results in a 400. - # Case: https://github.com/dandi/dandi-cli/issues/1662 - r.status_code == 400 - and "was not read from or written to within the timeout period" in r.text - ) - - @dataclass class EntryUploadTracker: """ @@ -1098,6 +1349,9 @@ class EntryUploadTracker: :meta private: """ + #: Whether the Zarr is uploaded via S3 multipart upload, which selects how + #: its entries are digested: `dandietag_nocache` if so, else `md5file_nocache`. + multipart: bool = False total_size: int = 0 digested_entries: list[UploadItem] = field(default_factory=list) fresh_entries: list[LocalZarrEntry] = field(default_factory=list) @@ -1109,12 +1363,16 @@ def register(self, e: LocalZarrEntry, digest: str | None = None) -> None: self.fresh_entries.append(e) self.total_size += e.size - @staticmethod - def _mkitem(e: LocalZarrEntry) -> UploadItem: + def _mkitem(self, e: LocalZarrEntry) -> UploadItem: # Avoid heavy import by importing within function: - from dandi.support.digests import md5file_nocache + from dandi.support.digests import dandietag_nocache, md5file_nocache - digest = md5file_nocache(e.filepath) + # Dispatch to the digest matching the Zarr's upload scheme. + digest = ( + dandietag_nocache(e.filepath) + if self.multipart + else md5file_nocache(e.filepath) + ) return UploadItem.from_entry(e, digest) def get_items(self, jobs: int = 5) -> Generator[UploadItem, None, None]: @@ -1166,13 +1424,6 @@ def from_entry(cls, e: LocalZarrEntry, digest: str) -> UploadItem: else: content_type = None size = pre_upload_size_check(e.filepath) - if size > S3_MAX_SINGLE_PART_UPLOAD: - raise ValueError( - f"Zarr chunk {e.filepath} is {size / 1024**3:.2f} GiB," - f" exceeding the S3 single-part upload limit of" - f" {S3_MAX_SINGLE_PART_UPLOAD / 1024**3:.0f} GiB." - f" Multipart upload for zarr chunks is not yet supported." - ) return cls( entry_path=str(e), filepath=e.filepath, @@ -1183,6 +1434,15 @@ def from_entry(cls, e: LocalZarrEntry, digest: str) -> UploadItem: @property def base64_digest(self) -> str: + # An entry of a multipart Zarr is digested with its S3 multipart ETag + # (``-``), which is not a plain MD5 and so has no base64 MD5 + # representation. Such entries are uploaded via multipart upload and do + # not go through the single-part path that needs this header. + if "-" in self.digest: + raise ValueError( + f"{self.entry_path}: digest {self.digest!r} is a multipart" + f" ETag, which has no base64 MD5 representation" + ) return b64encode(bytes.fromhex(self.digest)).decode("us-ascii") def upload_request(self) -> dict[str, str | None]: @@ -1190,12 +1450,20 @@ def upload_request(self) -> dict[str, str | None]: def _cmp_digests( - asset_path: str, local_entry: LocalZarrEntry, remote_digest: str + asset_path: str, + local_entry: LocalZarrEntry, + remote_digest: str, + multipart: bool = False, ) -> tuple[LocalZarrEntry, str, bool]: # Avoid heavy import by importing within function: - from dandi.support.digests import md5file_nocache + from dandi.support.digests import dandietag_nocache, md5file_nocache - local_digest = md5file_nocache(local_entry.filepath) + # Dispatch to the digest matching the Zarr's upload scheme. + local_digest = ( + dandietag_nocache(local_entry.filepath) + if multipart + else md5file_nocache(local_entry.filepath) + ) if local_digest != remote_digest: lgr.debug( "%s: Path %s in Zarr differs from local file; re-uploading", diff --git a/dandi/support/digests.py b/dandi/support/digests.py index 7a69a1629..f50265dc1 100644 --- a/dandi/support/digests.py +++ b/dandi/support/digests.py @@ -19,8 +19,10 @@ from dataclasses import dataclass, field import hashlib import logging +import os import os.path from pathlib import Path +import re from dandischema.digests.dandietag import DandiETag from fscacher import PersistentCache @@ -28,6 +30,7 @@ from zarr_checksum.tree import ZarrChecksumTree from .threaded_walk import threaded_walk +from ..consts import S3_MAX_SINGLE_PART_UPLOAD from ..utils import Hasher, exclude_from_zarr lgr = logging.getLogger("dandi.support.digests") @@ -83,12 +86,24 @@ def __call__(self, fpath: str | Path) -> dict[str, str]: @checksums.memoize_path def get_digest(filepath: str | Path, digest: str = "sha256") -> str: + """ + Compute the digest of ``filepath`` under the named algorithm. + + Besides the hashlib algorithms, ``digest`` may be ``"dandi-etag"`` or one + of the two Zarr checksums. A Zarr's checksum is an aggregate over its + entries' S3 ETags, which differ by the scheme the Zarr was uploaded with, + so the two schemes have to be named apart: ``"zarr-checksum"`` is the + checksum of a single-part Zarr and ``"zarr-checksum-multipart"`` that of a + multipart one, the scheme `dandi upload` uses for new Zarrs. + """ if digest == "dandi-etag": s = get_dandietag(filepath).as_str() assert isinstance(s, str) return s elif digest == "zarr-checksum": return get_zarr_checksum(Path(filepath)) + elif digest == "zarr-checksum-multipart": + return get_zarr_multipart_checksum(Path(filepath)) else: return Digester([digest])(filepath)[digest] @@ -98,43 +113,126 @@ def get_dandietag(filepath: str | Path) -> DandiETag: return DandiETag.from_file(filepath) -def get_zarr_checksum(path: Path, known: dict[str, str] | None = None) -> str: +#: Pattern of an S3 multipart ETag: an MD5 hex digest, a hyphen, and the number +#: of parts (e.g. ``d41d8cd98f00b204e9800998ecf8427e-3``). +_MULTIPART_ETAG_RE = re.compile(r"[0-9a-f]{32}-[1-9][0-9]*\Z") + + +def is_multipart_etag(digest: str) -> bool: + """ + Return whether ``digest`` is an S3 multipart ETag (``-``) rather + than a plain MD5 digest. An entry of a multipart Zarr is stored under such + an ETag (see `dandietag_nocache`), so this distinguishes the digests of a + multipart Zarr's entries from a single-part Zarr's. """ - Compute the Zarr checksum for a file or directory tree. + return bool(_MULTIPART_ETAG_RE.match(digest)) - If the digests for any files in the Zarr are already known, they can be - passed in the ``known`` argument, which must be a `dict` mapping - slash-separated paths relative to the root of the Zarr to hex digests. + +def zarr_has_oversized_entry(path: Path) -> bool: """ - if path.is_file(): - s = get_digest(path, "md5") - assert isinstance(s, str) - return s - if known is None: - known = {} + Return whether the Zarr at ``path`` contains any entry larger than + `S3_MAX_SINGLE_PART_UPLOAD`. Such a Zarr must be uploaded via S3 multipart + upload, since S3 rejects single-part PUTs above that size. + """ + for dirpath, dirnames, filenames in os.walk(path): + dp = Path(dirpath) + dirnames[:] = [d for d in dirnames if not exclude_from_zarr(dp / d)] + for fn in filenames: + fp = dp / fn + if exclude_from_zarr(fp): + continue + if os.path.getsize(fp) > S3_MAX_SINGLE_PART_UPLOAD: + return True + return False + + +def md5file_nocache(filepath: str | Path) -> str: + """ + Compute the plain MD5 digest of a file, bypassing the fscacher cache (which + has been shown to slow things down for the large numbers of files typically + present in Zarrs). + + This is the digest of an entry of a **single-part** Zarr, which S3 stores + under its plain MD5 ETag. For the multipart counterpart, see + `dandietag_nocache`. + """ + return Digester(["md5"])(filepath)["md5"] + + +def dandietag_nocache(filepath: str | Path) -> str: + """ + Compute the S3 multipart ETag (a.k.a. DANDI etag) of a file, bypassing the + fscacher cache (cf. `md5file_nocache`). + + This is the digest of an entry of a **multipart** Zarr, which S3 stores + under its multipart ETag; multipart is the scheme `dandi upload` uses for + new Zarrs. For the single-part counterpart, see `md5file_nocache`. + """ + s = DandiETag.from_file(filepath).as_str() + assert isinstance(s, str) + return s - def digest_file(f: Path) -> tuple[Path, str, int]: - assert known is not None + +def _zarr_checksum( + path: Path, known: dict[str, str], digest_file: Callable[[Path], str] +) -> str: + """ + Compute a Zarr checksum for the directory tree ``path``, digesting each + entry not already present in ``known`` with ``digest_file``. The two public + entry points — `get_zarr_checksum` (single-part) and + `get_zarr_multipart_checksum` (multipart) — differ only in that per-entry + digest function. + + :meta private: + """ + + def digest_entry(f: Path) -> tuple[Path, str, int]: relpath = f.relative_to(path).as_posix() try: dgst = known[relpath] except KeyError: - dgst = md5file_nocache(f) + dgst = digest_file(f) return (f, dgst, os.path.getsize(f)) zcc = ZarrChecksumTree() - for p, digest, size in threaded_walk(path, digest_file, exclude=exclude_from_zarr): + for p, digest, size in threaded_walk(path, digest_entry, exclude=exclude_from_zarr): zcc.add_leaf(p.relative_to(path), size, digest) return str(zcc.process()) -def md5file_nocache(filepath: str | Path) -> str: +def get_zarr_checksum(path: Path, known: dict[str, str] | None = None) -> str: """ - Compute the MD5 digest of a file without caching with fscacher, which has - been shown to slow things down for the large numbers of files typically - present in Zarrs + Compute the **single-part** Zarr checksum for a file or directory tree: + every entry is digested with its plain MD5, as S3 stores it for a + single-part upload. This is the checksum of a single-part Zarr; for the + multipart counterpart, see `get_zarr_multipart_checksum`. + + If the digests for any files in the Zarr are already known, they can be + passed in the ``known`` argument, which must be a `dict` mapping + slash-separated paths relative to the root of the Zarr to hex digests. """ - return Digester(["md5"])(filepath)["md5"] + if path.is_file(): + s = get_digest(path, "md5") + assert isinstance(s, str) + return s + return _zarr_checksum(path, known or {}, md5file_nocache) + + +def get_zarr_multipart_checksum(path: Path, known: dict[str, str] | None = None) -> str: + """ + Compute the **multipart** Zarr checksum for a file or directory tree: every + entry is digested with its S3 multipart ETag, as S3 stores it for a + multipart upload (the scheme `dandi upload` uses for new Zarrs). This is + the checksum of a multipart Zarr; for the single-part counterpart, see + `get_zarr_checksum`. + + If the digests for any files in the Zarr are already known, they can be + passed in the ``known`` argument, which must be a `dict` mapping + slash-separated paths relative to the root of the Zarr to hex digests. + """ + if path.is_file(): + return dandietag_nocache(path) + return _zarr_checksum(path, known or {}, dandietag_nocache) def checksum_zarr_dir( diff --git a/dandi/support/tests/test_digests.py b/dandi/support/tests/test_digests.py index af37214ea..603fc68bb 100644 --- a/dandi/support/tests/test_digests.py +++ b/dandi/support/tests/test_digests.py @@ -15,7 +15,15 @@ from pytest_mock import MockerFixture from .. import digests -from ..digests import Digester, checksum_zarr_dir, get_zarr_checksum +from ..digests import ( + Digester, + checksum_zarr_dir, + dandietag_nocache, + get_dandietag, + get_zarr_checksum, + get_zarr_multipart_checksum, + md5file_nocache, +) def test_digester(tmp_path): @@ -78,11 +86,17 @@ def test_get_zarr_checksum(mocker: MockerFixture, tmp_path: Path) -> None: assert ( get_zarr_checksum(tmp_path / "file1.txt") == "d0aa42f003e36c1ecaf9aa8f20b6f1ad" ) + # get_zarr_checksum is the single-part codepath: entries digested with MD5. assert get_zarr_checksum(tmp_path) == "25627e0fc7c609d10100d020f7782a25-8--197" assert get_zarr_checksum(sub1) == "64af93ad7f8d471c00044d1ddbd4c0ba-4--97" - assert get_zarr_checksum(empty) == "481a2f77ab786a0f45aafd5db0971caa-0--0" + # get_zarr_multipart_checksum is the distinct multipart codepath: entries + # are digested differently, so it yields a different checksum for the same + # content. (An empty Zarr's checksum does not depend on the scheme.) + assert get_zarr_multipart_checksum(tmp_path) != get_zarr_checksum(tmp_path) + assert get_zarr_multipart_checksum(empty) == get_zarr_checksum(empty) + spy = mocker.spy(digests, "md5file_nocache") assert ( get_zarr_checksum( @@ -157,3 +171,25 @@ def test_checksum_zarr_dir( checksum: str, ) -> None: assert checksum_zarr_dir(files=files, directories=directories) == checksum + + +@pytest.mark.ai_generated +def test_md5file_nocache_single_part(tmp_path: Path) -> None: + """An entry of a single-part Zarr is digested with its plain MD5.""" + f = tmp_path / "sample.txt" + f.write_bytes(b"123") + assert md5file_nocache(f) == "202cb962ac59075b964b07152d234b70" + + +@pytest.mark.ai_generated +def test_dandietag_nocache_multipart(tmp_path: Path) -> None: + """ + An entry of a multipart Zarr is digested with its multipart ETag rather than + its MD5, as that is what the server ends up storing for it. This is the + distinct multipart codepath, separate from `md5file_nocache`. + """ + f = tmp_path / "sample.txt" + f.write_bytes(b"123") + digest = dandietag_nocache(f) + assert digest != md5file_nocache(f) + assert digest == get_dandietag(f).as_str() diff --git a/dandi/tests/test_files.py b/dandi/tests/test_files.py index bc73aece2..0fae4c9e7 100644 --- a/dandi/tests/test_files.py +++ b/dandi/tests/test_files.py @@ -31,6 +31,8 @@ dandi_file, find_dandi_files, ) +from ..files.zarr import UploadItem +from ..support.digests import dandietag_nocache, md5file_nocache lgr = get_logger() @@ -472,30 +474,30 @@ def test_upload_zarr(new_dandiset, tmp_path): _ZARR_PROPERTIES_EXPECTED = { "2": { "total_size": 1516, - "total_digest": "4313ab36412db2981c3ed391b38604d6-5--1516", + "total_digest": "8411ffbee3d86259ddbbf9d0d9c754bb-5--1516", "entries": [ - (".zgroup", 24, "e20297935e73dd0154104d4ea53040ab"), - ("arr_0", 746, "51c74ec257069ce3a555bdddeb50230a-2--746"), - ("arr_0/.zarray", 315, "9e30a0a1a465e24220d4132fdd544634"), - ("arr_0/0", 431, "ed4e934a474f1d2096846c6248f18c00"), - ("arr_1", 746, "7b99a0ad9bd8bb3331657e54755b1a31-2--746"), - ("arr_1/.zarray", 315, "9e30a0a1a465e24220d4132fdd544634"), - ("arr_1/0", 431, "fba4dee03a51bde314e9713b00284a93"), + (".zgroup", 24, "c42a84c3473618a4013cbb106ada1c14-1"), + ("arr_0", 746, "02d454d2efbae1c8359399af1006bec4-2--746"), + ("arr_0/.zarray", 315, "980c7afd2e491fd1448dec190f96ed66-1"), + ("arr_0/0", 431, "41a358317ae6e63108b0df7a6d210334-1"), + ("arr_1", 746, "f0c0530e7ad17b8aab0a7136ce784d93-2--746"), + ("arr_1/.zarray", 315, "980c7afd2e491fd1448dec190f96ed66-1"), + ("arr_1/0", 431, "5f534328920f7ee3ac417fbf4ac1826e-1"), ], }, "3": { "total_size": 3935, - "total_digest": "00157f091c9a6295e89eb3c4c2efaeff-5--3935", + "total_digest": "5dda89edeb0a78f05e5a990a55059223-5--3935", "entries": [ - ("arr_0", 2192, "ae16256ae750e4303674ccf1e23fa3c6-2--2192"), - ("arr_0/c", 1573, "93912a45f2107a08090f7b283297d662-1--1573"), - ("arr_0/c/0", 1573, "6c237f8d2d4a41bc1e26e31518dafd9e"), - ("arr_0/zarr.json", 619, "850fae056c97aa9c76df0a52411f4086"), - ("arr_1", 1677, "debc9ca4b2184a6ef1a3d6fcf7d79fd9-2--1677"), - ("arr_1/c", 1058, "2642f5d2df2cddf469313abd9910b371-1--1058"), - ("arr_1/c/0", 1058, "084d662af7251a807649fb48edc36e95"), - ("arr_1/zarr.json", 619, "850fae056c97aa9c76df0a52411f4086"), - ("zarr.json", 66, "457126c0639af2eba0140851c39c1aad"), + ("arr_0", 2192, "0db93f2603c99fd46e40e88e2fc55513-2--2192"), + ("arr_0/c", 1573, "5174f547f2b6124aa6c10b4085eba421-1--1573"), + ("arr_0/c/0", 1573, "0f9a883ef281510b0a1c003dc7da5ea1-1"), + ("arr_0/zarr.json", 619, "8dec6b2c24e4625f746482abbd0d2028-1"), + ("arr_1", 1677, "0be9308adf874a99bca107ac507ec7ec-2--1677"), + ("arr_1/c", 1058, "8ca53b9c7ba6f1b52b1962f9de3a3a50-1--1058"), + ("arr_1/c/0", 1058, "e8f546d8c4808292e5c7d2d1b066ff24-1"), + ("arr_1/zarr.json", 619, "8dec6b2c24e4625f746482abbd0d2028-1"), + ("zarr.json", 66, "ad2687bd67b5a2a11bba074b0a9bdf22-1"), ], }, } @@ -561,6 +563,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_zarr_upload_item_single_part(tmp_path: Path) -> None: + """ + An entry of a single-part Zarr carries a plain MD5 digest, which has a + base64 representation used by the single-part path as the Content-MD5 + header. + """ + zarr_path = tmp_path / "example.zarr" + zarr.save(zarr_path, np.arange(1000), np.arange(1000, 0, -1)) + zf = dandi_file(zarr_path) + assert isinstance(zf, ZarrAsset) + entry = next(e for e in zf.iterfiles() if e.is_file()) + item = UploadItem.from_entry(entry, md5file_nocache(entry.filepath)) + assert item.base64_digest + + +@pytest.mark.ai_generated +def test_zarr_upload_item_multipart(tmp_path: Path) -> None: + """ + Every entry of a multipart Zarr is uploaded via multipart upload and carries + a multipart ETag rather than an MD5 digest, regardless of the entry's size. + """ + zarr_path = tmp_path / "example.zarr" + zarr.save(zarr_path, np.arange(1000), np.arange(1000, 0, -1)) + zf = dandi_file(zarr_path) + assert isinstance(zf, ZarrAsset) + entry = next(e for e in zf.iterfiles() if e.is_file()) + item = UploadItem.from_entry(entry, dandietag_nocache(entry.filepath)) + # A multipart ETag is not a hex digest, so it has no base64 MD5 form. + with pytest.raises(ValueError, match="multipart ETag"): + item.base64_digest + + 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)) diff --git a/dandi/tests/test_upload.py b/dandi/tests/test_upload.py index 8d1c372d7..90a1ce8ce 100644 --- a/dandi/tests/test_upload.py +++ b/dandi/tests/test_upload.py @@ -4,7 +4,7 @@ from datetime import datetime, timezone import os from pathlib import Path -from shutil import copyfile, rmtree +from shutil import copyfile, copytree, rmtree from typing import Any from unittest.mock import Mock from urllib.parse import urlparse @@ -22,7 +22,9 @@ from .test_helpers import assert_dirtrees_eq from ..consts import ( DOWNLOAD_SUFFIX, + S3_MAX_SINGLE_PART_UPLOAD, ZARR_MIME_TYPE, + ZARR_UPLOAD_TYPE_SINGLEPART, EmbargoStatus, SyncMode, dandiset_metadata_file, @@ -590,9 +592,12 @@ def test_zarr_upload_400_timeout_retry( original_request = RESTFullAPIClient.request def mock_request(self, method, path, **kwargs): - # Track attempts for each request + # Track upload attempts for each request. Only PUTs are counted: a + # multipart upload also POSTs to the object's URL to complete it, which + # is not another attempt at uploading the bytes. urlpath = urlparse(path).path if path.startswith("http") else path - request_attempts[urlpath] += 1 + if method == "PUT": + request_attempts[urlpath] += 1 # Simulate 400 timeout on first attempt for files containing "arr_0" if method == "PUT" and "arr_0" in path and request_attempts[urlpath] == 1: @@ -1090,3 +1095,117 @@ def test_upload_modified_zarr_reports_descriptive_message( assert "adding " in msg, msg assert "modifying " in msg, msg assert "deleting " in msg, msg + + +#: Size of a Zarr entry that S3 cannot accept as a single-part PUT, and which +#: therefore can only be uploaded via S3 multipart upload. +OVERSIZED_ENTRY_SIZE = S3_MAX_SINGLE_PART_UPLOAD + 1024**2 + +#: These tests write an oversized entry to disk and upload it, which costs +#: several GiB and several minutes, and they need an archive that supports +#: multipart Zarr upload, so they are opt-in. +oversized_zarr = pytest.mark.skipif( + not os.environ.get("DANDI_TESTS_OVERSIZED_ZARR"), + reason=( + "Set DANDI_TESTS_OVERSIZED_ZARR=1 to run tests that write a >5 GiB file" + " against a multipart-capable archive" + ), +) + + +@pytest.fixture(scope="module") +def oversized_zarr_path(tmp_path_factory: pytest.TempPathFactory) -> Path: + """ + A valid Zarr whose single chunk is just over `S3_MAX_SINGLE_PART_UPLOAD`. + + The array is created without data and its one chunk is then written + directly, so that the entry never has to be held in memory. The chunk is + filled with random bytes, both because the uncompressed array stores them + verbatim and so that the digests under test are of content that cannot be + trivially deduplicated. + """ + path = tmp_path_factory.mktemp("oversized") / "oversized.zarr" + root = zarr.open_group(str(path), mode="w") + root.create_dataset( + "arr", + shape=(OVERSIZED_ENTRY_SIZE,), + chunks=(OVERSIZED_ENTRY_SIZE,), + dtype="u1", + compressor=None, + ) + chunk = path / "arr" / "0" + with chunk.open("wb") as fp: + remaining = OVERSIZED_ENTRY_SIZE + while remaining > 0: + block = os.urandom(min(64 * 1024**2, remaining)) + fp.write(block) + remaining -= len(block) + assert chunk.stat().st_size == OVERSIZED_ENTRY_SIZE + return path + + +@pytest.mark.ai_generated +@oversized_zarr +def test_upload_zarr_oversized_entry( + new_dandiset: SampleDandiset, oversized_zarr_path: Path +) -> None: + """ + An entry too large for a single-part S3 PUT uploads via S3 multipart + upload, and the Zarr the archive ingests has the checksum computed locally + for a multipart Zarr. + """ + # Avoid heavy import by importing within function: + from ..support.digests import get_zarr_multipart_checksum, is_multipart_etag + + copytree(oversized_zarr_path, new_dandiset.dspath / "oversized.zarr") + new_dandiset.upload() + + (asset,) = new_dandiset.dandiset.get_assets() + assert isinstance(asset, RemoteZarrAsset) + assert asset.path == "oversized.zarr" + + # The oversized entry is stored under a multipart ETag, which is what makes + # it a multipart upload rather than a single-part PUT that happened to work. + entries = {str(e): e for e in asset.iterfiles()} + oversized = entries["arr/0"] + assert oversized.size == OVERSIZED_ENTRY_SIZE + assert is_multipart_etag(oversized.digest.value) + # S3 records the number of parts after the hyphen; more than one of them is + # what distinguishes a genuine multipart upload from a single-part PUT. + assert int(oversized.digest.value.split("-")[1]) > 1 + + # The archive's own checksum of what it ingested must match the checksum we + # compute locally for a multipart Zarr. + assert asset.get_digest().value == get_zarr_multipart_checksum( + new_dandiset.dspath / "oversized.zarr" + ) + + +@pytest.mark.ai_generated +@oversized_zarr +def test_upload_zarr_oversized_entry_to_singlepart_zarr( + new_dandiset: SampleDandiset, oversized_zarr_path: Path +) -> None: + """ + An oversized entry cannot be uploaded to a Zarr that was created as + single-part, since a Zarr's entries all have to use the scheme it was + created with. The upload is refused up front rather than failing partway + through. + """ + copytree(oversized_zarr_path, new_dandiset.dspath / "oversized.zarr") + + # Mint the Zarr ahead of the upload as a single-part one, so that `dandi + # upload` finds and reuses it instead of creating a multipart Zarr. This + # stands in for a Zarr created before the archive supported multipart + # upload. + new_dandiset.client.post( + "/zarr/", + json={ + "name": "oversized.zarr", + "dandiset": new_dandiset.dandiset_id, + "upload_type": ZARR_UPLOAD_TYPE_SINGLEPART, + }, + ) + + with pytest.raises(UploadError, match="requires multipart upload"): + new_dandiset.upload()