[CI] Add unit test workflows for Windows (#804)

Co-authored-by: Yohann Benchetrit <[email protected]>

.circleci/config.yml

@@ -913,6 +913,88 @@ jobs:
       - store_test_results:
           path: test-results
 
+  unittest_windows_optdepts_cpu:
+    <<: *binary_common
+    executor:
+      name: windows-cpu
+    steps:
+      - checkout
+      - designate_upload_channel
+      - run:
+          name: Generate cache key
+          # This will refresh cache on Sundays, nightly build should generate new cache.
+          command: echo "$(date +"%Y-%U")" > .circleci-weekly
+      - restore_cache:
+          keys:
+            - env-v2-windows-{{ arch }}-py<< parameters.python_version >>-{{ checksum ".circleci/unittest/windows_optdepts/scripts/environment.yml" }}-{{ checksum ".circleci-weekly" }}
+
+      - run:
+          name: Setup
+          command: .circleci/unittest/windows_optdepts/scripts/setup_env.sh
+      - save_cache:
+          key: env-v2-windows-{{ arch }}-py<< parameters.python_version >>-{{ checksum ".circleci/unittest/windows_optdepts/scripts/environment.yml" }}-{{ checksum ".circleci-weekly" }}
+
+          paths:
+            - conda
+            - env
+      - run:
+          name: Install torchrl
+          command: .circleci/unittest/windows_optdepts/scripts/install.sh
+      - run:
+          name: Run tests
+          command: .circleci/unittest/windows_optdepts/scripts/run_test.sh
+      - run:
+          name: Post process
+          command: .circleci/unittest/windows_optdepts/scripts/post_process.sh
+      - store_test_results:
+          path: test-results
+
+  unittest_windows_optdepts_gpu:
+    <<: *binary_common
+    executor:
+      name: windows-gpu
+    environment:
+      CUDA_VERSION: "11.6"
+      PYTHON_VERSION: << parameters.python_version >>
+    steps:
+      - checkout
+      - designate_upload_channel
+      - run:
+          name: Generate cache key
+          # This will refresh cache on Sundays, nightly build should generate new cache.
+          command: echo "$(date +"%Y-%U")" > .circleci-weekly
+      - restore_cache:
+          keys:
+            - env-v1-windows-{{ arch }}-py<< parameters.python_version >>-{{ checksum ".circleci/unittest/windows_optdepts/scripts/environment.yml" }}-{{ checksum ".circleci-weekly" }}
+
+      - run:
+          name: Setup
+          command: .circleci/unittest/windows_optdepts/scripts/setup_env.sh
+      - save_cache:
+          key: env-v1-windows-{{ arch }}-py<< parameters.python_version >>-{{ checksum ".circleci/unittest/windows_optdepts/scripts/environment.yml" }}-{{ checksum ".circleci-weekly" }}
+
+          paths:
+            - conda
+            - env
+      - run:
+          name: Install CUDA
+          command: packaging/windows/internal/cuda_install.bat
+      - run:
+          name: Update CUDA driver
+          command: packaging/windows/internal/driver_update.bat
+      - run:
+          name: Install torchrl
+          command: .circleci/unittest/windows_optdepts/scripts/install.sh
+      - run:
+          name: Run tests
+          command: .circleci/unittest/windows_optdepts/scripts/run_test.sh
+      - run:
+          name: Post process
+          command: .circleci/unittest/windows_optdepts/scripts/post_process.sh
+      - store_test_results:
+          path: test-results
+
+
 workflows:
   lint:
     jobs:
@@ -1173,3 +1255,36 @@ workflows:
           cu_version: cu113
           name: unittest_linux_examples_gpu_py3.9
           python_version: '3.9'
+
+      - unittest_windows_optdepts_cpu:
+          cu_version: cpu
+          name: unittest_windows_optdepts_cpu_py3.7
+          python_version: '3.7'
+      - unittest_windows_optdepts_cpu:
+          cu_version: cpu
+          name: unittest_windows_optdepts_cpu_py3.8
+          python_version: '3.8'
+      - unittest_windows_optdepts_cpu:
+          cu_version: cpu
+          name: unittest_windows_optdepts_cpu_py3.9
+          python_version: '3.9'
+      - unittest_windows_optdepts_cpu:
+          cu_version: cpu
+          name: unittest_windows_optdepts_cpu_py3.10
+          python_version: '3.10'
+      - unittest_windows_optdepts_gpu:
+          cu_version: cu116
+          name: unittest_windows_optdepts_gpu_py3.7
+          python_version: '3.7'
+      - unittest_windows_optdepts_gpu:
+          cu_version: cu116
+          name: unittest_windows_optdepts_gpu_py3.8
+          python_version: '3.8'
+      - unittest_windows_optdepts_gpu:
+          cu_version: cu116
+          name: unittest_windows_optdepts_gpu_py3.9
+          python_version: '3.9'
+      - unittest_windows_optdepts_gpu:
+          cu_version: cu116
+          name: unittest_windows_optdepts_gpu_py3.10
+          python_version: '3.10'

.circleci/unittest/windows_optdepts/scripts/environment.yml

@@ -0,0 +1,17 @@
+channels:
+  - pytorch
+  - defaults
+dependencies:
+  - pip
+  - pip:
+    - hypothesis
+    - future
+    - cloudpickle
+    - pytest
+    - pytest-cov
+    - pytest-mock
+    - pytest-instafail
+    - expecttest
+    - pyyaml
+    - scipy
+    - coverage

.circleci/unittest/windows_optdepts/scripts/install.sh

@@ -0,0 +1,61 @@
+#!/usr/bin/env bash
+
+unset PYTORCH_VERSION
+# For unittest, nightly PyTorch is used as the following section,
+# so no need to set PYTORCH_VERSION.
+# In fact, keeping PYTORCH_VERSION forces us to hardcode PyTorch version in config.
+
+set -ex
+
+this_dir="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )"
+
+eval "$(./conda/Scripts/conda.exe 'shell.bash' 'hook')"
+conda activate ./env
+
+# TODO, refactor the below logic to make it easy to understand how to get correct cuda_version.
+if [ "${CU_VERSION:-}" == cpu ] ; then
+    cudatoolkit="cpuonly"
+    version="cpu"
+else
+    if [[ ${#CU_VERSION} -eq 4 ]]; then
+        CUDA_VERSION="${CU_VERSION:2:1}.${CU_VERSION:3:1}"
+    elif [[ ${#CU_VERSION} -eq 5 ]]; then
+        CUDA_VERSION="${CU_VERSION:2:2}.${CU_VERSION:4:1}"
+    fi
+
+    cuda_toolkit_pckg="cudatoolkit"
+    if [[ $CUDA_VERSION == 11.6 || $CUDA_VERSION == 11.7 ]]; then
+        cuda_toolkit_pckg="pytorch-cuda"
+    fi
+
+    echo "Using CUDA $CUDA_VERSION as determined by CU_VERSION"
+    version="$(python -c "print('.'.join(\"${CUDA_VERSION}\".split('.')[:2]))")"
+    cudatoolkit="${cuda_toolkit_pckg}=${version}"
+fi
+
+
+# submodules
+git submodule sync && git submodule update --init --recursive
+
+printf "Installing PyTorch with %s\n" "${cudatoolkit}"
+conda install -y -c "pytorch-${UPLOAD_CHANNEL}" -c nvidia "pytorch-${UPLOAD_CHANNEL}"::pytorch[build="*${version}*"] "${cudatoolkit}"
+
+torch_cuda=$(python -c "import torch; print(torch.cuda.is_available())")
+echo torch.cuda.is_available is $torch_cuda
+
+if [ ! -z "${CUDA_VERSION:-}" ] ; then
+    if [ "$torch_cuda" == "False" ]; then
+        echo "torch with cuda installed but torch.cuda.is_available() is False"
+        exit 1
+    fi
+fi
+
+pip install pip --upgrade
+
+# install tensordict
+pip3 install git+https://github.com/pytorch-labs/tensordict
+
+source "$this_dir/set_cuda_envs.sh"
+
+printf "* Installing torchrl\n"
+pip3 install -e .

.circleci/unittest/windows_optdepts/scripts/install_conda.bat

@@ -0,0 +1 @@
+start /wait "" "%miniconda_exe%" /S /InstallationType=JustMe /RegisterPython=0 /AddToPath=0 /D=%tmp_conda%

.circleci/unittest/windows_optdepts/scripts/post_process.sh

@@ -0,0 +1,6 @@
+#!/usr/bin/env bash
+
+set -e
+
+eval "$(./conda/Scripts/conda.exe 'shell.bash' 'hook')"
+conda activate ./env

.circleci/unittest/windows_optdepts/scripts/run_test.sh

@@ -0,0 +1,12 @@
+#!/usr/bin/env bash
+
+set -e
+
+eval "$(./conda/Scripts/conda.exe 'shell.bash' 'hook')"
+conda activate ./env
+
+this_dir="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )"
+source "$this_dir/set_cuda_envs.sh"
+
+python -m torch.utils.collect_env
+pytest --junitxml=test-results/junit.xml -v --durations 20

.circleci/unittest/windows_optdepts/scripts/set_cuda_envs.sh

@@ -0,0 +1,48 @@
+#!/usr/bin/env bash
+set -ex
+
+echo CU_VERSION is "${CU_VERSION}"
+echo CUDA_VERSION is "${CUDA_VERSION}"
+
+# Currenly, CU_VERSION and CUDA_VERSION are not consistent.
+# to understand this code, see https://github.com/pytorch/vision/issues/4443
+version="cpu"
+if [[ ! -z "${CUDA_VERSION}" ]] ; then
+    version="$CUDA_VERSION"
+else
+    if [[ ${#CU_VERSION} -eq 5 ]]; then
+        version="${CU_VERSION:2:2}.${CU_VERSION:4:1}"
+    fi
+fi
+
+# Don't use if [[ "$version" == "cpu" ]]; then exit 0 fi.
+# It would exit the shell. One result is cpu tests would not run if the shell exit.
+# Unless there's an error, Don't exit.
+if [[ "$version" != "cpu" ]]; then
+    # set cuda envs
+    export PATH="/c/Program Files/NVIDIA GPU Computing Toolkit/CUDA/v${version}/bin:/c/Program Files/NVIDIA GPU Computing Toolkit/CUDA/v${version}/libnvvp:$PATH"
+    export CUDA_PATH_V${version/./_}="C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v${version}"
+    export CUDA_PATH="C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v${version}"
+
+    if  [ ! -d "$CUDA_PATH" ]; then
+        echo "$CUDA_PATH" does not exist
+        exit 1
+    fi
+
+    if [ ! -f "${CUDA_PATH}\include\nvjpeg.h" ]; then
+        echo "nvjpeg does not exist"
+        exit 1
+    fi
+
+    # check cuda driver version
+    for path in '/c/Program Files/NVIDIA Corporation/NVSMI/nvidia-smi.exe' /c/Windows/System32/nvidia-smi.exe; do
+        if [[ -x "$path" ]]; then
+            "$path" || echo "true";
+            break
+        fi
+    done
+
+    which nvcc
+    nvcc --version
+    env | grep CUDA
+fi

.circleci/unittest/windows_optdepts/scripts/setup_env.sh

@@ -0,0 +1,39 @@
+#!/usr/bin/env bash
+
+# This script is for setting up environment in which unit test is ran.
+# To speed up the CI time, the resulting environment is cached.
+#
+# Do not install PyTorch and torchvision here, otherwise they also get cached.
+
+set -e
+
+this_dir="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )"
+root_dir="$(git rev-parse --show-toplevel)"
+conda_dir="${root_dir}/conda"
+env_dir="${root_dir}/env"
+
+cd "${root_dir}"
+
+# 1. Install conda at ./conda
+if [ ! -d "${conda_dir}" ]; then
+    printf "* Installing conda\n"
+    export tmp_conda="$(echo $conda_dir | tr '/' '\\')"
+    export miniconda_exe="$(echo $root_dir | tr '/' '\\')\\miniconda.exe"
+    curl --output miniconda.exe https://repo.anaconda.com/miniconda/Miniconda3-latest-Windows-x86_64.exe -O
+    "$this_dir/install_conda.bat"
+    unset tmp_conda
+    unset miniconda_exe
+fi
+
+eval "$(${conda_dir}/Scripts/conda.exe 'shell.bash' 'hook')"
+
+# 2. Create test environment at ./env
+if [ ! -d "${env_dir}" ]; then
+    printf "* Creating a test environment\n"
+    conda create --prefix "${env_dir}" -y python="$PYTHON_VERSION"
+fi
+conda activate "${env_dir}"
+
+# 3. Install Conda dependencies
+printf "* Installing dependencies (except PyTorch)\n"
+conda env update --file "${this_dir}/environment.yml" --prune

.circleci/unittest/windows_optdepts/scripts/vc_env_helper.bat

@@ -0,0 +1,39 @@
+@echo on
+
+set VC_VERSION_LOWER=16
+set VC_VERSION_UPPER=17
+
+for /f "usebackq tokens=*" %%i in (`"%ProgramFiles(x86)%\Microsoft Visual Studio\Installer\vswhere.exe" -legacy -products * -version [%VC_VERSION_LOWER%^,%VC_VERSION_UPPER%^) -property installationPath`) do (
+    if exist "%%i" if exist "%%i\VC\Auxiliary\Build\vcvarsall.bat" (
+        set "VS15INSTALLDIR=%%i"
+        set "VS15VCVARSALL=%%i\VC\Auxiliary\Build\vcvarsall.bat"
+        goto vswhere
+    )
+)
+
+:vswhere
+if "%VSDEVCMD_ARGS%" == "" (
+    call "%VS15VCVARSALL%" x64 || exit /b 1
+) else (
+    call "%VS15VCVARSALL%" x64 %VSDEVCMD_ARGS% || exit /b 1
+)
+
+@echo on
+
+set DISTUTILS_USE_SDK=1
+
+set args=%1
+shift
+:start
+if [%1] == [] goto done
+set args=%args% %1
+shift
+goto start
+
+:done
+if "%args%" == "" (
+    echo Usage: vc_env_helper.bat [command] [args]
+    echo e.g. vc_env_helper.bat cl /c test.cpp
+)
+
+%args% || exit /b 1

packaging/windows/internal/driver_update.bat

@@ -0,0 +1,25 @@
+set "DRIVER_DOWNLOAD_LINK=https://ossci-windows.s3.amazonaws.com/461.09-data-center-tesla-desktop-winserver-2019-2016-international.exe"
+curl --retry 3 -kL %DRIVER_DOWNLOAD_LINK% --output 461.09-data-center-tesla-desktop-winserver-2019-2016-international.exe
+if errorlevel 1 exit /b 1
+
+start /wait 461.09-data-center-tesla-desktop-winserver-2019-2016-international.exe -s -noreboot
+if errorlevel 1 exit /b 1
+
+del 461.09-data-center-tesla-desktop-winserver-2019-2016-international.exe || ver > NUL
+
+setlocal EnableDelayedExpansion
+set NVIDIA_GPU_EXISTS=0
+for /F "delims=" %%i in ('wmic path win32_VideoController get name') do (
+    set GPUS=%%i
+    if not "x!GPUS:NVIDIA=!" == "x!GPUS!" (
+        SET NVIDIA_GPU_EXISTS=1
+        goto gpu_check_end
+    )
+)
+:gpu_check_end
+endlocal & set NVIDIA_GPU_EXISTS=%NVIDIA_GPU_EXISTS%
+
+if "%NVIDIA_GPU_EXISTS%" == "0" (
+    echo "CUDA Driver installation Failed"
+    exit /b 1
+)