fix: preserve native cleanup callbacks during interpreter shutdown

CHANGELOG.md

@@ -7,6 +7,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+- fix: preserve native cleanup callbacks during interpreter shutdown by @emme1t in #2397
+
 ## [0.3.36]
 
 - feat: update llama.cpp to ggml-org/llama.cpp@0c1e57098

llama_cpp/_internals.py

@@ -69,10 +69,11 @@ def __init__(
         self.model = model
         self.vocab = vocab
 
-        def free_model():
+        # Keep the native function alive through interpreter shutdown.
+        def free_model(model_free=llama_cpp.llama_model_free):
             if self.model is None:
                 return
-            llama_cpp.llama_model_free(self.model)
+            model_free(self.model)
             self.model = None
 
         self._exit_stack.callback(free_model)
@@ -269,10 +270,10 @@ def __init__(
         self.memory = llama_cpp.llama_get_memory(self.ctx)
         self.sampler = None  # LlamaContext doesn't manage samplers directly, but some cleanup code expects this attribute
 
-        def free_ctx():
+        def free_ctx(ctx_free=llama_cpp.llama_free):
             if self.ctx is None:
                 return
-            llama_cpp.llama_free(self.ctx)
+            ctx_free(self.ctx)
             self.ctx = None
 
         self._exit_stack.callback(free_ctx)
@@ -484,10 +485,10 @@ def __init__(
         self.batch = batch
         self.sampler = None  # LlamaBatch doesn't use samplers, but some cleanup code expects this attribute
 
-        def free_batch():
+        def free_batch(batch_free=llama_cpp.llama_batch_free):
             if self.batch is None:
                 return
-            llama_cpp.llama_batch_free(self.batch)
+            batch_free(self.batch)
             self.batch = None
 
         self._exit_stack.callback(free_batch)
@@ -681,12 +682,15 @@ def __init__(self):
         self.custom_samplers: List[Tuple[int, CustomSampler]] = []
         self._exit_stack = ExitStack()
 
-        def free_sampler():
+        def free_sampler(
+            sampler_free=llama_cpp.llama_sampler_free,
+            chain_remove=llama_cpp.llama_sampler_chain_remove,
+        ):
             if self.sampler is not None:
                 # NOTE: Must remove custom samplers before free or llama.cpp will try to free them
                 for i, _ in reversed(self.custom_samplers):
-                    llama_cpp.llama_sampler_chain_remove(self.sampler, i)
-                llama_cpp.llama_sampler_free(self.sampler)
+                    chain_remove(self.sampler, i)
+                sampler_free(self.sampler)
                 self.sampler = None
 
         self._exit_stack.callback(free_sampler)
@@ -879,9 +883,9 @@ def clone(self):
         new_sampler.custom_samplers = []
         new_sampler._exit_stack = ExitStack()
 
-        def free_sampler():
+        def free_sampler(sampler_free=llama_cpp.llama_sampler_free):
             if new_sampler.sampler is not None:
-                llama_cpp.llama_sampler_free(new_sampler.sampler)
+                sampler_free(new_sampler.sampler)
                 new_sampler.sampler = None
 
         new_sampler._exit_stack.callback(free_sampler)

llama_cpp/llama.py

@@ -443,10 +443,10 @@ def __init__(
                     f"Failed to initialize LoRA adapter from lora path: {self.lora_path}"
                 )
 
-            def free_lora_adapter():
+            def free_lora_adapter(adapter_free=llama_cpp.llama_adapter_lora_free):
                 if self._lora_adapter is None:
                     return
-                llama_cpp.llama_adapter_lora_free(self._lora_adapter)
+                adapter_free(self._lora_adapter)
                 self._lora_adapter = None
 
             self._stack.callback(free_lora_adapter)

tests/test_llama_cleanup.py

@@ -0,0 +1,103 @@
+import subprocess
+import sys
+import textwrap
+
+import pytest
+
+
+@pytest.mark.parametrize("clear_globals", ["none", "bindings", "modules"])
+@pytest.mark.parametrize("cleanup", ["close", "gc"])
+def test_cleanup_after_module_teardown(clear_globals, cleanup):
+    # Isolate cleared module globals and native resource lifetimes from pytest.
+    script = textwrap.dedent(
+        """
+        import ctypes
+        import gc
+        import sys
+        from collections import Counter
+
+        import llama_cpp.llama as llama_module
+        import llama_cpp.llama_cpp as bindings
+        import llama_cpp._internals as internals
+
+        calls = []
+        errors = []
+        sys.unraisablehook = errors.append
+
+        def record(name, function):
+            def wrapped(*args):
+                calls.append((name, args[1] if name == 'remove' else None))
+                return function(*args)
+            return wrapped
+
+        free_functions = {
+            'llama_model_free': 'model',
+            'llama_free': 'context',
+            'llama_batch_free': 'batch',
+            'llama_sampler_free': 'sampler',
+            'llama_sampler_chain_remove': 'remove',
+            'llama_adapter_lora_free': 'adapter',
+        }
+        # Exercise the adapter cleanup path without downloading a LoRA model.
+        bindings.llama_adapter_lora_init = lambda *args: ctypes.pointer(ctypes.c_void_p(1))
+        bindings.llama_set_adapters_lora = lambda *args: 0
+        bindings.llama_adapter_lora_free = lambda *args: None
+        for function, name in free_functions.items():
+            setattr(bindings, function, record(name, getattr(bindings, function)))
+
+        model = llama_module.Llama(
+            model_path='./vendor/llama.cpp/models/ggml-vocab-llama-spm.gguf',
+            vocab_only=True,
+            lora_path='test-adapter',
+            verbose=False,
+        )
+        sampler = internals.LlamaSampler()
+        sampler.add_greedy()
+        clone = sampler.clone()
+        sampler.add_custom(lambda candidates: None)
+        sampler.add_custom(lambda candidates: None)
+
+        clear_globals, cleanup = sys.argv[1:]
+        if clear_globals == 'bindings':
+            for function in free_functions:
+                setattr(bindings, function, None)
+        elif clear_globals == 'modules':
+            internals.llama_cpp = None
+            llama_module.llama_cpp = None
+
+        if cleanup == 'close':
+            clone.close()
+            sampler.close()
+            model.close()
+            # Explicit close and subsequent finalization must release once.
+            clone.close()
+            sampler.close()
+            model.close()
+            assert clone.sampler is None
+            assert sampler.sampler is None
+            assert model._model.model is None
+            assert model._ctx.ctx is None
+            assert model._batch.batch is None
+            assert model._lora_adapter is None
+
+        del clone, sampler, model
+        gc.collect()
+
+        assert not errors, [str(error.exc_value) for error in errors]
+        assert Counter(name for name, _ in calls) == {
+            'model': 1, 'context': 1, 'batch': 1,
+            'sampler': 2, 'remove': 2, 'adapter': 1,
+        }, calls
+        assert [index for name, index in calls if name == 'remove'] == [2, 1]
+        assert calls.index(('context', None)) < calls.index(('model', None))
+        assert calls.index(('adapter', None)) < calls.index(('model', None))
+        """
+    )
+    result = subprocess.run(
+        [sys.executable, "-c", script, clear_globals, cleanup],
+        capture_output=True,
+        text=True,
+        timeout=60,
+    )
+    assert result.returncode == 0, result.stdout + result.stderr
+    assert "Exception ignored" not in result.stderr