torchft/.github/workflows/unittest.yaml at main · meta-pytorch/torchft · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
name: Unit Tests

on:
  push:
    branches:
      - main
  pull_request:

jobs:
  unittest:
    strategy:
      fail-fast: false
      matrix:
        include:
          - runs-on: "linux.2xlarge"
            gpu-arch-type: "cpu"
            gpu-arch-version: ""
            torch-version: "nightly"
          - runs-on: "linux.g5.12xlarge.nvidia.gpu"
            gpu-arch-type: "cuda"
            gpu-arch-version: "12.4"
            torch-version: "nightly"

    uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main
    with:
      timeout: 120
      runner: ${{ matrix.runs-on }}
      gpu-arch-type: ${{ matrix.gpu-arch-type }}
      gpu-arch-version: ${{ matrix.gpu-arch-version }}
      script: |
        set -ex

        # install python and protobuf
        conda create -n venv python=3.12 libprotobuf -y
        conda activate venv
        python -m pip install --upgrade pip

        # install recent version of Rust via rustup
        curl https://sh.rustup.rs -sSf | sh -s -- --default-toolchain=stable --profile=default -y
        . "$HOME/.cargo/env"

        # Optionally install torch nightly, pulls latest CUDA from pip otherwise
        if [ "${{ matrix.torch-version }}" = "nightly" ]; then
          pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu128
        fi
        if [ "${{ matrix.torch-version }}" = "test" ]; then
          pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/test/cu128
        fi

        # Install dependencies
        pip install -e .[dev] -v

        # Run tests
        pytest -v
        cargo test -v