mirror of https://github.com/vllm-project/vllm.git
109 lines
3.5 KiB
Python
109 lines
3.5 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Ensure we perform lazy loading in vllm/__init__.py.
|
|
i.e: appears only within the ``if typing.TYPE_CHECKING:`` guard,
|
|
**except** for a short whitelist.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import ast
|
|
import pathlib
|
|
import sys
|
|
from collections.abc import Iterable
|
|
from typing import Final
|
|
|
|
REPO_ROOT: Final = pathlib.Path(__file__).resolve().parent.parent
|
|
INIT_PATH: Final = REPO_ROOT / "vllm" / "__init__.py"
|
|
|
|
# If you need to add items to whitelist, do it here.
|
|
ALLOWED_IMPORTS: Final[frozenset[str]] = frozenset({
|
|
"vllm.env_override",
|
|
})
|
|
ALLOWED_FROM_MODULES: Final[frozenset[str]] = frozenset({
|
|
".version",
|
|
})
|
|
|
|
|
|
def _is_internal(name: str | None, *, level: int = 0) -> bool:
|
|
if level > 0:
|
|
return True
|
|
if name is None:
|
|
return False
|
|
return name.startswith("vllm.") or name == "vllm"
|
|
|
|
|
|
def _fail(violations: Iterable[tuple[int, str]]) -> None:
|
|
print("ERROR: Disallowed eager imports in vllm/__init__.py:\n",
|
|
file=sys.stderr)
|
|
for lineno, msg in violations:
|
|
print(f" Line {lineno}: {msg}", file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
|
|
def main() -> None:
|
|
source = INIT_PATH.read_text(encoding="utf-8")
|
|
tree = ast.parse(source, filename=str(INIT_PATH))
|
|
|
|
violations: list[tuple[int, str]] = []
|
|
|
|
class Visitor(ast.NodeVisitor):
|
|
|
|
def __init__(self) -> None:
|
|
super().__init__()
|
|
self._in_type_checking = False
|
|
|
|
def visit_If(self, node: ast.If) -> None:
|
|
guard_is_type_checking = False
|
|
test = node.test
|
|
if isinstance(test, ast.Attribute) and isinstance(
|
|
test.value, ast.Name):
|
|
guard_is_type_checking = (test.value.id == "typing"
|
|
and test.attr == "TYPE_CHECKING")
|
|
elif isinstance(test, ast.Name):
|
|
guard_is_type_checking = test.id == "TYPE_CHECKING"
|
|
|
|
if guard_is_type_checking:
|
|
prev = self._in_type_checking
|
|
self._in_type_checking = True
|
|
for child in node.body:
|
|
self.visit(child)
|
|
self._in_type_checking = prev
|
|
for child in node.orelse:
|
|
self.visit(child)
|
|
else:
|
|
self.generic_visit(node)
|
|
|
|
def visit_Import(self, node: ast.Import) -> None:
|
|
if self._in_type_checking:
|
|
return
|
|
for alias in node.names:
|
|
module_name = alias.name
|
|
if _is_internal(
|
|
module_name) and module_name not in ALLOWED_IMPORTS:
|
|
violations.append((
|
|
node.lineno,
|
|
f"import '{module_name}' must be inside typing.TYPE_CHECKING", # noqa: E501
|
|
))
|
|
|
|
def visit_ImportFrom(self, node: ast.ImportFrom) -> None:
|
|
if self._in_type_checking:
|
|
return
|
|
module_as_written = ("." * node.level) + (node.module or "")
|
|
if _is_internal(
|
|
node.module, level=node.level
|
|
) and module_as_written not in ALLOWED_FROM_MODULES:
|
|
violations.append((
|
|
node.lineno,
|
|
f"from '{module_as_written}' import ... must be inside typing.TYPE_CHECKING", # noqa: E501
|
|
))
|
|
|
|
Visitor().visit(tree)
|
|
|
|
if violations:
|
|
_fail(violations)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|