electrai/.github/workflows/gpu-e2e.yml at main · Quantum-Accelerators/electrai · 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
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
# GPU E2E Training Test
#
# Runs deterministic e2e training tests on EC2 GPU instances.
#
# Can also be run in `update_expected` mode, which generates expected values for {linux, linux-gpu}.
# These are the authoritative expected values for Linux since CI runs on EC2.
# (Note: ubuntu-latest in gen-expected.yml produces slightly different values)
#
# To regenerate expected values:
#   gh workflow run gpu-e2e.yml -f update_expected=true
#   gh run download <run-id> -n expected-linux
#   cp expected_values.json tests/
# (expected-linux contains both linux + linux-gpu values)

name: GPU E2E Training Test

on:
  pull_request:
    branches: [main]
  workflow_dispatch:
    inputs:
      instance_type:
        description: 'EC2 instance type'
        default: 'g6.xlarge'
        type: string
      epochs:
        description: 'Number of training epochs'
        default: '5'
        type: string
      update_expected:
        description: 'Update expected values (for first run or after changes)'
        default: false
        type: boolean
      debug:
        description: 'Debug mode (false/true/N minutes)'
        default: 'false'
        type: string

permissions:
  id-token: write  # Required for AWS OIDC
  contents: read

jobs:
  ec2:
    if: github.event_name != 'pull_request' || github.event.pull_request.head.repo.full_name == github.repository
    uses: Open-Athena/ec2-gha/.github/workflows/runner.yml@v2
    with:
      ec2_instance_type: ${{ inputs.instance_type || 'g6.xlarge' }}
      ec2_image_id: ami-0365bff494b18bf93  # Deep Learning OSS Nvidia Driver AMI GPU PyTorch 2.7 (Ubuntu 22.04)
      aws_tags: '[{"Key": "project", "Value": "electrai"}]'
      debug: ${{ inputs.debug || 'false' }}
    secrets:
      GH_SA_TOKEN: ${{ secrets.GH_SA_TOKEN }}

  gpu-test:
    needs: ec2
    runs-on: ${{ needs.ec2.outputs.id }}
    steps:
      - name: Checkout
        uses: actions/checkout@v6

      - name: Check GPU
        run: nvidia-smi

      - name: Install uv
        uses: astral-sh/setup-uv@v7

      - name: Install dependencies
        run: uv sync

      - name: Run e2e training test (GPU)
        env:
          EPOCHS: ${{ inputs.epochs || '5' }}
          UPDATE_EXPECTED: ${{ inputs.update_expected && '--update-expected' || '' }}
        run: |
          # Run with GPU acceleration
          uv run python scripts/e2e_train.py \
            --gpu \
            --epochs "$EPOCHS" \
            --verbose \
            $UPDATE_EXPECTED

      - name: Upload GPU expected values
        if: ${{ inputs.update_expected }}
        uses: actions/upload-artifact@v7
        with:
          name: expected-linux-gpu
          path: tests/expected_values.json

      - name: Run e2e training test (CPU baseline)
        env:
          EPOCHS: ${{ inputs.epochs || '5' }}
          UPDATE_EXPECTED: ${{ inputs.update_expected && '--update-expected' || '' }}
        run: |
          # Run CPU test to verify determinism
          uv run python scripts/e2e_train.py --epochs "$EPOCHS" --verbose $UPDATE_EXPECTED

      - name: Upload CPU expected values
        if: ${{ inputs.update_expected }}
        uses: actions/upload-artifact@v7
        with:
          name: expected-linux
          path: tests/expected_values.json

      - name: Show GPU memory usage
        if: always()
        run: nvidia-smi