Fix Docker Numba cache warmup

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"]

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

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

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
-

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