[Rocm6.4] Update triton requirements and conditionalize python_version < '3.13' only for PyTorch <= 2.5 (#68) AND add in conda to dockerfile for torchvision

common/install_conda_docker.sh

@@ -0,0 +1,20 @@
+#!/bin/bash
+# Script used only in CD pipeline
+set -ex
+
+# Anaconda
+# Latest anaconda is using openssl-3 which is incompatible with all currently published versions of git
+# Which are using openssl-1.1.1, see https://anaconda.org/anaconda/git/files?version=2.40.1 for example
+MINICONDA_URL=https://repo.anaconda.com/miniconda/Miniconda3-py311_23.5.2-0-Linux-x86_64.sh
+wget -q $MINICONDA_URL
+# NB: Manually invoke bash per https://github.com/conda/conda/issues/10431
+bash $(basename "$MINICONDA_URL") -b -p /opt/conda
+rm $(basename "$MINICONDA_URL")
+export PATH=/opt/conda/bin:$PATH
+# See https://github.com/pytorch/builder/issues/1473
+# Pin conda to 23.5.2 as it's the last one compatible with openssl-1.1.1
+conda install -y conda=23.5.2 conda-build anaconda-client git ninja
+# The cmake version here needs to match with the minimum version of cmake
+# supported by PyTorch (3.18). There is only 3.18.2 on anaconda
+/opt/conda/bin/pip3 install cmake==3.18.2
+conda remove -y --force patchelf

manywheel/Dockerfile_2_28

@@ -123,6 +123,10 @@ COPY --from=jni                /usr/local/include/jni.h              /usr/local/
 FROM common as cpu_final
 ARG BASE_CUDA_VERSION=11.8
 ARG DEVTOOLSET_VERSION=11
+# Install Anaconda
+ADD ./common/install_conda_docker.sh install_conda.sh
+RUN bash ./install_conda.sh && rm install_conda.sh
+ENV PATH /opt/conda/bin:$PATH
 # Ensure the expected devtoolset is used
 ENV PATH=/opt/rh/gcc-toolset-${DEVTOOLSET_VERSION}/root/usr/bin:$PATH
 ENV LD_LIBRARY_PATH=/opt/rh/gcc-toolset-${DEVTOOLSET_VERSION}/root/usr/lib64:/opt/rh/gcc-toolset-${DEVTOOLSET_VERSION}/root/usr/lib:$LD_LIBRARY_PATH

manywheel/build_common.sh

@@ -73,6 +73,20 @@ if [[ "$BUILD_LIGHTWEIGHT" == "1" ]]; then
     build_version="${build_version}.lw"
 fi
 
+if [[ -z "$PYTORCH_ROOT" ]]; then
+    echo "Need to set PYTORCH_ROOT env variable"
+    exit 1
+fi
+# Always append the pytorch commit to the build_version by querying Git
+pushd "$PYTORCH_ROOT"
+PYTORCH_COMMIT=$(git rev-parse HEAD)
+popd
+
+# Append the commit as a ".git<shortsha>" suffix
+# This yields versions like: torch-2.5.0+rocm6.2.0.lw.gitabcd1234
+short_commit=$(echo "$PYTORCH_COMMIT" | cut -c1-8)
+build_version="${build_version}.git${short_commit}"
+
 echo "Final build_version: $build_version"
 
 export PYTORCH_BUILD_VERSION=$build_version
@@ -119,10 +133,7 @@ fi
 ########################################################
 # Compile wheels as well as libtorch
 #######################################################
-if [[ -z "$PYTORCH_ROOT" ]]; then
-    echo "Need to set PYTORCH_ROOT env variable"
-    exit 1
-fi
+
 pushd "$PYTORCH_ROOT"
 python setup.py clean
 retry pip install -r requirements.txt

manywheel/build_rocm.sh

@@ -312,18 +312,18 @@ if [ ${PYTORCH_VERSION%%\.*} -ge 2 ]; then
     if [[ $(uname) == "Linux" ]] && [[ "$DESIRED_PYTHON" != "3.12" || $(ver $PYTORCH_VERSION) -ge $(ver 2.4) ]]; then
 	# Triton commit got unified in PyTorch 2.5
 	if [[ $(ver $PYTORCH_VERSION) -ge $(ver 2.5) ]]; then
-            TRITON_SHORTHASH=$(cut -c1-10 $PYTORCH_ROOT/.ci/docker/ci_commit_pins/triton.txt)
+            TRITON_SHORTHASH=$(cut -c1-8 $PYTORCH_ROOT/.ci/docker/ci_commit_pins/triton.txt)
 	else
-            TRITON_SHORTHASH=$(cut -c1-10 $PYTORCH_ROOT/.ci/docker/ci_commit_pins/triton-rocm.txt)
+            TRITON_SHORTHASH=$(cut -c1-8 $PYTORCH_ROOT/.ci/docker/ci_commit_pins/triton-rocm.txt)
 	fi
         TRITON_VERSION=$(cat $PYTORCH_ROOT/.ci/docker/triton_version.txt)
 	# Only linux Python < 3.13 are supported wheels for triton
-	TRITON_CONSTRAINT="platform_system == 'Linux' and platform_machine == 'x86_64' and python_version < '3.13'"
+	TRITON_CONSTRAINT="platform_system == 'Linux' and platform_machine == 'x86_64'$(if [[ $(ver "$PYTORCH_VERSION") -le $(ver "2.5") ]]; then echo " and python_version < '3.13'"; fi)"
 
         if [[ -z "$PYTORCH_EXTRA_INSTALL_REQUIREMENTS" ]]; then
-            export PYTORCH_EXTRA_INSTALL_REQUIREMENTS="pytorch-triton-rocm==${TRITON_VERSION}+${ROCM_VERSION_WITH_PATCH}.${TRITON_SHORTHASH}; ${TRITON_CONSTRAINT}"
+            export PYTORCH_EXTRA_INSTALL_REQUIREMENTS="pytorch-triton-rocm==${TRITON_VERSION}+${ROCM_VERSION_WITH_PATCH}.git${TRITON_SHORTHASH}; ${TRITON_CONSTRAINT}"
         else
-            export PYTORCH_EXTRA_INSTALL_REQUIREMENTS="${PYTORCH_EXTRA_INSTALL_REQUIREMENTS} | pytorch-triton-rocm==${TRITON_VERSION}+${ROCM_VERSION_WITH_PATCH}.${TRITON_SHORTHASH}; ${TRITON_CONSTRAINT}"
+            export PYTORCH_EXTRA_INSTALL_REQUIREMENTS="${PYTORCH_EXTRA_INSTALL_REQUIREMENTS} | pytorch-triton-rocm==${TRITON_VERSION}+${ROCM_VERSION_WITH_PATCH}.git${TRITON_SHORTHASH}; ${TRITON_CONSTRAINT}"
         fi
     fi
 fi