From 31ba776fb3410f73ca62649c2ffa0aac91c7996f Mon Sep 17 00:00:00 2001 From: Arjan Zijderveld <5286904+arjanz@users.noreply.github.com> Date: Fri, 16 Oct 2026 12:32:52 +0200 Subject: [PATCH] Fix Docker Numba cache warmup and enforce cache-only smoke check --- Dockerfile | 4 + README.md | 7 + .../numba/interpreter_numba_aot_ffi.py | 224 +----------------- scripts/build_numba_aot.sh | 33 +-- scripts/check_numba_cache.py | 54 +++++ test/test_numba_aot_build.py | 116 +++++++++ 6 files changed, 197 insertions(+), 241 deletions(-) create mode 100644 scripts/check_numba_cache.py create mode 100644 test/test_numba_aot_build.py diff --git a/Dockerfile b/Dockerfile index 1193ba47..76e60651 100644 --- a/Dockerfile +++ b/Dockerfile @@ -49,4 +49,8 @@ RUN python -m compileall -b -d /app ./pyjamaz && \ chmod -R 777 /app/numba-cache && \ chmod -R 777 /app/pyjamaz/pvm/interpreters/numba +# Verify a fresh process can execute the PVM from cache after source stripping. +# This runs for each target architecture and must succeed before image export. +RUN python ./scripts/check_numba_cache.py + ENTRYPOINT ["./scripts/docker-entrypoint.sh"] diff --git a/README.md b/README.md index 7333f052..9c6924e2 100644 --- a/README.md +++ b/README.md @@ -28,6 +28,13 @@ docker buildx inspect --bootstrap docker buildx build --platform linux/amd64,linux/arm64 -t jamdottech/pyjamaz -t jamdottech/pyjamaz:vX.Y.Z-gpX.Y.Z --push . ``` +The Docker build warms the Numba cache from the interpreter's canonical +signatures and stops on compilation errors. After removing Python sources, it +runs `scripts/check_numba_cache.py` in a fresh process with project recompilation +blocked by the cache-only guard. The check executes a small PVM program and verifies +its registers, gas, program counter, exit status, and interpreter cache hit. +Both target architectures must pass this check before an image is published. + ## JAM prize M1 accounts * Polkadot: 146CmUoArEi1E2AogKCU5gkhBSN6BLDzxecFSCDAgVyEshra * Kusama: DBPAKp9B2gpBr9YmvqXyewYkR4ZwTn9uJDt64zGcVjZeiGi diff --git a/pyjamaz/pvm/interpreters/numba/interpreter_numba_aot_ffi.py b/pyjamaz/pvm/interpreters/numba/interpreter_numba_aot_ffi.py index 3b452448..4b081e7c 100644 --- a/pyjamaz/pvm/interpreters/numba/interpreter_numba_aot_ffi.py +++ b/pyjamaz/pvm/interpreters/numba/interpreter_numba_aot_ffi.py @@ -1,217 +1,15 @@ -# Minimal AOT wrapper that pre compiled interpreter_numba_jit +"""Warm the disk cache used by the NUMBA_AOT interpreter. -import pyjamaz.pvm.interpreters.numba.interpreter_numba_aot_patch +The interpreter and its helpers declare explicit, cached @njit signatures. +Importing them eagerly compiles those canonical signatures; keeping a second +list here lets the build drift from the runtime interface. +""" -import pyjamaz.pvm.interpreters.numba.interpreter_numba_aot_patch # apply Numba cache/locator patches +from . import interpreter_numba_aot_patch # noqa: F401; apply cache/CPU patches first +from .interpreter_numba_jit import invoke_native -from numba.core import config as nb_config -nb_config.DISABLE_JIT = False -import numpy as np - -from numba import types, int32, int64, uint32, uint8, uint64, boolean -from numba.pycc import CC - - -# ---- Export container ---- -cc = CC('interpreter_numba_aot') - - -u8_array_1d = types.Array(uint8, 1, 'C') -u8_array_list = types.ListType(u8_array_1d) -u64_array_1d = types.Array(uint64, 1, 'C') -u64_array_list = types.ListType(u64_array_1d) -int32_array_1d = types.Array(int32, 1, 'C') - - -# -# # Note: keep signature the same as invoke_jit -# @cc.export( -# 'invoke_native', -# int32( -# uint32, # pc -# int64, # gas -# uint32, # inst_nr -# uint32, # skip_len -# -# uint8[::1], # code -# uint32, # code_size -# int32[::1], # inst_arg_len_array -# int32[::1], # pc_to_inst_index -# int32[::1], # opcode_scheme_array (len 256) -# int32[::1], # jump_table_array -# -# uint64[::1], # mem_section_starts -# uint64[::1], # mem_section_ends -# u8_array_list, # section_arrays : List[uint8[:]] -# u64_array_list, # acl_bitmaps -# int32[::1], # section_access -# uint64[::1], # heap_info (len 3) -# -# uint64[::1], # reg (len 13) -# -# boolean, # logging_enabled -# types.ListType(types.unicode_type), # opcode_names list -# -# uint64[::1], # registers_out -# int64[::1], # state_out -# int64[::1], # heap_grew_out -# ) -# ) -# def invoke_native( -# pc_start, -# gas_start, -# inst_start, -# initial_skip_len, -# -# code, -# code_size, -# inst_arg_len, -# pc_to_inst_index, -# opcode_scheme, -# jump_table, -# -# mem_section_starts, -# mem_section_ends, -# section_arrays, -# acl_bitmaps, -# section_access, -# heap_info, # [current_heap_end, next_section_start, mem_writable_value] -# -# registers_in, -# -# logging, -# opcode_names, -# -# registers_out, -# state_out, -# heap_grew_out): -# -# # Delegate to the JIT implementation. Ensure dtypes at the call-site match this signature. -# return invoke_native_jit( -# pc_start, -# gas_start, -# inst_start, -# initial_skip_len, -# -# code, -# code_size, -# inst_arg_len, -# pc_to_inst_index, -# opcode_scheme, -# jump_table, -# -# mem_section_starts, -# mem_section_ends, -# section_arrays, -# acl_bitmaps, -# section_access, -# heap_info, # [current_heap_end, next_section_start, mem_writable_value] -# -# registers_in, -# -# logging, -# opcode_names, -# -# registers_out, -# state_out, -# heap_grew_out -# ) - - -if __name__ == '__main__': - from .interpreter_numba_jit import ( - invoke_native as invoke_native_jit, - sync_state_and_return, - sbrk_jit, - branch_jit, - djump_jit, - log - ) - - nb_config.DISABLE_JIT = False - - sig = int32( - uint32, # pc - int64, # gas - uint32, # inst_nr - uint32, # skip_len - - uint8[::1], # code - uint32, # code_size - int32[::1], # inst_arg_len_array - int32[::1], # pc_to_inst_index - int32[::1], # opcode_scheme_array (len 256) - int32[::1], # jump_table_array - - uint64[::1], # mem_section_starts - uint64[::1], # mem_section_ends - u8_array_list, # section_arrays : List[uint8[:]] - u64_array_list, # acl_bitmaps - int32[::1], # section_access - uint64[::1], # heap_info (len 3) - - uint64[::1], # reg (len 13) - - boolean, # logging_enabled - types.ListType(types.unicode_type), # opcode_names list - - uint64[::1], # registers_out - int64[::1], # state_out - int64[::1], # heap_grew_out - ) - - invoke_native_jit.compile(sig) - - sig = uint32(uint64[::1], uint64[::1], int64[::1], int64, int64, int64, int64, int64, uint32, uint32) - sync_state_and_return.compile(sig) - - sig = types.Tuple((uint64, int64))( - uint64, - uint64, - uint64, - int64, - u8_array_list, - uint64[::1], - u64_array_list - ) - sbrk_jit.compile(sig) - - sig = int32(uint32, int64, boolean, int32[::1]) - branch_jit.compile(sig) - - sig = int32(uint32, int32[::1], uint32, int32[::1]) - djump_jit.compile(sig) - - sig = types.void( - types.ListType(types.unicode_type), # opcode_names list - int64[::1], # local_state - uint64[::1], # regs - types.optional(u8_array_list), # mem - types.optional(uint64[::1]), # mem_starts - types.optional(uint64[::1]) - ) - log.compile(sig) - - from pyjamaz.pvm.interpreters.numba import defs as d - d.umul64wide_jit.compile(types.UniTuple(uint64, 2)(uint64, uint64)) - d.imul64wide_jit.compile(types.UniTuple(uint64, 2)(int64, int64)) - d.smul_u64wide_jit.compile(types.UniTuple(uint64, 2)(int64, uint64)) - d.rori64_jit.compile(uint64(uint64, uint64)) - d.roli64_jit.compile(uint64(uint64, uint64)) - d.rori32_jit.compile(uint32(uint32, uint32)) - d.roli32_jit.compile(uint32(uint32, uint32)) - d.pvm_smod_jit.compile(int64(int64, int64)) - d.pvm_rtz_div_jit.compile(int64(int64, int64)) - d.pvm_X_jit.compile(uint64(uint64, uint64)) - d.pvm_Z_jit.compile(int64(uint64, uint64)) - d.count_leading_zeroes_jit.compile(uint64(uint64, uint8)) - d.count_trailing_zeroes_jit.compile(uint64(uint64, uint8)) - d.reverse_bytes_jit.compile(uint64(uint64)) - d.riscv_div_jit.compile(int64(int64, int64)) - d.pvm_Z_inv_jit.compile(uint64(int64, uint8)) - d.read_uint_jit.compile(uint64(uint8[::1], uint32, uint8)) - d.mem_write_jit.compile(types.Tuple((int32, uint64))(uint64,uint64,uint8, uint64[::1],uint64[::1],u8_array_list,int32[::1])) - d.mem_read_jit.compile(types.Tuple((int32, uint64))(uint64,uint8,uint64[::1],uint64[::1],u8_array_list,int32[::1])) - - #cc.compile() +if __name__ == "__main__": + if not invoke_native.nopython_signatures: + raise RuntimeError("Numba interpreter cache warmup produced no compiled signatures") + print("Numba interpreter cache warmed using canonical @njit signatures") diff --git a/scripts/build_numba_aot.sh b/scripts/build_numba_aot.sh index 23db0fab..90a8c553 100755 --- a/scripts/build_numba_aot.sh +++ b/scripts/build_numba_aot.sh @@ -1,5 +1,5 @@ #!/usr/bin/env bash -set -uo pipefail +set -euo pipefail echo " Building Numba AOT PVM interpreter " @@ -9,12 +9,11 @@ if [[ "${DEBUG:-0}" == "1" ]]; then fi SCRIPT_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )/.." -NUMBA_DIR="$SCRIPT_DIR/pyjamaz/pvm/interpreters/numba" +cd "$SCRIPT_DIR" # Match CPU between build & run export NUMBA_CPU_NAME=generic -# Where to write/read the cache -#export NUMBA_CACHE_DIR=/app/numba-cache +export NUMBA_CPU_FEATURES="" # Cache-only, but allow warmup compiles export NUMBA_CACHE_ONLY=1 export NUMBA_CACHE_WARMUP=1 @@ -23,32 +22,10 @@ unset NUMBA_DISABLE_JIT export PVM_INTERPRETER=NUMBA_AOT_COMPILE echo "Script dir: $SCRIPT_DIR" -echo "Numba dir : $NUMBA_DIR" echo echo "Compiling interpreter_numba_aot..." echo "Started at: $(date)" -if python -m pyjamaz.pvm.interpreters.numba.interpreter_numba_aot_ffi; then - echo "✓ AOT build completed successfully" -else - echo "⚠️ AOT build failed (continuing shell session)" -fi +python -m pyjamaz.pvm.interpreters.numba.interpreter_numba_aot_ffi +echo "✓ AOT cache warmup completed successfully" echo "Finished at: $(date)" - - -unset NUMBA_CPU_NAME -unset NUMBA_CACHE_ONLY -unset NUMBA_CACHE_WARMUP -unset NUMBA_DISABLE_JIT -#unset PVM_INTERPRETER - - -echo -echo "Listing built files..." -found="$(find "$NUMBA_DIR" -maxdepth 1 -type f -name 'interpreter_numba_aot_ffi*' -print)" -if [ -n "$found" ]; then - ls -lh $found -else - echo "No artifacts built yet" -fi - diff --git a/scripts/check_numba_cache.py b/scripts/check_numba_cache.py new file mode 100644 index 00000000..e26902e3 --- /dev/null +++ b/scripts/check_numba_cache.py @@ -0,0 +1,54 @@ +#!/usr/bin/env python3 +"""Execute a small PVM program using only the image's prebuilt Numba cache.""" + +import os +from pathlib import Path +import sys + + +def main(): + # Set these before importing pyjamaz or Numba, including when the parent + # process has inherited build-time warmup or disabled-JIT settings. + os.environ.update({ + "PVM_INTERPRETER": "NUMBA_AOT", + "NUMBA_CACHE_ONLY": "1", + "NUMBA_CACHE_WARMUP": "0", + "NUMBA_CPU_NAME": "generic", + "NUMBA_CPU_FEATURES": "", + "NUMBA_DISABLE_JIT": "0", + }) + sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + + from jamcodec.base import JamBytes + + from pyjamaz.pvm import PVMInterpreter, PVMMemory + from pyjamaz.pvm.constants import ExitReason + from pyjamaz.pvm.interpreters.numba.interpreter_numba_jit import invoke_native + from pyjamaz.pvm.types import PVMCode, PVMProgram + + if PVMInterpreter.__module__ != "pyjamaz.pvm.interpreters.numba.interpreter_numba_jit": + raise RuntimeError("Cache smoke check requires the Numba PVM interpreter") + + # load_imm_64 r7, 0xDEADBEEFDEADBEEF; trap. Include the complete PVM + # code blob (header, instruction bytes, and opcode bitmask) so this check + # also works in the final image, which does not contain test fixtures. + code_blob = bytes.fromhex("00 00 0b 14 07 ef be ad de ef be ad de 00 01 04") + code = PVMCode.from_jam_bytes(JamBytes(code_blob)) + pvm = PVMInterpreter(PVMProgram(code, [0] * 13, PVMMemory())) + pvm.invoke(0, 10000) + + expected_registers = [0] * 13 + expected_registers[7] = 0xDEADBEEFDEADBEEF + actual = (pvm.get_registers(), int(pvm.pc), int(pvm.gas), pvm.status) + expected = (expected_registers, 10, 9998, ExitReason.panic.value) + if actual != expected: + raise RuntimeError(f"Cached PVM execution mismatch: {actual!r} != {expected!r}") + + cache_hits = sum(invoke_native.stats.cache_hits.values()) + if cache_hits < 1: + raise RuntimeError("PVM interpreter did not load an invoke_native overload from cache") + print(f"Numba cache-only PVM smoke check passed ({cache_hits} interpreter cache hit(s))") + + +if __name__ == "__main__": + main() diff --git a/test/test_numba_aot_build.py b/test/test_numba_aot_build.py new file mode 100644 index 00000000..6806db87 --- /dev/null +++ b/test/test_numba_aot_build.py @@ -0,0 +1,116 @@ +import os +from pathlib import Path +import shlex +import shutil +import subprocess +import sys + + +ROOT = Path(__file__).resolve().parents[1] + + +def run_build_step(args, app, env): + result = subprocess.run( + args, cwd=app, env=env, capture_output=True, text=True, timeout=300 + ) + assert result.returncode == 0, result.stdout + result.stderr + return result + + +def test_aot_build_propagates_compiler_failure(tmp_path): + # Inject a compiler failure without importing the PVM or changing its cache. + bin_dir = tmp_path / "bin" + bin_dir.mkdir() + python = bin_dir / "python" + python.write_text("#!/bin/sh\nexit 42\n") + python.chmod(0o755) + env = dict(os.environ, PATH=str(bin_dir) + os.pathsep + os.environ["PATH"]) + + result = subprocess.run( + ["bash", str(ROOT / "scripts/build_numba_aot.sh")], + cwd=tmp_path, env=env, capture_output=True, text=True, timeout=30, + ) + + assert result.returncode == 42, result.stdout + result.stderr + assert "AOT cache warmup completed successfully" not in result.stdout + + +def test_cache_only_smoke_after_docker_source_stripping(tmp_path): + app = tmp_path / "app" + app.mkdir() + shutil.copytree( + ROOT / "pyjamaz", app / "pyjamaz", + ignore=shutil.ignore_patterns( + "local_settings.py", "__pycache__", "*.pyc", "*.pyo", "*.nbc", "*.nbi" + ), + ) + scripts = app / "scripts" + scripts.mkdir() + for name in ("build_numba_aot.sh", "check_numba_cache.py"): + shutil.copy2(ROOT / "scripts" / name, scripts / name) + + # The build script invokes `python`; keep it in this test's environment. + bin_dir = tmp_path / "bin" + bin_dir.mkdir() + python = bin_dir / "python" + python.write_text(f"#!/bin/sh\nexec {shlex.quote(sys.executable)} \"$@\"\n") + python.chmod(0o755) + env = {key: value for key, value in os.environ.items() if not key.startswith("NUMBA_")} + for key in ("JAM_FUZZ", "JAM_FUZZ_SPEC", "PYTHONOPTIMIZE"): + env.pop(key, None) + cache = tmp_path / "numba-cache" + cache.mkdir() + env.update( + PATH=str(bin_dir) + os.pathsep + os.environ["PATH"], + PYTHONPATH=str(app), + PYTHONDONTWRITEBYTECODE="1", + NUMBA_CACHE_DIR=str(cache), + NUMBA_CPU_NAME="generic", + NUMBA_CPU_FEATURES="", + NUMBA_CACHE_ONLY="1", + NUMBA_CACHE_WARMUP="0", + PVM_INTERPRETER="NUMBA_AOT", + ) + run_build_step(["bash", str(scripts / "build_numba_aot.sh")], app, env) + assert list(cache.rglob("*.nbc")), "Warmup did not populate the clean cache" + + # Reproduce the Docker layout: legacy bytecode everywhere, retaining only + # the Numba sources needed by its cache locator and code fingerprints. + run_build_step( + [sys.executable, "-m", "compileall", "-q", "-b", "-d", "/app", "./pyjamaz"], + app, env, + ) + numba_dir = app / "pyjamaz/pvm/interpreters/numba" + for source in (app / "pyjamaz").rglob("*.py"): + if not source.is_relative_to(numba_dir): + source.unlink() + shutil.rmtree(numba_dir / "__pycache__", ignore_errors=True) + assert not (app / "pyjamaz/__init__.py").exists() + assert (app / "pyjamaz/__init__.pyc").is_file() + assert not (numba_dir / "__pycache__").exists() + assert all(source.is_relative_to(numba_dir) for source in (app / "pyjamaz").rglob("*.py")) + + # A fresh process must resolve the stripped copy, even with an editable + # installation of pyjamaz available in the parent test environment. + smoke = ( + "from pathlib import Path; import pyjamaz; import runpy; " + f"expected = Path({str(app / 'pyjamaz/__init__.pyc')!r}); " + "origin = Path(pyjamaz.__file__).resolve(); " + "assert origin == expected.resolve(), (origin, expected); " + f"runpy.run_path({str(scripts / 'check_numba_cache.py')!r}, run_name='__main__')" + ) + args = [sys.executable, "-c", smoke] + result = run_build_step(args, app, env) + assert "Numba cache-only PVM smoke check passed" in result.stdout + + # Prove the smoke check cannot silently recompile when an image cache is + # missing, rather than merely succeeding via a JIT fallback. + empty_cache = tmp_path / "empty-cache" + empty_cache.mkdir() + env["NUMBA_CACHE_DIR"] = str(empty_cache) + result = subprocess.run( + args, cwd=app, env=env, capture_output=True, text=True, timeout=300 + ) + assert result.returncode != 0, result.stdout + result.stderr + assert "NUMBA cache-only: cached overload not found" in result.stderr + assert not list(empty_cache.rglob("*.nbc"))