[ci] fix slow tests (vllm-project#10698)

Signed-off-by: youkaichao <youkaichao@gmail.com>
Signed-off-by: Andrew Feldman <afeldman@neuralmagic.com>

Author: youkaichao

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`.")

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`.")

tests/vllm_test_utils/vllm_test_utils/blame.py

@@ -46,8 +46,8 @@ def _trace_calls(frame, event, arg=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)