[CI] Fix benchmark workflows

.github/workflows/benchmarks.yml

@@ -24,27 +24,36 @@ jobs:
     name: CPU Pytest benchmark
     runs-on: ubuntu-20.04
     steps:
-      - uses: actions/checkout@v3
-      - uses: actions/setup-python@v4
+      - name: Who triggered this?
+        run: |
+          echo "Action triggered by ${{ github.event.pull_request.html_url }}"
+      - name: Checkout
+        uses: actions/checkout@v3
+        with:
+          fetch-depth: 50 # this is to make sure we obtain the target base commit
+      - name: Python Setup
+        uses: actions/setup-python@v4
         with:
-          python-version: 3.9
+          python-version: '3.10'
       - name: Setup Environment
         run: |
+          python3.10 -m venv ./py310
+          source ./py310/bin/activate
+
           python3 -m pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/cpu -U 
           python3 -m pip install git+https://github.com/pytorch/tensordict 
           python3 setup.py develop
           python3 -m pip install pytest pytest-benchmark
           python3 -m pip install "gym[accept-rom-license,atari]"
           python3 -m pip install "dm_control" "mujoco"
-          export TD_GET_DEFAULTS_TO_NONE=1
-      - name: Run benchmarks
-        run: |
+
           cd benchmarks/
           export TORCHDYNAMO_INLINE_INBUILT_NN_MODULES=1
-          python -m pytest --benchmark-json output.json
+          export TD_GET_DEFAULTS_TO_NONE=1
+          python3 -m pytest -vvv --rank 0 --benchmark-json output.json --ignore test_collectors_benchmark.py
       - name: Store benchmark results
-        if: ${{ github.ref == 'refs/heads/main' || github.event_name == 'workflow_dispatch' }}
         uses: benchmark-action/github-action-benchmark@v1
+        if: ${{ github.ref == 'refs/heads/main' || github.event_name == 'workflow_dispatch' }}
         with:
           name: CPU Benchmark Results
           tool: 'pytest'
@@ -68,48 +77,73 @@ jobs:
       image: nvidia/cuda:12.3.0-base-ubuntu22.04
       options: --gpus all
     steps:
-      - name: Install deps
+      - name: Set GITHUB_BRANCH environment variable
         run: |
-          export TZ=Europe/London
-          export DEBIAN_FRONTEND=noninteractive  # tzdata bug
-          apt-get update -y
-          apt-get install software-properties-common -y
-          add-apt-repository ppa:git-core/candidate -y
-          apt-get update -y
-          apt-get upgrade -y
-          apt-get -y install libglu1-mesa libgl1-mesa-glx libosmesa6 gcc curl g++ unzip wget libglfw3-dev libgles2-mesa-dev libglew-dev sudo git cmake libz-dev
+          if [ "${{ github.event_name }}" == "push" ]; then
+            export GITHUB_BRANCH=${{ github.event.branch }}
+          elif [ "${{ github.event_name }}" == "pull_request" ]; then
+            export GITHUB_BRANCH=${{ github.event.pull_request.head.ref }}
+          else
+            echo "Unsupported event type"
+            exit 1
+          fi
+          echo "GITHUB_BRANCH=$GITHUB_BRANCH" >> $GITHUB_ENV
+      - name: Who triggered this?
+        run: |
+          echo "Action triggered by ${{ github.event.pull_request.html_url }}"
       - name: Check ldd --version
         run: ldd --version
       - name: Checkout
         uses: actions/checkout@v3
+        with:
+          fetch-depth: 50 # this is to make sure we obtain the target base commit
       - name: Python Setup
         uses: actions/setup-python@v4
         with:
