jax gate: CI runs an unpinned jax that no local environment has

[oshaughnessy-junior]

Corrected. This issue opened claiming two local failures were caused by jax
version drift. Neither was. The version spread below is real and still worth
fixing; it explained neither failure. See the comments for the two corrections
and what the causes turned out to be.

The version spread

jax-ile-check installs jax and numpyro unpinned, deliberately, so the gate
sees what a user gets. At the #214 merge the runner resolved jax 0.10.2, and
no environment on the CIT home had it:

environment	jax
CI jax-ile-check	0.10.2
~/.cache/jaxci_venv	0.9.2
CVMFS igwn	0.7.1
conda coag_gpu / gwkokab / junior_rift / my_claude	0.7.2 / 0.9.2 / 0.9.0 / 0.9.0
conda asimov	0.4.23

PyPI's latest was already 0.11.1, so an unpinned local build does not reproduce
CI either. To mirror a run you must pin to what that run resolved.

Partly addressed. PR #293 makes jax-ile-check upload its resolved stack as
an artifact, so re-pinning is mechanical. A CI-matched venv now exists at
/scratch/richard.oshaughnessy/envs/jaxci-py311, rebuilt or re-pinned by
build_jaxci.sh in the same directory.

The two failures, correctly attributed

Both were chased as version drift and neither was.

1. test_laplace_high_amplitude_accuracy_and_trend fails on CIT on the branch
   and on pristine base, on both jax versions, with and without XLA_FLAGS.
   Its file is in the gate's EXCLUDED list, so CI never runs it and a green
   CI run says nothing about it. No green run of it exists on any machine. The
   assertion compares two values at the double-precision noise floor.
2. The whole-suite abort is the 25 GiB per-user memory cgroup, shared across
   sessions. dmesg names the kill by PID. Seen as exit 134, 137 and 139; the
   signal varies, the cause does not. Chunking the roster fixes it.

Still open

• Container image and a blocking lane that uses it.
• Demote the unpinned lane to a non-blocking canary, after the container
  lane exists and not before, or the repo has no jax gate at all in between.
• Run the EXCLUDED suite on a schedule. It had silently drifted out of sync
  with the driver it pins, and only a by-hand run during jax ile: persist and transfer compilation caches safely #214 caught it.
• Decide what ~/.cache/jaxci_venv is now that a CI-matched venv exists.

[oshaughnessy-junior]

PR #293 lands the first half of option 4: the gate's own comment now tells a reader how to separate an environment red from a branch red, without naming versions that will age. The rest of option 4 (deciding what jaxci_venv is, and refreshing or relabelling it) is not done.

[oshaughnessy-junior]

Path forward

The container is the right durable answer, and it needs one thing alongside it.

Why the container wins. It pins jax, jaxlib, CUDA, numpyro and lal as one
artifact, and the same image runs in CI and on CIT. A version string cannot do
that: the abort around test_jax_slowrot.py is native jaxlib, so "jax 0.9.2"
names neither the wheel nor the CUDA it was built against. Under an image, "which
jax did this run" becomes a digest, and a local red is reproducible in CI by
construction.

What it does not fix. Test only in a pinned image and you stop hearing about
upstream breakage until an upgrade forces it, all at once. That is exactly what
the unpinned policy in ci.yml exists to prevent, and dropping it would trade a
loud problem for a quiet one.

So: two lanes.

• Blocking: the container, digest-pinned. Every PR passes it, and anyone can run
  the same image on CIT.
• Canary: the present unpinned install, continue-on-error, nightly and on
  rift_O4d rather than per-PR. Its red means "bump the image", not "this PR is
  broken".

Whoever owns the image owns the bump cadence. A digest-pinned gate is stale by
default, and a canary nobody acts on becomes noise.

Cheap now, independent of the container.

• Attach the resolved versions as a CI artifact. The gate prints jax already;
  a pip freeze makes an old failure reproducible instead of archaeology.
• Settle what jaxci_venv is. Refresh it against the same resolution, or stop
  calling it the CI-equivalent environment. This is the item that would have
  saved the two runs described above, and it needs no CI change.

Keep the collection floor as it is. The count was 621 on both jax versions
while pass/fail differed, so the floor is the version-robust half of this gate.
It should not move into the container work.

Split, so this is not done twice. The container and the blocking lane belong
to the agent already on it. The canary lane, the version capture and the
jaxci_venv decision are separable and can proceed in parallel. This issue is
the join point; link the container PR here when it opens.

Scope note. jax-ile-check runs JAX_PLATFORMS=cpu, so a CPU image is
enough for this gate. GPU images and their /scratch placement are a separate,
larger question that containers/ already addresses.

[oshaughnessy-junior]

Sequencing: the canary lane must come after the container, not before

Recording this so nobody picks up option 2 next and makes things worse.

The unpinned lane described as "the canary" is currently the ONLY jax gate.
Marking it continue-on-error today would leave jax_ile with no blocking test
at all, until the container lane exists to replace it. The order has to be:

1. Container image plus a blocking lane that uses it. Nothing changes about the
   present job while this is built.
