[CI] Add gating mechanism to ensure PyTorch wheels pass tests before uploading to S3 release bucket

[araravik-psd]

• Purpose
  We aim to introduce gating checks for the PyTorch build and release process by refactoring the current workflow to ensure that tests are executed and pass successfully before artifacts are uploaded to the release S3 bucket.

• Proposed Approach
  The idea is to reuse the existing S3 bucket but omit a versioned or release-specific index during initial upload. This prevents direct exposure of unvalidated artifacts as release assets. Artifacts will only be promoted to the release path if all validation steps succeed.

Refer to the existing workflow:

build_portable_linux_pytorch_wheels.yml

• Current Workflow Design
  Job 1: Checks out the PyTorch source repositories, builds the wheels, and uploads artifacts directly to the release S3 bucket.

Job 2: Dynamically determines the test targets to run.

Job 3: Executes PyTorch test suite on appropriate targets.

• Proposed Workflow Refactor
  Job 1: Clone the PyTorch source repositories and build the wheels. Instead of uploading directly to the release path, the artifacts are temporarily stored in a staging directory within the same S3 bucket.

Job 2: Dynamically determine the appropriate test targets by running a configuration script.

Job 3: Execute tests against the generated wheels across the identified targets.

Job 4: Upon successful completion of all test jobs, promote the staged artifacts by copying them to the official release directory in the S3 bucket.

[ScottTodd]

I'm skeptical of using a subdirectory in the existing "nightly" bucket and then moving. That has the potential to leave us with files in an unstable state and may make it harder to have a predictable way to install releases that are either still being tested or failed the tests.

Here are some alternatives to consider:

A) Upload to "nightly staging" then copy to "nightly"

Always upload to a separate "nightly staging" bucket. After tests pass, copy to the "nightly" bucket.

There would be no destructive moves but files would be stored twice. We could have separate index pages for each, with the ability to perform uninterrupted bisects across "nightly staging" and we could always copy a release across manually if we want to override the CI results.

B) Upload to "dev" then copy to "nightly"

Always upload to the same "dev" bucket as test runs from workflow_dispatch testing, pull_request presubmit runs, and push postsubmit runs. After tests pass, copy to the "nightly" bucket.

Similar to above, but we'd have one less index page and would instead mix versions in the dev bucket. Pip install from the dev bucket would then automatically use "nightly" versions or explicit commit hash versioned "dev" versions... I think.

[ScottTodd]

C) Upload to only one location, filter in the index

We could keep all the files in the bucket and modify https://github.com/ROCm/TheRock/blob/main/build_tools/packaging/python/generate_release_index.py to only list files that passed testing. Perhaps by checking for the existence of a .tests_passed file or the absence of a .tests_pending file? Seems messy to maintain.

[marbre]

To maybe make things clearer

• At the moment packages get uploaded to
  • v2/gfx110X-dgpu
  • v2/gfx1151
  • v2/gfx120X-all
  • v2/gfx94X-dcgpu
  • v2/gfx950-dcgpu
• For each of the directories above an index is created.

A) Upload to "nightly staging" then copy to "nightly"

I don't think there is a strict need to have a separate bucket but we should not have a subdirectory under v2. Instead we could have v2-staging at the root level of the bucket.

The only real benefit I see is that we could disable Object Lock if the staging bucket is separate.

B) Upload to "dev" then copy to "nightly"

In the longer term we want to split out how nightly releases are generated and one cannot grantee that staging files get compromised if stored in a bucket that is still writable from the dev workflows. Furthermore, the combined index might be troublesome. dev version can indeed be considered newer as nightly releases with is a downside of the local version identifiers we use.

C) Upload to only one location, filter in the index

This could be messy to maintain but might indeed a clean solution. Btw, https://github.com/ROCm/TheRock/blob/main/build_tools/packaging/python/generate_release_index.py isn't used any longer and we switched to https://github.com/ROCm/TheRock/blob/main/build_tools/third_party/s3_management/manage.py.

---

Other things to consider

Beside the options discussed above, we also need to take into account that if we want to test torch wheels, and in particular its installation!, the way a user do, we must to be able to install via pip install -i ${INDEX} torch==${VERSION}. I see the following options:

I)

