.circleci: Add Python 3.9, CUDA 11.2 to CI (#3341)

(cherry picked from commit 2f40a483d73018ae6e1488a484c5927f2b309969)
Signed-off-by: Eli Uriegas <eliuriegas@fb.com>

.circleci/config.yml.in

@@ -73,6 +73,10 @@ binary_common: &binary_common
       description: "Wheel only: what docker image to use"
       type: string
       default: "pytorch/manylinux-cuda101"
+    conda_docker_image:
+      description: "Conda only: what docker image to use"
+      type: string
+      default: "pytorch/conda-builder:cpu"
   environment:
     PYTHON_VERSION: << parameters.python_version >>
     PYTORCH_VERSION: << parameters.pytorch_version >>
@@ -189,7 +193,7 @@ jobs:
   binary_linux_conda:
     <<: *binary_common
     docker:
-      - image: "pytorch/conda-cuda"
+      - image: "<< parameters.conda_docker_image >>"
     resource_class: 2xlarge+
     steps:
       - checkout_merge
@@ -411,7 +415,11 @@ jobs:
             set -x
             eval "$('/C/tools/miniconda3/Scripts/conda.exe' 'shell.bash' 'hook')"
             conda env remove -n python${PYTHON_VERSION} || true
-            conda create -yn python${PYTHON_VERSION} python=${PYTHON_VERSION}
+            CONDA_CHANNEL_FLAGS=""
+            if [[ "${PYTHON_VERSION}" = 3.9 ]]; then
+              CONDA_CHANNEL_FLAGS="-c=conda-forge"
+            fi
+            conda create ${CONDA_CHANNEL_FLAGS} -yn python${PYTHON_VERSION} python=${PYTHON_VERSION}
             conda activate python${PYTHON_VERSION}
             conda install Pillow
             conda install -v -y -c pytorch-nightly pytorch
@@ -436,7 +444,11 @@ jobs:
           command: |
             set -x
             eval "$('/C/tools/miniconda3/Scripts/conda.exe' 'shell.bash' 'hook')"
-            conda env remove -n python${PYTHON_VERSION} || true
+            CONDA_CHANNEL_FLAGS=""
+            if [[ "${PYTHON_VERSION}" = 3.9 ]]; then
+              CONDA_CHANNEL_FLAGS="-c=conda-forge"
+            fi
+            conda create ${CONDA_CHANNEL_FLAGS} -yn python${PYTHON_VERSION} python=${PYTHON_VERSION}
             conda create -yn python${PYTHON_VERSION} python=${PYTHON_VERSION}
             conda activate python${PYTHON_VERSION}
             pip install $(ls ~/workspace/torchvision*.whl) --pre -f https://download.pytorch.org/whl/nightly/torch_nightly.html
@@ -513,7 +525,7 @@ jobs:
             - env
       - run:
           name: Install torchvision
-          command: docker run -t --gpus all -v $PWD:$PWD -w $PWD -e UPLOAD_CHANNEL "${image_name}" .circleci/unittest/linux/scripts/install.sh
+          command: docker run -t --gpus all -v $PWD:$PWD -w $PWD -e UPLOAD_CHANNEL -e CU_VERSION "${image_name}" .circleci/unittest/linux/scripts/install.sh
       - run:
           name: Run tests
           command: docker run -t --gpus all -v $PWD:$PWD -w $PWD "${image_name}" .circleci/unittest/linux/scripts/run_test.sh

.circleci/regenerate.py

@@ -19,7 +19,8 @@ import yaml
 import os.path
 
 
-PYTHON_VERSIONS = ["3.6", "3.7", "3.8"]
+PYTHON_VERSIONS = ["3.6", "3.7", "3.8", "3.9"]
+CUDA_VERSION = ["10.1", "10.2", "11.2"]
 
 
 def build_workflows(prefix='', filter_branch=None, upload=False, indentation=6, windows_latest_only=False):
@@ -27,8 +28,8 @@ def build_workflows(prefix='', filter_branch=None, upload=False, indentation=6,
     for btype in ["wheel", "conda"]:
         for os_type in ["linux", "macos", "win"]:
             python_versions = PYTHON_VERSIONS
-            cu_versions_dict = {"linux": ["cpu", "cu92", "cu101", "cu102", "cu110"],
-                                "win": ["cpu", "cu101", "cu102", "cu110"],
+            cu_versions_dict = {"linux": ["cpu", "cu101", "cu102", "cu112"],
+                                "win": ["cpu", "cu101", "cu102", "cu112"],
                                 "macos": ["cpu"]}
             cu_versions = cu_versions_dict[os_type]
             for python_version in python_versions:
@@ -99,6 +100,7 @@ manylinux_images = {
     "cu101": "pytorch/manylinux-cuda101",
     "cu102": "pytorch/manylinux-cuda102",
     "cu110": "pytorch/manylinux-cuda110",
+    "cu112": "pytorch/manylinux-cuda112",
 }
 
 
@@ -109,6 +111,14 @@ def get_manylinux_image(cu_version):
     return f"pytorch/manylinux-cuda{cu_suffix}"
 
 
+def get_conda_image(cu_version):
+    if cu_version == "cpu":
+        return "pytorch/conda-builder:cpu"
+    if cu_version.startswith('cu'):
+        cu_suffix = cu_version[len('cu'):]
+    return f"pytorch/conda-builder:cuda{cu_suffix}"
+
+
 def generate_base_workflow(base_workflow_name, python_version, cu_version,
                            unicode, os_type, btype, *, filter_branch=None):
 
@@ -123,6 +133,7 @@ def generate_base_workflow(base_workflow_name, python_version, cu_version,
 
     if os_type != "win":
         d["wheel_docker_image"] = get_manylinux_image(cu_version)
+        d["conda_docker_image"] = get_conda_image(cu_version)
 
     if filter_branch is not None:
         d["filters"] = {

.circleci/unittest/linux/scripts/environment.yml

 channels:
   - pytorch
   - defaults
+  # using conda-forge for python v3.9+
+  - conda-forge
 dependencies:
-  - numpy
   - pytest
   - pytest-cov
   - codecov
   - pip
   - libpng
   - jpeg
-  - ffmpeg=4.2
   - ca-certificates
   - pip:
     - future

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

@@ -22,8 +22,9 @@ else
     version="$(python -c "print('.'.join(\"${CUDA_VERSION}\".split('.')[:2]))")"
     cudatoolkit="cudatoolkit=${version}"
 fi
+
 printf "Installing PyTorch with %s\n" "${cudatoolkit}"
-conda install -y -c "pytorch-${UPLOAD_CHANNEL}" pytorch "${cudatoolkit}"
+conda install -y -c "pytorch-${UPLOAD_CHANNEL}" -c conda-forge pytorch "${cudatoolkit}"
 
 printf "* Installing torchvision\n"
 python setup.py develop

.circleci/unittest/linux/scripts/setup_env.sh

@@ -36,4 +36,13 @@ conda activate "${env_dir}"
 
 # 3. Install Conda dependencies
 printf "* Installing dependencies (except PyTorch)\n"
+NUMPY_MIN_VER="1.11"
+FFMPEG_PIN="=4.2"
+if [[ "${PYTHON_VERSION}" = "3.9" ]]; then
+    NUMPY_MIN_VER="1.20"
+    FFMPEG_PIN=">=4.2"
+fi
+
+conda install -y -c conda-forge "numpy >=${NUMPY_MIN_VER}"
+conda install -y -c pytorch "ffmpeg${FFMPEG_PIN}"
 conda env update --file "${this_dir}/environment.yml" --prune

.circleci/unittest/windows/scripts/environment.yml

 channels:
   - pytorch
   - defaults
+  # use conda-forge for python v3.9+
+  - conda-forge
 dependencies:
-  - numpy
   - pytest
   - pytest-cov
   - codecov

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

@@ -24,8 +24,9 @@ else
     version="$(python -c "print('.'.join(\"${CUDA_VERSION}\".split('.')[:2]))")"
     cudatoolkit="cudatoolkit=${version}"
 fi
+
 printf "Installing PyTorch with %s\n" "${cudatoolkit}"
-conda install -y -c "pytorch-${UPLOAD_CHANNEL}" pytorch "${cudatoolkit}"
+conda install -y -c "pytorch-${UPLOAD_CHANNEL}" -c conda-forge pytorch "${cudatoolkit}"
 
 printf "* Installing torchvision\n"
 "$this_dir/vc_env_helper.bat" python setup.py develop

.circleci/unittest/windows/scripts/setup_env.sh

@@ -36,4 +36,9 @@ conda activate "${env_dir}"
 
 # 3. Install Conda dependencies
 printf "* Installing dependencies (except PyTorch)\n"
+NUMPY_MIN_VER="1.11"
+if [[ "${PYTHON_VERSION}" = "3.9" ]]; then
+    NUMPY_MIN_VER="1.20"
+fi
+conda install -y -c conda-forge "numpy >=${NUMPY_MIN_VER}"
 conda env update --file "${this_dir}/environment.yml" --prune

packaging/pkg_helpers.bash

@@ -49,6 +49,17 @@ setup_cuda() {
 
   # Now work out the CUDA settings
   case "$CU_VERSION" in
+    cu112)
+      if [[ "$OSTYPE" == "msys" ]]; then
+        export CUDA_HOME="C:\\Program Files\\NVIDIA GPU Computing Toolkit\\CUDA\\v11.2"
+      else
+        export CUDA_HOME=/usr/local/cuda-11.2/
+      fi
+      export FORCE_CUDA=1
+      # Hard-coding gencode flags is temporary situation until
+      # https://github.com/pytorch/pytorch/pull/23408 lands
+      export NVCC_FLAGS="-gencode=arch=compute_35,code=sm_35 -gencode=arch=compute_50,code=sm_50 -gencode=arch=compute_60,code=sm_60 -gencode=arch=compute_70,code=sm_70 -gencode=arch=compute_75,code=sm_75 -gencode=arch=compute_80,code=sm_80 -gencode=arch=compute_86,code=sm_86 -gencode=arch=compute_50,code=compute_50"
+      ;;
     cu110)
       if [[ "$OSTYPE" == "msys" ]]; then
         export CUDA_HOME="C:\\Program Files\\NVIDIA GPU Computing Toolkit\\CUDA\\v11.0"
@@ -170,10 +181,13 @@ setup_wheel_python() {
   if [[ "$(uname)" == Darwin || "$OSTYPE" == "msys" ]]; then
     eval "$(conda shell.bash hook)"
     conda env remove -n "env$PYTHON_VERSION" || true
-    conda create -yn "env$PYTHON_VERSION" python="$PYTHON_VERSION"
+    if [[ "$PYTHON_VERSION" == 3.9 ]]; then
+      export CONDA_CHANNEL_FLAGS="${CONDA_CHANNEL_FLAGS} -c=conda-forge"
+    fi
+    conda create ${CONDA_CHANNEL_FLAGS} -yn "env$PYTHON_VERSION" python="$PYTHON_VERSION"
     conda activate "env$PYTHON_VERSION"
     # Install libpng from Anaconda (defaults)
-    conda install libpng jpeg -y
+    conda install ${CONDA_CHANNEL_FLAGS} -c conda-forge libpng "jpeg<=9b" -y
   else
     # Install native CentOS libJPEG, LAME, freetype and GnuTLS
     yum install -y libjpeg-turbo-devel lame freetype gnutls
@@ -189,6 +203,7 @@ setup_wheel_python() {
       3.6) python_abi=cp36-cp36m ;;
       3.7) python_abi=cp37-cp37m ;;
       3.8) python_abi=cp38-cp38 ;;
+      3.9) python_abi=cp39-cp39 ;;
       *)
         echo "Unrecognized PYTHON_VERSION=$PYTHON_VERSION"
         exit 1
@@ -263,6 +278,9 @@ setup_conda_pytorch_constraint() {
   if [[ "$OSTYPE" == msys && "$CU_VERSION" == cu92 ]]; then
     export CONDA_CHANNEL_FLAGS="${CONDA_CHANNEL_FLAGS} -c defaults -c numba/label/dev"
   fi
+  if [[ "$PYTHON_VERSION" == 3.9 ]]; then
+    export CONDA_CHANNEL_FLAGS="${CONDA_CHANNEL_FLAGS} -c=conda-forge"
+  fi
 }
 
 # Translate CUDA_VERSION into CUDA_CUDATOOLKIT_CONSTRAINT
@@ -272,6 +290,9 @@ setup_conda_cudatoolkit_constraint() {
     export CONDA_CUDATOOLKIT_CONSTRAINT=""
   else
     case "$CU_VERSION" in
+      cu112)
+        export CONDA_CUDATOOLKIT_CONSTRAINT="- cudatoolkit >=11.2,<11.3 # [not osx]"
+        ;;
       cu110)
         export CONDA_CUDATOOLKIT_CONSTRAINT="- cudatoolkit >=11.0,<11.1 # [not osx]"
         ;;
@@ -307,6 +328,9 @@ setup_conda_cudatoolkit_plain_constraint() {
     export CMAKE_USE_CUDA=0
   else
     case "$CU_VERSION" in
+      cu112)
+        export CONDA_CUDATOOLKIT_CONSTRAINT="cudatoolkit=11.2"
+        ;;
       cu102)
         export CONDA_CUDATOOLKIT_CONSTRAINT="cudatoolkit=10.2"
         ;;

packaging/torchvision/meta.yaml

@@ -9,12 +9,15 @@ requirements:
   build:
     - {{ compiler('c') }} # [win]
     - libpng
-    - jpeg
-    - ffmpeg =4.2  # [not win]
+    - jpeg <=9b
+    # NOTE: The only ffmpeg version that we build is actually 4.2
+    - ffmpeg >=4.2  # [not win]
 
   host:
     - python
     - setuptools
+    - numpy >=1.20 # [py>=39]
+    - numpy >=1.11 # [py!=39]
     {{ environ.get('CONDA_PYTORCH_BUILD_CONSTRAINT') }}
     {{ environ.get('CONDA_CUDATOOLKIT_CONSTRAINT') }}
     {{ environ.get('CONDA_CPUONLY_FEATURE') }}
@@ -22,10 +25,11 @@ requirements:
   run:
     - python
     - libpng
-    - ffmpeg =4.2  # [not win]
-    - jpeg
+    - ffmpeg >=4.2  # [not win]
+    - jpeg <=9b
     - pillow >=4.1.1
-    - numpy >=1.11
+    - numpy >=1.20 # [py>=39]
+    - numpy >=1.11 # [py!=39]
     {{ environ.get('CONDA_PYTORCH_CONSTRAINT') }}
     {{ environ.get('CONDA_CUDATOOLKIT_CONSTRAINT') }}
 
@@ -50,7 +54,8 @@ test:
   requires:
     - pytest
     - scipy
-    - av =8.0.1
+    - av >=8.0.1
+    - jpeg <=9b
     - ca-certificates
 
 

packaging/windows/internal/cuda_install.bat

@@ -19,6 +19,7 @@ if %CUDA_VER% EQU 100 goto cuda100
 if %CUDA_VER% EQU 101 goto cuda101
 if %CUDA_VER% EQU 102 goto cuda102
 if %CUDA_VER% EQU 110 goto cuda110
+if %CUDA_VER% EQU 112 goto cuda112
 
 echo CUDA %CUDA_VERSION_STR% is not supported
 exit /b 1
@@ -107,6 +108,23 @@ if not exist "%SRC_DIR%\temp_build\cudnn-11.0-windows-x64-v8.0.4.30.zip" (
 
 goto cuda_common
 
+:cuda112
+
+if not exist "%SRC_DIR%\temp_build\cuda_11.2.0_460.89_win10.exe" (
+    curl -k -L https://ossci-windows.s3.amazonaws.com/cuda_11.2.0_460.89_win10.exe --output "%SRC_DIR%\temp_build\cuda_11.2.0_460.89_win10.exe"
+    if errorlevel 1 exit /b 1
+    set "CUDA_SETUP_FILE=%SRC_DIR%\temp_build\cuda_11.2.0_460.89_win10.exe"
+    set "ARGS=nvcc_11.2 cuobjdump_11.2 nvprune_11.2 nvprof_11.2 cupti_11.2 cublas_11.2 cublas_dev_11.2 cudart_11.2 cufft_11.2 cufft_dev_11.2 curand_11.2 curand_dev_11.2 cusolver_11.2 cusolver_dev_11.2 cusparse_11.2 cusparse_dev_11.2 npp_11.2 npp_dev_11.2 nvrtc_11.2 nvrtc_dev_11.2 nvml_dev_11.2"
+)
+
+if not exist "%SRC_DIR%\temp_build\cudnn-11.2-windows-x64-v8.1.0.77.zip" (
+    curl -k -L http://s3.amazonaws.com/ossci-windows/cudnn-11.2-windows-x64-v8.1.0.77.zip --output "%SRC_DIR%\temp_build\cudnn-11.2-windows-x64-v8.1.0.77.zip"
+    if errorlevel 1 exit /b 1
+    set "CUDNN_SETUP_FILE=%SRC_DIR%\temp_build\cudnn-11.2-windows-x64-v8.1.0.77.zip"
+)
+
+goto cuda_common
+
 :cuda_common
 
 if not exist "%SRC_DIR%\temp_build\NvToolsExt.7z" (

test/common_utils.py

@@ -18,6 +18,10 @@ from collections import OrderedDict
 import numpy as np
 from PIL import Image
 
+IS_PY39 = sys.version_info.major == 3 and sys.version_info.minor == 9
+PY39_SEGFAULT_SKIP_MSG = "Segmentation fault with Python 3.9, see https://github.com/pytorch/vision/issues/3367"
+PY39_SKIP = unittest.skipIf(IS_PY39, PY39_SEGFAULT_SKIP_MSG)
+
 
 @contextlib.contextmanager
 def get_tmp_dir(src=None, **kwargs):

test/test_video.py

@@ -4,6 +4,7 @@ import contextlib
 import tempfile
 import unittest
 import random
+import sys
 
 import itertools
 
@@ -13,6 +14,7 @@ import numpy as np
 import torch
 import torchvision
 from torchvision.io import _HAS_VIDEO_OPT, VideoReader
+from common_utils import PY39_SKIP
 
 try:
     import av
@@ -278,6 +280,7 @@ def _template_read_video(video_object, s=0, e=None):
 
 @unittest.skipIf(_HAS_VIDEO_OPT is False, "Didn't compile with ffmpeg")
 class TestVideo(unittest.TestCase):
+    @PY39_SKIP
     @unittest.skipIf(av is None, "PyAV unavailable")
     def test_read_video_tensor(self):
         """
@@ -362,6 +365,7 @@ class TestVideo(unittest.TestCase):
                 config.duration, reader_md["video"]["duration"][0], delta=0.5
             )
 
+    @PY39_SKIP
     @unittest.skipIf(av is None, "PyAV unavailable")
     def test_video_reading_fn(self):
         """

test/test_video_reader.py

@@ -10,6 +10,7 @@ import torch
 import torchvision.io as io
 from numpy.random import randint
 from torchvision.io import _HAS_VIDEO_OPT
+from common_utils import PY39_SKIP
 
 
 try:
@@ -426,6 +427,7 @@ class TestVideoReader(unittest.TestCase):
                     audio_timebase_den,
                 )
 
+    @PY39_SKIP
     def test_read_video_from_file(self):
         """
         Test the case when decoder starts with a video file to decode frames.
@@ -471,6 +473,7 @@ class TestVideoReader(unittest.TestCase):
             # compare decoding results
             self.compare_decoding_result(tv_result, pyav_result, config)
 
+    @PY39_SKIP
     def test_read_video_from_file_read_single_stream_only(self):
         """
         Test the case when decoder starts with a video file to decode frames, and
@@ -779,6 +782,7 @@ class TestVideoReader(unittest.TestCase):
             self.assertEqual(tv_result[0].size(1), height)
             self.assertEqual(tv_result[0].size(2), width)
 
+    @PY39_SKIP
     def test_read_video_from_file_audio_resampling(self):
         """
         Test the case when decoder starts with a video file to decode frames, and
@@ -838,6 +842,7 @@ class TestVideoReader(unittest.TestCase):
                         delta=0.1 * asample_rate.item(),
                     )
 
+    @PY39_SKIP
     def test_compare_read_video_from_memory_and_file(self):
         """
         Test the case when video is already in memory, and decoder reads data in memory
@@ -904,6 +909,7 @@ class TestVideoReader(unittest.TestCase):
             # finally, compare results decoded from memory and file
             self.compare_decoding_result(tv_result_memory, tv_result_file)
 
+    @PY39_SKIP
     def test_read_video_from_memory(self):
         """
         Test the case when video is already in memory, and decoder reads data in memory
@@ -948,6 +954,7 @@ class TestVideoReader(unittest.TestCase):
             self.check_separate_decoding_result(tv_result, config)
             self.compare_decoding_result(tv_result, pyav_result, config)
 
+    @PY39_SKIP
     def test_read_video_from_memory_get_pts_only(self):
         """
         Test the case when video is already in memory, and decoder reads data in memory.
@@ -1017,6 +1024,7 @@ class TestVideoReader(unittest.TestCase):
             self.assertEqual(tv_result_pts_only[5].numel(), 0)
             self.compare_decoding_result(tv_result, tv_result_pts_only)
 
+    @PY39_SKIP
     def test_read_video_in_range_from_memory(self):
         """
         Test the case when video is already in memory, and decoder reads data in memory.
@@ -1192,6 +1200,7 @@ class TestVideoReader(unittest.TestCase):
             probe_result = scripted_fun(video_tensor)
             self.check_meta_result(probe_result, config)
 
+    @PY39_SKIP
     def test_read_video_from_memory_scripted(self):
         """
         Test the case when video is already in memory, and decoder reads data in memory