sha256 for binaries index (#6241)

clee2000 · web-flow · commit 441473d2a4c8 · 2025-02-03T13:34:48.000-08:00
Restores the sha256 hash for binaries if they exist. In pytorch/pytorch#144887 the sha256 of the binary is manually calculated and uploaded as metadata with the binary because the automatic one given by s3 is wrong (likely due to it being part of a multipart upload https://docs.aws.amazon.com/AmazonS3/latest/userguide/tutorial-s3-mpu-additional-checksums.html). This PR * Attempts to adds a check for multipart upload sha256's and does not show them if they are part of a multipart upload and * Surfaces the new metadata given by the above PR * Restores the sha256 that was previously removed due to being wrong I also went back and corrected the bad shas from the multi part upload (not that it would have mattered, as this PR wouldn't show them in the index). The semantics of `copy_from` which is used in managed.py if `--compute-sha256` is used is correct and doesn't have the multi part upload issue
diff --git a/s3_management/manage.py b/s3_management/manage.py
@@ -382,14 +382,7 @@ def to_simple_package_html(
         out.append('  <body>')
         out.append('    <h1>Links for {}</h1>'.format(package_name.lower().replace("_", "-")))
         for obj in sorted(self.gen_file_list(subdir, package_name)):
-
             maybe_fragment = f"#sha256={obj.checksum}" if obj.checksum else ""
-
-            # Temporary skip assigning sha256 to nightly index
-            # to be reverted on Jan 24, 2025.
-            if subdir is not None and "nightly" in subdir:
-                maybe_fragment = ""
-
             pep658_attribute = ""
             if obj.pep658:
                 pep658_sha = f"sha256={obj.pep658}"
@@ -528,6 +521,7 @@ def fetch_object_names(cls: Type[S3IndexType], prefix: str) -> List[str]:
 
     def fetch_metadata(self: S3IndexType) -> None:
         # Add PEP 503-compatible hashes to URLs to allow clients to avoid spurious downloads, if possible.
+        regex_multipart_upload = r"^[A-Za-z0-9+/=]+=-[0-9]+$"
         with concurrent.futures.ThreadPoolExecutor(max_workers=6) as executor:
             for idx, future in {
                 idx: executor.submit(
@@ -540,10 +534,17 @@ def fetch_metadata(self: S3IndexType) -> None:
                 if obj.size is None
             }.items():
                 response = future.result()
-                sha256 = (_b64 := response.get("ChecksumSHA256")) and base64.b64decode(_b64).hex()
+                raw = response.get("ChecksumSHA256")
+                if raw and match(regex_multipart_upload, raw):
+                    # Possibly part of a multipart upload, making the checksum incorrect
+                    print(f"WARNING: {self.objects[idx].orig_key} has bad checksum: {raw}")
+                    raw = None
+                sha256 = raw and base64.b64decode(raw).hex()
                 # For older files, rely on checksum-sha256 metadata that can be added to the file later
                 if sha256 is None:
                     sha256 = response.get("Metadata", {}).get("checksum-sha256")
+                if sha256 is None:
+                    sha256 = response.get("Metadata", {}).get("x-amz-meta-checksum-sha256")
                 self.objects[idx].checksum = sha256
                 if size := response.get("ContentLength"):
                     self.objects[idx].size = int(size)
