A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 5, 12:15 UTC
Automatically add links to API docs for matching strings in docs (#37434)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Harry Mellor committed
Apr 7, 2026 at 15:21 UTC
79df4a794d6f2c933cf824b95afa0e9ae42723a9
2 files changed
+168
docs/mkdocs/hooks/autoref_code.py
new
+167
@@ -0,0 +1,167 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""
4
+MkDocs hook to automatically convert inline code references to API doc links.
5
+
6
+For example, `WeightTransferConfig` becomes
7
+[`WeightTransferConfig`][vllm.config.WeightTransferConfig]
8
+
9
+This works with the `autorefs` plugin to create clickable cross-references
10
+to API documentation pages generated by `mkdocstrings`.
11
+
12
+The hook builds an index of all documented public Python names (classes and
13
+functions with docstrings) from the vllm package at startup using AST parsing,
14
+then substitutes matching inline code spans on each page. Names without
15
+docstrings are excluded because mkdocstrings will not generate a page for them.
16
+"""
17
+
18
+import ast
19
+import logging
20
+from pathlib import Path
21
+
22
+import regex as re
23
+from mkdocs.config.defaults import MkDocsConfig
24
+from mkdocs.structure.files import Files
25
+from mkdocs.structure.pages import Page
26
+
27
+logger = logging.getLogger("mkdocs")
28
+
29
+ROOT_DIR = Path(__file__).parent.parent.parent.parent.resolve()
30
+VLLM_DIR = ROOT_DIR / "vllm"
31
+
32
+# Maps short name -> qualified name (e.g. "ModelConfig" -> "vllm.config.ModelConfig")
33
+_name_index: dict[str, str] = {}
34
+
35
+# Fenced code block pattern (``` or ~~~, with optional language specifier).
36
+_FENCED_BLOCK = re.compile(
37
+ r"(?:^|\n)(?P<fence>`{3,}|~{3,})[^\n]*\n.*?(?:\n(?P=fence))", re.DOTALL
38
+)
39
+
40
+# Inline code that is NOT already part of a markdown link.
41
+# Matches `Name` but not [`Name`] and not [`Name`][...] or [`Name`](...).
42
+_INLINE_CODE = re.compile(
43
+ r"(?<!\[)" # not preceded by [
44
+ r"`(?P<name>[A-Za-z0-9_]*)`" # `UpperCamelCase` or `UPPER_SNAKE`
45
+ r"(?!\])" # not followed by ]
46
+)
47
+
48
+
49
+def _has_docstring(node: ast.AST) -> bool:
50
+ """Check if a class or function node has a docstring."""
51
+ if not isinstance(node, ast.ClassDef | ast.FunctionDef | ast.AsyncFunctionDef):
52
+ return False
53
+ return ast.get_docstring(node, clean=False) is not None
54
+
55
+
56
+def _module_path(filepath: Path) -> str:
57
+ """Convert a filesystem path to a dotted module path."""
58
+ rel = filepath.relative_to(ROOT_DIR)
59
+ parts = list(rel.with_suffix("").parts)
60
+ if parts[-1] == "__init__":
61
+ parts = parts[:-1]
62
+ return ".".join(parts)
63
+
64
+
65
+def _index_file(filepath: Path) -> dict[str, str]:
66
+ """Extract documented public names from a Python file using AST parsing.
67
+
68
+ Only classes and functions with docstrings are included, since
69
+ mkdocstrings won't generate a page for undocumented symbols.
70
+ """
71
+ names: dict[str, str] = {}
72
+ try:
73
+ source = filepath.read_text(encoding="utf-8")
74
+ tree = ast.parse(source, filename=str(filepath))
75
+ except (SyntaxError, UnicodeDecodeError):
76
+ return names
77
+
78
+ module = _module_path(filepath)
79
+
80
+ for node in ast.iter_child_nodes(tree):
81
+ if (
82
+ # Class definitions (with docstring)
83
+ isinstance(node, ast.ClassDef)
84
+ and not node.name.startswith("_")
85
+ and _has_docstring(node)
86
+ ) or (
87
+ # Function definitions (with docstring, only uppercase/CamelCase)
88
+ isinstance(node, ast.FunctionDef | ast.AsyncFunctionDef)
89
+ and not node.name.startswith("_")
90
+ and node.name[0].isupper()
91
+ and _has_docstring(node)
92
+ ):
93
+ names[node.name] = f"{module}.{node.name}"
94
+
95
+ return names
96
+
97
+
98
+def _build_index() -> dict[str, str]:
99
+ """Walk the vllm package and build a name -> qualified path index."""
100
+ index: dict[str, str] = {}
101
+ # Track conflicts: if multiple modules define the same name,
102
+ # prefer shallower modules (more likely to be the public API).
103
+ depth: dict[str, int] = {}
104
+
105
+ for filepath in sorted(VLLM_DIR.rglob("*.py")):
106
+ # Skip internal/private modules
107
+ if any(part.startswith("_") and part != "__init__" for part in filepath.parts):
108
+ continue
109
+ # Skip third-party vendored code
110
+ rel = filepath.relative_to(VLLM_DIR)
111
+ if rel.parts and rel.parts[0] in ("third_party", "vllm_flash_attn"):
112
+ continue
113
+
114
+ module_depth = len(filepath.relative_to(ROOT_DIR).parts)
115
+ file_names = _index_file(filepath)
116
+
117
+ for name, qualified in file_names.items():
118
+ if name not in index or module_depth < depth[name]:
119
+ index[name] = qualified
120
+ depth[name] = module_depth
121
+
122
+ return index
123
+
124
+
125
+def on_startup(*, command: str, dirty: bool) -> None:
126
+ """Build the name index once at startup."""
127
+ global _name_index
128
+ _name_index = _build_index()
129
+ logger.info("autoref_code: indexed %d names from vllm/", len(_name_index))
130
+
131
+
132
+def on_page_markdown(
133
+ markdown: str, *, page: Page, config: MkDocsConfig, files: Files
134
+) -> str:
135
+ """Replace inline code references with autoref links."""
136
+ if not _name_index:
137
+ return markdown
138
+
139
+ # Skip API reference pages to avoid circular/redundant links.
140
+ if page.file.src_path.startswith("api/"):
141
+ return markdown
142
+
143
+ # Step 1: Mask fenced code blocks so we don't touch code inside them.
144
+ masks: list[str] = []
145
+
146
+ def _mask_block(match: re.Match) -> str:
147
+ masks.append(match.group(0))
148
+ return f"\ue000CODEBLOCK{len(masks) - 1}\ue000"
149
+
150
+ masked = _FENCED_BLOCK.sub(_mask_block, markdown)
151
+
152
+ # Step 2: Replace inline code references.
153
+ def _replace(match: re.Match) -> str:
154
+ name = match.group("name")
155
+ qualified = _name_index.get(name)
156
+ if qualified is None:
157
+ return match.group(0)
158
+ logger.debug("autoref_code: linking `%s` to [%s]", name, qualified)
159
+ return f"[`{name}`][{qualified}]"
160
+
161
+ result = _INLINE_CODE.sub(_replace, masked)
162
+
163
+ # Step 3: Restore masked code blocks.
164
+ result = re.sub(
165
+ r"\ue000CODEBLOCK(\d+)\ue000", lambda m: masks[int(m.group(1))], result
166
+ )
167
+ return result
mkdocs.yaml
+1
@@ -54,6 +54,7 @@ hooks:
54
- docs/mkdocs/hooks/generate_argparse.py
55
- docs/mkdocs/hooks/generate_metrics.py
56
- docs/mkdocs/hooks/url_schemes.py
57
+ - docs/mkdocs/hooks/autoref_code.py
58
59
plugins:
60
- meta