Skip to content

JIT cache: key on generated C source (closes #121, #123) - #129

Merged
lmoresi merged 7 commits into
developmentfrom
bugfix/jit-c-cache
Apr 22, 2026
Merged

lmoresi merged 7 commits into
developmentfrom
bugfix/jit-c-cache

Conversation

@lmoresi

@lmoresi lmoresi commented Apr 22, 2026

Copy link
Copy Markdown
Member

Summary

Replaces the sympy-structural JIT cache key with a SHA-256 hash of the
generated C source plus an ABI salt, and persists .so files across
Python 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:

  1. Phase 1 — split _createext into generate_c_source (pure text
    emission) + compile_and_load (file I/O + cc + dlopen) + thin
    _createext wrapper. No behaviour change.
  2. Phase 2 — replace the integer _ext_dict key with a 16-char hex
    hash of the canonicalised C source + ABI salt. Drop _fn_cache.
    Fix _extract_constants so the constants[] index assignment uses
    expr.name (stable) rather than str(expr) (current value).
    Update fast-path 2 in _build() to refresh _last_jit_cache_key
    so a C → S → C transition can't reuse the wrong wired bundle.
  3. Phase 3 — new src/underworld3/utilities/_jit_cache.py with
    load_module / store_module and a {hash}.manifest.json
    sidecar. Atomic writes via os.replace. UW_JIT_CACHE_DIR /
    XDG_CACHE_HOME / ~/.cache/underworld3/jit resolution; UW_JIT_CACHE=0
    disables.
  4. Phase 4 — per-hash flock for inter-process safety,
    comm.allgather determinism check across MPI ranks, and a
    .env-fingerprint file written by ./uw build that wipes the cache
    when the compiler / Python ABI / pixi env shifts.
  5. Phase 5 — docs (docs/developer/subsystems/jit-cache.md) and a
    reproducer 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.
  • Benchmark 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.
  • Cross-session: deleting all *.so in the cache dir, re-running
    → first step 3.6 s; running again → first step 2.4 s (1.2 s saved
    by skipping cc).
  • Manual ABI-skew test: echo bogus > ~/.cache/underworld3/jit/.env-fingerprint
    followed by ./uw build correctly logs and wipes the cache.

Out of scope

  • True "rank-0 compiles, others wait on barrier" optimisation — every
    rank still performs the cold compile today; only rank 0 publishes the
    result. The disk cache amortises this on subsequent runs.
  • Cache eviction policy. Users rm -rf ~/.cache/underworld3/jit if it
    ever gets large.
  • Removing _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

lmoresi added 5 commits April 22, 2026 08:20
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)
Copilot AI review requested due to automatic review settings April 22, 2026 00:14

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread uw
Comment on lines +288 to +289
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)

Copilot AI Apr 22, 2026

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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).

Suggested change
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)

Copilot uses AI. Check for mistakes.
Comment on lines +336 to +339
# 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)

Copilot AI Apr 22, 2026

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copilot uses AI. Check for mistakes.
Comment on lines 1221 to 1223
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}")

Copilot AI Apr 22, 2026

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copilot uses AI. Check for mistakes.
Comment on lines +198 to +227
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

Copilot AI Apr 22, 2026

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copilot uses AI. Check for mistakes.
Comment on lines +198 to +207
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

Copilot AI Apr 22, 2026

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copilot uses AI. Check for mistakes.
Comment thread uw
Comment on lines +285 to +306
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"

Copilot AI Apr 22, 2026

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Suggested change
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

Copilot uses AI. Check for mistakes.
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)
@lmoresi
lmoresi merged commit 0dcd61f into development Apr 22, 2026
1 check passed
@lmoresi
lmoresi deleted the bugfix/jit-c-cache branch June 13, 2026 00:48
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.

2 participants