mirror of
https://github.com/NVIDIA/TensorRT-LLM.git
synced 2026-01-14 06:27:45 +08:00
81 lines
2.6 KiB
Python
81 lines
2.6 KiB
Python
import gc
|
||
import inspect
|
||
import types
|
||
import weakref
|
||
from contextlib import contextmanager
|
||
|
||
from tensorrt_llm.logger import logger
|
||
|
||
|
||
@contextmanager
|
||
def assert_resource_freed(object_creation_func, *args, **kwargs):
|
||
"""
|
||
Create a resource via object_creation_func(*args, **kwargs),
|
||
disable the generational GC to force pure refcount freeing,
|
||
yield the resource, then assert that it was destroyed when the context exits.
|
||
If it wasn’t freed, collect and report all remaining referrers.
|
||
"""
|
||
# Ensure a clean start
|
||
gc.collect()
|
||
gc_was_enabled = gc.isenabled()
|
||
gc.disable()
|
||
|
||
resource = object_creation_func(*args, **kwargs)
|
||
alive_ref = weakref.ref(resource)
|
||
|
||
try:
|
||
yield resource
|
||
|
||
finally:
|
||
# Drop our own strong reference
|
||
try:
|
||
del resource
|
||
except NameError:
|
||
pass
|
||
|
||
# If still alive, diagnose
|
||
leaked = alive_ref()
|
||
if leaked is not None:
|
||
# Restore GC so we can introspect
|
||
if gc_was_enabled:
|
||
gc.enable()
|
||
# Give GC a chance to finalize anything pending
|
||
gc.collect()
|
||
|
||
# Find all objects still referring to our instance
|
||
refs = gc.get_referrers(leaked)
|
||
# Filter out inspection internals (frames, tracebacks, the weakref itself, etc.)
|
||
filtered = []
|
||
for r in refs:
|
||
# skip the weakref container itself
|
||
if isinstance(r, dict) and any(
|
||
isinstance(v, weakref.ref) and v() is leaked
|
||
for v in r.values()):
|
||
continue
|
||
# skip our own local variables frame
|
||
if inspect.isframe(r):
|
||
continue
|
||
# skip the generator’s internal cell
|
||
if isinstance(r, types.CellType):
|
||
continue
|
||
filtered.append(r)
|
||
|
||
# Build a human‐readable report
|
||
report_lines = [
|
||
f" - {type(r).__name__} at 0x{id(r):x}: {repr(r)[:200]!r}"
|
||
for r in filtered
|
||
]
|
||
report = "\n".join(
|
||
report_lines) or " <no non‐internal referrers found>"
|
||
|
||
if filtered:
|
||
raise AssertionError(
|
||
"Resource was NOT freed upon context exit!\n"
|
||
f"{len(filtered)} referrer(s) still alive:\n{report}\n")
|
||
else:
|
||
logger.info("Resource was freed upon context exit.")
|
||
|
||
# Otherwise, restore GC state
|
||
if gc_was_enabled:
|
||
gc.enable()
|