[release only] Aarch64 builds inject triton 3.1.0

pytorch · Oct 15, 2024 · 3328ca2 · 3328ca2
1 parent fa48aaf
commit 3328ca2
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 2 deletions.
diff --git 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
@@ -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}/"