Skip to content

ci: install jax on the recovery job so the gradient-path fits run - #519

Merged
wshlavacek merged 1 commit into
mainfrom
ci/recovery-job-needs-jax
Jul 24, 2026
Merged

ci: install jax on the recovery job so the gradient-path fits run#519
wshlavacek merged 1 commit into
mainfrom
ci/recovery-job-needs-jax

Conversation

@wshlavacek

Copy link
Copy Markdown
Collaborator

Follow-up to #518 / #517.

The recovery job added in #518 was dispatched for the first time to size it (the point of shipping it dispatch-only). It came back failure at ~11.7 min: 125 passed, 2 failed — and the two failures were a config gap in the job, not a stochastic tolerance flake:

FAILED tests/test_real_world_examples.py::test_real_world_runs_through_bngsim[Rijal-2025/lacud5_ode]
FAILED tests/test_real_world_examples.py::test_real_world_runs_through_bngsim[Rijal-2025/five_dl1_ode]
  PybnfError: Gradient assembly needs jax to transform a log-scaled parameter's
  gradient into sampling space ... the optional 'jax' extra.

Those two real-world cases fit log-scaled parameters through jax-backed gradient assembly, and they raise rather than skip when jax is absent. The recovery job was built with setup-pybnf's default jax: 'false', so it failed deterministically — it will fail identically every run until jax is installed.

Fix

Add jax: 'true' to the recovery job — exactly as tests.yml's pytest-jax job already does for the same gradient path. A complete recovery environment provisions jax.

Verification

  • The two previously-failing cases pass locally in ~15s with jax present.
  • Full end-to-end confirmation: this branch was dispatched via workflow_dispatch on the branch ref to confirm the whole recovery tier goes green under jax before merge. (Result linked in comments.)

Sizing takeaways from the first dispatch (context for #517)

  • Recovery tier ≈ 12 min on a 2-core runner — the ">1 hr" concern was a dev-machine artifact; it's comfortably schedulable.
  • Zero stochastic tolerance flakes this run — the only failures were the deterministic jax gap fixed here.
  • Side observation: the pre-existing slow job ran 32 min, vs the ~8 min its comment claims. Not touched here; flagged for a possible follow-up.

The recovery tier includes gradient-path parameter-recovery fits whose
log-scaled parameters route the gradient into sampling space through jax
(test_gradient_*, and the Rijal-2025 real-world cases in
test_real_world_examples.py). Those tests RAISE PybnfError rather than skip when
jax is absent, so the recovery job -- added in #518 with setup-pybnf's default
jax: 'false' -- failed deterministically on its first dispatch: 2 hard failures
(lacud5_ode, five_dl1_ode), 125 passed, in ~11.7 min. Not a stochastic tolerance
flake; a missing optional dependency.

A complete recovery environment provisions jax, so install it here via the
composite action's jax input -- exactly as tests.yml's pytest-jax job already
does for the same gradient path. The two failing cases pass locally in ~15s with
jax present.

Follow-up to #518 / #517.
@wshlavacek
wshlavacek merged commit c032899 into main Jul 24, 2026
9 checks passed
@wshlavacek
wshlavacek deleted the ci/recovery-job-needs-jax branch July 24, 2026 22:06
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant