From e7390b27c705de25710bc7dec15b4f7ba91ffd7f Mon Sep 17 00:00:00 2001
From: Fangjun Kuang <csukuangfj@gmail.com>
Date: Fri, 18 Oct 2024 14:27:49 +0800
Subject: [PATCH] Publish pre-built wheels for torch 2.5.0 (#1308)

---
 scripts/github_actions/build-ubuntu-cpu.sh      | 12 ++++++++----
 scripts/github_actions/build-ubuntu-cuda.sh     | 13 +++++++++----
 scripts/github_actions/generate_build_matrix.py | 14 +++++++++++++-
 scripts/github_actions/install_torch.sh         | 17 +++++++++++++++++
 4 files changed, 47 insertions(+), 9 deletions(-)

diff --git a/scripts/github_actions/build-ubuntu-cpu.sh b/scripts/github_actions/build-ubuntu-cpu.sh
index 8dcae439b..0808dac1f 100755
--- a/scripts/github_actions/build-ubuntu-cpu.sh
+++ b/scripts/github_actions/build-ubuntu-cpu.sh
@@ -15,15 +15,19 @@ if [ -z $TORCH_VERSION ]; then
   exit 1
 fi
 
-echo "Installing ${PYTHON_VERSION}.3"
+INSTALLED_PYTHON_VERSION=${PYTHON_VERSION}.2
+if [[ $PYTHON_VERSION == "3.13" ]]; then
+  INSTALLED_PYTHON_VERSION=${PYTHON_VERSION}.0
+fi
+echo "Installing $INSTALLED_PYTHON_VERSION"
 
 yum -y install openssl-devel bzip2-devel libffi-devel xz-devel wget redhat-lsb-core
 
 if true; then
   echo "Installing ${PYTHON_VERSION}.2"
-  curl -O https://www.python.org/ftp/python/${PYTHON_VERSION}.2/Python-${PYTHON_VERSION}.2.tgz
-  tar xf Python-${PYTHON_VERSION}.2.tgz
-  pushd Python-${PYTHON_VERSION}.2
+  curl -O https://www.python.org/ftp/python/$INSTALLED_PYTHON_VERSION/Python-$INSTALLED_PYTHON_VERSION.tgz
+  tar xf Python-$INSTALLED_PYTHON_VERSION.tgz
+  pushd Python-$INSTALLED_PYTHON_VERSION
 
   PYTHON_INSTALL_DIR=$PWD/py-${PYTHON_VERSION}
 
diff --git a/scripts/github_actions/build-ubuntu-cuda.sh b/scripts/github_actions/build-ubuntu-cuda.sh
index fbc8d4cf8..3615e20c3 100755
--- a/scripts/github_actions/build-ubuntu-cuda.sh
+++ b/scripts/github_actions/build-ubuntu-cuda.sh
@@ -30,10 +30,15 @@ fi
 yum -y install openssl-devel bzip2-devel libffi-devel xz-devel wget redhat-lsb-core
 
 
-echo "Installing ${PYTHON_VERSION}.2"
-curl -O https://www.python.org/ftp/python/${PYTHON_VERSION}.2/Python-${PYTHON_VERSION}.2.tgz
-tar xf Python-${PYTHON_VERSION}.2.tgz
-pushd Python-${PYTHON_VERSION}.2
+INSTALLED_PYTHON_VERSION=${PYTHON_VERSION}.2
+if [[ $PYTHON_VERSION == "3.13" ]]; then
+  INSTALLED_PYTHON_VERSION=${PYTHON_VERSION}.0
+fi
+echo "Installing $INSTALLED_PYTHON_VERSION"
+
+curl -O https://www.python.org/ftp/python/$INSTALLED_PYTHON_VERSION/Python-$INSTALLED_PYTHON_VERSION.tgz
+tar xf Python-$INSTALLED_PYTHON_VERSION.tgz
+pushd Python-$INSTALLED_PYTHON_VERSION
 
 PYTHON_INSTALL_DIR=$PWD/py-${PYTHON_VERSION}
 
diff --git a/scripts/github_actions/generate_build_matrix.py b/scripts/github_actions/generate_build_matrix.py
index de966938f..09bd4e30f 100755
--- a/scripts/github_actions/generate_build_matrix.py
+++ b/scripts/github_actions/generate_build_matrix.py
@@ -272,13 +272,25 @@ def generate_build_matrix(
                 else ["11.8.0", "12.1.0", "12.4.0"]
             ),
         },
+        "2.5.0": {
+            # Only Linux supports python 3.13
+            "python-version": ["3.9", "3.10", "3.11", "3.12", "3.13"],
+            "cuda": (
+                ["11.8", "12.1", "12.4"]  # default 12.4
+                if not for_windows
+                else ["11.8.0", "12.1.0", "12.4.0"]
+            ),
+        },
         # https://github.com/Jimver/cuda-toolkit/blob/master/src/links/windows-links.ts
     }
     if test_only_latest_torch:
-        latest = "2.4.1"
+        latest = "2.5.0"
         matrix = {latest: matrix[latest]}
 
     if for_windows or for_macos:
+        if "2.5.0" in matrix:
+            matrix["2.5.0"]["python-version"].remove("3.13")
+
         if "1.13.0" in matrix:
             matrix["1.13.0"]["python-version"].remove("3.11")
 
diff --git a/scripts/github_actions/install_torch.sh b/scripts/github_actions/install_torch.sh
index 84e9edd5d..4d06b90cc 100755
--- a/scripts/github_actions/install_torch.sh
+++ b/scripts/github_actions/install_torch.sh
@@ -230,6 +230,23 @@ case ${torch} in
         ;;
     esac
     ;;
+  2.5.*)
+    case ${cuda} in
+      11.8)
+        package="torch==${torch}+cu118"
+        # https://download.pytorch.org/whl/nightly/torch/
+        url=https://download.pytorch.org/whl/torch/
+        ;;
+      12.1)
+        package="torch==${torch}+cu121"
+        url=https://download.pytorch.org/whl/torch/
+        ;;
+      12.4)
+        package="torch==${torch}+cu124"
+        url=https://download.pytorch.org/whl/torch/
+        ;;
+    esac
+    ;;
   *)
     echo "Unsupported PyTorch version: ${torch}"
     exit 1