[Feature] Benchmarking worflow (pytorch#1028)

ai4co · Apr 6, 2023 · 38496d2 · 38496d2
1 parent f167c55
commit 38496d2
Show file tree

Hide file tree

Showing 9 changed files with 595 additions and 4 deletions.
diff --git a/.github/workflows/benchmarks.yml b/.github/workflows/benchmarks.yml
@@ -0,0 +1,119 @@
+name: Continuous Benchmark
+on:
+  push:
+    branches:
+      - main
+  workflow_dispatch:
+
+permissions:
+  deployments: write
+  contents: write
+
+jobs:
+
+  benchmark_cpu:
+    name: CPU Pytest benchmark
+    runs-on: ubuntu-20.04
+    steps:
+      - uses: actions/checkout@v3
+      - uses: actions/setup-python@v4
+        with:
+          python-version: 3.8
+      - name: Setup Environment
+        run: |
+          python -m pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/cpu 
+          python -m pip install git+https://github.com/pytorch-labs/tensordict 
+          python -m pip install -e .
+          python -m pip install pytest pytest-benchmark
+          python -m pip install dm_control
+      - name: Run benchmarks
+        run: |
+          cd benchmarks/
+          python -m pytest --benchmark-json output.json
+      - name: Store benchmark results
+        uses: benchmark-action/github-action-benchmark@v1
+        with:
+          name: CPU Benchmark Results
+          tool: 'pytest'
+          output-file-path: benchmarks/output.json
+          fail-on-alert: true
+          alert-threshold: '200%'
+          alert-comment-cc-users: '@vmoens'
+          comment-on-alert: true
+          github-token: ${{ secrets.GITHUB_TOKEN }}
+          gh-pages-branch: gh-pages
+          auto-push: true
+
+
+  benchmark_gpu:
+    name: GPU Pytest benchmark
+    runs-on: ubuntu-20.04
+    strategy:
+      matrix:
+        include:
+         - os: linux.4xlarge.nvidia.gpu
+           python-version: 3.8
+    defaults:
+      run:
+        shell: bash -l {0}
+    container: nvidia/cuda:11.8.0-cudnn8-devel-ubuntu18.04
+    steps:
+    - name: Install deps
+      run: |
+        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
+      uses: actions/checkout@v3
+    - name: Update pip
+      run: |
+        apt-get install python3.8 python3-pip -y
+        pip3 install --upgrade pip
+    - 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/rl/rl
+    - name: setup Path
+      run: |
+        echo /usr/local/bin >> $GITHUB_PATH
+    - name: Setup Environment
+      run: |
+        python -m pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu118 
+        python -m pip install git+https://github.com/pytorch-labs/tensordict 
+        python -m pip install -e .
+        python -m pip install pytest pytest-benchmark
+        python -m pip install dm_control
+    - name: Run benchmarks
+      run: |
+        cd benchmarks/
+        python -m pytest --benchmark-json output.json
+    - name: Store benchmark results
+      uses: benchmark-action/github-action-benchmark@v1
+      with:
+        name: GPU Benchmark Results
+        tool: 'pytest'
+        output-file-path: benchmarks/output.json
+        fail-on-alert: true
+        alert-threshold: '200%'
+        alert-comment-cc-users: '@vmoens'
+        comment-on-alert: true
+        github-token: ${{ secrets.GITHUB_TOKEN }}
+        gh-pages-branch: gh-pages
+        auto-push: true
diff --git a/README.md b/README.md
@@ -1,5 +1,6 @@
 [![pytorch](https://circleci.com/gh/pytorch/rl.svg?style=shield)](https://circleci.com/gh/pytorch/rl)
 [![Documentation](https://img.shields.io/badge/Documentation-blue.svg)](https://pytorch.org/rl/)
+[![Benchmarks](https://img.shields.io/badge/Benchmarks-blue.svg)](https://pytorch.github.io/rl/dev/bench/)
 [![codecov](https://codecov.io/gh/pytorch/rl/branch/main/graph/badge.svg?token=HcpK1ILV6r)](https://codecov.io/gh/pytorch/rl)
 [![Twitter Follow](https://img.shields.io/twitter/follow/torchrl1?style=social)](https://twitter.com/torchrl1)
 [![Python 3.7, 3.8](https://img.shields.io/badge/python-3.7%20%7C%203.8%20%7C%203.9%20%7C%203.10-blue.svg)](https://www.python.org/downloads/)

diff --git a/benchmarks/benchmark_batched_envs.py b/benchmarks/benchmark_batched_envs.py
@@ -1,3 +1,7 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
 """
 Benchmarking different types of batched environments
 ====================================================

diff --git a/benchmarks/requirements.txt b/benchmarks/requirements.txt
@@ -0,0 +1,2 @@
+pytest-benchmark
+tenacity
diff --git a/benchmarks/test_collectors_benchmark.py b/benchmarks/test_collectors_benchmark.py
@@ -0,0 +1,176 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+import pytest
+import torch.cuda
+
+from torchrl.collectors import SyncDataCollector
+from torchrl.collectors.collectors import (
+    MultiaSyncDataCollector,
+    MultiSyncDataCollector,
+    RandomPolicy,
+)
+from torchrl.envs import EnvCreator, StepCounter, TransformedEnv
+from torchrl.envs.libs.dm_control import DMControlEnv
+
+
+def single_collector_setup():
+    device = "cuda:0" if torch.cuda.device_count() else "cpu"
+    env = TransformedEnv(DMControlEnv("cheetah", "run", device=device), StepCounter(50))
+    c = SyncDataCollector(
+        env,
+        RandomPolicy(env.action_spec),
+        total_frames=10_000,
+        frames_per_batch=100,
+        device=device,
+    )
+    c = iter(c)
+    for _ in c:
+        break
+    return ((c,), {})
+
+
+def sync_collector_setup():
+    device = "cuda:0" if torch.cuda.device_count() else "cpu"
+    env = EnvCreator(
+        lambda: TransformedEnv(
+            DMControlEnv("cheetah", "run", device=device), StepCounter(50)
+        )
+    )
+    c = MultiSyncDataCollector(
+        [env, env],
+        RandomPolicy(env().action_spec),
+        total_frames=10_000,
+        frames_per_batch=100,
+        device=device,
+    )
+    c = iter(c)
+    for _ in c:
+        break
+    return ((c,), {})
+
+
+def async_collector_setup():
+    device = "cuda:0" if torch.cuda.device_count() else "cpu"
+    env = EnvCreator(
+        lambda: TransformedEnv(
+            DMControlEnv("cheetah", "run", device=device), StepCounter(50)
+        )
+    )
+    c = MultiaSyncDataCollector(
+        [env, env],
+        RandomPolicy(env().action_spec),
+        total_frames=10_000,
+        frames_per_batch=100,
+        device=device,
+    )
+    c = iter(c)
+    for _ in c:
+        break
+    return ((c,), {})
+
+
+def single_collector_setup_pixels():
+    device = "cuda:0" if torch.cuda.device_count() else "cpu"
+    env = TransformedEnv(
+        DMControlEnv("cheetah", "run", device=device, from_pixels=True), StepCounter(50)
+    )
+    c = SyncDataCollector(
+        env,
+        RandomPolicy(env.action_spec),
+        total_frames=10_000,
+        frames_per_batch=100,
+        device=device,
+    )
+    c = iter(c)
+    for _ in c:
+        break
+    return ((c,), {})
+
+
+def sync_collector_setup_pixels():
+    device = "cuda:0" if torch.cuda.device_count() else "cpu"
+    env = EnvCreator(
+        lambda: TransformedEnv(
+            DMControlEnv("cheetah", "run", device=device, from_pixels=True),
+            StepCounter(50),
+        )
+    )
+    c = MultiSyncDataCollector(
+        [env, env],
+        RandomPolicy(env().action_spec),
+        total_frames=10_000,
+        frames_per_batch=100,
+        device=device,
+    )
+    c = iter(c)
+    for _ in c:
+        break
+    return ((c,), {})
+
+
+def async_collector_setup_pixels():
+    device = "cuda:0" if torch.cuda.device_count() else "cpu"
+    env = EnvCreator(
+        lambda: TransformedEnv(
+            DMControlEnv("cheetah", "run", device=device, from_pixels=True),
+            StepCounter(50),
+        )
+    )
+    c = MultiaSyncDataCollector(
+        [env, env],
+        RandomPolicy(env().action_spec),
+        total_frames=10_000,
+        frames_per_batch=100,
+        device=device,
+    )
+    c = iter(c)
+    for _ in c:
+        break
+    return ((c,), {})
+
+
+def execute_collector(c):
+    # will run for 9 iterations (1 during setup)
+    for _ in c:
+        continue
+
+
+def test_single(benchmark):
+    benchmark.pedantic(
+        execute_collector, setup=single_collector_setup, iterations=1, rounds=5
+    )
+
+
+def test_sync(benchmark):
+    benchmark.pedantic(
+        execute_collector, setup=sync_collector_setup, iterations=1, rounds=5
+    )
+
+
+def test_async(benchmark):
+    benchmark.pedantic(
+        execute_collector, setup=async_collector_setup, iterations=1, rounds=5
+    )
+
+
+@pytest.mark.skipif(not torch.cuda.device_count(), reason="no rendering without cuda")
+def test_single_pixels(benchmark):
+    benchmark.pedantic(
+        execute_collector, setup=single_collector_setup_pixels, iterations=1, rounds=5
+    )
+
+
+@pytest.mark.skipif(not torch.cuda.device_count(), reason="no rendering without cuda")
+def test_sync_pixels(benchmark):
+    benchmark.pedantic(
+        execute_collector, setup=sync_collector_setup_pixels, iterations=1, rounds=5
+    )
+
+
+@pytest.mark.skipif(not torch.cuda.device_count(), reason="no rendering without cuda")
+def test_async_pixels(benchmark):
+    benchmark.pedantic(
+        execute_collector, setup=async_collector_setup_pixels, iterations=1, rounds=5
+    )
diff --git a/benchmarks/test_envs_benchmark.py b/benchmarks/test_envs_benchmark.py
@@ -0,0 +1,56 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+import torch
+
+from torchrl.envs import ParallelEnv, SerialEnv, StepCounter, TransformedEnv
+from torchrl.envs.libs.dm_control import DMControlEnv
+
+
+def make_simple_env():
+    device = "cuda:0" if torch.cuda.device_count() else "cpu"
+    env = DMControlEnv("cheetah", "run", device=device)
+    env.rollout(3)
+    return ((env,), {})
+
+
+def make_transformed_env():
+    device = "cuda:0" if torch.cuda.device_count() else "cpu"
+    env = TransformedEnv(DMControlEnv("cheetah", "run", device=device), StepCounter(50))
+    env.rollout(3)
+    return ((env,), {})
+
+
+def make_serial_env():
+    device = "cuda:0" if torch.cuda.device_count() else "cpu"
+    env = SerialEnv(3, lambda: DMControlEnv("cheetah", "run", device=device))
+    env.rollout(3)
+    return ((env,), {})
+
+
+def make_parallel_env():
+    device = "cuda:0" if torch.cuda.device_count() else "cpu"
+    env = ParallelEnv(3, lambda: DMControlEnv("cheetah", "run", device=device))
+    env.rollout(3)
+    return ((env,), {})
+
+
+def execute_env(env):
+    env.rollout(1000, break_when_any_done=False)
+
+
+def test_simple(benchmark):
+    benchmark.pedantic(execute_env, setup=make_simple_env, iterations=1, rounds=5)
+
+
+def test_transformed(benchmark):
+    benchmark.pedantic(execute_env, setup=make_transformed_env, iterations=1, rounds=5)
+
+
+def test_serial(benchmark):
+    benchmark.pedantic(execute_env, setup=make_serial_env, iterations=1, rounds=5)
+
+
+def test_parallel(benchmark):
+    benchmark.pedantic(execute_env, setup=make_parallel_env, iterations=1, rounds=5)