diff --git a/tests/entrypoints/llm/test_lazy_outlines.py b/tests/entrypoints/llm/test_lazy_outlines.py index 81fb000d8ac5..2c53676c5f5d 100644 --- a/tests/entrypoints/llm/test_lazy_outlines.py +++ b/tests/entrypoints/llm/test_lazy_outlines.py @@ -1,6 +1,7 @@ import sys +from contextlib import nullcontext -from vllm_test_utils import blame +from vllm_test_utils import BlameResult, blame from vllm import LLM, SamplingParams from vllm.distributed import cleanup_dist_env_and_memory @@ -56,9 +57,20 @@ def test_lazy_outlines(sample_regex): """ # make sure outlines is not imported module_name = "outlines" - with blame(lambda: module_name in sys.modules) as result: + # In CI, we only check finally if the module is imported. + # If it is indeed imported, we can rerun the test with `use_blame=True`, + # which will trace every function call to find the first import location, + # and help find the root cause. + # We don't run it in CI by default because it is slow. + use_blame = False + context = blame( + lambda: module_name in sys.modules) if use_blame else nullcontext() + with context as result: run_normal() run_lmfe(sample_regex) - assert not result.found, ( - f"Module {module_name} is already imported, the" - f" first import location is:\n{result.trace_stack}") + if use_blame: + assert isinstance(result, BlameResult) + print(f"the first import location is:\n{result.trace_stack}") + assert module_name not in sys.modules, ( + f"Module {module_name} is imported. To see the first" + f" import location, run the test with `use_blame=True`.") diff --git a/tests/test_lazy_torch_compile.py b/tests/test_lazy_torch_compile.py index 4756fac8e2a8..b950877a4337 100644 --- a/tests/test_lazy_torch_compile.py +++ b/tests/test_lazy_torch_compile.py @@ -2,15 +2,27 @@ # The utility function cannot be placed in `vllm.utils` # this needs to be a standalone script import sys +from contextlib import nullcontext -from vllm_test_utils import blame +from vllm_test_utils import BlameResult, blame module_name = "torch._inductor.async_compile" -with blame(lambda: module_name in sys.modules) as result: +# In CI, we only check finally if the module is imported. +# If it is indeed imported, we can rerun the test with `use_blame=True`, +# which will trace every function call to find the first import location, +# and help find the root cause. +# We don't run it in CI by default because it is slow. +use_blame = False +context = blame( + lambda: module_name in sys.modules) if use_blame else nullcontext() +with context as result: import vllm # noqa -assert not result.found, (f"Module {module_name} is already imported, the" - f" first import location is:\n{result.trace_stack}") +if use_blame: + assert isinstance(result, BlameResult) + print(f"the first import location is:\n{result.trace_stack}") -print(f"Module {module_name} is not imported yet") +assert module_name not in sys.modules, ( + f"Module {module_name} is imported. To see the first" + f" import location, run the test with `use_blame=True`.") diff --git a/tests/vllm_test_utils/vllm_test_utils/blame.py b/tests/vllm_test_utils/vllm_test_utils/blame.py index ad23ab83c2d8..1ddd3471d357 100644 --- a/tests/vllm_test_utils/vllm_test_utils/blame.py +++ b/tests/vllm_test_utils/vllm_test_utils/blame.py @@ -46,8 +46,8 @@ def blame(func: Callable) -> Generator[BlameResult, None, None]: pass return _trace_calls - sys.settrace(_trace_calls) - - yield result - - sys.settrace(None) + try: + sys.settrace(_trace_calls) + yield result + finally: + sys.settrace(None)