@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 233 lines 8.32 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import itertools
4 import logging
5 from dataclasses import dataclass
6 from functools import cached_property
7 from pathlib import Path
8
9 import mkdocs_awesome_nav.nav.directory as _nav_dir
10 import mkdocs_gen_files
11 import regex as re
12
13 logger = logging.getLogger("mkdocs")
14
15 ROOT_DIR = Path(__file__).parent.parent.parent.parent
16 EXAMPLE_DIR = ROOT_DIR / "examples"
17
18
19 def title(text: str) -> str:
20 # Default title case
21 text = text.replace("_", " ").replace("/", " - ").title()
22 # Custom substitutions
23 subs = {
24 "io": "IO",
25 "rl": "RL",
26 "api(s?)": r"API\1",
27 "cli": "CLI",
28 "cpu": "CPU",
29 "ipc": "IPC",
30 "llm": "LLM",
31 "mae": "MAE",
32 "ner": "NER",
33 "tpu": "TPU",
34 "lora": "LoRA",
35 "nccl": "NCCL",
36 "rlhf": "RLHF",
37 "vllm": "vLLM",
38 "openai": "OpenAI",
39 "lmcache": "LMCache",
40 "multilora": "MultiLoRA",
41 "mlpspeculator": "MLPSpeculator",
42 r"fp\d+": lambda x: x.group(0).upper(), # e.g. fp16, fp32
43 r"int\d+": lambda x: x.group(0).upper(), # e.g. int8, int16
44 }
45 for pattern, repl in subs.items():
46 text = re.sub(rf"\b{pattern}\b", repl, text, flags=re.IGNORECASE)
47 return text
48
49
50 @dataclass
51 class Example:
52 """Example class for generating documentation content from a given path.
53
54 Attributes:
55 path (Path): The path to the main directory or file.
56 category (str): The category of the document.
57
58 Properties::
59 main_file() -> Path | None: Determines the main file in the given path.
60 other_files() -> list[Path]: Determines other files in the directory excluding
61 the main file.
62 title() -> str: Determines the title of the document.
63
64 Methods:
65 generate() -> str: Generates the documentation content.
66
67 """
68
69 path: Path
70 category: str
71
72 @cached_property
73 def main_file(self) -> Path | None:
74 """Determines the main file in the given path.
75
76 If path is a file, it returns the path itself. If path is a directory, it
77 searches for Markdown files (*.md) in the directory and returns the first one
78 found. If no Markdown files are found, it returns None."""
79 # Single file example
80 if self.path.is_file():
81 return self.path
82 # Multi file example with a README
83 if md_paths := list(self.path.glob("*.md")):
84 return md_paths[0]
85 # Multi file example without a README
86 return None
87
88 @cached_property
89 def other_files(self) -> list[Path]:
90 """Determine other files in the directory excluding the main file.
91
92 If path is a file, it returns an empty list. Otherwise, it returns every file
93 in the directory except the main file in a list."""
94 # Single file example
95 if self.path.is_file():
96 return []
97 # Multi file example
98 is_other_file = lambda file: file.is_file() and file != self.main_file
99 return sorted(file for file in self.path.rglob("*") if is_other_file(file))
100
101 @cached_property
102 def is_code(self) -> bool:
103 return self.main_file is not None and self.main_file.suffix != ".md"
104
105 @cached_property
106 def title(self) -> str:
107 # Generate title from filename if no main md file found
108 if self.main_file is None or self.is_code:
109 return title(self.path.stem)
110 # Specify encoding for building on Windows
111 with open(self.main_file, encoding="utf-8") as f:
112 first_line = f.readline().strip()
113 match = re.match(r"^#\s+(?P<title>.+)$", first_line)
114 if match:
115 return match.group("title")
116 raise ValueError(f"Title not found in {self.main_file}")
117
118 def fix_relative_links(self, content: str) -> str:
119 """Fix relative links in markdown content by converting them to gh-file
120 format.
121
122 Args:
123 content (str): The markdown content to process
124
125 Returns:
126 str: Content with relative links converted to gh-file format
127
128 """
129 # Regex to match markdown links [text](relative_path)
130 # This matches links that don't start with http, https, ftp, or #
131 link_pattern = r"\[([^\]]*)\]\((?!(?:https?|ftp)://|#)([^)]+)\)"
132
133 def replace_link(match):
134 link_text = match.group(1)
135 relative_path = match.group(2)
136
137 # Make relative to repo root
138 gh_file = (self.main_file.parent / relative_path).resolve()
139 gh_file = gh_file.relative_to(ROOT_DIR)
140
141 # Make GitHub URL
142 url = "https://github.com/vllm-project/vllm/"
143 url += "tree/main" if self.path.is_dir() else "blob/main"
144 gh_url = f"{url}/{gh_file}"
145
146 return f"[{link_text}]({gh_url})"
147
148 return re.sub(link_pattern, replace_link, content)
149
150 def generate(self) -> str:
151 content = f"# {self.title}\n\n"
152 url = "https://github.com/vllm-project/vllm/"
153 url += "tree/main" if self.path.is_dir() else "blob/main"
154 content += f"Source <{url}/{self.path.relative_to(ROOT_DIR)}>.\n\n"
155
156 # Use long code fence to avoid issues with
157 # included files containing code fences too
158 code_fence = "``````"
159
160 if self.main_file is not None:
161 # Single file example or multi file example with a README
162 if self.is_code:
163 content += (
164 f"{code_fence}{self.main_file.suffix[1:]}\n"
165 f'--8<-- "{self.main_file}"\n'
166 f"{code_fence}\n"
167 )
168 else:
169 with open(self.main_file, encoding="utf-8") as f:
170 # Skip the title from md snippets as it's been included above
171 main_content = f.readlines()[1:]
172 content += self.fix_relative_links("".join(main_content))
173 content += "\n"
174 else:
175 # Multi file example without a README
176 for file in self.other_files:
177 file_title = title(str(file.relative_to(self.path).with_suffix("")))
178 content += f"## {file_title}\n\n"
179 content += (
180 f'{code_fence}{file.suffix[1:]}\n--8<-- "{file}"\n{code_fence}\n\n'
181 )
182 return content
183
184 if not self.other_files:
185 return content
186
187 content += "## Example materials\n\n"
188 for file in self.other_files:
189 content += f'??? abstract "{file.relative_to(self.path)}"\n'
190 if file.suffix != ".md":
191 content += f" {code_fence}{file.suffix[1:]}\n"
192 content += f' --8<-- "{file}"\n'
193 if file.suffix != ".md":
194 content += f" {code_fence}\n"
195
196 return content
197
198
199 # Monkey-patch dirname_to_title in awesome-nav so that sub-directory names are
200 # title-cased (e.g. "Offline Inference" instead of "Offline inference").
201 _nav_dir.dirname_to_title = title
202 logger.info("Generating example documentation")
203 logger.debug("Root directory: %s", ROOT_DIR.resolve())
204 logger.debug("Example directory: %s", EXAMPLE_DIR.resolve())
205
206 categories = sorted(
207 p for p in EXAMPLE_DIR.iterdir() if p.is_dir() and not p.name.startswith(".")
208 )
209
210 examples = []
211 glob_patterns = ["*.py", "*.md", "*.sh"]
212 # Find categorised examples
213 for category in categories:
214 logger.info("Processing category: %s", category.stem)
215 globs = [category.glob(pattern) for pattern in glob_patterns]
216 for path in itertools.chain(*globs):
217 examples.append(Example(path, category.stem))
218 # Find examples in subdirectories
219 globs = [category.glob(f"*/{pattern}") for pattern in glob_patterns]
220 for path in itertools.chain(*globs):
221 examples.append(Example(path.parent, category.stem))
222
223 # Generate the example documentation
224 for example in sorted(examples, key=lambda e: e.path.stem):
225 doc_path = f"examples/{example.category}/{example.path.stem}.md"
226 with mkdocs_gen_files.open(doc_path, "w") as f:
227 f.write(example.generate())
228 if example.main_file is not None:
229 # Point the edit button at the example's source file
230 edit_path = Path("..") / example.main_file.relative_to(ROOT_DIR)
231 mkdocs_gen_files.set_edit_path(doc_path, str(edit_path))
232 logger.debug("Example generated: %s", doc_path)
233 logger.info("Total examples generated: %d", len(examples))