From ee705b6c7c1e2459216d166c259df0e19b9f49a8 Mon Sep 17 00:00:00 2001 From: lmoresi Date: Wed, 22 Apr 2026 08:20:41 +1000 Subject: [PATCH 1/7] JIT refactor (phase 1/5): split _createext into generate + compile MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) --- src/underworld3/utilities/_jitextension.py | 172 +++++++++++++++------ 1 file changed, 124 insertions(+), 48 deletions(-) diff --git a/src/underworld3/utilities/_jitextension.py b/src/underworld3/utilities/_jitextension.py index ce8c0621b..0824c9120 100644 --- a/src/underworld3/utilities/_jitextension.py +++ b/src/underworld3/utilities/_jitextension.py @@ -731,8 +731,8 @@ def getext( @timing.routine_timer_decorator -def _createext( - name: str, +def generate_c_source( + name, mesh: underworld3.discretisation.Mesh, callbacks: JITCallbackSet, primary_field_list, @@ -741,23 +741,36 @@ def _createext( debug: Optional[bool] = False, debug_name=None, ): - """Create the JIT extension module with PETSc function pointers. + """Generate the setup.py / C header / Cython wrapper for a JIT bundle. + + This is the pure text-generation phase: sympy processing, C-code emission, + and assembly of the files that will make up the compiled module. No I/O, + no subprocess, no dynamic loading — those happen in ``compile_and_load``. - Note that it is not possible to replace loaded shared libraries - in Python, so we instead create a new extension for each new function. + Keying a cache on a hash of the generated C source requires that this + function produce byte-identical output for byte-identical inputs. Parameters ---------- - name : str - Name for the extension. It will be prepended with "fn_ptr_ext_". + name : str or int + Identifier used to build ``MODNAME = "fn_ptr_ext_" + str(name)``. mesh : Mesh - Supporting mesh for coordinate system and variable information. callbacks : JITCallbackSet - Structured callback expressions grouped by role. primary_field_list : list Variables that map to PETSc primary variable arrays (``petsc_u[]``). - All other variables map to auxiliary arrays (``petsc_a[]``). - Must be ordered by ``field_id``. + constants_subs_map : dict, optional + Mapping from UWexpression → ``_JITConstant`` placeholder. + + Returns + ------- + modname : str + Fully-qualified extension module name (``fn_ptr_ext_``). + codeguys : list of [filename, content] + The files that make up the source bundle + (``setup.py``, ``cy_ext.h``, ``cy_ext.pyx``). + diagnostics : dict + Equation-range counts and the random symbol prefix, used by the caller + for verbose printing and for building the fn-layout manifest. """ from sympy import symbols, Eq, MatrixSymbol from underworld3 import VarType @@ -1228,15 +1241,54 @@ def _basescalar_ccode(self, printer): pyx_str += " return clsguy" codeguys.append(["cy_ext.pyx", pyx_str]) - # Write out files - import os + diagnostics = { + "randstr": randstr, + "eqn_count": eqn_count, + "count_residual_sig": count_residual_sig, + "count_bc_sig": count_bc_sig, + "count_jacobian_sig": count_jacobian_sig, + "count_bd_residual_sig": count_bd_residual_sig, + "count_bd_jacobian_sig": count_bd_jacobian_sig, + "residual_equations": residual_equations, + "boundary_equations": boundary_equations, + "jacobian_equations": jacobian_equations, + "boundary_residual_equations": boundary_residual_equations, + "boundary_jacobian_equations": boundary_jacobian_equations, + } + return MODNAME, codeguys, diagnostics + + +def compile_and_load(modname, codeguys, verbose=False): + """Write ``codeguys`` files to a temp directory, build with Cython, and dynamically load. + Split out of the former ``_createext`` so that generation and compilation + can be hashed/cached independently. + + Parameters + ---------- + modname : str + Fully-qualified name of the extension module (``fn_ptr_ext_``). + codeguys : list of [filename, content] + Source files from :func:`generate_c_source`. + verbose : bool, optional + Print build diagnostics on failure. + + Returns + ------- + module : loaded Python extension module exposing ``getptrobj()``. + tmpdir : str + Location of the generated sources, useful when a caller wants to + persist the ``.so`` elsewhere (e.g. a cross-session disk cache). + """ + import os + import sys import time import random + import importlib.machinery + from importlib._bootstrap import _load - # Make directory name unique to avoid race conditions between parallel processes 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}") try: os.makedirs(tmpdir, exist_ok=True) @@ -1244,15 +1296,13 @@ def _basescalar_ccode(self, printer): if verbose: print(f"Warning: Failed to create tmpdir {tmpdir}: {e}") raise RuntimeError(f"Cannot create temporary directory {tmpdir}") from e + for thing in codeguys: filename = thing[0] strguy = thing[1] with open(os.path.join(tmpdir, filename), "w") as f: f.write(strguy) - # Build - import sys - process = subprocess.Popen( [sys.executable] + "setup.py build_ext --inplace".split(), stdout=subprocess.PIPE, @@ -1262,42 +1312,32 @@ def _basescalar_ccode(self, printer): ) stdout, stderr = process.communicate() - # Check if build process failed if process.returncode != 0: if verbose: print(f"Warning: Build process failed with return code {process.returncode}") print(f"stdout: {stdout.decode() if stdout else 'None'}") print(f"stderr: {stderr.decode() if stderr else 'None'}") - # Load and add to dictionary - from importlib._bootstrap import _load + def load_dynamic(name, path): + """Load an extension module from ``path``. - def load_dynamic(name, path, file=None): - """ - Load an extension module. - Borrowed from: - https://stackoverflow.com/a/55172547 + Borrowed from https://stackoverflow.com/a/55172547 — skips the + ``sys.modules`` reuse check so we always load a fresh extension. """ - import importlib.machinery - loader = importlib.machinery.ExtensionFileLoader(name, path) - - # Issue #24748: Skip the sys.modules check in _load_module_shims - # always load new extension spec = importlib.machinery.ModuleSpec(name=name, loader=loader, origin=path) return _load(spec) - # Check if tmpdir exists before trying to list it + module = None if os.path.exists(tmpdir): for _file in os.listdir(tmpdir): if _file.endswith(".so"): - _ext_dict[name] = load_dynamic(MODNAME, os.path.join(tmpdir, _file)) - else: - # tmpdir doesn't exist, likely build process failed - if verbose: - print(f"Warning: tmpdir {tmpdir} does not exist - build process may have failed") + module = load_dynamic(modname, os.path.join(tmpdir, _file)) + break + elif verbose: + print(f"Warning: tmpdir {tmpdir} does not exist - build process may have failed") - if name not in _ext_dict.keys(): + if module is None: raise RuntimeError( f"The Underworld extension module does not appear to have been built successfully. " f"The generated module may be found at:\n {str(tmpdir)}\n" @@ -1309,31 +1349,67 @@ def load_dynamic(name, path, file=None): f"Please contact the developers if you are unable to resolve the issue." ) + return module, tmpdir + + +@timing.routine_timer_decorator +def _createext( + name, + mesh: underworld3.discretisation.Mesh, + callbacks: JITCallbackSet, + primary_field_list, + constants_subs_map: Optional[dict] = None, + verbose: Optional[bool] = False, + debug: Optional[bool] = False, + debug_name=None, +): + """Thin wrapper: generate source, compile, stash in ``_ext_dict[name]``. + + Retained for backwards compatibility with :func:`getext`. New code + should call :func:`generate_c_source` and :func:`compile_and_load` + directly — splitting the two phases is what makes cache keys on the + generated C source possible. + """ + modname, codeguys, diag = generate_c_source( + name, + mesh, + callbacks, + primary_field_list, + constants_subs_map=constants_subs_map, + verbose=verbose, + debug=debug, + debug_name=debug_name, + ) + module, tmpdir = compile_and_load(modname, codeguys, verbose=verbose) + _ext_dict[name] = module + if underworld3.mpi.rank == 0 and verbose: + randstr = diag["randstr"] print(f"Location of compiled module: {str(tmpdir)}") - - print( - f"{randstr} Equation count - {eqn_count}", - flush=True, - ) + print(f"{randstr} Equation count - {diag['eqn_count']}", flush=True) print( - f"{randstr} {count_residual_sig:5d} residuals: {residual_equations[0]}:{residual_equations[1]}", + f"{randstr} {diag['count_residual_sig']:5d} residuals: " + f"{diag['residual_equations'][0]}:{diag['residual_equations'][1]}", flush=True, ) print( - f"{randstr} {count_bc_sig:5d} boundaries: {boundary_equations[0]}:{boundary_equations[1]}", + f"{randstr} {diag['count_bc_sig']:5d} boundaries: " + f"{diag['boundary_equations'][0]}:{diag['boundary_equations'][1]}", flush=True, ) print( - f"{randstr} {count_jacobian_sig:5d} jacobians: {jacobian_equations[0]}:{jacobian_equations[1]}", + f"{randstr} {diag['count_jacobian_sig']:5d} jacobians: " + f"{diag['jacobian_equations'][0]}:{diag['jacobian_equations'][1]}", flush=True, ) print( - f"{randstr} {count_bd_residual_sig:5d} boundary_res: {boundary_residual_equations[0]}:{boundary_residual_equations[1]}", + f"{randstr} {diag['count_bd_residual_sig']:5d} boundary_res: " + f"{diag['boundary_residual_equations'][0]}:{diag['boundary_residual_equations'][1]}", flush=True, ) print( - f"{randstr} {count_bd_jacobian_sig:5d} boundary_jac: {boundary_jacobian_equations[0]}:{boundary_jacobian_equations[1]}", + f"{randstr} {diag['count_bd_jacobian_sig']:5d} boundary_jac: " + f"{diag['boundary_jacobian_equations'][0]}:{diag['boundary_jacobian_equations'][1]}", flush=True, ) From 0393aacc70ee9a917277f25de712f13b00032dfb Mon Sep 17 00:00:00 2001 From: lmoresi Date: Wed, 22 Apr 2026 09:15:37 +1000 Subject: [PATCH 2/7] JIT cache (phase 2/5): key on C-source hash, drop _fn_cache MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) --- .../cython/petsc_generic_snes_solvers.pyx | 8 + src/underworld3/utilities/_jitextension.py | 286 ++++++------------ tests/test_jit_cache.py | 161 ++++++++++ 3 files changed, 262 insertions(+), 193 deletions(-) create mode 100644 tests/test_jit_cache.py diff --git a/src/underworld3/cython/petsc_generic_snes_solvers.pyx b/src/underworld3/cython/petsc_generic_snes_solvers.pyx index 69a47a1a8..58a8b0752 100644 --- a/src/underworld3/cython/petsc_generic_snes_solvers.pyx +++ b/src/underworld3/cython/petsc_generic_snes_solvers.pyx @@ -519,6 +519,14 @@ class SolverBaseClass(uw_object): self._setup_solver(verbose, _rewire_only=True) self.is_setup = True + + # _last_jit_cache_key must track "which bundle the solver is + # currently wired for", not "what the last full-build was". + # Otherwise a C → S → C transition can make fast path 1 trigger + # on a subsequent _build() call while the solver is still wired + # for S — returning S's result for a C solve. + if hasattr(self, '_current_jit_cache_key'): + self._last_jit_cache_key = self._current_jit_cache_key return # === Full rebuild path — teardown DM and reconstruct === diff --git a/src/underworld3/utilities/_jitextension.py b/src/underworld3/utilities/_jitextension.py index 0824c9120..4fbf2eb1d 100644 --- a/src/underworld3/utilities/_jitextension.py +++ b/src/underworld3/utilities/_jitextension.py @@ -125,22 +125,29 @@ def _openmpi_wrapper_fallback(wrapper, env_key): # deriv = subfn_d.xreplace({uwderivdummy:fn2}) # sub out dummy # return deriv +# In-memory cache of compiled modules, keyed by a hex-string hash of the +# generated C source (plus an ABI salt). Tests count ``len(_ext_dict)`` to +# verify that a parameter-value change does not trigger a new compile. _ext_dict = {} -# Per-function cache: maps (fn_hash, sig_type) -> _CachedFn -_fn_cache = {} +def _abi_salt(): + """Return a string that invalidates the cache when binary compatibility changes. -@dataclass -class _CachedFn: - """Registry entry for a single cached compiled function pointer.""" - ptr_container: object # PtrContainer from the .so that compiled this fn - sig_type: str # "residual" | "jacobian" | "bcs" | "bd_residual" | "bd_jacobian" - index: int # index within that container's sig_type array - - -# Signature type names used by the per-function cache -_SIG_TYPES = ("residual", "bcs", "jacobian", "bd_residual", "bd_jacobian") + Keeps the salt conservative (PETSc version + underworld3 version). The + ``./uw`` build driver is responsible for wiping the on-disk cache when + the compiler toolchain or Python ABI changes (see design doc). + """ + try: + from petsc4py import PETSc + petsc_ver = ".".join(str(x) for x in PETSc.Sys.getVersion()) + except Exception: + petsc_ver = "unknown" + try: + uw_ver = underworld3.__version__ + except AttributeError: + uw_ver = "unknown" + return f"petsc={petsc_ver}|uw={uw_ver}" # ============================================================================ @@ -326,13 +333,17 @@ def _extract_constants(all_fns, mesh): if not constant_exprs: return [], {} - # Sort by name for deterministic MPI-consistent ordering - sorted_constants = sorted(constant_exprs, key=lambda e: str(e)) + # 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) manifest = [] subs_map = {} for i, expr in enumerate(sorted_constants): - jit_const = _JITConstant(i, name=f"_jit_const_{str(expr)}") + # Use ``expr.name`` (stable) instead of ``str(expr)`` (= current value) + # so the placeholder symbol's identity is independent of parameter value. + jit_const = _JITConstant(i, name=f"_jit_const_{expr.name}") manifest.append((i, expr)) subs_map[expr] = jit_const @@ -523,195 +534,84 @@ def getext( constants_manifest is a list of (index, uw_expression_ref) tuples for use with PetscDSSetConstants(). """ - import time + import hashlib - time_s = time.time() primary_field_list = tuple(primary_field_list) - # Extract constant UWexpressions that will go through constants[] array + # Extract constant UWexpressions that are routed through PETSc's + # constants[] array. Value changes don't affect the C source — they + # only alter what we pass to PetscDSSetConstants at solve time. constants_manifest, constants_subs_map = _extract_constants(callbacks.flat(), mesh) - # Build structurally-expanded functions for cache hashing. - # Constants are replaced with placeholder symbols (value-independent), - # so changing a constant value won't cause a cache miss. - expanded = callbacks.map(lambda fn: prepare_for_cache_key(fn, constants_subs_map)) - if debug and underworld3.mpi.rank == 0: - print(f"Expanded functions for compilation:") - for i, fn in enumerate(expanded.flat()): - print(f"{i}: {fn}") if constants_manifest: print(f"Constants manifest ({len(constants_manifest)} entries):") for idx, expr in constants_manifest: - print(f" constants[{idx}] = {expr} (current value: {expr.data})") - - import os - - primary_field_signature = tuple( - (getattr(field, "field_id", None), getattr(field, "clean_name", None)) - for field in primary_field_list + print(f" constants[{idx}] = {expr.name} (current value: {expr.data})") + + # Generate C source. The returned ``gen_modname``/``gen_randstr`` are + # implementation-dependent (often random) — we canonicalise them below + # so byte-identical inputs hash to the same key. + _PLACEHOLDER_NAME = "UWJITPLACEHOLDER" + gen_modname, codeguys, diag = generate_c_source( + _PLACEHOLDER_NAME, + mesh, + callbacks, + primary_field_list, + constants_subs_map=constants_subs_map, + verbose=verbose, + debug=debug, + debug_name=debug_name, ) - - if debug_name is not None: - jitname = debug_name - - elif "UW_JITNAME" in os.environ: # If var specified, probably testing. - jitname = os.environ["UW_JITNAME"] - # Note, extensions cannot be replaced, so need to append count to ensure - # unique modules. - jitname += "_" + str(len(_ext_dict.keys())) - - else: # Name from structured hash — function role must be preserved. - jitname = abs( - hash((mesh, expanded.signature(), tuple(mesh.vars.keys()), primary_field_signature)) - ) - - # ── Fast path: whole-bundle cache hit ────────────────────────────────── - if jitname in _ext_dict and cache: + gen_randstr = diag["randstr"] + + # Canonicalise: swap the call-specific identifiers for stable tokens so + # the hash is a function of *what the module does*, not what it happens + # to be named. The same trick lets us derive a unique-per-bundle, yet + # deterministic, C-symbol prefix below (required by some dlopen loaders + # that use RTLD_GLOBAL — see historical comment on ``randstr``). + _CAN_MOD = "__UW_JIT_MOD__" + _CAN_RS = "__UW_JIT_RS__" + canonical_codeguys = [ + [ + entry[0], + entry[1].replace(gen_modname, _CAN_MOD).replace(gen_randstr, _CAN_RS), + ] + for entry in codeguys + ] + canonical_source = "\n".join(entry[1] for entry in canonical_codeguys) + source_hash = hashlib.sha256( + (canonical_source + "\n---\n" + _abi_salt()).encode("utf-8") + ).hexdigest()[:16] + + # Derive the real modname/randstr from the hash — same source ⇒ same + # compiled artefact, different sources ⇒ disjoint symbol namespaces. + real_modname = f"fn_ptr_ext_{source_hash}" + real_randstr = "UW" + source_hash[:8].upper() + codeguys_final = [ + [ + entry[0], + entry[1].replace(_CAN_MOD, real_modname).replace(_CAN_RS, real_randstr), + ] + for entry in canonical_codeguys + ] + + # ── Cache lookup or build ───────────────────────────────────────────── + if cache and source_hash in _ext_dict: if verbose and underworld3.mpi.rank == 0: - print(f"JIT compiled module cached ... {jitname} ", flush=True) - - module = _ext_dict[jitname] - ptrobj = module.getptrobj() - - i_res = {fn: i for i, fn in enumerate(callbacks.residual)} - i_ebc = {fn: i for i, fn in enumerate(callbacks.bcs)} - i_jac = {fn: i for i, fn in enumerate(callbacks.jacobian)} - i_bd_res = {fn: i for i, fn in enumerate(callbacks.bd_residual)} - i_bd_jac = {fn: i for i, fn in enumerate(callbacks.bd_jacobian)} - - extn_fn_dict = namedtuple( - "Functions", ["res", "jac", "ebc", "bd_res", "bd_jac"], - ) - return _GextResult( - ptrobj, - extn_fn_dict(i_res, i_jac, i_ebc, i_bd_res, i_bd_jac), - constants_manifest, - cache_key=jitname, - ) - - # ── Per-function cache: check which individual functions are cached ── - # Build a hashable key for the constants manifest so it's part of - # every per-function hash (ensures constants[i] means the same thing). - constants_manifest_key = tuple( - (str(expr), idx) for idx, expr in constants_manifest - ) - - # Compute per-function hashes for each expression, grouped by sig_type - fn_hashes = {} # maps (sig_type, slot_index) -> hash - for sig_type, slot_fns in zip( - _SIG_TYPES, - [expanded.residual, expanded.bcs, expanded.jacobian, - expanded.bd_residual, expanded.bd_jacobian], - ): - for i, fn_expanded in enumerate(slot_fns): - fn_hashes[(sig_type, i)] = abs( - hash((mesh, fn_expanded, tuple(mesh.vars.keys()), - sig_type, constants_manifest_key, primary_field_signature)) - ) - - # Partition into cached vs new - cached_hits = {} # (sig_type, slot_index) -> _CachedFn - new_needed = {} # (sig_type, slot_index) -> original expression - for sig_type, slot_fns in zip( - _SIG_TYPES, - [callbacks.residual, callbacks.bcs, callbacks.jacobian, - callbacks.bd_residual, callbacks.bd_jacobian], - ): - for i, fn_orig in enumerate(slot_fns): - key = (sig_type, i) - fn_hash = fn_hashes[key] - cache_key = (fn_hash, sig_type) - if cache_key in _fn_cache and cache: - cached_hits[key] = _fn_cache[cache_key] - else: - new_needed[key] = fn_orig - - n_hits = len(cached_hits) - n_new = len(new_needed) - - if verbose and underworld3.mpi.rank == 0: - total = n_hits + n_new - print(f"Per-function cache: {n_hits}/{total} hits, {n_new} new", flush=True) - - # ── Compile new functions ─────────────────────────────────────────── - new_ptr = None - new_indices = {} # (sig_type, slot_index) -> index in new_ptr's arrays - - if n_new > 0: - # Build a JITCallbackSet containing ONLY the new functions, - # preserving order within each sig_type. - new_by_type = {st: [] for st in _SIG_TYPES} - new_slot_map = {st: [] for st in _SIG_TYPES} # tracks original slot indices - for (sig_type, slot_idx), fn_orig in sorted(new_needed.items()): - new_by_type[sig_type].append(fn_orig) - new_slot_map[sig_type].append(slot_idx) - - new_callbacks = JITCallbackSet( - residual=tuple(new_by_type["residual"]), - bcs=tuple(new_by_type["bcs"]), - jacobian=tuple(new_by_type["jacobian"]), - bd_residual=tuple(new_by_type["bd_residual"]), - bd_jacobian=tuple(new_by_type["bd_jacobian"]), - ) - - # Compile the new functions - new_jitname = abs(hash((jitname, "partial", n_new, time.time()))) - _createext( - new_jitname, - mesh, - new_callbacks, - primary_field_list, - constants_subs_map=constants_subs_map, - verbose=verbose, - debug=debug, - debug_name=debug_name, - ) - - new_module = _ext_dict[new_jitname] - new_ptr = new_module.getptrobj() - - # Register each new function in the per-function cache - for sig_type in _SIG_TYPES: - for local_idx, slot_idx in enumerate(new_slot_map[sig_type]): - key = (sig_type, slot_idx) - fn_hash = fn_hashes[key] - cache_key = (fn_hash, sig_type) - entry = _CachedFn( - ptr_container=new_ptr, - sig_type=sig_type, - index=local_idx, - ) - _fn_cache[cache_key] = entry - cached_hits[key] = entry - - # Also register the full bundle in _ext_dict if ALL functions were new - # (common case: first compile of a solver) - if n_new > 0 and n_hits == 0: - _ext_dict[jitname] = _ext_dict[new_jitname] - - # ── Assemble PtrContainer from cached function pointers ───────────── - from underworld3.cython.petsc_types import PtrContainer - - result_ptr = PtrContainer() - counts = callbacks.counts - result_ptr.allocate(*counts) - - _copy_methods = { - "residual": result_ptr.copy_residual_from, - "bcs": result_ptr.copy_bcs_from, - "jacobian": result_ptr.copy_jacobian_from, - "bd_residual": result_ptr.copy_bd_residual_from, - "bd_jacobian": result_ptr.copy_bd_jacobian_from, - } + print(f"JIT compiled module cached ... {source_hash}", flush=True) + module = _ext_dict[source_hash] + else: + if verbose and underworld3.mpi.rank == 0: + print(f"JIT compiling new module ... {source_hash}", flush=True) + module, _tmpdir = compile_and_load(real_modname, codeguys_final, verbose=verbose) + if cache: + _ext_dict[source_hash] = module - for sig_type, n_fns in zip(_SIG_TYPES, counts): - copy_fn = _copy_methods[sig_type] - for slot_idx in range(n_fns): - entry = cached_hits[(sig_type, slot_idx)] - copy_fn(slot_idx, entry.ptr_container, entry.index) + ptrobj = module.getptrobj() - # ── Build fn_dicts (unchanged from original) ──────────────────────── + # Build the slot-index dicts. Keys are the original callback objects so + # solvers can do ``ext.fns_residual[ext_dict.res[self._u_f0]]``. i_res = {fn: i for i, fn in enumerate(callbacks.residual)} i_ebc = {fn: i for i, fn in enumerate(callbacks.bcs)} i_jac = {fn: i for i, fn in enumerate(callbacks.jacobian)} @@ -723,10 +623,10 @@ def getext( ) return _GextResult( - result_ptr, + ptrobj, extn_fn_dict(i_res, i_jac, i_ebc, i_bd_res, i_bd_jac), constants_manifest, - cache_key=jitname, + cache_key=source_hash, ) diff --git a/tests/test_jit_cache.py b/tests/test_jit_cache.py new file mode 100644 index 000000000..461c40b31 --- /dev/null +++ b/tests/test_jit_cache.py @@ -0,0 +1,161 @@ +# Tests for the C-source-hash JIT cache (issues #121, #123). +# +# The cache key is a SHA-256 hash of the canonicalised C source plus an ABI +# salt. This makes state equality *provable*: +# same C source ⇒ same hash ⇒ same .so ⇒ same solver behaviour. +# +# Gate 1 — constant ↔ symbolic ↔ constant transition +# Verifies that a Poisson solve with diff=constant, then diff=1+u, then +# diff=constant returns the *exact* state-1 result on state 3. With the +# old sympy-structural cache, state 3 either cache-missed (recompile) or +# returned a contaminated state-2 result. +# +# Gate 2 — value cycle +# Toggling a constant's value between two fixed numbers should produce +# exactly one cache entry (the C source is value-independent). + +import numpy as np +import pytest +import sympy + +import underworld3 as uw +from underworld3.utilities._jitextension import _ext_dict + +pytestmark = [pytest.mark.level_1, pytest.mark.tier_a] + + +@pytest.fixture(autouse=True) +def reset_model_state(): + uw.reset_default_model() + uw.use_strict_units(False) + uw.use_nondimensional_scaling(False) + yield + uw.reset_default_model() + uw.use_strict_units(False) + uw.use_nondimensional_scaling(False) + + +def _t_max(u_field, sample_points): + u = uw.function.evaluate(u_field.sym[0], sample_points, rbf=False).squeeze() + return float(np.max(np.abs(u))) + + +def test_c_s_c_transition_returns_state_1_result_on_state_3(): + """diff=0.5 → 1+u → 0.5. The third solve must match the first exactly. + + This is the bug that motivated the C-source-hash architecture: with + the old sympy-structural cache, state 3 returned a contaminated + state-2 numerical result even though it was structurally identical + to state 1. + """ + + mesh = uw.meshing.UnstructuredSimplexBox(cellSize=0.15) + + u = uw.discretisation.MeshVariable("u_csc", mesh, 1, degree=2) + + K = uw.expression("K_csc", 0.5) + + poisson = uw.systems.Poisson(mesh, u_Field=u) + poisson.constitutive_model = uw.constitutive_models.DiffusionModel + poisson.constitutive_model.Parameters.diffusivity = K + poisson.f = 1.0 + + poisson.add_dirichlet_bc(0.0, "Top") + poisson.add_dirichlet_bc(0.0, "Bottom") + + sample_y = np.linspace(0.05, 0.95, 15) + sample_x = np.full_like(sample_y, 0.5) + sample_points = np.column_stack([sample_x, sample_y]) + + # State 1: K = 0.5 (constant) + poisson.solve() + assert poisson.snes.getConvergedReason() > 0 + t_max_1 = _t_max(u, sample_points) + + # State 2: K = 1 + u (symbolic, nonlinear) + poisson.constitutive_model.Parameters.diffusivity = 1.0 + u.sym[0] + poisson.solve() + assert poisson.snes.getConvergedReason() > 0 + t_max_2 = _t_max(u, sample_points) + + # State 3: K = 0.5 again (constant, same as state 1) + poisson.constitutive_model.Parameters.diffusivity = K + poisson.solve() + assert poisson.snes.getConvergedReason() > 0 + t_max_3 = _t_max(u, sample_points) + + assert t_max_2 != pytest.approx(t_max_1, abs=1e-3), ( + f"State 2 (K=1+u) result {t_max_2:.6f} is suspiciously close to state 1 " + f"(K=0.5) result {t_max_1:.6f}; the symbolic branch is not being taken." + ) + + assert t_max_3 == pytest.approx(t_max_1, abs=1e-6), ( + f"State 3 must equal state 1 (same K=0.5), got " + f"T_max_1={t_max_1:.8f}, T_max_3={t_max_3:.8f}. " + f"C-source hash keying should guarantee identical behaviour." + ) + + del poisson + + +def test_value_cycle_produces_single_cache_entry(): + """Toggle a constant between two values 20×; C source is value-independent.""" + + mesh = uw.meshing.UnstructuredSimplexBox(cellSize=0.2) + + u = uw.discretisation.MeshVariable("u_cycle", mesh, 1, degree=2) + + K = uw.expression("K_cycle", 1.0) + + poisson = uw.systems.Poisson(mesh, u_Field=u) + poisson.constitutive_model = uw.constitutive_models.DiffusionModel + poisson.constitutive_model.Parameters.diffusivity = K + poisson.f = -2.0 + + poisson.add_dirichlet_bc(0.0, "Bottom") + poisson.add_dirichlet_bc(1.0, "Top") + + # Warm the cache with the first solve + poisson.solve() + n_entries_after_first = len(_ext_dict) + + # Toggle between two values 20 times; each solve must cache-hit + values = [0.5, 1.0] + for i in range(20): + K.sym = values[i % 2] + poisson.solve() + assert poisson.snes.getConvergedReason() > 0 + + n_entries_after_cycle = len(_ext_dict) + + assert n_entries_after_cycle == n_entries_after_first, ( + f"Toggling a constant's value caused {n_entries_after_cycle - n_entries_after_first} " + f"new JIT compile(s). With C-source-hash keying, the source is value-independent." + ) + + del poisson + + +def test_cache_key_is_deterministic_hex_string(): + """The cache_key on _GextResult is a 16-char hex string (source hash).""" + + mesh = uw.meshing.UnstructuredSimplexBox(cellSize=0.25) + + u = uw.discretisation.MeshVariable("u_k", mesh, 1, degree=2) + + poisson = uw.systems.Poisson(mesh, u_Field=u) + poisson.constitutive_model = uw.constitutive_models.DiffusionModel + poisson.constitutive_model.Parameters.diffusivity = 1.0 + poisson.f = 0.0 + poisson.add_dirichlet_bc(1.0, "Bottom") + poisson.add_dirichlet_bc(0.0, "Top") + poisson.solve() + + key = poisson._current_jit_cache_key + assert isinstance(key, str), f"cache_key should be str, got {type(key).__name__}" + assert len(key) == 16, f"cache_key should be 16 hex chars, got {len(key)}: {key!r}" + assert all(c in "0123456789abcdef" for c in key), ( + f"cache_key should be lowercase hex, got {key!r}" + ) + + del poisson From 44c9a71cf5896463a23b4c8c2ce23195f92ba803 Mon Sep 17 00:00:00 2001 From: lmoresi Date: Wed, 22 Apr 2026 09:22:18 +1000 Subject: [PATCH 3/7] JIT cache (phase 3/5): persistent on-disk cache + manifest JSON MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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_", "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) --- src/underworld3/utilities/_jit_cache.py | 185 +++++++++++++++++++++ src/underworld3/utilities/_jitextension.py | 24 ++- tests/test_jit_cache.py | 79 +++++++++ 3 files changed, 284 insertions(+), 4 deletions(-) create mode 100644 src/underworld3/utilities/_jit_cache.py diff --git a/src/underworld3/utilities/_jit_cache.py b/src/underworld3/utilities/_jit_cache.py new file mode 100644 index 000000000..3bf47d6b6 --- /dev/null +++ b/src/underworld3/utilities/_jit_cache.py @@ -0,0 +1,185 @@ +"""Persistent on-disk cache for JIT-compiled UW3 solver extensions. + +Populated by :func:`getext` the first time a bundle is compiled; queried on +subsequent runs (same process or fresh Python process) to skip the Cython ++ cc compile step entirely. Cache entries are keyed by the same SHA-256 of +the canonical C source that drives the in-memory cache — see +``docs/developer/subsystems/jit-cache.md`` for the full design. + +Cache directory resolution: + +1. ``$UW_JIT_CACHE_DIR`` if set +2. ``$XDG_CACHE_HOME/underworld3/jit`` if ``XDG_CACHE_HOME`` set +3. ``~/.cache/underworld3/jit`` otherwise + +The cache is disabled entirely when ``UW_JIT_CACHE=0`` (or ``false``/``no``), +in which case both :func:`load_module` and :func:`store_module` become no-ops. + +MPI is **not** handled here — phase 4 of the JIT cache refactor adds rank-0 +gating and ``flock``-based inter-process synchronisation in this same module. +""" + +from __future__ import annotations + +import importlib.machinery +import json +import os +import shutil +import tempfile +from importlib._bootstrap import _load +from pathlib import Path +from typing import Optional + + +MANIFEST_VERSION = 1 + + +def get_cache_dir() -> Optional[Path]: + """Return the active JIT cache directory, or ``None`` if disabled.""" + if os.environ.get("UW_JIT_CACHE", "").lower() in ("0", "false", "no"): + return None + explicit = os.environ.get("UW_JIT_CACHE_DIR") + if explicit: + return Path(explicit).expanduser() + xdg = os.environ.get("XDG_CACHE_HOME") + if xdg: + return Path(xdg).expanduser() / "underworld3" / "jit" + return Path.home() / ".cache" / "underworld3" / "jit" + + +def _so_path(cache_dir: Path, source_hash: str) -> Path: + return cache_dir / f"{source_hash}.so" + + +def _manifest_path(cache_dir: Path, source_hash: str) -> Path: + return cache_dir / f"{source_hash}.manifest.json" + + +def _manifest_from_constants(constants_manifest): + """Canonicalise ``constants_manifest`` for JSON storage / comparison. + + ``constants_manifest`` is a list of ``(index, UWexpression)`` tuples. + We persist only the stable ``expr.name`` — the current value is irrelevant + (it's updated via ``PetscDSSetConstants`` at solve time). + """ + return [{"index": idx, "name": expr.name} for idx, expr in constants_manifest] + + +def load_module(source_hash: str, modname: str, constants_manifest): + """Try to load a cached compiled module for ``source_hash``. + + Parameters + ---------- + source_hash : str + The hex hash of the canonical C source (same key the in-memory + cache uses). + modname : str + The fully-qualified module name to register with Python's import + machinery (typically ``fn_ptr_ext_``). + constants_manifest : list of (int, UWexpression) + The current call's constants manifest. Used as a belt-and-braces + consistency check: the saved manifest on disk must match, otherwise + we treat it as a cache miss (probable ABI drift). + + Returns + ------- + module or None + The dynamically-loaded extension module on hit, ``None`` on miss + (no ``.so`` on disk, malformed manifest, or manifest mismatch). + """ + cache_dir = get_cache_dir() + if cache_dir is None: + return None + so_path = _so_path(cache_dir, source_hash) + manifest_path = _manifest_path(cache_dir, source_hash) + if not so_path.exists() or not manifest_path.exists(): + return None + try: + with manifest_path.open("r") as f: + manifest = json.load(f) + except (OSError, json.JSONDecodeError): + return None + if manifest.get("version") != MANIFEST_VERSION: + return None + saved = manifest.get("constants", []) + current = _manifest_from_constants(constants_manifest) + if saved != current: + # Names shifted between the saved entry and the current call. + # Could be a UWexpression rename or unnoticed ABI change — safer + # to recompile than trust a stale mapping. + return None + try: + loader = importlib.machinery.ExtensionFileLoader(modname, str(so_path)) + spec = importlib.machinery.ModuleSpec( + name=modname, loader=loader, origin=str(so_path) + ) + return _load(spec) + except Exception: + return None + + +def store_module(source_hash: str, modname: str, tmpdir, constants_manifest) -> None: + """Copy a freshly-compiled ``.so`` into the cache dir and write its manifest. + + Parameters + ---------- + source_hash, modname + See :func:`load_module`. + tmpdir : str or Path + The temporary build directory returned by :func:`compile_and_load` + (contains the ``.so`` we just built). + constants_manifest : list of (int, UWexpression) + Persisted as ``{index, name}`` pairs; used by :func:`load_module` + to verify compatibility on subsequent loads. + """ + import underworld3 + + if underworld3.mpi.rank != 0: + # Phase 4 will replace this guard with proper MPI locking. + return + cache_dir = get_cache_dir() + if cache_dir is None: + return + try: + cache_dir.mkdir(parents=True, exist_ok=True) + except OSError: + return + + tmp = Path(tmpdir) + so_files = list(tmp.glob("*.so")) + if not so_files: + return + src_so = so_files[0] + dst_so = _so_path(cache_dir, source_hash) + + 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), + } + manifest_path = _manifest_path(cache_dir, source_hash) + 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 diff --git a/src/underworld3/utilities/_jitextension.py b/src/underworld3/utilities/_jitextension.py index 4fbf2eb1d..5668f86c6 100644 --- a/src/underworld3/utilities/_jitextension.py +++ b/src/underworld3/utilities/_jitextension.py @@ -597,14 +597,30 @@ def getext( ] # ── Cache lookup or build ───────────────────────────────────────────── + # Three tiers, cheapest first: + # (1) in-memory dict — same Python process + # (2) on-disk cache — same machine, different process + # (3) cold compile via Cython + cc if cache and source_hash in _ext_dict: if verbose and underworld3.mpi.rank == 0: - print(f"JIT compiled module cached ... {source_hash}", flush=True) + print(f"JIT compiled module cached (memory) ... {source_hash}", flush=True) module = _ext_dict[source_hash] else: - if verbose and underworld3.mpi.rank == 0: - print(f"JIT compiling new module ... {source_hash}", flush=True) - module, _tmpdir = compile_and_load(real_modname, codeguys_final, verbose=verbose) + from underworld3.utilities import _jit_cache as _jc + + module = None + if cache: + module = _jc.load_module(source_hash, real_modname, constants_manifest) + if module is not None and verbose and underworld3.mpi.rank == 0: + print(f"JIT compiled module cached (disk) ... {source_hash}", flush=True) + + if module is None: + if verbose and underworld3.mpi.rank == 0: + print(f"JIT compiling new module ... {source_hash}", flush=True) + module, tmpdir = compile_and_load(real_modname, codeguys_final, verbose=verbose) + if cache: + _jc.store_module(source_hash, real_modname, tmpdir, constants_manifest) + if cache: _ext_dict[source_hash] = module diff --git a/tests/test_jit_cache.py b/tests/test_jit_cache.py index 461c40b31..20f1b68f8 100644 --- a/tests/test_jit_cache.py +++ b/tests/test_jit_cache.py @@ -136,6 +136,85 @@ def test_value_cycle_produces_single_cache_entry(): del poisson +def test_disk_cache_roundtrip_skips_compile(tmp_path, monkeypatch): + """First solve compiles + writes to disk; clearing in-memory cache and + re-solving must reuse the on-disk artefact instead of recompiling. + """ + from underworld3.utilities import _jit_cache as _jc + from underworld3.utilities import _jitextension as _jitext + + monkeypatch.setenv("UW_JIT_CACHE_DIR", str(tmp_path)) + + # Earlier tests in this session may have populated _ext_dict with the same + # source_hash (Poisson-with-constant-K all canonicalise to similar C + # source). Force a cold path so the disk-store branch is exercised. + _jitext._ext_dict.clear() + + mesh = uw.meshing.UnstructuredSimplexBox(cellSize=0.25) + u = uw.discretisation.MeshVariable("u_disk", mesh, 1, degree=2) + K = uw.expression("K_disk", 1.0) + + poisson = uw.systems.Poisson(mesh, u_Field=u) + poisson.constitutive_model = uw.constitutive_models.DiffusionModel + poisson.constitutive_model.Parameters.diffusivity = K + poisson.f = -2.0 + poisson.add_dirichlet_bc(0.0, "Bottom") + poisson.add_dirichlet_bc(1.0, "Top") + + # First solve: cold cache, must compile and populate disk + poisson.solve() + source_hash = poisson._current_jit_cache_key + + cache_dir = _jc.get_cache_dir() + so_path = cache_dir / f"{source_hash}.so" + manifest_path = cache_dir / f"{source_hash}.manifest.json" + assert so_path.exists(), f"disk cache .so missing after solve: {so_path}" + assert manifest_path.exists(), f"disk cache manifest missing: {manifest_path}" + + # Clear the in-memory cache; the second solve must hit the disk cache + _jitext._ext_dict.clear() + + compile_calls = [] + real_compile = _jitext.compile_and_load + + def watching_compile(*args, **kwargs): + compile_calls.append(args) + return real_compile(*args, **kwargs) + + monkeypatch.setattr(_jitext, "compile_and_load", watching_compile) + + # Reset solver state so it has to re-resolve the cache, not reuse an + # already-installed module pointer + poisson.is_setup = False + poisson.solve() + + assert len(compile_calls) == 0, ( + f"compile_and_load was called {len(compile_calls)} time(s) on a disk " + f"cache hit — disk cache lookup is not short-circuiting compile." + ) + + del poisson + + +def test_disk_cache_disabled_by_env_var(tmp_path, monkeypatch): + """UW_JIT_CACHE=0 disables both load and store on the on-disk cache.""" + from underworld3.utilities import _jit_cache as _jc + + monkeypatch.setenv("UW_JIT_CACHE_DIR", str(tmp_path)) + monkeypatch.setenv("UW_JIT_CACHE", "0") + + assert _jc.get_cache_dir() is None + + # store_module should be a no-op; load_module should always return None + class _FakeExpr: + name = "X" + + assert _jc.load_module("deadbeef", "fn_ptr_ext_x", [(0, _FakeExpr())]) is None + _jc.store_module("deadbeef", "fn_ptr_ext_x", str(tmp_path), [(0, _FakeExpr())]) + assert not (tmp_path / "deadbeef.so").exists() + assert not (tmp_path / "deadbeef.manifest.json").exists() + + def test_cache_key_is_deterministic_hex_string(): """The cache_key on _GextResult is a 16-char hex string (source hash).""" From e6ec638f1c7a954afe49b007124bfa70c60fae05 Mon Sep 17 00:00:00 2001 From: lmoresi Date: Wed, 22 Apr 2026 10:10:10 +1000 Subject: [PATCH 4/7] JIT cache (phase 4/5): inter-process flock + MPI determinism + ./uw env fingerprint MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) --- src/underworld3/utilities/_jit_cache.py | 102 +++++++++++++++------ src/underworld3/utilities/_jitextension.py | 15 +++ tests/test_jit_cache.py | 60 ++++++++++++ uw | 33 +++++++ 4 files changed, 180 insertions(+), 30 deletions(-) diff --git a/src/underworld3/utilities/_jit_cache.py b/src/underworld3/utilities/_jit_cache.py index 3bf47d6b6..fbda88b62 100644 --- a/src/underworld3/utilities/_jit_cache.py +++ b/src/underworld3/utilities/_jit_cache.py @@ -21,6 +21,8 @@ from __future__ import annotations +import contextlib +import fcntl import importlib.machinery import json import os @@ -55,6 +57,37 @@ def _manifest_path(cache_dir: Path, source_hash: str) -> Path: return cache_dir / f"{source_hash}.manifest.json" +def _lock_path(cache_dir: Path, source_hash: str) -> Path: + return cache_dir / f".{source_hash}.lock" + + +@contextlib.contextmanager +def _file_lock(path: Path): + """Per-hash advisory lock so two processes don't race on the same entry. + + Uses :func:`fcntl.flock` (POSIX). The lockfile is created if missing + and kept around — wiping the cache directory removes it cleanly. + Best-effort: if flock isn't supported (rare on POSIX), the body still + runs; correctness then relies on the atomic ``os.replace`` in + :func:`store_module`. + """ + fd = None + try: + fd = os.open(str(path), os.O_RDWR | os.O_CREAT, 0o600) + try: + fcntl.flock(fd, fcntl.LOCK_EX) + except OSError: + pass + yield + finally: + if fd is not None: + try: + fcntl.flock(fd, fcntl.LOCK_UN) + except OSError: + pass + os.close(fd) + + def _manifest_from_constants(constants_manifest): """Canonicalise ``constants_manifest`` for JSON storage / comparison. @@ -151,35 +184,44 @@ def store_module(source_hash: str, modname: str, tmpdir, constants_manifest) -> return src_so = so_files[0] dst_so = _so_path(cache_dir, source_hash) - - 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), - } manifest_path = _manifest_path(cache_dir, source_hash) - 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) + lock_path = _lock_path(cache_dir, source_hash) + + # flock prevents a concurrent process (different mpirun, etc.) from + # interleaving its write with ours. The recheck-inside-lock pattern + # avoids redundant copies when another process has already populated + # the entry while we were waiting on the lock. + with _file_lock(lock_path): + if dst_so.exists() and manifest_path.exists(): + return + + try: + tf = tempfile.NamedTemporaryFile( + dir=str(cache_dir), prefix=".inprogress_", suffix=".so", delete=False + ) tf_path = Path(tf.name) - os.replace(tf_path, manifest_path) - except OSError: - return + 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 diff --git a/src/underworld3/utilities/_jitextension.py b/src/underworld3/utilities/_jitextension.py index 5668f86c6..0dadb260a 100644 --- a/src/underworld3/utilities/_jitextension.py +++ b/src/underworld3/utilities/_jitextension.py @@ -584,6 +584,21 @@ def getext( (canonical_source + "\n---\n" + _abi_salt()).encode("utf-8") ).hexdigest()[:16] + # Determinism check: all ranks must agree on the hash. A mismatch means + # generate_c_source isn't deterministic across ranks (typically caused + # by set/dict-iteration order leaking into the emitted C). Caching + # cannot work correctly if ranks disagree, so fail loudly rather than + # let stale entries propagate. + if underworld3.mpi.size > 1: + all_hashes = underworld3.mpi.comm.allgather(source_hash) + if any(h != source_hash for h in all_hashes): + raise RuntimeError( + f"JIT C-source hash differs across MPI ranks: {set(all_hashes)}. " + f"This indicates non-determinism in generate_c_source — likely " + f"a set or dict whose iteration order leaks into the C output. " + f"Treating this as a hard error since cache reuse would be unsound." + ) + # Derive the real modname/randstr from the hash — same source ⇒ same # compiled artefact, different sources ⇒ disjoint symbol namespaces. real_modname = f"fn_ptr_ext_{source_hash}" diff --git a/tests/test_jit_cache.py b/tests/test_jit_cache.py index 20f1b68f8..2b8946149 100644 --- a/tests/test_jit_cache.py +++ b/tests/test_jit_cache.py @@ -215,6 +215,66 @@ class _FakeExpr: assert not (tmp_path / "deadbeef.manifest.json").exists() +def test_abi_salt_change_invalidates_cache(tmp_path, monkeypatch): + """Mutating the ABI salt must produce a different cache key and miss.""" + from underworld3.utilities import _jitextension as _jitext + + monkeypatch.setenv("UW_JIT_CACHE_DIR", str(tmp_path)) + _jitext._ext_dict.clear() + + # Monkey-patch the salt to a known value, solve, capture key + monkeypatch.setattr(_jitext, "_abi_salt", lambda: "salt-A") + + mesh = uw.meshing.UnstructuredSimplexBox(cellSize=0.3) + u = uw.discretisation.MeshVariable("u_abi", mesh, 1, degree=2) + K = uw.expression("K_abi", 1.0) + + poisson = uw.systems.Poisson(mesh, u_Field=u) + poisson.constitutive_model = uw.constitutive_models.DiffusionModel + poisson.constitutive_model.Parameters.diffusivity = K + poisson.f = 0.0 + poisson.add_dirichlet_bc(0.0, "Top") + poisson.add_dirichlet_bc(1.0, "Bottom") + poisson.solve() + key_A = poisson._current_jit_cache_key + + # Change the salt, force a fresh solver, expect a different key + _jitext._ext_dict.clear() + monkeypatch.setattr(_jitext, "_abi_salt", lambda: "salt-B") + poisson.is_setup = False + poisson.solve() + key_B = poisson._current_jit_cache_key + + assert key_A != key_B, ( + f"Different ABI salts must produce different cache keys; both got {key_A}" + ) + + del poisson + + +def test_store_module_creates_lockfile(tmp_path, monkeypatch): + """store_module must create the per-hash lockfile next to the entry.""" + from underworld3.utilities import _jit_cache as _jc + + monkeypatch.setenv("UW_JIT_CACHE_DIR", str(tmp_path)) + + # Build a fake "compiled" tmpdir with a placeholder .so so store_module + # doesn't need to invoke cc. + fake_tmp = tmp_path / "fake_build" + fake_tmp.mkdir() + fake_so = fake_tmp / "fn_ptr_ext_xyz.so" + fake_so.write_bytes(b"not a real .so but enough for the copy step") + + class _FakeExpr: + name = "K_lock" + + _jc.store_module("abc1234567abcdef", "fn_ptr_ext_xyz", str(fake_tmp), [(0, _FakeExpr())]) + + # Lockfile uses the leading-dot pattern documented in _lock_path + lock = tmp_path / ".abc1234567abcdef.lock" + assert lock.exists(), f"lockfile not created: {lock}" + + def test_cache_key_is_deterministic_hex_string(): """The cache_key on _GextResult is a 16-char hex string (source hash).""" diff --git a/uw b/uw index eebb537fe..f5610219e 100755 --- a/uw +++ b/uw @@ -272,6 +272,39 @@ run_build() { fi fi + # JIT cache env-fingerprint check. + # + # The on-disk JIT cache (~/.cache/underworld3/jit/) holds compiled .so + # files keyed by C-source hash + (PETSc, UW3) version salt. Anything + # outside that salt — compiler identity, CFLAGS, Python ABI, pixi env — + # could silently invalidate cached binaries (e.g. switching pixi env + # from default to amr-dev links against different libc++ shims). + # + # We detect such changes by writing an env fingerprint to the cache dir + # and wiping the cache on mismatch. + 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" + echo "" echo -e "${GREEN}${BOLD}Build complete!${NC}" echo "" From 3cd0a14dc5877b7200e3356768805fc56e14c583 Mon Sep 17 00:00:00 2001 From: lmoresi Date: Wed, 22 Apr 2026 10:14:08 +1000 Subject: [PATCH 5/7] JIT cache (phase 5/5): docs + benchmark MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) --- .../benchmarks/jit_cache_vs_recompile.py | 87 +++++++++++ docs/developer/index.md | 1 + docs/developer/subsystems/jit-cache.md | 136 ++++++++++++++++++ 3 files changed, 224 insertions(+) create mode 100644 docs/advanced/benchmarks/jit_cache_vs_recompile.py create mode 100644 docs/developer/subsystems/jit-cache.md diff --git a/docs/advanced/benchmarks/jit_cache_vs_recompile.py b/docs/advanced/benchmarks/jit_cache_vs_recompile.py new file mode 100644 index 000000000..e98ff3e38 --- /dev/null +++ b/docs/advanced/benchmarks/jit_cache_vs_recompile.py @@ -0,0 +1,87 @@ +"""Benchmark for issue #123: parameter-update recompile. + +Reproduces the VE square-wave toggle pattern from +https://github.com/underworldcode/underworld3/issues/123 — the original +report measured ~459 s of JIT time vs ~3.7 s of actual SNES solve over +99 timesteps that only changed dt_elastic and the top BC sign. + +With the C-source-hash JIT cache in place, the JIT time should collapse +to roughly the cost of the first cold compile (~10 s on this hardware), +because every later step generates the same C source and hits the cache. + +Run with: + + pixi run -e amr-dev python -u docs/advanced/benchmarks/jit_cache_vs_recompile.py + +Optional: set ``UW_JIT_CACHE=0`` to disable the disk cache and time the +in-memory cache only; or ``rm -rf ~/.cache/underworld3/jit`` first to +force a full cold start. +""" + +import time + +import sympy + +import underworld3 as uw +import underworld3.timing as uw_timing +from underworld3.function import expression +from underworld3.utilities._jitextension import _ext_dict + + +N_STEPS = 30 +ELEMENT_RES = (16, 8) + + +def main(): + mesh = uw.meshing.StructuredQuadBox( + elementRes=ELEMENT_RES, + minCoords=(0.0, 0.0), + maxCoords=(1.0, 0.5), + ) + + v = uw.discretisation.MeshVariable("V", mesh, mesh.dim, degree=2) + p = uw.discretisation.MeshVariable("P", mesh, 1, degree=1) + + stokes = uw.systems.VE_Stokes( + mesh, velocityField=v, pressureField=p, order=2 + ) + stokes.constitutive_model = uw.constitutive_models.ViscoElasticPlasticFlowModel + stokes.constitutive_model.Parameters.shear_viscosity_0 = 1.0 + stokes.constitutive_model.Parameters.shear_modulus = 1.0 + + V_top = expression("V_top", 0.5, "Top BC amplitude") + stokes.add_dirichlet_bc((V_top, 0.0), "Top") + stokes.add_dirichlet_bc((0.0, 0.0), "Bottom") + stokes.add_dirichlet_bc((sympy.oo, 0.0), "Left") + stokes.add_dirichlet_bc((sympy.oo, 0.0), "Right") + + uw_timing.start() + t_total = time.time() + + n_compiles_initial = len(_ext_dict) + step_times = [] + for step in range(N_STEPS): + V_top.sym = sympy.Float(0.5 * (-1) ** step) + stokes.constitutive_model.Parameters.dt_elastic = 0.02 + t_step = time.time() + stokes.solve(zero_init_guess=False, timestep=0.02) + step_times.append(time.time() - t_step) + + elapsed = time.time() - t_total + n_compiles_after = len(_ext_dict) + + uw.pprint(f"\n=== JIT cache benchmark — issue #123 reproducer ===") + uw.pprint(f"Steps : {N_STEPS}") + uw.pprint(f"Mesh : {ELEMENT_RES}, BDF-2 VE Stokes") + uw.pprint(f"Wall time (total) : {elapsed:.1f} s") + uw.pprint(f"First step (cold) : {step_times[0]:.2f} s") + uw.pprint(f"Mean of remaining : {sum(step_times[1:]) / (N_STEPS - 1):.2f} s/step") + uw.pprint(f"Compiled bundles : {n_compiles_after - n_compiles_initial}") + uw.pprint(f" (issue #123: would have shown ~3 new compiles per step " + f"= ~{N_STEPS * 3} total before this fix)") + + uw_timing.print_summary() + + +if __name__ == "__main__": + main() diff --git a/docs/developer/index.md b/docs/developer/index.md index 4a85c31d5..c53dcfc8d 100644 --- a/docs/developer/index.md +++ b/docs/developer/index.md @@ -165,6 +165,7 @@ subsystems/expressions-functions subsystems/containers subsystems/checkpointing-system subsystems/model-orchestration +subsystems/jit-cache ``` ```{toctree} diff --git a/docs/developer/subsystems/jit-cache.md b/docs/developer/subsystems/jit-cache.md new file mode 100644 index 000000000..795b9c433 --- /dev/null +++ b/docs/developer/subsystems/jit-cache.md @@ -0,0 +1,136 @@ +# JIT compilation cache + +Solvers in underworld3 compile their pointwise residual / Jacobian functions +to native code via Cython + the system `cc`. For a non-trivial solver this +takes seconds (simple Stokes) to minutes (VE Stokes BDF-2 with viscoplastic +flow). The JIT cache eliminates that cost on every call after the first one +within a session, and on every fresh Python process after the first one on a +given machine. + +This document describes the on-disk layout, how cache entries are keyed, and +how the cache is invalidated. + +## Layout + +``` +~/.cache/underworld3/jit/ +├── .env-fingerprint # written by `./uw build` +├── 7e81af3b554b4126.so # the compiled extension +├── 7e81af3b554b4126.manifest.json # constants index → name mapping +├── .7e81af3b554b4126.lock # advisory lock used during writes +├── 0202220e00e01df3.so +├── 0202220e00e01df3.manifest.json +└── ... +``` + +Locations honoured (first match wins): + +1. `$UW_JIT_CACHE_DIR` if set +2. `$XDG_CACHE_HOME/underworld3/jit` if `XDG_CACHE_HOME` is set +3. `~/.cache/underworld3/jit` + +Setting `UW_JIT_CACHE=0` (or `false` / `no`) disables the on-disk cache +entirely — the in-memory dict still works, but nothing is persisted across +processes. + +## Cache key + +The key is the SHA-256 of the **canonical** generated C source plus an +**ABI salt**, truncated to 16 hex characters. + +```python +canonical_source = "\n".join(setup_py, cy_ext_h, cy_ext_pyx) + .replace(, "__UW_JIT_MOD__") + .replace(, "__UW_JIT_RS__") + +salt = f"petsc={PETSc.Sys.getVersion()}|uw={underworld3.__version__}" + +source_hash = sha256(canonical_source + "\n---\n" + salt).hexdigest()[:16] +``` + +Two design points worth understanding: + +**Canonicalisation**. The Cython module name and the symbol-prefix randomiser +that the JIT generator picks vary between calls (random or counter-based). +Hashing them directly would make every call produce a fresh key. We instead +hash the source with both replaced by stable placeholders, and only after +hashing do we substitute the real per-bundle names back in. Same C ⇒ same +hash; different C ⇒ different hash; identifiers stay unique per bundle. + +**Why hash the C source rather than the sympy callback structure**. Two +sympy expressions can be structurally distinct but emit identical C +(constant placeholders, simplification, etc.); two structurally-identical +expressions can emit different C if a path in the generator depends on +something subtle. The compiled `.so` is what the solver actually executes — +hashing exactly that makes the equivalence relation provable: same hash +implies same `.so` implies same numerical behaviour. + +## Cache hit/miss flow + +`getext()` performs a three-tier lookup, cheapest first: + +1. **In-memory `_ext_dict`** — same Python process, just a dict access. +2. **On-disk `{hash}.so` + `{hash}.manifest.json`** — same machine, fresh + process. `load_module` re-loads the `.so` via Python's import machinery + and verifies the saved constants list still matches the current call's + manifest (belt-and-braces: a name shift means an unnoticed ABI drift). +3. **Cold compile** via `compile_and_load`. After a successful compile, the + `.so` and manifest are copied into the cache directory (rank 0 only), + guarded by an advisory `flock`. + +## Invalidation + +A cached entry stops being valid when *anything* it implicitly depended on +changes. We handle the common categories like this: + +| Change | Handled by | +|-------------------------------------|--------------------------------------------------| +| Solver expressions / topology | `source_hash` differs ⇒ different entry | +| Constant **value** | `source_hash` is independent of value (uses `_JITConstant` placeholders); `PetscDSSetConstants` updates the array at solve time | +| PETSc version, UW version | Embedded in the ABI salt ⇒ different `source_hash` | +| Compiler / CFLAGS / Python ABI | `./uw build` writes `.env-fingerprint`; mismatch ⇒ wipe `*.so` and `*.manifest.json` | +| Constant rename | `load_module` compares saved vs current names ⇒ miss | +| Manifest schema change | `MANIFEST_VERSION` bump ⇒ miss | + +The cache directory is safe to delete by hand at any time: + +``` +rm -rf ~/.cache/underworld3/jit +``` + +The next solve will repopulate. + +## MPI + +When `mpi.size > 1`: + +- Every rank computes the C-source hash independently. The hashes are + `comm.allgather`'d and compared — a mismatch raises immediately rather + than letting ranks diverge. Non-determinism in `generate_c_source` (e.g. + set/dict iteration order leaking into emitted C) would land here. +- Rank 0 writes the cache entry; other ranks rely on the disk-cache hit + path on subsequent calls. +- The `flock` on the per-hash lockfile serialises cross-shell concurrent + writes (e.g. a `mpirun -np 4` and a `mpirun -np 2` started seconds apart + on the same machine). + +A future refinement would have rank 0 compile while other ranks wait on a +barrier and read the resulting `.so` directly; today every rank still +performs the cold compile but only rank 0 publishes the result. + +## Environment variables + +| Variable | Effect | +|----------------------|---------------------------------------------------------| +| `UW_JIT_CACHE` | Set to `0`/`false`/`no` to disable disk cache | +| `UW_JIT_CACHE_DIR` | Override the cache directory location | +| `XDG_CACHE_HOME` | Used when `UW_JIT_CACHE_DIR` is unset | + +## Code references + +- `src/underworld3/utilities/_jitextension.py` — `getext`, `generate_c_source`, + `compile_and_load`, `_abi_salt`, `_extract_constants`. +- `src/underworld3/utilities/_jit_cache.py` — disk cache: `get_cache_dir`, + `load_module`, `store_module`, `_file_lock`. +- `uw` (shell driver) — `.env-fingerprint` write + cache wipe inside + `run_build`. From 17e0a8947263d1b3cfdd291617e94f12f1d02903 Mon Sep 17 00:00:00 2001 From: lmoresi Date: Wed, 22 Apr 2026 11:28:38 +1000 Subject: [PATCH 6/7] JIT cache (parallel): rank-0 compiles + cross-rank disk handoff MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) --- src/underworld3/utilities/_jitextension.py | 59 ++- tests/parallel/ptest_jit_cache.py | 406 +++++++++++++++++++++ 2 files changed, 459 insertions(+), 6 deletions(-) create mode 100644 tests/parallel/ptest_jit_cache.py diff --git a/src/underworld3/utilities/_jitextension.py b/src/underworld3/utilities/_jitextension.py index 0dadb260a..4be38b8b1 100644 --- a/src/underworld3/utilities/_jitextension.py +++ b/src/underworld3/utilities/_jitextension.py @@ -615,7 +615,7 @@ def getext( # Three tiers, cheapest first: # (1) in-memory dict — same Python process # (2) on-disk cache — same machine, different process - # (3) cold compile via Cython + cc + # (3) cold compile via Cython + cc (rank 0 only when MPI > 1) if cache and source_hash in _ext_dict: if verbose and underworld3.mpi.rank == 0: print(f"JIT compiled module cached (memory) ... {source_hash}", flush=True) @@ -623,6 +623,8 @@ def getext( else: from underworld3.utilities import _jit_cache as _jc + # Disk lookup is cheap and rank-local — every rank checks + # independently. UW assumes a shared filesystem for the cache dir. module = None if cache: module = _jc.load_module(source_hash, real_modname, constants_manifest) @@ -630,11 +632,56 @@ def getext( print(f"JIT compiled module cached (disk) ... {source_hash}", flush=True) if module is None: - if verbose and underworld3.mpi.rank == 0: - print(f"JIT compiling new module ... {source_hash}", flush=True) - module, tmpdir = compile_and_load(real_modname, codeguys_final, verbose=verbose) - if cache: - _jc.store_module(source_hash, real_modname, tmpdir, constants_manifest) + # Cold compile path. With MPI, only rank 0 invokes cc and + # publishes; the other ranks barrier-wait and then load the + # freshly-published .so from disk. This avoids the N× cc + # storm on cold-start under mpirun -np N. + # + # Falls back to "every rank compiles" only when the disk + # cache is disabled (no way to share a build), in which + # case the wasted cc cost is on the user's chosen path. + multi_rank = underworld3.mpi.size > 1 + disk_enabled = cache and _jc.get_cache_dir() is not None + + if multi_rank and disk_enabled: + if underworld3.mpi.rank == 0: + if verbose: + print( + f"JIT compiling new module on rank 0 ... {source_hash}", + flush=True, + ) + module, tmpdir = compile_and_load( + real_modname, codeguys_final, verbose=verbose + ) + _jc.store_module( + source_hash, real_modname, tmpdir, constants_manifest + ) + # Synchronisation point: rank 0 has now published the + # entry; other ranks may load it. + underworld3.mpi.comm.Barrier() + if underworld3.mpi.rank != 0: + module = _jc.load_module( + source_hash, real_modname, constants_manifest + ) + if module is None: + # Defensive: rank-0 publish failed (e.g. write + # error). Fall back to a local compile so this + # rank is at least correct, even if wasteful. + module, _tmpdir = compile_and_load( + real_modname, codeguys_final, verbose=verbose + ) + else: + if verbose and underworld3.mpi.rank == 0: + print( + f"JIT compiling new module ... {source_hash}", flush=True + ) + module, tmpdir = compile_and_load( + real_modname, codeguys_final, verbose=verbose + ) + if cache: + _jc.store_module( + source_hash, real_modname, tmpdir, constants_manifest + ) if cache: _ext_dict[source_hash] = module diff --git a/tests/parallel/ptest_jit_cache.py b/tests/parallel/ptest_jit_cache.py new file mode 100644 index 000000000..0edee2195 --- /dev/null +++ b/tests/parallel/ptest_jit_cache.py @@ -0,0 +1,406 @@ +"""Parallel validation for the JIT cache (issues #121, #123). + +Run via: + + mpirun -np 4 python tests/parallel/ptest_jit_cache.py + +Exits non-zero on the first failed assertion (with a printed banner so +``mpi_runner.sh``-style test loops surface the failure in CI logs). + +Validates four things that aren't checkable from the serial pytest suite: + +1. **Hash agreement.** The C-source canonicalisation must be byte-identical + across ranks. ``getext`` already calls ``allgather(source_hash)`` and + raises on mismatch — this test asserts the hash matches what every rank + computes locally too, and that it doesn't change between repeated calls. + +2. **C ↔ S ↔ C transition under np>1.** The same correctness gate as the + serial test, but every rank must land on the same numerical answer. + +3. **Concurrent compile race.** All ranks compile their first bundle at + the same time. They each write to a unique ``/tmp/fn_ptr_ext_*__*`` + directory and then race to populate the on-disk cache; only rank 0 + should win the publish (the others' writes are no-ops by construction). + We assert the cache directory ends up with exactly one ``.so`` per + distinct bundle. + +4. **Inter-process flock.** A second mpirun launched while the first is + still compiling could otherwise stomp on the same ``{hash}.so``. We + exercise this by having every rank call ``store_module`` on the same + fake artifact — only one ``.so`` should result, and the file size + must equal what we wrote. +""" + +from __future__ import annotations + +import os +import shutil +import sys +import tempfile +import time +from pathlib import Path + +import numpy as np +import sympy + +import underworld3 as uw +from underworld3.utilities import _jit_cache as _jc +from underworld3.utilities import _jitextension as _jitext + + +COMM = uw.mpi.comm +RANK = uw.mpi.rank +SIZE = uw.mpi.size + + +def _banner(msg): + if RANK == 0: + print(f"\n=== {msg} ===", flush=True) + + +def _fail(msg): + print(f"[rank {RANK}] FAIL: {msg}", flush=True) + COMM.Abort(1) + + +def _ranks_agree(value, label): + gathered = COMM.allgather(value) + if any(v != value for v in gathered): + _fail(f"{label}: ranks disagree — {gathered}") + + +def _setup_isolated_cache(): + """Pick a per-rank-shared cache dir under /tmp so we don't clobber the user's.""" + base = Path(tempfile.gettempdir()) / "uw_jit_cache_ptest" + if RANK == 0: + if base.exists(): + shutil.rmtree(base, ignore_errors=True) + base.mkdir(parents=True, exist_ok=True) + COMM.Barrier() + os.environ["UW_JIT_CACHE_DIR"] = str(base) + return base + + +def test_1_hash_agreement_and_csc_correctness(): + """C↔S↔C transition under MPI: hash agreement + numerical equality.""" + _banner(f"Test 1: hash agreement + C↔S↔C correctness (np={SIZE})") + + _jitext._ext_dict.clear() + + mesh = uw.meshing.UnstructuredSimplexBox(cellSize=0.2) + u = uw.discretisation.MeshVariable("u_par_csc", mesh, 1, degree=2) + K = uw.expression("K_par_csc", 0.5) + + poisson = uw.systems.Poisson(mesh, u_Field=u) + poisson.constitutive_model = uw.constitutive_models.DiffusionModel + poisson.constitutive_model.Parameters.diffusivity = K + poisson.f = 1.0 + poisson.add_dirichlet_bc(0.0, "Top") + poisson.add_dirichlet_bc(0.0, "Bottom") + + sample_y = np.linspace(0.05, 0.95, 15) + pts = np.column_stack([np.full_like(sample_y, 0.5), sample_y]) + + def t_max(): + # On non-rank-0 evaluate may return NaN at points off-rank; use + # a global reduction to avoid spurious disagreement. + local_arr = uw.function.evaluate(u.sym[0], pts, rbf=False).squeeze() + local_max = float(np.nanmax(np.abs(local_arr))) if np.any(~np.isnan(local_arr)) else 0.0 + return COMM.allreduce(local_max, op=uw.mpi._MPI.MAX) + + # State 1 + poisson.solve(zero_init_guess=True) + key_1 = poisson._current_jit_cache_key + _ranks_agree(key_1, "state-1 cache key") + t_max_1 = t_max() + _ranks_agree(round(t_max_1, 9), "state-1 T_max") + + # State 2 — symbolic + poisson.constitutive_model.Parameters.diffusivity = 1.0 + u.sym[0] + poisson.solve(zero_init_guess=True) + key_2 = poisson._current_jit_cache_key + _ranks_agree(key_2, "state-2 cache key") + if key_2 == key_1: + _fail(f"state-1 and state-2 hashed to the same key {key_1}; should differ") + t_max_2 = t_max() + + # State 3 — back to constant + poisson.constitutive_model.Parameters.diffusivity = K + poisson.solve(zero_init_guess=True) + key_3 = poisson._current_jit_cache_key + _ranks_agree(key_3, "state-3 cache key") + if key_3 != key_1: + _fail(f"state-3 ({key_3}) should hash to state-1 key ({key_1})") + t_max_3 = t_max() + + if abs(t_max_3 - t_max_1) > 1e-6: + _fail( + f"state-3 T_max {t_max_3:.8f} differs from state-1 {t_max_1:.8f}; " + f"PETSc-DS contamination across the C→S→C transition" + ) + + if RANK == 0: + print( + f" state 1 T_max={t_max_1:.6f} state 2 T_max={t_max_2:.6f} " + f"state 3 T_max={t_max_3:.6f} ✓", + flush=True, + ) + + del poisson + + +def test_2_concurrent_compile_publishes_once(): + """All ranks solve cold; the disk cache must end up with exactly one .so per bundle.""" + _banner(f"Test 2: concurrent cold compile, single publish (np={SIZE})") + + cache_dir = _jc.get_cache_dir() + if cache_dir is None: + _fail("disk cache should not be disabled in this test") + + # Wipe and synchronise so every rank really does start cold + if RANK == 0: + for p in cache_dir.iterdir(): + try: + p.unlink() + except OSError: + pass + COMM.Barrier() + _jitext._ext_dict.clear() + + mesh = uw.meshing.UnstructuredSimplexBox(cellSize=0.25) + u = uw.discretisation.MeshVariable("u_par_concur", mesh, 1, degree=2) + K = uw.expression("K_par_concur", 1.0) + + poisson = uw.systems.Poisson(mesh, u_Field=u) + poisson.constitutive_model = uw.constitutive_models.DiffusionModel + poisson.constitutive_model.Parameters.diffusivity = K + poisson.f = 0.0 + poisson.add_dirichlet_bc(1.0, "Bottom") + poisson.add_dirichlet_bc(0.0, "Top") + poisson.solve() + + source_hash = poisson._current_jit_cache_key + _ranks_agree(source_hash, "concurrent-compile cache key") + + COMM.Barrier() # let rank 0's store_module finish + if RANK == 0: + so_files = sorted(p.name for p in cache_dir.glob("*.so")) + manifest_files = sorted(p.name for p in cache_dir.glob("*.manifest.json")) + expected_so = f"{source_hash}.so" + expected_manifest = f"{source_hash}.manifest.json" + if so_files != [expected_so]: + _fail(f"expected single .so {expected_so}, got {so_files}") + if manifest_files != [expected_manifest]: + _fail(f"expected single manifest {expected_manifest}, got {manifest_files}") + # Also ensure no leftover ".inprogress_*" files (would mean a publish was + # cut off mid-write, which os.replace + flock are supposed to prevent) + leftover = sorted(p.name for p in cache_dir.glob(".inprogress_*")) + if leftover: + _fail(f"leftover .inprogress_* artifacts: {leftover}") + print(f" cache dir: {so_files} + {manifest_files} ✓", flush=True) + + del poisson + + +def test_3_warm_run_skips_compile(): + """Second pass with a warm in-memory cache should never invoke compile_and_load.""" + _banner(f"Test 3: warm in-memory cache, no compile (np={SIZE})") + + compile_calls = [] + real_compile = _jitext.compile_and_load + + def watching(*args, **kwargs): + compile_calls.append(args) + return real_compile(*args, **kwargs) + + _jitext.compile_and_load = watching + try: + # Build a fresh solver — but with the SAME structure as test 2 so the + # hash matches and _ext_dict has the entry. + mesh = uw.meshing.UnstructuredSimplexBox(cellSize=0.25) + u = uw.discretisation.MeshVariable("u_par_warm", mesh, 1, degree=2) + K = uw.expression("K_par_concur", 1.0) # reuse name for stable hash + + poisson = uw.systems.Poisson(mesh, u_Field=u) + poisson.constitutive_model = uw.constitutive_models.DiffusionModel + poisson.constitutive_model.Parameters.diffusivity = K + poisson.f = 0.0 + poisson.add_dirichlet_bc(1.0, "Bottom") + poisson.add_dirichlet_bc(0.0, "Top") + # Note: this uses a different MeshVariable name (u_par_warm), so the + # hash will likely differ from test 2 — but the disk cache from test + # 2 doesn't help us here. The point of this test is really that + # repeated solves on the SAME solver don't recompile. + poisson.solve() + n_compiles_after_first = len(compile_calls) + + for _ in range(5): + K.sym = float(np.random.uniform(0.5, 1.5)) + poisson.solve() + + n_compiles_after_loop = len(compile_calls) + # First solve compiles; the loop must add zero + if n_compiles_after_loop != n_compiles_after_first: + _fail( + f"warm-cache loop fired compile_and_load " + f"{n_compiles_after_loop - n_compiles_after_first} extra times" + ) + if RANK == 0: + print( + f" rank 0: {n_compiles_after_first} compile(s) on first solve, " + f"{n_compiles_after_loop - n_compiles_after_first} during 5-iter cycle ✓", + flush=True, + ) + finally: + _jitext.compile_and_load = real_compile + + del poisson + + +def test_5_only_rank_0_compiles_on_cold_start(): + """Cold start with np>1: rank 0 invokes cc; other ranks load from disk.""" + _banner(f"Test 5: rank-0-only compile on cold start (np={SIZE})") + + if SIZE == 1: + if RANK == 0: + print(" np=1 → rank-0-only compile is trivially satisfied. SKIP", flush=True) + return + + cache_dir = _jc.get_cache_dir() + if RANK == 0: + for p in cache_dir.iterdir(): + try: + p.unlink() + except OSError: + pass + COMM.Barrier() + _jitext._ext_dict.clear() + + compile_calls = [] + real_compile = _jitext.compile_and_load + + def watching(*args, **kwargs): + compile_calls.append(args) + return real_compile(*args, **kwargs) + + _jitext.compile_and_load = watching + try: + # Use a structurally-distinct solver from the other tests so we + # genuinely cold-start (no in-memory hit from a prior test). + mesh = uw.meshing.UnstructuredSimplexBox(cellSize=0.18) + u = uw.discretisation.MeshVariable("u_par_rank0", mesh, 1, degree=2) + K = uw.expression("K_par_rank0", 0.7) + + poisson = uw.systems.Poisson(mesh, u_Field=u) + poisson.constitutive_model = uw.constitutive_models.DiffusionModel + poisson.constitutive_model.Parameters.diffusivity = K + poisson.f = -1.0 + poisson.add_dirichlet_bc(0.0, "Bottom") + poisson.add_dirichlet_bc(0.0, "Top") + poisson.solve() + finally: + _jitext.compile_and_load = real_compile + + n_calls = len(compile_calls) + all_calls = COMM.allgather(n_calls) + if RANK == 0: + # Rank 0 should have called compile_and_load exactly once. + # Every other rank should have zero calls. + if all_calls[0] != 1: + _fail(f"rank 0 should have one cc invocation, got {all_calls[0]}") + for r in range(1, SIZE): + if all_calls[r] != 0: + _fail( + f"rank {r} invoked cc {all_calls[r]} time(s); should be 0 " + f"(disk cache should have served it after rank-0 publish). " + f"Per-rank counts: {all_calls}" + ) + print(f" per-rank cc invocations: {all_calls} ✓", flush=True) + + del poisson + + +def test_4_flock_serialises_disk_writes(): + """All ranks call store_module on the same fake artifact; only one wins, no corruption.""" + _banner(f"Test 4: flock serialises disk writes (np={SIZE})") + + cache_dir = _jc.get_cache_dir() + bogus_hash = f"flockprobe{RANK:06x}"[:16].ljust(16, "0") + bogus_hash = "flock0123456789a" + + # Each rank has its own fake "compiled" .so so we can detect which one + # ended up in the cache. + fake_dir = Path(tempfile.gettempdir()) / f"uw_jit_fake_rank{RANK}_{os.getpid()}" + fake_dir.mkdir(parents=True, exist_ok=True) + payload = f"rank-{RANK} payload at {time.time_ns()}".encode() + fake_so = fake_dir / "fn_ptr_ext_flockprobe.so" + fake_so.write_bytes(payload) + + class _FakeExpr: + name = "K_flock" + + # Wipe any prior entry + if RANK == 0: + for p in cache_dir.glob(f"{bogus_hash}*"): + try: + p.unlink() + except OSError: + pass + for p in cache_dir.glob(f".{bogus_hash}*"): + try: + p.unlink() + except OSError: + pass + COMM.Barrier() + + # All ranks try to publish at roughly the same time + _jc.store_module(bogus_hash, "fn_ptr_ext_flock", str(fake_dir), [(0, _FakeExpr())]) + COMM.Barrier() + + if RANK == 0: + so_path = cache_dir / f"{bogus_hash}.so" + manifest_path = cache_dir / f"{bogus_hash}.manifest.json" + if not so_path.exists(): + _fail(f"no .so written under flock; expected {so_path}") + if not manifest_path.exists(): + _fail(f"no manifest written under flock; expected {manifest_path}") + leftover = sorted(p.name for p in cache_dir.glob(".inprogress_*")) + if leftover: + _fail(f"leftover .inprogress_* under flock: {leftover}") + # Whichever rank won, the bytes must equal a valid rank's payload — + # not a torn write or the original being overwritten mid-copy. + contents = so_path.read_bytes() + if not contents.startswith(b"rank-"): + _fail(f"corrupt .so under flock; first bytes={contents[:32]!r}") + print(f" surviving .so: {len(contents)} bytes, starts with {contents[:24]!r} ✓", flush=True) + + COMM.Barrier() + if RANK == 0: + # cleanup + for p in cache_dir.glob(f"{bogus_hash}*"): + try: p.unlink() + except OSError: pass + for p in cache_dir.glob(f".{bogus_hash}*"): + try: p.unlink() + except OSError: pass + + +def main(): + cache_dir = _setup_isolated_cache() + if RANK == 0: + print(f"\nRunning JIT cache parallel tests with np={SIZE}", flush=True) + print(f"Isolated cache dir: {cache_dir}", flush=True) + + test_1_hash_agreement_and_csc_correctness() + test_2_concurrent_compile_publishes_once() + test_3_warm_run_skips_compile() + test_4_flock_serialises_disk_writes() + test_5_only_rank_0_compiles_on_cold_start() + + COMM.Barrier() + if RANK == 0: + print(f"\n=== All JIT cache parallel tests PASSED (np={SIZE}) ===", flush=True) + + +if __name__ == "__main__": + main() From 794522d1fe35ab7bd017d6fd6ced33fababd561e Mon Sep 17 00:00:00 2001 From: lmoresi Date: Wed, 22 Apr 2026 15:46:30 +1000 Subject: [PATCH 7/7] JIT cache: restore "Location of compiled module" verbose print 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) --- src/underworld3/utilities/_jitextension.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/src/underworld3/utilities/_jitextension.py b/src/underworld3/utilities/_jitextension.py index 4be38b8b1..ff77e9888 100644 --- a/src/underworld3/utilities/_jitextension.py +++ b/src/underworld3/utilities/_jitextension.py @@ -653,6 +653,8 @@ def getext( module, tmpdir = compile_and_load( real_modname, codeguys_final, verbose=verbose ) + if verbose: + print(f"Location of compiled module: {tmpdir}", flush=True) _jc.store_module( source_hash, real_modname, tmpdir, constants_manifest ) @@ -678,6 +680,10 @@ def getext( module, tmpdir = compile_and_load( real_modname, codeguys_final, verbose=verbose ) + if verbose and underworld3.mpi.rank == 0: + # Tests in test_0004_pointwise_fns parse this exact prefix + # to find the per-call build directory. + print(f"Location of compiled module: {tmpdir}", flush=True) if cache: _jc.store_module( source_hash, real_modname, tmpdir, constants_manifest