name: GraphTrainer 8 GPU H100 Integration Tests

on:
  push:
    branches: [ main ]
    paths:
      - 'torchtitan/experiments/graph_trainer/**'
      - '.github/workflows/integration_test_8gpu_graph_trainer_h100.yaml'
  pull_request:
    types: [labeled, synchronize]
    paths:
      - 'torchtitan/experiments/graph_trainer/**'
      - '.github/workflows/integration_test_8gpu_graph_trainer_h100.yaml'
  schedule:
    # Runs every 12 hours
    - cron: '0 */12 * * *'

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

defaults:
  run:
    shell: bash -l -eo pipefail {0}

permissions:
  id-token: write
  contents: read

jobs:
  # Step 1: Dynamically compute the matrix based on conditions
  set-matrix:
    # Skip scheduled runs on forks, where they would only fail and email the fork owner
    if: >-
      (github.repository_owner == 'pytorch' || github.event_name != 'schedule') &&
      (github.event_name != 'pull_request' || contains(github.event.pull_request.labels.*.name, 'ciflow/h100.8'))
    uses: ./.github/workflows/set-matrix.yaml
    with:
      runner-cuda: mt-l-bx86iamx-176-1800-h100-8
      is-experimental: true

  # Step 2: Use the dynamic matrix in the build-test job
  build-test:
    needs: set-matrix
    uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main
    strategy:
      fail-fast: false
      matrix: ${{ fromJSON(needs.set-matrix.outputs.matrix) }}
    with:
      runner: ${{ matrix.runner }}
      gpu-arch-type: ${{ matrix.gpu-arch-type }}
      gpu-arch-version: ${{ matrix.gpu-arch-version }}
      # Private-ECR image, tagged with the .ci/docker hash from set-matrix.yaml.
      docker-image: 308535385114.dkr.ecr.us-east-1.amazonaws.com/torchtitan/${{ matrix.docker-image }}:${{ needs.set-matrix.outputs.docker-hash }}
      repository: pytorch/torchtitan
      upload-artifact: outputs
      timeout: 45
      script: |
        set -eux

        # The generic Linux job chooses to use base env, not the one setup by the image
        CONDA_ENV=$(conda env list --json | jq -r ".envs | .[-1]")
        conda activate "${CONDA_ENV}"

        # The ARC H100 runner mounts the shared HF cache (/mnt/hf_cache) read-only.
        # The c4_test loader (torchtitan/hf_datasets/text_datasets.py) calls
        # datasets.load_dataset non-streaming, which os.makedirs the HF datasets
        # cache root before loading -- that write fails on the read-only mount and
        # crashes dataloader build for every debug-model flavor. Point the HF cache
        # at the writable per-job $RUNNER_TEMP so all test stages can build datasets.
        export HF_HOME="$RUNNER_TEMP/hf_home"
        export HF_DATASETS_CACHE="$RUNNER_TEMP/hf_home/datasets"

        # Log GPU info / driver version for debugging.
        DRIVER_VERSION=$(nvidia-smi --query-gpu=driver_version --format=csv,noheader | head -n 1 || true)
        echo "CUDA driver version: ${DRIVER_VERSION}"

        pip config --user set global.progress_bar off

        # Pre-install torch's pure-python deps from the in-cluster pypi-cache for speed.
        python -m pip install filelock typing-extensions "setuptools<82" sympy networkx jinja2 fsspec numpy
        # Clear PIP_EXTRA_INDEX_URL so the default cpu index can't supply a +cpu torch.
        python -m pip uninstall -y torch
        PIP_EXTRA_INDEX_URL= python -m pip install --pre \
          torch --index-url ${{ matrix.index-url }}
        python -m pip install git+https://github.com/meta-pytorch/autoparallel.git

        mkdir -p "$RUNNER_TEMP/artifacts-to-be-uploaded"

        # Install DeepEP for HybridEP integration test
        bash /install_deepep.sh

        # Disable Nvlink Sharp. The CI machine seems to be in unstable state to support
        # NVLS according to several CI runs.
        # DeepEP needs CUDA_HOME specified to JIT kernels.
        CUDA_HOME=/usr/local/cuda NCCL_NVLS_ENABLE=0 python -m torchtitan.experiments.graph_trainer.tests.integration_tests --test_suite graph_trainer_h100 --gpu_arch_type cuda $RUNNER_TEMP/artifacts-to-be-uploaded --ngpu 8

        # Run the MoE numerics tests
        NCCL_NVLS_ENABLE=0 pytest torchtitan/experiments/graph_trainer/tests/test_numerics.py::TestGraphTrainerNumerics -v -k "moe"

        # Run precompile integration tests (DSv3 with EP; Llama3 runs in default workflow)
        NCCL_NVLS_ENABLE=0 python -m torchtitan.experiments.graph_trainer.tests.run_precompile_tests $RUNNER_TEMP/artifacts-to-be-uploaded/precompile --ngpu 8 --test_name aot_fx_trace_deepseek_v3_precompile_fsdp_tp_ep

        # Run bitwise deterministic
        pytest torchtitan/experiments/graph_trainer/tests/test_bitwise_deterministic.py -v

        # AutoParallel depends on a separately installed package that can break
        # on PyTorch internal API movement. Keep it last so the graph_trainer
        # tests above still produce signal when that dependency is stale.
        NCCL_NVLS_ENABLE=0 python -m torchtitan.experiments.graph_trainer.tests.integration_tests --test_suite graph_trainer_autoparallel_h100 --gpu_arch_type cuda $RUNNER_TEMP/artifacts-to-be-uploaded/autoparallel --ngpu 4
        NCCL_NVLS_ENABLE=0 pytest torchtitan/experiments/graph_trainer/tests/test_numerics.py::TestGraphTrainerAutoParallelNumerics::test_deepseek_v3_aot_fx_trace_autoparallel_vs_eager -v

        rm -rf $RUNNER_TEMP/artifacts-to-be-uploaded/*/checkpoint