2. Once that lane is green and required, demote the unpinned job to
   continue-on-error, move it to nightly and rift_O4d, and keep it.

Doing 2 before 1 trades a reproducibility problem for an unguarded merge path.

Done

• PR jax gate: record the resolved stack, and say how to tell an environment red from a branch red #293: the gate's comment now explains how to separate an environment red
  from a branch red, without naming versions; and jax-ile-check uploads its
  resolved stack as an artifact, on failure as well as success.
• The version spread is recorded in this operator's host atlas, so an agent on
  CIT is told that ~/.cache/jaxci_venv is not a CI mirror before it spends two
  runs finding out.

Open, and not started

• The container and its blocking lane (item 1 above).
• The demotion of the unpinned lane (item 2), blocked on item 1.
• What ~/.cache/jaxci_venv should be. Refreshing it is not free: other agents
  run against it concurrently, so changing its jax under them would break work in
  flight. A separate, clearly-named environment is the safer shape.

[oshaughnessy-junior]

Correction: this is not a jax version difference

I attributed the two divergences above to jax version. That was wrong for the
first one, and the check that shows it is a venv built to match CI.

/scratch/richard.oshaughnessy/envs/jaxci-py311 now holds the resolved stack of
the green jax-ile-check run at c2e5b2b, on CI's Python 3.11: jax 0.10.2,
jaxlib 0.10.2, numpyro 0.21.0, numpy 2.4.6, scipy 1.17.1, numba 0.67.0,
lalsuite 7.26.15. In that environment
test_laplace_high_amplitude_accuracy_and_trend still fails on CIT, while the
same package set passes on the GitHub runner. It also fails with no XLA_FLAGS
set, and with each of the two flags I had been exporting, on identical numbers.
So the jax version, the package set and the flags are all ruled out.

What differs is the machine. CIT is an AMD EPYC 9475F; the runner is not. numpy
ships DYNAMIC_ARCH OpenBLAS and XLA emits host-CPU code, so both dispatch
different kernels per CPU.

The test is the actual defect

assert errs[0] < 0.15      # passes
assert errs[1] < 0.08      # passes
assert errs[1] < errs[0]   # this one

Measured here: errs[0] = 7.105427357601002e-14, errs[1] = 4.170033207628876e-10. The docstring records 1.8e-2 at scale 50 when the pin was
written, so the Laplace path has since become exact to ~1e-13 at both
amplitudes. Once both errors collapsed to the double-precision noise floor, the
O(1/A) trend stopped being measurable and the ordering became a coin flip that
lands differently on different CPUs. errs[1] is bit-stable across jax versions
and flags; errs[0] moves in its last bits.

That is a latent CI flake rather than an environment problem. It passes on the runner
today for no reason anyone controls. The two bound assertions still carry the
physics; the trend assertion no longer does.

What survives

The version spread in the issue body is still real, and a container still fixes
reproducibility of the package set. But it would NOT have fixed this failure, and
the plan should not claim otherwise. Cross-CPU numerics need either a tolerance
that reflects the noise floor or a reference the trend can be measured against.

PR #293's gate comment is unaffected: it names no versions and says to reproduce
on the pristine base in the same environment, the step that identified this.

[oshaughnessy-junior]

Second correction: I asserted a cause twice, and both times on a bad inference

Retracting the "host CPU" claim in my previous comment.

test_angle_marg_exact.py is in this gate's EXCLUDED list, by a deliberate and
documented decision. CI never runs it. My reasoning was "CI passes it, CIT fails
it, same package set, so it must be the CPU", and the first clause is false.
There was no divergence to explain.

The same mistake produced the jax-version claim before it. Both times I inferred
from a green CI run to a specific test passing in CI, without checking the
roster. That is the inference this gate exists to make impossible, and I made it
twice.

What is actually established

test_laplace_high_amplitude_accuracy_and_trend fails on CIT on the branch and
on pristine base, on jax 0.9.2 and 0.10.2, with no XLA_FLAGS and with either
thread-taming flag, on the same two numbers throughout. I have no green run of it
from any machine. Whether it fails everywhere or only on some CPUs is open, and I
am not going to guess a third time. The likely reading is that the Laplace path
became exact to ~1e-13, so an ordering assertion between two noise-floor values
stopped meaning anything.

The other failure was memory, and that one is settled

The whole-suite abort is the per-user memory cgroup, not jax and not a test. The
kernel names it:

Memory cgroup out of memory: Killed process 1209638 (python)
total-vm:21094896kB anon-rss:12422192kB UID:40428

PID 1209638 is the pytest process from that gate run; the limit is 25 GiB and it
is shared across every session this user has open. Splitting the roster helps and
one file per process is enough: all 13 files of the chunk that died pass
individually, and the other two chunks pass whole (243 and 240 tests). The gate
does run locally. It cannot run as one process.

A cost of the exclusion, found by accident

On pristine base, test_driver_labels_a_suspect_angle_grid_in_provenance asserts
a driver signature the driver had already changed. The test had drifted out of
sync with the code it pins and nothing caught it, because nothing runs the file.
PR #214 happened to fix it. That is an argument for running the excluded suite on
a schedule, which fits the canary lane in the plan above.