[CI] Fix windows build (#2760)

.github/scripts/version_script.bat

@@ -1,3 +1,53 @@
 @echo off
 set TORCHRL_BUILD_VERSION=0.7.0
 echo TORCHRL_BUILD_VERSION is set to %TORCHRL_BUILD_VERSION%
+
+@echo on
+
+set VC_VERSION_LOWER=17
+set VC_VERSION_UPPER=18
+if "%VC_YEAR%" == "2019" (
+    set VC_VERSION_LOWER=16
+    set VC_VERSION_UPPER=17
+)
+if "%VC_YEAR%" == "2017" (
+    set VC_VERSION_LOWER=15
+    set VC_VERSION_UPPER=16
+)
+
+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
+
+if "%CU_VERSION%" == "xpu" call "C:\Program Files (x86)\Intel\oneAPI\setvars.bat"
+
+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

.github/unittest/windows_optdepts/scripts/install.sh → .github/unittest/windows_optdepts/scripts/unittest.sh

@@ -1,16 +1,38 @@
 #!/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.
+# 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 -ex
+# =================================== Setup =================================================
 
 this_dir="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )"
+root_dir="$(git rev-parse --show-toplevel)"
+env_dir="${root_dir}/env"
+
+cd "${root_dir}"
+
+eval "$($(which conda) shell.bash hook)" && set -x
+
+# Create test environment at ./env
+printf "* Creating a test environment\n"
+conda create --name ci -y python="$PYTHON_VERSION"
+
+printf "* Activating the environment"
+conda activate ci
+
+printf "Python version"
+echo $(which python)
+echo $(python --version)
+echo $(conda info -e)
+
+
+python -m pip install hypothesis future cloudpickle pytest pytest-cov pytest-mock pytest-instafail pytest-rerunfailures expecttest pyyaml scipy coverage
+
+# =================================== Install =================================================
 
-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
@@ -98,3 +120,16 @@ python -c """
 from torchrl.data import ReplayBuffer
 print('successfully imported torchrl')
 """
+
+# =================================== Run =================================================
+
+source "$this_dir/set_cuda_envs.sh"
+
+# we don't use torchsnapshot
+export CKPT_BACKEND=torch
+export MAX_IDLE_COUNT=60
+export BATCHED_PIPE_TIMEOUT=60
+export LAZY_LEGACY_OP=False
+
+python -m torch.utils.collect_env
+pytest --junitxml=test-results/junit.xml -v --durations 200  --ignore test/test_distributed.py --ignore test/test_rlhf.py

.github/workflows/build-wheels-m1.yml

@@ -34,7 +34,7 @@ jobs:
           - repository: pytorch/rl
             smoke-test-script: test/smoke_test.py
             package-name: torchrl
-    name: pytorch/rl
+    name: ${{ matrix.repository }}
     uses: pytorch/test-infra/.github/workflows/build_wheels_macos.yml@main
     with:
       repository: ${{ matrix.repository }}

.github/workflows/test-windows-optdepts.yml

@@ -21,15 +21,23 @@ permissions:
 
 jobs:
   unittests-cpu:
+    strategy:
+      matrix:
+        python-version:
+          - "3.10"
+        runner: ["windows.4xlarge"]
+        gpu-arch-type: ["cpu"]
+      fail-fast: false
     uses: pytorch/test-infra/.github/workflows/windows_job.yml@main
     with:
-      runner: "windows.4xlarge"
       repository: pytorch/rl
+      runner: ${{ matrix.runner }}
       timeout: 40
+      test-infra-ref: main
       script: |
         set -euxo pipefail
 
-        export PYTHON_VERSION="3.9"
+        export PYTHON_VERSION=${{ matrix.python-version }}
         export CU_VERSION="cpu"
         export torch_cuda="False"
 
@@ -48,17 +56,7 @@ jobs:
         fi
         export TD_GET_DEFAULTS_TO_NONE=1
 
-        ## setup_env.sh
-        ./.github/unittest/windows_optdepts/scripts/setup_env.sh
-
-        ## install.sh
-        ./.github/unittest/windows_optdepts/scripts/install.sh
-
-        ## run_test.sh
-        ./.github/unittest/windows_optdepts/scripts/run_test.sh
-
-        ## post_process.sh
-        ./.github/unittest/windows_optdepts/scripts/post_process.sh
+        ./.github/unittest/windows_optdepts/scripts/unittest.sh
 
 #  unittests-gpu:
 #    uses: pytorch/test-infra/.github/workflows/windows_job.yml@main