name: Run 1xH100 Tests

on:
  push:
    branches:
      - main
      - 'gh/**'
  pull_request:
    branches:
      - main
      - 'gh/**'

concurrency:
  group: 1xH100_tests-${{ github.workflow }}-${{ github.ref == 'refs/heads/main' && github.run_number || github.ref }}
  cancel-in-progress: true

env:
  HF_TOKEN: ${{ secrets.HF_TOKEN }}

jobs:
  test:
    strategy:
      fail-fast: false
      matrix:
        include:
          - name: H100
            runs-on: mt-l-x86iamx-22-225-h100
            torch-spec: '--pre torch torchvision mslk --index-url https://download.pytorch.org/whl/nightly/cu130'
            gpu-arch-type: "cuda"
            gpu-arch-version: "13.0"
          - name: H100 2.13
            runs-on: mt-l-x86iamx-22-225-h100
            torch-spec: 'torch==2.13.0 torchvision==0.28.0 mslk==1.3.0 --index-url https://download.pytorch.org/whl/cu130'
            gpu-arch-type: "cuda"
            gpu-arch-version: "13.0"
    permissions:
      id-token: write
      contents: read
    uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main
    with:
      timeout: 90
      runner: ${{ matrix.runs-on }}
      gpu-arch-type: ${{ matrix.gpu-arch-type }}
      gpu-arch-version: ${{ matrix.gpu-arch-version }}
      submodules: recursive
      script: |
        conda create -n venv python=3.10 -y
        conda activate venv
        export PATH=/opt/rh/devtoolset-10/root/usr/bin/:$PATH
        python -m pip install --upgrade pip
        pip install uv
        # Pre-install torch/vision's pure-python deps from the in-cluster pypi-cache for speed.
        #
        # spmd-types is appended on purpose: the torch nightly hard-depends on it, but the pytorch
        # index only re-links to files.pythonhosted.org, which the OSDC H100 pool can't reach.
        # Installing it now via the reachable cache means the cache-bypassing torch install below
        # (--index-url only) never has to fetch it.
        #
        # TODO: remove once the OSDC pool can reach files.pythonhosted.org (or spmd-types is
        # mirrored on download.pytorch.org). Egress owned by PyTorch dev infra.
        pip install filelock typing-extensions "setuptools<82" sympy networkx jinja2 fsspec numpy pillow spmd-types
        # Clear PIP_EXTRA_INDEX_URL (the runner's default cpu /whl/cpu/) so it can't supply a
        # +cpu torch; the torch-spec's --index-url makes the literal nightly cu130 index the only source.
        PIP_EXTRA_INDEX_URL= pip install ${{ matrix.torch-spec }}
        pip install -U --pre flashinfer-python --index-url https://flashinfer.ai/whl/nightly/ --no-deps
        pip install apache-tvm-ffi
        pip install nvidia-ml-py
        uv pip install -r dev-requirements.txt
        # Editable install so build_ext --inplace drops the compiled extensions
        # (e.g. _C_cutlass_90a.abi3.so) into the in-tree torchao/ dir. pytest is run
        # from the repo root and, because test/ is a package with a conftest.py,
        # prepends the repo root to sys.path -- so `import torchao` resolves to the
        # source tree, which shadows the installed wheel. A non-editable `pip install .`
        # leaves that source tree without the .so, so compiled CUDA ops (e.g. the
        # sparse cutlass op) never register and their tests fail/skip. Editable keeps
        # the extensions available in the shadowing source tree.
        pip install -e . --no-build-isolation -vv
        pytest test/integration --verbose -s
        pytest test/dtypes/test_affine_quantized_float.py --verbose -s
        pytest test/quantization/quantize_/workflows/float8 --verbose -s
        python test/quantization/quantize_/workflows/int4/test_int4_tensor.py
        python test/quantization/quantize_/workflows/int4/test_int4_preshuffled_tensor.py
        ./test/float8/test_everything_single_gpu.sh
        pytest test/prototype/mx_formats/ --verbose -s
        pytest test/prototype/moe_training/test_fp8_grouped_mm.py --verbose -s
        pytest test/prototype/moe_training/test_mxfp8_grouped_mm.py --verbose -s
        pytest test/prototype/moe_training/test_training.py --verbose -s