Skip to content

[DRAFT] Test use uv for building wheels #6626

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.

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

Draft
wants to merge 7 commits into
base: main
Choose a base branch
from
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
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
72 changes: 15 additions & 57 deletions .github/actions/setup-binary-builds/action.yml
Original file line number Diff line number Diff line change
Expand Up @@ -84,77 +84,35 @@ runs:
# and "/" is not allowed in artifact names. //\//_ is to replace all forward slashes,
# not just the first one
echo "ARTIFACT_NAME=${REPOSITORY//\//_}_${REF//\//_}_${PYTHON_VERSION}_${CU_VERSION}_${ARCH}" >> "${GITHUB_ENV}"
- name: Setup miniconda (for pytorch_pkg_helpers)
if: ${{ inputs.setup-miniconda == 'true' }}
uses: conda-incubator/[email protected]
with:
miniconda-version: "latest"
python-version: 3.9
- name: Clean conda environment
shell: bash -l {0}
run: |
set -euxo pipefail
conda info | grep -i 'base environment'
conda clean --all --quiet --yes
- name: Reset channel priority

- name: Install uv environment
shell: bash -l {0}
run: |
set -euxo pipefail
conda config --set channel_priority false
curl -LsSf https://astral.sh/uv/install.sh | sh

- name: Generate file from pytorch_pkg_helpers
working-directory: ${{ inputs.repository }}
shell: bash -l {0}
run: |
set -euxo pipefail
CONDA_ENV="${RUNNER_TEMP}/pytorch_pkg_helpers_${GITHUB_RUN_ID}"
conda create \
--yes --quiet \
--prefix "${CONDA_ENV}" \
"python=3.9"
CONDA_ENV="${CONDA_ENV}"
CONDA_RUN="conda run -p ${CONDA_ENV}"
${CONDA_RUN} python -m pip install ${GITHUB_WORKSPACE}/test-infra/tools/pkg-helpers
uv venv --python 3.9
source .venv/bin/activate

uv pip install ${GITHUB_WORKSPACE}/test-infra/tools/pkg-helpers
BUILD_ENV_FILE="${RUNNER_TEMP}/build_env_${GITHUB_RUN_ID}"
${CONDA_RUN} python -m pytorch_pkg_helpers > "${BUILD_ENV_FILE}"
python -m pytorch_pkg_helpers > "${BUILD_ENV_FILE}"
cat "${BUILD_ENV_FILE}"
echo "BUILD_ENV_FILE=${BUILD_ENV_FILE}" >> "${GITHUB_ENV}"
- name: Setup conda environment for build

- name: Setup uv environment for build
shell: bash -l {0}
env:
PYTHON_VERSION: ${{ inputs.python-version }}
run: |
set -euxo pipefail
CONDA_ENV="${RUNNER_TEMP}/conda_environment_${GITHUB_RUN_ID}"
export CONDA_EXTRA_PARAM=""

if [[ "${PYTHON_VERSION:-}" == "3.13t" ]]; then
export PYTHON_VERSION=3.13
export CONDA_EXTRA_PARAM=" python-freethreading -c conda-forge"

# downgrade conda version for python 3.13t install.
# TODO: remove this once python 3.13t is fully supported on conda
# Please see : https://github.com/conda/conda/issues/14554
if [[ "$(uname)" == Darwin ]]; then
# required to be able to downgrade on MacOS arm64
conda install -y python=3.9
if [[ -n "$(conda list | grep conda-anaconda-telemetry)" ]]; then
conda uninstall -y conda-anaconda-telemetry conda-anaconda-tos
fi
conda install -y conda=24.7.1 conda-libmamba-solver=24.1.0
else
conda install -y conda=24.7.1
fi
fi

conda create \
--yes --quiet \
--prefix "${CONDA_ENV}" \
"python=${PYTHON_VERSION}" \
cmake=3.31.2 \
ninja=1.12.1 \
pkg-config=0.29 \
wheel=0.37 \
${CONDA_EXTRA_PARAM}

