Merge pull request #3 from samuelburnham/gpu-bench

Test GPU bench on AWS
samuelburnham · Oct 24, 2023 · 79f9c5c · 79f9c5c
2 parents aa63a65 + 2a14ceb
commit 79f9c5c
Showing 1 changed file with 22 additions and 5 deletions.
diff --git a/.github/workflows/bench_pr_comment.yml b/.github/workflows/bench_pr_comment.yml
@@ -14,7 +14,7 @@ concurrency:
 jobs:
   run-benchmark:
     name: run end2end benchmark
-    runs-on: ubuntu-benchmark-runner
+    runs-on: buildjet-16vcpu-ubuntu-2204
     if:
       github.event.issue.pull_request
       && github.event.issue.state == 'open'
@@ -41,28 +41,45 @@ jobs:
 
   gpu-benchmark:
     name: run fibonacci benchmark on GPU
-    runs-on: [self-hosted, gpu-bench]
+    runs-on: [self-hosted, gpu-bench-aws]
+    env:
+      NVIDIA_VISIBLE_DEVICES: all
+      NVIDIA_DRIVER_CAPABILITITES: compute,utility
+      EC_GPU_FRAMEWORK: cuda
     if:
       github.event.issue.pull_request
       && github.event.issue.state == 'open'
-      && contains(github.event.comment.body, '!benchmark')
+      && contains(github.event.comment.body, '!gpu-benchmark')
       && (github.event.comment.author_association == 'MEMBER' || github.event.comment.author_association == 'OWNER')
     steps:
       - uses: xt0rted/pull-request-comment-branch@v2
         id: comment-branch
-
       - uses: actions/checkout@v4
         if: success()
         with:
           ref: ${{ steps.comment-branch.outputs.head_ref }}
+          submodules: recursive
       # Set the Rust env vars
       - uses: actions-rs/toolchain@v1
       - uses: Swatinem/rust-cache@v2
+      # Check we have access to the machine's Nvidia drivers
+      - run: nvidia-smi
+      # The `compute`/`sm` number corresponds to the Nvidia GPU architecture
+      # In this case, the self-hosted machine uses the Ampere architecture, but we want this to be configurable
+      # See https://arnon.dk/matching-sm-architectures-arch-and-gencode-for-various-nvidia-cards/
+      - name: Set env for CUDA compute
+        run: echo "CUDA_ARCH=$(nvidia-smi --query-gpu=compute_cap --format=csv,noheader | sed 's/\.//g')" >> $GITHUB_ENV
+      - name: set env for EC_GPU
+        run: echo 'EC_GPU_CUDA_NVCC_ARGS=--fatbin --gpu-architecture=sm_${{ env.CUDA_ARCH }} --generate-code=arch=compute_${{ env.CUDA_ARCH }},code=sm_${{ env.CUDA_ARCH }}' >> $GITHUB_ENV
+      - run: echo "${{ env.EC_GPU_CUDA_NVCC_ARGS}}"
+      # Check that CUDA is installed with a driver-compatible version
+      # This must also be compatible with the GPU architecture, see above link
+      - run: nvcc --version
       - uses: boa-dev/criterion-compare-action@v3
         with:
           # Optional. Compare only this benchmark target
           benchName: "fibonacci"
           # Optional. Features activated in the benchmark
-          features: "cuda,opencl"
+          features: "cuda"
           # Needed. The name of the branch to compare with
           branchName: ${{ github.ref_name }}