From aec5fb5de2380a9084f4bb430eec4e576066a17c Mon Sep 17 00:00:00 2001
From: Andrey Talman <atalman@fb.com>
Date: Tue, 15 Oct 2024 15:33:00 -0400
Subject: [PATCH] [Release 2.5] Inject triton 3.1.0 into staged for pypi
 dependency for linux aarch64 builds (#2011)

---
 release/pypi/prep_binary_for_pypi.sh   | 9 ++++++++-
 release/pypi/upload_pypi_to_staging.sh | 2 +-
 2 files changed, 9 insertions(+), 2 deletions(-)

diff --git a/release/pypi/prep_binary_for_pypi.sh b/release/pypi/prep_binary_for_pypi.sh
index 154b22852..c1be4a862 100755
--- a/release/pypi/prep_binary_for_pypi.sh
+++ b/release/pypi/prep_binary_for_pypi.sh
@@ -55,7 +55,14 @@ for whl_file in "$@"; do
         find "${dist_info_folder}" -type f -exec sed -i "s!${version_with_suffix}!${version_no_suffix}!" {} \;
         # Moves distinfo from one with a version suffix to one without
         # Example: torch-1.8.0+cpu.dist-info => torch-1.8.0.dist-info
-        mv "${dist_info_folder}" "${dirname_dist_info_folder}/${basename_dist_info_folder/${version_with_suffix}/${version_no_suffix}}"
+        if [[ -n "${VERSION_SUFFIX}" ]]; then
+            mv "${dist_info_folder}" "${dirname_dist_info_folder}/${basename_dist_info_folder/${version_with_suffix}/${version_no_suffix}}"
+        fi
+        if [[ "${PLATFORM}" == "manylinux2014_aarch64" && ${PACKAGE_NAME} == "torch" ]]; then
+            echo "Injecting triton"
+            sed -i 's,nvidia-nvtx-cu12.*$,&\nRequires-Dist: triton==3.1.0 ; platform_system == "Linux" and platform_machine == "x86_64" and python_version < "3.13"\r,g' "${dist_info_folder}/METADATA"
+        fi
+
         cd "${whl_dir}"
 
         (
diff --git a/release/pypi/upload_pypi_to_staging.sh b/release/pypi/upload_pypi_to_staging.sh
index e481f9f9c..8dc33692b 100644
--- a/release/pypi/upload_pypi_to_staging.sh
+++ b/release/pypi/upload_pypi_to_staging.sh
@@ -58,7 +58,7 @@ for pkg in ${pkgs_to_promote}; do
         curl -fSL -o "${orig_pkg}" "https://download.pytorch.org${pkg}"
     )
 
-    if [[ -n "${VERSION_SUFFIX}" ]]; then
+    if [[ -n "${VERSION_SUFFIX}" ]] || [[ "${PLATFORM}" == "manylinux2014_aarch64" && "${PACKAGE_NAME}" == "torch" ]]; then
         OUTPUT_DIR="${output_tmp_dir}" bash "${DIR}/prep_binary_for_pypi.sh" "${orig_pkg}"
     else
         mv "${orig_pkg}" "${output_tmp_dir}/"