ci-gpu.yml

# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# GPU tests live in a separate workflow because NVIDIA self-hosted runners
# block pull_request events entirely. Keeping them here avoids a confusing
# "Skipped" entry with unresolved matrix names on every PR.

name: CI / GPU

on:
  workflow_dispatch:
  push:
    branches:
      - main
      - "pull-request/[0-9]+"
  merge_group:
    types:
      - checks_requested

concurrency:
  group: ${{ github.workflow }}-${{ github.ref }}
  cancel-in-progress: true

env:
  PIP_NO_CACHE_DIR: "1"
  PIP_DISABLE_PIP_VERSION_CHECK: "1"
  PIP_PREFER_BINARY: "1"

jobs:
  gpu-tests:
    runs-on: linux-amd64-gpu-rtxpro6000-latest-1
    # contents: write gives the token push access, which is required to see the
    # permanent-draft release that fetch-models downloads the weights from.
    permissions:
      contents: write
    container:
      image: ubuntu:24.04
      options: -u root --security-opt seccomp=unconfined --shm-size 16g
      env:
        NVIDIA_VISIBLE_DEVICES: ${{ env.NVIDIA_VISIBLE_DEVICES }}
    timeout-minutes: 45
    strategy:
      fail-fast: false
      matrix:
        python-version: ["3.11", "3.12", "3.13"]
        # cu128 = representative CUDA 12.x wheel; cu130 = CUDA 13.0 wheel.
        torch-cuda: ["cu128", "cu130"]
    name: "gpu / py${{ matrix.python-version }} / ${{ matrix.torch-cuda }}"
    steps:
      - name: Install system dependencies
        run: |
          export DEBIAN_FRONTEND=noninteractive
          apt-get update
          apt-get install -y git gcc

      - uses: actions/setup-python@v5
        with:
          python-version: ${{ matrix.python-version }}

      - uses: actions/checkout@v4

      - name: Fetch pre-trained models
        uses: ./.github/actions/fetch-models

      - name: Verify GPU
        run: nvidia-smi

      - name: Install dependencies and run tests
        run: bash code/scripts/check_python_compat.sh
        env:
          PYTHON_BIN: python
          MODE: train
          SKIP_TESTS: "0"
          REQUIRE_GPU: "1"
          TORCH_CUDA: ${{ matrix.torch-cuda }}
          VENV_DIR: .venv_train_${{ matrix.python-version }}_${{ matrix.torch-cuda }}
          REQ_FILE: code/requirements_public_gpu_${{ matrix.torch-cuda == 'cu130' && 'cu13' || 'cu12' }}.txt

      - name: Training + inference with LER check
        shell: bash
        run: |
          source .venv_train_${{ matrix.python-version }}_${{ matrix.torch-cuda }}/bin/activate
          bash code/scripts/smoke_run.sh 2>&1 | tee /tmp/ci_train.log
          r=${PIPESTATUS[0]}; [ $r -ne 0 ] && exit $r
          # 0.35: short run (16k samples, 4 epochs for stable LER across py versions)
          python code/scripts/check_ler_from_log.py /tmp/ci_train.log --max-ler 0.35
        env:
          EXPERIMENT_NAME: ci_short
          PREDECODER_TRAIN_SAMPLES: "16384"
          PREDECODER_VAL_SAMPLES: "2048"
          PREDECODER_TEST_SAMPLES: "2048"
          PREDECODER_TRAIN_EPOCHS: "4"

      - name: Training + inference with multi-worker DataLoader (num_workers=2)
        shell: bash
        run: |
          source .venv_train_${{ matrix.python-version }}_${{ matrix.torch-cuda }}/bin/activate
          bash code/scripts/smoke_run.sh 2>&1 | tee /tmp/ci_multiworker.log
          r=${PIPESTATUS[0]}; [ $r -ne 0 ] && exit $r
          python code/scripts/check_ler_from_log.py /tmp/ci_multiworker.log --max-ler 0.35
        env:
          EXPERIMENT_NAME: ci_multiworker
          PREDECODER_TRAIN_SAMPLES: "16384"
          PREDECODER_VAL_SAMPLES: "2048"
          PREDECODER_TEST_SAMPLES: "2048"
          PREDECODER_TRAIN_EPOCHS: "4"
          PREDECODER_INFERENCE_NUM_WORKERS: "2"

  # ---------------------------------------------------------------------------
  # Mid-tier (~5-10 min): extended training + inference with LER check.
  # Runs only after merge to main (not on PR branches) to save GPU time.
  # Single Python version — multi-version coverage is handled by gpu-tests.
  # ---------------------------------------------------------------------------
  mid-gpu-tests:
    if: github.ref == 'refs/heads/main'
    needs: gpu-tests
    runs-on: linux-amd64-gpu-rtxpro6000-latest-1
    container:
      image: ubuntu:24.04
      options: -u root --security-opt seccomp=unconfined --shm-size 16g
      env:
        NVIDIA_VISIBLE_DEVICES: ${{ env.NVIDIA_VISIBLE_DEVICES }}
    timeout-minutes: 40
    steps:
      - name: Install system dependencies
        run: |
          export DEBIAN_FRONTEND=noninteractive
          apt-get update
          apt-get install -y git gcc

      - uses: actions/setup-python@v5
        with:
          python-version: "3.13"

      - uses: actions/checkout@v4

      - name: Verify GPU
        run: nvidia-smi

      - name: Install Python dependencies
        run: |
          python -m venv .venv_mid
          . .venv_mid/bin/activate
          python -m pip install --upgrade pip setuptools wheel
          # TODO: matrix by CUDA major version [cu12, cu13]
          pip install -r code/requirements_public_train-cu12.txt

      - name: Mid-tier training + inference with LER check (32k train, 4 epochs)
        shell: bash
        run: |
          . .venv_mid/bin/activate
          bash code/scripts/smoke_run.sh 2>&1 | tee /tmp/ci_mid.log
          r=${PIPESTATUS[0]}; [ $r -ne 0 ] && exit $r
          # 0.2: mid-tier (32k/4 epochs); loosen if flaky
          python code/scripts/check_ler_from_log.py /tmp/ci_mid.log --max-ler 0.2
        env:
          EXPERIMENT_NAME: ci_mid
          PREDECODER_TRAIN_SAMPLES: "32768"
          PREDECODER_VAL_SAMPLES: "4096"
          PREDECODER_TEST_SAMPLES: "4096"
          PREDECODER_TRAIN_EPOCHS: "4"

      - name: HE compile tests (torch.compile + autotune on GPU)
        run: |
          . .venv_mid/bin/activate
          PYTHONPATH=code python -m unittest discover -s code/tests/mid -p "test_*.py" -v

  # ---------------------------------------------------------------------------
  # Multi-GPU tests: validates NCCL, DDP gradient sync, and per-rank data
  # generation across 2 GPUs.
  #
  # Runner requirement: a self-hosted runner with >=2 GPUs.
  # NVIDIA GHA runners follow the naming pattern
  #   linux-amd64-gpu-<model>-latest-<gpu-count>
  # so the 2-GPU variant of the existing rtxpro6000 runner would be:
  #   linux-amd64-gpu-rtxpro6000-latest-2
  # Confirm this label with your runner pool before enabling; if no 2-GPU
  # runner exists the job will queue indefinitely.
  #
  # Runs only after merge to main (not on PR branches) to conserve GPU quota.
  # ---------------------------------------------------------------------------
  multi-gpu-tests:
    if: github.ref == 'refs/heads/main'
    needs: gpu-tests
    runs-on: linux-amd64-gpu-rtxpro6000-latest-2
    container:
      image: ubuntu:24.04
      options: -u root --security-opt seccomp=unconfined --shm-size 16g
      env:
        NVIDIA_VISIBLE_DEVICES: ${{ env.NVIDIA_VISIBLE_DEVICES }}
    timeout-minutes: 35
    steps:
      - name: Install system dependencies
        run: |
          export DEBIAN_FRONTEND=noninteractive
          apt-get update
          apt-get install -y git python3 python3-pip python3-venv

      - uses: actions/checkout@v4

      - name: Verify 2 GPUs are visible
        run: |
          nvidia-smi
          count=$(nvidia-smi --query-gpu=name --format=csv,noheader | wc -l)
          echo "GPU count: ${count}"
          [ "${count}" -ge 2 ] || { echo "ERROR: expected >=2 GPUs, found ${count}"; exit 1; }

      - name: Install Python dependencies
        run: |
          python3 -m venv .venv_multigpu
          . .venv_multigpu/bin/activate
          python -m pip install --upgrade pip setuptools wheel
          # TODO: matrix by CUDA major version [cu12, cu13]
          pip install -r code/requirements_public_train-cu12.txt

      - name: Run multi-GPU unit tests
        run: |
          . .venv_multigpu/bin/activate
          PYTHONPATH=code python -m unittest discover \
            -s code/tests -p "test_multi_gpu.py" -v

      - name: Multi-GPU smoke training (2 GPUs, DDP)
        # smoke_run.sh hardcodes GPUS=1; call local_run.sh directly so we
        # can pass GPUS=2 and exercise the torch.distributed.run path.
        shell: bash
        run: |
          . .venv_multigpu/bin/activate
          export PREDECODER_TIMING_RUN=1
          export PREDECODER_DISABLE_SDR=1
          export PREDECODER_LER_FINAL_ONLY=1
          export PREDECODER_INFERENCE_NUM_SAMPLES=32
          export PREDECODER_INFERENCE_LATENCY_SAMPLES=0
          export PREDECODER_INFERENCE_MEAS_BASIS=both
          export PREDECODER_INFERENCE_NUM_WORKERS=0
          EXPERIMENT_NAME=ci_multi_gpu WORKFLOW=train GPUS=2 \
            bash code/scripts/local_run.sh 2>&1 | tee /tmp/ci_multigpu_train.log
          r=${PIPESTATUS[0]}; [ $r -ne 0 ] && exit $r
          EXPERIMENT_NAME=ci_multi_gpu WORKFLOW=inference GPUS=2 \
            bash code/scripts/local_run.sh 2>&1 | tee /tmp/ci_multigpu_infer.log
          r=${PIPESTATUS[0]}; [ $r -ne 0 ] && exit $r
          # [LER Validation] lines are emitted during training, not inference
          python code/scripts/check_ler_from_log.py /tmp/ci_multigpu_train.log --max-ler 0.35
        env:
          PREDECODER_TRAIN_SAMPLES: "16384"
          PREDECODER_VAL_SAMPLES: "2048"
          PREDECODER_TEST_SAMPLES: "2048"
          PREDECODER_TRAIN_EPOCHS: "4"

      - name: Multi-GPU smoke training with parallel spacelike HE (2 GPUs, DDP)
        # Additive coverage on top of the default-config multi-GPU smoke above.
        # Forces data.use_compile=True + data.use_parallel_spacelike=True so the
        # parallel + compiled spacelike HE path runs end-to-end under DDP on
        # 2 GPUs. Failure modes specific to this combination (per-rank device
        # pinning of the partition, torch.compile cache contention across
        # ranks, deadlocks during the compiled inner loop) surface as a
        # training crash here. The existing default-config step above is
        # intentionally left untouched so we do not regress on coverage of
        # the default path.
        shell: bash
        run: |
          . .venv_multigpu/bin/activate
          export PREDECODER_TIMING_RUN=1
          export PREDECODER_DISABLE_SDR=1
          export PREDECODER_LER_FINAL_ONLY=1
          export PREDECODER_INFERENCE_NUM_SAMPLES=32
          export PREDECODER_INFERENCE_LATENCY_SAMPLES=0
          export PREDECODER_INFERENCE_MEAS_BASIS=both
          export PREDECODER_INFERENCE_NUM_WORKERS=0
          EXPERIMENT_NAME=ci_multi_gpu_he WORKFLOW=train GPUS=2 \
            EXTRA_PARAMS="data.use_compile=True data.use_parallel_spacelike=True" \
            bash code/scripts/local_run.sh 2>&1 | tee /tmp/ci_multigpu_he_train.log
          r=${PIPESTATUS[0]}; [ $r -ne 0 ] && exit $r
          EXPERIMENT_NAME=ci_multi_gpu_he WORKFLOW=inference GPUS=2 \
            EXTRA_PARAMS="data.use_compile=True data.use_parallel_spacelike=True" \
            bash code/scripts/local_run.sh 2>&1 | tee /tmp/ci_multigpu_he_infer.log
          r=${PIPESTATUS[0]}; [ $r -ne 0 ] && exit $r
          python code/scripts/check_ler_from_log.py /tmp/ci_multigpu_he_train.log --max-ler 0.35
        env:
          PREDECODER_TRAIN_SAMPLES: "16384"
          PREDECODER_VAL_SAMPLES: "2048"
          PREDECODER_TEST_SAMPLES: "2048"
          PREDECODER_TRAIN_EPOCHS: "4"

  # ---------------------------------------------------------------------------
  # GPU coverage: captures GPU-specific code paths missed by the CPU coverage job
  # ---------------------------------------------------------------------------
  gpu-coverage:
    runs-on: linux-amd64-gpu-rtxpro6000-latest-1
    # contents: write — see gpu-tests: needed to read the draft release.
    permissions:
      contents: write
    container:
      image: ubuntu:24.04
      options: -u root --security-opt seccomp=unconfined --shm-size 16g
      env:
        NVIDIA_VISIBLE_DEVICES: ${{ env.NVIDIA_VISIBLE_DEVICES }}
    timeout-minutes: 20
    steps:
      - name: Install system dependencies
        run: |
          export DEBIAN_FRONTEND=noninteractive
          apt-get update
          apt-get install -y git python3 python3-pip python3-venv

      - uses: actions/checkout@v4

      - name: Fetch pre-trained models
        uses: ./.github/actions/fetch-models

      - name: Verify GPU
        run: nvidia-smi

      - name: Install Python dependencies
        run: |
          python3 -m venv .venv_gpu_cov
          . .venv_gpu_cov/bin/activate
          python -m pip install --upgrade pip setuptools wheel
          # TODO: matrix by CUDA major version [cu12, cu13]
          pip install -r code/requirements_public_train-cu12.txt
          pip install -r code/requirements_ci.txt

      - name: Run tests with GPU coverage
        run: |
          . .venv_gpu_cov/bin/activate
          PYTHONPATH=code coverage run -m unittest discover -s code/tests -p "test_*.py"
          coverage report
          coverage html -d htmlcov-gpu
          coverage xml -o coverage-gpu.xml

      - name: Upload GPU coverage artifacts
        if: always()
        uses: actions/upload-artifact@v4
        with:
          name: gpu-coverage-report
          path: |
            htmlcov-gpu/
            coverage-gpu.xml