Skip to content

Use cached PyTorch wheels on MacOS jobs #9484

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 10 commits into from
Mar 24, 2025
Merged

Use cached PyTorch wheels on MacOS jobs #9484

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
eaed605
Upload PyTorch MacOS wheel to S3
huydhn Mar 21, 2025
3071d26
No ACL
huydhn Mar 21, 2025
947c39c
Implement reading from the cache
huydhn Mar 21, 2025
9cbe26a
Fix typo
huydhn Mar 21, 2025
1576f32
Another attempt
huydhn Mar 21, 2025
fbd4949
Fix a bug
huydhn Mar 21, 2025
b47eb88
Merge branch 'main' into cache-macos-pytorch-build-artifact
huydhn Mar 21, 2025
bc6baf3
Fix the platform version
huydhn Mar 21, 2025
3b1938f
Address review comments
huydhn Mar 24, 2025
34fed00
Another tweak
huydhn Mar 24, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
44 changes: 39 additions & 5 deletions .ci/scripts/utils.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -60,12 +60,46 @@ install_pytorch_and_domains() {
# Fetch the target commit
pushd pytorch || return
git checkout "${TORCH_VERSION}"
git submodule update --init --recursive

export USE_DISTRIBUTED=1
# Then build and install PyTorch
python setup.py bdist_wheel
pip install "$(echo dist/*.whl)"
local system_name=$(uname)
if [[ "${system_name}" == "Darwin" ]]; then
local platform=$(python -c 'import sysconfig; import platform; v=platform.mac_ver()[0].split(".")[0]; platform=sysconfig.get_platform().split("-"); platform[1]=f"{v}_0"; print("_".join(platform))')
fi
local python_version=$(python -c 'import platform; v=platform.python_version_tuple(); print(f"{v[0]}{v[1]}")')
local torch_release=$(cat version.txt)
local torch_short_hash=${TORCH_VERSION:0:7}
local torch_wheel_path="cached_artifacts/pytorch/executorch/pytorch_wheels/${system_name}/${python_version}"
local torch_wheel_name="torch-${torch_release}%2Bgit${torch_short_hash}-cp${python_version}-cp${python_version}-${platform:-}.whl"

local cached_torch_wheel="https://gha-artifacts.s3.us-east-1.amazonaws.com/${torch_wheel_path}/${torch_wheel_name}"
# Cache PyTorch wheel is only needed on MacOS, Linux CI already has this as part
# of the Docker image
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Don't you need to set default value for TORCH_WHEEL_NOT_FOUND (to handle non Darwin case)

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

True, this function is currently used only on MacOS, but I remember reading that we can now build ExecuTorch on Windows too

local torch_wheel_not_found=0
if [[ "${system_name}" == "Darwin" ]]; then
pip install "${cached_torch_wheel}" || torch_wheel_not_found=1
else
torch_wheel_not_found=1
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

remove this else statement and just set the local torch_wheel_not_found=1?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Oh, I want it to default to 0 (found the wheel) because pip install "${cached_torch_wheel}" || torch_wheel_not_found=1 will set it to 1 when the pip command fails or not MacOS

fi

# Found no such wheel, we will build it from source then
if [[ "${torch_wheel_not_found}" == "1" ]]; then
echo "No cached wheel found, continue with building PyTorch at ${TORCH_VERSION}"

git submodule update --init --recursive
USE_DISTRIBUTED=1 python setup.py bdist_wheel
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

log that we're building from source

pip install "$(echo dist/*.whl)"

# Only AWS runners have access to S3
if command -v aws && [[ -z "${GITHUB_RUNNER:-}" ]]; then
for wheel_path in dist/*.whl; do
local wheel_name=$(basename "${wheel_path}")
echo "Caching ${wheel_name}"
aws s3 cp "${wheel_path}" "s3://gha-artifacts/${torch_wheel_path}/${wheel_name}"
done
fi
else
echo "Use cached wheel at ${cached_torch_wheel}"
fi

# Grab the pinned audio and vision commits from PyTorch
TORCHAUDIO_VERSION=$(cat .github/ci_commit_pins/audio.txt)
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We can also cache audio, vision, and other wheels, but the gain is probably smaller because it's fast to build them. This can come in subsequent PRs.

Expand Down
2 changes: 2 additions & 0 deletions .github/workflows/_unittest.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -49,4 +49,6 @@ jobs:
ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
script: |
set -eux
# This is needed to get the prebuilt PyTorch wheel from S3
${CONDA_RUN} --no-capture-output pip install awscli==1.37.21
.ci/scripts/unittest-macos.sh --build-tool "${{ inputs.build-tool }}" --build-mode "${{ inputs.build-mode }}" --editable "${{ inputs.editable }}"
2 changes: 1 addition & 1 deletion .github/workflows/trunk.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -228,7 +228,7 @@ jobs:
name: test-coreml-delegate
uses: pytorch/test-infra/.github/workflows/macos_job.yml@main
with:
runner: macos-13-xlarge
runner: macos-latest-xlarge
python-version: '3.11'
submodules: 'true'
ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
Expand Down
Loading