name: GraphTrainer 8 GPU Integration Tests

on:
  push:
    branches: [ main ]
    paths:
      - 'torchtitan/experiments/graph_trainer/**'
      - '.github/workflows/integration_test_8gpu_graph_trainer.yaml'
  pull_request:
    types: [opened, synchronize, reopened, ready_for_review]
    paths:
      - 'torchtitan/experiments/graph_trainer/**'
      - '.github/workflows/integration_test_8gpu_graph_trainer.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:
  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' || github.event.pull_request.draft == false)
    uses: ./.github/workflows/set-matrix.yaml
    with:
      gpu-arch: cuda

  build-test:
    needs: set-matrix
    if: ${{ needs.set-matrix.outputs.matrix != '' && fromJSON(needs.set-matrix.outputs.matrix).include[0] != null }}
    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 }}
      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}"

        # 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

        # torch nightly now depends on cuda-toolkit[...] extras (e.g.
        # nvidia-curand==10.4.0.35) whose wheels live on PyPI, not on the pytorch
        # nightly index. --index-url *replaces* the default PyPI index, so add it
        # back via --extra-index-url or pip cannot resolve those cuda deps.
        python -m pip install --force-reinstall --pre \
          torch torchvision \
          --index-url https://download.pytorch.org/whl/nightly/cu132 \
          --extra-index-url https://pypi.org/simple
        python -m pip install git+https://github.com/meta-pytorch/autoparallel.git

        sudo mkdir -p "$RUNNER_TEMP/artifacts-to-be-uploaded"
        sudo chown -R $(id -u):$(id -g) "$RUNNER_TEMP/artifacts-to-be-uploaded"

        python -m torchtitan.experiments.graph_trainer.tests.integration_tests --test_suite graph_trainer_default --gpu_arch_type cuda $RUNNER_TEMP/artifacts-to-be-uploaded --ngpu 8


        # Run the numerics unit tests. Standard MoE tests run in the H100 workflow.
        pytest torchtitan/experiments/graph_trainer/tests/test_numerics.py::TestSimpleFSDP -v
        pytest torchtitan/experiments/graph_trainer/tests/test_numerics.py::TestGraphTrainerNumerics -v -k "dense"

        # Run the graph passes unit tests
        pytest torchtitan/experiments/graph_trainer/tests/test_passes.py -v

        # Run kernel annotation profiler tests (skips if cuda-compat < 13.1)
        pytest torchtitan/experiments/graph_trainer/tests/test_profiler.py -v

        # Run the make_fx tracer unit tests
        pytest torchtitan/experiments/graph_trainer/tests/test_trace_module.py -v

        # Run precompile unit tests
        pytest torchtitan/experiments/graph_trainer/tests/test_precompile.py -v

        # Run precompile integration tests (Llama3 only; DSv3 runs in H100 workflow)
        python -m torchtitan.experiments.graph_trainer.tests.run_precompile_tests $RUNNER_TEMP/artifacts-to-be-uploaded/precompile --ngpu 8 --test_name aot_fx_trace_llama3_precompile_fsdp_tp

        # Run bitwise deterministic and SAC peak-memory guardrail tests
        pytest torchtitan/experiments/graph_trainer/tests/test_bitwise_deterministic.py -v
        pytest torchtitan/experiments/graph_trainer/tests/test_sac_peak_memory.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.
        python -m torchtitan.experiments.graph_trainer.tests.integration_tests --test_suite graph_trainer_autoparallel --gpu_arch_type cuda $RUNNER_TEMP/artifacts-to-be-uploaded/autoparallel --ngpu 4
        pytest torchtitan/experiments/graph_trainer/tests/test_numerics.py::TestGraphTrainerAutoParallelNumerics::test_llama3_aot_fx_trace_autoparallel_vs_eager -v

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