name: nv-a6000

on:
  pull_request:
    paths:
      - 'accelerator/cuda_accelerator.py'
      - 'deepspeed/inference/v2/**'
      - 'tests/unit/inference/v2/**'
      - '.github/workflows/nv-a6000.yml'
  workflow_dispatch:
    inputs:
      mii_branch:
        description: 'DeepSpeed-MII Branch'
        required: false
        default: 'main'
        type: string

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

jobs:
  unit-tests:
    name: nv-a6000 / inference and MII tests
    runs-on: [self-hosted, nvidia, a6000]
    container:
      image: nvcr.io/nvidia/pytorch:25.01-py3
      ports:
        - 80
      options: --gpus all --shm-size "8G"

    steps:
      - uses: actions/checkout@v4

      - name: Check container state
        run: |
          ldd --version
          nvcc --version
          nvidia-smi
          python -c "import torch; print('torch:', torch.__version__, torch)"
          python -c "import torch; print('CUDA available:', torch.cuda.is_available())"
      - name: Install transformers
        run: |
          git clone https://github.com/huggingface/transformers
          cd transformers
          # if you need to use an older transformers version temporarily in case of breakage
          # git checkout 981c276
          git rev-parse --short HEAD
          python -m pip install .
      - name: Install deepspeed
        run: |
          python -m pip install docutils==0.18.1 jinja2==3.0 urllib3==1.26.11 ninja
          python -m pip install .[dev,1bit,autotuning,inf]
          ds_report
      - name: Python environment
        run: |
          python -m pip list
      - name: Unit tests
        run: |
          unset TORCH_CUDA_ARCH_LIST # only jit compile for current arch
          cd tests
          python -m pytest --color=yes --durations=0 --verbose -rF -m 'inference_v2' unit/ --torch_ver="2.6" --cuda_ver="12"
          python -m pytest --color=yes --durations=0 --verbose -rF -m 'inference_v2_ops' unit/ --torch_ver="2.6" --cuda_ver="12"
      - name: MII unit tests
        run: |
          BRANCH="main"
          if [[ ! -z "${{ github.event.inputs.mii_branch }}" ]]; then
              BRANCH="${{ github.event.inputs.mii_branch }}"
          fi
          echo "Cloning DeepSpeed-MII branch: $BRANCH"
          git clone -b $BRANCH --depth=1 https://github.com/deepspeedai/DeepSpeed-MII.git
          cd DeepSpeed-MII
          pip install .[dev]
          cd tests
          python -m pytest --color=yes --durations=0 --verbose -rF ./