[ci] fix slow tests (#10698)

Signed-off-by: youkaichao <youkaichao@gmail.com>
2024-11-27 09:26:14 -08:00
parent 9e0a147d50
commit 308cc5e21e
3 changed files with 39 additions and 15 deletions
--- 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)