test-gpu.yaml

name: Rust GPU Tests

on:
  push:

concurrency:
  group: "${{ github.workflow }} @ ${{ github.event.pull_request.head.label || github.head_ref || github.ref }}"
  cancel-in-progress: true

jobs:
  e2e:
    runs-on: gpu
    steps:
      - name: Checkout code
        uses: actions/checkout@v4
      - name: Show errors inline
        uses: r7kamura/rust-problem-matchers@v1
      - name: Install apt dependencies
        run: sudo apt-get update && sudo apt-get install -y pkg-config libssl-dev
      - name: Install CUDA and NCCL dependencies
        env:
          DEBIAN_FRONTEND: noninteractive
        run: |
          wget -q https://developer.download.nvidia.com/compute/cuda/12.1.0/local_installers/cuda_12.1.0_530.30.02_linux.run
          sudo sh cuda_12.1.0_530.30.02_linux.run --silent --toolkit
          echo 'export PATH=/usr/local/cuda-12.1/bin:$PATH' >> ~/.bashrc
          echo 'export LD_LIBRARY_PATH=/usr/local/cuda-12.1/lib64:$LD_LIBRARY_PATH' >> ~/.bashrc
          source ~/.bashrc
          /usr/local/cuda-12.1/bin/nvcc --version
          sudo apt-get update
          sudo apt-get install -y libnccl2 libnccl-dev
        shell: bash
      - name: Install Rust nightly
        uses: dtolnay/rust-toolchain@master
        with:
          toolchain: nightly
      - name: E2E Tests
        run: cargo test --release e2e
        shell: bash
        env:
          NCCL_P2P_DIRECT_DISABLE: 1
          NCCL_NET: Socket