diff --git a/CHANGELOG.md b/CHANGELOG.md index 5c34e2729..a1db2a1f2 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/llama_cpp/_internals.py b/llama_cpp/_internals.py index b45d34b2c..3ae9f31a7 100644 --- a/llama_cpp/_internals.py +++ b/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) diff --git a/llama_cpp/llama.py b/llama_cpp/llama.py index 14e2f8500..5b3e6bd4d 100644 --- a/llama_cpp/llama.py +++ b/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) diff --git a/tests/test_llama_cleanup.py b/tests/test_llama_cleanup.py new file mode 100644 index 000000000..b03ea5aaf --- /dev/null +++ b/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