Improve S3 promotion: ordering, R2 grep fixes, and a per-target promoter (#8071)

## Summary

Three coordinated changes to the release promotion path.

### 1. Split `aws_promote` and reorder `s3_to_s3.sh`

`release/promote/common_utils.sh` now splits `aws_promote` into two
functions:

- `aws_promote` - server-side S3 copy only (fast).
- `aws_set_checksums` - the `manage_v2.py --set-checksum` loop, which
downloads every wheel from S3 to compute SHA256 (slow).

`release/promote/s3_to_s3.sh` now calls them in the new order: **S3 copy
→ R2 upload → S3 checksum recompute**. R2 previously waited on
the per-wheel SHA256 download loop because it ran *inside* the old
`aws_promote`; the new ordering unblocks R2 immediately after the
server-side S3 copy.

`aws_set_checksums` also skips automatically when the destination prefix
is not `whl` / `whl/test` (e.g. `libtorch`), since `manage_v2.py
--set-checksum` raises `ValueError: Prefix must be whl or whl/test` for
anything else.

### 2. Fix `r2_promote` grep bugs for glob package names and suffixes

`PACKAGE_INCLUDE_SUFFIX` (e.g. `*manylinux*`) and `PACKAGE_NAME` (e.g.
`libtorch-*`) are AWS S3 *globs* in `aws_promote`'s `--include` flag.
The old `r2_promote` passed them straight to `grep` where `*` is the
regex zero-or-more operator, so `triton-3.7.0*manylinux*` and
`libtorch-*-2.12.0*` matched nothing and R2 silently found 0 files. Both
now have their `*` expanded to `.*` before being plugged into `grep -E`,
so the same expression that selects files for `aws_promote` also selects
them for `r2_promote`.

### 3. New `--promote-from` flag in `update_dependencies.py`

`s3_management/update_dependencies.py` gains a `promote_target` function
and `--promote-from` CLI flag for promoting per-package wheel artifacts
across prefixes (e.g. `whl/test/cu132 -> whl/cu132`). For each entry in
`PACKAGES_PER_PROJECT`, it copies `.whl`, `.whl.metadata`, `.tar.gz`,
and `.tgz` files that exist at the source, matched by distribution name
(PEP 503 normalized), and skips any that already exist at the
destination. Also accepts `xpu` as a valid target name.

Example usage:

```bash
python3 s3_management/update_dependencies.py \
    --promote-from whl/test \
    --prefix whl \
    --target cu132 \
    --dry-run
```

## Test plan

- [x] Verified the new `promote_target` finds and copies the right files
for `cu132` in dry-run mode.
- [x] Verified `r2_promote` now matches both `*manylinux*`-filtered
wheels and `libtorch-*` globs after the regex fix.
- [x] Verified `aws_set_checksums` skips with a clean log line for the
libtorch promotion (prefix `libtorch` instead of `whl`).
- [x] Local syntax checks (`bash -n`, `python3 -c "ast.parse(...)"`)
pass for all three files.
- [ ] End-to-end run with `DRY_RUN=disabled` on a real release.

Authored with Claude Code.

Author: atalman

release/promote/common_utils.sh

@@ -58,29 +58,50 @@ aws_promote() {
     )
     # ^ We grep for package_name-.*pytorch_version to avoid any situations where domain libraries have
     #   the same version on our S3 buckets
+}
 
