Skip to content

Commit d920d67

Browse files
atalmanatalman
authored
Add post-promotion SHA256 recompute for prod whl/ release packages (#8199)
## Summary Adds post-promotion SHA256 checksum recompute for prod `whl/` release packages so wheels published to download.pytorch.org always carry `x-amz-meta-checksum-sha256` metadata. Only wheels **missing** a checksum are computed — existing checksums are never recomputed or clobbered, and only the S3 `pytorch` bucket is written (no R2). ## Changes ### `s3_management/manage_v2.py` - `--recompute-sha256-pattern` is now **optional** (`nargs="?"`): - With a value (e.g. `rocm7.1`) → scans only that subfolder under the prefix (unchanged behavior). - **Bare, with no value** → scans **all** accelerator subfolders (`cu*`/`rocm*`/`cpu`/`xpu`) under the prefix in one invocation. Nested channels (`whl/nightly`, `whl/test`) are excluded because their names don't match the accelerator patterns. - New helper `list_accelerator_subdirs()` enumerates those subfolders via a delimiter listing. - Combine with `--package-name` / `--package-version` to scope to a single release. ```bash # Recompute one release across all prod accelerator subfolders: python s3_management/manage_v2.py whl --recompute-sha256-pattern \ --package-name torch --package-version 2.12.1 ``` ### `.github/workflows/release-post-promotion.yml` (new) - `workflow_dispatch` with inputs `project` (choice, default `torch`) and `version` (optional override of `PYTORCH_VERSION`). - Reads `torch`, `torchvision`, `triton`, `triton_rocm`, `triton_xpu` versions from `release/release_versions.sh` (the triton variants share `TRITON_VERSION`) and recomputes SHA256 across all prod `whl/` accelerator subfolders. - Same OIDC setup as `release-download-pytorch-org.yml`: `id-token: write`, `environment: promote-env`, role `gha_workflow_promote_wheels`, region `us-east-1`. ### `.github/workflows/release-download-pytorch-org.yml` - Added a step that recomputes SHA256 for the promoted package after a real (non-dryrun) promotion, resolving its version via `${PACKAGE^^}_VERSION` from `release_versions.sh`. ## Test Plan - `manage_v2.py` parses cleanly; argparse resolves the three cases correctly: bare flag → scan-all, explicit value → single subfolder, omitted → no-op. - Verified `main` already provides the required version vars in `release/release_versions.sh`.
1 parent 0876801 commit d920d67

3 files changed

Lines changed: 179 additions & 44 deletions

File tree

‎.github/workflows/release-download-pytorch-org.yml‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -99,3 +99,23 @@ jobs:
9999
100100
# shellcheck disable=SC2086
101101
promote_s3 ${PACKAGE} whl "${!version}"
102+
103+
- name: Recompute SHA256 on whl/ (prod) for promoted package
104+
if: ${{ inputs.dryrun == 'disabled' }}
105+
shell: bash
106+
env:
107+
PACKAGE: ${{ inputs.package || 'torchvision' }}
108+
run: |
109+
set -ex
110+
# Init release versions variables
111+
source ./release/release_versions.sh
112+
# shellcheck disable=SC2086
113+
version="${PACKAGE^^}_VERSION"
114+
# Backfill SHA256 metadata for wheels missing it (skips existing).
115+
# Bare --recompute-sha256-pattern scans all prod accelerator
116+
# subfolders (cu*/rocm*/cpu/xpu) for this package/version.
117+
# shellcheck disable=SC2086
118+
python s3_management/manage_v2.py whl \
119+
--recompute-sha256-pattern \
120+
--package-name "${PACKAGE}" \
121+
--package-version "${!version}"

‎.github/workflows/release-post-promotion.yml‎

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,70 @@
1+
name: Release post-promotion SHA256 recompute
2+
3+
on:
4+
workflow_dispatch:
5+
inputs:
6+
project:
7+
description: "Release project to recompute checksums for"
8+
required: true
9+
type: choice
10+
default: torch
11+
options:
12+
- torch
13+
version:
14+
description: "Override torch version (defaults to release_versions.sh PYTORCH_VERSION)"
15+
required: false
16+
type: string
17+
18+
permissions:
19+
id-token: write
20+
contents: read
21+
22+
jobs:
23+
recompute-sha256:
24+
runs-on: ubuntu-latest
25+
timeout-minutes: 120
26+
environment: promote-env
27+
container:
28+
image: pytorch/almalinux-builder:cpu
29+
steps:
30+
- uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
31+
- name: Configure aws credentials (pytorch account)
32+
uses: aws-actions/configure-aws-credentials@50ac8dd1e1b10d09dac7b8727528b91bed831ac0 # v3.0.2
33+
with:
34+
role-to-assume: arn:aws:iam::749337293305:role/gha_workflow_promote_wheels
35+
aws-region: us-east-1
36+
- name: Recompute SHA256 on whl/ (prod) for torch release packages
37+
shell: bash
38+
env:
39+
PROJECT: ${{ inputs.project || 'torch' }}
40+
PYTORCH_VERSION: ${{ inputs.version }}
41+
run: |
42+
set -ex
43+
# Install requirements
44+
pip install awscli==1.32.18
45+
pip install -r s3_management/requirements.txt
46+
47+
# Init release version variables. PYTORCH_VERSION is honored as an
48+
# override when provided via the workflow input (release_versions.sh
49+
# uses ${PYTORCH_VERSION:-...}); otherwise the file default is used.
50+
source ./release/release_versions.sh
51+
52+
# Recompute missing SHA256 across all prod accelerator subfolders
53+
# (cu*/rocm*/cpu/xpu) for a package/version. Bare
54+
# --recompute-sha256-pattern scans every subfolder; existing
55+
# checksums are skipped.
56+
recompute() {
57+
echo "=== Recomputing SHA256 for $1 $2 on whl/ (prod) ==="
58+
python s3_management/manage_v2.py whl \
59+
--recompute-sha256-pattern \
60+
--package-name "$1" \
61+
--package-version "$2"
62+
}
63+
64+
# The torch project covers the core PyTorch wheel set. triton,
65+
# triton_rocm and triton_xpu all share TRITON_VERSION.
66+
recompute torch "${PYTORCH_VERSION}"
67+
recompute torchvision "${TORCHVISION_VERSION}"
68+
recompute triton "${TRITON_VERSION}"
69+
recompute triton_rocm "${TRITON_VERSION}"
70+
recompute triton_xpu "${TRITON_VERSION}"

‎s3_management/manage_v2.py‎

Lines changed: 89 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -27,8 +27,11 @@
2727
# the digest as S3 object metadata (x-amz-meta-checksum-sha256).
2828
# - --set-checksum: compute and set SHA256 metadata for a specific
2929
# package/version combination (requires --package-name and --package-version).
30-
# - --recompute-sha256-pattern PATTERN: compute SHA256 for all .whl files
31-
# matching PATTERN under the given prefix that are missing checksums.
30+
# - --recompute-sha256-pattern [PATTERN]: compute SHA256 for .whl files that
31+
# are missing checksums. With PATTERN, only that subfolder under the prefix
32+
# is scanned; with no value, all accelerator subfolders (cu*/rocm*/cpu/xpu)
33+
# under the prefix are scanned (nightly/test excluded). Combine with
34+
# --package-name / --package-version to scope to a single release.
3235
# - --recompute-missing-sha256: scan the entire prefix for .whl files that
3336
# are missing x-amz-meta-checksum-sha256 metadata and compute/set it.
3437
# Example: python s3_management/manage_v2.py channel --recompute-missing-sha256
@@ -60,6 +63,10 @@
6063
#
6164
# # Recompute SHA256 for a specific subdir pattern:
6265
# python s3_management/manage_v2.py whl/test --recompute-sha256-pattern rocm6.4
66+
#
67+
# # Recompute SHA256 for one release across all prod accelerator subfolders:
68+
# python s3_management/manage_v2.py whl --recompute-sha256-pattern \
69+
# --package-name torch --package-version 2.12.1
6370

6471
import argparse
6572
import base64
@@ -1686,21 +1693,42 @@ def set_checksum_metadata(prefix: str, package_name: str, version: str) -> None:
16861693
_compute_and_set_checksums(matching_objects)
16871694

16881695

1696+
def list_accelerator_subdirs(prefix: str) -> List[str]:
1697+
"""List immediate accelerator subdirectories under a prefix as scan prefixes.
1698+
1699+
Returns scan prefixes like 'whl/cu126/', 'whl/rocm7.1/', 'whl/cpu/' for
1700+
every immediate subdirectory whose name matches ACCEPTED_SUBDIR_PATTERNS.
1701+
Nested channels such as whl/nightly and whl/test are excluded because their
1702+
names do not match the accelerator patterns.
1703+
"""
1704+
scan_prefixes: List[str] = []
1705+
paginator = CLIENT.get_paginator("list_objects_v2")
1706+
for page in paginator.paginate(
1707+
Bucket=BUCKET.name, Prefix=f"{prefix}/", Delimiter="/"
1708+
):
1709+
for common_prefix in page.get("CommonPrefixes", []):
1710+
name = common_prefix["Prefix"].rstrip("/").split("/")[-1]
1711+
if any(match(f"{pat}$", name) for pat in ACCEPTED_SUBDIR_PATTERNS):
1712+
scan_prefixes.append(common_prefix["Prefix"])
1713+
return scan_prefixes
1714+
1715+
16891716
def recompute_sha256_for_pattern(
16901717
prefix: str,
1691-
pattern: str,
1718+
pattern: Optional[str] = None,
16921719
package_name: Optional[str] = None,
16931720
version: Optional[str] = None,
16941721
) -> None:
16951722
"""Compute SHA256 checksums for objects matching a pattern that don't have checksums.
16961723
16971724
Args:
1698-
prefix: The S3 prefix to search in (e.g., "whl/test")
1699-
pattern: The pattern to match against object keys (e.g., "rocm6.4")
1725+
prefix: The S3 prefix to search in (e.g., "whl", "whl/test")
1726+
pattern: The subfolder to scan under the prefix (e.g., "rocm6.4"). When
1727+
empty/None, every accelerator subfolder (cu*/rocm*/cpu/xpu) under the
1728+
prefix is scanned (nightly/test channels are excluded).
17001729
package_name: Optional package name to filter (e.g., "torch", "torchvision")
17011730
version: Optional version to filter (e.g., "2.5.0", "2.5.0+rocm7.1")
17021731
"""
1703-
print(f"INFO: Searching in '{prefix}' for objects matching pattern '{pattern}'")
17041732
normalized_package = None
17051733
if package_name:
17061734
print(f"INFO: Filtering by package name: '{package_name}'")
@@ -1710,44 +1738,54 @@ def recompute_sha256_for_pattern(
17101738
if version:
17111739
print(f"INFO: Filtering by version: '{version}'")
17121740

1741+
# Determine which prefixes to scan. A pattern selects a single subfolder;
1742+
# no pattern scans all accelerator subfolders under the prefix.
1743+
if pattern:
1744+
print(f"INFO: Searching in '{prefix}' for objects matching pattern '{pattern}'")
1745+
scan_prefixes = [f"{prefix}/{pattern}/"]
1746+
else:
1747+
scan_prefixes = list_accelerator_subdirs(prefix)
1748+
names = ", ".join(p.rstrip("/").split("/")[-1] for p in scan_prefixes)
1749+
print(
1750+
f"INFO: No pattern given; scanning all accelerator subfolders under "
1751+
f"'{prefix}/': {names or '(none found)'}"
1752+
)
1753+
17131754
# Find all matching objects
17141755
matching_objects = []
1756+
for scan_prefix in scan_prefixes:
1757+
print(f"INFO: Scanning prefix '{scan_prefix}'...")
1758+
for obj in BUCKET.objects.filter(Prefix=scan_prefix):
1759+
key = obj.key
1760+
# Only process wheel files
1761+
if key.endswith(".whl"):
1762+
basename = path.basename(key).lower()
1763+
# If package_name is specified, filter by it
1764+
if normalized_package:
1765+
# Wheel filename format: {package}-{version}-...
1766+
if not basename.startswith(f"{normalized_package}-"):
1767+
continue
17151768

1716-
# Construct the scan prefix by combining prefix and pattern
1717-
scan_prefix = f"{prefix}/{pattern}/"
1718-
print(f"INFO: Scanning prefix '{scan_prefix}'...")
1719-
1720-
for obj in BUCKET.objects.filter(Prefix=scan_prefix):
1721-
key = obj.key
1722-
# Only process wheel files
1723-
if key.endswith(".whl"):
1724-
basename = path.basename(key).lower()
1725-
# If package_name is specified, filter by it
1726-
if normalized_package:
1727-
# Wheel filename format: {package}-{version}-...
1728-
if not basename.startswith(f"{normalized_package}-"):
1729-
continue
1730-
1731-
# If version is specified, filter by it
1732-
if version:
1733-
# Check for version pattern in the filename
1734-
# Handle both URL-encoded (+) and regular versions
1735-
# Also handle local version specifiers (e.g., 2.9.1+rocm6.4)
1736-
version_encoded = version.replace("+", "%2B").lower()
1737-
version_lower = version.lower()
1738-
# Version can be followed by - (exact match) or + or %2B (local version)
1739-
version_match = (
1740-
f"-{version_encoded}-" in basename
1741-
or f"-{version_lower}-" in basename
1742-
or f"-{version_encoded}+" in basename
1743-
or f"-{version_lower}+" in basename
1744-
or f"-{version_encoded}%2b" in basename
1745-
or f"-{version_lower}%2b" in basename
1746-
)
1747-
if not version_match:
1748-
continue
1769+
# If version is specified, filter by it
1770+
if version:
1771+
# Check for version pattern in the filename
1772+
# Handle both URL-encoded (+) and regular versions
1773+
# Also handle local version specifiers (e.g., 2.9.1+rocm6.4)
1774+
version_encoded = version.replace("+", "%2B").lower()
1775+
version_lower = version.lower()
1776+
# Version can be followed by - (exact match) or + or %2B (local version)
1777+
version_match = (
1778+
f"-{version_encoded}-" in basename
1779+
or f"-{version_lower}-" in basename
1780+
or f"-{version_encoded}+" in basename
1781+
or f"-{version_lower}+" in basename
1782+
or f"-{version_encoded}%2b" in basename
1783+
or f"-{version_lower}%2b" in basename
1784+
)
1785+
if not version_match:
1786+
continue
17491787

1750-
matching_objects.append(key)
1788+
matching_objects.append(key)
17511789

17521790
if not matching_objects:
17531791
filters = []
@@ -1756,7 +1794,8 @@ def recompute_sha256_for_pattern(
17561794
if version:
17571795
filters.append(f"version '{version}'")
17581796
filter_msg = f" for {', '.join(filters)}" if filters else ""
1759-
print(f"WARNING: No matching objects found for pattern '{pattern}'{filter_msg}")
1797+
scope = f"pattern '{pattern}'" if pattern else f"all subfolders of '{prefix}'"
1798+
print(f"WARNING: No matching objects found for {scope}{filter_msg}")
17601799
return
17611800

17621801
print(f"INFO: Found {len(matching_objects)} matching wheel files")
@@ -1864,9 +1903,15 @@ def create_parser() -> argparse.ArgumentParser:
18641903
parser.add_argument(
18651904
"--recompute-sha256-pattern",
18661905
type=str,
1906+
nargs="?",
1907+
const="",
1908+
default=None,
18671909
metavar="PATTERN",
1868-
help="Compute SHA256 checksums for objects matching this pattern that don't already have "
1869-
"checksums (e.g., 'whl/test/rocm7.1'). Objects with existing checksums are skipped.",
1910+
help="Compute SHA256 checksums for .whl objects that don't already have them. "
1911+
"With a PATTERN (e.g. 'rocm7.1') only that subfolder under the prefix is scanned. "
1912+
"Passed with no value, all accelerator subfolders (cu*/rocm*/cpu/xpu) under the "
1913+
"prefix are scanned (nightly/test excluded). Combine with --package-name / "
1914+
"--package-version to scope to a single release. Existing checksums are skipped.",
18701915
)
18711916
parser.add_argument(
18721917
"--recompute-missing-sha256",
@@ -1895,7 +1940,7 @@ def main() -> None:
18951940
return
18961941

18971942
# Handle --recompute-sha256-pattern command
1898-
if args.recompute_sha256_pattern:
1943+
if args.recompute_sha256_pattern is not None:
18991944
recompute_sha256_for_pattern(
19001945
args.prefix,
19011946
args.recompute_sha256_pattern,

0 commit comments

Comments
 (0)