Make the numba wrapper functions disk-cacheable (and faster to call)

[velochy]

What / why

This addresses the slow-start problem reported in #115, which was closed at the time with
"I don't think we can do much about the numba compilation times in nutpie" — it turns out we
can, for every process after the first.

The numba-backend logp and expand wrappers close over the inner pytensor dispatcher and over
the shared-variable arrays. Numba refuses to disk-cache functions that reference such "dynamic
globals" (nutpie currently suppresses exactly that warning), so the whole wrapper chain is
retyped, relowered and re-optimised by LLVM in every process that compiles a model — and
because the wrapper inlines everything up to the dispatcher call, that compilation traverses
the entire model graph. For a large hierarchical model this costs minutes of compile time and
a GB-scale arena that numba never releases, in every worker process, on every run, even when
the inner functions themselves come straight from PyTensor's numba disk cache.

Three changes make the chain cacheable while keeping its structure as-is:

1. Call the inner function by address. It is compiled up front for the wrapper's argument
   types; its C-wrapper address is stored in user_data next to the shared-array pointers and
   rebuilt into a first-class function value by a small intrinsic. (Same pattern PyTensor uses
   for its cached LAPACK calls.)
2. Zero-size stand-ins for the prototype tuple. The extraction function's shared_tuple
   is only ever used for its numba types — every slot is overwritten via tuple_setitem
   before use. Zero-size arrays have identical types and are small enough for numba to embed
   as constants, where the real shared arrays are dynamic globals.
3. Positional shared-variable keys (s0000, …) instead of uuid4().hex, so the record
   field names — and with them the cache keys — are identical across processes.

The cfuncs are then built with cache=True by default (cache=False in compile_kwargs
opts out), and the warning suppression is removed: if cacheability regresses, the warning is
the only way to notice.

Results

8.5k-parameter hierarchical model, ~3000 apply nodes (numba backend, same machine, RSS deltas
of the compiling process):

	before	after
first-ever compile	1113 s / +4.2 GB	440 s / +0.7 GB
warm-cache compile	425 s / +2.7 GB	262 s / +0.7 GB
logp+grad call	29.0 ms	21.6 ms
expand call	18.1 ms	13.0 ms

Small model (toy hierarchical logistic): 32 s compile in every process before; 42 s once,
10.5 s in later processes after.

Per-call cost on small standard models: the indirect call adds a fixed ~1-1.5 us to the logp
path, visible only where a gradient evaluation itself is a few microseconds (radon-style
hierarchical regression: 7.9 -> 9.4 us). An LKJ/MvNormal model (22 us evals) is at parity,
and the large model above is 26% faster. So: cheap-eval models pay a small per-call premium,
compute-dense models gain.

The call-time improvement was unexpected but is reproducible: compiling the model function as
its own unit (instead of inlining ~3000 nodes into one giant cfunc) appears to give LLVM a
better optimisation unit, and the shared arguments are now typed readonly.

Draws are bit-identical to the current wrappers (same seeds, checksums compared on both the
toy and the large model). with_data works unchanged — shapes still come from user_data
at call time, nothing shape-specific is compiled in.

Notes for review

• numba.experimental.function_type._get_wrapper_address is a private API (it is the
  mechanism behind first-class function support). If that's a concern, the address could be
  obtained from the dispatcher's compile result instead — happy to change it.
• The first-class call site passes a plain Tuple where the dispatcher's *args signature
  says StarArgTuple; the element types and ABI are identical. The function type is declared
  with the plain form for that reason.
• The inner functions are compiled with readonly array arguments because the prototype-tuple
  elements (constants) are typed readonly; the carray views passed at runtime convert
  safely.
