[Core] Fix circular reference which leaked llm instance in local dev env (#4737)

Storing exception frame is extremely prone to circular refernece because it contains the reference to objects.

When tensorizer is not installed, it leaks llm instance because error frame has references to various modules which cause circular reference problem.

I also found spec decoding has a circular reference issue, and I solved it using weakref.proxy.
This commit is contained in:
SangBin Cho
2024-05-10 23:54:32 +09:00
committed by GitHub
parent dac6a3f6ed
commit 6a0f617210
4 changed files with 22 additions and 7 deletions

View File

@@ -3,9 +3,12 @@
Run `pytest tests/basic_correctness/test_basic_correctness.py`.
"""
import os
import weakref
import pytest
from vllm import LLM
MODELS = [
"facebook/opt-125m",
"meta-llama/Llama-2-7b-hf",
@@ -13,6 +16,16 @@ MODELS = [
VLLM_ATTENTION_BACKEND = "VLLM_ATTENTION_BACKEND"
def test_vllm_gc_ed():
"""Verify vllm instance is GC'ed when it is deleted"""
llm = LLM("facebook/opt-125m")
weak_llm = weakref.ref(llm)
del llm
# If there's any circular reference to vllm, this fails
# because llm instance is not GC'ed.
assert weak_llm() is None
@pytest.mark.parametrize("model", MODELS)
@pytest.mark.parametrize("dtype", ["half"])
@pytest.mark.parametrize("max_tokens", [5])