lazy_imports.py 1.49 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3

4
5
6
7
# Description: Test the lazy import module
# The utility function cannot be placed in `vllm.utils`
# this needs to be a standalone script
import sys
youkaichao's avatar
youkaichao committed
8
from contextlib import nullcontext
9

youkaichao's avatar
youkaichao committed
10
from vllm_test_utils import BlameResult, blame
11

12
13
14
15
16
17
18
19
20
21
22
# List of modules that should not be imported too early.
# Lazy import `torch._inductor.async_compile` to avoid creating
# too many processes before we set the number of compiler threads.
# Lazy import `cv2` to avoid bothering users who only use text models.
# `cv2` can easily mess up the environment.
module_names = ["torch._inductor.async_compile", "cv2"]


def any_module_imported():
    return any(module_name in sys.modules for module_name in module_names)

23

youkaichao's avatar
youkaichao committed
24
25
26
27
28
29
# 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
30
context = blame(any_module_imported) if use_blame else nullcontext()
youkaichao's avatar
youkaichao committed
31
with context as result:
32
33
    import vllm  # noqa

youkaichao's avatar
youkaichao committed
34
35
36
if use_blame:
    assert isinstance(result, BlameResult)
    print(f"the first import location is:\n{result.trace_stack}")
37

38
39
assert not any_module_imported(), (
    f"Some the modules in {module_names} are imported. To see the first"
youkaichao's avatar
youkaichao committed
40
    f" import location, run the test with `use_blame=True`.")