• Cache artifacts land next to compile_pymc.py (__pycache__/*.nbi/.nbc) via numba's
  normal locator, plus PyTensor's own cache for the inner functions. Nothing new to manage.
• The jax backend is untouched.

Test plan

• tests/test_pymc.py::test_wrapper_functions_cacheable — no dynamic globals in either
  wrapper (the property the disk cache depends on), plus a sampling smoke test.
• tests/test_pymc.py::test_wrapper_with_data — with_data correctness against a
  recomputed deterministic.
• Existing suite: test_order_shared and friends pass unchanged.
• Cross-model cache safety, tested out-of-tree: eight-schools (centered and non-centered —
  deliberately identical n_dim, record layout, wrapper source and C signature), a radon-style
  hierarchical model and an LKJ/MvNormal model compiled in one process and re-compiled in a
  fresh process in reverse order against the same caches: all posteriors bit-identical to
  the uncached reference. With equal layouts the wrapper is genuinely model-agnostic (the
  inner function's address arrives via user_data at call time), so cache sharing between
  such models is correct by construction; differing layouts change the record dtype and the
  dim guards reject loudly.

🤖 Generated with Claude Code

[velochy]

Real Margus here: This is AI generated, but I did hold it's hand, forcing it to find a simplest possible solution, testing with the standard model library etc.

I'm not that well versed in nutpie internals (hence the AI use), but as you can see, this has a significant impact on both speed and memory for the larger models I am running, so I am invested in getting this in. And I'm hoping I'm not the only one for whom this would be a real and significant benefit.

[aseyboldt]

This certainly looks very interesting!
I want to have a closer look, not sure I can get this done before stancon, but I'll try. Thanks for looking into it.

[velochy]

Also probably worth explaining the weird use case with PyMC. We are sampling multiple pymc models that depend on eachother. Numba has war-crime level atrocious memory handling when it comes to compile cache, so during compilation it creates gigabytes of meta-info and it keeps it around with no way to free it. (Supposedly not fully their fault, comes with the LLVM compiler they use)

So the way to run memory-efficiently is the following:

• Have a master control thread
• Fork a separate process that does compilation - peaks memory, but frees it all once done. Disk cache quite small.
• Fork a new separate process for pm.sample - the cache gets loaded again, but only the needed parts. Fork needed here so it would not accumulate over the sequential model runs.

Ridiculous, but it works. The only way to actually properly fix it seems deep in the numba internals with an issue from 10 years ago flagging it already so I'm guessing unlikely to ever be handled, so I'm commiting to the workaround.

The approach works without this PR, but the patch here makes the compiled cache use a lot more fast and efficient, so if it makes sense in the nutpie context as well, I'd be very happy to get it included :)

[ricardoV94]

your bot may be over-reading into the warning message. Usually I see cache disabled because we have large constant arrays, not "dynamic globals" (same warning message iirc). I don't know what dynamic globals are, because AFAICT everything but constants are explicit inputs the way we define numba functions. And I don't see what you're doing here to handle constants differently

[ricardoV94]

or your "fix" was exactly that, to make constants explicit inputs? My understanding is a bit vague

[velochy]

Closing this — you were right to push back, and re-testing against a fair baseline shows the premise is wrong.

Two errors in what I posted:

1. The baseline was unfair. nutpie's numba backend never asks for caching (numba.cfunc(c_sig, **kwargs) with no cache key), so my "before" numbers were stock-with-caching-off versus my branch-with-caching-on, and I attributed the whole gap to the restructuring. Simply passing cache=True through makes the stock wrapper cache fine — no warning, cache files written, faster second process. On the 8.5k-parameter model that motivated this:

maxdiff (8.5k params, ~3000 apply nodes)	cold	warm
stock + cache=True	267 s / +1.46 GB	222 s / +1.53 GB
this branch	213 s / +1.48 GB	243 s / +1.49 GB

Run-to-run noise. The restructuring buys nothing.

2. Wrong culprit, exactly as you suspected. The "dynamic globals" I chased were in the PyTensor-generated function, not nutpie's wrapper, and the cause was large constant arrays — make_constant_array in numba/core/base.py bakes arrays by address above size_limit = 10**6, which sets has_dynamic_globals and silently disables that function's disk cache (PyTensor suppresses the refusal warning, so nothing surfaces). In my case they were projection matrices in my own model code; making them pytensor.shared fixed it on my side. nutpie's cfuncs report has_dynamic_globals=False in every configuration I have since tested, including with 1.9 MB of shared data.

To your second question: for the shared arrays the answer is "they already were explicit inputs" — they arrive as pointers in user_data, and only the typing prototype tuple held real arrays. So there was nothing to fix there either.

The only defensible remnant is that cache=True isn't the default, which is worth maybe 45 s on a 220 s compile and no memory — happy to send that one-liner (plus dropping the now-misleading warning suppression) if you think it is worth having, but it does not need this patch. Sorry for the noise, and thanks for the sanity check.

[ricardoV94]

@velochy we've talked about promoting large constants to explicit inputs (pytensor C codegen aready works like that), and it makes sense since numba isn't interning these anyway and I suppose llvmlite is not looking at them either. So no loss, and still cacheability. It's just more backend work in pytensor/NumbaLinker. A quick band-aid is to explicitly promote large constants to shared variables and then compiling that function instead.

[ricardoV94]

The only defensible remnant is that cache=True isn't the default, which is worth maybe 45 s on a 220 s compile and no memory

Which cache=True isn't default? pytensor uses config.numba__cache and that's true by default. nutpie does something on top of it? Just remember caching is not trivial, much less in numba :)

[velochy]

A quick band-aid is to explicitly promote large constants to shared variables and then compiling that function instead.

Exactly what I did, actually, and that created a big win in both warm compile time and memory use that was reported initially. So yes, the potential impact of doing that automagically is pretty significant.

Since pytensor/link/numba/dispatch/basic.py suppresses the Cannot cache compiled function "numba_funcified_fgraph…" warning, this was actually pretty hard to discover. So you might just start by letting that pass with a helpful "solution is to make your constants into shared variables". As by now, most developers reportedly use AI, this alone will likely be the 80% solution :)

Which cache=True isn't default? pytensor uses config.numba__cache and that's true by default. nutpie does something on top of it? Just remember caching is not trivial, much less in numba :)

Not pytensor's — config.numba__cache is on and works; the funcified fgraph does come from disk. I meant nutpie's own wrapper functions: _compile_pymc_model_numba calls numba.cfunc(c_sig, **kwargs) with no cache key, so the extraction fn and the C callback that wrap the pytensor function are recompiled in every process even when the inner function loads from cache. Today you can switch it on per call — pm.sample(..., compile_kwargs={"cache": True}) reaches numba.cfunc(..., cache=True). Saves ~40s on a 260s compile for me - not nothing, but not something I'd push hard for to be included, esp considering I can set it myself via compiler_kwargs

[ricardoV94]

pm.sample(..., compile_kwargs={"cache": True}) I didn't know about this numba kwarg :D

"solution is to make your constants into shared variables".

I don't know how actionable that is for the average person who'd see the warning. I leaned negative E[utility] and suppressed.

[ricardoV94]

The fact you wasted a lot of time on misreading the message confirms my guess, although I didn't think of changing/expanding the message.