Test use uv for building wheels

Author: atalman

.github/actions/setup-binary-builds/action.yml

@@ -84,75 +84,34 @@ 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
+
+          python -m 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
-            fi
-            conda install -y conda=24.7.1 conda-libmamba-solver=24.1.0
-          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

.github/workflows/build_wheels_linux.yml

@@ -216,7 +216,7 @@ jobs:
           # shellcheck disable=SC1090
           source "${BUILD_ENV_FILE}"
           # 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
@@ -232,8 +232,8 @@ 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
+          export PYTORCH_VERSION="$(uv pip show torch | grep ^Version: | sed 's/Version: *//' | sed 's/+.\+//')"
+          uv python -m pip install build==1.2.2
           echo "Successfully installed Python build package"
           ${CONDA_RUN} ${{ inputs.build-command }}
       - name: Build the wheel (setup-py)
@@ -243,10 +243,10 @@ 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
+          export PYTORCH_VERSION="$(uv pip show torch | grep ^Version: | sed 's/Version: *//' | sed 's/+.\+//')"
+          uv python setup.py clean
           echo "Successfully ran `python setup.py clean`"
-          ${CONDA_RUN} python setup.py bdist_wheel
+          uv python setup.py bdist_wheel
       - name: Repair Manylinux_2_28 Wheel
         shell: bash -l {0}
         env:
@@ -280,7 +280,7 @@ jobs:
           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}"
@@ -295,13 +295,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}"
+              uv python -c "import ${PACKAGE_NAME}"
             else
-              ${CONDA_RUN} python -c "import ${PACKAGE_NAME}; print('package version is ', ${PACKAGE_NAME}.__version__)"
+              uv 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}"
+            uv python "${{ inputs.repository }}/${SMOKE_TEST_SCRIPT}"
           fi
       # NB: Only upload to GitHub after passing smoke tests