pytorch · thejcannon · Jul 5, 2023 · Jul 12, 2023 · Oct 27, 2023 · Oct 27, 2023
diff --git a/manywheel/upload.sh b/manywheel/upload.sh
@@ -33,12 +33,17 @@ for cuda_ver in "${CUDA_VERSIONS[@]}"; do
 
     # Upload the wheels to s3
     if [[ -d "$wheel_dir" ]]; then
+        pushd "$wheel_dir"
+        find . -type f -exec sh -c 'unzip -j {} -d . "*.dist-info/METADATA" && mv METADATA {}.metadata' \;
         echo "Uploading all of: $(ls $wheel_dir) to $s3_wheel_dir"
-        ls "$wheel_dir" | xargs -I {} aws s3 cp "$wheel_dir"/{} "$s3_wheel_dir" --acl public-read
+        ls . | xargs -I {} aws s3 cp {} "$s3_wheel_dir" --acl public-read
+        popd
     fi
 
     if [[ -d "$libtorch_dir" ]]; then
+        pushd "$libtorch_dir"
         echo "Uploading all of: $(ls $libtorch_dir) to $s3_libtorch_dir"
-        ls "$libtorch_dir" | xargs -I {} aws s3 cp "$libtorch_dir"/{} "$s3_libtorch_dir" --acl public-read
+        ls . | xargs -I {} aws s3 cp {} "$s3_libtorch_dir" --acl public-read
+        popd
     fi
 done
diff --git a/s3_management/manage.py b/s3_management/manage.py
@@ -14,7 +14,6 @@
 
 
 S3 = boto3.resource('s3')
-CLIENT = boto3.client('s3')
 BUCKET = S3.Bucket('pytorch')
 
 ACCEPTED_FILE_EXTENSIONS = ("whl", "zip", "tar.gz")
@@ -121,8 +120,9 @@ def between_bad_dates(package_build_time: datetime):
 
 
 class S3Index:
-    def __init__(self: S3IndexType, objects: List[str], prefix: str) -> None:
+    def __init__(self: S3IndexType, objects: List[str], whls_with_metadata: Set[str], prefix: str) -> None:
         self.objects = objects
+        self.whls_with_metadata = set(whls_with_metadata)
         self.prefix = prefix.rstrip("/")
         self.html_name = PREFIXES_WITH_HTML[self.prefix]
         # should dynamically grab subdirectories like whl/test/cu101
@@ -255,7 +255,13 @@ 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)):
-            out.append(f'    <a href="/{obj}">{path.basename(obj).replace("%2B","+")}</a><br/>')
+            attributes = []
+            if obj in self.whls_with_metadata:
+                # Serve the PEP 658 metadata attributes.
+                # For extra juiciness, we should expose the sha256, instead of "true".
+                attributes += 'data-dist-info-metadata="true"'
+            attributes = " ".join(attributes)
+            out.append(f'    <a href="/{obj}">{path.basename(obj).replace("%2B","+")} {attributes}</a><br/>')
         # Adding html footer
         out.append('  </body>')
         out.append('</html>')
@@ -338,6 +344,7 @@ def save_pep503_htmls(self) -> None:
     @classmethod
     def from_S3(cls: Type[S3IndexType], prefix: str) -> S3IndexType:
         objects = []
+        whls_with_metadata = {}
         prefix = prefix.rstrip("/")
         for obj in BUCKET.objects.filter(Prefix=prefix):
             is_acceptable = any([path.dirname(obj.key) == prefix] + [
@@ -346,11 +353,14 @@ def from_S3(cls: Type[S3IndexType], prefix: str) -> S3IndexType:
                     path.dirname(obj.key)
                 )
                 for pattern in ACCEPTED_SUBDIR_PATTERNS
-            ]) and obj.key.endswith(ACCEPTED_FILE_EXTENSIONS)
-            if is_acceptable:
-                sanitized_key = obj.key.replace("+", "%2B")
+            ])
+            sanitized_key = obj.key.replace("+", "%2B")
+            if obj.key.endswith(ACCEPTED_FILE_EXTENSIONS) and is_acceptable:
                 objects.append(sanitized_key)
-        return cls(objects, prefix)
+            if obj.key.endswith(".whl.metadata"):
+                whls_with_metadata.append(sanitized_key[:-9])
+
+        return cls(objects, whls_with_metadata, prefix)
 
 def create_parser() -> argparse.ArgumentParser:
     parser = argparse.ArgumentParser("Manage S3 HTML indices for PyTorch")

diff --git a/wheel/upload.sh b/wheel/upload.sh
@@ -18,14 +18,18 @@ fi
 
 # Upload wheels to s3
 if [[ -d "$MAC_WHEEL_FINAL_FOLDER" ]]; then
+    pushd "$MAC_WHEEL_FINAL_FOLDER"
     s3_dir="s3://pytorch/whl/${PIP_UPLOAD_FOLDER}cpu/"
+    find . -type f -exec sh -c 'unzip -j {} -d . "*.dist-info/METADATA" && mv METADATA {}.metadata' \;
     echo "Uploading all of: $(ls $MAC_WHEEL_FINAL_FOLDER) to $s3_dir"
-    ls "$MAC_WHEEL_FINAL_FOLDER" | xargs -I {} aws s3 cp "$MAC_WHEEL_FINAL_FOLDER"/{} "$s3_dir" --acl public-read
+    ls . | xargs -I {} aws s3 cp {} "$s3_dir" --acl public-read
 fi
 
 # Upload libtorch packages to s3
 if [[ -d "$MAC_LIBTORCH_FINAL_FOLDER" ]]; then
+    pushd "$MAC_LIBTORCH_FINAL_FOLDER"
     s3_dir="s3://pytorch/libtorch/${PIP_UPLOAD_FOLDER}cpu/"
     echo "Uploading all of: $(ls $MAC_LIBTORCH_FINAL_FOLDER) to $s3_dir"
-    ls "$MAC_LIBTORCH_FINAL_FOLDER" | xargs -I {} aws s3 cp "$MAC_LIBTORCH_FINAL_FOLDER"/{} "$s3_dir" --acl public-read
+    ls . | xargs -I {} aws s3 cp {} "$s3_dir" --acl public-read
+    popd
 fi
diff --git a/windows/upload/upload.sh b/windows/upload/upload.sh
@@ -24,16 +24,19 @@ popd
 
 pushd winwheels/whl
 if [[ "$package_name" == pytorch ]]; then
-    find . -name "*torch-*.whl" | cut -f 2- -d'/' | xargs -I {} aws s3 cp {} s3://pytorch/whl/{}  --acl public-read
+    whl_name="*torch-*.whl*"
 elif [[ "$package_name" == torchvision ]]; then
-    find . -name "*torchvision*.whl" | cut -f 2- -d'/' | xargs -I {} aws s3 cp {} s3://pytorch/whl/{}  --acl public-read
+    whl_name="*torchvision*.whl*"
 fi
+find . -type f -name "$whl_name" -exec sh -c 'unzip -j {} -d . "*.dist-info/METADATA" && mv METADATA {}.metadata' \;
+find . -name "$whl_name" | cut -f 2- -d'/' | xargs -I {} aws s3 cp {} s3://pytorch/whl/{}  --acl public-read
 popd
 
 
 if [[ "$package_name" == pytorch ]]; then
     pushd winwheels/libtorch
     find . -name "*.zip" |  cut -f 2- -d'/' | xargs -I {} aws s3 cp {} s3://pytorch/libtorch/{}  --acl public-read
+    popd
 fi
 
 # then run