feat(xla): test marker, extra + CPU-PJRT CI job #1
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| name: XLA tests Workflow | |
| # Hardware-free validation of the XLA code paths (grid_sample gather branch, all_gather device | |
| # probe, XLAPrecision selection, bool-mask dtype, keypoint host-sync gates, ...). torch_xla reports | |
| # device.type == "xla" on ANY PJRT backend, so PJRT_DEVICE=CPU exercises every device-gated branch | |
| # without a TPU. Real-TPU-only concerns (MXU/TFLOPS, bf16-MXU convergence parity) stay on Colab/Kaggle. | |
| # | |
| # torch_xla ships Linux x86_64 wheels only -> ubuntu runner. Its minor must match torch's minor; | |
| # floating both to the latest in torch>=2.2,<3 keeps them aligned (torch_xla 2.9 <-> torch 2.9). | |
| # | |
| # Trigger is manual + PR for now; the job self-passes while no @pytest.mark.xla tests exist yet | |
| # (pytest exit code 5 = "no tests collected"). Flip the guard off once Phase 1 lands xla tests. | |
| on: | |
| workflow_dispatch: | |
| pull_request: | |
| branches: ["main", "release/*", "develop"] | |
| paths: | |
| - "src/rfdetr/**" | |
| - "tests/**" | |
| - "pyproject.toml" | |
| - ".github/workflows/ci-tests-xla.yml" | |
| permissions: | |
| contents: read | |
| defaults: | |
| run: | |
| shell: bash | |
| env: | |
| UV_TORCH_BACKEND: "cpu" # CPU-only torch; works with 'uv pip' | |
| PJRT_DEVICE: "CPU" # run XLA on the CPU PJRT plugin (no TPU, no libtpu) | |
| concurrency: | |
| group: pytest-xla-${{ github.ref }} | |
| cancel-in-progress: ${{ github.event_name == 'pull_request' }} | |
| jobs: | |
| run-xla-tests: | |
| name: XLA CPU-PJRT testing | |
| timeout-minutes: 25 | |
| runs-on: ubuntu-latest | |
| steps: | |
| - name: 📥 Checkout the repository | |
| uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # v6.0.1 | |
| - name: 🐍 Install uv and set Python version 3.11 | |
| uses: astral-sh/setup-uv@61cb8a9741eeb8a550a1b8544337180c0fc8476b # v7.2.0 | |
| with: | |
| version: "0.9.26" | |
| python-version: "3.11" | |
| activate-environment: true | |
| - name: 🚀 Install Packages (torch + torch_xla, CPU) | |
| timeout-minutes: 8 | |
| # [xla] pulls torch_xla; UV_TORCH_BACKEND=cpu keeps torch CPU-only. Both float to the latest | |
| # matching minor within torch>=2.2,<3, so torch_xla and torch stay ABI-aligned. | |
| run: | | |
| uv pip install -e ".[train,xla,cli]" --group tests | |
| - name: 🔎 Verify torch_xla runtime | |
| run: | | |
| uv run --no-sync python -c "import torch, torch_xla, torch_xla.core.xla_model as xm; d = xm.xla_device(); print('torch', torch.__version__, '| torch_xla', torch_xla.__version__, '| device', d); assert str(d).startswith('xla')" | |
| - name: 🧪 Run the XLA-marked Tests | |
| run: | | |
| set +e | |
| uv run --no-sync pytest src/ tests/ \ | |
| -n 1 -m "xla and not gpu" \ | |
| --ignore=tests/run_smoke_all_models.py \ | |
| --ignore=tests/legacy/test_checkpoint_compat.py \ | |
| --timeout=420 \ | |
| --durations=50 | |
| rc=$? | |
| set -e | |
| if [ "$rc" -eq 5 ]; then | |
| echo "No @pytest.mark.xla tests collected yet (Phase 1 not landed) — passing." | |
| exit 0 | |
| fi | |
| exit $rc | |
| - name: Minimize uv cache | |
| continue-on-error: true | |
| run: uv cache prune --ci |