[CI] Fix wheels

ghstack-source-id: 0f26021
Pull Request resolved: #2876

Author: vmoens

.github/scripts/pre-build-script-win.sh

@@ -0,0 +1,7 @@
+#!/bin/bash
+
+pip install --upgrade setuptools
+
+export TORCHRL_BUILD_VERSION=0.7.0
+
+${CONDA_RUN} pip install git+https://github.com/pytorch/tensordict.git -U

.github/scripts/pre-build-script.sh

@@ -0,0 +1,5 @@
+#!/bin/bash
+
+pip install --upgrade setuptools
+
+${CONDA_RUN} pip install git+https://github.com/pytorch/tensordict.git -U

.github/scripts/td_script.sh

@@ -1,5 +1,25 @@
 #!/bin/bash
 
 export TORCHRL_BUILD_VERSION=0.7.0
+pip install --upgrade setuptools
 
-${CONDA_RUN} pip install git+https://github.com/pytorch/tensordict.git -U
+# Check if ARCH is set to aarch64
+ARCH=${ARCH:-}  # This sets ARCH to an empty string if it's not defined
+
+if pip list | grep -q torch; then
+    echo "Torch is installed."
+    if [[ "$ARCH" == "aarch64" ]]; then
+        ${CONDA_RUN} pip install git+https://github.com/pytorch/tensordict.git -U --no-deps
+    else
+        ${CONDA_RUN} pip install tensordict-nightly -U
+    fi
+elif [[ -n "${SMOKE_TEST_SCRIPT:-}" ]]; then
+    ${CONDA_RUN} ${PIP_INSTALL_TORCH}
+    if [[ "$ARCH" == "aarch64" ]]; then
+        ${CONDA_RUN} pip install git+https://github.com/pytorch/tensordict.git -U --no-deps
+    else
+        ${CONDA_RUN} pip install tensordict-nightly -U
+    fi
+else
+    echo "Torch is not installed - tensordict will be installed later."
+fi

.github/scripts/version_script.bat

@@ -36,6 +36,9 @@ if "%CU_VERSION%" == "xpu" call "C:\Program Files (x86)\Intel\oneAPI\setvars.bat
 
 set DISTUTILS_USE_SDK=1
 
+:: Upgrade setuptools before installing PyTorch
+pip install --upgrade setuptools==72.1.0 || exit /b 1
+
 set args=%1
 shift
 :start

.github/workflows/build-wheels-aarch64-linux.yml

@@ -35,6 +35,7 @@ jobs:
           - repository: pytorch/rl
             smoke-test-script: test/smoke_test.py
             package-name: torchrl
+            pre-script: .github/scripts/pre-build-script.sh
     name: pytorch/rl
     uses: pytorch/test-infra/.github/workflows/build_wheels_linux.yml@main
     with:

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

@@ -34,6 +34,7 @@ jobs:
           - repository: pytorch/rl
             smoke-test-script: test/smoke_test.py
             package-name: torchrl
+            pre-script: .github/scripts/pre-build-script.sh
     name: pytorch/rl
     uses: pytorch/test-infra/.github/workflows/build_wheels_linux.yml@main
     with:

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

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

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

@@ -32,7 +32,7 @@ jobs:
       matrix:
         include:
           - repository: pytorch/rl
-            pre-script: .github/scripts/td_script.sh
+            pre-script: .github/scripts/pre-build-script-win.sh
             env-script: .github/scripts/version_script.bat
             post-script: "python packaging/wheel/relocate.py"
             smoke-test-script: test/smoke_test.py