From ae91f105582e6c3be3b47ba633fb9f704aa87b18 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Mon, 3 Aug 2026 14:53:29 -0500 Subject: [PATCH] Persist the numba compile of the built-in UDF binops to disk floordiv, rfloordiv, absfirst, abssecond, and rpow are UDFs, so every process that touched one paid its numba compile again. They are plain module-level functions, which is exactly the case numba can cache, so that cost belongs once per machine rather than once per process. Thread the flag through as a private keyword-only `_cache` on register_new into _build, where it reaches numba.njit alongside the existing error_model. Keeping it private avoids widening the public API for something no caller outside this module should set. Ops registered by users stay uncached: numba keys a cache entry on a stable on-disk source, which lambdas and interactively defined functions do not have. Measured here, fresh process, medians of 3, load average ~4 on a shared machine so read these as indicative. First-touch floordiv alone: 434 ms uncached, 347 ms warm (-20%). First-touch of all five: 1337 ms uncached, 733 ms warm (-45%). Writing a cold cache costs about 107 ms once for all five (1444 ms on the run that populates it). What stays: the remaining ~733 ms is numba/LLVM startup and per-process object linking, which no on-disk cache removes. In-process dispatch once an op is built is untouched. The first process on a machine is slightly slower, not faster. Read-only installs are safe, verified rather than assumed. With graphblas/core/operator and its __pycache__ chmod'd a-w, the run wrote zero files in-tree and 66 .nbi/.nbc files under ~/Library/Caches/numba, with empty stderr under `python -W always`. The NumbaWarning in numba/core/caching.py is a source-content check, not a filesystem one. docs/user_guide/udf.rst gains a "Compilation caching" section covering the cache location, the read-only fallback, and why user ops are excluded. --- docs/user_guide/udf.rst | 24 ++++++++++++++++++++++ graphblas/core/operator/binary.py | 33 ++++++++++++++++++++++--------- graphblas/tests/test_op.py | 19 ++++++++++++++++++ 3 files changed, 67 insertions(+), 9 deletions(-) diff --git a/docs/user_guide/udf.rst b/docs/user_guide/udf.rst index 670cbc6e3..2f5bdcc01 100644 --- a/docs/user_guide/udf.rst +++ b/docs/user_guide/udf.rst @@ -188,3 +188,27 @@ time: # numba.njit has run. The compile happens on first lookup (``unary.heavy_op[int]``). + +Compilation caching +------------------- + +A handful of built-in operators are themselves UDFs (``binary.floordiv``, +``rfloordiv``, ``absfirst``, ``abssecond``, and ``rpow``). Because these are +plain module-level functions, python-graphblas compiles them with Numba's +``cache=True``, so the Numba half of the compile is paid once per machine +instead of once per Python process. Numba writes the cache (``.nbi`` index and +``.nbc`` data files) into the ``__pycache__`` directory next to the operator +module. A warm cache shaves roughly 20% off a single op's first-touch build and +about 45% when several are used; the rest is Numba/LLVM startup and per-process +object linking, which caching cannot remove. + +If that ``__pycache__`` directory is not writable (a read-only install, for +example), Numba silently falls back to its per-user cache directory +(``~/.cache/numba`` on Linux, ``~/Library/Caches/numba`` on macOS; set +``NUMBA_CACHE_DIR`` to override). No warning is emitted for the fallback. + +Operators you create with ``register_new`` or ``register_anonymous`` are never +passed ``cache=True``: Numba can only cache a function with a stable on-disk +source, which a lambda or an interactively defined function does not have. A +registered operator's compile is paid once per process, and ``lazy=True`` is +the tool for deferring it. diff --git a/graphblas/core/operator/binary.py b/graphblas/core/operator/binary.py index bdcfccc71..01c790acd 100644 --- a/graphblas/core/operator/binary.py +++ b/graphblas/core/operator/binary.py @@ -519,7 +519,7 @@ class BinaryOp(OpBase): } @classmethod - def _build(cls, name, func, *, is_udt=False, anonymous=False): + def _build(cls, name, func, *, is_udt=False, anonymous=False, cache=False): if not isinstance(func, FunctionType): raise TypeError(f"UDF argument must be a function, not {type(func)}") if name is None: @@ -532,7 +532,12 @@ def _build(cls, name, func, *, is_udt=False, anonymous=False): # setting it only on the wrapper leaves ``x // 0`` raising # ZeroDivisionError inside a cfunc, where Numba prints the traceback # and returns, handing GraphBLAS an element it never wrote. - binary_udf = numba.njit(func, error_model="numpy") + # + # ``cache`` is only True for the module-level built-in UDFs. Numba can + # only key a cache entry on a stable on-disk source, which lambdas and + # interactively defined functions do not have, so ops registered by + # users stay uncached. + binary_udf = numba.njit(func, error_model="numpy", cache=cache) new_type_obj = cls(name, func, anonymous=anonymous, is_udt=is_udt, numba_func=binary_udf) return_types = {} nt = numba.types @@ -730,7 +735,9 @@ def register_anonymous(cls, func, name=None, *, parameterized=False, is_udt=Fals return cls._build(name, func, anonymous=True, is_udt=is_udt) @classmethod - def register_new(cls, name, func, *, parameterized=False, is_udt=False, lazy=False): + def register_new( + cls, name, func, *, parameterized=False, is_udt=False, lazy=False, _cache=False + ): """Register a new BinaryOp and save it to ``graphblas.binary`` namespace. Parameters @@ -797,13 +804,14 @@ def register_new(cls, name, func, *, parameterized=False, is_udt=False, lazy=Fal "func": func, "parameterized": parameterized, "is_udt": is_udt, + "_cache": _cache, }, ) elif parameterized: binary_op = ParameterizedBinaryOp(name, func, is_udt=is_udt) setattr(module, funcname, binary_op) else: - binary_op = cls._build(name, func, is_udt=is_udt) + binary_op = cls._build(name, func, is_udt=is_udt, cache=_cache) setattr(module, funcname, binary_op) # Also save it to `graphblas.op` if not yet defined opmodule, funcname = cls._remove_nesting(name, module=op, modname="op", strict=False) @@ -856,13 +864,20 @@ def _initialize(cls): if _supports_udfs: # Add floordiv # cdiv truncates towards 0, while floordiv truncates towards -inf - BinaryOp.register_new("floordiv", _floordiv, lazy=True) # cast to integer - BinaryOp.register_new("rfloordiv", _rfloordiv, lazy=True) # cast to integer + # cache=True persists the numba compilation of these built-in UDFs to + # disk so it is paid once per machine, not once per process. It shaves + # ~20% off a single op's first-touch build and ~45% when several are + # used (the per-process object-link and LLVM-init cost is not cached). + # Only module-level built-ins get cache=True (see _build). + BinaryOp.register_new("floordiv", _floordiv, lazy=True, _cache=True) # cast to integer + BinaryOp.register_new( + "rfloordiv", _rfloordiv, lazy=True, _cache=True + ) # cast to integer # For aggregators - BinaryOp.register_new("absfirst", _absfirst, lazy=True) - BinaryOp.register_new("abssecond", _abssecond, lazy=True) - BinaryOp.register_new("rpow", _rpow, lazy=True) + BinaryOp.register_new("absfirst", _absfirst, lazy=True, _cache=True) + BinaryOp.register_new("abssecond", _abssecond, lazy=True, _cache=True) + BinaryOp.register_new("rpow", _rpow, lazy=True, _cache=True) # For algorithms binary._delayed["binom"] = (_register_binom, {}) # Lazy with custom creation diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index 71543251d..4361b8dc2 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -3808,3 +3808,22 @@ def test_initialize_does_not_build_lazy_udfs(): still_lazy = lines.get("still lazy", "").split() assert still_lazy == ["floordiv", "rfloordiv", "absfirst", "abssecond", "rpow"], report + + +@pytest.mark.skipif("not supports_udfs") +def test_builtin_udfs_are_disk_cached(): + # The built-in UDF binops are module-level functions, so numba can persist + # their compilation across processes. Anything a user registers cannot be: + # numba keys a cache entry on a stable on-disk source, which a lambda or an + # interactively defined function does not have. + from numba.core.caching import NullCache + + for name in ["floordiv", "rfloordiv", "absfirst", "abssecond", "rpow"]: + numba_func = getattr(binary, name)._numba_func + assert not isinstance(numba_func._cache, NullCache), name + + def _uncached_probe(x, y): + return x + y + + user_op = BinaryOp.register_anonymous(_uncached_probe) + assert isinstance(user_op._numba_func._cache, NullCache)