wehub-resource-sync
262 lines
10 KiB
Python
262 lines
10 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
import ast
|
|
import logging
|
|
import re
|
|
import sys
|
|
from argparse import SUPPRESS, Action, ArgumentParser, HelpFormatter, _ArgumentGroup
|
|
from collections.abc import Iterable
|
|
from importlib.machinery import ModuleSpec
|
|
from pathlib import Path
|
|
from typing import Literal
|
|
from unittest.mock import MagicMock
|
|
|
|
from pydantic_core import core_schema
|
|
|
|
|
|
class _FlexibleArgumentParser(ArgumentParser):
|
|
"""Fallback parser for docs when vllm is unavailable.
|
|
|
|
Accepts the 'deprecated' kwarg used by vllm CLI and emits warnings
|
|
if a deprecated argument is actually provided.
|
|
"""
|
|
|
|
_deprecated: set[Action] = set()
|
|
_deprecated_warned: set[str] = set()
|
|
|
|
if sys.version_info < (3, 13):
|
|
|
|
def parse_known_args(self, args=None, namespace=None):
|
|
namespace, args = super().parse_known_args(args, namespace)
|
|
for action in _FlexibleArgumentParser._deprecated:
|
|
if (
|
|
hasattr(namespace, dest := action.dest)
|
|
and getattr(namespace, dest) != action.default
|
|
and dest not in _FlexibleArgumentParser._deprecated_warned
|
|
):
|
|
_FlexibleArgumentParser._deprecated_warned.add(dest)
|
|
logger.warning("argument '%s' is deprecated", dest)
|
|
return namespace, args
|
|
|
|
def add_argument(self, *args, **kwargs):
|
|
deprecated = kwargs.pop("deprecated", False)
|
|
action = super().add_argument(*args, **kwargs)
|
|
if deprecated:
|
|
_FlexibleArgumentParser._deprecated.add(action)
|
|
return action
|
|
|
|
class _FlexibleArgumentGroup(_ArgumentGroup):
|
|
def add_argument(self, *args, **kwargs):
|
|
deprecated = kwargs.pop("deprecated", False)
|
|
action = super().add_argument(*args, **kwargs)
|
|
if deprecated:
|
|
_FlexibleArgumentParser._deprecated.add(action)
|
|
return action
|
|
|
|
def add_argument_group(self, *args, **kwargs):
|
|
group = self._FlexibleArgumentGroup(self, *args, **kwargs)
|
|
self._action_groups.append(group)
|
|
return group
|
|
|
|
|
|
logger = logging.getLogger("mkdocs")
|
|
|
|
# Define root and doc output directories
|
|
ROOT_DIR = Path(__file__).parent.parent.parent.parent
|
|
ARGPARSE_DOC_DIR = ROOT_DIR / "docs/generated/argparse_omni"
|
|
|
|
# Ensure the repo root is in sys.path for dynamic imports
|
|
sys.path.insert(0, str(ROOT_DIR))
|
|
|
|
|
|
class PydanticMagicMock(MagicMock):
|
|
"""`MagicMock` that's able to generate pydantic-core schemas and avoids _mock_name=None errors."""
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
name = kwargs.pop("name", None)
|
|
super().__init__(*args, **kwargs)
|
|
self.__spec__ = ModuleSpec(name, None)
|
|
self._mock_name = name or "mock"
|
|
|
|
def __get_pydantic_core_schema__(self, source_type, handler):
|
|
return core_schema.any_schema()
|
|
|
|
|
|
# --- Static extraction for CLI argument docs ---
|
|
|
|
|
|
def extract_omni_serve_subparser_init():
|
|
"""
|
|
Statically parse vllm_omni/entrypoints/cli/serve.py to extract the subparser_init method
|
|
and return a callable that adds arguments to a parser. This avoids import and mock issues.
|
|
"""
|
|
serve_path = ROOT_DIR / "vllm_omni" / "entrypoints" / "cli" / "serve.py"
|
|
with open(serve_path, encoding="utf-8") as f:
|
|
source = f.read()
|
|
tree = ast.parse(source, filename=str(serve_path))
|
|
# Find class OmniServeCommand
|
|
for node in tree.body:
|
|
if isinstance(node, ast.ClassDef) and node.name == "OmniServeCommand":
|
|
for item in node.body:
|
|
if isinstance(item, ast.FunctionDef) and item.name == "subparser_init":
|
|
# We'll exec this function body in a dummy context
|
|
func_code = ast.Module(body=[item], type_ignores=[])
|
|
code = compile(func_code, filename=str(serve_path), mode="exec")
|
|
# Prepare dummy context
|
|
local_vars = {}
|
|
|
|
# Provide a dummy subparsers with add_parser returning our parser
|
|
class DummySubparsers:
|
|
def add_parser(self, name, **kwargs):
|
|
return _FlexibleArgumentParser(prog=name)
|
|
|
|
dummy_subparsers = DummySubparsers()
|
|
# Provide globals for exec
|
|
# Extract DESCRIPTION from the source file (assume it's a triple-quoted string at the top)
|
|
m = re.search(r'DESCRIPTION\s*=\s*([ru]?""".*?""")', source, re.DOTALL)
|
|
DESCRIPTION = m.group(1)[3:-3] if m else ""
|
|
exec_globals = {
|
|
"_FlexibleArgumentParser": _FlexibleArgumentParser,
|
|
"FlexibleArgumentParser": _FlexibleArgumentParser,
|
|
"TrackingArgumentParser": _FlexibleArgumentParser, # only needs add_argument, not tracking
|
|
"TrackingNamespace": None, # type annotation only
|
|
"make_arg_parser": lambda parser: parser, # no-op for doc
|
|
"_ensure_vllm_platform": lambda: None, # no-op for doc
|
|
"VLLM_SUBCMD_PARSER_EPILOG": "",
|
|
"logger": logger,
|
|
"DummySubparsers": DummySubparsers,
|
|
"argparse": __import__("argparse"),
|
|
"json": __import__("json"),
|
|
"DESCRIPTION": DESCRIPTION,
|
|
}
|
|
exec(code, exec_globals, local_vars)
|
|
# Get the function
|
|
subparser_init = local_vars["subparser_init"]
|
|
|
|
# Return a callable that mimics OmniServeCommand().subparser_init
|
|
def parser_factory():
|
|
class DummySelf:
|
|
name = "serve"
|
|
|
|
return subparser_init(DummySelf(), dummy_subparsers)
|
|
|
|
return parser_factory
|
|
raise RuntimeError("Could not statically extract OmniServeCommand.subparser_init")
|
|
|
|
|
|
OmniServeCommand = type("OmniServeCommand", (), {"subparser_init": staticmethod(extract_omni_serve_subparser_init())})
|
|
|
|
|
|
class MarkdownFormatter(HelpFormatter):
|
|
"""Custom formatter that generates markdown for argument groups."""
|
|
|
|
def __init__(self, prog: str, starting_heading_level: int = 3):
|
|
super().__init__(prog, max_help_position=sys.maxsize, width=sys.maxsize)
|
|
|
|
self._section_heading_prefix = "#" * starting_heading_level
|
|
self._argument_heading_prefix = "#" * (starting_heading_level + 1)
|
|
self._markdown_output = []
|
|
|
|
def start_section(self, heading: str):
|
|
if heading not in {"positional arguments", "options"}:
|
|
heading_md = f"\n{self._section_heading_prefix} {heading}\n\n"
|
|
self._markdown_output.append(heading_md)
|
|
|
|
def end_section(self):
|
|
pass
|
|
|
|
def add_text(self, text: str):
|
|
if text:
|
|
self._markdown_output.append(f"{text.strip()}\n\n")
|
|
|
|
def add_usage(self, usage, actions, groups, prefix=None):
|
|
pass
|
|
|
|
def add_arguments(self, actions: Iterable[Action]):
|
|
for action in actions:
|
|
if len(action.option_strings) == 0 or "--help" in action.option_strings:
|
|
continue
|
|
|
|
option_strings = f"`{'`, `'.join(action.option_strings)}`"
|
|
heading_md = f"{self._argument_heading_prefix} {option_strings}\n\n"
|
|
self._markdown_output.append(heading_md)
|
|
|
|
if choices := action.choices:
|
|
choices = f"`{'`, `'.join(str(c) for c in choices)}`"
|
|
self._markdown_output.append(f"Possible choices: {choices}\n\n")
|
|
elif (metavar := action.metavar) and isinstance(metavar, (list, tuple)):
|
|
metavar = f"`{'`, `'.join(str(m) for m in metavar)}`"
|
|
self._markdown_output.append(f"Possible choices: {metavar}\n\n")
|
|
|
|
if action.help:
|
|
self._markdown_output.append(f"{action.help}\n\n")
|
|
|
|
if (default := action.default) != SUPPRESS:
|
|
# Make empty string defaults visible
|
|
if default == "":
|
|
default = '""'
|
|
self._markdown_output.append(f"Default: `{default}`\n\n")
|
|
|
|
def format_help(self):
|
|
"""Return the formatted help as markdown."""
|
|
return "".join(self._markdown_output)
|
|
|
|
|
|
# Function to create parser using subparser_init style CLI class
|
|
|
|
|
|
def create_parser_subparser_init(subcmd_class):
|
|
"""
|
|
Create an argparse parser using subparser_init style CLI class, with MarkdownFormatter.
|
|
Compatible with both AST-extracted static parser_factory and real class instance method.
|
|
"""
|
|
|
|
class DummySubparsers:
|
|
def add_parser(self, name, **kwargs):
|
|
return _FlexibleArgumentParser(prog=name)
|
|
|
|
dummy_subparsers = DummySubparsers()
|
|
# If subparser_init is a zero-arg static function (AST-extracted), call it directly
|
|
subparser_init = getattr(subcmd_class, "subparser_init", None)
|
|
if subparser_init is not None:
|
|
import inspect
|
|
|
|
sig = inspect.signature(subparser_init)
|
|
if len(sig.parameters) == 0:
|
|
parser = subparser_init()
|
|
else:
|
|
parser = subcmd_class().subparser_init(dummy_subparsers)
|
|
else:
|
|
parser = subcmd_class().subparser_init(dummy_subparsers)
|
|
parser.formatter_class = MarkdownFormatter
|
|
return parser
|
|
|
|
|
|
def on_startup(command: Literal["build", "gh-deploy", "serve"], dirty: bool):
|
|
"""
|
|
Entry point for doc generation. Builds doc directory and outputs markdown for each CLI command.
|
|
"""
|
|
logger.info("Generating vllm-omni argparse documentation")
|
|
logger.debug("Root directory: %s", ROOT_DIR.resolve())
|
|
logger.debug("Output directory: %s", ARGPARSE_DOC_DIR.resolve())
|
|
if not ARGPARSE_DOC_DIR.exists():
|
|
ARGPARSE_DOC_DIR.mkdir(parents=True)
|
|
|
|
# Register all CLI parsers; you can easily add more commands here
|
|
parsers = {
|
|
"omni_serve": create_parser_subparser_init(OmniServeCommand),
|
|
# "another_cmd": create_parser_subparser_init(AnotherCommandClass),
|
|
}
|
|
|
|
for stem, parser in parsers.items():
|
|
doc_path = ARGPARSE_DOC_DIR / f"{stem}.inc.md"
|
|
with open(doc_path, "w", encoding="utf-8") as f:
|
|
f.write(super(type(parser), parser).format_help())
|
|
logger.info("Argparse generated: %s", doc_path.relative_to(ROOT_DIR))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
logging.basicConfig(level=logging.INFO)
|
|
on_startup("build", False)
|