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
huydhn 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 8 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
36 changes: 32 additions & 4 deletions .ci/scripts/utils.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -62,10 +62,38 @@ install_pytorch_and_domains() {
git checkout "${TORCH_VERSION}"
git submodule update --init --recursive
Copy link
Contributor

@mergennachin mergennachin Mar 24, 2025

Choose a reason for hiding this comment

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

can move this command (cloning all submodules) when we haven't found the cache entry?

Copy link
Contributor

@mergennachin mergennachin Mar 24, 2025

Choose a reason for hiding this comment

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

i think this will reduce even further

Copy link
Contributor Author

@huydhn huydhn Mar 24, 2025

Choose a reason for hiding this comment

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

Yeah, good catch, I only need the version.txt from PyTorch


export USE_DISTRIBUTED=1
# Then build and install PyTorch
python setup.py bdist_wheel
pip install "$(echo dist/*.whl)"
SYSTEM_NAME=$(uname)
if [[ "${SYSTEM_NAME}" == "Darwin" ]]; then
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
PYTHON_VERSION=$(python -c 'import platform; v=platform.python_version_tuple(); print(f"{v[0]}{v[1]}")')
TORCH_RELEASE=$(cat version.txt)
TORCH_SHORT_HASH=${TORCH_VERSION:0:7}
TORCH_WHEEL_PATH="cached_artifacts/pytorch/executorch/pytorch_wheels/${SYSTEM_NAME}/${PYTHON_VERSION}"
TORCH_WHEEL_NAME="torch-${TORCH_RELEASE}%2Bgit${TORCH_SHORT_HASH}-cp${PYTHON_VERSION}-cp${PYTHON_VERSION}-${PLATFORM:-}.whl"
Copy link
Contributor Author

@huydhn huydhn Mar 22, 2025 •
edited
Loading

Choose a reason for hiding this comment

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

I'm open to suggestions here on how to figure out the name of the required PyTorch wheel. When building from source, its name is like torch-2.7.0a0+git295f2ed-cp311-cp311-macosx_14_0_arm64.whl. Maybe there is a way to get this from PyTorch setup.py without actually building the wheel.

Copy link
Contributor

@mergennachin mergennachin Mar 24, 2025

Choose a reason for hiding this comment

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

should we have local variables instead of global env variables

local system_name, torch_release etc?

Copy link
Contributor Author

@huydhn huydhn Mar 24, 2025

Choose a reason for hiding this comment

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

Just FYI, after updating these variables with local, I look around and find some feature request on shellcheck about this koalaman/shellcheck#468, but it hasn't been implemented yet (probably not anytime soon)


CACHE_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

@mergennachin mergennachin Mar 24, 2025

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

@huydhn huydhn Mar 24, 2025

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

if [[ "${SYSTEM_NAME}" == "Darwin" ]]; then
pip install "${CACHE_TORCH_WHEEL}" || TORCH_WHEEL_NOT_FOUND=1
Copy link
Contributor

@mergennachin mergennachin Mar 24, 2025

Choose a reason for hiding this comment

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

can you log when no cache is found, and log the wheel name?

fi

# Found no such wheel, we will build it from source then
if [[ "${TORCH_WHEEL_NOT_FOUND:-0}" == "1" ]]; then
USE_DISTRIBUTED=1 python setup.py bdist_wheel
Copy link
Contributor

@mergennachin mergennachin Mar 24, 2025

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
WHEEL_NAME=$(basename "${WHEEL_PATH}")
Copy link
Contributor

@mergennachin mergennachin Mar 24, 2025

Choose a reason for hiding this comment

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

log the name of the wheel that's being uploaded

aws s3 cp "${WHEEL_PATH}" "s3://gha-artifacts/${TORCH_WHEEL_PATH}/${WHEEL_NAME}"
done
fi
else
echo "Use cached wheel at ${CACHE_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

@huydhn huydhn Mar 22, 2025

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