validate_config.py 5.3 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Ensures all fields in a config dataclass have default values
and that each field has a docstring.
"""

import ast
import inspect
import sys
11
from itertools import pairwise
12

13
14
import regex as re

15
16
17
18
19
20
21
22
23
24
25
26
27
28

def get_attr_docs(cls_node: ast.ClassDef) -> dict[str, str]:
    """
    Get any docstrings placed after attribute assignments in a class body.

    Adapted from https://davidism.com/attribute-docstrings/
    https://davidism.com/mit-license/
    """

    out = {}

    # Consider each pair of nodes.
    for a, b in pairwise(cls_node.body):
        # Must be an assignment then a constant string.
29
30
31
32
33
34
        if (
            not isinstance(a, (ast.Assign, ast.AnnAssign))
            or not isinstance(b, ast.Expr)
            or not isinstance(b.value, ast.Constant)
            or not isinstance(b.value.value, str)
        ):
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
            continue

        doc = inspect.cleandoc(b.value.value)

        # An assignment can have multiple targets (a = b = v), but an
        # annotated assignment only has one target.
        targets = a.targets if isinstance(a, ast.Assign) else [a.target]

        for target in targets:
            # Must be assigning to a plain name.
            if not isinstance(target, ast.Name):
                continue

            out[target.id] = doc

    return out


class ConfigValidator(ast.NodeVisitor):
54
    def __init__(self): ...
55
56

    def visit_ClassDef(self, node):
57
58
59
60
61
62
63
64
65
        # Validate classes with a @config decorator
        decorators = set()
        for decorator in node.decorator_list:
            if isinstance(decorator, ast.Call):
                decorator = decorator.func
            if isinstance(decorator, ast.Name) and decorator.id == "config":
                decorators.add(decorator.id)

        if decorators == {"config"}:
66
            validate_class(node)
67
68
        elif "config" in decorators:
            fail(f"config decorator for {node.name} should be used alone", node)
69
70
71
72
73
74
75
76
77
78

        self.generic_visit(node)


def validate_class(class_node: ast.ClassDef):
    attr_docs = get_attr_docs(class_node)

    for stmt in class_node.body:
        # A field is defined as a class variable that has a type annotation.
        if isinstance(stmt, ast.AnnAssign):
79
            # Skip ClassVar and InitVar
80
            # see https://docs.python.org/3/library/dataclasses.html#class-variables
81
            # and https://docs.python.org/3/library/dataclasses.html#init-only-variables
82
83
84
85
86
            if (
                isinstance(stmt.annotation, ast.Subscript)
                and isinstance(stmt.annotation.value, ast.Name)
                and stmt.annotation.value.id in {"ClassVar", "InitVar"}
            ):
87
88
89
90
91
92
93
                continue

            if isinstance(stmt.target, ast.Name):
                field_name = stmt.target.id
                if stmt.value is None:
                    fail(
                        f"Field '{field_name}' in {class_node.name} must have "
94
95
96
                        "a default value.",
                        stmt,
                    )
97
98
99
100

                if field_name not in attr_docs:
                    fail(
                        f"Field '{field_name}' in {class_node.name} must have "
101
102
103
                        "a docstring.",
                        stmt,
                    )
104

105
106
107
108
109
110
                if (
                    isinstance(stmt.annotation, ast.Subscript)
                    and isinstance(stmt.annotation.value, ast.Name)
                    and stmt.annotation.value.id == "Union"
                    and isinstance(stmt.annotation.slice, ast.Tuple)
                ):
111
112
                    args = stmt.annotation.slice.elts
                    literal_args = [
113
114
115
116
117
                        arg
                        for arg in args
                        if isinstance(arg, ast.Subscript)
                        and isinstance(arg.value, ast.Name)
                        and arg.value.id == "Literal"
118
119
120
121
122
123
124
                    ]
                    if len(literal_args) > 1:
                        fail(
                            f"Field '{field_name}' in {class_node.name} must "
                            "use a single "
                            "Literal type. Please use 'Literal[Literal1, "
                            "Literal2]' instead of 'Union[Literal1, Literal2]'"
125
126
127
                            ".",
                            stmt,
                        )
128
129
130
131
132
133
134
135


def validate_ast(tree: ast.stmt):
    ConfigValidator().visit(tree)


def validate_file(file_path: str):
    try:
136
        print(f"Validating {file_path} config dataclasses ", end="")
137
138
139
140
141
142
143
        with open(file_path, encoding="utf-8") as f:
            source = f.read()

        tree = ast.parse(source, filename=file_path)
        validate_ast(tree)
    except ValueError as e:
        print(e)
144
        raise SystemExit(1) from e
145
146
147
148
149
150
151
152
153
154
    else:
        print("✅")


def fail(message: str, node: ast.stmt):
    raise ValueError(f"❌ line({node.lineno}): {message}")


def main():
    for filename in sys.argv[1:]:
155
156
157
158
159
160
161
        # Only run for Python files in vllm/ or tests/
        if not re.match(r"^(vllm|tests)/.*\.py$", filename):
            continue
        # Only run if the file contains @config
        with open(filename, encoding="utf-8") as f:
            if "@config" in f.read():
                validate_file(filename)
162
163
164
165


if __name__ == "__main__":
    main()