Conversation
JAX 0.4.38 silently no-ops the persistent compilation cache when any jax.Array is materialized before the cache dir is set (jax-ml/jax#25768), so every process paid full cold compile (~7 min GPU / ~20 min/cell CPU). JAX 0.5.x has a worse bug: cache hits silently return stale .at[].set() updates inside lax.fori_loop / lax.cond (jax-ml/jax#31733), corrupting the env step — reproduced here with 17/35 blue-Remove tests failing on a warm cache under 0.5.1. Bug is fixed from 0.7.1 onward; we land on 0.10. Distrax is dropped because its transitive tensorflow-probability<=0.25 imports jax.interpreters.xla.pytype_aval_mappings, removed in JAX 0.7. The 5 callsites only used distrax.Categorical for sampling / log_prob / entropy, all easily expressed on jax.random.categorical and jax.nn.log_softmax — see src/jaxborg/policies/categorical.py (31 lines, flax struct so it's jit/vmap/scan-compatible). Dropping distrax also removes TFP, gast, and decorator from the dep graph. Knock-on bumps from JAX 0.7+ requiring ml-dtypes>=0.5 (numpy 2 C ABI): * numpy 1.26.4 -> 2.3.5 (overrides cyborg's pin; cyborg runtime is numpy-2-compatible empirically) * scipy 1.12 -> 1.17 (overrides jaxmarl's `scipy<=1.12` precautionary upper bound; cf. jaxmarl commit 3ffa5b8f and issue #175) * torch 2.2 -> 2.10 (overrides cyborg's pin; needed for numpy-2 init) Also retires a flaky test that relied on accidental RNG alignment between CybORG's numpy and JAX's threefry: tests/subsystems/test_fsm_red_agent.py::test_fsm_hidden_state_applies_after_completion_step walked both pipelines forward on seed=0 and asserted they hit the same FSM sequence — a parity-by-coincidence that breaks any time JAX's PRNG layout changes. Restructured as a pure-state check of the two-stage delayed-update mechanism (fsm_red_schedule_post_step_update stages red_fsm_delayed_states; fsm_red_apply_delayed_update commits it on the next step), which is the actual invariant the test name promises. Verification: * fast suite: 772/772, cold 154s -> warm 115s (cache hit, correct) * full suite incl. slow: 1214 passed, 100 skipped, 3 xfailed (xfails are pre-existing, unrelated) * cache miscompile probe: blue_remove.py 35/35 cold and 35/35 warm against the same cache dir — under 0.5.1 the same probe gave 17/35 failures on warm
…erage Remove three xfail tests in test_fsm_red_env_differential.py that depended on retired CybORG green/replay tape infrastructure. Their first two checks (red_4 known-hosts parity and red_4 action-selection parity) are already covered by test_red_policy_matches_cyborg_multistep across 200 steps x 5 seeds. The third (end-state host_compromised/red_privilege parity under FSM red + green phish) had no equivalent — existing green-sync tests use SleepAgent for red, so no exploit/privesc chains fire. Add TestFsmRedGreenSyncParity::test_no_critical_state_diffs_over_10_steps which closes that gap via CC4DifferentialHarness(FSM red + EnterpriseGreen + sync_green_rng=True) and asserts at least one privesc fired so the test can't pass on a degenerate trajectory. Also adds seed=0 to the existing red_policy_parity parametrize to preserve the original tests' seed.
The 4h→6h→10h timeout bumps were chasing a non-existent ceiling: GitHub hard-cancels any single job at 360 min regardless of timeout-minutes, so timeout: 600 on main is a no-op. The slow gate needs to scale below that line, not above it. Switch test-slow to an 8-way matrix using pytest-split. 545 slow items split 69×7 + 62; each shard runs well under an hour, wall-clock for the gate drops from ~3.5h to ~1h once durations are populated. .test_durations is held in the GH Actions cache (actions/cache@v4), not git — committing it would thrash history every run since the floats churn. Each shard restores the most recent cached durations, runs with --store-durations, and saves its own updated copy under a unique key. Over runs, fresh entries from every group accumulate. Sized for the cold-start case where no durations exist yet: - 8 shards (vs strictly-needed 4) buys headroom against round-robin clustering of heavy parametrized tests - timeout-minutes: 240 — first-run safety net only; steady-state is well under an hour. Still under the 360 hard cap. Adds pytest-split to the dev dep group; uv.lock regenerated.
The shard 8/8 hang was a JAX-array-as-int hot loop in _global_host_to_relative_slot / _abs_subnet_to_relative — pure-Python host-side helpers that were doing int(jax_array[i, j]) per element of small const lookup tables. Each access round-tripped through full JAX primitive dispatch (slice_p.bind → apply_primitive); on cold-cache CI runners the JAX persistent-cache lock contention from the also-running xdist workers turned that into an unbounded stall. Switched to one np.asarray + np.flatnonzero per call. Module already declared itself JAX-free at the top — this aligns with that intent. Workflow: - drop addopts via -o "addopts=" so -n auto doesn't stack on top of sharding (xdist + pytest-split contend on the JAX cache lock) - fix the durations cache: previous strategy had each shard overwrite .test_durations with only its 1/N slice, so prefix-restore picked up partial files that never accumulated. Now each shard writes .test_durations_cache/group-N.json, and a merge step rebuilds the full .test_durations before pytest runs - timeout 240 → 60 (hang root cause is fixed; no first-run cliff to budget for) Verified locally: test_native_generic_exploit_respects_blocked_scan_- source_route_matches_cyborg now finishes in 3m4s (was unbounded).
First run after the cache schema change has empty .test_durations_cache, so pytest-split round-robins by test count and heavy parametrized tests cluster into 1-2 shards that exceed 60 min serially. Subsequent runs (with partials cached) all fit under 30 min.
Per-shard cache keys could never warm up: actions/cache restores at most one cache per step, so each shard merged from a single shard's 2.4 KB partial instead of all eight. pytest-split fell back to round-robin and shard 2/8 hit the 120 min timeout while shard 3/8 finished in 59 s. Replace per-shard caches with a single aggregate refreshed by a new aggregate-durations job that merges every shard's uploaded .test_durations (plus the prior aggregate so timed-out shards keep their last-known durations) and saves under slow-durations-aggregated-<run_id>. Each shard restores the most recent aggregate via the slow-durations-aggregated- prefix. Bump test-slow timeout 120 -> 180 min for cold-cache headroom on the first run before the aggregate exists.
Two fixes for the slow-test sharding bootstrap: 1. upload-artifact@v4 excludes hidden files by default, so `.test_durations` was silently dropped. The aggregate-durations job consequently found 0 artifacts and never populated the cache. Add `include-hidden-files: true`. 2. Even with #1 fixed, the very first run on a new cache scope (this PR, main after merge, future PRs) still cold-starts with round-robin sharding and hits the per-shard timeout. Add a conditional "seed from URL" step that pulls a known-good `.test_durations` from `vars.SLOW_DURATIONS_SEED_URL` when the cache is empty. Subsequent runs hit the cache and skip the download, and the aggregate job keeps the cache refreshed.
…teps to mlflow conf
…eval, and fixed seed ranges for train/eval for deterministic runs
… hosts, matching actual implementation
…; minote update to eval param naming for consistency across scripts
…mant red overtaking blue
…ests run on a change
…y in IPPO/MAPPO agents
… game partial observability intact
…te referenced files to recipe path
… against same reds as standard trained
Coordinate initial and automatic resets across parallel JAX environments. Require diversified banks to cover NUM_ENVS, preserve singleton controls, and support both Blue-only and joint trainers.
Extend Blue inputs from 210 to 402 with host presence, completed Analyse evidence, and persistent alert history. Match JAX and native CybORG memory semantics, propagate dimensions through policies and checkpoints, and enable the contract for dual-team cotraining recipes.
Evaluate final Blue policies against stock FSM Red for 100 native episodes. Allow each post-training job to select CPU or CUDA, and record simulator provenance and score semantics alongside the existing held-out JAX evaluations.
Use 16 whole-sequence minibatches for GRU and LSTM cotraining, yielding 64 optimizer steps per rollout. Remove unused CleanRL blocks and document that earlier eight-minibatch runs require rerunning for matched comparisons.
…g w/o replacement
… update yaml scripts for comparison runs
…iate/most recent recipes
… topology; also updated invalid scripted-red settings and MAPPO checkpoint lookup
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
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary