mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 12:54:11 +08:00
50d418f69f
Previously we already replaced most use of `python setup.py develop/install`. This PR also replaces the use of `setup.py bdist_wheel` with the modern `python -m build --wheel` alternative. Pull Request resolved: https://github.com/pytorch/pytorch/pull/156712 Approved by: https://github.com/atalman ghstack dependencies: #156711
108 lines
3.6 KiB
Bash
Executable File
108 lines
3.6 KiB
Bash
Executable File
#!/bin/bash
|
|
|
|
set -ex
|
|
|
|
mkdir -p /opt/triton
|
|
if [ -z "${TRITON}" ] && [ -z "${TRITON_CPU}" ]; then
|
|
echo "TRITON and TRITON_CPU are not set. Exiting..."
|
|
exit 0
|
|
fi
|
|
|
|
source "$(dirname "${BASH_SOURCE[0]}")/common_utils.sh"
|
|
|
|
get_pip_version() {
|
|
conda_run pip list | grep -w $* | head -n 1 | awk '{print $2}'
|
|
}
|
|
|
|
if [ -n "${XPU_VERSION}" ]; then
|
|
TRITON_REPO="https://github.com/intel/intel-xpu-backend-for-triton"
|
|
TRITON_TEXT_FILE="triton-xpu"
|
|
elif [ -n "${TRITON_CPU}" ]; then
|
|
TRITON_REPO="https://github.com/triton-lang/triton-cpu"
|
|
TRITON_TEXT_FILE="triton-cpu"
|
|
else
|
|
TRITON_REPO="https://github.com/triton-lang/triton"
|
|
TRITON_TEXT_FILE="triton"
|
|
fi
|
|
|
|
# The logic here is copied from .ci/pytorch/common_utils.sh
|
|
TRITON_PINNED_COMMIT=$(get_pinned_commit ${TRITON_TEXT_FILE})
|
|
|
|
if [ -n "${UBUNTU_VERSION}" ];then
|
|
apt update
|
|
apt-get install -y gpg-agent
|
|
fi
|
|
|
|
# Keep the current cmake and numpy version here, so we can reinstall them later
|
|
CMAKE_VERSION=$(get_pip_version cmake)
|
|
NUMPY_VERSION=$(get_pip_version numpy)
|
|
|
|
if [ -z "${MAX_JOBS}" ]; then
|
|
export MAX_JOBS=$(nproc)
|
|
fi
|
|
|
|
# Git checkout triton
|
|
mkdir /var/lib/jenkins/triton
|
|
chown -R jenkins /var/lib/jenkins/triton
|
|
chgrp -R jenkins /var/lib/jenkins/triton
|
|
pushd /var/lib/jenkins/
|
|
|
|
as_jenkins git clone --recursive ${TRITON_REPO} triton
|
|
cd triton
|
|
as_jenkins git checkout ${TRITON_PINNED_COMMIT}
|
|
as_jenkins git submodule update --init --recursive
|
|
|
|
# Old versions of python have setup.py in ./python; newer versions have it in ./
|
|
if [ ! -f setup.py ]; then
|
|
cd python
|
|
fi
|
|
|
|
pip_install pybind11==3.0.1
|
|
|
|
# TODO: remove patch setup.py once we have a proper fix for https://github.com/triton-lang/triton/issues/4527
|
|
as_jenkins sed -i -e 's/https:\/\/tritonlang.blob.core.windows.net\/llvm-builds/https:\/\/oaitriton.blob.core.windows.net\/public\/llvm-builds/g' setup.py
|
|
|
|
if [ -n "${UBUNTU_VERSION}" ] && [ -n "${GCC_VERSION}" ] && [[ "${GCC_VERSION}" == "7" ]]; then
|
|
# Triton needs at least gcc-9 to build
|
|
apt-get install -y g++-9
|
|
|
|
CXX=g++-9 conda_run python -m build --wheel --no-isolation
|
|
elif [ -n "${UBUNTU_VERSION}" ] && [ -n "${CLANG_VERSION}" ]; then
|
|
# Triton needs <filesystem> which surprisingly is not available with clang-9 toolchain
|
|
add-apt-repository -y ppa:ubuntu-toolchain-r/test
|
|
apt-get install -y g++-9
|
|
|
|
CXX=g++-9 conda_run python -m build --wheel --no-isolation
|
|
else
|
|
conda_run python -m build --wheel --no-isolation
|
|
fi
|
|
|
|
# Copy the wheel to /opt for multi stage docker builds
|
|
cp dist/*.whl /opt/triton
|
|
# Install the wheel for docker builds that don't use multi stage
|
|
pip_install dist/*.whl
|
|
|
|
# TODO: This is to make sure that the same cmake and numpy version from install conda
|
|
# script is used. Without this step, the newer cmake version (3.25.2) downloaded by
|
|
# triton build step via pip will fail to detect conda MKL. Once that issue is fixed,
|
|
# this can be removed.
|
|
#
|
|
# The correct numpy version also needs to be set here because conda claims that it
|
|
# causes inconsistent environment. Without this, conda will attempt to install the
|
|
# latest numpy version, which fails ASAN tests with the following import error: Numba
|
|
# needs NumPy 1.20 or less.
|
|
# Note that we install numpy with pip as conda might not have the version we want
|
|
if [ -n "${CMAKE_VERSION}" ]; then
|
|
pip_install "cmake==${CMAKE_VERSION}"
|
|
fi
|
|
if [ -n "${NUMPY_VERSION}" ]; then
|
|
pip_install "numpy==${NUMPY_VERSION}"
|
|
fi
|
|
|
|
# IMPORTANT: helion needs to be installed without dependencies.
|
|
# It depends on torch and triton. We don't want to install
|
|
# triton and torch from production on Docker CI images
|
|
if [[ "$ANACONDA_PYTHON_VERSION" != 3.9* ]]; then
|
|
pip_install helion --no-deps
|
|
fi
|