Conversation
cl126162
force-pushed
the
stack-03-phase2-kwargs-config-fixes
branch
from
September 7, 2026 06:30
fc0968e to
1492e1f
Compare
cl126162
force-pushed
the
stack-04-phase3-hf-mixin-refactor
branch
from
September 7, 2026 06:31
e55ed6b to
225cd4b
Compare
cl126162
force-pushed
the
stack-03-phase2-kwargs-config-fixes
branch
from
September 7, 2026 10:14
1492e1f to
e084eb0
Compare
Finding 3.1: jaxfluids, jax, maia and nek each carried a private copy of
_setup_environment_data / _resolve_configuration_file /
_find_configuration_file, byte-identical except the ~/.cache namespace
string (jaxfluidsgym/jaxgym/maiagym/nekgym). Bug fixes had to be applied
four times (Tasks 2.5/2.6 already were). Extracts the shared logic into
hydrogym/hf_env_mixin.py:
- HFEnvConfigMixin with the three resolution methods, parameterized by
the class-level HF_CACHE_NAMESPACE; method bodies copied VERBATIM from
the JAX-Fluids backend (the first migrant) so behavior is unchanged.
- ConfigError moves to hf_env_mixin; jaxfluids env_core imports it under
the same name, so existing `except ConfigError` sites keep working
(same exception object, pinned by test).
- SOLVER_TYPE class attr (used for HFDataManager's fallback_profile --
jaxfluids previously hardcoded the literal "JAXFLUIDS" inline; the
value is unchanged).
- JAXFluidsFlowEnv now (HFEnvConfigMixin, JAXFluidsEnv); its local
copies of the three methods and its local ConfigError class are
deleted. _init_from_hf keeps its JAX-Fluids-specific error message.
Migration order follows the audit: JAX-Fluids first (smallest blast
radius); jax (3.2), nek (3.3), maia (3.4) follow in separate commits.
Nek's _resolve_configuration_file is a behavioral variant (returns None
instead of raising; different auto-detect list) -- its migration will not
be a pure find-replace.
Validation evidence (Safety Protocol rule 3):
- New test/test_hf_env_mixin.py (18 tests) in the full-gpu-stack
container (hydrogym-full-gpu-stack-audit, /opt/venvs/ml): 18 passed.
Covers cache hit/miss/failure paths, namespace parameterization, all
4 config-file resolution cases plus error paths, auto-detect priority
order, and pins the migration (subclass uses inherited methods, no
local overrides, namespace/profile values, ConfigError identity).
- Same file in the firedrake-test container (brave_shaw, no jaxfluids
installed): 16 passed, 2 skipped (importorskip guards fire).
- JAX-Fluids GPU smoke (examples/jaxfluids/test_jaxfluids_env.py,
Nozzle2D_coarse) post-change, run from host via test_gpu_solvers.sh:
"TIMEOUT (saw step/reset output)" -- the expected outcome for that
script (1000 steps hardcoded, uniform timeout cap; the runner reports
TIMEOUT separately from FAIL). Step output matches the pre-change
baseline log (08:20 run) step-for-step: "Env reset." then env_step=10
(sim_step=1000, sim_time=1.500e-04) and env_step=20 (2000, 3.000e-04);
wct/step/cell 86ns post vs 90ns pre -- within run-to-run noise. This
run exercised the migrated path end-to-end including a fresh HF
download (use_clean_cache=True).
- ruff check + ruff format --check + isort --check-only: clean on all
three files.
third_party/ submodules untouched.
(cherry picked from commit 1cab7e61602a33768c843f864b0254f624698d3c)
NekEnv's private copy of _setup_environment_data is deleted; the shared
mixin now provides it. This was the migration the audit flagged as "less
than a pure find-replace" -- handled as follows:
- _setup_environment_data: behavior identical, including the "[NEK] "
print prefix. The mixin gained a LOG_PREFIX class attr (default "")
so Nek's tagged log lines survive byte-for-byte while the other
backends keep their bare messages ("" -> identical f-string output).
- _resolve_configuration_file: NOT migrated. Nek keeps its own
simplified resolver as a documented divergence from the mixin
(None -> only environment_config.yaml/config.yaml; missing paths
return None and let _init_from_hf raise, instead of the mixin's
more specific immediate ConfigError). Preserving that behavior
matters more than deleting 20 lines; the divergence is now spelled
out in the docstring instead of being silent.
- ConfigError: moved to hf_env_mixin; `hydrogym.nek.ConfigError` and
`from hydrogym.nek.env import ConfigError` resolve to the same
shared exception object (pinned by test).
- NekEnv is now (HFEnvConfigMixin, gym.Env) with
HF_CACHE_NAMESPACE="nekgym", LOG_PREFIX="[NEK] ", SOLVER_TYPE
unchanged.
Validation evidence (Safety Protocol rule 3):
- Unit tests in the full-gpu-stack container (mpi4py available so
hydrogym.nek.env imports for real): mixin suite (incl. new
TestNekMigration + [NEK]-prefix byte-for-byte print test) plus
test_nek_mpi_bind_to / test_nek_runtime_overrides /
test_nek_reward_aggregation / test_nek_dead_kwargs /
test_hf_data_manager_token_revision = 60 passed, 2 skipped.
- Nek MPMD smoke in the nek5000-test container (peaceful_lovelace):
fresh workspace, TCFmini_3D_Re180, 3 solver steps, nproc=10,
"run successful" in the log, EXIT=0 -- the migrated
_setup_environment_data ran live inside a real MPMD launch
(env download/cache-hit path + config resolution + solver).
(cherry picked from commit 4fca0fded3cfeef5264a021b52646bb9584f6e2b)
MaiaFlowEnv's private copies of _setup_environment_data /
_resolve_configuration_file / _find_configuration_file are deleted; the
shared mixin now provides them (bodies were byte-identical to the mixin's
already -- namespace "maiagym", bare log messages). MaiaFlowEnv is now
(HFEnvConfigMixin, gym.Env) with HF_CACHE_NAMESPACE="maiagym" and
SOLVER_TYPE="MAIA_LB" unchanged. ConfigError moves to hf_env_mixin;
`from hydrogym.maia.env_core import ConfigError` (used by all the maia
envs) resolves to the same shared exception object (pinned by test).
Unused glob/Path imports dropped from maia env_core.
This completes the HF-mixin migrations (3.1 jaxfluids -> 3.2 jax ->
3.3 nek -> 3.4 maia): all four backends now share one copy of the
environment-data download / config-file resolution logic.
Validation evidence (Safety Protocol rule 3):
- Unit tests in the full-gpu-stack container (mpi4py available):
test_hf_env_mixin.py (25 tests, incl. new TestMaiaMigration) +
test_hf_data_manager_token_revision.py + test_substep_naming.py =
43 passed, 2 skipped.
- Real MAIA MPMD run as the merge gate (the audit requires a live
mpirun, not just review): maia_gpu smoke (Cylinder_2D_Re200,
prepare_workspace + mpirun -np 1 python test_maia_env.py : -np 1
maia properties_run.toml) post-change: PASS with step/reset output,
"Test completed successfully" in .devcontainer/test-logs/maia_gpu.log.
- ruff check + ruff format --check + isort --check-only: clean.
(cherry picked from commit 01326861d167bfb0de7f2a1f48fe82b3ddd70416)
Second backend onto the shared mixin (after JAX-Fluids, 1cab7e6): delete JAXFlowEnv's private copies of _setup_environment_data / _resolve_configuration_file / _find_configuration_file and the local ConfigError, inherit HFEnvConfigMixin, set HF_CACHE_NAMESPACE = "jaxgym" (SOLVER_TYPE = "JAX" was already a class attr). _update_configuration_paths and the rest of the module are untouched. test_jax_env_core.py's fake_home fixture now patches pathlib.Path.home directly, since env_core no longer imports Path after dropping the local copies. Note: no shipped JAX environment uses JAXFlowEnv (both KolmogorovFlow and ChannelFlowSpectralEnv extend the HF-free JAXFlowEnvBase), so this migration carries zero behavior risk for the shipped envs; the class is pinned by the migration tests in test_hf_env_mixin.py (TestJaxMigration) and test_jax_env_core.py. Validation: - ml container (hydrogym-full-gpu-stack-audit): 31 passed across test_hf_env_mixin.py, test_jax_env_core.py, test_substep_naming.py, test_hf_data_manager_token_revision.py (incl. the 12 drift guards). - GPU smokes (host, .devcontainer/scripts/test_gpu_solvers.sh): jax_kolmogorov PASS (~45s, step output matches baseline); jax_channel FAILED once on a transient cuBlas allocation failure during jit compilation (GPU memory contention with the desktop stack), PASSED on isolated retry in 17s with correct step output. - ruff + isort clean. Co-Authored-By: Claude Code <noreply@anthropic.com> (cherry picked from commit cf106e51b37f6b20cf35e9401ef15f0ad211b829)
New hydrogym/core_external.py hosts both MPMD world-split strategies
verbatim, one named function each, plus the audit's requested mixin:
- mpi_split(comm_world, nproc=None, controller_rank=0, intercomm_tag=99,
log_prefix="[MPI_SPLIT] ") -- the Nek rank-color split
(Split + Create_intercomm), previously a private copy in nek/env.py.
- split_comm_by_appnum(comm_world) -- the MAIA APPNUM split
(Split + group translation + Allreduce root discovery), previously
inline in MaiaInterface.init_comm; init_comm now delegates and keeps
only its attribute-assignment contract.
- ExternalProcessEnvMixin with _split_mpmd_comm() and class-level knobs
(CONTROLLER_RANK / INTERCOMM_TAG / MPI_SPLIT_LOG_PREFIX).
The audit frames this as deduplication, but the two are genuinely
different protocols (not copies), so nothing was merged into one;
tag-protocol code stays untouched per backend. Scope discrepancy is
recorded here rather than silently resolved.
NekEnv adopts the mixin: its call sites now use
self._split_mpmd_comm(comm_world, nproc=self.nproc). The module-level
`mpi_split` name is preserved as a re-export so `hydrogym.nek.mpi_split`
and the monkeypatch idiom keep working (the three tests that stub the
split now patch hydrogym.core_external.mpi_split with **kw-tolerant
lambdas, since the mixin passes the new knobs through).
TWO LATENT BUGS FOUND BY THE NEW MPI-TIER SMOKE (test/mpmd_smoke_split.py,
both fixed; identical results on the production paths, which mask them):
1. mpi_split hardcoded remote_leader=1 on BOTH sides. Correct only for
the controller (whose remote leader is world rank 1; in production the
workers are the Nek5000 Fortran binary with its own handshake, so the
Python color-1 branch never ran). A Python worker would contact
itself and leave the controller hanging in Create_intercomm forever --
reproduced live. Now remote_leader=1 - color.
2. The APPNUM root translation went world group -> app group, giving
MPI_UNDEFINED (-1) for every app except app 0; the value was only ever
right because the Python controller is always app 0. Reproduced live
(remote_root=-1 for a Python app 1). Now app group -> world group.
Validation:
- Unit tier (ml container): 65 passed, 2 skipped across
test_core_external.py (new: error paths, mixin delegation/knobs,
backend wiring, re-export identity, split-fix pins),
test_hf_env_mixin.py, test_jax_env_core.py, test_nek_reward_aggregation.py,
test_nek_mpi_bind_to.py, test_hf_data_manager_token_revision.py.
- MPI tier, no solver binaries (ml container):
mpirun -np 2 python test/mpmd_smoke_split.py nek -> exit 0
(intercomm handshake + ping-pong + nproc-mismatch error path)
mpirun -np 1 ... appnum : -np 1 ... appnum -> exit 0
(APPNUM split + remote-root exchange)
python test/mpmd_smoke_split.py nek-error -> clean refusal
- Nek5000 merge gate (peaceful_lovelace, real MPMD run, nproc=10):
test_nek_direct.py --steps 3 -> EXIT=0, "run successful".
- MAIA merge gate (host GPU script): maia_gpu PASS ("Test completed
successfully", step/reset output seen).
- ruff + isort clean (also fixes 4 pre-existing E501s in
test_hf_data_manager_token_revision.py that the file touch surfaced).
Co-Authored-By: Claude Code <noreply@anthropic.com>
(cherry picked from commit 8f3e10d0cd799fb74edb144eae61f049e2a38b1a)
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
cl126162
force-pushed
the
stack-04-phase3-hf-mixin-refactor
branch
from
September 7, 2026 10:14
225cd4b to
252517c
Compare
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
2 participants
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.
Fourth PR in the sequential `hydrogym-audit-v2` stack. Targets #285 (Phase 2).
Scope (Phase 3 — shared-mixin extraction, Tasks 3.1-3.5)
Each backend independently reimplemented the same HuggingFace Hub
data-staging logic (with subtle drift between them) and, for the two
MPMD backends, the same MPI comm-split logic. This phase extracts both
into shared mixins and migrates every backend onto them:
Verified
codespell all clean
HFEnvConfigMixin` verified working; `core_external.py` syntax-checked
(its `mpi4py` import is an optional extra, not installed in the plain
lint venv used here — exercised for real via the actual MPMD backends
in the devcontainer verification passes documented in
`HYDROGYM_ENGINEERING_AUDIT_v2.md`)
🤖 Generated with Claude Code