JIT cache: key on generated C source (closes #121, #123) - #129
Conversation
Break the monolithic _createext() into three clearly-scoped functions so that generation and compilation can be hashed/cached independently — the prerequisite for keying the JIT cache on generated C source (see ~/.claude/plans/jit-cache-c-source-hash.md). - generate_c_source(...) → (modname, codeguys, diagnostics) Pure text phase: sympy processing, c/cy source emission. No I/O. - compile_and_load(modname, codeguys) → (module, tmpdir) Writes files, runs `setup.py build_ext`, dynamically loads the .so. - _createext(...) stays as a thin wrapper that stores the loaded module in _ext_dict[name], preserving backwards compatibility with getext(). Behaviour is byte-equivalent to the old implementation. All existing JIT-related tests continue to pass: - tests/test_1001_poisson_constants.py (6/6) - tests/test_1000_poissonCart.py (7/7) - tests/test_1010_stokesCart.py (6/6) Part of the JIT cache refactor that unifies issues #121 (persistent cross-session cache) and #123 (per-value recompile). Underworld development team with AI support from Claude Code (claude.com/claude-code)
The core of the refactor. Two complementary changes make the JIT cache
correct by construction for constant ↔ symbolic ↔ constant transitions:
1. Cache key is now sha256(canonical C source + ABI salt) rather than
a hash of the sympy callback structure. This means same compiled code
⇒ same cache entry, regardless of how many times a UWexpression's
value changed in between.
- Replaces _ext_dict (integer-keyed) with the same name but str keys.
- Drops _fn_cache / _CachedFn / _SIG_TYPES — the bundle-level hash
subsumes the partial-reuse benefit, and disk persistence (phase 3)
will amortise the one-extra-compile cost across sessions.
- _abi_salt() hashes (PETSc version, underworld3 version). Heavier
toolchain fingerprinting is handled by ./uw build in phase 4.
2. Fix _extract_constants so placeholder names and sort order are based
on expr.name (stable) rather than str(expr) (= current value). With
the old str() key, changing one constant's value could re-order the
constants[] indices, silently remapping which constant each UWexpression
referred to.
3. Fix a latent _build() fast-path bug exposed by (1): fast path 2
(function rewire in place) wasn't updating _last_jit_cache_key, so
a later call could see _current == _last (using a stale _last from
an earlier full build) and bypass rewire, leaving the solver wired
for a different bundle than the caller expects. Added the update.
Net result — gate tests in tests/test_jit_cache.py:
- C → S → C transition: state-3 returns exactly state-1's result
(bug that motivated this refactor).
- 20× toggle between two constant values: exactly 1 cache entry
(value changes don't invalidate the C source).
- cache_key is a 16-char hex string (deterministic form).
Regression:
tests/ -m "level_1 and tier_a" → 52 passed, 3 skipped (MPI), 0 failed.
Part of the JIT cache refactor that unifies issues #121 (persistent
cross-session cache) and #123 (per-value recompile).
Underworld development team with AI support from Claude Code (claude.com/claude-code)
Skip the cc compile entirely on second-and-subsequent runs of the same solver setup, even from a fresh Python process. Resolves #121. New module: src/underworld3/utilities/_jit_cache.py - get_cache_dir(): resolves UW_JIT_CACHE_DIR / XDG_CACHE_HOME / ~/.cache/underworld3/jit, with UW_JIT_CACHE=0 disabling the cache. - load_module(source_hash, modname, constants_manifest): returns the dynamically-loaded extension on hit, or None on miss (no .so on disk, malformed/version-mismatched manifest, or constants-name shift). - store_module(source_hash, modname, tmpdir, constants_manifest): copies the freshly-compiled .so to {hash}.so and writes {hash}.manifest.json. Uses NamedTemporaryFile + os.replace for atomicity so a half-written entry never gets loaded by a concurrent reader. Hook in getext (src/underworld3/utilities/_jitextension.py): Three-tier lookup, cheapest first — in-memory dict, disk cache, cold compile. On cold-compile we now also persist the .so so the next process gets a free pass. Manifest JSON format: { "version": 1, "modname": "fn_ptr_ext_<hash>", "constants": [{"index": 0, "name": "K_diff"}, ...] } Belt-and-braces consistency: load_module compares the saved constants list (by name and index) against the current call's manifest. A mismatch means an unnoticed ABI drift or expression rename — recompile rather than risk a stale mapping. MPI is intentionally not handled here yet — phase 4 of this refactor adds rank-0 gating and flock-based inter-process synchronisation. New gate tests in tests/test_jit_cache.py: - test_disk_cache_roundtrip_skips_compile: solve, clear in-memory cache, re-solve; assert compile_and_load was NOT called. - test_disk_cache_disabled_by_env_var: UW_JIT_CACHE=0 makes both load and store no-ops. Regression: tests/ -m "level_1 and tier_a" → 54 passed, 3 skipped (MPI), 0 failed. Underworld development team with AI support from Claude Code (claude.com/claude-code)
…nv fingerprint
Three independent safeguards on the on-disk cache:
1. Per-hash flock (fcntl) in store_module
Two processes (different mpirun, different shells) racing on the same
uncached entry now serialise on a {hash}.lock file instead of
interleaving partial writes. Recheck-inside-lock pattern avoids the
redundant copy when the other process won the race.
2. comm.allgather determinism check in getext
When MPI size > 1, every rank computes the canonical C-source hash
independently and asserts they all agree. A mismatch means
generate_c_source is non-deterministic across ranks (typically a set
or dict whose iteration order leaks into emitted C) — failing loudly
here is much better than letting ranks load disjoint cache entries
and produce silent disagreements at solve time.
3. ./uw env fingerprint (~/.cache/underworld3/jit/.env-fingerprint)
The conservative ABI salt in _abi_salt() only covers PETSc + UW
versions. A change to the compiler family, CFLAGS, Python ABI, or
pixi env can still invalidate cached .so files. ./uw build now
computes a fingerprint over those four and wipes the cache when it
doesn't match the saved one. Verified manually:
$ echo bogus > ~/.cache/underworld3/jit/.env-fingerprint
$ ./uw build # → "Compiler / Python / pixi env changed; wiping JIT cache..."
New gate tests:
- test_abi_salt_change_invalidates_cache: monkey-patch _abi_salt to
two different strings, assert different cache keys.
- test_store_module_creates_lockfile: assert the .lock sentinel is
written next to the entry.
Regression: tests/ -m "level_1 and tier_a" → 56 passed, 3 skipped (MPI),
0 failed.
Underworld development team with AI support from Claude Code (claude.com/claude-code)
Adds docs/developer/subsystems/jit-cache.md describing the architecture, invalidation matrix, MPI semantics, and supported environment variables. Linked into the developer subsystems toctree. Adds docs/advanced/benchmarks/jit_cache_vs_recompile.py — the issue #123 reproducer (VE Stokes BDF-2 with toggling Top BC sign and dt_elastic). Measured locally on a 16x8 mesh, 30 steps: Wall time (total) : 20.1 s First step (cold) : 4.96 s Mean of remaining : 0.52 s/step Compiled bundles : 1 Issue #123 reported the original code spending ~459 s of JIT vs ~3.7 s of SNES solve over 99 steps; the equivalent here would have been ~90 recompiles instead of 1. Cross-session win is also visible: a second process with a warm disk cache cuts first-step time from 3.6 s to 2.4 s (no cc invocation). Underworld development team with AI support from Claude Code (claude.com/claude-code)
There was a problem hiding this comment.
Pull request overview
This PR overhauls Underworld3’s JIT caching to use a content-addressable key derived from the generated C source (plus an ABI salt), adds a persistent on-disk cache for compiled extensions, and updates solver wiring so constant/value-only changes avoid unnecessary recompiles—addressing the cold-start and “recompile on parameter update” regressions from #121 and #123.
Changes:
- Replace SymPy-structural cache keying with a SHA-256–based hash of canonicalised generated C source + ABI salt, and ensure solver fast paths track the active wired bundle.
- Add persistent disk cache (
{hash}.so+{hash}.manifest.json) with atomic writes and per-hash advisory locking. - Add tests and developer documentation for cache behaviour, invalidation, and determinism expectations.
Reviewed changes
Copilot reviewed 8 out of 8 changed files in this pull request and generated 6 comments.
Show a summary per file
| File | Description |
|---|---|
uw |
Adds build-time env fingerprinting intended to invalidate/wipe the disk JIT cache when toolchain/ABI changes. |
tests/test_jit_cache.py |
New test suite covering cache key determinism, constant/value cycling, and disk round-trips. |
src/underworld3/utilities/_jitextension.py |
Refactors JIT generation/compile/load, computes canonical-source hash key, removes per-function cache, and integrates disk cache tier. |
src/underworld3/utilities/_jit_cache.py |
New persistent cache implementation: directory resolution, module load/store, lockfiles, and manifest verification. |
src/underworld3/cython/petsc_generic_snes_solvers.pyx |
Fixes _last_jit_cache_key tracking on rewire path to prevent wrong-bundle reuse across C→S→C transitions. |
docs/developer/subsystems/jit-cache.md |
New design/ops documentation for cache keying, layout, invalidation, and MPI considerations. |
docs/developer/index.md |
Adds the new JIT cache subsystem doc to the developer docs toctree. |
docs/advanced/benchmarks/jit_cache_vs_recompile.py |
Adds a benchmark reproducer for the parameter-update recompilation issue. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| cc_id=$(${CC:-cc} --version 2>/dev/null | head -1) | ||
| python_id=$($PIXI run -e "$env" python -c "import sys; print(sys.implementation.name, sys.version)" 2>/dev/null) |
There was a problem hiding this comment.
cc_id=$(${CC:-cc} --version ...) will break if CC contains spaces or additional flags (e.g. CC='ccache clang'), because the shell will try to execute only the first token. If this fingerprint is meant to reflect the actual compiler invocation, consider handling a multi-word CC (or falling back to $PIXI run ... python -c to query the compiler used by the build).
| cc_id=$(${CC:-cc} --version 2>/dev/null | head -1) | |
| python_id=$($PIXI run -e "$env" python -c "import sys; print(sys.implementation.name, sys.version)" 2>/dev/null) | |
| local -a cc_cmd pixi_cmd | |
| read -r -a cc_cmd <<< "${CC:-cc}" | |
| read -r -a pixi_cmd <<< "${PIXI:-pixi}" | |
| cc_id=$("${cc_cmd[@]}" --version 2>/dev/null | head -1) | |
| python_id=$("${pixi_cmd[@]}" run -e "$env" python -c "import sys; print(sys.implementation.name, sys.version)" 2>/dev/null) |
| # Sort by the user-given symbol name, not ``str(expr)`` — ``__str__`` on a | ||
| # UWexpression returns the current *value*, which shuffles the index | ||
| # assignment whenever a value changes. ``.name`` is stable. | ||
| sorted_constants = sorted(constant_exprs, key=lambda e: e.name) |
There was a problem hiding this comment.
constant_exprs is a set(), and sorting only by expr.name can still be non-deterministic if two distinct UWexpressions share the same .name (notably possible for ephemeral expressions created with _unique_name_generation=True). In that case, ties preserve the original set iteration order, which can differ between processes/runs and undermine the MPI determinism guarantees this cache relies on. Consider sorting with a deterministic tie-breaker (e.g. (expr.name, expr.expression_number) or another stable per-expression id) so constants[] indices are reproducible.
| unique_suffix = f"{os.getpid()}_{int(time.time() * 1000)}_{random.randint(1000, 9999)}" | ||
| tmpdir = os.path.join("/tmp", f"{MODNAME}_{unique_suffix}") | ||
| tmpdir = os.path.join("/tmp", f"{modname}_{unique_suffix}") | ||
|
|
There was a problem hiding this comment.
compile_and_load() writes build artifacts to a hard-coded /tmp/... path. On some systems (macOS sandboxing, restricted HPC nodes, custom TMPDIR, containers without /tmp), this can fail even though a usable temp directory exists. Using tempfile.mkdtemp() / TemporaryDirectory() (which respects TMPDIR/TMP/TEMP) would make JIT compilation more robust and avoids having to roll a custom unique-suffix scheme.
| try: | ||
| tf = tempfile.NamedTemporaryFile( | ||
| dir=str(cache_dir), prefix=".inprogress_", suffix=".so", delete=False | ||
| ) | ||
| tf_path = Path(tf.name) | ||
| tf.close() | ||
| shutil.copy2(src_so, tf_path) | ||
| os.replace(tf_path, dst_so) | ||
| except OSError: | ||
| return | ||
|
|
||
| manifest = { | ||
| "version": MANIFEST_VERSION, | ||
| "modname": modname, | ||
| "constants": _manifest_from_constants(constants_manifest), | ||
| } | ||
| try: | ||
| with tempfile.NamedTemporaryFile( | ||
| dir=str(cache_dir), | ||
| prefix=".inprogress_", | ||
| suffix=".json", | ||
| delete=False, | ||
| mode="w", | ||
| encoding="utf-8", | ||
| ) as tf: | ||
| json.dump(manifest, tf) | ||
| tf_path = Path(tf.name) | ||
| os.replace(tf_path, manifest_path) | ||
| except OSError: | ||
| return |
There was a problem hiding this comment.
If the .so copy succeeds but manifest writing fails (or vice versa), the cache entry can be left half-populated (e.g. orphaned .so without a manifest). That causes perpetual cache misses and can leak files over time. Consider cleaning up the already-written counterpart on failure, or writing both to temp paths first and only committing (via os.replace) once both writes have succeeded.
| try: | ||
| tf = tempfile.NamedTemporaryFile( | ||
| dir=str(cache_dir), prefix=".inprogress_", suffix=".so", delete=False | ||
| ) | ||
| tf_path = Path(tf.name) | ||
| tf.close() | ||
| shutil.copy2(src_so, tf_path) | ||
| os.replace(tf_path, dst_so) | ||
| except OSError: | ||
| return |
There was a problem hiding this comment.
On exceptions during the in-progress tempfile writes/copies, the .inprogress_* files are not removed (the except OSError: return paths exit without unlinking). Over time this can accumulate junk in the cache dir and may confuse future tooling. Consider explicitly deleting any created temp files in the exception paths (best-effort) before returning.
| local jit_cache_dir="${UW_JIT_CACHE_DIR:-$HOME/.cache/underworld3/jit}" | ||
| local fp_file="$jit_cache_dir/.env-fingerprint" | ||
| local cc_id python_id pixi_env_id | ||
| cc_id=$(${CC:-cc} --version 2>/dev/null | head -1) | ||
| python_id=$($PIXI run -e "$env" python -c "import sys; print(sys.implementation.name, sys.version)" 2>/dev/null) | ||
| pixi_env_id="$env" | ||
| local fp | ||
| fp=$(printf '%s\nCC=%s\nCFLAGS=%s\nPYTHON=%s\nPIXI_ENV=%s\n' \ | ||
| "fingerprint-v1" "$cc_id" "${CFLAGS:-}" "$python_id" "$pixi_env_id" \ | ||
| | shasum -a 256 | cut -c1-32) | ||
| if [ -f "$fp_file" ]; then | ||
| local prev_fp | ||
| prev_fp=$(cat "$fp_file" 2>/dev/null) | ||
| if [ "$prev_fp" != "$fp" ]; then | ||
| echo " Compiler / Python / pixi env changed; wiping JIT cache..." | ||
| find "$jit_cache_dir" -maxdepth 1 \ | ||
| \( -name '*.so' -o -name '*.manifest.json' -o -name '.*.lock' \ | ||
| -o -name '.inprogress_*' \) -delete 2>/dev/null | ||
| fi | ||
| fi | ||
| mkdir -p "$jit_cache_dir" | ||
| echo "$fp" > "$fp_file" |
There was a problem hiding this comment.
The env-fingerprint logic hard-codes the cache dir as $HOME/.cache/underworld3/jit (unless UW_JIT_CACHE_DIR is set), but the Python-side cache resolver also honours XDG_CACHE_HOME and can be disabled via UW_JIT_CACHE=0. As-is, ./uw build may wipe/write a different directory than the one actually used at runtime, leaving stale cached .so files in the real cache dir (or creating a cache dir even when the disk cache is disabled). Consider using the same resolution rules as underworld3.utilities._jit_cache.get_cache_dir() and skipping this block when UW_JIT_CACHE disables the cache.
| local jit_cache_dir="${UW_JIT_CACHE_DIR:-$HOME/.cache/underworld3/jit}" | |
| local fp_file="$jit_cache_dir/.env-fingerprint" | |
| local cc_id python_id pixi_env_id | |
| cc_id=$(${CC:-cc} --version 2>/dev/null | head -1) | |
| python_id=$($PIXI run -e "$env" python -c "import sys; print(sys.implementation.name, sys.version)" 2>/dev/null) | |
| pixi_env_id="$env" | |
| local fp | |
| fp=$(printf '%s\nCC=%s\nCFLAGS=%s\nPYTHON=%s\nPIXI_ENV=%s\n' \ | |
| "fingerprint-v1" "$cc_id" "${CFLAGS:-}" "$python_id" "$pixi_env_id" \ | |
| | shasum -a 256 | cut -c1-32) | |
| if [ -f "$fp_file" ]; then | |
| local prev_fp | |
| prev_fp=$(cat "$fp_file" 2>/dev/null) | |
| if [ "$prev_fp" != "$fp" ]; then | |
| echo " Compiler / Python / pixi env changed; wiping JIT cache..." | |
| find "$jit_cache_dir" -maxdepth 1 \ | |
| \( -name '*.so' -o -name '*.manifest.json' -o -name '.*.lock' \ | |
| -o -name '.inprogress_*' \) -delete 2>/dev/null | |
| fi | |
| fi | |
| mkdir -p "$jit_cache_dir" | |
| echo "$fp" > "$fp_file" | |
| if [ "${UW_JIT_CACHE:-1}" != "0" ]; then | |
| local jit_cache_dir | |
| if [ -n "${UW_JIT_CACHE_DIR:-}" ]; then | |
| jit_cache_dir="$UW_JIT_CACHE_DIR" | |
| elif [ -n "${XDG_CACHE_HOME:-}" ]; then | |
| jit_cache_dir="$XDG_CACHE_HOME/underworld3/jit" | |
| else | |
| jit_cache_dir="$HOME/.cache/underworld3/jit" | |
| fi | |
| local fp_file="$jit_cache_dir/.env-fingerprint" | |
| local cc_id python_id pixi_env_id | |
| cc_id=$(${CC:-cc} --version 2>/dev/null | head -1) | |
| python_id=$($PIXI run -e "$env" python -c "import sys; print(sys.implementation.name, sys.version)" 2>/dev/null) | |
| pixi_env_id="$env" | |
| local fp | |
| fp=$(printf '%s\nCC=%s\nCFLAGS=%s\nPYTHON=%s\nPIXI_ENV=%s\n' \ | |
| "fingerprint-v1" "$cc_id" "${CFLAGS:-}" "$python_id" "$pixi_env_id" \ | |
| | shasum -a 256 | cut -c1-32) | |
| if [ -f "$fp_file" ]; then | |
| local prev_fp | |
| prev_fp=$(cat "$fp_file" 2>/dev/null) | |
| if [ "$prev_fp" != "$fp" ]; then | |
| echo " Compiler / Python / pixi env changed; wiping JIT cache..." | |
| find "$jit_cache_dir" -maxdepth 1 \ | |
| \( -name '*.so' -o -name '*.manifest.json' -o -name '.*.lock' \ | |
| -o -name '.inprogress_*' \) -delete 2>/dev/null | |
| fi | |
| fi | |
| mkdir -p "$jit_cache_dir" | |
| echo "$fp" > "$fp_file" | |
| fi |
The original Phase-2 design had every MPI rank invoke cc independently on cold-start, which is N× the work and N× the disk hammer. Since UW3 assumes a shared filesystem for the cache directory, we can do better: - Rank 0 invokes compile_and_load once and publishes to the disk cache. - Other ranks barrier-wait, then load the freshly-published .so. - If the disk cache is disabled (UW_JIT_CACHE=0), every rank still compiles independently — no way to share the artefact otherwise. - Defensive fallback: if rank 0's publish failed (e.g. disk full), the other ranks fall back to a local compile so the run still completes correctly. Also adds tests/parallel/ptest_jit_cache.py — five gates that exercise the cache under MPI: 1. allgather hash agreement + C↔S↔C correctness 2. Concurrent cold start publishes exactly one .so per bundle 3. Warm in-memory cache: no recompile across solver instances 4. flock serialises cross-rank store_module calls (no torn writes) 5. NEW: rank-0-only compile — per-rank cc counts assert [1, 0, 0, 0] All five pass under both np=2 and np=4. Note: the docs/advanced/benchmarks/jit_cache_vs_recompile.py VE Stokes benchmark deadlocks at np=4 with the 16x8 mesh. Verified that this deadlock is **pre-existing on origin/development** (reproduced with both _jitextension.py and petsc_generic_snes_solvers.pyx checked out from origin/development) — unrelated to the JIT cache work. Filing a separate issue. Smaller meshes and Poisson under np=4 are unaffected. Underworld development team with AI support from Claude Code (claude.com/claude-code)
CI's tests/test_0004_pointwise_fns.py greps the captured verbose stdout for the prefix "Location of compiled module: " and uses the path to verify the build directory exists. The phase-2 getext() rewrite stopped emitting that line (it was only printed by the legacy _createext() wrapper, which getext no longer calls). Restored: any cold-compile path in getext() now emits the same line on rank 0 when verbose=True, matching the pre-refactor format byte-for-byte so the parsing in test_0004_pointwise_fns.py works unchanged. Underworld development team with AI support from Claude Code (claude.com/claude-code)
Summary
Replaces the sympy-structural JIT cache key with a SHA-256 hash of the
generated C source plus an ABI salt, and persists
.sofiles acrossPython sessions on disk. Resolves both the per-value-recompile bug
(#123) and the multi-minute cold-start cost on real models (#121).
Why a single PR
Both issues stem from the same root cause — sympy structural equality
is not the same as compiled-code equality. Hashing the actual generated
C source makes state equality provable: same C source ⇒ same
.so⇒same solver behaviour, by construction. Cross-session persistence is a
natural extension once the in-process cache uses a stable, content-addressable
key.
What's in the PR
Phase-by-phase commits are intentionally bisectable:
_createextintogenerate_c_source(pure textemission) +
compile_and_load(file I/O + cc + dlopen) + thin_createextwrapper. No behaviour change._ext_dictkey with a 16-char hexhash of the canonicalised C source + ABI salt. Drop
_fn_cache.Fix
_extract_constantsso the constants[] index assignment usesexpr.name(stable) rather thanstr(expr)(current value).Update fast-path 2 in
_build()to refresh_last_jit_cache_keyso a C → S → C transition can't reuse the wrong wired bundle.
src/underworld3/utilities/_jit_cache.pywithload_module/store_moduleand a{hash}.manifest.jsonsidecar. Atomic writes via
os.replace.UW_JIT_CACHE_DIR/XDG_CACHE_HOME/~/.cache/underworld3/jitresolution;UW_JIT_CACHE=0disables.
flockfor inter-process safety,comm.allgatherdeterminism check across MPI ranks, and a.env-fingerprintfile written by./uw buildthat wipes the cachewhen the compiler / Python ABI / pixi env shifts.
docs/developer/subsystems/jit-cache.md) and areproducer benchmark for Parameter update triggers full JIT recompilation despite JITConstant cache mechanism #123.
Test plan
tests/test_jit_cache.py(new): 7 gates covering C↔S↔C transition,value cycling, disk roundtrip, env-var disable, ABI-salt invalidation,
lockfile creation, hash format.
tests/test_1001_poisson_constants.py(existing): all 6 still pass.pytest tests/ -m "level_1 and tier_a"→ 56 passed, 3 skipped (MPI), 0 failed.docs/advanced/benchmarks/jit_cache_vs_recompile.py:30-step VE Stokes BDF-2 with toggling Top BC sign + dt_elastic →
1 compile total (vs the ~90 implied by issue Parameter update triggers full JIT recompilation despite JITConstant cache mechanism #123's profile),
~0.5 s per warm step.
*.soin the cache dir, re-running→ first step 3.6 s; running again → first step 2.4 s (1.2 s saved
by skipping cc).
echo bogus > ~/.cache/underworld3/jit/.env-fingerprintfollowed by
./uw buildcorrectly logs and wipes the cache.Out of scope
rank still performs the cold compile today; only rank 0 publishes the
result. The disk cache amortises this on subsequent runs.
rm -rf ~/.cache/underworld3/jitif itever gets large.
_JITConstant. The placeholder substitution itself is fine;only the naming was buggy.
Closes #121
Closes #123
Underworld development team with AI support from Claude Code