-    # After copying, explicitly set SHA256 checksums for wheels that don't have them
-    # This ensures checksums are preserved even if --metadata-directive COPY fails to copy them
-    if [[ $DRY_RUN = "disabled" ]]; then
-        echo "+ Setting SHA256 checksums for copied wheels..."
-        dest_prefix="${PYTORCH_S3_TO#s3://pytorch/}"
-        dest_prefix="${dest_prefix%/}"
-
-        script_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
-        manage_v2_script="${script_dir}/../../s3_management/manage_v2.py"
-
-        if [[ -f "${manage_v2_script}" ]]; then
-            echo "+ Running: python ${manage_v2_script} ${dest_prefix} --set-checksum --package-name ${package_name} --package-version ${pytorch_version}"
-            python "${manage_v2_script}" "${dest_prefix}" \
-                --set-checksum \
-                --package-name "${package_name}" \
-                --package-version "${pytorch_version}" || {
-                echo "- WARNING: Failed to set SHA256 checksums, but copy succeeded"
-            }
-        else
-            echo "- WARNING: manage_v2.py not found at ${manage_v2_script}, skipping checksum computation"
-        fi
+aws_set_checksums() {
+    # Re-derive SHA256 checksum metadata on the S3 destination wheels.
+    # Runs as the final step of promotion so faster operations (S3 copy, R2
+    # upload) are not blocked waiting on this download-heavy loop.
+    package_name=$1
+    pytorch_version=$(get_pytorch_version)
+    DRY_RUN=${DRY_RUN:-enabled}
+
+    if [[ $DRY_RUN != "disabled" ]]; then
+        echo "+ DRY RUN: skipping SHA256 recomputation for ${package_name}"
+        return 0
+    fi
+
+    echo "=-=-=-= Setting SHA256 checksums for ${package_name} v${pytorch_version} on S3 =-=-=-="
+    dest_prefix="${PYTORCH_S3_TO#s3://pytorch/}"
+    dest_prefix="${dest_prefix%/}"
+
+    # manage_v2.py --set-checksum only supports whl and whl/test prefixes
+    # (it raises ValueError on anything else). Skip for libtorch/etc.
+    case "${dest_prefix}" in
+        whl|whl/test) ;;
+        *)
+            echo "+ Skipping SHA256 recomputation: dest prefix '${dest_prefix}' is not whl/whl-test; manage_v2.py only supports those."
+            return 0
+            ;;
+    esac
+
+    script_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+    manage_v2_script="${script_dir}/../../s3_management/manage_v2.py"
+
+    if [[ ! -f "${manage_v2_script}" ]]; then
+        echo "- WARNING: manage_v2.py not found at ${manage_v2_script}, skipping checksum computation"
+        return 0
     fi