-          python-version: 3.9
+          python-version: '3.10'
+      - name: Setup Environment
+        run: |
+          export TZ=Europe/London
+          export DEBIAN_FRONTEND=noninteractive  # tzdata bug
+          apt-get update -y
+          apt-get install software-properties-common -y
+          add-apt-repository ppa:git-core/candidate -y
+          apt-get update -y
+          apt-get upgrade -y
+          apt-get -y install libglu1-mesa libgl1-mesa-glx libosmesa6 gcc curl g++ unzip wget libglfw3-dev libgles2-mesa-dev libglew-dev sudo git cmake libz-dev libpython3.10-dev
       - name: Setup git
         run: git config --global --add safe.directory /__w/rl/rl
       - name: setup Path
         run: |
           echo /usr/local/bin >> $GITHUB_PATH
-      - name: Setup Environment
+      - name: Setup benchmarks
         run: |
-          python3 -m pip install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cu121 -U
-          python3 -m pip install git+https://github.com/pytorch/tensordict 
-          python3 setup.py develop
-          python3 -m pip install pytest pytest-benchmark
-          python3 -m pip install "gym[accept-rom-license,atari]"
-          python3 -m pip install "dm_control" "mujoco"
-          export TD_GET_DEFAULTS_TO_NONE=1
-      - name: check GPU presence
+          echo "BASE_SHA=$(echo ${{ github.event.pull_request.base.sha }} | cut -c1-8)" >> $GITHUB_ENV
+          echo "HEAD_SHA=$(echo ${{ github.event.pull_request.head.sha }} | cut -c1-8)" >> $GITHUB_ENV
+          echo "BASELINE_JSON=$(mktemp)" >> $GITHUB_ENV
+          echo "CONTENDER_JSON=$(mktemp)" >> $GITHUB_ENV
+          echo "PR_COMMENT=$(mktemp)" >>  $GITHUB_ENV
+      - name: Run
         run: |
