mirror of
https://github.com/NVIDIA/TensorRT-LLM.git
synced 2026-01-14 06:27:45 +08:00
341 lines
12 KiB
Python
341 lines
12 KiB
Python
import logging
|
|
import re
|
|
from dataclasses import dataclass
|
|
from itertools import chain, groupby
|
|
from pathlib import Path
|
|
from typing import Optional
|
|
|
|
|
|
def underline(title: str, character: str = "=") -> str:
|
|
return f"{title}\n{character * len(title)}"
|
|
|
|
|
|
def generate_title(filename: str) -> str:
|
|
with open(filename) as f:
|
|
# fine the first line that contains '###'
|
|
for line in f:
|
|
if '###' in line:
|
|
title = line[3:].strip()
|
|
break
|
|
assert title is not None, f"No title found in {filename}"
|
|
return underline(title)
|
|
|
|
|
|
@dataclass
|
|
class DocMeta:
|
|
title: str
|
|
order: int
|
|
section: str
|
|
filename: Path
|
|
|
|
|
|
def extract_meta_info(filename: str) -> Optional[DocMeta]:
|
|
"""Extract metadata from file following the pattern ### :[a-zA-Z_]+[0-9]* <value>"""
|
|
metadata_pattern = re.compile(r'^### :([a-zA-Z_]+[0-9]*)\s+(.+)$')
|
|
|
|
with open(filename) as f:
|
|
metadata = DocMeta(title="",
|
|
order=0,
|
|
section="",
|
|
filename=Path(filename))
|
|
|
|
for line in f:
|
|
line = line.strip()
|
|
match = metadata_pattern.match(line)
|
|
if match:
|
|
key = match.group(1).strip()
|
|
value = match.group(2).strip()
|
|
setattr(metadata, key, value)
|
|
elif not line.startswith('###'):
|
|
continue
|
|
if metadata.title == "":
|
|
return None
|
|
return metadata
|
|
|
|
|
|
# NOTE: Update here to keep consistent with the examples
|
|
LLMAPI_SECTIONS = ["Basics", "Customization", "Slurm"]
|
|
|
|
|
|
def generate_examples():
|
|
root_dir = Path(__file__).parent.parent.parent.resolve()
|
|
ignore_list = {
|
|
'__init__.py', 'quickstart_example.py', 'quickstart_advanced.py',
|
|
'quickstart_multimodal.py', 'star_attention.py'
|
|
}
|
|
doc_dir = root_dir / "docs/source/examples"
|
|
|
|
def collect_script_paths(examples_subdir: str) -> list[Path]:
|
|
"""Collect Python and shell script paths from an examples subdirectory."""
|
|
script_dir = root_dir / f"examples/{examples_subdir}"
|
|
script_paths = list(
|
|
chain(script_dir.glob("*.py"), script_dir.glob("*.sh")))
|
|
return [
|
|
path for path in sorted(script_paths)
|
|
if path.name not in ignore_list
|
|
]
|
|
|
|
# Collect source paths for LLMAPI examples
|
|
llmapi_script_paths = collect_script_paths("llm-api")
|
|
llmapi_doc_paths = [
|
|
doc_dir / f"{path.stem}.rst" for path in llmapi_script_paths
|
|
]
|
|
llmapi_script_base_url = "https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/llm-api"
|
|
|
|
# Collect source paths for trtllm-serve examples
|
|
serve_script_paths = collect_script_paths("serve")
|
|
serve_doc_paths = [
|
|
doc_dir / f"{path.stem}.rst" for path in serve_script_paths
|
|
]
|
|
serve_script_base_url = "https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/serve"
|
|
|
|
def _get_lines_without_metadata(filename: str) -> str:
|
|
"""Get line ranges that exclude metadata lines.
|
|
Returns a string like "5-10,15-20" for use in :lines: directive.
|
|
"""
|
|
with open(filename) as f:
|
|
metadata_pattern = re.compile(r'^### :([a-zA-Z_]+[0-9]*)\s+(.+)$')
|
|
all_lines = f.readlines()
|
|
|
|
# Find line numbers that are NOT metadata (1-indexed)
|
|
content_lines = []
|
|
for line_num, line in enumerate(all_lines, 1):
|
|
line_stripped = line.strip()
|
|
# Include line if it's not empty and not metadata
|
|
if not metadata_pattern.match(line_stripped):
|
|
content_lines.append(line_num)
|
|
|
|
if not content_lines:
|
|
return "" # No content lines found
|
|
|
|
# Group consecutive line numbers into ranges
|
|
ranges = []
|
|
start = content_lines[0]
|
|
end = start
|
|
|
|
for line_num in content_lines[1:]:
|
|
if line_num == end + 1:
|
|
# Consecutive line, extend current range
|
|
end = line_num
|
|
else:
|
|
# Gap found, close current range and start new one
|
|
if start == end:
|
|
ranges.append(str(start))
|
|
else:
|
|
ranges.append(f"{start}-{end}")
|
|
start = line_num
|
|
end = line_num
|
|
|
|
# Add the final range
|
|
if start == end:
|
|
ranges.append(str(start))
|
|
else:
|
|
ranges.append(f"{start}-{end}")
|
|
|
|
return ",".join(ranges)
|
|
|
|
# Generate the example docs for each example script
|
|
def write_scripts(base_url: str,
|
|
example_script_paths: list[Path],
|
|
doc_paths: list[Path],
|
|
extra_content="") -> list[DocMeta]:
|
|
metas = []
|
|
for script_path, doc_path in zip(example_script_paths, doc_paths):
|
|
if script_path.name in ignore_list:
|
|
logging.warning(f"Ignoring file: {script_path.name}")
|
|
continue
|
|
script_url = f"{base_url}/{script_path.name}"
|
|
|
|
# Determine language based on file extension
|
|
language = "python" if script_path.suffix == ".py" else "bash"
|
|
|
|
# Make script_path relative to doc_path and call it include_path
|
|
include_path = '../../..' / script_path.relative_to(root_dir)
|
|
|
|
# Extract metadata from the script file
|
|
if meta := extract_meta_info(str(script_path)):
|
|
title = underline(meta.title)
|
|
else:
|
|
logging.warning(
|
|
f"No metadata found for {script_path.name}, using filename as title"
|
|
)
|
|
title = script_path.stem.replace('_', ' ').title()
|
|
meta = DocMeta(title=title,
|
|
order=0,
|
|
section="",
|
|
filename=script_path)
|
|
title = underline(title)
|
|
metas.append(meta)
|
|
|
|
# Get line ranges excluding metadata
|
|
lines_without_metadata = _get_lines_without_metadata(
|
|
str(script_path))
|
|
|
|
# Build literalinclude directive
|
|
literalinclude_lines = [f".. literalinclude:: {include_path}"]
|
|
if lines_without_metadata:
|
|
literalinclude_lines.append(
|
|
f" :lines: {lines_without_metadata}")
|
|
literalinclude_lines.extend(
|
|
[f" :language: {language}", f" :linenos:"])
|
|
|
|
content = (f"{title}\n"
|
|
f"{extra_content}"
|
|
f"Source {script_url}.\n\n"
|
|
f"{chr(10).join(literalinclude_lines)}\n")
|
|
with open(doc_path, "w+") as f:
|
|
logging.warning(f"Writing {doc_path}")
|
|
f.write(content)
|
|
|
|
return metas
|
|
|
|
def write_index(metas: list[DocMeta], doc_template_path: Path,
|
|
doc_path: Path, example_name: str,
|
|
section_order: list[str]):
|
|
"""Write the index file for the examples.
|
|
|
|
Args:
|
|
metas: The metadata for the examples.
|
|
doc_template_path: The path to the template file.
|
|
doc_path: The path to the output file.
|
|
example_name: The name of the examples.
|
|
section_order: The order of sections to display.
|
|
|
|
The template file is expected to have the following placeholders:
|
|
- %EXAMPLE_DOCS%: The documentation for the examples.
|
|
- %EXAMPLE_NAME%: The name of the examples.
|
|
"""
|
|
with open(doc_template_path) as f:
|
|
template_content = f.read()
|
|
|
|
# Sort metadata by section order and example order
|
|
sort_key = lambda x: (section_order.index(x.section)
|
|
if section_order and x.section in section_order
|
|
else 0, int(x.order))
|
|
metas.sort(key=sort_key)
|
|
|
|
content = []
|
|
for section, group in groupby(metas, key=lambda x: x.section):
|
|
if section_order and section not in section_order:
|
|
raise ValueError(
|
|
f"Section '{section}' not in section_order {section_order}")
|
|
|
|
group_list = list(group)
|
|
content.extend([
|
|
section, "_" * len(section), "", ".. toctree::",
|
|
" :maxdepth: 2", ""
|
|
])
|
|
|
|
for meta in group_list:
|
|
content.append(f" {meta.filename.stem}")
|
|
content.append("")
|
|
|
|
example_docs = "\n".join(content)
|
|
|
|
# Replace placeholders and write to file
|
|
output_content = template_content.replace("%EXAMPLE_DOCS%",
|
|
example_docs).replace(
|
|
"%EXAMPLE_NAME%",
|
|
example_name)
|
|
with open(doc_path, "w") as f:
|
|
f.write(output_content)
|
|
|
|
# Generate the toctree for LLMAPI example scripts
|
|
llmapi_metas = write_scripts(llmapi_script_base_url, llmapi_script_paths,
|
|
llmapi_doc_paths)
|
|
write_index(metas=llmapi_metas,
|
|
doc_template_path=doc_dir / "llm_examples_index.template.rst_",
|
|
doc_path=doc_dir / "llm_api_examples.rst",
|
|
example_name="LLM Examples",
|
|
section_order=LLMAPI_SECTIONS)
|
|
|
|
# Generate the toctree for trtllm-serve example scripts
|
|
serve_extra_content = (
|
|
"Refer to the `trtllm-serve documentation "
|
|
"<https://nvidia.github.io/TensorRT-LLM/commands/trtllm-serve.html>`_ "
|
|
"for starting a server.\n\n")
|
|
serve_metas = write_scripts(serve_script_base_url, serve_script_paths,
|
|
serve_doc_paths, serve_extra_content)
|
|
write_index(metas=serve_metas,
|
|
doc_template_path=doc_dir / "llm_examples_index.template.rst_",
|
|
doc_path=doc_dir / "trtllm_serve_examples.rst",
|
|
example_name="Online Serving Examples",
|
|
section_order=[])
|
|
|
|
|
|
def extract_all_and_eval(file_path):
|
|
''' Extract the __all__ variable from a Python file.
|
|
This is a trick to make the CI happy even the tensorrt_llm lib is not available.
|
|
NOTE: This requires the __all__ variable to be defined at the end of the file.
|
|
'''
|
|
with open(file_path, 'r') as file:
|
|
content = file.read()
|
|
|
|
lines = content.split('\n')
|
|
filtered_line_begin = 0
|
|
|
|
for i, line in enumerate(lines):
|
|
if line.startswith("__all__"):
|
|
filtered_line_begin = i
|
|
break
|
|
|
|
code_to_eval = '\n'.join(lines[filtered_line_begin:])
|
|
|
|
local_vars = {}
|
|
exec(code_to_eval, {}, local_vars)
|
|
return local_vars
|
|
|
|
|
|
def get_pydantic_methods() -> list[str]:
|
|
from pydantic import BaseModel
|
|
|
|
class Dummy(BaseModel):
|
|
pass
|
|
|
|
methods = set(
|
|
[method for method in dir(Dummy) if not method.startswith('_')])
|
|
methods.discard("__init__")
|
|
return list(methods)
|
|
|
|
|
|
def generate_llmapi():
|
|
root_dir = Path(__file__).parent.parent.parent.resolve()
|
|
|
|
# Set up destination paths
|
|
doc_dir = root_dir / "docs/source/llm-api"
|
|
doc_dir.mkdir(exist_ok=True)
|
|
doc_path = doc_dir / "reference.rst"
|
|
|
|
llmapi_all_file = root_dir / "tensorrt_llm/llmapi/__init__.py"
|
|
public_classes_names = extract_all_and_eval(llmapi_all_file)['__all__']
|
|
|
|
content = underline("API Reference", "-") + "\n\n"
|
|
for cls_name in public_classes_names:
|
|
cls_name = cls_name.strip()
|
|
options = [
|
|
" :members:",
|
|
" :undoc-members:",
|
|
" :show-inheritance:",
|
|
" :special-members: __init__",
|
|
" :member-order: groupwise",
|
|
]
|
|
|
|
options.append(" :inherited-members:")
|
|
if cls_name in ["TorchLlmArgs", "TrtLlmArgs"]:
|
|
# exclude tons of methods from Pydantic
|
|
options.append(
|
|
f" :exclude-members: {','.join(get_pydantic_methods())}")
|
|
|
|
content += f".. autoclass:: tensorrt_llm.llmapi.{cls_name}\n"
|
|
content += "\n".join(options) + "\n\n"
|
|
|
|
with open(doc_path, "w+") as f:
|
|
f.write(content)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import os
|
|
path = os.environ["TEKIT_ROOT"] + "/examples/llm-api/llm_inference.py"
|
|
#print(extract_meta_info(path))
|
|
generate_examples()
|