Update

[ghstack-poisoned]

Author: Vincent Moens

.github/workflows/benchmarks.yml

@@ -22,17 +22,64 @@ jobs:
   benchmark_cpu:
     name: CPU Pytest benchmark
     runs-on: linux.g5.4xlarge.nvidia.gpu
+    defaults:
+      run:
+        shell: bash -l {0}
+    container:
+      image: nvidia/cuda:12.3.0-base-ubuntu22.04
+      options: --gpus all
     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: Install deps
+        run: |
+          export TZ=Europe/London
+          export DEBIAN_FRONTEND=noninteractive  # tzdata bug
+          apt-get update -y
+          apt-get update -y
+          apt-get upgrade -y
+          apt-get -y install gcc curl g++ unzip wget sudo git cmake
+      - 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.10'
+      - name: Setup conda
+        run: |
+          rm -rf $HOME/miniconda
+          wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O ~/miniconda.sh
+          bash ~/miniconda.sh -b -p $HOME/miniconda
+      - name: setup Path
+        run: |
+          echo "$HOME/miniconda/bin" >> $GITHUB_PATH
+          echo "CONDA=$HOME/miniconda" >> $GITHUB_PATH
+      - name: create and activate conda env
+        run: |
+          $HOME/miniconda/bin/conda create --name build_binary python=${{ matrix.python-version }}
+          $HOME/miniconda/bin/conda info
+          $HOME/miniconda/bin/activate build_binary
+      - name: Setup git
+        run: git config --global --add safe.directory /__w/tensordict/tensordict
+      - name: setup Path
+        run: |
+          echo /usr/local/bin >> $GITHUB_PATH
       - name: Setup Environment
         run: |
-          apt-get -y install gcc curl g++ unzip wget sudo git cmake
-          pip install ninja -U
-          pip install -e .
-          pip install pytest pytest-benchmark
+          python -m pip install ninja pybind11 -U
+          python -m pip install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cu121
+          python -m pip install -e .
+          python -m pip install pytest pytest-benchmark
+      - name: check GPU presence
+        run: |
+          python -c """import torch
+          assert torch.cuda.device_count()
+          """
       - name: Run benchmarks
         run: |
           cd benchmarks/
@@ -44,7 +91,7 @@ jobs:
         uses: benchmark-action/github-action-benchmark@v1
         if: ${{ github.ref == 'refs/heads/main' || github.event_name == 'workflow_dispatch' }}
         with:
-          name: CPU Benchmark Results
+          name: GPU Benchmark Results
           tool: 'pytest'
           output-file-path: benchmarks/output.json
           fail-on-alert: true
@@ -55,7 +102,6 @@ jobs:
           gh-pages-branch: gh-pages
           auto-push: true
 
-
   benchmark_gpu:
     name: GPU Pytest benchmark
     runs-on: linux.g5.4xlarge.nvidia.gpu
@@ -108,7 +154,7 @@ jobs:
           echo /usr/local/bin >> $GITHUB_PATH
       - name: Setup Environment
         run: |
-          python -m pip install ninja -U
+          python -m pip install ninja pybind11 -U
           python -m pip install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cu121
           python -m pip install -e .
           python -m pip install pytest pytest-benchmark

.github/workflows/benchmarks_pr.yml

@@ -12,13 +12,30 @@ concurrency:
   cancel-in-progress: true
 
 jobs:
+
   benchmark_cpu:
     name: CPU Pytest benchmark
     runs-on: linux.g5.4xlarge.nvidia.gpu
+    defaults:
+      run:
+        shell: bash -l {0}
+    container:
+      image: nvidia/cuda:12.3.0-base-ubuntu22.04
+      options: --gpus all
     steps:
       - 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 update -y
+          apt-get upgrade -y
+          apt-get -y install gcc curl g++ unzip wget sudo git cmake
+      - name: Check ldd --version
+        run: ldd --version
       - name: Checkout
         uses: actions/checkout@v3
         with:
@@ -27,12 +44,36 @@ jobs:
         uses: actions/setup-python@v4
         with:
           python-version: '3.10'
+      - name: Setup conda
+        run: |
+          rm -rf $HOME/miniconda
+          wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O ~/miniconda.sh
+          bash ~/miniconda.sh -b -p $HOME/miniconda
+      - name: setup Path
+        run: |
+          echo "$HOME/miniconda/bin" >> $GITHUB_PATH
+          echo "CONDA=$HOME/miniconda" >> $GITHUB_PATH
+      - name: create and activate conda env
+        run: |
+          $HOME/miniconda/bin/conda create --name build_binary python=${{ matrix.python-version }}
+          $HOME/miniconda/bin/conda info
+          $HOME/miniconda/bin/activate build_binary
+      - name: Setup git
+        run: git config --global --add safe.directory /__w/tensordict/tensordict
+      - name: setup Path
+        run: |
+          echo /usr/local/bin >> $GITHUB_PATH
       - name: Setup Environment
         run: |
-          apt-get -y install gcc curl g++ unzip wget sudo git cmake
-          python -m pip install ninja -U
-          pip install -e .
-          pip install pytest pytest-benchmark
+          python -m pip install ninja pybind11 -U
+          python -m pip install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cu121
+          python -m pip install -e .
+          python -m pip install pytest pytest-benchmark
+      - 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
@@ -61,7 +102,7 @@ jobs:
           comparison-benchmark-metric: 'ops'
           comparison-higher-is-better: true
           comparison-threshold: 5
-          benchmark-title: 'Result of CPU Benchmark Tests'
+          benchmark-title: 'Result of GPU Benchmark Tests'
 
   benchmark_gpu:
     name: GPU Pytest benchmark
@@ -115,7 +156,7 @@ jobs:
           echo /usr/local/bin >> $GITHUB_PATH
       - name: Setup Environment
         run: |
-          python -m pip install ninja -U
+          python -m pip install ninja pybind11 -U
           python -m pip install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cu121
           python -m pip install -e .
           python -m pip install pytest pytest-benchmark