-          python -c """import torch
+          python3.10 -m venv --system-site-packages ./py310
+          source ./py310/bin/activate
+          export PYTHON_INCLUDE_DIR=/usr/include/python3.10
+
+          python3.10 -m pip install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cu124 -U
+          python3.10 -m pip install cmake ninja pytest pytest-benchmark mujoco dm_control "gym[accept-rom-license,atari]"
+          python3.10 -m pip install git+https://github.com/pytorch/tensordict 
+          python3.10 setup.py develop
+          # python3.10 -m pip install git+https://github.com/pytorch/rl@$GITHUB_BRANCH
+
+          # test import
+          python3 -c """import torch
           assert torch.cuda.device_count()
           """
-      - name: Run benchmarks
-        run: |
+
           cd benchmarks/
           export TORCHDYNAMO_INLINE_INBUILT_NN_MODULES=1
-          python3 -m pytest --benchmark-json output.json
+          export TD_GET_DEFAULTS_TO_NONE=1
+          python3 -m pytest -vvv --rank 0 --benchmark-json output.json --ignore test_collectors_benchmark.py
       - name: Store benchmark results
         uses: benchmark-action/github-action-benchmark@v1
         if: ${{ github.ref == 'refs/heads/main' || github.event_name == 'workflow_dispatch' }}

.github/workflows/benchmarks_pr.yml

@@ -1,5 +1,4 @@
 name: Continuous Benchmark (PR)
-
 on:
   pull_request:
 
@@ -12,6 +11,7 @@ concurrency:
   cancel-in-progress: true
 
 jobs:
+
   benchmark_cpu:
     name: CPU Pytest benchmark
     runs-on: ubuntu-20.04
@@ -26,28 +26,30 @@ jobs:
       - name: Python Setup
         uses: actions/setup-python@v4
         with:
-          python-version: 3.9
-      - name: Setup Environment
-        run: |
-          python3 -m pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/cpu -U 
-          python3 -m pip install git+https://github.com/pytorch/tensordict 
-          python3 setup.py develop
-          python3 -m pip install pytest pytest-benchmark
-          python3 -m pip install "gym[accept-rom-license,atari]"
-          python3 -m pip install "dm_control" "mujoco"
-          export TD_GET_DEFAULTS_TO_NONE=1
+          python-version: '3.10'
       - name: Setup benchmarks
         run: |
           echo "BASE_SHA=$(echo ${{ github.event.pull_request.base.sha }} | cut -c1-8)" >> $GITHUB_ENV
           echo "HEAD_SHA=$(echo ${{ github.event.pull_request.head.sha }} | cut -c1-8)" >> $GITHUB_ENV
           echo "BASELINE_JSON=$(mktemp)" >> $GITHUB_ENV
           echo "CONTENDER_JSON=$(mktemp)" >> $GITHUB_ENV
           echo "PR_COMMENT=$(mktemp)" >>  $GITHUB_ENV
-      - name: Run benchmarks
+      - name: Setup Environment and tests
         run: |
+          python3.10 -m venv ./py310
+          source ./py310/bin/activate
+
+          python3 -m pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/cpu -U 
+          python3 -m pip install git+https://github.com/pytorch/tensordict 
+          python3 setup.py develop
+          python3 -m pip install pytest pytest-benchmark
+          python3 -m pip install "gym[accept-rom-license,atari]"
+          python3 -m pip install "dm_control" "mujoco"
+
           cd benchmarks/
           export TORCHDYNAMO_INLINE_INBUILT_NN_MODULES=1
-          RUN_BENCHMARK="pytest --rank 0 --benchmark-json "
+          export TD_GET_DEFAULTS_TO_NONE=1
+          RUN_BENCHMARK="python3 -m pytest -vvv --rank 0 --ignore test_collectors_benchmark.py --benchmark-json "
           git checkout ${{ github.event.pull_request.base.sha }}
           $RUN_BENCHMARK ${{ env.BASELINE_JSON }}
           git checkout ${{ github.event.pull_request.head.sha }}
@@ -71,22 +73,23 @@ jobs:
       run:
         shell: bash -l {0}
     container:
-      image: nvidia/cuda:12.3.0-base-ubuntu22.04
+      image: nvidia/cuda:12.4.1-cudnn-runtime-ubuntu22.04
       options: --gpus all
     steps:
+      - name: Set GITHUB_BRANCH environment variable
+        run: |
+          if [ "${{ github.event_name }}" == "push" ]; then
+            export GITHUB_BRANCH=${{ github.event.branch }}
+          elif [ "${{ github.event_name }}" == "pull_request" ]; then
+            export GITHUB_BRANCH=${{ github.event.pull_request.head.ref }}
+          else
+            echo "Unsupported event type"
+            exit 1
+          fi
+          echo "GITHUB_BRANCH=$GITHUB_BRANCH" >> $GITHUB_ENV
       - name: Who triggered this?
         run: |
           echo "Action triggered by ${{ github.event.pull_request.html_url }}"
-      - name: Install deps
-        run: |
-          export TZ=Europe/London
-          export DEBIAN_FRONTEND=noninteractive  # tzdata bug
-          apt-get update -y
-          apt-get install software-properties-common -y
-          add-apt-repository ppa:git-core/candidate -y
-          apt-get update -y
-          apt-get upgrade -y
-          apt-get -y install libglu1-mesa libgl1-mesa-glx libosmesa6 gcc curl g++ unzip wget libglfw3-dev libgles2-mesa-dev libglew-dev sudo git cmake libz-dev
       - name: Check ldd --version
         run: ldd --version
       - name: Checkout
@@ -96,38 +99,50 @@ jobs:
       - name: Python Setup
         uses: actions/setup-python@v4
         with:
-          python-version: 3.9
+          python-version: '3.10'
+      - name: Setup Environment
+        run: |
+          export TZ=Europe/London
+          export DEBIAN_FRONTEND=noninteractive  # tzdata bug
+          apt-get update -y
+          apt-get install software-properties-common -y
+          add-apt-repository ppa:git-core/candidate -y
+          apt-get update -y
+          apt-get upgrade -y
+          apt-get -y install libglu1-mesa libgl1-mesa-glx libosmesa6 gcc curl g++ unzip wget libglfw3-dev libgles2-mesa-dev libglew-dev sudo git cmake libz-dev libpython3.10-dev
       - name: Setup git
         run: git config --global --add safe.directory /__w/rl/rl
       - name: setup Path
         run: |
           echo /usr/local/bin >> $GITHUB_PATH
-      - name: Setup Environment
-        run: |
-          python3 -m pip install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cu121 -U
-          python3 -m pip install git+https://github.com/pytorch/tensordict 
-          python3 setup.py develop
-          python3 -m pip install pytest pytest-benchmark
-          python3 -m pip install "gym[accept-rom-license,atari]"
-          python3 -m pip install "dm_control" "mujoco"
-          export TD_GET_DEFAULTS_TO_NONE=1
-      - name: check GPU presence
-        run: |
-          python -c """import torch
-          assert torch.cuda.device_count()
-          """
       - name: Setup benchmarks
         run: |
           echo "BASE_SHA=$(echo ${{ github.event.pull_request.base.sha }} | cut -c1-8)" >> $GITHUB_ENV
           echo "HEAD_SHA=$(echo ${{ github.event.pull_request.head.sha }} | cut -c1-8)" >> $GITHUB_ENV
           echo "BASELINE_JSON=$(mktemp)" >> $GITHUB_ENV
           echo "CONTENDER_JSON=$(mktemp)" >> $GITHUB_ENV
           echo "PR_COMMENT=$(mktemp)" >>  $GITHUB_ENV
-      - name: Run benchmarks
+      - name: Run
         run: |
+          python3.10 -m venv --system-site-packages ./py310
+          source ./py310/bin/activate
+          export PYTHON_INCLUDE_DIR=/usr/include/python3.10
+
+          python3.10 -m pip install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cu124 -U
+          python3.10 -m pip install cmake ninja pytest pytest-benchmark mujoco dm_control "gym[accept-rom-license,atari]"
+          python3.10 -m pip install git+https://github.com/pytorch/tensordict 
+          python3.10 setup.py develop
+          # python3.10 -m pip install git+https://github.com/pytorch/rl@$GITHUB_BRANCH
+
+          # test import
+          python3 -c """import torch
+          assert torch.cuda.device_count()
+          """
+
           cd benchmarks/
           export TORCHDYNAMO_INLINE_INBUILT_NN_MODULES=1
-          RUN_BENCHMARK="pytest --rank 0 --benchmark-json "
+          export TD_GET_DEFAULTS_TO_NONE=1
+          RUN_BENCHMARK="python3 -m pytest -vvv --rank 0 --ignore test_collectors_benchmark.py --benchmark-json "
           git checkout ${{ github.event.pull_request.base.sha }}
           $RUN_BENCHMARK ${{ env.BASELINE_JSON }}
           git checkout ${{ github.event.pull_request.head.sha }}

setup.py

@@ -152,11 +152,15 @@ def get_extensions():
     }
     sources = list(extension_sources)
 
+    include_dirs = [this_dir]
+    python_include_dir = os.getenv("PYTHON_INCLUDE_DIR")
+    if python_include_dir is not None:
+        include_dirs.append(python_include_dir)
     ext_modules = [
         extension(
             "torchrl._torchrl",
             sources,
-            include_dirs=[this_dir],
+            include_dirs=include_dirs,
             extra_compile_args=extra_compile_args,
             extra_link_args=extra_link_args,
         )

test/test_utils.py

@@ -191,7 +191,7 @@ def test_implement_for_check_versions(
     [
         ("0.27.0", None, "1.0.0"),
         ("0.27.2", None, "1.0.0"),
-        ("1.0.1", "1.0.0", None),
+        # ("1.0.1", "1.0.0", None),
     ],
 )
 @pytest.mark.parametrize(

torchrl/objectives/common.py

@@ -26,7 +26,7 @@
 
 try:
     from torch.compiler import is_dynamo_compiling
-except ModuleNotFoundError:
+except ImportError:
     from torch._dynamo import is_compiling as is_dynamo_compiling