Fix for the cuda 11.6 and usage of cudatoolkit (#2501)

Summary:
Fix for the cuda 11.6 and usage of cudatoolkit

Pull Request resolved: https://github.com/pytorch/audio/pull/2501

Reviewed By: mthrok

Differential Revision: D37388598

Pulled By: atalman

fbshipit-source-id: 41add7ad6fbb3d156cc1270625dc085c62f7a531

.circleci/config.yml

@@ -476,7 +476,11 @@ jobs:
           command: |
             set -x
             source /usr/local/etc/profile.d/conda.sh && conda activate python${PYTHON_VERSION}
-            conda install -v -y -c pytorch-${UPLOAD_CHANNEL} pytorch cudatoolkit=${CU_VERSION:2:2}.${CU_VERSION:4} -c conda-forge
+            if [[ "$CU_VERSION" == cu116 ]]; then
+              conda install -v -y -c pytorch-${UPLOAD_CHANNEL} -c nvidia/label/cuda-11.6.2 pytorch cuda=${CU_VERSION:2:2}.${CU_VERSION:4}
+            else
+              conda install -v -y -c pytorch-${UPLOAD_CHANNEL} pytorch cudatoolkit=${CU_VERSION:2:2}.${CU_VERSION:4} -c conda-forge
+            fi
             conda install -v -y -c file://$HOME/workspace/conda-bld torchaudio
       - checkout
       - run:
@@ -559,7 +563,12 @@ jobs:
             conda create -yn python${PYTHON_VERSION} python=${PYTHON_VERSION}
             conda activate python${PYTHON_VERSION}
             # Include numpy and cudatoolkit in the install conda-forge chanell is used for cudatoolkit
-            conda install -v -y -c pytorch-${UPLOAD_CHANNEL} pytorch numpy cudatoolkit=${CU_VERSION:2:2}.${CU_VERSION:4} -c conda-forge
+
+            if [[ "$CU_VERSION" == cu116 ]]; then
+              conda install -v -y -c pytorch-${UPLOAD_CHANNEL} -c nvidia/label/cuda-11.6.2 pytorch numpy ffmpeg cuda=${CU_VERSION:2:2}.${CU_VERSION:4}
+            else
+              conda install -v -y -c pytorch-${UPLOAD_CHANNEL} pytorch numpy ffmpeg cudatoolkit=${CU_VERSION:2:2}.${CU_VERSION:4} -c conda-forge
+            fi
             # Install from torchaudio file
             conda install -v -y $(ls ~/workspace/conda-bld/win-64/torchaudio*.tar.bz2)
       - run:
@@ -569,7 +578,7 @@ jobs:
             conda activate python${PYTHON_VERSION}
             # Install sound backend
             pip install PySoundFile
-            conda install 'ffmpeg<5'
+            # conda install 'ffmpeg<5'
             ./test/smoke_test/run_smoke_test.sh
 
   smoke_test_windows_pip:

.circleci/config.yml.in

@@ -476,7 +476,11 @@ jobs:
           command: |
             set -x
             source /usr/local/etc/profile.d/conda.sh && conda activate python${PYTHON_VERSION}
-            conda install -v -y -c pytorch-${UPLOAD_CHANNEL} pytorch cudatoolkit=${CU_VERSION:2:2}.${CU_VERSION:4} -c conda-forge
+            if [[ "$CU_VERSION" == cu116 ]]; then
+              conda install -v -y -c pytorch-${UPLOAD_CHANNEL} -c nvidia/label/cuda-11.6.2 pytorch cuda=${CU_VERSION:2:2}.${CU_VERSION:4}
+            else
+              conda install -v -y -c pytorch-${UPLOAD_CHANNEL} pytorch cudatoolkit=${CU_VERSION:2:2}.${CU_VERSION:4} -c conda-forge
+            fi
             conda install -v -y -c file://$HOME/workspace/conda-bld torchaudio
       - checkout
       - run:
@@ -559,7 +563,12 @@ jobs:
             conda create -yn python${PYTHON_VERSION} python=${PYTHON_VERSION}
             conda activate python${PYTHON_VERSION}
             # Include numpy and cudatoolkit in the install conda-forge chanell is used for cudatoolkit
-            conda install -v -y -c pytorch-${UPLOAD_CHANNEL} pytorch numpy cudatoolkit=${CU_VERSION:2:2}.${CU_VERSION:4} -c conda-forge
+
+            if [[ "$CU_VERSION" == cu116 ]]; then
+              conda install -v -y -c pytorch-${UPLOAD_CHANNEL} -c nvidia/label/cuda-11.6.2 pytorch numpy ffmpeg cuda=${CU_VERSION:2:2}.${CU_VERSION:4}
+            else
+              conda install -v -y -c pytorch-${UPLOAD_CHANNEL} pytorch numpy ffmpeg cudatoolkit=${CU_VERSION:2:2}.${CU_VERSION:4} -c conda-forge
+            fi
             # Install from torchaudio file
             conda install -v -y $(ls ~/workspace/conda-bld/win-64/torchaudio*.tar.bz2)
       - run:
@@ -569,7 +578,7 @@ jobs:
             conda activate python${PYTHON_VERSION}
             # Install sound backend
             pip install PySoundFile
-            conda install 'ffmpeg<5'
+            # conda install 'ffmpeg<5'
             ./test/smoke_test/run_smoke_test.sh
 
   smoke_test_windows_pip:

.circleci/unittest/linux/scripts/install.sh

@@ -32,7 +32,14 @@ if [ -z "${CUDA_VERSION:-}" ] ; then
     version="cpu"
 else
     version="$(python -c "print('.'.join(\"${CUDA_VERSION}\".split('.')[:2]))")"
-    cudatoolkit="cudatoolkit=${version}"
+
+    export CUDATOOLKIT_CHANNEL="nvidia"
+    cuda_toolkit_pckg="cudatoolkit"
+    if [[ "$CU_VERSION" == cu116 ]]; then
+        export CUDATOOLKIT_CHANNEL="nvidia/label/cuda-11.6.2"
+        cuda_toolkit_pckg="cuda"
+    fi
+    cudatoolkit="${cuda_toolkit_pckg}=${version}"
 fi
 
 printf "Installing PyTorch with %s\n" "${cudatoolkit}"
@@ -51,7 +58,12 @@ printf "Installing PyTorch with %s\n" "${cudatoolkit}"
       pytorch_build="pytorch[build="*${version}*"]"
     fi
     set -x
-    conda install ${CONDA_CHANNEL_FLAGS:-} -y -c "pytorch-${UPLOAD_CHANNEL}" $MKL_CONSTRAINT "pytorch-${UPLOAD_CHANNEL}::${pytorch_build}" ${cudatoolkit}
+
+    if [[ -z "$cudatoolkit" ]]; then
+        conda install ${CONDA_CHANNEL_FLAGS:-} -y -c "pytorch-${UPLOAD_CHANNEL}" $MKL_CONSTRAINT "pytorch-${UPLOAD_CHANNEL}::${pytorch_build}"
+    else
+        conda install ${CONDA_CHANNEL_FLAGS:-} -y -c "pytorch-${UPLOAD_CHANNEL}" $MKL_CONSTRAINT "pytorch-${UPLOAD_CHANNEL}::${pytorch_build}" "$CUDATOOLKIT_CHANNEL::${cudatoolkit}"
+    fi
 )
 
 # 2. Install torchaudio

.circleci/unittest/windows/scripts/install.sh

@@ -26,7 +26,13 @@ if [ -z "${CUDA_VERSION:-}" ] ; then
     version="cpu"
 else
     version="$(python -c "print('.'.join(\"${CUDA_VERSION}\".split('.')[:2]))")"
-    cudatoolkit="cudatoolkit=${version}"
+
+    cuda_toolkit_pckg="cudatoolkit"
+    if [[ "$CU_VERSION" == cu116 ]]; then
+        cuda_toolkit_pckg="cuda"
+    fi
+
+    cudatoolkit="${cuda_toolkit_pckg}=${version}"
 fi
 printf "Installing PyTorch with %s\n" "${cudatoolkit}"
 conda install -y -c "pytorch-${UPLOAD_CHANNEL}" -c conda-forge "pytorch-${UPLOAD_CHANNEL}"::pytorch[build="*${version}*"] "${cudatoolkit}" 'mkl=2020.4' pytest

packaging/build_conda.sh

@@ -18,11 +18,16 @@ setup_visual_studio_constraint
 export CUDATOOLKIT_CHANNEL="nvidia"
 # NOTE: This is needed because `cudatoolkit=11.5` has a dependency on conda-forge
 #       See: https://github.com/pytorch/audio/pull/2224#issuecomment-1049185550
+export CUDA116_CUDA_DEPENDENCY=""
 if [[ ${CU_VERSION} = "cu116" ]]; then
-    CONDA_CHANNEL_FLAGS="${CONDA_CHANNEL_FLAGS} -c conda-forge"
+    export CUDATOOLKIT_CHANNEL="nvidia/label/cuda-11.6.2"
+    export CUDA116_CUDA_DEPENDENCY="cuda"
+elif [[ ! -z ${CU_VERSION} ]]; then
+    export CUDA116_CUDA_DEPENDENCY="cudatoolkit"
 fi
+
 # NOTE: There are some dependencies that are not available for macOS on Python 3.10 without conda-forge
 if [[ ${OSTYPE} =~ darwin* ]] && [[ ${PYTHON_VERSION} = "3.10" ]]; then
     CONDA_CHANNEL_FLAGS="${CONDA_CHANNEL_FLAGS} -c conda-forge"
 fi
-conda build -c defaults -c $CUDATOOLKIT_CHANNEL ${CONDA_CHANNEL_FLAGS:-} --no-anaconda-upload --python "$PYTHON_VERSION" packaging/torchaudio
+conda build -c defaults -c $CUDATOOLKIT_CHANNEL ${CONDA_CHANNEL_FLAGS:-}  --no-anaconda-upload --python "$PYTHON_VERSION" $CUDA116_CUDA_DEPENDENCY packaging/torchaudio

packaging/pkg_helpers.bash

@@ -238,7 +238,7 @@ setup_conda_cudatoolkit_constraint() {
   else
     case "$CU_VERSION" in
       cu116)
-        export CONDA_CUDATOOLKIT_CONSTRAINT="- cudatoolkit >=11.6,<11.7 # [not osx]"
+        export CONDA_CUDATOOLKIT_CONSTRAINT=""
         ;;
       cu113)
         export CONDA_CUDATOOLKIT_CONSTRAINT="- cudatoolkit >=11.3,<11.4 # [not osx]"