echo "CONDA_ENV=${CONDA_ENV}" >> "${GITHUB_ENV}"
echo "CONDA_RUN=conda run -p ${CONDA_ENV}" >> "${GITHUB_ENV}"
uv venv py${PYTHON_VERSION} --python ${PYTHON_VERSION}
source ./py${PYTHON_VERSION}/bin/activate
uv pip install cmake==3.31.2 ninja wheel==0.37 setuptools
echo "VIRTUAL_ENV=${VIRTUAL_ENV}" >> "${GITHUB_ENV}"
25 changes: 15 additions & 10 deletions .github/workflows/build_wheels_linux.yml
Original file line number Diff line number Diff line change
Expand Up @@ -215,8 +215,9 @@ jobs:
set -euxo pipefail
# shellcheck disable=SC1090
source "${BUILD_ENV_FILE}"
source ${VIRTUAL_ENV}/bin/activate
# shellcheck disable=SC2086
${CONDA_RUN} ${PIP_INSTALL_TORCH} ${{ inputs.pip-install-torch-extra-args }}
uv ${PIP_INSTALL_TORCH} ${{ inputs.pip-install-torch-extra-args }}
- name: Run Pre-Script with Caching
if: ${{ inputs.pre-script != '' }}
uses: ./test-infra/.github/actions/run-script-with-cache
Expand All @@ -232,8 +233,9 @@ jobs:
run: |
set -euxo pipefail
source "${BUILD_ENV_FILE}"
export PYTORCH_VERSION="$(${CONDA_RUN} pip show torch | grep ^Version: | sed 's/Version: *//' | sed 's/+.\+//')"
${CONDA_RUN} python -m pip install build==1.2.2
source ${VIRTUAL_ENV}/bin/activate
export PYTORCH_VERSION="$(uv pip show torch | grep ^Version: | sed 's/Version: *//' | sed 's/+.\+//')"
uv pip install build==1.2.2
echo "Successfully installed Python build package"
${CONDA_RUN} ${{ inputs.build-command }}
- name: Build the wheel (setup-py)
Expand All @@ -243,10 +245,11 @@ jobs:
run: |
set -euxo pipefail
source "${BUILD_ENV_FILE}"
export PYTORCH_VERSION="$(${CONDA_RUN} pip show torch | grep ^Version: | sed 's/Version: *//' | sed 's/+.\+//')"
${CONDA_RUN} python setup.py clean
source ${VIRTUAL_ENV}/bin/activate
export PYTORCH_VERSION="$(uv pip show torch | grep ^Version: | sed 's/Version: *//' | sed 's/+.\+//')"
python setup.py clean
echo "Successfully ran `python setup.py clean`"
${CONDA_RUN} python setup.py bdist_wheel
python setup.py bdist_wheel
- name: Repair Manylinux_2_28 Wheel
shell: bash -l {0}
env:
Expand All @@ -255,6 +258,7 @@ jobs:
run: |
set -euxo pipefail
source "${BUILD_ENV_FILE}"
source ${VIRTUAL_ENV}/bin/activate
for pkg in ${{ inputs.repository }}/dist/*-linux_*.whl; do
# if the glob didn't match anything
if [[ ! -e $pkg ]]; then
Expand All @@ -277,10 +281,11 @@ jobs:
run: |
set -euxo pipefail
source "${BUILD_ENV_FILE}"
source ${VIRTUAL_ENV}/bin/activate
WHEEL_NAME=$(ls "${{ inputs.repository }}/dist/")
echo "$WHEEL_NAME"

${CONDA_RUN} pip install "${{ inputs.repository }}/dist/$WHEEL_NAME"
uv pip install "${{ inputs.repository }}/dist/$WHEEL_NAME"
# Checking that we have a pinned version of torch in our dependency tree
(
pushd "${RUNNER_TEMP}"
Expand All @@ -295,13 +300,13 @@ jobs:
# Special case for torchrec temporarily since __version__ does not
# work correctly on main in torchrec. This block will be
# removed once we fix it.
${CONDA_RUN} python -c "import ${PACKAGE_NAME}"
python -c "import ${PACKAGE_NAME}"
else
${CONDA_RUN} python -c "import ${PACKAGE_NAME}; print('package version is ', ${PACKAGE_NAME}.__version__)"
python -c "import ${PACKAGE_NAME}; print('package version is ', ${PACKAGE_NAME}.__version__)"
fi
else
echo "${{ inputs.repository }}/${SMOKE_TEST_SCRIPT} found"
${CONDA_RUN} python "${{ inputs.repository }}/${SMOKE_TEST_SCRIPT}"
python "${{ inputs.repository }}/${SMOKE_TEST_SCRIPT}"
fi
# NB: Only upload to GitHub after passing smoke tests

Expand Down
Loading