The ROCm packages are additionally uploaded to the staging directory / bucket alongside the torch wheels.

II)

The ROCm packages get added to the index in the staging directory. Technically there is no need to store the packages in the same bucket / location. This is kind of related to C) but maintenance should be easier.

III)

If it is acceptable to don't have a single pip install command we could split the installation and pip install ROCm from the release bucket and afterwards install torch from the staging bucket.

At least for II and III this assumes that ROCm packages are always pushed to the nightlies bucket and we want to release them even though torch tests might have failed.

[ScottTodd]

I)

The ROCm packages are additionally uploaded to the staging directory / bucket alongside the torch wheels.

Yes, I'm working under that assumption. We upload all files to staging then only promote them as a full set if our acceptance bar is met. If PyTorch tests failed, the ROCm packages should not be promoted either. We can make that line blurrier for flaky tests and projects other than PyTorch eventually, but let's start there.

---

Another option to consider is using "prerelease" versions of packages and modifying them with a script like https://github.com/iree-org/iree/blob/main/build_tools/python_deploy/promote_whl_from_rc_to_final.py. Then we could switch from installing "passing" nightly releases to "all" nightly releases by just adding --pre to a pip install command.

I like the v2-staging directory idea!

[marbre]

Another option to consider is using "prerelease" versions of packages and modifying them with a script like https://github.com/iree-org/iree/blob/main/build_tools/python_deploy/promote_whl_from_rc_to_final.py. Then we could switch from installing "passing" nightly releases to "all" nightly releases by just adding --pre to a pip install command.

At least all our ROCm packages are release candidates and therefore pre-releases.

[ScottTodd]

As part of this, we'll need the v2-staging subdirectory in each bucket that it exists in to have a complete index, which would either be standalone or reference the v2 index. This will avoid errors like https://github.com/ROCm/TheRock/actions/runs/16454515470/job/46574141517#step:5:43

++ Exec [/home/runner/_work/TheRock/TheRock]$ /home/runner/_work/TheRock/TheRock/.venv/bin/python -m pip install --index-url=https://d25kgig7rdsyks.cloudfront.net/v2-staging/gfx94X-dcgpu torch==2.7.0a0+rocm7.0.0.dev0.515115ea2cb85a0b71b5507ce56a627d14c7ae73
Looking in indexes: https://d25kgig7rdsyks.cloudfront.net/v2-staging/gfx94X-dcgpu
Collecting torch==2.7.0a0+rocm7.0.0.dev0.515115ea2cb85a0b71b5507ce56a627d14c7ae73
  Downloading https://d25kgig7rdsyks.cloudfront.net/v2-staging/gfx94X-dcgpu/torch-2.7.0a0%2Brocm7.0.0.dev0.515115ea2cb85a0b71b5507ce56a627d14c7ae73-cp312-cp312-linux_x86_64.whl (632.0 MB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 632.0/632.0 MB 87.8 MB/s eta 0:00:00
INFO: pip is looking at multiple versions of torch to determine which version is compatible with other requirements. This could take a while.
ERROR: Could not find a version that satisfies the requirement filelock (from torch) (from versions: none)
ERROR: No matching distribution found for filelock

For now I ran https://github.com/ROCm/TheRock/blob/main/build_tools/third_party/s3_management/manage.py and https://github.com/ROCm/TheRock/blob/main/build_tools/third_party/s3_management/update_dependencies.py locally with v2 changed to v2-staging and therock-nightly-python changed to therock-dev-python.

[araravik-psd]

@stellaraccident @amd-chrissosa , please let us know if the expectations documenting below are in order for gating torch wheels before we merge #1110 .

Documenting the expectations with gating torch builds using this PR #1110 after discussions with @marbre and @ScottTodd

1. v2-staging will have all the torch wheels for different architectures built, but they might not be tested
2. Promotion from v2-staging to v2 will happen only if the wheels have been tested
3. Promotion from v2-staging to v2 will not happen if our CI does provide runners for running tests on built wheels
4. v2 will have all tested torch wheels for different architectures
5. Previously uploaded torch wheels to v2 are not copied to v2-staging and are not deleted from v2. Upload to v2-staging and conditional promotion from v2-staging to v2 will happen once this PR is merged