+
+    echo "+ Running: python ${manage_v2_script} ${dest_prefix} --set-checksum --package-name ${package_name} --package-version ${pytorch_version}"
+    python "${manage_v2_script}" "${dest_prefix}" \
+        --set-checksum \
+        --package-name "${package_name}" \
+        --package-version "${pytorch_version}" || {
+        echo "- WARNING: Failed to set SHA256 checksums, but copy succeeded"
+    }
 }
 
 r2_promote() {
@@ -105,11 +126,19 @@ r2_promote() {
     echo "=-=-=-= Promoting ${package_name} v${pytorch_version} to R2 =-=-=-="
     echo "+ R2 destination: ${r2_dest}"
 
+    # PACKAGE_NAME and PACKAGE_INCLUDE_SUFFIX are consumed as AWS S3 globs by
+    # aws_promote (e.g. libtorch-*, *manylinux*). Convert their '*' wildcards
+    # into '.*' so the same expressions work as a grep -E regex here.
+    local pkg_regex="${package_name//\*/.*}"
+    local include_glob="${PACKAGE_INCLUDE_SUFFIX:-*}"
+    local include_regex="${include_glob//\*/.*}"
+    local match_pattern="${pkg_regex}-${pytorch_version}${include_regex}"
+
     if [[ $DRY_RUN = "enabled" ]]; then
         echo "+ DRY RUN: Would copy matching files from ${PYTORCH_S3_FROM} to R2 ${r2_dest}"
         # List what would be copied
         ${AWS} s3 ls "${PYTORCH_S3_FROM/\/$//}/" --recursive \
-            | grep "${package_name}-${pytorch_version}${PACKAGE_INCLUDE_SUFFIX:-}" || true
+            | grep -E "${match_pattern}" || true
         return 0
     fi
 
@@ -129,7 +158,7 @@ r2_promote() {
     local s3_from_path="${PYTORCH_S3_FROM/\/$//}"
     local file_list="${tmp_dir}/file_list.txt"
     ${AWS} s3 ls "${s3_from_path}/" --recursive \
-        | grep "${package_name}-${pytorch_version}${PACKAGE_INCLUDE_SUFFIX:-}" \
+        | grep -E "${match_pattern}" \
         | awk '{print $NF}' > "${file_list}" || true
 
     local total_files

release/promote/s3_to_s3.sh

@@ -25,5 +25,13 @@ else
     echo "+ R2_ONLY=true, skipping S3-to-S3 promotion"
 fi
 
-# Also promote to R2 (Cloudflare) if credentials are available
+# Promote to R2 (Cloudflare) before the slow SHA256 recomputation step so R2
+# is not blocked waiting on per-wheel downloads on the S3 destination.
 r2_promote "${PACKAGE_NAME}"
+
+# Finally, recompute SHA256 checksum metadata on the S3 destination wheels.
+# This is the slowest step (downloads every wheel from S3) and runs last so
+# it does not delay the R2 upload above.
+if [[ "${R2_ONLY}" != "true" ]]; then
+    aws_set_checksums "${PACKAGE_NAME}"
+fi

s3_management/update_dependencies.py

@@ -14,6 +14,7 @@
 VALID_TARGET_PATTERNS = [
     r"^cu[0-9]+$",  # CUDA: cu118, cu121, cu126, cu128, cu129, cu130, cuXYZ
     r"^rocm[0-9]+\.[0-9]+$",  # ROCm: rocm5.7, rocm6.0, rocm6.4, rocm7.1, rocm7.2
+    r"^xpu$",  # Intel XPU
 ]
 
 # Cloudflare R2 configuration for writing indexes
@@ -1082,6 +1083,162 @@ def copy_target(
     return True
 
 
+def promote_target(
+    source_prefix: str,
+    target_prefix: str,
+    target: str,
+    *,
+    dry_run: bool = False,
+) -> bool:
+    """
+    Promote a target across prefixes by copying packages from PACKAGES_PER_PROJECT.
+
+    Scan files directly in {source_prefix}/{target}/ (not in subdirectories)
+    and copy any wheel/tarball whose distribution name belongs to
+    PACKAGES_PER_PROJECT into {target_prefix}/{target}/. PEP 658 metadata
+    sidecars (.whl.metadata) accompanying matched wheels are also copied.
+    Used for promotion flows such as whl/test/cu132 -> whl/cu132, where the
+    destination target directory does not yet exist. Only binary artifacts
+    and their metadata are copied (.whl, .whl.metadata, .tar.gz, .tgz);
+    index.html files and per-package subdirectories are intentionally skipped
+    because the destination index is regenerated by other tooling.
+
+    Args:
+        source_prefix: Source prefix (e.g., "whl/test")
+        target_prefix: Destination prefix (e.g., "whl")
+        target: The target name (e.g., "cu132")
+        dry_run: If True, don't actually copy anything
+
+    Returns:
+        True if at least one object was (or would be) copied, False otherwise.
+    """
+    if not is_valid_target(target):
+        print(f"ERROR: Invalid target name '{target}'")
+        print(f"Valid patterns: {VALID_TARGET_PATTERNS}")
+        return False
+
+    source_target_path = f"{source_prefix}/{target}"
+    dest_target_path = f"{target_prefix}/{target}"
+
+    # Check source exists
+    if not dry_run and not target_exists(source_prefix, target):
+        print(f"ERROR: Source '{source_target_path}' does not exist in S3.")
+        return False
+
+    # Warn if destination already exists but continue anyway
+    if not dry_run and target_exists(target_prefix, target):
+        print(
+            f"WARNING: Destination '{dest_target_path}' already exists, "
+            "will copy into it."
+        )
+
+    print(
+        f"{'[DRY RUN] ' if dry_run else ''}"
+        f"Promoting {source_target_path} -> {dest_target_path}"
+    )
+
+    copied_count = 0
+    skipped_count = 0
+    existing_count = 0
+    artifact_suffixes = (".whl", ".whl.metadata", ".tar.gz", ".tgz")
+    matched_packages: set[str] = set()
+
+    def s3_object_exists(key: str) -> bool:
+        try:
+            CLIENT.head_object(Bucket="pytorch", Key=key)
+            return True
+        except CLIENT.exceptions.ClientError as exc:
+            if exc.response["Error"]["Code"] in ("404", "NoSuchKey"):
+                return False
+            raise
+
+    # Normalized lookup set: lowercase, "_" -> "-" (PEP 503 / wheel filename convention)
+    allowed_packages = {name.lower().replace("_", "-") for name in PACKAGES_PER_PROJECT}
+
+    def package_in_allowlist(filename: str) -> str:
+        """Return the canonical package name if filename belongs to an allowed
+        package, else empty string. Wheel and sdist filenames are
+        '<distribution>-<version>...' with the distribution normalized to use
+        underscores; split on the first '-' to recover it."""
+        dist, sep, _ = filename.partition("-")
+        if not sep:
+            return ""
+        canonical = dist.lower().replace("_", "-")
+        return canonical if canonical in allowed_packages else ""
+
+    paginator = CLIENT.get_paginator("list_objects_v2")
+    # Delimiter='/' restricts the listing to objects directly under the prefix
+    # (no recursion into per-package subdirectories).
+    for page in paginator.paginate(
+        Bucket="pytorch",
+        Prefix=f"{source_target_path}/",
+        Delimiter="/",
+    ):
+        for obj in page.get("Contents", []):
+            source_key = obj["Key"]
+            filename = source_key.rsplit("/", 1)[-1]
+
+            if not filename.endswith(artifact_suffixes):
+                skipped_count += 1
+                continue
+
+            canonical = package_in_allowlist(filename)
+            if not canonical:
+                skipped_count += 1
+                continue
+
+            matched_packages.add(canonical)
+            dest_key = f"{dest_target_path}/{filename}"
+
+            # Skip if the destination already has this artifact
+            if s3_object_exists(dest_key):
+                existing_count += 1
+                print(f"  Skipping (already exists at destination): {dest_key}")
+                continue
+
+            if dry_run:
+                print(f"  [DRY RUN] Would copy: {source_key} -> {dest_key}")
+            else:
+                print(f"  Copying: {source_key} -> {dest_key}")
+                CLIENT.copy_object(
+                    Bucket="pytorch",
+                    CopySource={"Bucket": "pytorch", "Key": source_key},
+                    Key=dest_key,
+                    ACL="public-read",
+                )
+
+                # Also copy to R2 if configured (cross-service copy)
+                if R2_BUCKET:
+                    print(f"  Copying to R2: {source_key} -> {dest_key}")
+                    response = CLIENT.get_object(Bucket="pytorch", Key=source_key)
+                    body = response["Body"].read()
+                    # Set ContentType from the filename rather than trusting the
+                    # source object's metadata: wheels and tarballs are binary
+                    # and must not be served as text/html. PEP 658 .whl.metadata
+                    # sidecars are RFC 822 text.
+                    if filename.endswith(".whl.metadata"):
+                        content_type = "text/plain"
+                    else:
+                        content_type = "application/octet-stream"
+                    R2_BUCKET.Object(key=dest_key).put(
+                        ACL="public-read",
+                        ContentType=content_type,
+                        CacheControl="no-cache,no-store,must-revalidate",
+                        Body=body,
+                    )
+
+            copied_count += 1
+
+    print(
+        f"{'[DRY RUN] ' if dry_run else ''}"
+        f"Promotion complete: {len(matched_packages)} packages matched, "
+        f"{copied_count} artifacts copied "
+        f"({existing_count} already existed at destination, "
+        f"{skipped_count} files skipped)"
+    )
+    return copied_count > 0
+
+
 def create_target(
     prefix: str,
     target: str,
@@ -1188,6 +1345,15 @@ def main() -> None:
         help="Initialize a new target by copying all folders and index.html from an existing source target "
         "(e.g., --init-from cu130 --target cu132 copies whl/nightly/cu130/* -> whl/nightly/cu132/*)",
     )
+    parser.add_argument(
+        "--promote-from",
+        type=str,
+        help="Promote a target across prefixes by copying every package in "
+        "PACKAGES_PER_PROJECT that exists under <promote-from>/<target>/ into "
+        "<prefix>/<target>/ (e.g., --promote-from whl/test --prefix whl "
+        "--target cu132 copies whl/test/cu132/<pkg>/ -> whl/cu132/<pkg>/ for "
+        "every package listed in PACKAGES_PER_PROJECT).",
+    )
     parser.add_argument(
         "--prefix",
         type=str,
@@ -1211,6 +1377,20 @@ def main() -> None:
         )
         return
 
+    # Handle promote-from mode (promote target across prefixes)
+    if args.promote_from:
+        if not args.target:
+            print("ERROR: --target is required when using --promote-from")
+            return
+
+        promote_target(
+            args.promote_from,
+            args.prefix,
+            args.target,
+            dry_run=args.dry_run,
+        )
+        return
+
     # Handle target creation mode
     if args.create_target:
         if not args.target: