name: TorchFT CPU Unit Test

on:
  push:
    branches: [ main ]
    paths:
      - 'torchtitan/experiments/torchft/**'
      - 'torchtitan/components/checkpointer/**'
      - '.github/workflows/unit_test_cpu_torchft.yaml'
  pull_request:
    paths:
      - 'torchtitan/experiments/torchft/**'
      - 'torchtitan/components/checkpointer/**'
      - '.github/workflows/unit_test_cpu_torchft.yaml'

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

permissions:
  id-token: write
  contents: read

jobs:
  set-matrix:
    uses: ./.github/workflows/set-matrix.yaml

  build-test:
    needs: set-matrix
    uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main
    with:
      runner: mt-l-x86iavx512-8-64
      docker-image: 308535385114.dkr.ecr.us-east-1.amazonaws.com/torchtitan/torchtitan-ubuntu-22.04-clang12:${{ needs.set-matrix.outputs.docker-hash }}
      repository: pytorch/torchtitan
      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}"

        pip config --user set global.progress_bar off

        python -m pip install torchft-nightly
        pip install --force-reinstall --pre torch --index-url https://download.pytorch.org/whl/nightly/cpu

        USE_CPP=0 python -m pip install --pre torchao --index-url https://download.pytorch.org/whl/nightly/cpu

        pytest torchtitan/experiments/torchft/tests/ -v
