From 15a448a9d3f3235513803fb3c1c991fe2c0eb94b Mon Sep 17 00:00:00 2001 From: Wing Lian Date: Sun, 4 Oct 2026 17:43:43 -0400 Subject: [PATCH] ci: sdist manifest check, GPU-only unit tests, deps-hash image cache, multigpu path filter, testmon arms (#4107) --- .github/CONTRIBUTING.md | 44 + .github/workflows/docker-e2e-jobs.yml | 132 +++ .github/workflows/gpu-label.yml | 14 +- .github/workflows/multi-gpu-e2e-jobs.yml | 72 ++ .github/workflows/testmon-baseline.yml | 83 ++ .github/workflows/testmon-cpu.yml | 122 +++ .github/workflows/tests-nightly.yml | 8 +- .github/workflows/tests.yml | 203 +++-- .gitignore | 2 + cicd/Dockerfile-uv.jinja | 8 +- cicd/cicd.sh | 20 +- cicd/cicd_cuda_kernels.sh | 8 +- cicd/cicd_selected.sh | 37 + cicd/cicd_testmon.sh | 61 ++ cicd/deps_hash.py | 126 +++ cicd/e2e_selected.py | 20 + cicd/e2e_testmon.py | 29 + cicd/multigpu.py | 63 ++ cicd/multigpu.sh | 12 +- cicd/multigpu_selected.sh | 32 + cicd/nf4.sh | 2 +- cicd/nvfp4.sh | 17 +- cicd/select_e2e_tests.py | 852 ++++++++++++++++++ cicd/single_gpu.py | 62 ++ pyproject.toml | 2 + scripts/check_installed_axolotl.py | 109 +++ scripts/check_sdist_manifest.py | 146 +++ src/axolotl/integrations/liger/models/base.py | 2 +- .../kimi_linear/tokenization_kimi.py | 2 +- .../monkeypatch/attention/flex_attn.py | 7 +- .../gradient_checkpointing/offload_cpu.py | 2 + .../gradient_checkpointing/offload_disk.py | 2 + src/axolotl/utils/schemas/validation.py | 55 +- tests/conftest.py | 71 ++ ...native_nvfp4_deepspeed_lora_merge_aware.py | 2 + ...ve_nvfp4_deepspeed_lora_merge_lifecycle.py | 2 + .../solo/test_native_nvfp4_dynamic_ste.py | 2 + .../test_native_nvfp4_fsdp2_lora_parity.py | 2 + ...t_native_nvfp4_fsdp2_merge_aware_resume.py | 2 + .../solo/test_native_nvfp4_fsdp2_recipe.py | 2 + .../solo/test_native_nvfp4_tp_hf_load.py | 2 + .../solo/test_native_nvfp4_tp_lora_parity.py | 2 + .../solo/test_sonicmoe_nvfp4_fsdp2.py | 2 + .../solo/test_torchao_lora_deepspeed.py | 2 + .../test_torchao_lora_deepspeed_checkpoint.py | 2 + ...torchao_lora_deepspeed_zero3_checkpoint.py | 2 + tests/e2e/utils.py | 19 +- .../kernels/test_merge_aware_linear.py | 1 + .../test_expert_parallel_adapter_resume.py | 2 + tests/integrations/test_gemma4_moe.py | 2 + tests/integrations/test_sonicmoe.py | 1 + .../test_sonicmoe_fp4_cute_prep.py | 1 + tests/kernels/test_rms_norm_gated.py | 4 +- .../loaders/test_nf4_adapter_init_dequant.py | 2 + .../test_checkpoint_activation_offload.py | 2 + tests/monkeypatch/test_lora_fp32_gradients.py | 1 + tests/monkeypatch/test_lora_kernels_nvfp4.py | 1 + .../test_offload_amp_device_type.py | 2 + tests/monkeypatch/test_sdpa_varlen.py | 9 +- tests/scripts/__init__.py | 0 tests/scripts/test_check_sdist_manifest.py | 64 ++ tests/scripts/test_deps_hash.py | 135 +++ tests/scripts/test_select_e2e_tests.py | 404 +++++++++ tests/test_gpu_marker.py | 51 ++ tests/test_lora.py | 3 + .../validation/test_config_validators.py | 28 + .../test_flash_attn_availability.py | 116 ++- .../schemas/validation/test_moe_quant.py | 1 + 68 files changed, 3184 insertions(+), 116 deletions(-) create mode 100644 .github/workflows/testmon-baseline.yml create mode 100644 .github/workflows/testmon-cpu.yml create mode 100755 cicd/cicd_selected.sh create mode 100644 cicd/cicd_testmon.sh create mode 100644 cicd/deps_hash.py create mode 100644 cicd/e2e_selected.py create mode 100644 cicd/e2e_testmon.py create mode 100755 cicd/multigpu_selected.sh create mode 100644 cicd/select_e2e_tests.py create mode 100644 scripts/check_installed_axolotl.py create mode 100644 scripts/check_sdist_manifest.py create mode 100644 tests/scripts/__init__.py create mode 100644 tests/scripts/test_check_sdist_manifest.py create mode 100644 tests/scripts/test_deps_hash.py create mode 100644 tests/scripts/test_select_e2e_tests.py create mode 100644 tests/test_gpu_marker.py diff --git a/.github/CONTRIBUTING.md b/.github/CONTRIBUTING.md index 60da5223b..88b4e7e98 100644 --- a/.github/CONTRIBUTING.md +++ b/.github/CONTRIBUTING.md @@ -141,6 +141,50 @@ The label handler uses `pull_request_target` with trusted inline code only and m Outside of PRs, the `docker-e2e-tests` suite runs on merges to `main`, and the multi-GPU suite runs on its semi-weekly schedule or manual dispatch. +##### Selected e2e arms (experimental) + +Alongside the full suites, every PR run selects the e2e test files its diff can affect +and runs just those in two extra arms, `docker-e2e-tests-selected` (single GPU) and +`multigpu-selected`. They gate nothing today; they exist to be compared against the +full suites. The `select-e2e` job writes the selection and the reason for every file +to its job summary and uploads it as the `e2e-selection` artifact. + +`cicd/select_e2e_tests.py` derives the selection from the tree, so new features need +no registration: + +- A changed module selects the tests whose config sets a key it reads, names it as a + config value (a `plugins` entry, a dataset `type`, an `rl` method, a `base_model` + containing a model-support package name), or imports it directly. +- Plugins, prompt strategies and other packages reached by config value are + "registry" entries. A member no test names selects nothing. +- A module that reads a key most e2e configs set (`base_model`, `sequence_len`, + `micro_batch_size`, ...) is core and runs the whole scope, as does a module with no + derivable edges, a deleted module, a console-script entry point, `pyproject.toml`, + anything under `cicd/` or `.github/workflows/`, a `conftest.py`, and any error + inside the selector. The selector can over-select, never silently under-select. +- Worker scripts and helpers next to a test (`_*_worker.py`, parity probes) count as + part of the tests that name or import them. A test with no visible config at all + rides along with every subset. + +When a module is enabled by one config key but reached through a hub that reads +ubiquitous keys, declare the key so the selector can narrow instead of running +everything: + +```python +__ci_config_keys__ = ("activation_offloading",) +``` + +The selector then picks exactly the tests that set any declared key. A declared +name that is not a config field runs the whole scope, so typos fail safe. + +Put `[test all]` in a commit message to force the full scope. Preview locally with: + +```bash +python cicd/select_e2e_tests.py --base origin/main --explain \ + --exclude tests/e2e/multigpu --exclude tests/e2e/kernels # single-GPU scope +python cicd/select_e2e_tests.py --base origin/main --explain --scope tests/e2e/multigpu +``` + CPU NF4 tests marked `nf4_distributed` run in a dedicated job with its own timeout, across the same PyTorch versions as the main CPU matrix. The source, sdist, nightly, and single-GPU general suites exclude this marker while keeping the quick NF4 tests. diff --git a/.github/workflows/docker-e2e-jobs.yml b/.github/workflows/docker-e2e-jobs.yml index 283ae4243..3517d8b65 100644 --- a/.github/workflows/docker-e2e-jobs.yml +++ b/.github/workflows/docker-e2e-jobs.yml @@ -87,6 +87,8 @@ jobs: echo "MODAL_IMAGE_BUILDER_VERSION=2024.10" >> $GITHUB_ENV echo "N_GPUS=${{ matrix.num_gpus }}" >> $GITHUB_ENV echo "E2E_DOCKERFILE=${{ matrix.dockerfile || 'Dockerfile-uv.jinja'}}" >> $GITHUB_ENV + - name: Compute deps hash + run: python cicd/deps_hash.py - name: Run tests job on Modal env: CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }} @@ -138,6 +140,8 @@ jobs: echo "N_GPUS=${{ matrix.num_gpus }}" >> $GITHUB_ENV echo "GPU_TYPE=${{ matrix.gpu_type || 'L40S'}}" >> $GITHUB_ENV echo "E2E_DOCKERFILE=${{ matrix.dockerfile || 'Dockerfile-uv.jinja'}}" >> $GITHUB_ENV + - name: Compute deps hash + run: python cicd/deps_hash.py - name: Run tests job on Modal env: CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }} @@ -192,12 +196,138 @@ jobs: echo "N_GPUS=${{ matrix.num_gpus }}" >> $GITHUB_ENV echo "GPU_TYPE=${{ matrix.gpu_type || 'L40S'}}" >> $GITHUB_ENV echo "E2E_DOCKERFILE=${{ matrix.dockerfile || 'Dockerfile-uv.jinja'}}" >> $GITHUB_ENV + - name: Compute deps hash + run: python cicd/deps_hash.py - name: Run tests job on Modal env: CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }} run: | modal run -m cicd.e2e_cuda_kernels + docker-e2e-tests-testmon: + if: > + github.repository_owner == 'axolotl-ai-cloud' && + needs.gate-skip-e2e.outputs.skip != 'true' + # this job needs to be run on self-hosted GPU runners... + runs-on: [self-hosted, modal] + timeout-minutes: 120 + # Only run the remainder of the matrix if the first e2e check passed; + # this is to save on wasted compute costs for known failures that get caught in the first run + needs: [gate-skip-e2e, docker-e2e-tests-1st] + + strategy: + fail-fast: false + matrix: + include: + - cuda: 130 + cuda_version: 13.0.0 + python_version: "3.12" + pytorch: 2.14.0 + num_gpus: 1 + axolotl_extras: + steps: + - name: Checkout + uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + with: + persist-credentials: false + - name: Install Python + uses: actions/setup-python@5fda3b95a4ea91299a34e894583c3862153e4b97 # v7.0.0 + with: + python-version: "3.12" + - name: Install Modal + run: | + python -m pip install --upgrade pip + pip install modal==1.3.0.post1 jinja2 + - name: Update env vars + run: | + echo "BASE_TAG=main-base-py${{ matrix.python_version }}-cu${{ matrix.cuda }}-${{ matrix.pytorch }}" >> $GITHUB_ENV + echo "PYTORCH_VERSION=${{ matrix.pytorch}}" >> $GITHUB_ENV + echo "AXOLOTL_ARGS=${{ matrix.axolotl_args}}" >> $GITHUB_ENV + echo "AXOLOTL_EXTRAS=${{ matrix.axolotl_extras}}" >> $GITHUB_ENV + echo "CUDA=${{ matrix.cuda }}" >> $GITHUB_ENV + echo "MODAL_IMAGE_BUILDER_VERSION=2024.10" >> $GITHUB_ENV + echo "N_GPUS=${{ matrix.num_gpus }}" >> $GITHUB_ENV + echo "GPU_TYPE=${{ matrix.gpu_type || 'L40S'}}" >> $GITHUB_ENV + echo "E2E_DOCKERFILE=${{ matrix.dockerfile || 'Dockerfile-uv.jinja'}}" >> $GITHUB_ENV + - name: Compute deps hash + run: python cicd/deps_hash.py + - name: Run testmon e2e job on Modal + env: + CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }} + run: | + modal run -m cicd.e2e_testmon + + docker-e2e-tests-selected: + if: > + github.repository_owner == 'axolotl-ai-cloud' && + github.event_name == 'pull_request' && + needs.gate-skip-e2e.outputs.skip != 'true' + runs-on: [self-hosted, modal] + timeout-minutes: 120 + needs: [gate-skip-e2e] + strategy: + fail-fast: false + matrix: + include: + - cuda: 130 + cuda_version: 13.0.0 + python_version: "3.12" + pytorch: 2.14.0 + axolotl_extras: + num_gpus: 1 + steps: + - name: Download e2e selection + uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1 + continue-on-error: true + with: + name: e2e-selection + path: ${{ runner.temp }}/e2e-selection + - name: Read selection + id: selection + run: | + dir="$RUNNER_TEMP/e2e-selection" + mode=$(cat "$dir/single.mode" 2>/dev/null || echo all) + echo "mode=$mode" >> "$GITHUB_OUTPUT" + echo "selected-e2e (single): mode=$mode" + if [ "$mode" = "subset" ]; then + echo "E2E_SELECTED_TESTS=$(paste -sd, "$dir/single.txt")" >> "$GITHUB_ENV" + cat "$dir/single.txt" + fi + - name: Checkout + if: steps.selection.outputs.mode == 'subset' + uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + with: + persist-credentials: false + - name: Install Python + if: steps.selection.outputs.mode == 'subset' + uses: actions/setup-python@5fda3b95a4ea91299a34e894583c3862153e4b97 # v7.0.0 + with: + python-version: "3.12" + - name: Install Modal + if: steps.selection.outputs.mode == 'subset' + run: | + python -m pip install --upgrade pip + pip install modal==1.3.0.post1 jinja2 + - name: Update env vars + if: steps.selection.outputs.mode == 'subset' + run: | + echo "BASE_TAG=main-base-py${{ matrix.python_version }}-cu${{ matrix.cuda }}-${{ matrix.pytorch }}" >> $GITHUB_ENV + echo "PYTORCH_VERSION=${{ matrix.pytorch}}" >> $GITHUB_ENV + echo "AXOLOTL_ARGS=${{ matrix.axolotl_args}}" >> $GITHUB_ENV + echo "AXOLOTL_EXTRAS=${{ matrix.axolotl_extras}}" >> $GITHUB_ENV + echo "CUDA=${{ matrix.cuda }}" >> $GITHUB_ENV + echo "N_GPUS=${{ matrix.num_gpus }}" >> $GITHUB_ENV + echo "GPU_TYPE=${{ matrix.gpu_type || 'L40S' }}" >> $GITHUB_ENV + echo "E2E_DOCKERFILE=${{ matrix.dockerfile || 'Dockerfile-uv.jinja'}}" >> $GITHUB_ENV + echo "MODAL_IMAGE_BUILDER_VERSION=2024.10" >> $GITHUB_ENV + - name: Compute deps hash + if: steps.selection.outputs.mode == 'subset' + run: python cicd/deps_hash.py + - name: Run selected tests on Modal + if: steps.selection.outputs.mode == 'subset' + run: | + modal run -m cicd.e2e_selected + docker-e2e-cleanup: runs-on: [self-hosted, modal] timeout-minutes: 90 @@ -236,6 +366,8 @@ jobs: echo "CUDA=${{ matrix.cuda }}" >> $GITHUB_ENV echo "MODAL_IMAGE_BUILDER_VERSION=2024.10" >> $GITHUB_ENV echo "N_GPUS=${{ matrix.num_gpus }}" >> $GITHUB_ENV + - name: Compute deps hash + run: python cicd/deps_hash.py - name: Run tests job on Modal run: | modal run -m cicd.cleanup diff --git a/.github/workflows/gpu-label.yml b/.github/workflows/gpu-label.yml index 9c5df3f52..ba5401d1d 100644 --- a/.github/workflows/gpu-label.yml +++ b/.github/workflows/gpu-label.yml @@ -42,14 +42,22 @@ jobs: continue; } - // Wait for CPU tests to finish, but leave a running GPU rerun alone. + // The gate job runs after the CPU jobs and re-reads labels, so an + // in-flight run that has not evaluated it yet will see this label. let status = run.status; for (let attempt = 0; status !== 'completed' && attempt < 360; attempt++) { const jobs = await github.paginate(github.rest.actions.listJobsForWorkflowRun, { ...context.repo, run_id: run.id, per_page: 100, }); - if (jobs.some(job => job.name.startsWith('docker-e2e /') || - job.name.startsWith('multi-gpu-e2e /'))) break; + const gate = jobs.find(job => job.name.startsWith('check-gpu-label /')); + const gateDone = gate !== undefined && gate.status === 'completed'; + if (!gate || (!gateDone && gate.status !== 'in_progress')) { + core.info(`${workflow}: run ${run.id} has not evaluated the GPU label gate yet; it will pick up the label.`); + return; + } + // A gate already in progress may have read the labels before this one was added. + if (gateDone && jobs.some(job => job.name.startsWith('docker-e2e /') || + job.name.startsWith('multi-gpu-e2e /'))) break; await new Promise(resolve => setTimeout(resolve, 10000)); const { data: current } = await github.rest.actions.getWorkflowRun({ ...context.repo, run_id: run.id, diff --git a/.github/workflows/multi-gpu-e2e-jobs.yml b/.github/workflows/multi-gpu-e2e-jobs.yml index 595140486..33d8a7f5b 100644 --- a/.github/workflows/multi-gpu-e2e-jobs.yml +++ b/.github/workflows/multi-gpu-e2e-jobs.yml @@ -77,8 +77,80 @@ jobs: echo "MULTIGPU_TEST_SUITE=${{ matrix.test_suite }}" >> $GITHUB_ENV echo "NF4_TEST_STACK=${{ matrix.nf4_test_stack }}" >> $GITHUB_ENV echo "E2E_DOCKERFILE=${{ matrix.dockerfile || 'Dockerfile-uv.jinja'}}" >> $GITHUB_ENV + - name: Compute deps hash + run: python cicd/deps_hash.py - name: Run tests job on Modal env: CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }} run: | modal run -m cicd.multigpu + + multigpu-selected: + if: > + github.repository_owner == 'axolotl-ai-cloud' && + github.event_name == 'pull_request' + runs-on: [self-hosted, modal] + timeout-minutes: 120 + strategy: + fail-fast: false + matrix: + include: + - cuda: 130 + cuda_version: 13.0.0 + python_version: "3.12" + pytorch: 2.14.0 + axolotl_extras: + num_gpus: 2 + steps: + - name: Download e2e selection + uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1 + continue-on-error: true + with: + name: e2e-selection + path: ${{ runner.temp }}/e2e-selection + - name: Read selection + id: selection + run: | + dir="$RUNNER_TEMP/e2e-selection" + mode=$(cat "$dir/multigpu.mode" 2>/dev/null || echo all) + echo "mode=$mode" >> "$GITHUB_OUTPUT" + echo "selected-e2e (multigpu): mode=$mode" + if [ "$mode" = "subset" ]; then + echo "E2E_SELECTED_TESTS=$(paste -sd, "$dir/multigpu.txt")" >> "$GITHUB_ENV" + cat "$dir/multigpu.txt" + fi + - name: Checkout + if: steps.selection.outputs.mode == 'subset' + uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + with: + persist-credentials: false + - name: Install Python + if: steps.selection.outputs.mode == 'subset' + uses: actions/setup-python@5fda3b95a4ea91299a34e894583c3862153e4b97 # v7.0.0 + with: + python-version: "3.12" + - name: Install Modal + if: steps.selection.outputs.mode == 'subset' + run: | + python -m pip install --upgrade pip + pip install modal==1.3.0.post1 jinja2 + - name: Update env vars + if: steps.selection.outputs.mode == 'subset' + run: | + echo "BASE_TAG=main-base-py${{ matrix.python_version }}-cu${{ matrix.cuda }}-${{ matrix.pytorch }}" >> $GITHUB_ENV + echo "PYTORCH_VERSION=${{ matrix.pytorch}}" >> $GITHUB_ENV + echo "AXOLOTL_ARGS=${{ matrix.axolotl_args}}" >> $GITHUB_ENV + echo "AXOLOTL_EXTRAS=${{ matrix.axolotl_extras}}" >> $GITHUB_ENV + echo "CUDA=${{ matrix.cuda }}" >> $GITHUB_ENV + echo "N_GPUS=${{ matrix.num_gpus }}" >> $GITHUB_ENV + echo "GPU_TYPE=${{ matrix.gpu_type || 'H100' }}" >> $GITHUB_ENV + echo "E2E_DOCKERFILE=${{ matrix.dockerfile || 'Dockerfile-uv.jinja'}}" >> $GITHUB_ENV + echo "MULTIGPU_TEST_SUITE=multigpu_selected" >> $GITHUB_ENV + echo "NF4_TEST_STACK=" >> $GITHUB_ENV + - name: Compute deps hash + if: steps.selection.outputs.mode == 'subset' + run: python cicd/deps_hash.py + - name: Run selected tests on Modal + if: steps.selection.outputs.mode == 'subset' + run: | + modal run -m cicd.multigpu diff --git a/.github/workflows/testmon-baseline.yml b/.github/workflows/testmon-baseline.yml new file mode 100644 index 000000000..a8113f95a --- /dev/null +++ b/.github/workflows/testmon-baseline.yml @@ -0,0 +1,83 @@ +name: testmon-baseline + +# Main commits mostly land with [skip ci], so the baselines are refreshed here instead. +on: + schedule: + - cron: "0 3 * * *" + workflow_dispatch: + +permissions: + contents: read + +jobs: + main-changed: + runs-on: ubuntu-latest + outputs: + run: ${{ steps.check.outputs.run }} + steps: + - uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9.0.0 + id: check + with: + script: | + if (context.eventName === 'workflow_dispatch') { + core.setOutput('run', 'true'); + return; + } + const { data: commit } = await github.rest.repos.getCommit({ ...context.repo, ref: 'main' }); + const committed = new Date(commit.commit.committer.date); + const ageHours = (Date.now() - committed.getTime()) / 3600000; + core.info(`main head ${commit.sha} committed ${ageHours.toFixed(1)}h ago`); + core.setOutput('run', String(ageHours < 24)); + + cpu: + needs: main-changed + if: needs.main-changed.outputs.run == 'true' + uses: ./.github/workflows/testmon-cpu.yml + + single-gpu: + needs: main-changed + if: github.repository_owner == 'axolotl-ai-cloud' && needs.main-changed.outputs.run == 'true' + runs-on: [self-hosted, modal] + timeout-minutes: 120 + env: + TRANSFORMERS_IS_CI: "yes" + UV_SYSTEM_PYTHON: "1" + strategy: + fail-fast: false + matrix: + include: + - cuda: 130 + cuda_version: 13.0.0 + python_version: "3.12" + pytorch: 2.14.0 + num_gpus: 1 + axolotl_extras: + steps: + - name: Checkout + uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + with: + persist-credentials: false + - name: Install Python + uses: actions/setup-python@5fda3b95a4ea91299a34e894583c3862153e4b97 # v7.0.0 + with: + python-version: "3.12" + - name: Install Modal + run: | + python -m pip install --upgrade pip + pip install modal==1.3.0.post1 jinja2 + - name: Update env vars + run: | + echo "BASE_TAG=main-base-py${{ matrix.python_version }}-cu${{ matrix.cuda }}-${{ matrix.pytorch }}" >> $GITHUB_ENV + echo "PYTORCH_VERSION=${{ matrix.pytorch}}" >> $GITHUB_ENV + echo "AXOLOTL_ARGS=${{ matrix.axolotl_args}}" >> $GITHUB_ENV + echo "AXOLOTL_EXTRAS=${{ matrix.axolotl_extras}}" >> $GITHUB_ENV + echo "CUDA=${{ matrix.cuda }}" >> $GITHUB_ENV + echo "MODAL_IMAGE_BUILDER_VERSION=2024.10" >> $GITHUB_ENV + echo "N_GPUS=${{ matrix.num_gpus }}" >> $GITHUB_ENV + echo "GPU_TYPE=${{ matrix.gpu_type || 'L40S'}}" >> $GITHUB_ENV + echo "E2E_DOCKERFILE=${{ matrix.dockerfile || 'Dockerfile-uv.jinja'}}" >> $GITHUB_ENV + - name: Compute deps hash + run: python cicd/deps_hash.py + - name: Run testmon e2e job on Modal + run: | + modal run -m cicd.e2e_testmon diff --git a/.github/workflows/testmon-cpu.yml b/.github/workflows/testmon-cpu.yml new file mode 100644 index 000000000..688f6b88c --- /dev/null +++ b/.github/workflows/testmon-cpu.yml @@ -0,0 +1,122 @@ +name: testmon-cpu + +on: + workflow_call: + +permissions: + contents: read + +env: + TRANSFORMERS_IS_CI: "yes" + UV_SYSTEM_PYTHON: "1" + PYTHON_VERSION: "3.12" + PYTORCH_VERSION: "2.14.0" + +jobs: + pytest-testmon: + name: PyTest testmon (py3.12, torch 2.14.0) + runs-on: ubuntu-latest + timeout-minutes: 45 + + steps: + - name: cleanup node + run: | + sudo rm -rf /usr/share/dotnet /usr/local/lib/android /opt/ghc /opt/hostedtoolcache/CodeQL + + - name: Check out repository code + uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + with: + persist-credentials: false + + - name: Restore Cache from S3 + id: hf-cache-restore-s3 + run: | + mkdir -p ~/.cache/huggingface/hub + curl -L https://axolotl-ci.b-cdn.net/hf-cache.tar.zst | tar -xpf - -C ~/.cache/huggingface/hub/ --use-compress-program unzstd --strip-components=1 + ls -ltr ~/.cache/huggingface/hub/ + + - name: Setup Python + uses: actions/setup-python@5fda3b95a4ea91299a34e894583c3862153e4b97 # v7.0.0 + with: + python-version: ${{ env.PYTHON_VERSION }} + + - name: Install uv + uses: astral-sh/setup-uv@bec219d24cd3e171d82865faccec33120bb574f4 # v10.1.0 + with: + enable-cache: true + prune-cache: true + cache-dependency-glob: "pyproject.toml" + + - name: Install PyTorch + run: | + uv pip install torch==${PYTORCH_VERSION} torchvision + uv pip freeze | grep -E "^(torch|torchvision)==" > /tmp/torch-pin.txt + + - name: Install dependencies + run: | + uv pip install --no-build-isolation -e '.[modal,baseten,ringmaster]' --override /tmp/torch-pin.txt + python scripts/cutcrossentropy_install.py --uv | sh + uv pip install black mypy pre-commit types-requests quartodoc jupyter blobfile tiktoken \ + codecov-cli pytest pytest-retry pytest-sugar pytest-testmon pytest-xdist tbparse + + - name: Make sure PyTorch version wasn't clobbered + run: | + python -c "import torch; assert '${PYTORCH_VERSION}' in torch.__version__, f'Expected torch ${PYTORCH_VERSION} but got {torch.__version__}'" + + - name: Ensure axolotl CLI was installed + run: | + axolotl --help + + - name: Pre-Download dataset fixture + run: | + hf download --repo-type=dataset axolotl-ai-internal/axolotl-oss-dataset-fixtures + + - name: Show HF cache + run: hf cache ls + + # testmon stores the package set and Python version itself and starts over when they change, + # so one baseline per arm is enough; the key only needs the commit. + - name: Restore testmon baseline + uses: actions/cache/restore@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0 + with: + path: .testmon + key: testmon-cpu-${{ github.sha }} + restore-keys: | + testmon-cpu- + + - name: Run tests with testmon selection (xdist -n4 kept, pytest-testmon 2.2.0 syncs workers) + timeout-minutes: 30 + shell: bash + run: | + mkdir -p .testmon + tm() { + local name=$1; shift + TESTMON_DATAFILE=".testmon/${name}.testmondata" pytest --testmon-forceselect "$@" | tee -a testmon-run.log + } + tm core -v --durations=10 -n4 --dist loadfile --ignore=tests/utils/ --ignore=tests/integrations/ --ignore=tests/e2e/ --ignore=tests/patched/ --ignore=tests/cli/ --ignore=tests/monkeypatch/ tests/ + tm monkeypatch -v --durations=10 -m "not slow and not nf4_distributed" tests/monkeypatch/ + tm patched -v --durations=10 tests/patched/ + tm cli -v --durations=10 tests/cli/ + tm utils -v --durations=10 tests/utils/ + tm integrations -v --durations=10 \ + --ignore=tests/integrations/kernels/ \ + --ignore=tests/integrations/monkeypatch/test_tiled_mlp_moe.py \ + --ignore=tests/integrations/test_gemma4_moe.py \ + --ignore=tests/integrations/test_scattermoe_lora.py \ + --ignore=tests/integrations/test_scattermoe_lora_kernels.py \ + --ignore=tests/integrations/test_scattermoe_multi_lora.py \ + --ignore=tests/integrations/test_sonicmoe_multi_lora.py \ + tests/integrations/ + + - name: testmon selection summary + if: always() + run: | + grep -hE "^testmon:|estmon savings|deselected" testmon-run.log || true + + # Main is the only writer: PR runs read the baseline and discard their own datafile. + - name: Save testmon baseline + if: github.ref == 'refs/heads/main' && github.event_name != 'pull_request' + uses: actions/cache/save@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0 + with: + path: .testmon + key: testmon-cpu-${{ github.sha }} diff --git a/.github/workflows/tests-nightly.yml b/.github/workflows/tests-nightly.yml index e29c4d0ae..6d3ad381b 100644 --- a/.github/workflows/tests-nightly.yml +++ b/.github/workflows/tests-nightly.yml @@ -103,7 +103,7 @@ jobs: uv pip install --no-build-isolation -e . --override /tmp/torch-pin.txt python scripts/cutcrossentropy_install.py --uv | sh uv pip install black mypy pre-commit types-requests quartodoc jupyter blobfile tiktoken \ - codecov codecov-cli pytest pytest-cov pytest-retry pytest-sugar pytest-xdist tbparse + codecov-cli pytest pytest-cov pytest-retry pytest-sugar pytest-xdist tbparse - name: Override with nightly HF packages run: | @@ -197,6 +197,8 @@ jobs: echo "N_GPUS=${{ matrix.num_gpus }}" >> $GITHUB_ENV echo "E2E_DOCKERFILE=${{ matrix.dockerfile || 'Dockerfile-uv.jinja'}}" >> $GITHUB_ENV echo "NIGHTLY_BUILD=${{ matrix.nightly_build }}" >> $GITHUB_ENV + - name: Compute deps hash + run: python cicd/deps_hash.py - name: Run tests job on Modal env: CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }} @@ -253,6 +255,8 @@ jobs: echo "NF4_TEST_STACK=${{ matrix.nf4_test_stack }}" >> $GITHUB_ENV echo "NIGHTLY_BUILD=${{ matrix.nightly_build }}" >> $GITHUB_ENV echo "E2E_DOCKERFILE=${{ matrix.dockerfile || 'Dockerfile-uv.jinja'}}" >> $GITHUB_ENV + - name: Compute deps hash + run: python cicd/deps_hash.py - name: Run tests job on Modal env: CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }} @@ -307,6 +311,8 @@ jobs: echo "N_GPUS=${{ matrix.num_gpus }}" >> $GITHUB_ENV echo "E2E_DOCKERFILE=${{ matrix.dockerfile || 'Dockerfile-uv.jinja'}}" >> $GITHUB_ENV echo "NIGHTLY_BUILD=${{ matrix.nightly_build }}" >> $GITHUB_ENV + - name: Compute deps hash + run: python cicd/deps_hash.py - name: Run tests job on Modal env: CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }} diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 98e142821..9b17b9941 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -9,10 +9,8 @@ on: - "**.py" - "pyproject.toml" - ".github/workflows/*.yml" - - "cicd/cicd.sh" - - "cicd/cicd_cuda_kernels.sh" - - "cicd/prepare_hf_cache.sh" - - "cicd/Dockerfile-uv.jinja" + - "cicd/*.sh" + - "cicd/*.jinja" pull_request: # GPU tests are label-gated, so adding the label re-runs this workflow via # gpu-label.yml after the initial CPU run has completed. @@ -21,10 +19,8 @@ on: - "**.py" - "pyproject.toml" - ".github/workflows/*.yml" - - "cicd/cicd.sh" - - "cicd/cicd_cuda_kernels.sh" - - "cicd/prepare_hf_cache.sh" - - "cicd/Dockerfile-uv.jinja" + - "cicd/*.sh" + - "cicd/*.jinja" workflow_dispatch: # Cancel jobs on the same ref if a new one is triggered @@ -90,6 +86,8 @@ jobs: python_version: ["3.12", "3.14"] pytorch_version: ["2.13.0", "2.14.0"] timeout-minutes: 45 + env: + AXOLOTL_CI_ENFORCE_GPU_MARKER: "1" steps: - name: cleanup node @@ -130,7 +128,7 @@ jobs: uv pip install --no-build-isolation -e '.[modal,baseten,ringmaster]' --override /tmp/torch-pin.txt python scripts/cutcrossentropy_install.py --uv | sh uv pip install black mypy pre-commit types-requests quartodoc jupyter blobfile tiktoken \ - codecov codecov-cli pytest pytest-cov pytest-retry pytest-sugar pytest-xdist tbparse + codecov-cli pytest pytest-cov pytest-retry pytest-sugar pytest-xdist tbparse - name: Make sure PyTorch version wasn't clobbered run: | @@ -181,18 +179,19 @@ jobs: flags: unittests,pytorch-${{ matrix.pytorch_version }} fail_ci_if_error: false - pytest-sdist: - name: PyTest SDist + pytest-testmon: + if: ${{ !github.event.pull_request.draft }} + needs: [prime-cdn-s3-cache, pre-commit] + uses: ./.github/workflows/testmon-cpu.yml + + sdist: + name: SDist manifest and install runs-on: ubuntu-latest if: ${{ !github.event.pull_request.draft }} needs: [prime-cdn-s3-cache, pre-commit] - strategy: - max-parallel: 2 - fail-fast: false - matrix: - python_version: ["3.12", "3.14"] - pytorch_version: ["2.13.0", "2.14.0"] timeout-minutes: 45 + env: + PYTORCH_VERSION: "2.14.0" steps: - name: cleanup node @@ -214,7 +213,7 @@ jobs: - name: Setup Python uses: actions/setup-python@5fda3b95a4ea91299a34e894583c3862153e4b97 # v7.0.0 with: - python-version: ${{ matrix.python_version }} + python-version: "3.12" - name: Install uv uses: astral-sh/setup-uv@bec219d24cd3e171d82865faccec33120bb574f4 # v10.1.0 @@ -223,34 +222,31 @@ jobs: prune-cache: true cache-dependency-glob: "pyproject.toml" + - name: Check sdist manifest + run: | + python scripts/check_sdist_manifest.py --out-dir dist + - name: Install PyTorch run: | - uv pip install torch==${{ matrix.pytorch_version }} torchvision + uv pip install torch==${PYTORCH_VERSION} torchvision uv pip freeze | grep -E "^(torch|torchvision)==" > /tmp/torch-pin.txt - - name: Install dependencies + - name: Install axolotl from the sdist-built wheel run: | - uv pip install packaging setuptools_scm build wheel psutil - python -m build --no-isolation --sdist - for sdist in dist/axolotl*.tar.gz; do - uv pip install --no-build-isolation "${sdist}[modal,baseten,ringmaster]" --override /tmp/torch-pin.txt - done + wheel=$(ls dist/axolotl-*.whl) + uv pip install --no-build-isolation "${wheel}[modal,baseten,ringmaster]" --override /tmp/torch-pin.txt python scripts/cutcrossentropy_install.py --uv | sh - uv pip install black mypy pre-commit types-requests quartodoc jupyter blobfile tiktoken \ - codecov codecov-cli pytest pytest-cov pytest-retry pytest-sugar pytest-xdist tbparse + uv pip install pytest pytest-retry pytest-sugar pytest-xdist tbparse - name: Make sure PyTorch version wasn't clobbered run: | - python -c "import torch; assert '${{ matrix.pytorch_version }}' in torch.__version__, f'Expected torch ${{ matrix.pytorch_version }} but got {torch.__version__}'" + python -c "import torch; assert '${PYTORCH_VERSION}' in torch.__version__, f'Expected torch ${PYTORCH_VERSION} but got {torch.__version__}'" - name: Ensure axolotl CLI was installed - run: | - axolotl --help + run: axolotl --help - name: Check generated CLI config options - if: ${{ matrix.python_version == '3.12' && matrix.pytorch_version == '2.14.0' }} - run: | - axolotl generate-cli-config-options --check + run: axolotl generate-cli-config-options --check - name: Verify agent docs are discoverable run: | @@ -262,41 +258,134 @@ jobs: axolotl agent-docs sft | grep -q "SFT" python -c "from axolotl.cli.agent_docs import get_doc, list_topics; assert len(list_topics()) >= 5; assert 'GRPO' in get_doc('grpo')" - - name: Show HF cache - run: hf cache ls - - - name: Run tests - timeout-minutes: 30 + - name: Import every axolotl module and resolve entry points + env: + AXOLOTL_DO_NOT_TRACK: "1" run: | - pytest -v --durations=10 -n4 --dist loadfile --ignore=tests/utils/ --ignore=tests/integrations/ --ignore=tests/e2e/ --ignore=tests/patched/ --ignore=tests/cli/ --ignore=tests/monkeypatch/ tests/ --cov=axolotl --cov-report=xml - pytest -v --durations=10 -m "not slow and not nf4_distributed" tests/monkeypatch/ --cov=axolotl --cov-append --cov-report=xml - pytest -v --durations=10 tests/patched/ --cov=axolotl --cov-append --cov-report=xml - pytest -v --durations=10 tests/cli/ --cov=axolotl --cov-append --cov-report=xml - pytest -v --durations=10 tests/utils/ --cov=axolotl --cov-append --cov-report=xml - pytest -v --durations=10 \ - --ignore=tests/integrations/kernels/ \ - --ignore=tests/integrations/monkeypatch/test_tiled_mlp_moe.py \ - --ignore=tests/integrations/test_gemma4_moe.py \ - --ignore=tests/integrations/test_scattermoe_lora.py \ - --ignore=tests/integrations/test_scattermoe_lora_kernels.py \ - --ignore=tests/integrations/test_scattermoe_multi_lora.py \ - --ignore=tests/integrations/test_sonicmoe_multi_lora.py \ - tests/integrations/ --cov=axolotl --cov-append --cov-report=xml + python scripts/check_installed_axolotl.py --import-walk --entry-points --require-site-packages - - name: Show HF cache - run: hf cache ls + - name: Pre-Download dataset fixture + run: | + hf download --repo-type=dataset axolotl-ai-internal/axolotl-oss-dataset-fixtures + + - name: Run CLI tests against the installed wheel + timeout-minutes: 20 + run: | + pytest -v --durations=10 tests/cli/ + + select-e2e: + name: Select e2e tests + runs-on: ubuntu-latest + if: ${{ !github.event.pull_request.draft }} + timeout-minutes: 15 + steps: + - name: Check out repository code + uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + with: + persist-credentials: false + # the PR merge commit plus both parents, so the diff is exactly the PR + fetch-depth: 2 + + - name: Setup Python + uses: actions/setup-python@5fda3b95a4ea91299a34e894583c3862153e4b97 # v7.0.0 + with: + python-version: "3.12" + + # Preview only: the selected-e2e GPU arms run this subset alongside the full suites. + # Anything the selector cannot reason about resolves to "all", and those arms then skip. + - name: Select impacted e2e test files + continue-on-error: true + run: | + out="$RUNNER_TEMP/e2e-selection" + mkdir -p "$out" + args=() + if [ "$GITHUB_EVENT_NAME" = "pull_request" ]; then + args=(--merge-commit) + fi + python cicd/select_e2e_tests.py "${args[@]}" --name single --explain \ + --exclude tests/e2e/multigpu --exclude tests/e2e/kernels \ + --output "$out/single.txt" + python cicd/select_e2e_tests.py "${args[@]}" --name multigpu --explain \ + --scope tests/e2e/multigpu \ + --output "$out/multigpu.txt" + + - name: Upload selection + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: e2e-selection + path: ${{ runner.temp }}/e2e-selection + if-no-files-found: warn + retention-days: 7 + + multigpu-paths: + name: Multi-GPU path filter + runs-on: ubuntu-latest + if: ${{ !github.event.pull_request.draft }} + outputs: + run_multigpu: ${{ steps.filter.outputs.run_multigpu }} + steps: + - uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9.0.0 + id: filter + with: + script: | + // trailing '/' = prefix, '*' = glob within one path segment, otherwise exact path + const PATTERNS = [ + 'src/axolotl/core/', + 'src/axolotl/loaders/', + 'src/axolotl/model_support/', + 'src/axolotl/monkeypatch/', + 'src/axolotl/kernels/', + 'src/axolotl/utils/distributed.py', + 'src/axolotl/utils/schemas/', + 'src/axolotl/utils/model_shard_quant.py', + 'src/axolotl/integrations/expert_parallel/', + 'src/axolotl/integrations/context_parallel/', + 'src/axolotl/integrations/cut_cross_entropy/', + 'src/axolotl/integrations/kernels/', + 'src/axolotl/train.py', + 'src/axolotl/cli/', + 'deepspeed_configs/', + 'tests/e2e/multigpu/', + 'cicd/', + '.github/workflows/multi-gpu-e2e*.yml', + '.github/workflows/tests.yml', + 'pyproject.toml', + 'requirements*.txt', + 'scripts/cutcrossentropy_install.py', + 'cicd/Dockerfile*', + ]; + if (context.eventName !== 'pull_request') { + core.info(`event ${context.eventName}: multigpu enabled`); + core.setOutput('run_multigpu', 'true'); + return; + } + const escape = (s) => s.replace(/[.+?^${}()|[\]\\]/g, '\\$&'); + const toRegex = (p) => new RegExp('^' + p.split('*').map(escape).join('[^/]*') + '$'); + const matches = (file) => PATTERNS.some((p) => + p.endsWith('/') ? file.startsWith(p) : p.includes('*') ? toRegex(p).test(file) : file === p); + const files = await github.paginate(github.rest.pulls.listFiles, { + ...context.repo, + pull_number: context.payload.pull_request.number, + per_page: 100, + }); + const hits = files.flatMap((f) => [f.filename, f.previous_filename]).filter((name) => name && matches(name)); + core.info(`${files.length} changed file(s); ${hits.length} matched multigpu paths`); + for (const hit of hits) core.info(` ${hit}`); + core.setOutput('run_multigpu', String(hits.length > 0)); check-gpu-label: + # Deferred so a label added while CPU jobs run is seen without a rerun + needs: [pre-commit, nf4-distributed, pytest-testmon, sdist] uses: ./.github/workflows/gpu-pr-gate.yml docker-e2e: - needs: [pre-commit, nf4-distributed, pytest, pytest-sdist, check-gpu-label] + needs: [pre-commit, nf4-distributed, pytest-testmon, sdist, check-gpu-label] if: needs.check-gpu-label.outputs.enabled == 'true' uses: ./.github/workflows/docker-e2e.yml secrets: inherit multi-gpu-e2e: - needs: [pre-commit, nf4-distributed, pytest, pytest-sdist, check-gpu-label] - if: github.event_name == 'pull_request' && needs.check-gpu-label.outputs.enabled == 'true' + needs: [pre-commit, nf4-distributed, pytest-testmon, sdist, check-gpu-label, multigpu-paths] + if: github.event_name == 'pull_request' && needs.check-gpu-label.outputs.enabled == 'true' && needs.multigpu-paths.outputs.run_multigpu == 'true' uses: ./.github/workflows/multi-gpu-e2e.yml secrets: inherit diff --git a/.gitignore b/.gitignore index f62ac905d..f1f4ef33a 100644 --- a/.gitignore +++ b/.gitignore @@ -196,3 +196,5 @@ out/ # scm auto-versioning src/axolotl/_version.py +.testmon/ +.ci_source_refreshed diff --git a/cicd/Dockerfile-uv.jinja b/cicd/Dockerfile-uv.jinja index 16d9cdaa5..7eed30fcc 100644 --- a/cicd/Dockerfile-uv.jinja +++ b/cicd/Dockerfile-uv.jinja @@ -1,12 +1,12 @@ FROM axolotlai/axolotl-base-uv:{{ BASE_TAG }} +ENV AXOLOTL_DEPS_HASH="{{ DEPS_HASH }}" ENV TORCH_CUDA_ARCH_LIST="7.0 7.5 8.0 8.6 9.0+PTX" ENV AXOLOTL_EXTRAS="{{ AXOLOTL_EXTRAS }}" ENV AXOLOTL_ARGS="{{ AXOLOTL_ARGS }}" ENV CUDA="{{ CUDA }}" ENV PYTORCH_VERSION="{{ PYTORCH_VERSION }}" -ENV GITHUB_REF="{{ GITHUB_REF }}" -ENV GITHUB_SHA="{{ GITHUB_SHA }}" +ENV IMAGE_REF="{{ IMAGE_REF }}" ENV NIGHTLY_BUILD="{{ NIGHTLY_BUILD }}" ENV HF_HOME="{{ HF_HOME }}" @@ -16,7 +16,7 @@ RUN git clone --depth=1 https://github.com/axolotl-ai-cloud/axolotl.git WORKDIR /workspace/axolotl -RUN git fetch origin +$GITHUB_REF && \ +RUN git fetch origin +$IMAGE_REF && \ git checkout FETCH_HEAD RUN uv pip install packaging==26.0 setuptools==78.1.1 @@ -43,7 +43,7 @@ RUN python scripts/cutcrossentropy_install.py --uv | sh # So we can test the Docker image RUN uv pip install black mypy pre-commit types-requests quartodoc jupyter blobfile tiktoken \ - codecov codecov-cli pytest pytest-cov pytest-retry pytest-sugar pytest-xdist tbparse + codecov-cli pytest pytest-cov pytest-retry pytest-sugar pytest-testmon pytest-xdist tbparse # fix so that git fetch/pull from remote works RUN git config remote.origin.fetch "+refs/heads/*:refs/remotes/origin/*" && \ diff --git a/cicd/cicd.sh b/cicd/cicd.sh index 8e3b284b3..95f9be87d 100755 --- a/cicd/cicd.sh +++ b/cicd/cicd.sh @@ -13,7 +13,7 @@ bash ./cicd/prepare_hf_cache.sh env -u CODECOV_TOKEN python -c "from kernels import get_kernel; get_kernel(\"kernels-community/flash-attn2\", version=3, trust_remote_code=True)" # Run unit tests with initial coverage report -pytest -v --durations=10 -n8 -m "not slow and not nf4_distributed" \ +pytest -v --durations=10 -n8 --maxfail=10 -m "gpu and not slow and not nf4_distributed" \ --ignore=tests/e2e/ \ --ignore=tests/integrations/ \ --ignore=tests/patched/ \ @@ -22,7 +22,7 @@ pytest -v --durations=10 -n8 -m "not slow and not nf4_distributed" \ --cov=axolotl # Run patched training tests with coverage append -pytest --full-trace -vvv --durations=10 \ +pytest --full-trace -vvv --durations=10 --maxfail=10 \ /workspace/axolotl/tests/e2e/patched \ --cov=axolotl \ --cov-append @@ -30,25 +30,25 @@ pytest --full-trace -vvv --durations=10 \ # Run solo tests with coverage append # test_rm_lora is run in its own process below (it fails on py3.11 when sharing # the solo process with other tests; isolating it avoids cross-test state). -pytest -v --durations=10 -n1 \ +pytest -v --durations=10 --maxfail=10 -n1 \ --ignore=tests/e2e/solo/test_reward_model_smollm2.py \ /workspace/axolotl/tests/e2e/solo/ \ --cov=axolotl \ --cov-append # Run reward-model test isolated in its own process -pytest -v --durations=10 -s \ +pytest -v --durations=10 --maxfail=10 -s \ /workspace/axolotl/tests/e2e/solo/test_reward_model_smollm2.py \ --cov=axolotl \ --cov-append # Run E2E integration tests with coverage append -pytest -v --durations=10 \ +pytest -v --durations=10 --maxfail=10 \ /workspace/axolotl/tests/e2e/integrations/ \ --cov=axolotl \ --cov-append -pytest -v --durations=10 -n8 --dist loadfile \ +pytest -v --durations=10 -n8 --dist loadfile --maxfail=10 -m gpu \ --ignore=tests/e2e/kernels/ \ --ignore=tests/integrations/kernels/ \ --ignore=tests/integrations/monkeypatch/test_tiled_mlp_moe.py \ @@ -61,12 +61,8 @@ pytest -v --durations=10 -n8 --dist loadfile \ --cov=axolotl \ --cov-append -pytest -v --durations=10 /workspace/axolotl/tests/cli \ - --cov=axolotl \ - --cov-append - # Run remaining e2e tests with coverage append and final report -pytest -v --durations=10 \ +pytest -v --durations=10 --maxfail=10 \ --ignore=tests/e2e/kernels/ \ --ignore=tests/e2e/solo/ \ --ignore=tests/e2e/patched/ \ @@ -78,4 +74,4 @@ pytest -v --durations=10 \ --cov-append \ --cov-report=xml:e2e-coverage.xml -codecov upload-process -t $CODECOV_TOKEN -f e2e-coverage.xml -F e2e,pytorch-${PYTORCH_VERSION} || true +codecovcli upload-process -t $CODECOV_TOKEN -f e2e-coverage.xml -F e2e,pytorch-${PYTORCH_VERSION} || true diff --git a/cicd/cicd_cuda_kernels.sh b/cicd/cicd_cuda_kernels.sh index 77d16930a..17117f34c 100755 --- a/cicd/cicd_cuda_kernels.sh +++ b/cicd/cicd_cuda_kernels.sh @@ -6,11 +6,11 @@ python -c "import torch; assert '$PYTORCH_VERSION' in torch.__version__, f'Expec bash ./cicd/prepare_hf_cache.sh # Keep process-wide model patching isolated from kernel parity tests. -pytest -v --durations=10 \ +pytest -v --durations=10 --maxfail=10 \ /workspace/axolotl/tests/e2e/kernels/lora_patching/ \ --cov=axolotl -pytest -v --durations=10 \ +pytest -v --durations=10 --maxfail=10 \ --ignore=tests/e2e/kernels/lora_patching/ \ /workspace/axolotl/tests/e2e/kernels/ \ /workspace/axolotl/tests/integrations/kernels/ \ @@ -23,10 +23,10 @@ pytest -v --durations=10 \ --cov=axolotl \ --cov-append -pytest -v --durations=10 -m slow \ +pytest -v --durations=10 --maxfail=10 -m slow \ /workspace/axolotl/tests/monkeypatch/test_fla_mamba.py \ --cov=axolotl \ --cov-append \ --cov-report=xml:e2e-kernel-coverage.xml -codecov upload-process -t "$CODECOV_TOKEN" -f e2e-kernel-coverage.xml -F e2e,kernels,pytorch-${PYTORCH_VERSION} || true +codecovcli upload-process -t "$CODECOV_TOKEN" -f e2e-kernel-coverage.xml -F e2e,kernels,pytorch-${PYTORCH_VERSION} || true diff --git a/cicd/cicd_selected.sh b/cicd/cicd_selected.sh new file mode 100755 index 000000000..0744c94d9 --- /dev/null +++ b/cicd/cicd_selected.sh @@ -0,0 +1,37 @@ +#!/bin/bash +# Run only the single-GPU e2e test files the selector picked (E2E_SELECTED_TESTS, comma-separated). +set -eo pipefail + +python -c "import torch; assert '$PYTORCH_VERSION' in torch.__version__, f'Expected torch $PYTORCH_VERSION but got {torch.__version__}'" + +: "${E2E_SELECTED_TESTS:?no selected tests}" +IFS=',' read -ra SELECTED <<< "$E2E_SELECTED_TESTS" + +bash ./cicd/prepare_hf_cache.sh +env -u CODECOV_TOKEN python -c "from kernels import get_kernel; get_kernel(\"kernels-community/flash-attn2\", version=3, trust_remote_code=True)" + +PATCHED=() SOLO=() REWARD=() INTEGRATIONS=() REST=() +for f in "${SELECTED[@]}"; do + case "$f" in + tests/e2e/multigpu/*|tests/e2e/kernels/*) ;; + tests/e2e/patched/*) PATCHED+=("$f") ;; + tests/e2e/solo/test_reward_model_smollm2.py) REWARD+=("$f") ;; + tests/e2e/solo/*) SOLO+=("$f") ;; + tests/e2e/integrations/*) INTEGRATIONS+=("$f") ;; + tests/e2e/*) REST+=("$f") ;; + esac +done + +run() { + local -n files=$1; shift + [ ${#files[@]} -eq 0 ] && return 0 + echo "=== ${#files[@]} selected file(s): $*" + pytest --durations=10 --maxfail=10 "$@" "${files[@]}" +} + +# same process layout as cicd.sh, minus coverage +run PATCHED --full-trace -vvv +run SOLO -v -n1 +run REWARD -v -s +run INTEGRATIONS -v +run REST -v diff --git a/cicd/cicd_testmon.sh b/cicd/cicd_testmon.sh new file mode 100644 index 000000000..d0d951119 --- /dev/null +++ b/cicd/cicd_testmon.sh @@ -0,0 +1,61 @@ +#!/bin/bash +set -e + +python -c "import torch; assert '$PYTORCH_VERSION' in torch.__version__, f'Expected torch $PYTORCH_VERSION but got {torch.__version__}'" + +bash ./cicd/prepare_hf_cache.sh + +env -u CODECOV_TOKEN python -c "from kernels import get_kernel; get_kernel(\"kernels-community/flash-attn2\", version=3, trust_remote_code=True)" + +TESTMON_BASELINE_DIR="/workspace/data/testmon/${BASE_TAG:?}/single-gpu" +TESTMON_DIR="/workspace/axolotl/.testmon" +mkdir -p "$TESTMON_DIR" +if [ -d "$TESTMON_BASELINE_DIR" ]; then + cp -a "$TESTMON_BASELINE_DIR"/. "$TESTMON_DIR"/ + echo "testmon: restored baseline from $TESTMON_BASELINE_DIR" +else + echo "testmon: no baseline at $TESTMON_BASELINE_DIR, running everything" +fi + +tm() { + local name=$1; shift + TESTMON_DATAFILE="$TESTMON_DIR/${name}.testmondata" pytest --testmon-forceselect --maxfail=10 "$@" +} + +# testmon only traces the pytest process; these files run the code under test in child processes +SUBPROCESS_TESTS="/workspace/axolotl/tests/e2e/test_evaluate.py /workspace/axolotl/tests/e2e/test_preprocess.py /workspace/axolotl/tests/e2e/test_qwen.py /workspace/axolotl/tests/e2e/integrations/test_kd.py /workspace/axolotl/tests/e2e/test_quantization.py" + +tm patched --full-trace -vvv --durations=10 /workspace/axolotl/tests/e2e/patched + +tm solo -v --durations=10 -n1 \ + --ignore=tests/e2e/solo/test_reward_model_smollm2.py \ + /workspace/axolotl/tests/e2e/solo/ + +tm solo-reward-model -v --durations=10 -s \ + /workspace/axolotl/tests/e2e/solo/test_reward_model_smollm2.py + +tm integrations -v --durations=10 \ + --ignore=tests/e2e/integrations/test_kd.py \ + /workspace/axolotl/tests/e2e/integrations/ + +tm e2e -v --durations=10 \ + --ignore=tests/e2e/kernels/ \ + --ignore=tests/e2e/solo/ \ + --ignore=tests/e2e/patched/ \ + --ignore=tests/e2e/multigpu/ \ + --ignore=tests/e2e/integrations/ \ + --ignore=tests/cli \ + --ignore=tests/e2e/test_evaluate.py \ + --ignore=tests/e2e/test_preprocess.py \ + --ignore=tests/e2e/test_qwen.py \ + --ignore=tests/e2e/test_quantization.py \ + /workspace/axolotl/tests/e2e/ + +# shellcheck disable=SC2086 +pytest -v --durations=10 --maxfail=10 $SUBPROCESS_TESTS + +if [ "${GITHUB_REF:-}" = "refs/heads/main" ]; then + mkdir -p "$TESTMON_BASELINE_DIR" + cp -a "$TESTMON_DIR"/. "$TESTMON_BASELINE_DIR"/ + echo "testmon: saved baseline to $TESTMON_BASELINE_DIR" +fi diff --git a/cicd/deps_hash.py b/cicd/deps_hash.py new file mode 100644 index 000000000..b1cebd9b0 --- /dev/null +++ b/cicd/deps_hash.py @@ -0,0 +1,126 @@ +#!/usr/bin/env python +"""Hash the inputs that change the CI Modal image, so edits elsewhere reuse it. + +Only the dependency-bearing parts of pyproject.toml count: a pytest marker or a ruff +rule must not force a rebuild of a 150 s image. The hash also picks IMAGE_REF, the git +ref the image installs dependencies from: main whenever the branch's dependency inputs +match main's, so every such PR shares one image, and the branch's own ref otherwise. +The container checks out the exact commit under test at start-up either way. +""" + +from __future__ import annotations + +import datetime as dt +import hashlib +import json +import os +import subprocess # nosec B404 +import sys +import tomllib +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parents[1] +DEFAULT_DOCKERFILE = "Dockerfile-uv.jinja" +INSTALL_SCRIPTS = ("scripts/cutcrossentropy_install.py",) +MAIN_REF = "refs/heads/main" + + +def dependency_sections(pyproject: dict) -> dict: + project = pyproject.get("project", {}) + return { + "build-system": pyproject.get("build-system"), + "requires-python": project.get("requires-python"), + "dependencies": project.get("dependencies"), + "optional-dependencies": project.get("optional-dependencies"), + "dependency-groups": pyproject.get("dependency-groups"), + "tool.uv": pyproject.get("tool", {}).get("uv"), + "tool.setuptools": pyproject.get("tool", {}).get("setuptools"), + } + + +def _read(repo_root: Path, rel: str, rev: str | None) -> bytes: + if rev is None: + return (repo_root / rel).read_bytes() + return subprocess.check_output( # nosec B603 B607 + ["git", "show", f"{rev}:{rel}"], cwd=repo_root, stderr=subprocess.DEVNULL + ) + + +def compute( + repo_root: Path, + env: dict[str, str], + today: dt.date | None = None, + rev: str | None = None, +) -> str: + """Hash of the image inputs at the working tree, or at ``rev`` when given.""" + digest = hashlib.sha256() + sections = dependency_sections( + tomllib.loads(_read(repo_root, "pyproject.toml", rev).decode()) + ) + digest.update(json.dumps(sections, sort_keys=True).encode()) + dockerfiles = [DEFAULT_DOCKERFILE] + if env.get("E2E_DOCKERFILE") and env["E2E_DOCKERFILE"] != DEFAULT_DOCKERFILE: + dockerfiles.append(env["E2E_DOCKERFILE"]) + for rel in [*INSTALL_SCRIPTS, *(f"cicd/{name}" for name in dockerfiles)]: + digest.update(rel.encode()) + digest.update(_read(repo_root, rel, rev)) + digest.update( + f"extras={env.get('AXOLOTL_EXTRAS', '')} args={env.get('AXOLOTL_ARGS', '')} " + f"base={env.get('BASE_TAG', '')}".encode() + ) + day = today or dt.datetime.now(dt.UTC).date() + if env.get("NIGHTLY_BUILD") == "true": + digest.update(day.isoformat().encode()) + else: + # unpinned ranges and the moving base tag drift; bound a shared image's age to a week + digest.update("{0}-W{1:02d}".format(*day.isocalendar()[:2]).encode()) + return digest.hexdigest()[:16] + + +def image_ref( + repo_root: Path, env: dict[str, str], deps_hash: str, today: dt.date | None = None +) -> str: + """main when this tree's dependency inputs match main's, else the ref under test.""" + github_ref = env.get("GITHUB_REF", MAIN_REF) + if github_ref == MAIN_REF: + return MAIN_REF + subprocess.run( # nosec B603 B607 + [ + "git", + "fetch", + "-q", + "--depth=1", + "origin", + f"+{MAIN_REF}:refs/remotes/origin/main", + ], + cwd=repo_root, + check=False, + capture_output=True, + ) + try: + main_hash = compute(repo_root, env, today, rev="origin/main") + except ( + subprocess.CalledProcessError, + OSError, + tomllib.TOMLDecodeError, + UnicodeDecodeError, + ): + return github_ref + return MAIN_REF if main_hash == deps_hash else github_ref + + +def main() -> int: + env = dict(os.environ) + deps_hash = compute(REPO_ROOT, env) + ref = image_ref(REPO_ROOT, env, deps_hash) + github_env = os.environ.get("GITHUB_ENV") + if github_env: + with open(github_env, "a", encoding="utf-8") as fh: + fh.write(f"DEPS_HASH={deps_hash}\nIMAGE_REF={ref}\n") + print(f"DEPS_HASH={deps_hash}") + print(f"IMAGE_REF={ref}") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/cicd/e2e_selected.py b/cicd/e2e_selected.py new file mode 100644 index 000000000..c81001f32 --- /dev/null +++ b/cicd/e2e_selected.py @@ -0,0 +1,20 @@ +"""Modal app to run the selector-picked single-GPU e2e test files.""" + +from .single_gpu import GPU_CONFIG, VOLUME_CONFIG, app, cicd_image, run_cmd + + +@app.function( + image=cicd_image, + gpu=GPU_CONFIG, + timeout=120 * 60, + cpu=8.0, + memory=131072, + volumes=VOLUME_CONFIG, +) +def cicd_pytest_selected(): + run_cmd("bash ./cicd/cicd_selected.sh", "/workspace/axolotl") + + +@app.local_entrypoint() +def main(): + cicd_pytest_selected.remote() diff --git a/cicd/e2e_testmon.py b/cicd/e2e_testmon.py new file mode 100644 index 000000000..d4a79b45a --- /dev/null +++ b/cicd/e2e_testmon.py @@ -0,0 +1,29 @@ +"""Modal app to run the single-GPU e2e suite under pytest-testmon.""" + +import os + +import modal + +from .single_gpu import GPU_CONFIG, VOLUME_CONFIG, app, cicd_image, run_cmd + +TESTMON_MOUNT = "/workspace/data/testmon" +testmon_volume = modal.Volume.from_name("axolotl-ci-testmon", create_if_missing=True) + + +@app.function( + image=cicd_image, + gpu=GPU_CONFIG, + timeout=120 * 60, + cpu=8.0, + memory=131072, + volumes={**VOLUME_CONFIG, TESTMON_MOUNT: testmon_volume}, +) +def cicd_pytest_testmon(): + run_cmd("bash ./cicd/cicd_testmon.sh", "/workspace/axolotl") + if os.environ.get("GITHUB_REF") == "refs/heads/main": + testmon_volume.commit() + + +@app.local_entrypoint() +def main(): + cicd_pytest_testmon.remote() diff --git a/cicd/multigpu.py b/cicd/multigpu.py index c23c580e2..8b960a850 100644 --- a/cicd/multigpu.py +++ b/cicd/multigpu.py @@ -25,8 +25,11 @@ df_args = { "AXOLOTL_ARGS": os.environ.get("AXOLOTL_ARGS", ""), "PYTORCH_VERSION": os.environ.get("PYTORCH_VERSION", "2.13.0"), "BASE_TAG": os.environ.get("BASE_TAG", "main-base-py3.12-cu130-2.13.0"), + "DEPS_HASH": os.environ.get("DEPS_HASH", "dev"), "CUDA": os.environ.get("CUDA", "130"), "GITHUB_REF": os.environ.get("GITHUB_REF", "refs/heads/main"), + "IMAGE_REF": os.environ.get("IMAGE_REF") + or os.environ.get("GITHUB_REF", "refs/heads/main"), "GITHUB_SHA": os.environ.get("GITHUB_SHA", ""), "NIGHTLY_BUILD": os.environ.get("NIGHTLY_BUILD", ""), "MULTIGPU_TEST_SUITE": os.environ.get("MULTIGPU_TEST_SUITE", "multigpu"), @@ -35,6 +38,7 @@ df_args = { "HF_HOME": "/workspace/data/huggingface-cache/hub", "PYTHONUNBUFFERED": os.environ.get("PYTHONUNBUFFERED", "1"), "DEEPSPEED_LOG_LEVEL": os.environ.get("DEEPSPEED_LOG_LEVEL", "WARNING"), + "E2E_SELECTED_TESTS": os.environ.get("E2E_SELECTED_TESTS", ""), } dockerfile_contents = df_template.render(**df_args) @@ -62,9 +66,67 @@ GPU_TYPE = os.environ.get("GPU_TYPE", "H100") GPU_CONFIG = f"{GPU_TYPE}:{N_GPUS}" +REFRESH_MARKER = ".ci_source_refreshed" + + +def _git(run_folder: str, *args: str, check: bool = True, capture: bool = False): + import subprocess # nosec + + return subprocess.run( # nosec + ["git", *args], + cwd=run_folder, + check=check, + capture_output=capture, + text=True, + ) + + +def checkout_source(run_folder: str, ref: str, sha: str) -> None: + _git(run_folder, "fetch", "--depth=1", "origin", f"+{ref}") + ref_head = _git(run_folder, "rev-parse", "FETCH_HEAD", capture=True).stdout.strip() + target = sha + if ( + _git(run_folder, "cat-file", "-e", f"{sha}^{{commit}}", check=False).returncode + != 0 + ): + if ( + _git( + run_folder, "fetch", "--depth=1", "origin", sha, check=False + ).returncode + != 0 + ): + # refs/pull/N/merge is regenerated when the base moves; the old merge SHA is then unreachable + if not ref.startswith("refs/pull/"): + raise RuntimeError(f"{sha} not reachable from {ref}") + print(f"WARNING: {sha} not fetchable, testing tip of {ref}", flush=True) + target = ref_head + _git(run_folder, "checkout", "-f", target) + _git(run_folder, "log", "-1", "--oneline") + + +def refresh_source(run_folder: str) -> None: + import subprocess # nosec + + sha = os.environ.get("GITHUB_SHA", "") + if not sha: + return + marker = pathlib.Path(run_folder) / REFRESH_MARKER + if marker.exists() and marker.read_text().strip() == sha: + return + checkout_source(run_folder, os.environ.get("GITHUB_REF", "refs/heads/main"), sha) + subprocess.run( # nosec + ["uv", "pip", "install", "--no-build-isolation", "--no-deps", "-e", "."], + cwd=run_folder, + check=True, + ) + marker.write_text(sha) + + def run_cmd(cmd: str, run_folder: str): import subprocess # nosec + refresh_source(run_folder) + # Propagate errors from subprocess. if exit_code := subprocess.call(cmd.split(), cwd=run_folder): # nosec exit(exit_code) @@ -88,6 +150,7 @@ def cicd_pytest(): "multigpu": "./cicd/multigpu.sh", "nf4": "bash ./cicd/nf4.sh", "nvfp4": "bash ./cicd/nvfp4.sh", + "multigpu_selected": "bash ./cicd/multigpu_selected.sh", } run_cmd(scripts[suite], "/workspace/axolotl") diff --git a/cicd/multigpu.sh b/cicd/multigpu.sh index 84a7cbec0..a5b921bb5 100755 --- a/cicd/multigpu.sh +++ b/cicd/multigpu.sh @@ -24,29 +24,29 @@ done env -u CODECOV_TOKEN python -c "from kernels import get_kernel; get_kernel(\"kernels-community/flash-attn2\", version=3, trust_remote_code=True)" # Only run two tests at a time to avoid OOM on GPU (with coverage collection) -pytest -v --durations=10 -n2 --maxfail=3 \ +pytest -v --durations=10 -n2 --maxfail=10 \ --ignore=/workspace/axolotl/tests/e2e/multigpu/solo/ \ --ignore=/workspace/axolotl/tests/e2e/multigpu/patched/ \ /workspace/axolotl/tests/e2e/multigpu/ \ --cov=axolotl -# Run solo tests with coverage append -pytest -v --durations=10 -n1 \ +# Run solo tests with coverage append; the NVFP4 suite has its own SM100 job +pytest -v --durations=10 -n1 --maxfail=10 -m "not nvfp4 and not slow" \ /workspace/axolotl/tests/e2e/multigpu/solo/ \ --cov=axolotl \ --cov-append -pytest -v --durations=10 -n1 \ +pytest -v --durations=10 -n1 --maxfail=10 \ /workspace/axolotl/tests/integrations/test_expert_parallel_nvfp4_wrap.py \ --cov=axolotl \ --cov-append -pytest -v --durations=10 -n1 /workspace/axolotl/tests/e2e/multigpu/patched/ \ +pytest -v --durations=10 -n1 --maxfail=10 /workspace/axolotl/tests/e2e/multigpu/patched/ \ --cov=axolotl \ --cov-append \ --cov-report=xml:multigpu-coverage.xml # Upload coverage to Codecov if CODECOV_TOKEN is available if [ -n "$CODECOV_TOKEN" ]; then - codecov upload-process -t "${CODECOV_TOKEN}" -f multigpu-coverage.xml -F multigpu,docker-tests,pytorch-${PYTORCH_VERSION} || true + codecovcli upload-process -t "${CODECOV_TOKEN}" -f multigpu-coverage.xml -F multigpu,docker-tests,pytorch-${PYTORCH_VERSION} || true fi diff --git a/cicd/multigpu_selected.sh b/cicd/multigpu_selected.sh new file mode 100755 index 000000000..be3ea9db6 --- /dev/null +++ b/cicd/multigpu_selected.sh @@ -0,0 +1,32 @@ +#!/bin/bash +# Run only the multi-GPU e2e test files the selector picked (E2E_SELECTED_TESTS, comma-separated). +set -eo pipefail + +python -c "import torch; assert '$PYTORCH_VERSION' in torch.__version__, f'Expected torch $PYTORCH_VERSION but got {torch.__version__}'" + +: "${E2E_SELECTED_TESTS:?no selected tests}" +IFS=',' read -ra SELECTED <<< "$E2E_SELECTED_TESTS" + +bash ./cicd/prepare_hf_cache.sh +env -u CODECOV_TOKEN python -c "from kernels import get_kernel; get_kernel(\"kernels-community/flash-attn2\", version=3, trust_remote_code=True)" + +SOLO=() PATCHED=() REST=() +for f in "${SELECTED[@]}"; do + case "$f" in + tests/e2e/multigpu/solo/*) SOLO+=("$f") ;; + tests/e2e/multigpu/patched/*) PATCHED+=("$f") ;; + tests/e2e/multigpu/*) REST+=("$f") ;; + esac +done + +run() { + local -n files=$1; shift + [ ${#files[@]} -eq 0 ] && return 0 + echo "=== ${#files[@]} selected file(s): $*" + pytest -v --durations=10 --maxfail=10 "$@" "${files[@]}" +} + +# same process layout as multigpu.sh, minus coverage; NVFP4 has its own SM100 job +run REST -n2 +run SOLO -n1 -m "not nvfp4 and not slow" +run PATCHED -n1 diff --git a/cicd/nf4.sh b/cicd/nf4.sh index 501280066..455ed2b22 100644 --- a/cicd/nf4.sh +++ b/cicd/nf4.sh @@ -44,6 +44,6 @@ assert not skipped, ( PY if [ -n "${CODECOV_TOKEN:-}" ]; then - codecov upload-process -t "$CODECOV_TOKEN" -f nf4-coverage.xml \ + codecovcli upload-process -t "$CODECOV_TOKEN" -f nf4-coverage.xml \ -F nf4,multigpu,docker-tests,pytorch-${PYTORCH_VERSION} || true fi diff --git a/cicd/nvfp4.sh b/cicd/nvfp4.sh index 14581539b..3b48a23af 100644 --- a/cicd/nvfp4.sh +++ b/cicd/nvfp4.sh @@ -17,19 +17,8 @@ for name in ("torch", "transformers", "accelerate", "torchao", "deepspeed", "pef print(f"{name}=={importlib.metadata.version(name)}", flush=True) NVFP4_ENV -pytest -v --durations=10 -n0 \ - tests/e2e/multigpu/solo/test_native_nvfp4_dynamic_ste.py \ - tests/e2e/multigpu/solo/test_sonicmoe_nvfp4_fsdp2.py \ - tests/e2e/multigpu/solo/test_native_nvfp4_fsdp2_lora_parity.py \ - tests/e2e/multigpu/solo/test_native_nvfp4_fsdp2_merge_aware_resume.py \ - tests/e2e/multigpu/solo/test_native_nvfp4_fsdp2_recipe.py \ - tests/e2e/multigpu/solo/test_native_nvfp4_tp_hf_load.py \ - tests/e2e/multigpu/solo/test_native_nvfp4_tp_lora_parity.py \ - tests/e2e/multigpu/solo/test_torchao_lora_deepspeed.py \ - tests/e2e/multigpu/solo/test_torchao_lora_deepspeed_checkpoint.py \ - tests/e2e/multigpu/solo/test_torchao_lora_deepspeed_zero3_checkpoint.py \ - tests/e2e/multigpu/solo/test_native_nvfp4_deepspeed_lora_merge_aware.py \ - tests/e2e/multigpu/solo/test_native_nvfp4_deepspeed_lora_merge_lifecycle.py \ +pytest -v --durations=10 -n0 -m "nvfp4 and not slow" \ + tests/e2e/multigpu/solo/ \ --junitxml=nvfp4-junit.xml --cov=axolotl --cov-report=xml:nvfp4-coverage.xml python - <<'NVFP4_JUNIT' @@ -42,6 +31,6 @@ assert not skipped, f"NVFP4 GPU CI unexpectedly skipped {len(skipped)} tests: {s NVFP4_JUNIT if [ -n "${CODECOV_TOKEN:-}" ]; then - codecov upload-process -t "$CODECOV_TOKEN" -f nvfp4-coverage.xml \ + codecovcli upload-process -t "$CODECOV_TOKEN" -f nvfp4-coverage.xml \ -F nvfp4,multigpu,docker-tests,pytorch-${PYTORCH_VERSION} || true fi diff --git a/cicd/select_e2e_tests.py b/cicd/select_e2e_tests.py new file mode 100644 index 000000000..1a298f163 --- /dev/null +++ b/cicd/select_e2e_tests.py @@ -0,0 +1,852 @@ +#!/usr/bin/env python +"""Select the e2e test files a diff can affect, for the experimental selected-e2e CI arms. + +Every edge is derived from the tree, so new features need no registration: + +* The config-key universe is every class-body annotated field under ``src/axolotl`` + that some in-scope e2e test sets. A module "reads" a key when its AST references + the attribute, subscript or ``.get`` by that name. +* Keys set by most in-scope tests carry no selection signal. A changed module that + reads any such key is core and runs the whole scope. +* Otherwise a changed module selects the tests that import it directly, that name + its module path, file stem or package in a config string, or that set a key it + reads. A module with no edges at all runs the whole scope. +* A module may declare ``__ci_config_keys__ = ("activation_offloading",)`` when the + graph cannot see what enables it; the selector then picks exactly the tests that set + those keys. A declared key that is not a config field runs the whole scope. +* Changed test files select themselves. Any other change under ``tests/``, a + deleted module, a build or CI file, or an exception in the selector itself runs + the whole scope. Nothing here can select fewer tests than the graph supports. +""" + +from __future__ import annotations + +import argparse +import ast +import fnmatch +import os +import re +import subprocess # nosec B404 +import sys +import tomllib +import traceback +from collections import defaultdict +from dataclasses import dataclass, field +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parents[1] +SRC_ROOT = "src/axolotl" +UBIQUITY_THRESHOLD = 0.5 +SOURCE_UBIQUITY_THRESHOLD = 0.2 +RUN_ALL_GLOBS = ( + "pyproject.toml", + "setup.py", + "setup.cfg", + "requirements*.txt", + "cicd/*", + "cicd/**/*", + ".github/workflows/*", + "scripts/*install*", + "tests/conftest.py", + f"{SRC_ROOT}/__init__.py", +) +FORCE_ALL_TOKENS = ("[test all]", "[no filter]") +# files that cannot reach a test at runtime; anything else unmodeled runs the whole scope +INERT_GLOBS = ( + "docs/*", + "docs/**/*", + "examples/*", + "examples/**/*", + "*.md", + "**/*.md", + "*.qmd", + "**/*.qmd", + "**/LICENSE", + "LICENSE", + ".gitignore", + ".pre-commit-config.yaml", + "*.png", + "**/*.png", + "*.svg", + "**/*.svg", + "*.jpg", + "**/*.jpg", +) +# Stems and package names too generic to identify a feature from a config string. +DECLARED_KEYS_NAME = "__ci_config_keys__" +GENERIC_NAMES = frozenset( + { + "__init__", + "args", + "base", + "cli", + "common", + "config", + "constants", + "core", + "helpers", + "loss", + "main", + "model", + "models", + "plugin", + "train", + "trainer", + "types", + "utils", + } +) +MIN_STRING_LEN = 3 +FUZZY_MIN_LEN = 5 +MIN_REGISTRY_ENTRIES = 3 +REGISTRY_NAMED_FRACTION = 1 / 3 + + +@dataclass +class Selection: + mode: str + tests: list[str] = field(default_factory=list) + reason: str = "" + explain: dict[str, set[str]] = field(default_factory=dict) + + +def _git(repo: Path, *args: str) -> str: + return subprocess.check_output( # nosec B603 B607 + ["git", *args], cwd=repo, text=True, stderr=subprocess.STDOUT + ) + + +def _safe_parse(path: Path) -> ast.AST | None: + try: + return ast.parse(path.read_text(encoding="utf-8")) + except (SyntaxError, UnicodeDecodeError, OSError): + return None + + +class Selector: + def __init__( + self, repo: Path, scope: str, exclude: tuple[str, ...] = (), head: str = "HEAD" + ): + self.repo = repo + self.head = head + self._edge_cache: dict[str, tuple[set[str] | None, str] | None] = {} + self._registry_cache: dict[str, bool] = {} + self._value_cache: dict[str, set[str]] = {} + self.scope = scope.rstrip("/") + self.exclude = tuple(e.rstrip("/") for e in exclude) + + # -- repo state -------------------------------------------------------- + + def test_files(self) -> list[str]: + out = [] + for path in sorted((self.repo / self.scope).rglob("test_*.py")): + rel = path.relative_to(self.repo).as_posix() + if any(rel.startswith(f"{e}/") for e in self.exclude): + continue + out.append(rel) + return out + + def source_files(self) -> list[str]: + return sorted( + p.relative_to(self.repo).as_posix() + for p in (self.repo / SRC_ROOT).rglob("*.py") + ) + + @staticmethod + def module_name(rel: str) -> str: + parts = Path(rel).with_suffix("").parts + if parts[0] == "src": + parts = parts[1:] + if parts[-1] == "__init__": + parts = parts[:-1] + return ".".join(parts) + + def _diff_base(self, base: str | None, merge_commit: bool) -> str: + if merge_commit: + return f"{self.head}^1" + assert base + return _git(self.repo, "merge-base", base, self.head).strip() + + def changed( + self, base: str | None, merge_commit: bool + ) -> tuple[list[str], list[str]]: + rev = self._diff_base(base, merge_commit) + out = _git(self.repo, "diff", "--name-status", "--no-renames", rev, self.head) + changed: list[str] = [] + deleted: list[str] = [] + for line in out.splitlines(): + status, _, path = line.partition("\t") + (deleted if status.startswith("D") else changed).append(path.strip()) + return changed, deleted + + def commit_messages(self, base: str | None, merge_commit: bool) -> str: + rev = self._diff_base(base, merge_commit) + return _git(self.repo, "log", "--format=%B", f"{rev}..{self.head}") + + # -- AST extraction ---------------------------------------------------- + + @staticmethod + def _config_keys_and_values(tree: ast.AST) -> tuple[set[str], set[str]]: + """Top-level keys of dict literals, and every string constant in the file. + + Configs are assembled from helper arguments and fixtures as well as literals, so + any string may be a config value; nested dict keys are not config keys. + """ + keys: set[str] = set() + values: set[str] = set() + nested: set[int] = set() + for node in ast.walk(tree): + if isinstance(node, ast.Dict): + for child in ast.walk(node): + if child is not node and isinstance(child, ast.Dict): + nested.add(id(child)) + for node in ast.walk(tree): + if isinstance(node, ast.Constant) and isinstance(node.value, str): + values.add(node.value) + elif isinstance(node, ast.Dict) and id(node) not in nested: + for key in node.keys: + if isinstance(key, ast.Constant) and isinstance(key.value, str): + keys.add(key.value) + return keys, values - keys + + @staticmethod + def _import_targets(node: ast.Import | ast.ImportFrom, package: str) -> set[str]: + if isinstance(node, ast.Import): + return {alias.name for alias in node.names} + if node.level == 0: + base = node.module or "" + else: + parts = package.split(".") + parts = parts[: len(parts) - (node.level - 1)] if node.level > 1 else parts + base = ".".join(parts + ([node.module] if node.module else [])) + if not base: + return set() + return {base} | {f"{base}.{alias.name}" for alias in node.names} + + def _imports(self, tree: ast.AST, package: str = "") -> set[str]: + mods: set[str] = set() + for node in ast.walk(tree): + if isinstance(node, (ast.Import, ast.ImportFrom)): + mods |= self._import_targets(node, package) + return mods + + @staticmethod + def _declared_keys(tree: ast.AST) -> tuple[str, ...] | None: + for node in getattr(tree, "body", []): + if not isinstance(node, ast.Assign): + continue + if not any( + isinstance(t, ast.Name) and t.id == DECLARED_KEYS_NAME + for t in node.targets + ): + continue + keys: list[str] = [] + elts = ( + node.value.elts if isinstance(node.value, (ast.Tuple, ast.List)) else [] + ) + for elt in elts: + if not (isinstance(elt, ast.Constant) and isinstance(elt.value, str)): + return () + keys.append(elt.value) + return tuple(keys) + return None + + def _scan_source( + self, tree: ast.AST, package: str + ) -> tuple[set[str], set[str], set[str]]: + """One walk: class fields, imports, and referenced names.""" + fields: set[str] = set() + imports: set[str] = set() + names: set[str] = set() + for node in ast.walk(tree): + if isinstance(node, ast.ClassDef): + for stmt in node.body: + if isinstance(stmt, ast.AnnAssign) and isinstance( + stmt.target, ast.Name + ): + fields.add(stmt.target.id) + elif isinstance(node, (ast.Import, ast.ImportFrom)): + imports |= self._import_targets(node, package) + elif isinstance(node, ast.Attribute): + names.add(node.attr) + elif isinstance(node, ast.Subscript) and isinstance( + node.slice, ast.Constant + ): + if isinstance(node.slice.value, str): + names.add(node.slice.value) + elif isinstance(node, ast.Call): + func = node.func + is_get = isinstance(func, ast.Attribute) and func.attr in {"get", "pop"} + is_getattr = isinstance(func, ast.Name) and func.id in { + "getattr", + "hasattr", + } + if is_get and node.args and isinstance(node.args[0], ast.Constant): + if isinstance(node.args[0].value, str): + names.add(node.args[0].value) + if ( + is_getattr + and len(node.args) > 1 + and isinstance(node.args[1], ast.Constant) + ): + if isinstance(node.args[1].value, str): + names.add(node.args[1].value) + return fields, imports, names + + @staticmethod + def _referenced_names(tree: ast.AST) -> set[str]: + names = set() + for node in ast.walk(tree): + if isinstance(node, ast.Attribute): + names.add(node.attr) + elif isinstance(node, ast.Subscript) and isinstance( + node.slice, ast.Constant + ): + if isinstance(node.slice.value, str): + names.add(node.slice.value) + elif isinstance(node, ast.Call): + func = node.func + is_get = isinstance(func, ast.Attribute) and func.attr in {"get", "pop"} + is_getattr = isinstance(func, ast.Name) and func.id in { + "getattr", + "hasattr", + } + if is_get and node.args and isinstance(node.args[0], ast.Constant): + if isinstance(node.args[0].value, str): + names.add(node.args[0].value) + if ( + is_getattr + and len(node.args) > 1 + and isinstance(node.args[1], ast.Constant) + ): + if isinstance(node.args[1].value, str): + names.add(node.args[1].value) + return names + + # -- graph ------------------------------------------------------------- + + def build(self) -> None: + tests = self.test_files() + sources = self.source_files() + conftests = { + p.relative_to(self.repo).as_posix(): _safe_parse(p) + for p in (self.repo / self.scope).rglob("conftest.py") + } + + # helper modules next to the tests (worker scripts, parity probes) carry configs too + helpers: dict[str, tuple[set[str], set[str], set[str]]] = {} + for path in (self.repo / self.scope).rglob("*.py"): + rel = path.relative_to(self.repo).as_posix() + if path.name.startswith("test_") or path.name == "__init__.py": + continue + tree = _safe_parse(path) + if tree is None: + continue + hkeys, hstrings = self._config_keys_and_values(tree) + helpers[rel] = (hkeys, hstrings, self._imports(tree)) + + def helper_refs(rel: str, strings: set[str], imports: set[str]) -> set[str]: + """Helpers this file names by filename or stem, or imports by module.""" + here = Path(rel).parent + found = set() + for helper in helpers: + hpath = Path(helper) + hmod = self.module_name(helper) + if ( + hpath.name in strings + or hpath.stem in strings + or helper in strings + or any(imp == hmod or imp.startswith(hmod + ".") for imp in imports) + or ( + hpath.parent == here + and any(imp.endswith("." + hpath.stem) for imp in imports) + ) + ): + found.add(helper) + return found + + test_keys: dict[str, set[str]] = {} + test_strings: dict[str, set[str]] = {} + test_imports: dict[str, set[str]] = {} + helper_tests: dict[str, set[str]] = defaultdict(set) + for rel in tests: + tree = _safe_parse(self.repo / rel) + keys: set[str] = set() + strings: set[str] = set() + imports: set[str] = set() + if tree is not None: + keys, strings = self._config_keys_and_values(tree) + imports = self._imports(tree) + # fixtures merge into the config, so a conftest on the path contributes its keys + for cpath, ctree in conftests.items(): + if ctree is not None and rel.startswith(cpath[: -len("conftest.py")]): + ckeys, cstrings = self._config_keys_and_values(ctree) + keys |= ckeys + strings |= cstrings + seen: set[str] = set() + frontier = helper_refs(rel, strings, imports) + while frontier: + helper = frontier.pop() + if helper in seen: + continue + seen.add(helper) + hkeys, hstrings, himports = helpers[helper] + keys |= hkeys + strings |= hstrings + imports |= himports + frontier |= helper_refs(helper, hstrings, himports) - seen + test_keys[rel] = keys + test_strings[rel] = strings + test_imports[rel] = imports + for helper in seen: + helper_tests[helper].add(rel) + + all_fields: set[str] = set() + src_trees = {} + src_names: dict[str, set[str]] = {} + src_declared: dict[str, tuple[str, ...]] = {} + by_module = {self.module_name(rel): rel for rel in sources} + importers: dict[str, set[str]] = defaultdict(set) + for rel in sources: + tree = _safe_parse(self.repo / rel) + src_trees[rel] = tree + if tree is None: + continue + module = self.module_name(rel) + package = ( + module if rel.endswith("__init__.py") else module.rpartition(".")[0] + ) + fields, imports, names = self._scan_source(tree, package) + all_fields |= fields + src_names[rel] = names + declared = self._declared_keys(tree) + if declared is not None: + src_declared[rel] = declared + for imp in imports: + target = by_module.get(imp) + if target is None and "." in imp: + target = by_module.get(imp.rsplit(".", 1)[0]) + if target is not None and target != rel: + importers[target].add(rel) + self.importers = importers + used_keys = set().union(*test_keys.values()) if test_keys else set() + universe = all_fields & used_keys + readers: dict[str, int] = defaultdict(int) + for names in src_names.values(): + for key in names & universe: + readers[key] += 1 + universe -= { + k + for k, n in readers.items() + if n / max(len(src_names), 1) >= SOURCE_UBIQUITY_THRESHOLD + } + + opaque_exempt: set[str] = set() + configured = [t for t in tests if test_keys[t] & universe] + counts: dict[str, int] = defaultdict(int) + for test in configured: + for key in test_keys[test] & universe: + counts[key] += 1 + ubiquitous = { + k + for k, c in counts.items() + if configured and c / len(configured) >= UBIQUITY_THRESHOLD + } + # a test with no visible config can be reached by anything, so it rides along with every subset + self.opaque_tests = [ + t for t in tests if t not in configured and t not in opaque_exempt + ] + + self.entry_point_dirs = self._entry_point_dirs(by_module) + self.script_modules = self._script_modules(by_module) + self.tests = tests + self.helper_tests = helper_tests + self.test_keys = test_keys + self.test_strings = test_strings + self.test_imports = test_imports + self.src_trees = src_trees + self.src_names = src_names + self.src_declared = src_declared + self.all_fields = all_fields + self.universe = universe + self.ubiquitous = ubiquitous + + @staticmethod + def _norm(text: str) -> str: + return re.sub(r"[^a-z0-9]", "", text.lower()) + + def _value_edges(self, rel: str) -> set[str]: + """Tests whose config values name this module: a dotted path, a stem or package name.""" + if rel in self._value_cache: + return self._value_cache[rel] + hits = self._compute_value_edges(rel) + self._value_cache[rel] = hits + return hits + + def _compute_value_edges(self, rel: str) -> set[str]: + module = self.module_name(rel) + path = Path(rel) + is_package = path.name == "__init__.py" + own = path.parent.name if is_package else path.stem.lstrip("_") + exact: set[str] = set() + fuzzy: set[str] = set() + for name in (own, path.parent.name): + if name in GENERIC_NAMES or name == "axolotl" or len(name) < MIN_STRING_LEN: + continue + exact.add(name) + # repo ids like org/Model-Name reach model support packages + if name == own and len(self._norm(name)) >= FUZZY_MIN_LEN: + fuzzy.add(self._norm(name)) + hits = set() + for test, values in self.test_strings.items(): + if ( + values & exact + or module in values + # a plugin value is a class path inside its package + or (is_package and any(v.startswith(module + ".") for v in values)) + ): + hits.add(test) + continue + if fuzzy and any( + f in self._norm(v) for v in values if "/" in v for f in fuzzy + ): + hits.add(test) + return hits + + def _import_edges(self, rel: str) -> set[str]: + module = self.module_name(rel) + hits = set() + for test, imports in self.test_imports.items(): + if any(imp == module or imp.startswith(module + ".") for imp in imports): + hits.add(test) + return hits + + def _pyproject(self) -> dict: + try: + with open(self.repo / "pyproject.toml", "rb") as fh: + return tomllib.load(fh).get("project", {}) + except (OSError, tomllib.TOMLDecodeError): + return {} + + def _script_modules(self, by_module: dict[str, str]) -> set[str]: + """Modules behind console scripts: every subprocess-driven test starts there.""" + return { + by_module[target.split(":")[0]] + for target in self._pyproject().get("scripts", {}).values() + if target.split(":")[0] in by_module + } + + def _entry_point_dirs(self, by_module: dict[str, str]) -> set[str]: + """Package directories registered through pyproject entry points (plugins, model support).""" + groups = self._pyproject().get("entry-points", {}) + dirs = set() + for group, entries in groups.items(): + if not group.startswith("axolotl"): + continue + for target in entries.values(): + module = target.split(":")[0] + while module and module not in by_module: + module = module.rpartition(".")[0] + if module: + dirs.add(str(Path(by_module[module]).parent)) + return dirs + + def _is_registry_dir(self, directory: str) -> bool: + """A directory most of whose entries are named by config values (strategies, plugins, model support).""" + if directory in self._registry_cache: + return self._registry_cache[directory] + if directory == SRC_ROOT: + return False + entries: dict[str, bool] = {} + for rel in self.src_trees: + if not rel.startswith(directory + "/"): + continue + rest = rel[len(directory) + 1 :] + entry = rest.split("/", 1)[0] + if entry == "__init__.py": + continue + entries[entry] = entries.get(entry, False) or bool(self._value_edges(rel)) + named = sum(entries.values()) + by_value = ( + named >= MIN_REGISTRY_ENTRIES + and named / len(entries) >= REGISTRY_NAMED_FRACTION + ) + declared = sum( + 1 for entry in entries if f"{directory}/{entry}" in self.entry_point_dirs + ) + result = by_value or declared >= MIN_REGISTRY_ENTRIES + self._registry_cache[directory] = result + return result + + def _registry_entry(self, rel: str) -> str | None: + """The registry entry holding this module: the file itself or its package directory.""" + directory = str(Path(rel).parent) + if self._is_registry_dir(directory): + return rel + if self._is_registry_dir(str(Path(directory).parent)): + return directory + return None + + def _is_registry_leaf(self, rel: str, entry: str) -> bool: + """Reached only by config value: nothing outside the registry entry imports any of it.""" + members = [m for m in self.src_trees if m == entry or m.startswith(entry + "/")] + return not any( + imp + for member in members + for imp in self.importers.get(member, ()) + if imp != entry and not imp.startswith(entry + "/") + ) + + def _own_edges(self, rel: str) -> tuple[set[str] | None, str] | None: + """Edges this module carries itself; None when it has none and importers decide.""" + if rel in self._edge_cache: + return self._edge_cache[rel] + result = self._compute_own_edges(rel) + self._edge_cache[rel] = result + return result + + def _compute_own_edges(self, rel: str) -> tuple[set[str] | None, str] | None: + if rel in self.script_modules: + return None, "console script entry point" + declared = self.src_declared.get(rel) + if declared is not None: + unknown = [k for k in declared if k not in self.all_fields] + if not declared or unknown: + return ( + None, + f"{DECLARED_KEYS_NAME} names no config field: {unknown or declared}", + ) + hits = {t for t, keys in self.test_keys.items() if keys & set(declared)} + return hits, f"declared keys {list(declared)}" + names = self.src_names.get(rel) + if names is None: + tree = _safe_parse(self.repo / rel) + if tree is None: + return None, "unparseable module" + names = self._referenced_names(tree) + read = names & self.universe + discriminating = read - self.ubiquitous + key_hits = {t for t, keys in self.test_keys.items() if keys & discriminating} + imported = self._import_edges(rel) + named = self._value_edges(rel) + entry = self._registry_entry(rel) + if entry is not None: + if not named: + # members of a plugin or strategy package are reached through the package name + init = f"{entry}/__init__.py" + if init != rel and init in self.src_trees: + named = self._value_edges(init) + if named: + return ( + named | imported | key_hits, + f"named by config values; keys={sorted(discriminating)[:4]}", + ) + if self._is_registry_leaf(rel, entry): + return ( + imported | key_hits, + f"registry leaf no config value names; keys={sorted(discriminating)[:4]}", + ) + core = read & self.ubiquitous + if core: + return None, f"reads ubiquitous key(s) {sorted(core)[:3]}" + hits = named | imported | key_hits + if hits or discriminating: + return hits, f"keys={sorted(discriminating)[:4]}" + return None + + def impact_of_source(self, rel: str) -> tuple[set[str] | None, str]: + """Return (tests, reason); tests=None means the whole scope. + + A module without edges of its own is reached through whatever imports it, so + the walk climbs importers until every branch ends at a module with edges. + """ + own = self._own_edges(rel) + if own is not None: + return own + hits: set[str] = set() + seen = {rel} + frontier = [rel] + explained = [] + while frontier: + module = frontier.pop() + parents = self.importers.get(module, set()) - seen + if not parents and module == rel: + return None, "no derivable edges and no importers" + for parent in parents: + seen.add(parent) + edges = self._own_edges(parent) + if edges is None: + frontier.append(parent) + continue + tests, reason = edges + if tests is None: + return None, f"via {parent}: {reason}" + hits |= tests + explained.append(parent) + if not explained: + return None, "no derivable edges through importers" + return hits, f"via {sorted(explained)[:3]}" + + # -- decision ---------------------------------------------------------- + + @staticmethod + def _inert(path: str) -> bool: + return any(fnmatch.fnmatch(path, g) for g in INERT_GLOBS) + + def select(self, base: str | None, merge_commit: bool) -> Selection: + if not merge_commit and not base: + return Selection("all", self.test_files(), "no base ref") + try: + changed, deleted = self.changed(base, merge_commit) + except subprocess.CalledProcessError as err: + return Selection( + "all", self.test_files(), f"diff failed: {err.output.strip()}" + ) + messages = self.commit_messages(base, merge_commit).lower() + if any(tok in messages for tok in FORCE_ALL_TOKENS): + return Selection( + "all", self.test_files(), "commit message forces a full run" + ) + + self.build() + run_all: list[str] = [] + for path in changed + deleted: + if any(fnmatch.fnmatch(path, g) for g in RUN_ALL_GLOBS): + run_all.append(f"{path} matches a run-all pattern") + for path in deleted: + if path.startswith(SRC_ROOT + "/") and path.endswith(".py"): + run_all.append(f"deleted module {path}") + elif path.startswith("tests/") and not Path(path).name.startswith("test_"): + run_all.append(f"deleted test support {path}") + elif not self._inert(path) and not path.startswith( + ("tests/", SRC_ROOT + "/") + ): + run_all.append(f"deleted unmodeled file {path}") + + selected: dict[str, set[str]] = defaultdict(set) + in_scope = set(self.tests) + for path in changed: + if path in in_scope: + selected[path].add("changed test file") + continue + if path.startswith("tests/"): + if path.endswith(".py") and Path(path).name.startswith("test_"): + continue # a test outside this scope + if path in self.helper_tests: + for test in self.helper_tests[path]: + selected[test].add(f"{path} (helper it references)") + continue + run_all.append(f"{path} is shared test support") + continue + if path.startswith(SRC_ROOT + "/") and path.endswith(".py"): + hits, reason = self.impact_of_source(path) + if hits is None: + run_all.append(f"{path}: {reason}") + continue + for test in hits: + selected[test].add(f"{path} ({reason})") + continue + if not self._inert(path): + # package data, deepspeed configs, anything the graph does not model + run_all.append(f"{path} is not modeled") + if run_all: + return Selection("all", self.tests, "; ".join(run_all), dict(selected)) + if selected: + for test in self.opaque_tests: + selected[test].add("no visible config; rides along with every subset") + tests = sorted(selected) + if not tests: + return Selection("none", [], "no in-scope test is reachable from the diff") + if len(tests) == len(self.tests): + return Selection( + "all", tests, "every in-scope test is reachable", dict(selected) + ) + return Selection( + "subset", + tests, + f"{len(tests)} of {len(self.tests)} test files", + dict(selected), + ) + + +def _write_outputs(name: str, sel: Selection) -> None: + output = os.environ.get("GITHUB_OUTPUT") + if output: + with open(output, "a", encoding="utf-8") as fh: + fh.write(f"{name}_mode={sel.mode}\n") + fh.write(f"{name}_count={len(sel.tests)}\n") + fh.write( + f"{name}_tests={','.join(sel.tests) if sel.mode == 'subset' else ''}\n" + ) + summary = os.environ.get("GITHUB_STEP_SUMMARY") + if summary: + with open(summary, "a", encoding="utf-8") as fh: + fh.write( + f"### e2e selection `{name}`: **{sel.mode}** ({len(sel.tests)} files)\n\n" + ) + fh.write(f"{sel.reason}\n\n") + if sel.mode == "subset": + fh.writelines(f"- `{t}`\n" for t in sel.tests) + fh.write("\n") + + +def main() -> int: + parser = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + parser.add_argument("--scope", default="tests/e2e") + parser.add_argument("--exclude", action="append", default=[]) + parser.add_argument( + "--base", help="base ref; the diff is merge-base(base, HEAD)..HEAD" + ) + parser.add_argument( + "--merge-commit", + action="store_true", + help="HEAD is a PR merge commit; diff HEAD^1..HEAD", + ) + parser.add_argument( + "--head", + default="HEAD", + help="revision to diff (the working tree is still what gets parsed)", + ) + parser.add_argument("--name", default="e2e", help="prefix for GitHub outputs") + parser.add_argument( + "--output", type=Path, help="write the selected test files here" + ) + parser.add_argument("--explain", action="store_true") + parser.add_argument("--repo-root", type=Path, default=REPO_ROOT) + args = parser.parse_args() + + selector = Selector(args.repo_root, args.scope, tuple(args.exclude), args.head) + try: + sel = selector.select(args.base, args.merge_commit) + except Exception: # pylint: disable=broad-except + traceback.print_exc() + sel = Selection( + "all", selector.test_files(), "selector raised; running everything" + ) + + print(f"### MODE: {sel.mode} ({sel.reason}) ###") + print(f"### {len(sel.tests)} test file(s) in scope {args.scope} ###") + for test in sel.tests: + print(f" {test}") + if args.explain: + for why in sorted(sel.explain.get(test, ())): + print(f" <- {why}") + if args.explain and hasattr(selector, "ubiquitous"): + print(f"### ubiquitous keys (no signal): {sorted(selector.ubiquitous)}") + registries = sorted( + d for d, is_reg in selector._registry_cache.items() if is_reg + ) # pylint: disable=protected-access + print(f"### registry directories: {registries}") + print(f"### opaque tests (always in a subset): {selector.opaque_tests}") + if args.output: + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text("".join(f"{t}\n" for t in sel.tests), encoding="utf-8") + args.output.with_suffix(".mode").write_text(f"{sel.mode}\n", encoding="utf-8") + _write_outputs(args.name, sel) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/cicd/single_gpu.py b/cicd/single_gpu.py index 7252c5ff5..43f5cba90 100644 --- a/cicd/single_gpu.py +++ b/cicd/single_gpu.py @@ -24,14 +24,18 @@ df_args = { "AXOLOTL_ARGS": os.environ.get("AXOLOTL_ARGS", ""), "PYTORCH_VERSION": os.environ.get("PYTORCH_VERSION", "2.13.0"), "BASE_TAG": os.environ.get("BASE_TAG", "main-base-py3.12-cu130-2.13.0"), + "DEPS_HASH": os.environ.get("DEPS_HASH", "dev"), "CUDA": os.environ.get("CUDA", "130"), "GITHUB_REF": os.environ.get("GITHUB_REF", "refs/heads/main"), + "IMAGE_REF": os.environ.get("IMAGE_REF") + or os.environ.get("GITHUB_REF", "refs/heads/main"), "GITHUB_SHA": os.environ.get("GITHUB_SHA", ""), "NIGHTLY_BUILD": os.environ.get("NIGHTLY_BUILD", ""), "CODECOV_TOKEN": os.environ.get("CODECOV_TOKEN", ""), "HF_HOME": "/workspace/data/huggingface-cache/hub", "PYTHONUNBUFFERED": os.environ.get("PYTHONUNBUFFERED", "1"), "DEEPSPEED_LOG_LEVEL": os.environ.get("DEEPSPEED_LOG_LEVEL", "WARNING"), + "E2E_SELECTED_TESTS": os.environ.get("E2E_SELECTED_TESTS", ""), } dockerfile_contents = df_template.render(**df_args) @@ -60,9 +64,67 @@ GPU_TYPE = os.environ.get("GPU_TYPE", "L40S") GPU_CONFIG = f"{GPU_TYPE}:{N_GPUS}" +REFRESH_MARKER = ".ci_source_refreshed" + + +def _git(run_folder: str, *args: str, check: bool = True, capture: bool = False): + import subprocess # nosec + + return subprocess.run( # nosec + ["git", *args], + cwd=run_folder, + check=check, + capture_output=capture, + text=True, + ) + + +def checkout_source(run_folder: str, ref: str, sha: str) -> None: + _git(run_folder, "fetch", "--depth=1", "origin", f"+{ref}") + ref_head = _git(run_folder, "rev-parse", "FETCH_HEAD", capture=True).stdout.strip() + target = sha + if ( + _git(run_folder, "cat-file", "-e", f"{sha}^{{commit}}", check=False).returncode + != 0 + ): + if ( + _git( + run_folder, "fetch", "--depth=1", "origin", sha, check=False + ).returncode + != 0 + ): + # refs/pull/N/merge is regenerated when the base moves; the old merge SHA is then unreachable + if not ref.startswith("refs/pull/"): + raise RuntimeError(f"{sha} not reachable from {ref}") + print(f"WARNING: {sha} not fetchable, testing tip of {ref}", flush=True) + target = ref_head + _git(run_folder, "checkout", "-f", target) + _git(run_folder, "log", "-1", "--oneline") + + +def refresh_source(run_folder: str) -> None: + import subprocess # nosec + + sha = os.environ.get("GITHUB_SHA", "") + if not sha: + return + marker = pathlib.Path(run_folder) / REFRESH_MARKER + if marker.exists() and marker.read_text().strip() == sha: + return + checkout_source(run_folder, os.environ.get("GITHUB_REF", "refs/heads/main"), sha) + subprocess.run( # nosec + ["uv", "pip", "install", "--no-build-isolation", "--no-deps", "-e", "."], + cwd=run_folder, + check=True, + ) + marker.write_text(sha) + + def run_cmd(cmd: str, run_folder: str): import subprocess # nosec + refresh_source(run_folder) + sp_env = os.environ.copy() sp_env["AXOLOTL_DATASET_NUM_PROC"] = "8" diff --git a/pyproject.toml b/pyproject.toml index 7c6089077..cda60d886 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -260,6 +260,8 @@ markers = [ "distributed_cpu: CPU multi-process (gloo) tests unrelated to NF4 loading", "perf: marks tests as performance/memory regression guards (Blackwell only, large E)", "integration: marks tests as integration tests (real model modeling paths)", + "gpu: tests that require a CUDA device; auto-applied from CUDA skipif reasons (see tests/conftest.py)", + "nvfp4: NVFP4 multi-GPU tests that need SM100+ GPUs; run by cicd/nvfp4.sh", ] # UV specific configuration diff --git a/scripts/check_installed_axolotl.py b/scripts/check_installed_axolotl.py new file mode 100644 index 000000000..31e383124 --- /dev/null +++ b/scripts/check_installed_axolotl.py @@ -0,0 +1,109 @@ +"""Smoke-check an installed axolotl: import every module and resolve the entry points.""" + +import argparse +import importlib +import importlib.metadata as md +import importlib.util +import pkgutil +import sys + +ENTRY_POINT_GROUPS = ( + "axolotl.plugins", + "axolotl.model_support", + "axolotl.cloud_providers", + "axolotl.cli_commands", +) +REQUIRED_ENTRY_POINTS = {"axolotl.plugins": 2, "axolotl.cloud_providers": 3} +# trl removed its vLLM weight-sync server; these are being replaced by the native vLLM path +KNOWN_BROKEN_MODULES = { + "axolotl.scripts.vllm_serve_lora", + "axolotl.scripts.vllm_worker_ext", +} + + +def import_walk() -> bool: + import axolotl + + optional: list[tuple[str, str]] = [] + failed: list[tuple[str, str]] = [] + known: list[tuple[str, str]] = [] + names = [ + m.name + for m in pkgutil.walk_packages( + axolotl.__path__, "axolotl.", onerror=lambda _name: None + ) + ] + for name in names: + try: + importlib.import_module(name) + except Exception as exc: # noqa: BLE001 + top = (getattr(exc, "name", None) or "").partition(".")[0] + absent = ( + isinstance(exc, ModuleNotFoundError) + and top + and not top.startswith("axolotl") + and importlib.util.find_spec(top) is None + ) + bucket = ( + known + if name in KNOWN_BROKEN_MODULES + else optional + if absent + else failed + ) + bucket.append((name, f"{type(exc).__name__}: {exc}")) + ok = len(names) - len(optional) - len(failed) - len(known) + print( + f"import-walk: ok={ok} optional={len(optional)} " + f"known-broken={len(known)} failed={len(failed)}" + ) + for name, why in optional: + print(f" optional: {name}: {why}") + for name, why in known: + print(f" known-broken: {name}: {why}") + for name, why in failed: + print(f" FAILED: {name}: {why}") + return not failed + + +def check_entry_points() -> bool: + ok = True + for group in ENTRY_POINT_GROUPS: + eps = [ + ep for ep in md.distribution("axolotl").entry_points if ep.group == group + ] + print(f"{group}: {len(eps)}") + for ep in eps: + try: + ep.load() + except Exception as exc: # noqa: BLE001 + ok = False + print(f" FAILED: {ep.value}: {type(exc).__name__}: {exc}") + required = REQUIRED_ENTRY_POINTS.get(group, 0) + if len(eps) < required: + ok = False + print(f" FAILED: expected at least {required} entry points in {group}") + return ok + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--import-walk", action="store_true") + parser.add_argument("--entry-points", action="store_true") + parser.add_argument("--require-site-packages", action="store_true") + args = parser.parse_args() + + ok = True + if args.require_site_packages: + import axolotl + + assert "site-packages" in axolotl.__file__, axolotl.__file__ + if args.import_walk: + ok = import_walk() and ok + if args.entry_points: + ok = check_entry_points() and ok + sys.exit(0 if ok else 1) + + +if __name__ == "__main__": + main() diff --git a/scripts/check_sdist_manifest.py b/scripts/check_sdist_manifest.py new file mode 100644 index 000000000..527de7e8c --- /dev/null +++ b/scripts/check_sdist_manifest.py @@ -0,0 +1,146 @@ +#!/usr/bin/env python +"""Check that a wheel built from the sdist ships every git-tracked file under src/axolotl.""" + +import argparse +import fnmatch +import os +import shutil +import subprocess # nosec B404 +import sys +import tarfile +import tempfile +import zipfile +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parents[1] +PACKAGE_DIR = "src/axolotl" +WHEEL_PREFIX = "axolotl/" +SDIST_REQUIRED = ("AGENTS.md", "README.md", "LICENSE", "VERSION") +SDIST_REQUIRED_GLOBS = ("docs/agents/*.md",) +GIT_LOCATION_VARS = ( + "GIT_DIR", + "GIT_WORK_TREE", + "GIT_INDEX_FILE", + "GIT_COMMON_DIR", + "GIT_OBJECT_DIRECTORY", + "GIT_ALTERNATE_OBJECT_DIRECTORIES", + "GIT_NAMESPACE", +) + + +def tracked_package_files(repo_root: Path) -> set[str]: + env = {k: v for k, v in os.environ.items() if k not in GIT_LOCATION_VARS} + out = subprocess.check_output( # nosec B603 B607 + ["git", "ls-files", PACKAGE_DIR], cwd=repo_root, text=True, env=env + ) + files = set() + for line in out.splitlines(): + rel = line[len(PACKAGE_DIR) + 1 :] + parts = rel.split("/") + if "__pycache__" in parts or "tests" in parts[:-1] or rel.endswith(".pyc"): + continue + files.add(rel) + return files + + +def wheel_package_files(wheel: Path) -> set[str]: + with zipfile.ZipFile(wheel) as zf: + return { + n[len(WHEEL_PREFIX) :] for n in zf.namelist() if n.startswith(WHEEL_PREFIX) + } + + +def compare(tracked: set[str], shipped: set[str]) -> tuple[list[str], list[str]]: + return sorted(tracked - shipped), sorted(shipped - tracked) + + +def check_sdist_contents(sdist: Path) -> list[str]: + """Return the required sdist entries that are absent.""" + with tarfile.open(sdist) as tf: + names = tf.getnames() + top = names[0].split("/")[0] + problems = [ + f"{top}/{name}" for name in SDIST_REQUIRED if f"{top}/{name}" not in names + ] + for pattern in SDIST_REQUIRED_GLOBS: + if not fnmatch.filter(names, f"{top}/{pattern}"): + problems.append(f"{top}/{pattern}") + return problems + + +def report(missing: list[str], extra: list[str], tracked: int, shipped: int) -> int: + for path in missing: + print(f"MISSING (tracked, not in wheel): {path}") + for path in extra: + print(f"extra (in wheel, not git-tracked): {path}") + print( + f"sdist-manifest: tracked={tracked} wheel={shipped} " + f"missing={len(missing)} extra={len(extra)}" + ) + return 1 if missing else 0 + + +def build(kind: str, source: Path, out_dir: Path) -> Path: + if shutil.which("uv"): + cmd = ["uv", "build", f"--{kind}", "--out-dir", str(out_dir), str(source)] + else: + cmd = [ + sys.executable, + "-m", + "build", + f"--{kind}", + "--outdir", + str(out_dir), + str(source), + ] + subprocess.run(cmd, check=True) # nosec B603 + pattern = "*.tar.gz" if kind == "sdist" else "*.whl" + produced = list(out_dir.glob(pattern)) + assert len(produced) == 1, f"expected one {pattern} in {out_dir}, got {produced}" + return produced[0] + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--sdist", type=Path, help="check this prebuilt sdist") + parser.add_argument("--keep", action="store_true", help="keep the temp dir") + parser.add_argument("--out-dir", type=Path, help="copy sdist and wheel here") + parser.add_argument("--repo-root", type=Path, default=REPO_ROOT) + args = parser.parse_args() + + work = Path(tempfile.mkdtemp(prefix="sdist-check-")) + rc = 0 + try: + sdist = args.sdist or build("sdist", args.repo_root, work / "dist") + problems = check_sdist_contents(sdist) + for problem in problems: + print(f"MISSING from sdist: {problem}") + with tarfile.open(sdist) as tf: + tf.extractall(work / "src", filter="data") + (sdist_dir,) = [p for p in (work / "src").iterdir() if p.is_dir()] + wheel = build("wheel", sdist_dir, work / "wheel") + missing, extra = compare( + tracked_package_files(args.repo_root), wheel_package_files(wheel) + ) + rc = report( + missing, + extra, + len(tracked_package_files(args.repo_root)), + len(wheel_package_files(wheel)), + ) + if problems: + rc = 1 + if args.out_dir: + args.out_dir.mkdir(parents=True, exist_ok=True) + shutil.copy2(sdist, args.out_dir) + shutil.copy2(wheel, args.out_dir) + finally: + if args.keep: + print(f"kept: {work}") + else: + shutil.rmtree(work, ignore_errors=True) + sys.exit(rc) + + +if __name__ == "__main__": + main() diff --git a/src/axolotl/integrations/liger/models/base.py b/src/axolotl/integrations/liger/models/base.py index a9dbe9412..9c91578e8 100644 --- a/src/axolotl/integrations/liger/models/base.py +++ b/src/axolotl/integrations/liger/models/base.py @@ -5,8 +5,8 @@ Generic FLCE patch for untested models similar to Llama from typing import Optional, Tuple, Union import torch +from liger_kernel.transformers.fsdp import _FSDPForwardRedirection from liger_kernel.transformers.model.loss_utils import LigerForCausalLMLoss -from liger_kernel.transformers.trainer.orpo_trainer import _FSDPForwardRedirection from liger_kernel.utils import PEFT_AVAILABLE from peft.utils import ModulesToSaveWrapper from torch.distributed.fsdp import FullyShardedDataParallel diff --git a/src/axolotl/model_support/kimi_linear/tokenization_kimi.py b/src/axolotl/model_support/kimi_linear/tokenization_kimi.py index 83f7ab4ae..fe3ec75ad 100644 --- a/src/axolotl/model_support/kimi_linear/tokenization_kimi.py +++ b/src/axolotl/model_support/kimi_linear/tokenization_kimi.py @@ -23,7 +23,7 @@ from typing import ( import tiktoken from tiktoken.load import load_tiktoken_bpe from tokenizers import AddedToken -from transformers.models.gpt2.tokenization_gpt2 import bytes_to_unicode +from transformers.convert_slow_tokenizer import bytes_to_unicode from transformers.tokenization_utils import PreTrainedTokenizer logger = getLogger(__name__) diff --git a/src/axolotl/monkeypatch/attention/flex_attn.py b/src/axolotl/monkeypatch/attention/flex_attn.py index 0071c4223..ab8c50ce9 100644 --- a/src/axolotl/monkeypatch/attention/flex_attn.py +++ b/src/axolotl/monkeypatch/attention/flex_attn.py @@ -5,7 +5,7 @@ import sys import torch import transformers from packaging import version -from transformers.utils.import_utils import _torch_version, is_torch_less_or_equal +from transformers.utils.import_utils import is_torch_less_or_equal from axolotl.utils.logging import get_logger @@ -55,7 +55,10 @@ def patch_flex_wrapper(**flex_attn_compile_kwargs): # In PyTorch 2.6.0, there's a known issue with flex attention compilation which may # cause errors. The suggested fix is to compile with "max-autotune-no-cudagraphs" # see https://github.com/pytorch/pytorch/issues/146260 for training - elif version.parse(_torch_version).base_version == "2.6.0" and training: + elif ( + version.parse(torch.__version__).base_version == "2.6.0" + and training + ): self._compiled_flex_attention = torch.compile( flex_attention, dynamic=False, mode="max-autotune-no-cudagraphs" ) diff --git a/src/axolotl/monkeypatch/gradient_checkpointing/offload_cpu.py b/src/axolotl/monkeypatch/gradient_checkpointing/offload_cpu.py index 143015dc1..594082600 100644 --- a/src/axolotl/monkeypatch/gradient_checkpointing/offload_cpu.py +++ b/src/axolotl/monkeypatch/gradient_checkpointing/offload_cpu.py @@ -1,5 +1,7 @@ """CPU offloaded checkpointing""" +__ci_config_keys__ = ("activation_offloading",) + # Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); diff --git a/src/axolotl/monkeypatch/gradient_checkpointing/offload_disk.py b/src/axolotl/monkeypatch/gradient_checkpointing/offload_disk.py index 7386c17d2..9b5ba2b5a 100644 --- a/src/axolotl/monkeypatch/gradient_checkpointing/offload_disk.py +++ b/src/axolotl/monkeypatch/gradient_checkpointing/offload_disk.py @@ -2,6 +2,8 @@ DISCO - DIsk-based Storage and Checkpointing with Optimized prefetching """ +__ci_config_keys__ = ("activation_offloading",) + # Copyright 2025 Axolotl AI. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); diff --git a/src/axolotl/utils/schemas/validation.py b/src/axolotl/utils/schemas/validation.py index 267a276c2..300e00078 100644 --- a/src/axolotl/utils/schemas/validation.py +++ b/src/axolotl/utils/schemas/validation.py @@ -2,6 +2,7 @@ import fnmatch import json +import os import tempfile from collections.abc import Callable from pathlib import Path @@ -161,16 +162,62 @@ def _flash_attn_kernel_failure(attn_implementation: str) -> str | None: repo_id = FLASH_ATTN_KERNEL_FALLBACK.get(attn_implementation) if repo_id is None: return None + version: int | None = None try: + # resolving the pinned version lists hub refs, so it fails offline too + version = get_attn_kernel_version(repo_id) # Direct kernels calls do not read Transformers' allow_all_hub_kernels flag. - get_kernel( - repo_id, version=get_attn_kernel_version(repo_id), trust_remote_code=True - ) + get_kernel(repo_id, version=version, trust_remote_code=True) except Exception as err: # noqa: BLE001 - return f"{type(err).__name__}: {err}" + try: + _get_kernel_from_cache(repo_id, version) + except Exception: # noqa: BLE001 + return f"{type(err).__name__}: {err}" + LOG.warning( + "kernels hub lookup for %s failed (%s: %s); using the cached build", + repo_id, + type(err).__name__, + err, + ) return None +def _get_kernel_from_cache(repo_id: str, version: int | None): + """Load a hub kernel from a cached snapshot without touching the network.""" + from huggingface_hub import constants + from huggingface_hub.file_download import repo_folder_name + from kernels import get_local_kernel + + cache_dir = Path(os.environ.get("KERNELS_CACHE") or constants.HF_HUB_CACHE) + # get_kernel(version=N) downloads by commit and writes no refs entry, so the + # kernels offline resolver cannot map the version; read the snapshots directly + candidates: list[Path] = [] + for repo_type in ("kernel", "model"): + repo_dir = cache_dir / repo_folder_name(repo_id=repo_id, repo_type=repo_type) + ref = repo_dir / "refs" / f"v{version}" + if version is not None and ref.is_file(): + pinned = repo_dir / "snapshots" / ref.read_text().strip() + if pinned.is_dir(): + candidates.append(pinned) + snapshots = repo_dir / "snapshots" + if snapshots.is_dir(): + candidates.extend( + sorted( + (p for p in snapshots.iterdir() if (p / "build").is_dir()), + key=lambda p: p.stat().st_mtime, + reverse=True, + ) + ) + errors = [] + for snapshot in dict.fromkeys(candidates): + try: + return get_local_kernel(snapshot, trust_remote_code=True) + except Exception as err: # noqa: BLE001 + errors.append(f"{snapshot.name}: {type(err).__name__}: {err}") + detail = f": {'; '.join(errors)}" if errors else "" + raise FileNotFoundError(f"no loadable cached snapshot of {repo_id}{detail}") + + SUPPORTED_METRICS = {"sacrebleu", "comet", "ter", "chrf", "perplexity"} FSDP1_MIGRATION_DOCS = "docs/multi-gpu.qmd#sec-migrate-fsdp1-fsdp2" diff --git a/tests/conftest.py b/tests/conftest.py index 5910be771..de7be6b21 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -5,6 +5,7 @@ import functools import importlib import logging import os +import re import shutil import sys import tempfile @@ -79,6 +80,41 @@ _CUDA_FATAL_MARKERS = ( ) _cuda_context_poisoned = False +GPU_SKIP_REASON_RE = re.compile(r"\b(?:cuda|gpu)s?\b", re.IGNORECASE) +_GPU_MARKER_PROPERTY = "axolotl_gpu_marker_missing" +_gpu_marker_missing: dict[str, str] = {} + + +def is_gpu_skip_reason(reason) -> bool: + return bool(reason) and GPU_SKIP_REASON_RE.search(str(reason)) is not None + + +def skip_marker_text(mark) -> str: + parts = [str(mark.kwargs.get("reason", ""))] + parts.extend(arg for arg in mark.args if isinstance(arg, str)) + return " ".join(parts) + + +def item_is_cuda_gated(item) -> bool: + return any( + mark.name in ("skipif", "skip") and is_gpu_skip_reason(skip_marker_text(mark)) + for mark in item.iter_markers() + ) + + +def skip_reason_from_report(report) -> str: + longrepr = getattr(report, "longrepr", None) + if isinstance(longrepr, tuple) and len(longrepr) == 3: + return str(longrepr[2]) + return str(longrepr or "") + + +@pytest.hookimpl(tryfirst=True) +def pytest_collection_modifyitems(config, items): # pylint: disable=unused-argument + for item in items: + if item.get_closest_marker("gpu") is None and item_is_cuda_gated(item): + item.add_marker(pytest.mark.gpu) + @pytest.hookimpl(hookwrapper=True) def pytest_runtest_makereport(item, call): # pylint: disable=unused-argument @@ -88,6 +124,41 @@ def pytest_runtest_makereport(item, call): # pylint: disable=unused-argument if report.failed and call.excinfo is not None: if any(marker in str(call.excinfo.value) for marker in _CUDA_FATAL_MARKERS): _cuda_context_poisoned = True + if ( + report.skipped + and not _cuda_context_poisoned + and item.get_closest_marker("gpu") is None + ): + reason = skip_reason_from_report(report) + if is_gpu_skip_reason(reason): + report.user_properties.append((_GPU_MARKER_PROPERTY, reason)) + + +def pytest_runtest_logreport(report): + for key, reason in report.user_properties: + if key == _GPU_MARKER_PROPERTY: + _gpu_marker_missing[report.nodeid] = reason + + +def pytest_sessionfinish(session, exitstatus): # pylint: disable=unused-argument + if hasattr(session.config, "workerinput") or not _gpu_marker_missing: + return + reporter = session.config.pluginmanager.get_plugin("terminalreporter") + lines = [ + f"[gpu-marker] {nodeid}: {reason}" + for nodeid, reason in sorted(_gpu_marker_missing.items()) + ] + header = ( + f"[gpu-marker] {len(lines)} CUDA-gated test(s) skipped without @pytest.mark.gpu" + ) + if reporter is not None: + reporter.write_sep("-", header) + for line in lines: + reporter.write_line(line) + else: + print(header, *lines, sep="\n") + if os.environ.get("AXOLOTL_CI_ENFORCE_GPU_MARKER") == "1": + session.exitstatus = 1 def pytest_runtest_setup(item): diff --git a/tests/e2e/multigpu/solo/test_native_nvfp4_deepspeed_lora_merge_aware.py b/tests/e2e/multigpu/solo/test_native_nvfp4_deepspeed_lora_merge_aware.py index 2de66800d..5eb3df736 100644 --- a/tests/e2e/multigpu/solo/test_native_nvfp4_deepspeed_lora_merge_aware.py +++ b/tests/e2e/multigpu/solo/test_native_nvfp4_deepspeed_lora_merge_aware.py @@ -9,6 +9,8 @@ from pathlib import Path import pytest import torch +pytestmark = pytest.mark.nvfp4 + def _require_dynamic_nvfp4(): if any(torch.cuda.get_device_capability(index)[0] < 10 for index in range(2)): diff --git a/tests/e2e/multigpu/solo/test_native_nvfp4_deepspeed_lora_merge_lifecycle.py b/tests/e2e/multigpu/solo/test_native_nvfp4_deepspeed_lora_merge_lifecycle.py index 6a95eb640..f3c19ada9 100644 --- a/tests/e2e/multigpu/solo/test_native_nvfp4_deepspeed_lora_merge_lifecycle.py +++ b/tests/e2e/multigpu/solo/test_native_nvfp4_deepspeed_lora_merge_lifecycle.py @@ -9,6 +9,8 @@ from pathlib import Path import pytest import torch +pytestmark = pytest.mark.nvfp4 + def _require_dynamic_nvfp4(): if any(torch.cuda.get_device_capability(index)[0] < 10 for index in range(2)): diff --git a/tests/e2e/multigpu/solo/test_native_nvfp4_dynamic_ste.py b/tests/e2e/multigpu/solo/test_native_nvfp4_dynamic_ste.py index 40031deaf..c6d7b1524 100644 --- a/tests/e2e/multigpu/solo/test_native_nvfp4_dynamic_ste.py +++ b/tests/e2e/multigpu/solo/test_native_nvfp4_dynamic_ste.py @@ -8,6 +8,8 @@ from pathlib import Path import pytest +pytestmark = pytest.mark.nvfp4 + @pytest.mark.skipif(not __import__("torch").cuda.is_available(), reason="requires CUDA") def test_native_nvfp4_dynamic_ste_non_target_lora_gradients(tmp_path): diff --git a/tests/e2e/multigpu/solo/test_native_nvfp4_fsdp2_lora_parity.py b/tests/e2e/multigpu/solo/test_native_nvfp4_fsdp2_lora_parity.py index 0710cc520..ab06bbd57 100644 --- a/tests/e2e/multigpu/solo/test_native_nvfp4_fsdp2_lora_parity.py +++ b/tests/e2e/multigpu/solo/test_native_nvfp4_fsdp2_lora_parity.py @@ -8,6 +8,8 @@ from pathlib import Path import pytest +pytestmark = pytest.mark.nvfp4 + @pytest.mark.parametrize("offload", [False, True]) @pytest.mark.parametrize("dynamic_activation", [False, True]) diff --git a/tests/e2e/multigpu/solo/test_native_nvfp4_fsdp2_merge_aware_resume.py b/tests/e2e/multigpu/solo/test_native_nvfp4_fsdp2_merge_aware_resume.py index 2042429a1..bdba481e0 100644 --- a/tests/e2e/multigpu/solo/test_native_nvfp4_fsdp2_merge_aware_resume.py +++ b/tests/e2e/multigpu/solo/test_native_nvfp4_fsdp2_merge_aware_resume.py @@ -8,6 +8,8 @@ from pathlib import Path import pytest +pytestmark = pytest.mark.nvfp4 + @pytest.mark.parametrize( ("cpu_ram_efficient", "cpu_offload"), diff --git a/tests/e2e/multigpu/solo/test_native_nvfp4_fsdp2_recipe.py b/tests/e2e/multigpu/solo/test_native_nvfp4_fsdp2_recipe.py index 153296eda..621709031 100644 --- a/tests/e2e/multigpu/solo/test_native_nvfp4_fsdp2_recipe.py +++ b/tests/e2e/multigpu/solo/test_native_nvfp4_fsdp2_recipe.py @@ -8,6 +8,8 @@ from pathlib import Path import pytest +pytestmark = pytest.mark.nvfp4 + @pytest.mark.skipif(not __import__("torch").cuda.is_available(), reason="requires CUDA") def test_native_nvfp4_fsdp2_recipe(tmp_path): diff --git a/tests/e2e/multigpu/solo/test_native_nvfp4_tp_hf_load.py b/tests/e2e/multigpu/solo/test_native_nvfp4_tp_hf_load.py index 5e34aae31..f97e0caa8 100644 --- a/tests/e2e/multigpu/solo/test_native_nvfp4_tp_hf_load.py +++ b/tests/e2e/multigpu/solo/test_native_nvfp4_tp_hf_load.py @@ -7,6 +7,8 @@ from pathlib import Path import pytest +pytestmark = pytest.mark.nvfp4 + @pytest.mark.skipif( not __import__("torch").cuda.is_available(), reason="requires two CUDA GPUs" diff --git a/tests/e2e/multigpu/solo/test_native_nvfp4_tp_lora_parity.py b/tests/e2e/multigpu/solo/test_native_nvfp4_tp_lora_parity.py index 4e755ee1d..220bb6a4c 100644 --- a/tests/e2e/multigpu/solo/test_native_nvfp4_tp_lora_parity.py +++ b/tests/e2e/multigpu/solo/test_native_nvfp4_tp_lora_parity.py @@ -8,6 +8,8 @@ from pathlib import Path import pytest +pytestmark = pytest.mark.nvfp4 + @pytest.mark.parametrize("dtype", ["float32", "bfloat16"]) @pytest.mark.skipif( diff --git a/tests/e2e/multigpu/solo/test_sonicmoe_nvfp4_fsdp2.py b/tests/e2e/multigpu/solo/test_sonicmoe_nvfp4_fsdp2.py index 5f607b0dd..666a45c2a 100644 --- a/tests/e2e/multigpu/solo/test_sonicmoe_nvfp4_fsdp2.py +++ b/tests/e2e/multigpu/solo/test_sonicmoe_nvfp4_fsdp2.py @@ -8,6 +8,8 @@ from pathlib import Path import pytest +pytestmark = pytest.mark.nvfp4 + def test_sonicmoe_grouped_lora_dispatch_keeps_scaling_separate(): from ._sonicmoe_nvfp4_fsdp2_worker import _grouped_lora diff --git a/tests/e2e/multigpu/solo/test_torchao_lora_deepspeed.py b/tests/e2e/multigpu/solo/test_torchao_lora_deepspeed.py index 158088e34..9796744f2 100644 --- a/tests/e2e/multigpu/solo/test_torchao_lora_deepspeed.py +++ b/tests/e2e/multigpu/solo/test_torchao_lora_deepspeed.py @@ -8,6 +8,8 @@ from pathlib import Path import pytest import torch +pytestmark = pytest.mark.nvfp4 + @pytest.mark.parametrize("stage", [1, 2]) @pytest.mark.skipif( diff --git a/tests/e2e/multigpu/solo/test_torchao_lora_deepspeed_checkpoint.py b/tests/e2e/multigpu/solo/test_torchao_lora_deepspeed_checkpoint.py index 599d41e08..5fbe5937b 100644 --- a/tests/e2e/multigpu/solo/test_torchao_lora_deepspeed_checkpoint.py +++ b/tests/e2e/multigpu/solo/test_torchao_lora_deepspeed_checkpoint.py @@ -6,6 +6,8 @@ from pathlib import Path import pytest import torch +pytestmark = pytest.mark.nvfp4 + @pytest.mark.skipif(torch.cuda.device_count() < 2, reason="two CUDA devices required") @pytest.mark.parametrize("stage", [1, 2]) diff --git a/tests/e2e/multigpu/solo/test_torchao_lora_deepspeed_zero3_checkpoint.py b/tests/e2e/multigpu/solo/test_torchao_lora_deepspeed_zero3_checkpoint.py index cc495754a..047472572 100644 --- a/tests/e2e/multigpu/solo/test_torchao_lora_deepspeed_zero3_checkpoint.py +++ b/tests/e2e/multigpu/solo/test_torchao_lora_deepspeed_zero3_checkpoint.py @@ -8,6 +8,8 @@ from pathlib import Path import pytest import torch +pytestmark = pytest.mark.nvfp4 + def _dynamic_nvfp4_supported(): return ( diff --git a/tests/e2e/utils.py b/tests/e2e/utils.py index f49143f0b..be0b34a2d 100644 --- a/tests/e2e/utils.py +++ b/tests/e2e/utils.py @@ -50,9 +50,26 @@ def flash_attn_available() -> bool: try: from transformers.utils import is_flash_attn_2_available - return is_flash_attn_2_available(kernels_fallback_ok=True) + if is_flash_attn_2_available(kernels_fallback_ok=True): + return True + except Exception: # pylint: disable=broad-except + pass + # transformers' probe lists hub refs over the network and caches a False on + # any error; the validator's check also accepts a cached snapshot + try: + from transformers.integrations.hub_kernels import allow_all_hub_kernels + + from axolotl.utils.schemas.validation import _flash_attn_kernel_failure + + with allow_all_hub_kernels(): + available = _flash_attn_kernel_failure("flash_attention_2") is None except Exception: # pylint: disable=broad-except return False + if available: + cache_clear = getattr(is_flash_attn_2_available, "cache_clear", None) + if cache_clear is not None: + cache_clear() + return available # pytest.mark works on plain (non-TestCase) classes; unittest.skipUnless does not. diff --git a/tests/integrations/kernels/test_merge_aware_linear.py b/tests/integrations/kernels/test_merge_aware_linear.py index 0566e2df7..29bad3a69 100644 --- a/tests/integrations/kernels/test_merge_aware_linear.py +++ b/tests/integrations/kernels/test_merge_aware_linear.py @@ -297,6 +297,7 @@ def test_mixed_merge_aware_qkv_keeps_other_projections_optimized(monkeypatch): ) +@pytest.mark.gpu @pytest.mark.parametrize("autocast", [False, True]) @pytest.mark.parametrize("device", ["cpu", "cuda"]) def test_fp32_adapter_matches_export_effective_weight(autocast, device): diff --git a/tests/integrations/test_expert_parallel_adapter_resume.py b/tests/integrations/test_expert_parallel_adapter_resume.py index 610c8855e..14d2fac14 100644 --- a/tests/integrations/test_expert_parallel_adapter_resume.py +++ b/tests/integrations/test_expert_parallel_adapter_resume.py @@ -10,6 +10,8 @@ from peft import LoraConfig, PeftModel, get_peft_model from axolotl.integrations.expert_parallel.shard import ep_local_adapter_dir from axolotl.utils.dict import DictDefault +pytestmark = pytest.mark.gpu + E_GLOBAL, EP_SIZE, HIDDEN, RANK = 8, 2, 6, 2 diff --git a/tests/integrations/test_gemma4_moe.py b/tests/integrations/test_gemma4_moe.py index e3408c782..c3923af73 100644 --- a/tests/integrations/test_gemma4_moe.py +++ b/tests/integrations/test_gemma4_moe.py @@ -19,6 +19,8 @@ import torch import torch.nn.functional as F from torch import nn +pytestmark = pytest.mark.gpu + # ============================================================================ # Gemma4 reference implementation (extracted from transformers) # ============================================================================ diff --git a/tests/integrations/test_sonicmoe.py b/tests/integrations/test_sonicmoe.py index b73195987..ad1bfcd4b 100644 --- a/tests/integrations/test_sonicmoe.py +++ b/tests/integrations/test_sonicmoe.py @@ -206,6 +206,7 @@ class TestExpertsClassMetadata: Verify our forward respects these without an actual CUDA kernel call. """ + @pytest.mark.gpu def test_non_gated_requires_relu2(self): # non-gated experts are supported only with relu²; any other act must fail loudly. from axolotl.integrations.kernels.libs.sonicmoe.experts import ( diff --git a/tests/integrations/test_sonicmoe_fp4_cute_prep.py b/tests/integrations/test_sonicmoe_fp4_cute_prep.py index 64abbc3f7..efa92548f 100644 --- a/tests/integrations/test_sonicmoe_fp4_cute_prep.py +++ b/tests/integrations/test_sonicmoe_fp4_cute_prep.py @@ -179,6 +179,7 @@ def test_gate_up_interleave_perm(): assert torch.equal(wi[inv], w) +@pytest.mark.gpu def test_fp4_cute_unavailable_without_sm100(): from axolotl.integrations.kernels.libs.sonicmoe.fp4_cute import fp4_cute_available diff --git a/tests/kernels/test_rms_norm_gated.py b/tests/kernels/test_rms_norm_gated.py index 10b379184..fc488fde5 100644 --- a/tests/kernels/test_rms_norm_gated.py +++ b/tests/kernels/test_rms_norm_gated.py @@ -10,10 +10,12 @@ import torch.nn.functional as F pytest.importorskip("triton", reason="triton required for fused kernels") +pytestmark = pytest.mark.gpu + if not torch.cuda.is_available(): pytest.skip("CUDA required for fused kernel tests", allow_module_level=True) -from axolotl.kernels.rms_norm_gated import FusedRMSNormGated +from axolotl.kernels.rms_norm_gated import FusedRMSNormGated # noqa: E402 class EagerRMSNormGated(torch.nn.Module): diff --git a/tests/loaders/test_nf4_adapter_init_dequant.py b/tests/loaders/test_nf4_adapter_init_dequant.py index 38c9039b1..942599fe9 100644 --- a/tests/loaders/test_nf4_adapter_init_dequant.py +++ b/tests/loaders/test_nf4_adapter_init_dequant.py @@ -8,6 +8,8 @@ import torch from axolotl.utils.dict import DictDefault from axolotl.utils.nf4 import BnbNF4Parametrization, TorchaoNF4Parametrization +pytestmark = pytest.mark.gpu + def _save_llama(path): from transformers import LlamaConfig, LlamaForCausalLM diff --git a/tests/monkeypatch/test_checkpoint_activation_offload.py b/tests/monkeypatch/test_checkpoint_activation_offload.py index 856449287..22974038c 100644 --- a/tests/monkeypatch/test_checkpoint_activation_offload.py +++ b/tests/monkeypatch/test_checkpoint_activation_offload.py @@ -1,5 +1,6 @@ from functools import partial +import pytest import torch from torch import nn from torch.utils.checkpoint import checkpoint @@ -16,6 +17,7 @@ class TinyCheckpointLayer(GradientCheckpointingLayer): return hidden_states * hidden_states +@pytest.mark.gpu def test_checkpoint_offload_marks_non_reentrant_checkpoint_input(): device = "cuda" if torch.cuda.is_available() else "cpu" layer = TinyCheckpointLayer() diff --git a/tests/monkeypatch/test_lora_fp32_gradients.py b/tests/monkeypatch/test_lora_fp32_gradients.py index 9af29840f..2076d1e91 100644 --- a/tests/monkeypatch/test_lora_fp32_gradients.py +++ b/tests/monkeypatch/test_lora_fp32_gradients.py @@ -184,6 +184,7 @@ def test_fused_groups_keep_fp32_parameter_gradients(kind, monkeypatch): assert relative_error < 0.03 +@pytest.mark.gpu @pytest.mark.parametrize("loaded", [False, True]) @pytest.mark.parametrize("enabled", [False, True]) def test_loader_precision_for_new_and_saved_adapters(tmp_path, loaded, enabled): diff --git a/tests/monkeypatch/test_lora_kernels_nvfp4.py b/tests/monkeypatch/test_lora_kernels_nvfp4.py index fc0c530e3..684047570 100644 --- a/tests/monkeypatch/test_lora_kernels_nvfp4.py +++ b/tests/monkeypatch/test_lora_kernels_nvfp4.py @@ -8,6 +8,7 @@ from test_lora_fp32_gradients import projection from torch import nn +@pytest.mark.gpu @pytest.mark.parametrize("kind", ["linear", "qkv", "swiglu", "geglu"]) @pytest.mark.parametrize("device", ["cpu", "cuda"]) @pytest.mark.parametrize("fp32", [False, True]) diff --git a/tests/monkeypatch/test_offload_amp_device_type.py b/tests/monkeypatch/test_offload_amp_device_type.py index 9742c56dd..6287d030c 100644 --- a/tests/monkeypatch/test_offload_amp_device_type.py +++ b/tests/monkeypatch/test_offload_amp_device_type.py @@ -1,5 +1,6 @@ """AMP decorators in the gradient-checkpointing offload paths.""" +import pytest import torch from axolotl.kernels.utils import torch_amp_custom_fwd @@ -23,6 +24,7 @@ def test_amp_decorators_resolve_a_valid_device_type(): torch.testing.assert_close(double(torch.ones(3)), torch.full((3,), 2.0)) +@pytest.mark.gpu def test_cpu_offloaded_checkpointer_matches_baseline(): """Forward/backward roundtrip on whichever device this test runs on.""" device = "cuda" if torch.cuda.is_available() else "cpu" diff --git a/tests/monkeypatch/test_sdpa_varlen.py b/tests/monkeypatch/test_sdpa_varlen.py index f047ecc29..15df99722 100644 --- a/tests/monkeypatch/test_sdpa_varlen.py +++ b/tests/monkeypatch/test_sdpa_varlen.py @@ -10,9 +10,12 @@ from axolotl.monkeypatch.attention.sdpa_varlen import ( varlen_available, ) -pytestmark = pytest.mark.skipif( - not varlen_available(), reason="torch.nn.attention.varlen needs torch >= 2.10" -) +pytestmark = [ + pytest.mark.skipif( + not varlen_available(), reason="torch.nn.attention.varlen needs torch >= 2.10" + ), + pytest.mark.gpu, +] requires_cuda = pytest.mark.skipif( not torch.cuda.is_available(), reason="varlen_attn needs CUDA" diff --git a/tests/scripts/__init__.py b/tests/scripts/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/scripts/test_check_sdist_manifest.py b/tests/scripts/test_check_sdist_manifest.py new file mode 100644 index 000000000..c7c6a8bc1 --- /dev/null +++ b/tests/scripts/test_check_sdist_manifest.py @@ -0,0 +1,64 @@ +"""Unit tests for scripts/check_sdist_manifest.py.""" + +import importlib.util +import io +import tarfile +import zipfile +from pathlib import Path + +SCRIPT = Path(__file__).resolve().parents[2] / "scripts" / "check_sdist_manifest.py" +_spec = importlib.util.spec_from_file_location("check_sdist_manifest", SCRIPT) +assert _spec is not None and _spec.loader is not None +mod = importlib.util.module_from_spec(_spec) +_spec.loader.exec_module(mod) + + +def test_compare_reports_missing_and_extra(): + assert mod.compare({"a.py", "b/c.yaml"}, {"a.py", "d.py"}) == ( + ["b/c.yaml"], + ["d.py"], + ) + + +def test_compare_clean(): + assert mod.compare({"a.py"}, {"a.py"}) == ([], []) + + +def test_wheel_package_files_strips_prefix(tmp_path): + wheel = tmp_path / "axolotl-1.0-py3-none-any.whl" + with zipfile.ZipFile(wheel, "w") as zf: + zf.writestr("axolotl/x.py", "") + zf.writestr("axolotl/sub/y.yaml", "") + zf.writestr("axolotl-1.0.dist-info/RECORD", "") + assert mod.wheel_package_files(wheel) == {"x.py", "sub/y.yaml"} + + +def test_tracked_package_files_filters(monkeypatch): + out = ( + "src/axolotl/a.py\n" + "src/axolotl/__pycache__/a.pyc\n" + "src/axolotl/x/tests/t.py\n" + "src/axolotl/b.yaml\n" + ) + monkeypatch.setattr(mod.subprocess, "check_output", lambda *a, **k: out) + assert mod.tracked_package_files(Path(".")) == {"a.py", "b.yaml"} + + +def test_report_fails_only_on_missing(capsys): + assert mod.report(["gone.yaml"], [], 2, 1) == 1 + assert "missing=1" in capsys.readouterr().out + assert mod.report([], ["untracked.py"], 1, 2) == 0 + + +def test_check_sdist_contents(tmp_path): + sdist = tmp_path / "pkg-1.tar.gz" + with tarfile.open(sdist, "w:gz") as tf: + info = tarfile.TarInfo("pkg-1/AGENTS.md") + tf.addfile(info, io.BytesIO(b"")) + problems = mod.check_sdist_contents(sdist) + assert problems == [ + "pkg-1/README.md", + "pkg-1/LICENSE", + "pkg-1/VERSION", + "pkg-1/docs/agents/*.md", + ] diff --git a/tests/scripts/test_deps_hash.py b/tests/scripts/test_deps_hash.py new file mode 100644 index 000000000..58ad806c3 --- /dev/null +++ b/tests/scripts/test_deps_hash.py @@ -0,0 +1,135 @@ +"""Unit tests for cicd/deps_hash.py.""" + +import datetime as dt +import importlib.util +import shutil +import subprocess # nosec B404 +import sys +from pathlib import Path + +SCRIPT = Path(__file__).resolve().parents[2] / "cicd" / "deps_hash.py" +_spec = importlib.util.spec_from_file_location("deps_hash", SCRIPT) +assert _spec is not None and _spec.loader is not None +mod = importlib.util.module_from_spec(_spec) +sys.modules["deps_hash"] = mod +_spec.loader.exec_module(mod) + +REPO = SCRIPT.parents[1] + + +def _copy_inputs(tmp_path: Path) -> Path: + root = tmp_path / "repo" + for rel in ( + "pyproject.toml", + "scripts/cutcrossentropy_install.py", + "cicd/Dockerfile-uv.jinja", + ): + (root / rel).parent.mkdir(parents=True, exist_ok=True) + shutil.copy(REPO / rel, root / rel) + return root + + +TODAY = dt.date(2026, 10, 4) + + +def test_marker_or_ruff_edits_keep_the_hash(tmp_path): + root = _copy_inputs(tmp_path) + before = mod.compute(root, {}, today=TODAY) + pyproject = root / "pyproject.toml" + pyproject.write_text( + pyproject.read_text().replace( + "[tool.pytest.ini_options]", + "[tool.pytest.ini_options]\nxfail_strict = true", + ) + ) + assert mod.compute(root, {}, today=TODAY) == before + + +def test_dependency_edits_change_the_hash(tmp_path): + root = _copy_inputs(tmp_path) + before = mod.compute(root, {}, today=TODAY) + pyproject = root / "pyproject.toml" + pyproject.write_text( + pyproject.read_text().replace( + "dependencies = [", 'dependencies = [\n "left-pad==1.0",', 1 + ) + ) + assert mod.compute(root, {}, today=TODAY) != before + + +def test_dockerfile_extras_and_nightly_date_participate(tmp_path): + root = _copy_inputs(tmp_path) + base = mod.compute(root, {}, today=TODAY) + assert mod.compute(root, {"AXOLOTL_EXTRAS": "vllm"}, today=TODAY) != base + (root / "cicd/Dockerfile-uv.jinja").write_text("FROM scratch\n") + assert mod.compute(root, {}, today=TODAY) != base + day1 = mod.compute(root, {"NIGHTLY_BUILD": "true"}, today=dt.date(2026, 1, 1)) + day2 = mod.compute(root, {"NIGHTLY_BUILD": "true"}, today=dt.date(2026, 1, 2)) + assert day1 != day2 + + +def test_shared_image_expires_weekly(tmp_path): + root = _copy_inputs(tmp_path) + same_week = mod.compute(root, {}, today=dt.date(2026, 10, 5)) == mod.compute( + root, {}, today=dt.date(2026, 10, 9) + ) + next_week = mod.compute(root, {}, today=dt.date(2026, 10, 5)) == mod.compute( + root, {}, today=dt.date(2026, 10, 12) + ) + assert same_week and not next_week + + +def _git(repo: Path, *args: str) -> None: + subprocess.run( # nosec B603 B607 + ["git", *args], + cwd=repo, + check=True, + capture_output=True, + env={ + "GIT_AUTHOR_NAME": "t", + "GIT_AUTHOR_EMAIL": "t@t", + "GIT_COMMITTER_NAME": "t", + "GIT_COMMITTER_EMAIL": "t@t", + "HOME": str(repo), + "PATH": "/usr/bin:/bin", + }, + ) + + +def _git_repo_with_main(tmp_path: Path) -> Path: + root = _copy_inputs(tmp_path) + _git(root, "init", "-q", "-b", "main") + _git(root, "add", "-A") + _git(root, "commit", "-qm", "base") + # origin/main is what image_ref compares against; point it at the local main + _git(root, "update-ref", "refs/remotes/origin/main", "HEAD") + _git(root, "checkout", "-qb", "feature") + return root + + +def test_image_ref_is_main_when_dependencies_match(tmp_path): + root = _git_repo_with_main(tmp_path) + env = {"GITHUB_REF": "refs/pull/1/merge"} + pyproject = root / "pyproject.toml" + pyproject.write_text(pyproject.read_text() + "\n# docs only\n") + deps_hash = mod.compute(root, env, today=TODAY) + assert mod.image_ref(root, env, deps_hash, today=TODAY) == "refs/heads/main" + + +def test_image_ref_is_the_branch_when_dependencies_differ(tmp_path): + root = _git_repo_with_main(tmp_path) + env = {"GITHUB_REF": "refs/pull/1/merge"} + pyproject = root / "pyproject.toml" + pyproject.write_text( + pyproject.read_text().replace( + "dependencies = [", 'dependencies = [\n "left-pad==1.0",', 1 + ) + ) + deps_hash = mod.compute(root, env, today=TODAY) + assert mod.image_ref(root, env, deps_hash, today=TODAY) == "refs/pull/1/merge" + + +def test_image_ref_falls_back_to_the_branch_without_main(tmp_path): + root = _copy_inputs(tmp_path) + env = {"GITHUB_REF": "refs/pull/1/merge"} + assert mod.image_ref(root, env, "abc", today=TODAY) == "refs/pull/1/merge" diff --git a/tests/scripts/test_select_e2e_tests.py b/tests/scripts/test_select_e2e_tests.py new file mode 100644 index 000000000..e1000b5bf --- /dev/null +++ b/tests/scripts/test_select_e2e_tests.py @@ -0,0 +1,404 @@ +"""Unit tests for cicd/select_e2e_tests.py against synthetic repositories.""" + +import importlib.util +import subprocess # nosec B404 +import sys +from pathlib import Path + +import pytest + +SCRIPT = Path(__file__).resolve().parents[2] / "cicd" / "select_e2e_tests.py" +_spec = importlib.util.spec_from_file_location("select_e2e_tests", SCRIPT) +assert _spec is not None and _spec.loader is not None +mod = importlib.util.module_from_spec(_spec) +sys.modules["select_e2e_tests"] = ( + mod # dataclasses resolve postponed annotations via sys.modules +) +_spec.loader.exec_module(mod) + +SCHEMA = """ +class Cfg: + base_model: str + sequence_len: int + lora_r: int + gradient_checkpointing: bool +""" +BUILDER = "def build(cfg):\n return cfg.base_model, cfg.sequence_len\n" +ADAPTER = ( + "from axolotl.utils.collators import pad\n\ndef load(cfg):\n return cfg.lora_r\n" +) +COLLATORS = "def pad(x):\n return x\n" +OFFLOAD = "def offload(cfg):\n return cfg.gradient_checkpointing\n" +KD_ARGS = "class KDArgs:\n kd_temperature: float\n" +KD_PLUGIN = "def hook(cfg):\n return cfg.kd_temperature, cfg.sequence_len\n" +STRATEGY = "def load(cfg):\n return cfg.sequence_len\n" + + +def _cfg_test(name: str, extra: str = "", values: str = "") -> str: + return f''' +from axolotl.utils.dict import DictDefault + +def test_{name}(): + cfg = DictDefault({{ + "base_model": "org/{name}-model", + "sequence_len": 512, + "datasets": [{{"path": "ds", "type": "{values or "alpaca"}"}}], + {extra} + }}) + assert cfg +''' + + +def _git(repo: Path, *args: str) -> None: + subprocess.run( # nosec B603 B607 + ["git", *args], + cwd=repo, + check=True, + capture_output=True, + env={ + "GIT_AUTHOR_NAME": "t", + "GIT_AUTHOR_EMAIL": "t@t", + "GIT_COMMITTER_NAME": "t", + "GIT_COMMITTER_EMAIL": "t@t", + "HOME": str(repo), + "PATH": "/usr/bin:/bin", + }, + ) + + +def _write(repo: Path, files: dict[str, str]) -> None: + for rel, content in files.items(): + path = repo / rel + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(content) + + +@pytest.fixture(name="repo") +def _repo(tmp_path: Path) -> Path: + repo = tmp_path / "repo" + _write( + repo, + { + "pyproject.toml": "[project]\nname='axolotl'\n", + "docs/readme.md": "hi\n", + "src/axolotl/__init__.py": "", + "src/axolotl/utils/__init__.py": "", + "src/axolotl/utils/dict.py": "class DictDefault(dict):\n pass\n", + "src/axolotl/utils/collators.py": COLLATORS, + "src/axolotl/utils/schemas/__init__.py": "", + "src/axolotl/utils/schemas/config.py": SCHEMA, + "src/axolotl/core/__init__.py": "", + "src/axolotl/core/builder.py": BUILDER, + "src/axolotl/loaders/__init__.py": "", + "src/axolotl/loaders/adapter.py": ADAPTER, + "src/axolotl/monkeypatch/__init__.py": "", + "src/axolotl/monkeypatch/offload.py": OFFLOAD, + "src/axolotl/prompt_strategies/__init__.py": "", + "src/axolotl/prompt_strategies/alpaca.py": STRATEGY, + "src/axolotl/prompt_strategies/chat_template.py": STRATEGY, + "src/axolotl/prompt_strategies/sharegpt.py": STRATEGY, + "src/axolotl/prompt_strategies/completion.py": STRATEGY, + "src/axolotl/integrations/__init__.py": "", + "src/axolotl/integrations/kd/__init__.py": "", + "src/axolotl/integrations/kd/args.py": KD_ARGS, + "src/axolotl/integrations/kd/plugin.py": KD_PLUGIN, + "src/axolotl/integrations/liger/__init__.py": "", + "src/axolotl/integrations/liger/plugin.py": "def hook(cfg):\n return 1\n", + "src/axolotl/integrations/cce/__init__.py": "", + "src/axolotl/integrations/cce/plugin.py": "def hook(cfg):\n return 1\n", + "tests/conftest.py": "import pytest\n", + "tests/e2e/test_sft.py": _cfg_test("sft", values="chat_template"), + "tests/e2e/test_lora.py": _cfg_test( + "lora", '"lora_r": 8, "gradient_checkpointing": True,' + ), + "tests/e2e/test_kd.py": _cfg_test( + "kd", + '"plugins": ["axolotl.integrations.kd.KDPlugin"], "kd_temperature": 1.0,', + ), + "tests/e2e/test_liger.py": _cfg_test( + "liger", '"plugins": ["axolotl.integrations.liger.LigerPlugin"],' + ), + "tests/e2e/multigpu/test_fsdp.py": _cfg_test("fsdp"), + "tests/e2e/test_completion.py": _cfg_test( + "completion", values="completion" + ), + "tests/e2e/test_cce.py": _cfg_test( + "cce", '"plugins": ["axolotl.integrations.cce.CCEPlugin"],' + ), + }, + ) + _git(repo, "init", "-q", "-b", "main") + _git(repo, "add", "-A") + _git(repo, "commit", "-qm", "base") + _git(repo, "checkout", "-qb", "feature") + return repo + + +def _select( + repo: Path, + change: dict[str, str], + delete: tuple[str, ...] = (), + message: str = "change", +): + _write(repo, change) + for rel in delete: + (repo / rel).unlink() + _git(repo, "add", "-A") + _git(repo, "commit", "-qm", message) + selector = mod.Selector(repo, "tests/e2e", exclude=("tests/e2e/multigpu",)) + return selector.select("main", merge_commit=False) + + +def test_feature_key_selects_only_tests_setting_it(repo): + sel = _select(repo, {"src/axolotl/loaders/adapter.py": ADAPTER + "\n# touched\n"}) + assert sel.mode == "subset" + assert sel.tests == ["tests/e2e/test_lora.py"] + + +def test_ubiquitous_key_reader_runs_everything(repo): + sel = _select(repo, {"src/axolotl/core/builder.py": BUILDER + "\n# touched\n"}) + assert sel.mode == "all" + assert "ubiquitous" in sel.reason + + +def test_edgeless_module_inherits_from_importer(repo): + sel = _select(repo, {"src/axolotl/utils/collators.py": COLLATORS + "\n# touched\n"}) + assert sel.mode == "subset" + assert sel.tests == ["tests/e2e/test_lora.py"] + + +def test_plugin_named_by_config_value_wins_over_ubiquity(repo): + sel = _select( + repo, {"src/axolotl/integrations/kd/plugin.py": KD_PLUGIN + "\n# touched\n"} + ) + assert sel.mode == "subset" + assert sel.tests == ["tests/e2e/test_kd.py"] + + +def test_brand_new_plugin_needs_no_registration(repo): + sel = _select( + repo, + { + "src/axolotl/integrations/newthing/__init__.py": "", + "src/axolotl/integrations/newthing/args.py": "class NewArgs:\n newthing_alpha: float\n", + "src/axolotl/integrations/newthing/plugin.py": "def hook(cfg):\n return cfg.newthing_alpha, cfg.base_model\n", + "tests/e2e/test_newthing.py": _cfg_test( + "newthing", + '"plugins": ["axolotl.integrations.newthing.NewPlugin"], "newthing_alpha": 0.5,', + ), + }, + ) + assert sel.mode == "subset" + assert sel.tests == ["tests/e2e/test_newthing.py"] + + +def test_registry_leaf_nobody_names_selects_nothing(repo): + sel = _select( + repo, {"src/axolotl/prompt_strategies/sharegpt.py": STRATEGY + "\n# touched\n"} + ) + assert sel.mode == "none" + + +def test_registry_leaf_named_by_dataset_type(repo): + sel = _select( + repo, + {"src/axolotl/prompt_strategies/chat_template.py": STRATEGY + "\n# touched\n"}, + ) + assert sel.mode == "subset" + assert sel.tests == ["tests/e2e/test_sft.py"] + + +def test_changed_test_file_selects_itself(repo): + sel = _select(repo, {"tests/e2e/test_liger.py": _cfg_test("liger", '"x": 1,')}) + assert sel.mode == "subset" + assert sel.tests == ["tests/e2e/test_liger.py"] + + +def test_out_of_scope_test_change_selects_nothing(repo): + sel = _select( + repo, {"tests/e2e/multigpu/test_fsdp.py": _cfg_test("fsdp", '"x": 1,')} + ) + assert sel.mode == "none" + + +def test_docs_only_selects_nothing(repo): + sel = _select(repo, {"docs/readme.md": "changed\n"}) + assert sel.mode == "none" + + +@pytest.mark.parametrize( + "change", + [ + {"pyproject.toml": "[project]\nname='axolotl'\nversion='1'\n"}, + {"tests/conftest.py": "import pytest\n# touched\n"}, + {"tests/e2e/utils.py": "def helper():\n return 1\n"}, + ], +) +def test_shared_files_run_everything(repo, change): + sel = _select(repo, change) + assert sel.mode == "all" + + +def test_deleted_module_runs_everything(repo): + sel = _select(repo, {}, delete=("src/axolotl/monkeypatch/offload.py",)) + assert sel.mode == "all" + assert "deleted module" in sel.reason + + +def test_orphan_module_runs_everything(repo): + sel = _select( + repo, {"src/axolotl/monkeypatch/offload.py": OFFLOAD + "\n# touched\n"} + ) + # reads a discriminating key, so it selects the tests setting it even without importers + assert sel.mode == "subset" + assert sel.tests == ["tests/e2e/test_lora.py"] + sel = _select( + repo, {"src/axolotl/monkeypatch/orphan.py": "def f():\n return 1\n"} + ) + assert sel.mode == "all" + assert "no importers" in sel.reason + + +def test_commit_tag_forces_full_run(repo): + sel = _select(repo, {"docs/readme.md": "x\n"}, message="docs [test all]") + assert sel.mode == "all" + + +def test_run_all_reason_lists_every_trigger(repo): + sel = _select( + repo, + { + "pyproject.toml": "[project]\nname='axolotl'\nversion='1'\n", + "src/axolotl/core/builder.py": BUILDER + "#\n", + }, + ) + assert sel.mode == "all" + assert "pyproject.toml" in sel.reason and "builder.py" in sel.reason + + +def test_declared_keys_are_authoritative(repo): + sel = _select( + repo, + { + "src/axolotl/monkeypatch/offload.py": '__ci_config_keys__ = ("gradient_checkpointing",)\n' + + "def offload(cfg):\n return cfg.base_model\n" + }, + ) + assert sel.mode == "subset" + assert sel.tests == ["tests/e2e/test_lora.py"] + + +def test_declared_unknown_key_runs_everything(repo): + sel = _select( + repo, + { + "src/axolotl/monkeypatch/offload.py": '__ci_config_keys__ = ("no_such_key",)\n' + + OFFLOAD + }, + ) + assert sel.mode == "all" + assert "no_such_key" in sel.reason + + +def _grow_base(repo: Path, files: dict[str, str]) -> None: + """Add files to the base branch so they are not part of the feature diff.""" + _write(repo, files) + _git(repo, "add", "-A") + _git(repo, "commit", "-qm", "base grows") + _git(repo, "branch", "-f", "main", "HEAD") + + +WORKER_TEST = ( + "from pathlib import Path\n\n" + "def test_launch():\n" + ' worker = Path(__file__).with_name("_fast_worker.py")\n' + " assert worker.exists()\n" +) + + +def test_worker_script_config_counts_for_its_test(repo): + _grow_base( + repo, + { + "tests/e2e/_fast_worker.py": _cfg_test( + "worker", '"gradient_checkpointing": True,' + ), + "tests/e2e/test_worker_launcher.py": WORKER_TEST, + }, + ) + sel = _select( + repo, {"src/axolotl/monkeypatch/offload.py": OFFLOAD + "\n# touched\n"} + ) + assert sel.mode == "subset" + assert sel.tests == ["tests/e2e/test_lora.py", "tests/e2e/test_worker_launcher.py"] + + +def test_changed_worker_selects_the_tests_that_launch_it(repo): + _grow_base( + repo, + { + "tests/e2e/_fast_worker.py": _cfg_test("worker"), + "tests/e2e/test_worker_launcher.py": WORKER_TEST, + }, + ) + sel = _select(repo, {"tests/e2e/_fast_worker.py": _cfg_test("worker", '"x": 1,')}) + assert sel.mode == "subset" + assert sel.tests == ["tests/e2e/test_worker_launcher.py"] + + +def test_opaque_test_rides_along_with_every_subset(repo): + _grow_base( + repo, {"tests/e2e/test_opaque.py": "def test_opaque():\n assert True\n"} + ) + sel = _select(repo, {"src/axolotl/loaders/adapter.py": ADAPTER + "\n# touched\n"}) + assert sel.mode == "subset" + assert sel.tests == ["tests/e2e/test_lora.py", "tests/e2e/test_opaque.py"] + + +def test_opaque_test_does_not_turn_none_into_subset(repo): + _grow_base( + repo, {"tests/e2e/test_opaque.py": "def test_opaque():\n assert True\n"} + ) + sel = _select(repo, {"docs/readme.md": "changed again\n"}) + assert sel.mode == "none" + + +def test_deleted_helper_runs_everything(repo): + _grow_base( + repo, + { + "tests/e2e/_fast_worker.py": _cfg_test("worker"), + "tests/e2e/test_worker_launcher.py": WORKER_TEST, + }, + ) + sel = _select(repo, {}, delete=("tests/e2e/_fast_worker.py",)) + assert sel.mode == "all" + assert "deleted test support" in sel.reason + + +@pytest.mark.parametrize( + "path", + [ + "src/axolotl/utils/chat_templates/templates/new.jinja", + "deepspeed_configs/zero2.json", + "cicd/anything.txt", + ], +) +def test_unmodeled_files_run_everything(repo, path): + sel = _select(repo, {path: "data\n"}) + assert sel.mode == "all" + assert "not modeled" in sel.reason or "run-all" in sel.reason + + +@pytest.mark.parametrize( + "path", + [ + "docs/guide.qmd", + "examples/llama/lora.yml", + "src/axolotl/integrations/kd/README.md", + "LICENSE", + ], +) +def test_inert_files_select_nothing(repo, path): + sel = _select(repo, {path: "text\n"}) + assert sel.mode == "none" diff --git a/tests/test_gpu_marker.py b/tests/test_gpu_marker.py new file mode 100644 index 000000000..ad7172214 --- /dev/null +++ b/tests/test_gpu_marker.py @@ -0,0 +1,51 @@ +"""Tests for the gpu-marker helpers in tests/conftest.py.""" + +import pytest + +from tests.conftest import is_gpu_skip_reason, skip_marker_text + + +@pytest.mark.parametrize( + "reason", + [ + "CUDA required", + "requires two CUDA GPUs", + "FLA kernels require CUDA", + "Requires CUDA GPU", + "needs cuda", + "requires at least 16 GiB free GPU memory", + "Need >=2 GPUs for FSDP2 multi-rank tests", + "varlen_attn needs CUDA", + ], +) +def test_gpu_reasons_match(reason): + assert is_gpu_skip_reason(reason) + + +@pytest.mark.parametrize( + "reason", + [ + "swanlab package not installed", + "xformers not available", + "torchao required", + "deep_ep not installed", + "Not running in CI cache preload", + "cudatoolkit missing", + "gpuless", + "", + None, + ], +) +def test_non_gpu_reasons_do_not_match(reason): + assert not is_gpu_skip_reason(reason) + + +def test_skip_marker_text_reads_reason_and_string_conditions(): + mark = pytest.mark.skipif("not torch.cuda.is_available()", reason="needs it").mark + text = skip_marker_text(mark) + assert "cuda" in text and "needs it" in text + assert ( + skip_marker_text(pytest.mark.skipif(True, reason="CUDA required").mark) + == "CUDA required" + ) + assert skip_marker_text(pytest.mark.skip(reason="flaky").mark) == "flaky" diff --git a/tests/test_lora.py b/tests/test_lora.py index 50cbea9bc..30cc95a00 100644 --- a/tests/test_lora.py +++ b/tests/test_lora.py @@ -2,6 +2,8 @@ tests for loading loras """ +import pytest + from axolotl.loaders import ModelLoader, load_tokenizer from axolotl.utils.config import normalize_config, validate_config from axolotl.utils.dict import DictDefault @@ -22,6 +24,7 @@ minimal_config = DictDefault( ) +@pytest.mark.gpu class TestLoRALoad: """ Test class for loading LoRA weights diff --git a/tests/utils/schemas/validation/test_config_validators.py b/tests/utils/schemas/validation/test_config_validators.py index d989f8e4e..4e00d59c8 100644 --- a/tests/utils/schemas/validation/test_config_validators.py +++ b/tests/utils/schemas/validation/test_config_validators.py @@ -428,10 +428,38 @@ class TestFlashAttnAvailabilityMessage: raise RuntimeError("HTTP 429 Too Many Requests") monkeypatch.setattr(kernels, "get_kernel", failing_get_kernel) + from axolotl.utils.schemas import validation + + def no_cached_build(*_, **__): + raise FileNotFoundError("no loadable cached snapshot") + + monkeypatch.setattr(validation, "_get_kernel_from_cache", no_cached_build) cfg = min_base_cfg | DictDefault(attn_implementation="flash_attention_2") with pytest.raises(ValueError, match="HTTP 429 Too Many Requests"): validate_config(cfg) + def test_hub_failure_falls_back_to_cached_build(self, min_base_cfg, monkeypatch): + import kernels + import torch + from transformers import utils as transformers_utils + + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr( + transformers_utils, "is_flash_attn_2_available", lambda **_: False + ) + + def failing_get_kernel(*_, **__): + raise RuntimeError("HTTP 429 Too Many Requests") + + monkeypatch.setattr(kernels, "get_kernel", failing_get_kernel) + from axolotl.utils.schemas import validation + + monkeypatch.setattr( + validation, "_get_kernel_from_cache", lambda *_, **__: object() + ) + cfg = min_base_cfg | DictDefault(attn_implementation="flash_attention_2") + validate_config(cfg) + def test_successful_lookup_overrides_transformers_verdict( self, min_base_cfg, monkeypatch ): diff --git a/tests/utils/schemas/validation/test_flash_attn_availability.py b/tests/utils/schemas/validation/test_flash_attn_availability.py index 314c37d7f..41422845f 100644 --- a/tests/utils/schemas/validation/test_flash_attn_availability.py +++ b/tests/utils/schemas/validation/test_flash_attn_availability.py @@ -20,6 +20,8 @@ class TestFlashAttnAvailabilityValidator: import kernels import transformers.utils + import axolotl.utils.schemas.validation as validation + monkeypatch.setattr( transformers.utils, "is_flash_attn_2_available", lambda **_: value ) @@ -28,11 +30,12 @@ class TestFlashAttnAvailabilityValidator: ) if not value: # a hub lookup that succeeds overrides transformers' verdict, so - # "unavailable" has to fail there as well + # "unavailable" has to fail there as well, online and from the cache def no_build(*_, **__): raise FileNotFoundError("Cannot find a build variant") monkeypatch.setattr(kernels, "get_kernel", no_build) + monkeypatch.setattr(validation, "_get_kernel_from_cache", no_build) def test_fa2_unavailable_raises(self, min_base_cfg, monkeypatch): self._force_availability(monkeypatch, False) @@ -82,6 +85,8 @@ class TestFlashAttnAvailabilityValidator: import transformers.integrations.hub_kernels as hub_kernels import transformers.utils + import axolotl.utils.schemas.validation as validation + checker = Mock(return_value=False) monkeypatch.setattr( transformers.utils, f"is_flash_attn_{attn_version}_available", checker @@ -96,6 +101,11 @@ class TestFlashAttnAvailabilityValidator: return object() monkeypatch.setattr(kernels, "get_kernel", get_kernel) + + def no_cached_build(*_, **__): + raise FileNotFoundError("no loadable cached snapshot") + + monkeypatch.setattr(validation, "_get_kernel_from_cache", no_cached_build) attn_implementation = f"flash_attention_{attn_version}" cfg = min_base_cfg | DictDefault(attn_implementation=attn_implementation) if has_build: @@ -106,3 +116,107 @@ class TestFlashAttnAvailabilityValidator: with pytest.raises(ValueError, match="Cannot find a build variant"): validate_config(cfg) checker.cache_clear.assert_not_called() + + @pytest.mark.parametrize("cached", [True, False]) + def test_hub_network_failure_falls_back_to_cache( + self, min_base_cfg, monkeypatch, cached + ): + import kernels + import transformers.integrations.hub_kernels as hub_kernels + + import axolotl.utils.schemas.validation as validation + + self._force_availability(monkeypatch, False) + monkeypatch.setattr(hub_kernels, "get_attn_kernel_version", lambda _: 1) + + def online(*_, **__): + raise ConnectionError("Connection reset by peer") + + def from_cache(repo_id, version): + assert version == 1 + if not cached: + raise FileNotFoundError("Cannot find a local snapshot") + return object() + + monkeypatch.setattr(kernels, "get_kernel", online) + monkeypatch.setattr(validation, "_get_kernel_from_cache", from_cache) + cfg = min_base_cfg | DictDefault(attn_implementation="flash_attention_2") + if cached: + validated = validate_config(cfg) + assert validated.attn_implementation == "flash_attention_2" + else: + with pytest.raises(ValueError, match="Connection reset by peer"): + validate_config(cfg) + + @pytest.mark.parametrize("cached", [True, False]) + def test_version_lookup_failure_falls_back_to_cache( + self, min_base_cfg, monkeypatch, cached + ): + """Resolving the pinned version lists hub refs, so it fails offline as well.""" + import kernels + import transformers.integrations.hub_kernels as hub_kernels + + import axolotl.utils.schemas.validation as validation + + self._force_availability(monkeypatch, False) + + def offline(_): + raise ConnectionError("Name or service not known") + + def never(*_, **__): + raise AssertionError("get_kernel must not run without a version") + + def from_cache(repo_id, version): + assert version is None + if not cached: + raise FileNotFoundError("Cannot find a local snapshot") + return object() + + monkeypatch.setattr(hub_kernels, "get_attn_kernel_version", offline) + monkeypatch.setattr(kernels, "get_kernel", never) + monkeypatch.setattr(validation, "_get_kernel_from_cache", from_cache) + cfg = min_base_cfg | DictDefault(attn_implementation="flash_attention_2") + if cached: + assert validate_config(cfg).attn_implementation == "flash_attention_2" + else: + with pytest.raises(ValueError, match="Name or service not known"): + validate_config(cfg) + + def test_cache_lookup_prefers_version_ref_then_newest_build( + self, monkeypatch, tmp_path + ): + import kernels + from huggingface_hub import constants + + import axolotl.utils.schemas.validation as validation + + repo = tmp_path / "kernels--kernels-community--flash-attn2" + for sha in ("aaa", "bbb"): + (repo / "snapshots" / sha / "build").mkdir(parents=True) + (repo / "snapshots" / "ccc").mkdir() # no build dir: never a candidate + monkeypatch.setattr(constants, "HF_HUB_CACHE", str(tmp_path)) + monkeypatch.delenv("KERNELS_CACHE", raising=False) + loaded = [] + monkeypatch.setattr( + kernels, + "get_local_kernel", + lambda path, **_: loaded.append(path.name) or object(), + ) + + validation._get_kernel_from_cache("kernels-community/flash-attn2", 3) + assert loaded[-1] in {"aaa", "bbb"} + + (repo / "refs").mkdir() + (repo / "refs" / "v3").write_text("bbb") + validation._get_kernel_from_cache("kernels-community/flash-attn2", 3) + assert loaded[-1] == "bbb" + + def test_cache_lookup_without_snapshot_raises(self, monkeypatch, tmp_path): + from huggingface_hub import constants + + import axolotl.utils.schemas.validation as validation + + monkeypatch.setattr(constants, "HF_HUB_CACHE", str(tmp_path)) + monkeypatch.delenv("KERNELS_CACHE", raising=False) + with pytest.raises(FileNotFoundError, match="no loadable cached snapshot"): + validation._get_kernel_from_cache("kernels-community/flash-attn2", 3) diff --git a/tests/utils/schemas/validation/test_moe_quant.py b/tests/utils/schemas/validation/test_moe_quant.py index 904b922d9..58e48929a 100644 --- a/tests/utils/schemas/validation/test_moe_quant.py +++ b/tests/utils/schemas/validation/test_moe_quant.py @@ -162,6 +162,7 @@ class TestPeftPatchIdempotency: patch_peft_target_parameters_matching._axolotl_patched = False +@pytest.mark.gpu class TestMoeAdapterTrainMergeRoundtrip: """E2E: train adapter on quantized MoE experts, then merge onto plain model.