From 12ba846698e2d0dfaf5ff4be2f07ecf2b9c5d7bf Mon Sep 17 00:00:00 2001 From: Johnny Date: Fri, 25 Apr 2025 17:52:14 +0200 Subject: [PATCH 1/5] Update setup.py --- setup.py | 31 ++++++++++++++++++++++++++++--- 1 file changed, 28 insertions(+), 3 deletions(-) diff --git a/setup.py b/setup.py index 5da43ce..9025744 100644 --- a/setup.py +++ b/setup.py @@ -119,12 +119,15 @@ def get_minference_version() -> str: return str(version) -def get_platform(): +def get_arch(): """ - Returns the platform name as used in wheel filenames. + Returns the system aarch for the current system. """ if sys.platform.startswith("linux"): - return f"linux_{platform.uname().machine}" + if platform.machine() == "x86_64": + return "x86_64" + if platform.machine() == "arm64" or platform.machine() == "aarch64": + return "aarch64" elif sys.platform == "darwin": mac_version = ".".join(platform.mac_ver()[0].split(".")[:2]) return f"macosx_{mac_version}_x86_64" @@ -134,6 +137,28 @@ def get_platform(): raise ValueError("Unsupported platform: {}".format(sys.platform)) +def get_system() -> str: + """ + Returns the system name as used in wheel filenames. + """ + if platform.system() == "Windows": + return "win" + elif platform.system() == "Darwin": + mac_version = ".".join(platform.mac_ver()[0].split(".")[:1]) + return f"macos_{mac_version}" + elif platform.system() == "Linux": + return "linux" + else: + raise ValueError("Unsupported system: {}".format(platform.system())) + + +def get_platform() -> str: + """ + Returns the platform name as used in wheel filenames. + """ + return f"{get_system()}_{get_arch()}" + + def get_wheel_url(): # Determine the version numbers that will be used to determine the correct wheel # We're using the CUDA version used to build torch, not the one currently installed From 1b7f6ddb13132e35bf237692367d32fbcb3d4ca7 Mon Sep 17 00:00:00 2001 From: johnnynunez Date: Fri, 25 Apr 2025 18:45:31 +0200 Subject: [PATCH 2/5] Feature(build): enhance CUDA architecture support and improve installation scripts --- .github/workflows/release.yml | 46 ++++++++++---------- .github/workflows/scripts/build.sh | 2 +- .github/workflows/scripts/cuda-install.sh | 31 ++++++++++++- .github/workflows/scripts/pytorch-install.sh | 4 +- .github/workflows/unittest.yml | 4 +- 5 files changed, 59 insertions(+), 28 deletions(-) diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 30a490b..ddf6a41 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -42,25 +42,27 @@ jobs: wheel: name: Build Wheel - runs-on: ${{ matrix.os }} needs: release - + runs-on: ${{ matrix.os }} strategy: fail-fast: false matrix: - os: ['ubuntu-20.04'] - python-version: ['3.9', '3.10', '3.11', '3.12', '3.13'] - pytorch-version: ['2.2.2', '2.3.1', '2.4.0', '2.5.1', '2.6.0'] - cuda-version: ['12.4.0'] - exclude: - # see https://github.com/pytorch/pytorch/blob/main/RELEASE.md#release-compatibility-matrix - # Pytorch < 2.5 does not support Python 3.13 - - pytorch-version: '2.2.2' - python-version: '3.13' - - pytorch-version: '2.3.1' - python-version: '3.13' - - pytorch-version: '2.4.0' - python-version: '3.13' + # Using ubuntu-22.04 instead of 24.04 for more compatibility (glibc). Ideally we'd use the + # manylinux docker image, but I haven't figured out how to install CUDA on manylinux. + os: [ ubuntu-22.04, ubuntu-22.04-arm ] + python-version: [ '3.9', '3.10', '3.11', '3.12', '3.13' ] + torch-version: [ '2.4.0', '2.5.1', '2.6.0', '2.7.0' ] + cuda-version: [ '12.4.1' ] + # We need separate wheels that either uses C++11 ABI (-D_GLIBCXX_USE_CXX11_ABI) or not. + # Pytorch wheels currently don't use it, but nvcr images have Pytorch compiled with C++11 ABI. + # Without this we get import error (undefined symbol: _ZN3c105ErrorC2ENS_14SourceLocationESs) + # when building without C++11 ABI and using it on nvcr images. + cxx11_abi: [ 'FALSE', 'TRUE' ] + exclude: + # see https://github.com/pytorch/pytorch/blob/main/RELEASE.md#release-compatibility-matrix + # Pytorch < 2.5 does not support Python 3.13 + - torch-version: '2.4.0' + python-version: '3.13' steps: - name: Checkout @@ -90,7 +92,7 @@ jobs: - name: Install CUDA ${{ matrix.cuda-version }} if: ${{ matrix.cuda-version != 'cpu' }} - uses: Jimver/cuda-toolkit@v0.2.19 + uses: Jimver/cuda-toolkit@v0.2.23 id: cuda-toolkit with: cuda: ${{ matrix.cuda-version }} @@ -142,10 +144,10 @@ jobs: strategy: fail-fast: false matrix: - os: ['ubuntu-20.04'] - python-version: ['3.10'] - pytorch-version: ['2.3.0'] # Must be the most recent version that meets requirements-cuda.txt. - cuda-version: ['12.2.2'] + os: ['ubuntu-latest'] + python-version: ['3.12'] + pytorch-version: ['2.7.0'] # Must be the most recent version that meets requirements-cuda.txt. + cuda-version: [ '12.4.1' ] steps: - name: Checkout @@ -163,7 +165,7 @@ jobs: bash -x .github/workflows/scripts/env.sh - name: Set up Python - uses: actions/setup-python@v4 + uses: actions/setup-python@v5 with: python-version: ${{ matrix.python-version }} @@ -175,7 +177,7 @@ jobs: - name: Install CUDA ${{ matrix.cuda-version }} if: ${{ matrix.cuda-version != 'cpu' }} - uses: Jimver/cuda-toolkit@v0.2.14 + uses: Jimver/cuda-toolkit@v0.2.23 id: cuda-toolkit with: cuda: ${{ matrix.cuda-version }} diff --git a/.github/workflows/scripts/build.sh b/.github/workflows/scripts/build.sh index f830091..f4a0e26 100644 --- a/.github/workflows/scripts/build.sh +++ b/.github/workflows/scripts/build.sh @@ -14,7 +14,7 @@ $python_executable -m pip install flash_attn triton # Limit the number of parallel jobs to avoid OOM export MAX_JOBS=1 # Make sure release wheels are built for the following architectures -export TORCH_CUDA_ARCH_LIST="7.0 7.5 8.0 8.6 8.9 9.0+PTX" +export TORCH_CUDA_ARCH_LIST="7.0 7.5 8.0 8.6 8.9 9.0 10.0 12.0+PTX" # Build if [ "$3" = sdist ]; then diff --git a/.github/workflows/scripts/cuda-install.sh b/.github/workflows/scripts/cuda-install.sh index 312c6e8..325de75 100644 --- a/.github/workflows/scripts/cuda-install.sh +++ b/.github/workflows/scripts/cuda-install.sh @@ -5,8 +5,37 @@ cuda_version=$(echo $1 | tr "." "-") # Removes '-' and '.' ex: ubuntu-20.04 -> ubuntu2004 OS=$(echo $2 | tr -d ".\-") +ARCH=$(uname -m) +ARCH_TYPE=$ARCH + +# Detectar si es Tegra +if [[ "$ARCH" == "aarch64" ]]; then + if uname -a | grep -qi tegra; then + ARCH_TYPE="tegra-aarch64" + fi +fi + +echo "Detected architecture: ${ARCH_TYPE}" + # Installs CUDA -wget -nv https://developer.download.nvidia.com/compute/cuda/repos/${OS}/x86_64/cuda-keyring_1.1-1_all.deb +if [[ "$ARCH_TYPE" == "tegra-aarch64" ]]; then + # Jetson (Tegra) + wget -nv \ + https://developer.download.nvidia.com/compute/cuda/repos/${OS}/arm64/cuda-${DISTRO}.pin \ + -O /etc/apt/preferences.d/cuda-repository-pin-600 + +elif [[ "$ARCH_TYPE" == "tegra-aarch64" ]]; then + # Jetson (Tegra) + wget -nv \ + https://developer.download.nvidia.com/compute/cuda/repos/${OS}/arm64/cuda-${DISTRO}.pin \ + -O /etc/apt/preferences.d/cuda-repository-pin-600 +else + # ARM64 SBSA (Grace) + wget -nv \ + https://developer.download.nvidia.com/compute/cuda/repos/${OS}/sbsa/cuda-${DISTRO}.pin \ + -O /etc/apt/preferences.d/cuda-repository-pin-600 +fi + sudo dpkg -i cuda-keyring_1.1-1_all.deb rm cuda-keyring_1.1-1_all.deb sudo apt -qq update diff --git a/.github/workflows/scripts/pytorch-install.sh b/.github/workflows/scripts/pytorch-install.sh index 2c24ca8..0178a14 100644 --- a/.github/workflows/scripts/pytorch-install.sh +++ b/.github/workflows/scripts/pytorch-install.sh @@ -20,8 +20,8 @@ pip install typing-extensions==4.12.2 echo $MATRIX_CUDA_VERSION echo $MATRIX_TORCH_VERSION export TORCH_CUDA_VERSION=$(python -c "from os import environ as env; \ -minv = {'2.2': 118, '2.3': 118, '2.4': 118, '2.5': 118, '2.6': 118}[env['MATRIX_TORCH_VERSION']]; \ -maxv = {'2.2': 121, '2.3': 121, '2.4': 124, '2.5': 124, '2.6': 124}[env['MATRIX_TORCH_VERSION']]; \ +minv = {'2.4': 118, '2.5': 118, '2.6': 118, '2.7': 118}[env['MATRIX_TORCH_VERSION']]; \ +maxv = {'2.4': 124, '2.5': 124, '2.6': 126, '2.7': 128}[env['MATRIX_TORCH_VERSION']]; \ print(max(min(int(env['MATRIX_CUDA_VERSION']), maxv), minv))" \ ) if [[ ${pytorch_version} == *"dev"* ]]; then diff --git a/.github/workflows/unittest.yml b/.github/workflows/unittest.yml index 7cc737f..2cee0e2 100644 --- a/.github/workflows/unittest.yml +++ b/.github/workflows/unittest.yml @@ -20,8 +20,8 @@ jobs: strategy: fail-fast: false matrix: - os: [ubuntu-latest, macos-latest, windows-2019] - python-version: ["3.9", "3.10", "3.11"] + os: [ubuntu-latest, ubuntu-24.04-arm, macos-latest, windows-latest, windows-11-arm] + python-version: ["3.9", "3.10", "3.11", "3.12", "3.13"] exclude: - os: macos-latest python-version: '3.9' From 3dde686fc84efeb117a87fde6e3c338ad1cc0b5e Mon Sep 17 00:00:00 2001 From: johnnynunez Date: Fri, 25 Apr 2025 18:55:56 +0200 Subject: [PATCH 3/5] testing --- .github/workflows/release.yml | 6 ------ 1 file changed, 6 deletions(-) diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index ddf6a41..cfdf534 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -3,12 +3,6 @@ # For more information see: https://help.github.com/en/actions/language-and-framework-guides/using-python-with-github-actions#publishing-to-package-registries name: Release - -on: - push: - tags: - - v* - # Needed to create release and upload assets permissions: contents: write From a2f411db87fc531cba03cc4bc3226833a108731c Mon Sep 17 00:00:00 2001 From: johnnynunez Date: Fri, 25 Apr 2025 18:57:03 +0200 Subject: [PATCH 4/5] testing --- .github/workflows/release.yml | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index cfdf534..22091eb 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -3,6 +3,15 @@ # For more information see: https://help.github.com/en/actions/language-and-framework-guides/using-python-with-github-actions#publishing-to-package-registries name: Release + +on: + push: + branches: + - main + pull_request: + branches: + - main + # Needed to create release and upload assets permissions: contents: write From 773213a26151525ae95c1d78911d29e3d5cec701 Mon Sep 17 00:00:00 2001 From: Johnny Date: Fri, 25 Apr 2025 19:20:55 +0200 Subject: [PATCH 5/5] Update release.yml --- .github/workflows/release.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 22091eb..1a572c9 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -54,7 +54,7 @@ jobs: # manylinux docker image, but I haven't figured out how to install CUDA on manylinux. os: [ ubuntu-22.04, ubuntu-22.04-arm ] python-version: [ '3.9', '3.10', '3.11', '3.12', '3.13' ] - torch-version: [ '2.4.0', '2.5.1', '2.6.0', '2.7.0' ] + pytorch-version: [ '2.4.0', '2.5.1', '2.6.0', '2.7.0' ] cuda-version: [ '12.4.1' ] # We need separate wheels that either uses C++11 ABI (-D_GLIBCXX_USE_CXX11_ABI) or not. # Pytorch wheels currently don't use it, but nvcr images have Pytorch compiled with C++11 ABI. @@ -64,7 +64,7 @@ jobs: exclude: # see https://github.com/pytorch/pytorch/blob/main/RELEASE.md#release-compatibility-matrix # Pytorch < 2.5 does not support Python 3.13 - - torch-version: '2.4.0' + - pytorch-version: '2.4.0' python-version: '3.13' steps: