Skip to content
CI

fix: one reader for the fused-VJP switch, make the batched knob reachable (F13, F38) #680

Sign in to view logs

fix: one reader for the fused-VJP switch, make the batched knob reachable (F13, F38)

fix: one reader for the fused-VJP switch, make the batched knob reachable (F13, F38) #680

Workflow file for this run

.github/workflows/ci.yml at b4f7925
name: CI
on:
push:
branches: [main, master]
pull_request:
jobs:
lint:
# Format/import checks are Python-version independent -- run once.
runs-on: ubuntu-latest
steps:
- name: Checkout
uses: actions/checkout@v4
- name: Setup Python
uses: actions/setup-python@v5
with:
python-version: "3.13"
cache: pip
- name: Install package + dev tools
run: |
python -m pip install --upgrade pip
pip install git+https://github.com/TobiBu/yggdrax.git
pip install -e ".[dev]"
- name: Black
run: black --check .
- name: isort
run: isort --check-only .
test-full:
# Complete suite (including the compile-bound `slow` tests) on the dev
# Python version, with coverage. This is the comprehensive gate.
#
# SHARDED, because the single job was memory-bound, not time-bound. Measured
# on its exact command: peak RSS 15.70 GB before the phase-5 tests and 16.07
# GB after -- against a 16 GB hosted runner. At 98% of capacity any addition
# OOM-kills the runner, which surfaces as "the runner has received a
# shutdown signal" / exit 143 partway through, NOT as a test failure or a
# timeout, and so reads like flaky infrastructure.
#
# A worker's RSS grows across the tests it runs (JAX holds compiled
# executables in-process), so cutting the tests per job is what actually
# bounds the peak -- unlike `-n 1`, which would halve concurrent memory but
# roughly double wall-clock into the cap.
#
# Shard by MEMORY, not by test count: the two are wildly disproportionate
# here, because a handful of float64 FMM gradient tests dominate. Measured
# per group at `-n 2` (no coverage):
#
# tests/integration 15:50 12.10 GB <- 125 tests, the hog
# tests/unit 12:16 5.84 GB <- 421 tests
# tests/test_*.py 3:04 5.18 GB <- 89 tests
# tests/characterization 1:11 3.63 GB <- 13 tests
#
# Those are the paths and counts AS MEASURED; the 20 root-level test files
# have since moved into the tiers, so `tests/test_*.py` is empty and its 89
# cases now sit under integration/distributed/unit. The relative picture is
# what this table is for -- integration dominates and needs its own shard --
# and that has not changed. The live counts are on the matrix below.
#
# A count-balanced 2-way split (unit | everything else) was tried first and
# does NOT work: it leaves 15.21 GB on one side, still at the ceiling. The
# integration tests need a shard to themselves.
#
# The memory fix proper is in `tests/conftest.py` (registering
# `test_mutual_fmm.py` for `jax.clear_caches()`), which takes
# `tests/integration` from 12.68 GB to 5.47 GB with coverage. This sharding
# is what pays for that fix's +66% wall cost, and it keeps any single
# shard's peak far enough below 16 GB that the next heavy gradient test does
# not put us straight back here.
#
# Rejected: dropping branch coverage. Measured at 9 MB of 12.7 GB (0.07%) --
# see the note in pyproject.toml's `[tool.coverage.run]`.
name: test-full (${{ matrix.shard }})
runs-on: ubuntu-latest
# Hard cap so a memory-bomb / hung-worker never burns the 6h workflow
# ceiling again (an OOM-killed xdist worker makes the job hang otherwise).
#
# 50 -> 60. The cap's job is to catch the ~2.6x CPU slowdown that is the whole
# reason the JAX ceiling is <0.11, so it has to sit BELOW that and ABOVE the
# real spread. Both bounds are now measured rather than assumed:
#
# warm on main (merge of #77) 26 min
# PR branches, Tier 1 (10 PRs) 30, 36, 39, 42, 42, 43, 43, 44, 45 min
# 2.6x the warm time ~68 min
#
# A PR branch busts the exact .jax_cache key (it is keyed on
# hashFiles('jaccpot/**/*.py'), so ANY .py edit misses) and runs off the
# prefix restore, which is why every PR sits 4-19 min above main. 60 leaves
# 15 min over the worst observed and still trips well before 68, so the
# regression detector is intact. Do not raise it past ~65 without replacing
# the detector with something explicit.
#
# That spread was measured on the UNSHARDED job. This one is sharded (below),
# so every shard runs a strict subset and sits under those numbers -- the long
# pole was integration at 33.8 min. 60 is therefore headroom here rather than a
# fitted bound, and the detector still holds where it matters: 2.6 x 33.8 is
# ~88 min, comfortably past the cap.
timeout-minutes: 60
strategy:
# Both shards always report: one shard's failure must not mask the other's
# result, since they cover disjoint halves of the suite.
fail-fast: false
matrix:
include:
# Split by path so the set is a partition with no test left behind:
# `tests/experimental` is excluded by the marker and `tests/perf` (3
# cases) by the default addopts, and everything else lives under
# exactly one of these entries or the `test-mac-runtime` job -- EXCEPT
# `tests/distributed`, which is deliberately unsharded; see below.
# Verified by collection count, with `--ignore=tests/unit/runtime` as
# the step applies it and `-m "not experimental"` applied throughout.
# Re-measured 2026-08-21:
# integration 165 + mutual-static-device 14 + unit 755 + rest 34 = 968
# + tests/unit/runtime 273 = 1241
# plus tests/distributed 25 (NOT run here) = 1266, the whole collectable
# universe.
#
# This is the check paying for itself, not a formality. `rest` used to be
# `tests/characterization tests/test_*.py`; the move of the 20 root-level
# test files into tiers emptied that glob and created `tests/distributed`,
# which belonged to no shard. Unsharded on main the directory still ran,
# so nothing was red -- it would simply have stopped being tested here.
# Re-measure when a test directory appears or moves. The previous figures
# (173 / 682 / 59 / 213 = 1127, and before that 173 / 641 / 53 / 202 =
# 1069) are what the re-measurement above replaces -- the partition was
# exact each time, only the totals had drifted.
#
# WHY `tests/distributed` IS NOT IN ANY SHARD. It cannot run here. Every
# file in it skips on `device_count() < 2` and GitHub's runners have no
# GPU, so all 25 cases collect and skip: the shard paid collection and
# import cost to assert nothing. Worse, it read as covered -- audit row
# F34 recorded `distributed/fmm.py` at 19% for exactly this reason, and
# when the tier was finally executed on two cards (2026-08-21) it failed
# 5 of 25, four of them one real cross-domain far-field defect that had
# been latent for as long as the directory had been "in CI".
#
# Removing it is therefore honest, not a reduction in coverage: there was
# none. Do not add it back to buy a green tick. It is run deliberately,
# on two cards, per `docs/refactor_audit_2026-08.md` row F34; and the
# parts of it that need no collectives have been lifted into
# `tests/integration/test_distributed_cross_domain_far_extents.py`, which
# DOES run here on one device -- that is where to put anything new that
# can be made single-device, rather than here.
# `integration` MINUS the new static-shape/device module, which gets its
# own shard below. This shard's composition is otherwise exactly main's,
# deliberately: it is memory-bound (see the note above), and bisecting
# the OOM under a 15 GB cgroup showed it lands in
# `test_mutual_fmm.py`'s own gradient/rollout tests at positions 28-38
# of the run -- so the pile is main's, and the fix is to keep NEW cases
# out of it rather than to regroup what is already here.
#
# An earlier attempt moved all the mutual files into one shard instead.
# That concentrated those hogs into a single worker and OOM-killed at
# both `-n 2` and `-n 1`. Do not group them.
- shard: integration
paths: tests/integration
extra: --ignore=tests/integration/test_mutual_fmm_static_device.py
- shard: mutual-static-device
paths: tests/integration/test_mutual_fmm_static_device.py
extra: ""
- shard: unit
paths: tests/unit
extra: ""
- shard: rest
paths: tests/characterization
extra: ""
env:
JAX_ENABLE_X64: "1"
# Persist JAX's compiled executables across runs so the heavy FMM
# compiles become warm-cache hits (conftest enables the disk cache only
# when this dir is set).
JACCPOT_TEST_JAX_CACHE_DIR: ${{ github.workspace }}/.jax_cache
steps:
- name: Checkout
uses: actions/checkout@v4
- name: Setup Python
uses: actions/setup-python@v5
with:
python-version: "3.13"
cache: pip
- name: Restore JAX compilation cache
uses: actions/cache@v4
with:
path: ${{ github.workspace }}/.jax_cache
# Recompiles are keyed on source; bust the cache when jaccpot changes.
# Keyed per shard as well: the two shards compile disjoint sets and
# would otherwise race to save the same key, so one would always lose.
key: jax-compile-py313-${{ matrix.shard }}-${{ hashFiles('jaccpot/**/*.py') }}
restore-keys: |
jax-compile-py313-${{ matrix.shard }}-
- name: Install package + dev tools
run: |
python -m pip install --upgrade pip
pip install git+https://github.com/TobiBu/yggdrax.git
pip install -e ".[dev]"
- name: Suite with coverage
# `-n 2`: the memory-limited runner OOMs with more float64-FMM workers
# (overrides pyproject's `-n auto`). `-m "not experimental"` matches the
# default addopts; all `slow` tests DO run here.
#
# `--dist loadgroup`: `-n 2` bounds the worker COUNT but not what those
# two workers hold at once, and tests/integration/test_mutual_fmm.py has
# several block-step tests whose peak RSS is measured at 1.3-1.9 GB each
# (see the `xdist_group` note there). Any two of them scheduled together
# is ~2.9-3.4 GB of solver state plus coverage, which is what crashed
# `gw1` intermittently -- on main at be172a52 and on PR #122, both times
# reported as the *test that was unlucky* rather than the heaviest one.
# `loadgroup` keeps `load`'s scheduling for everything else and only
# pins same-group tests to one worker, so the heavy ones serialise with
# each other and nothing else changes.
#
# Measured cost of serialising them, this shard end to end on one box:
# 12:03 -> 13:27, +84 s (+11.6%), identical 171 passed / 2 skipped. On the
# runner's 36 min that is ~4 min against a 60 min cap.
#
# `--ignore=tests/unit/runtime`: the Dehnen-MAC suite there is 94 cases
# that each build a solver and run a full FMM solve (26-95 s apiece on
# CPU). It has the `test-mac-runtime` job below to itself, so the `unit`
# shard must exclude it or the work runs twice -- and the shard would
# absorb ~20 min it has no cap headroom for. The ignore is harmless on
# the other two shards, whose paths never reach it.
#
# Branch coverage is NOT set here on purpose: `[tool.coverage.run]` in
# pyproject.toml sets `branch = true`, and its note records that
# `--cov-branch` on the command line cannot override that either way, so
# passing the flag would be inert decoration.
run: |
pytest -n 2 --dist loadgroup -m "not experimental" ${{ matrix.paths }} \
${{ matrix.extra }} \
--ignore=tests/unit/runtime \
--cov=jaccpot \
--cov-report=xml --cov-report=term:skip-covered
- name: Upload coverage XML artifact
# Per-shard name: each shard covers only the code its half of the suite
# reaches, so these are partial reports and must not overwrite each
# other. There is no `fail_under` gate, so a partial report fails
# nothing; combine them downstream if a total is ever wanted.
uses: actions/upload-artifact@v4
with:
name: coverage-xml-${{ matrix.shard }}
path: coverage.xml
test-mac-runtime:
# The Dehnen mass-dependent MAC suite (tests/unit/runtime), split out of
# test-full. Every case builds a solver and runs a full FMM solve: 26-95 s
# each on CPU, which is 20+ minutes this job absorbs so that test-full keeps
# both its ~36 min profile and the meaning of its cap.
#
# Collected case count, `-m "not experimental"`: 94 when this job was written,
# 213 on main, and 256 on this branch -- the eq (16b) estimator, the cache-key
# identity, the split-build and large-N-lane criterion tests, and the refuted
# `dehnen_theta` mode add 43. The cap below is NOT raised to fit them: it is a
# regression detector, and widening it to absorb new tests is what would blind
# it (the same argument test-full's `timeout-minutes` note makes). If this job
# starts timing out, the answer is another split by path, not a bigger number.
#
# Splitting also halves the COLD-cache cost per job, which matters more than it
# looks: the .jax_cache key is hashFiles('jaccpot/**/*.py'), so any runtime
# change starts cold, and a cancelled job never saves the cache -- so one
# over-budget run makes the next one cold too, indefinitely. Two smaller jobs
# each finish and each save.
runs-on: ubuntu-latest
timeout-minutes: 50
env:
JAX_ENABLE_X64: "1"
JACCPOT_TEST_JAX_CACHE_DIR: ${{ github.workspace }}/.jax_cache
steps:
- name: Checkout
uses: actions/checkout@v4
- name: Setup Python
uses: actions/setup-python@v5
with:
python-version: "3.13"
cache: pip
- name: Restore JAX compilation cache
uses: actions/cache@v4
with:
path: ${{ github.workspace }}/.jax_cache
# Separate key from test-full's: the two jobs compile disjoint sets, so
# sharing one key would have them overwrite each other's entries.
key: jax-compile-macrt-py313-${{ hashFiles('jaccpot/**/*.py') }}
restore-keys: |
jax-compile-macrt-py313-
- name: Install package + dev tools
run: |
python -m pip install --upgrade pip
pip install git+https://github.com/TobiBu/yggdrax.git
pip install -e ".[dev]"
- name: Dehnen MAC runtime suite
# Same `-n 2` memory ceiling as test-full, and `slow` runs here too: this
# job exists precisely to run the compile-bound cases.
run: |
pytest tests/unit/runtime -n 2 -m "not experimental" \
--cov=jaccpot --cov-branch \
--cov-report=xml:coverage-mac-runtime.xml \
--cov-report=term:skip-covered
- name: Upload coverage XML artifact
uses: actions/upload-artifact@v4
with:
name: coverage-xml-mac-runtime
path: coverage-mac-runtime.xml
test-smoke:
# Fast (non-slow) suite on the other supported Python versions for
# version-compatibility. The ~50 compile-bound `slow` tests run only in
# test-full (on 3.13); keeping them out of the matrix is where most of the
# CI wall-clock reduction comes from.
runs-on: ubuntu-latest
# 25 -> 30 -> 45. The 30 was still too tight, and the Tier 1 refactor made that
# unambiguous: it cancelled twice at exactly the cap (PRs #85, #86) while the
# 3.11 leg of the same runs passed with 43 SECONDS to spare. Measured, per leg:
#
# warm on main (merge of #77) 3.11 24 min 3.12 19 min
# Tier 1 PR branches (10 PRs) 3.11 21-29.5 3.12 20-30 (2 cancelled at 30)
#
# The spread is not tracking diff size -- #81 (7 renamed constants) ran 21/20
# and #88 (docstrings only) ran 29.5/28 -- so it is cache-restore luck plus
# runner speed, and no amount of tightening will make it predictable. Note the
# cause is NOT JAX's cache key including source metadata:
# `compilation_cache_include_metadata_in_key` is already False. It is that
# actions/cache keys on hashFiles('jaccpot/**/*.py'), so every PR misses the
# exact key and restores whatever the prefix last saved.
#
# 45 keeps the regression detector: 2.6x the 24 min warm time is ~62 min, so a
# JAX-0.11-class slowdown still trips this. And it breaks the self-perpetuating
# part -- a cancelled job never SAVES its cache, so the two cancellations left
# the next runs colder, which is exactly how a too-tight cap sustains itself.
#
# 45 -> 75. The 24 min warm time the paragraph above reasons from is gone: the
# 3.11 leg now lands within a minute of the cap on essentially every run.
# Measured across six recent runs (main and PR branches):
#
# 3.11 44.2 44.9 44.3 35.9 3.12 32.0 30.6 37.2 32.5
#
# PR #125 is what this cost: BOTH legs were cancelled at 45m03s with the tests
# themselves GREEN -- `914 passed, 60 skipped in 2578.14s (0:42:58)` -- and the
# job killed during the post-job cache save. A red X on a passing suite, and by
# the paragraph above, a cache that never got written.
#
# Be honest about what 75 buys: it restores headroom, NOT the detector. At a
# ~45 min warm baseline a 2.6x slowdown is ~117 min, so no cap that is also a
# sane wall-clock can still catch that class of regression here -- the detector
# was lost when the suite grew, not when this number changed. The job now
# collects 980 non-slow cases, of which tests/unit is 835.
#
# So the real fix is the one test-mac-runtime's note already prescribes for
# exactly this situation: "another split by path, not a bigger number". Smoke
# is still one un-sharded job per Python version while test-full has been split
# four ways. Splitting it (tests/unit/runtime | tests/unit | rest) would put the
# longest leg near 30 min and make a 45-50 cap meaningful again. That is a
# bigger change than unblocking the queue, so it is written down here rather
# than done silently alongside it.
timeout-minutes: 75
strategy:
fail-fast: false
matrix:
python-version: ["3.11", "3.12"]
env:
JAX_ENABLE_X64: "1"
JACCPOT_TEST_JAX_CACHE_DIR: ${{ github.workspace }}/.jax_cache
steps:
- name: Checkout
uses: actions/checkout@v4
- name: Setup Python
uses: actions/setup-python@v5
with:
python-version: ${{ matrix.python-version }}
cache: pip
- name: Restore JAX compilation cache
uses: actions/cache@v4
with:
path: ${{ github.workspace }}/.jax_cache
# Keyed per Python version as well as per source hash: the matrix legs
# compile with different interpreters, and the existing keys' note about
# disjoint sets overwriting each other applies between them too.
key: jax-compile-smoke-py${{ matrix.python-version }}-${{ hashFiles('jaccpot/**/*.py') }}
# The prefix fallback is what makes a library change only PARTIALLY cold
# instead of fully: without it, any edit under jaccpot/ would put this job
# back to a 28-minute cold run, and a cancelled job never saves a cache --
# so one over-budget run would keep the next one cold indefinitely.
restore-keys: |
jax-compile-smoke-py${{ matrix.python-version }}-
- name: Install package + dev tools
run: |
python -m pip install --upgrade pip
pip install git+https://github.com/TobiBu/yggdrax.git
pip install -e ".[dev]"
- name: Smoke tests (fast subset)
# Override the default addopts marker to also drop `slow`.
run: pytest -n 2 -m "not slow and not experimental"
test-distributed-mutual:
# The ONLY job that runs the distributed mutual FMM, because it is the only
# one with more than one device.
#
# `tests/integration/test_mutual_distributed.py` is guarded by
# `skipif(device_count() < 2)`, and every other job leaves JAX on its single
# default CPU device -- so before this job existed the suite was collected and
# skipped on every single run, and a green CI said nothing at all about the
# lane. That is not hypothetical: the cross-domain ownership rule was wrong for
# a merged LET coarse tree (two devices numbering the same remote node
# differently, so they disagreed about who owns a pair), and it was caught by
# hand locally, not here. The force was 8.6e-3 wrong while global momentum
# stayed at 3.4e-17 and the cross-pair count stayed exactly right, so nothing
# cheaper than this suite would have noticed.
#
# `--xla_force_host_platform_device_count=4` fakes four CPU devices, which is
# enough for the collectives to be real: the halo import and the reverse halo
# both go through a ragged all-to-all, and the pair partition needs at least
# two domains to have anything to partition. It is NOT a substitute for real
# multi-GPU -- NCCL is not exercised here -- but it is what makes the
# correctness criteria (force vs an exact direct sum, GLOBAL momentum, overflow
# reduced across devices) run at all.
#
# Its own job rather than an env var on an existing one: the flag changes the
# device topology for the whole process, so folding it into a shared tier would
# silently re-target every other test in that tier.
#
# Small on purpose -- 16 particles per device, seconds of compute -- so it adds
# no memory pressure to a workflow already sharded against a 16 GB ceiling.
runs-on: ubuntu-latest
# Measured 6:09 for the four tests serially on a fast workstation; hosted
# runners are the slower side of that, and every test compiles its own
# shard_map program. 35 leaves room for a 3-5x runner penalty without being
# so wide it stops detecting a real hang.
timeout-minutes: 35
env:
JAX_ENABLE_X64: "1"
JAX_PLATFORMS: cpu
XLA_FLAGS: --xla_force_host_platform_device_count=4
steps:
- name: Checkout
uses: actions/checkout@v4
- name: Setup Python
uses: actions/setup-python@v5
with:
python-version: "3.13"
cache: pip
- name: Install package + dev tools
run: |
python -m pip install --upgrade pip
pip install git+https://github.com/TobiBu/yggdrax.git
pip install -e ".[dev]"
- name: Distributed mutual FMM suite
# `-n 0`: four host devices live in ONE process and each test builds its
# own mesh, so serial keeps the device set unambiguous.
#
# `-rs` prints skip reasons, which the next step needs in the log: a
# module-level skip means the installed yggdrax predates the LET reverse
# halo, and that has to be readable rather than inferred.
run: |
set -o pipefail
pytest tests/integration/test_mutual_distributed.py -n 0 -rs -v \
| tee pytest-distributed-mutual.txt
- name: Fail if the suite was skipped rather than run
# A skipped suite PASSES, and a green-but-vacuous job is worse than a red
# one -- it is the exact state this job was added to end. Reuses the log
# from the step above rather than running the suite twice.
#
# Keyed on `-rs`'s SKIPPED lines, not on parsing a "N passed" summary:
# `addopts` already carries `-q`, so the summary line's presence depends on
# how the two verbosity flags happen to cancel, and a guard that silently
# stops matching is the failure mode being guarded against. Whether the
# tests PASSED is the previous step's exit code, not this step's business.
run: |
if grep -qE "^SKIPPED|no tests ran" pytest-distributed-mutual.txt; then
echo "::error::the distributed mutual suite was skipped, not run -- see the reason in the previous step"
exit 1
fi
benchmark-guard:
runs-on: ubuntu-latest
env:
JAX_ENABLE_X64: "1"
steps:
- name: Checkout
uses: actions/checkout@v4
- name: Setup Python
uses: actions/setup-python@v5
with:
python-version: "3.11"
cache: pip
- name: Install package + dev tools
run: |
python -m pip install --upgrade pip
pip install git+https://github.com/TobiBu/yggdrax.git
pip install -e ".[dev]"
- name: Run benchmark guard
run: |
python -m bench.ci_benchmark_guard \
--n 384 \
--runs 2 \
--warmup 1 \
--preset fast \
--basis solidfmm \
--theta 0.6 \
--leaf-size 16 \
--max-order 4 \
--dtype float32