ci.yml

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