0
tools/docs_codegen/__init__.py
Normal file
0
tools/docs_codegen/__init__.py
Normal file
93
tools/docs_codegen/cli.py
Normal file
93
tools/docs_codegen/cli.py
Normal file
@@ -0,0 +1,93 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import sys
|
||||
from typing import TextIO
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Make `python3 tools/docs_codegen/cli.py ...` importable regardless of the
|
||||
# launch directory by putting the repo root (this file's parents[2]) on the path.
|
||||
from pathlib import Path
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
|
||||
|
||||
from tools.docs_codegen.errors import DocsCodegenError
|
||||
|
||||
|
||||
def build_arg_parser() -> argparse.ArgumentParser:
|
||||
"""Build the CLI argument parser (mutually exclusive ``--doc`` / ``--block`` selection)."""
|
||||
arg_parser = argparse.ArgumentParser(description="Generate shell code blocks from model-code directives.")
|
||||
selection_group = arg_parser.add_mutually_exclusive_group()
|
||||
selection_group.add_argument(
|
||||
"--doc",
|
||||
dest="doc_path",
|
||||
help="Generate all blocks from one repository-relative markdown path.",
|
||||
)
|
||||
selection_group.add_argument(
|
||||
"--block",
|
||||
dest="block_ref",
|
||||
help="Block reference in '<doc_path>::<block_name>' form.",
|
||||
)
|
||||
_add_generate_flags(arg_parser)
|
||||
|
||||
return arg_parser
|
||||
|
||||
|
||||
def main(
|
||||
argv: list[str] | None = None,
|
||||
*,
|
||||
stdout: TextIO | None = None,
|
||||
stderr: TextIO | None = None,
|
||||
) -> int:
|
||||
"""CLI entry point; returns a process exit code (0 ok, 1 on a known generation error)."""
|
||||
stdout = stdout or sys.stdout
|
||||
stderr = stderr or sys.stderr
|
||||
args = build_arg_parser().parse_args(argv)
|
||||
|
||||
try:
|
||||
return _handle_generate(args, stdout=stdout)
|
||||
except DocsCodegenError as exc:
|
||||
print(exc, file=stderr)
|
||||
return 1
|
||||
|
||||
|
||||
def _add_generate_flags(arg_parser: argparse.ArgumentParser) -> None:
|
||||
"""Register the ``--stdout`` and ``--dry-run`` generation flags."""
|
||||
arg_parser.add_argument("--stdout", action="store_true", help="Print generated content after the output path.")
|
||||
arg_parser.add_argument("--dry-run", action="store_true", help="Generate content without writing files.")
|
||||
|
||||
|
||||
def _handle_generate(args: argparse.Namespace, *, stdout: TextIO) -> int:
|
||||
"""Run generation for all blocks, one document, or one block per the parsed args."""
|
||||
from tools.docs_codegen.generator import create_default_generator_service
|
||||
|
||||
service = create_default_generator_service()
|
||||
if args.doc_path is not None:
|
||||
generated_artifacts = service.generate_document(args.doc_path, dry_run=args.dry_run)
|
||||
elif args.block_ref is not None:
|
||||
doc_path, block_name = _parse_block_ref(args.block_ref)
|
||||
generated_artifacts = [service.generate_block(doc_path, block_name, dry_run=args.dry_run)]
|
||||
else:
|
||||
generated_artifacts = service.generate_all(dry_run=args.dry_run)
|
||||
|
||||
for output_path, script in generated_artifacts:
|
||||
print(output_path, file=stdout)
|
||||
if args.stdout:
|
||||
print(script.content.rstrip(), file=stdout)
|
||||
|
||||
return 0
|
||||
|
||||
|
||||
def _parse_block_ref(block_ref: str) -> tuple[str, str]:
|
||||
"""Split a ``<doc_path>::<block_name>`` reference into its two parts."""
|
||||
if "::" not in block_ref:
|
||||
raise DocsCodegenError("block reference must use '<doc_path>::<block_name>'")
|
||||
|
||||
doc_path, block_name = block_ref.rsplit("::", 1)
|
||||
if not doc_path or not block_name:
|
||||
raise DocsCodegenError("block reference must use '<doc_path>::<block_name>'")
|
||||
return doc_path, block_name
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
512
tools/docs_codegen/converters.py
Normal file
512
tools/docs_codegen/converters.py
Normal file
@@ -0,0 +1,512 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import shlex
|
||||
from abc import ABC, abstractmethod
|
||||
from collections import OrderedDict
|
||||
from collections.abc import Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
|
||||
from tools.docs_codegen.errors import make_docs_codegen_error
|
||||
from tools.docs_codegen.scanner import ModelCodeBlock
|
||||
from tools.docs_codegen.utils import (
|
||||
ScalarValue,
|
||||
parse_command_tokens,
|
||||
render_cli_command,
|
||||
require_indexed_mapping,
|
||||
require_mapping,
|
||||
require_mapping_list,
|
||||
require_node_field,
|
||||
require_non_empty_string,
|
||||
require_scalar_mapping,
|
||||
substitute_template_positionals,
|
||||
trim_blank_edges,
|
||||
)
|
||||
from tools.docs_codegen.yaml_loader import LoadedYaml
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class GeneratedScript:
|
||||
"""A converter output ready to be persisted as an artifact."""
|
||||
|
||||
content: str
|
||||
language: str = "shell"
|
||||
|
||||
|
||||
class BaseConverter(ABC):
|
||||
"""Minimal contract shared by all converter plugins."""
|
||||
|
||||
name: str
|
||||
|
||||
@abstractmethod
|
||||
def convert(self, loaded_yaml: LoadedYaml, *, block: ModelCodeBlock) -> GeneratedScript:
|
||||
"""Convert one loaded YAML document into one generated artifact."""
|
||||
|
||||
|
||||
def build_default_converters() -> dict[str, BaseConverter]:
|
||||
"""Instantiate the built-in converters keyed by their ``converter_tag`` name."""
|
||||
converters: dict[str, BaseConverter] = {}
|
||||
for converter in (
|
||||
SingleNodeConverter(),
|
||||
MultiNodeConverter(),
|
||||
ExternalDpTemplateConverter(),
|
||||
ExternalDpLaunchConverter(),
|
||||
ExternalDpProxyConverter(),
|
||||
):
|
||||
converters[converter.name] = converter
|
||||
return converters
|
||||
|
||||
|
||||
def get_converter(
|
||||
converters: Mapping[str, BaseConverter],
|
||||
tag: str,
|
||||
*,
|
||||
block: ModelCodeBlock | None = None,
|
||||
) -> BaseConverter:
|
||||
"""Look up a converter by ``converter_tag``, raising a helpful error if unknown."""
|
||||
converter = converters.get(tag)
|
||||
if converter is None:
|
||||
supported = ", ".join(sorted(converters))
|
||||
raise make_docs_codegen_error(
|
||||
f"converter_tag '{tag}' is not registered; supported converters: {supported}",
|
||||
block=block,
|
||||
converter_tag=tag,
|
||||
)
|
||||
return converter
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Shell Rendering Helpers
|
||||
# ============================================================================
|
||||
|
||||
|
||||
def _join_shell_sections(*sections: Sequence[str]) -> str:
|
||||
"""Concatenate line groups, trimming each and separating them with one blank line."""
|
||||
rendered_lines: list[str] = []
|
||||
for section in sections:
|
||||
normalized = trim_blank_edges(section)
|
||||
if not normalized:
|
||||
continue
|
||||
if rendered_lines and rendered_lines[-1] != "":
|
||||
rendered_lines.append("")
|
||||
rendered_lines.extend(normalized)
|
||||
return "\n".join(rendered_lines).rstrip() + "\n"
|
||||
|
||||
|
||||
def _render_env_export_lines(
|
||||
envs: Mapping[str, ScalarValue],
|
||||
*,
|
||||
defaults: Mapping[str, ScalarValue] | None = None,
|
||||
) -> list[str]:
|
||||
"""Render ``envs`` (with optional ``defaults`` overrides) as ``export NAME=value`` lines."""
|
||||
# Keys are already normalized to ``str`` by require_scalar_mapping upstream.
|
||||
exports: OrderedDict[str, ScalarValue] = OrderedDict(envs)
|
||||
if defaults is not None:
|
||||
exports.update(defaults)
|
||||
return [f"export {name}={_quote_env_value(value)}" for name, value in exports.items()]
|
||||
|
||||
|
||||
def _format_vllm_serve_command(tokens: Sequence[str], *, block: ModelCodeBlock) -> list[str]:
|
||||
"""Render ``vllm serve <model> ...`` as backslash-continued, one-option-per-line shell."""
|
||||
if len(tokens) < 3 or tokens[0] != "vllm" or tokens[1] != "serve":
|
||||
raise make_docs_codegen_error(
|
||||
"generated command must start with 'vllm serve <model>'",
|
||||
block=block,
|
||||
)
|
||||
|
||||
model = _quote_cli_arg(tokens[2])
|
||||
command_lines = [f"vllm serve {model}"]
|
||||
option_lines: list[str] = []
|
||||
token_index = 3
|
||||
|
||||
while token_index < len(tokens):
|
||||
token = tokens[token_index]
|
||||
if not token.startswith("-"):
|
||||
raise make_docs_codegen_error(
|
||||
f"generated command contains an unsupported positional argument '{token}'",
|
||||
block=block,
|
||||
)
|
||||
|
||||
if token_index + 1 < len(tokens) and not tokens[token_index + 1].startswith("-"):
|
||||
value = tokens[token_index + 1]
|
||||
if token == "--kv-transfer-config":
|
||||
stripped = value.strip()
|
||||
if stripped.startswith(("{", "[")) and stripped.endswith(("}", "]")):
|
||||
try:
|
||||
parsed = json.loads(stripped)
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
else:
|
||||
if isinstance(parsed, (dict, list)):
|
||||
value = json.dumps(parsed, indent=4, ensure_ascii=False)
|
||||
|
||||
option_lines.append(f"{token} {_quote_cli_arg(value)}")
|
||||
token_index += 2
|
||||
continue
|
||||
|
||||
option_lines.append(token)
|
||||
token_index += 1
|
||||
|
||||
if not option_lines:
|
||||
return command_lines
|
||||
|
||||
command_lines[0] = command_lines[0] + " \\"
|
||||
for index, line in enumerate(option_lines):
|
||||
suffix = " \\" if index < len(option_lines) - 1 else ""
|
||||
indented_line = line.replace("\n", "\n ")
|
||||
command_lines.append(f" {indented_line}{suffix}")
|
||||
|
||||
return command_lines
|
||||
|
||||
|
||||
def _build_shell_script(
|
||||
envs: Mapping[str, ScalarValue],
|
||||
command_tokens: Sequence[str],
|
||||
*,
|
||||
block: ModelCodeBlock,
|
||||
env_defaults: Mapping[str, ScalarValue] | None = None,
|
||||
) -> GeneratedScript:
|
||||
"""Assemble a script from env exports followed by the ``vllm serve`` command."""
|
||||
content = _join_shell_sections(
|
||||
_render_env_export_lines(envs, defaults=env_defaults),
|
||||
_format_vllm_serve_command(command_tokens, block=block),
|
||||
)
|
||||
return GeneratedScript(content=content)
|
||||
|
||||
|
||||
# docs_codegen emits a *copy-pasteable* script, so unlike the e2e runtime's
|
||||
# format_server_cmd() (tests/e2e/nightly/multi_node/external_dp/scripts/utils.py),
|
||||
# which shlex-quotes everything for a one-off *log* line, we need two
|
||||
# context-specific quoters that both keep ``$VAR`` / ``${VAR}`` / ``$1`` as live
|
||||
# shell expansions the reader can still edit.
|
||||
|
||||
|
||||
def _quote_env_value(value: ScalarValue) -> str:
|
||||
"""Quote a value for an ``export NAME=value`` line.
|
||||
|
||||
Wraps in *double* quotes (which still expand ``$``-references) only when the
|
||||
value carries whitespace or shell metacharacters; plain values and bare
|
||||
``$``-expansions are emitted unquoted.
|
||||
"""
|
||||
if value is None:
|
||||
return ""
|
||||
if isinstance(value, str):
|
||||
text = value
|
||||
else:
|
||||
text = str(value)
|
||||
|
||||
needs_quote = text != "" and (
|
||||
any(char.isspace() for char in text) or any(char in text for char in "'\"`;|&<>*?[]{}")
|
||||
)
|
||||
if not needs_quote:
|
||||
return text
|
||||
escaped = text.replace("\\", "\\\\").replace('"', '\\"').replace("`", "\\`")
|
||||
return f'"{escaped}"'
|
||||
|
||||
|
||||
def _quote_cli_arg(token: str) -> str:
|
||||
"""Quote a single ``vllm serve`` argument token.
|
||||
|
||||
Uses ``shlex.quote`` (single quotes ⇒ fully literal) for whitespace, embedded
|
||||
double quotes (JSON), and JSON-like ``{...}`` / ``[...]`` containers so a
|
||||
space-free ``--profiler-config {"a":"b"}`` value is not mangled by the shell.
|
||||
Plain shell expansions like ``$SERVER_PORT`` / ``${NODE_0_IP}`` start with
|
||||
``$`` and are intentionally left unquoted so they stay live.
|
||||
"""
|
||||
if not token:
|
||||
return '""'
|
||||
needs_quote = any(char.isspace() for char in token) or '"' in token or (token[:1] in "{[" and token[-1:] in "}]")
|
||||
if needs_quote:
|
||||
return shlex.quote(token)
|
||||
return token
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Single Node Converter
|
||||
# ============================================================================
|
||||
|
||||
SINGLE_NODE_DEFAULT_SERVER_PORT = "8000"
|
||||
SINGLE_NODE_AUTO_SERVER_PORT = "DEFAULT_PORT"
|
||||
|
||||
|
||||
def _resolve_single_node_server_port(envs: Mapping[str, ScalarValue]) -> ScalarValue:
|
||||
"""Pick the SERVER_PORT export value, mapping the ``DEFAULT_PORT`` sentinel to ``8000``."""
|
||||
server_port = envs.get("SERVER_PORT")
|
||||
if server_port is None or server_port == SINGLE_NODE_AUTO_SERVER_PORT:
|
||||
return SINGLE_NODE_DEFAULT_SERVER_PORT
|
||||
return server_port
|
||||
|
||||
|
||||
def _convert_single_node_case(
|
||||
loaded_yaml: LoadedYaml,
|
||||
*,
|
||||
block: ModelCodeBlock,
|
||||
) -> GeneratedScript:
|
||||
"""Render ``test_cases[case_index]`` into env exports + a ``vllm serve`` command."""
|
||||
test_case = require_indexed_mapping(
|
||||
loaded_yaml.yaml_root,
|
||||
collection_name="test_cases",
|
||||
option_name="case_index",
|
||||
block=block,
|
||||
default_index=0,
|
||||
)
|
||||
envs = require_scalar_mapping(test_case.get("envs"), field_name="envs", block=block)
|
||||
model = require_non_empty_string(test_case.get("model"), field_name="model", block=block)
|
||||
server_cmd = parse_command_tokens(test_case.get("server_cmd"), field_name="server_cmd", block=block)
|
||||
server_cmd_extra = []
|
||||
if test_case.get("server_cmd_extra") is not None:
|
||||
server_cmd_extra = parse_command_tokens(
|
||||
test_case.get("server_cmd_extra"),
|
||||
field_name="server_cmd_extra",
|
||||
block=block,
|
||||
)
|
||||
return _build_shell_script(
|
||||
envs,
|
||||
["vllm", "serve", model, *server_cmd, *server_cmd_extra],
|
||||
block=block,
|
||||
env_defaults={"SERVER_PORT": _resolve_single_node_server_port(envs)},
|
||||
)
|
||||
|
||||
|
||||
class SingleNodeConverter(BaseConverter):
|
||||
"""Render a single-node ``vllm serve`` script from ``test_cases[case_index]``."""
|
||||
|
||||
name = "single_node"
|
||||
|
||||
def convert(self, loaded_yaml: LoadedYaml, *, block: ModelCodeBlock) -> GeneratedScript:
|
||||
return _convert_single_node_case(loaded_yaml, block=block)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Multi Node Converter
|
||||
# ============================================================================
|
||||
|
||||
|
||||
def _convert_multi_node_host(
|
||||
loaded_yaml: LoadedYaml,
|
||||
*,
|
||||
block: ModelCodeBlock,
|
||||
) -> GeneratedScript:
|
||||
"""Render ``deployment[host_index]`` into env exports + its complete ``vllm serve`` command."""
|
||||
deployment_item = require_indexed_mapping(
|
||||
loaded_yaml.yaml_root,
|
||||
collection_name="deployment",
|
||||
option_name="host_index",
|
||||
block=block,
|
||||
)
|
||||
envs = require_scalar_mapping(deployment_item.get("envs"), field_name="envs", block=block)
|
||||
server_cmd = parse_command_tokens(deployment_item.get("server_cmd"), field_name="server_cmd", block=block)
|
||||
return _build_shell_script(envs, server_cmd, block=block)
|
||||
|
||||
|
||||
class MultiNodeConverter(BaseConverter):
|
||||
"""Render one host's ``vllm serve`` script from ``deployment[host_index]``."""
|
||||
|
||||
name = "multi_node"
|
||||
|
||||
def convert(self, loaded_yaml: LoadedYaml, *, block: ModelCodeBlock) -> GeneratedScript:
|
||||
return _convert_multi_node_host(loaded_yaml, block=block)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# External DP Converters
|
||||
#
|
||||
# These read the external-DP YAML schema directly (``model`` / ``routing`` /
|
||||
# ``config`` / ``templates``) used by
|
||||
# tests/e2e/nightly/multi_node/external_dp/config/*.yaml. They are tightly
|
||||
# coupled to that schema by design.
|
||||
# ============================================================================
|
||||
|
||||
LAUNCH_ONLINE_DP_SCRIPT = "launch_online_dp.py"
|
||||
PROXY_SCRIPT = "load_balance_proxy_server_example.py"
|
||||
ROUTING_DISAGGREGATED_PREFILL = "disaggregated_prefill"
|
||||
|
||||
# Mirror tests/e2e/nightly/multi_node/external_dp/scripts/external_dp_config.py
|
||||
# (proxy runs on node 0, port 1999); these are not part of the YAML.
|
||||
EXTERNAL_DP_PROXY_NODE_INDEX = 0
|
||||
EXTERNAL_DP_PROXY_PORT = 1999
|
||||
|
||||
# Maps external-DP ``${VAR}`` template variables to the positional shell
|
||||
# parameters that ``launch_online_dp.py`` forwards to ``run_dp_template.sh``
|
||||
# (``$1=visible_devices`` ... ``$7=tp_size``). Used so generated template
|
||||
# snippets read like the hand-written ``run_dp_template.sh`` instead of leaking
|
||||
# raw ``${DP_SIZE}`` placeholders.
|
||||
RUN_DP_TEMPLATE_POSITIONALS: dict[str, str] = {
|
||||
"VISIBLE_DEVICES": "$1",
|
||||
"PORT": "$2",
|
||||
"DP_SIZE": "$3",
|
||||
"DP_RANK": "$4",
|
||||
"DP_ADDRESS": "$5",
|
||||
"DP_RPC_PORT": "$6",
|
||||
"TP_SIZE": "$7",
|
||||
}
|
||||
|
||||
# Ordered (config[] field, launch_online_dp.py flag) pairs; preserves CLI flag order.
|
||||
LAUNCH_FIELD_FLAGS: tuple[tuple[str, str], ...] = (
|
||||
("dp_size", "--dp-size"),
|
||||
("tp_size", "--tp-size"),
|
||||
("dp_size_local", "--dp-size-local"),
|
||||
("dp_rank_start", "--dp-rank-start"),
|
||||
("dp_address", "--dp-address"),
|
||||
("dp_rpc_port", "--dp-rpc-port"),
|
||||
("port_start", "--vllm-start-port"),
|
||||
)
|
||||
|
||||
|
||||
def _node_ip_placeholder(node_index: int) -> str:
|
||||
"""Return the ``${NODE_<i>_IP}`` shell placeholder for a node index."""
|
||||
return f"${{NODE_{node_index}_IP}}"
|
||||
|
||||
|
||||
# ----------------------------------------------------------------------------
|
||||
# Template converter (per node): env exports + ``vllm serve`` command.
|
||||
# ----------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _convert_external_dp_template(loaded_yaml: LoadedYaml, *, block: ModelCodeBlock) -> GeneratedScript:
|
||||
"""Render ``templates[host_index]`` into per-node env exports + ``vllm serve`` command.
|
||||
|
||||
``${VAR}`` template variables are rewritten to the ``$1``..``$7`` positional
|
||||
parameters that ``run_dp_template.sh`` expects.
|
||||
"""
|
||||
template = require_indexed_mapping(
|
||||
loaded_yaml.yaml_root,
|
||||
collection_name="templates",
|
||||
option_name="host_index",
|
||||
block=block,
|
||||
)
|
||||
model = require_non_empty_string(loaded_yaml.yaml_root.get("model"), field_name="model", block=block)
|
||||
|
||||
raw_envs = require_scalar_mapping(template.get("envs"), field_name="envs", block=block)
|
||||
envs = {
|
||||
key: (
|
||||
substitute_template_positionals(value, positionals=RUN_DP_TEMPLATE_POSITIONALS)
|
||||
if isinstance(value, str)
|
||||
else value
|
||||
)
|
||||
for key, value in raw_envs.items()
|
||||
}
|
||||
|
||||
raw_server_cmd = parse_command_tokens(
|
||||
template.get("server_cmd_template"),
|
||||
field_name="server_cmd_template",
|
||||
block=block,
|
||||
)
|
||||
server_cmd = [
|
||||
substitute_template_positionals(token, positionals=RUN_DP_TEMPLATE_POSITIONALS) for token in raw_server_cmd
|
||||
]
|
||||
|
||||
return _build_shell_script(envs, ["vllm", "serve", model, *server_cmd], block=block)
|
||||
|
||||
|
||||
class ExternalDpTemplateConverter(BaseConverter):
|
||||
"""Render one external-DP node's env exports + ``vllm serve`` command from ``templates``."""
|
||||
|
||||
name = "external_dp_template"
|
||||
|
||||
def convert(self, loaded_yaml: LoadedYaml, *, block: ModelCodeBlock) -> GeneratedScript:
|
||||
return _convert_external_dp_template(loaded_yaml, block=block)
|
||||
|
||||
|
||||
# ----------------------------------------------------------------------------
|
||||
# Launch converter (whole cluster): one ``python launch_online_dp.py`` line per
|
||||
# config node, single-line, separated by a blank line.
|
||||
# ----------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _convert_external_dp_launch(loaded_yaml: LoadedYaml, *, block: ModelCodeBlock) -> GeneratedScript:
|
||||
"""Render one ``python launch_online_dp.py ...`` line per ``config`` node."""
|
||||
nodes = require_mapping_list(loaded_yaml.yaml_root, collection_name="config", block=block, non_empty=True)
|
||||
commands: list[str] = []
|
||||
for node_index, node in enumerate(nodes):
|
||||
options = [
|
||||
(flag, [str(require_node_field(node, field, node_index=node_index, block=block))])
|
||||
for field, flag in LAUNCH_FIELD_FLAGS
|
||||
]
|
||||
commands.append(render_cli_command(["python", LAUNCH_ONLINE_DP_SCRIPT], options, multiline=False).rstrip())
|
||||
return GeneratedScript(content="\n\n".join(commands) + "\n")
|
||||
|
||||
|
||||
class ExternalDpLaunchConverter(BaseConverter):
|
||||
"""Render the cluster-wide ``launch_online_dp.py`` commands, one per ``config`` node."""
|
||||
|
||||
name = "external_dp_launch"
|
||||
|
||||
def convert(self, loaded_yaml: LoadedYaml, *, block: ModelCodeBlock) -> GeneratedScript:
|
||||
return _convert_external_dp_launch(loaded_yaml, block=block)
|
||||
|
||||
|
||||
# ----------------------------------------------------------------------------
|
||||
# Proxy converter (whole cluster): the load-balance proxy launch command.
|
||||
# ----------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _expand_proxy_group(
|
||||
indices: object,
|
||||
nodes: list[dict],
|
||||
*,
|
||||
group_name: str,
|
||||
block: ModelCodeBlock,
|
||||
) -> tuple[list[str], list[str]]:
|
||||
"""Expand a routing group's node indices into per-rank ``(hosts, ports)`` lists."""
|
||||
if not isinstance(indices, list) or not indices:
|
||||
raise make_docs_codegen_error(
|
||||
f"routing.groups.{group_name} must be a non-empty list",
|
||||
block=block,
|
||||
)
|
||||
hosts: list[str] = []
|
||||
ports: list[str] = []
|
||||
for raw_index in indices:
|
||||
node_index = int(raw_index)
|
||||
if node_index < 0 or node_index >= len(nodes):
|
||||
raise make_docs_codegen_error(
|
||||
f"routing.groups.{group_name} index {node_index} is out of range for 'config' with {len(nodes)} items",
|
||||
block=block,
|
||||
)
|
||||
node = nodes[node_index]
|
||||
dp_size_local = int(require_node_field(node, "dp_size_local", node_index=node_index, block=block))
|
||||
port_start = int(require_node_field(node, "port_start", node_index=node_index, block=block))
|
||||
for local_rank in range(dp_size_local):
|
||||
hosts.append(_node_ip_placeholder(node_index))
|
||||
ports.append(str(port_start + local_rank))
|
||||
return hosts, ports
|
||||
|
||||
|
||||
def _convert_external_dp_proxy(loaded_yaml: LoadedYaml, *, block: ModelCodeBlock) -> GeneratedScript:
|
||||
"""Render the load-balance proxy command from the ``routing`` groups."""
|
||||
nodes = require_mapping_list(loaded_yaml.yaml_root, collection_name="config", block=block, non_empty=True)
|
||||
routing = require_mapping(loaded_yaml.yaml_root.get("routing"), field_name="routing", block=block)
|
||||
|
||||
routing_type = routing.get("type")
|
||||
if routing_type != ROUTING_DISAGGREGATED_PREFILL:
|
||||
raise make_docs_codegen_error(
|
||||
f"converter_tag 'external_dp_proxy' only supports routing.type "
|
||||
f"'{ROUTING_DISAGGREGATED_PREFILL}', got {routing_type!r}",
|
||||
block=block,
|
||||
)
|
||||
|
||||
groups = require_mapping(routing.get("groups"), field_name="routing.groups", block=block)
|
||||
prefiller_hosts, prefiller_ports = _expand_proxy_group(
|
||||
groups.get("prefiller"), nodes, group_name="prefiller", block=block
|
||||
)
|
||||
decoder_hosts, decoder_ports = _expand_proxy_group(groups.get("decoder"), nodes, group_name="decoder", block=block)
|
||||
|
||||
options = [
|
||||
("--host", [_node_ip_placeholder(EXTERNAL_DP_PROXY_NODE_INDEX)]),
|
||||
("--port", [str(EXTERNAL_DP_PROXY_PORT)]),
|
||||
("--prefiller-hosts", prefiller_hosts),
|
||||
("--prefiller-ports", prefiller_ports),
|
||||
("--decoder-hosts", decoder_hosts),
|
||||
("--decoder-ports", decoder_ports),
|
||||
]
|
||||
content = render_cli_command(["python", PROXY_SCRIPT], options, multiline=True, expand_values=True)
|
||||
return GeneratedScript(content=content)
|
||||
|
||||
|
||||
class ExternalDpProxyConverter(BaseConverter):
|
||||
"""Render the disaggregated-prefill load-balance proxy launch command."""
|
||||
|
||||
name = "external_dp_proxy"
|
||||
|
||||
def convert(self, loaded_yaml: LoadedYaml, *, block: ModelCodeBlock) -> GeneratedScript:
|
||||
return _convert_external_dp_proxy(loaded_yaml, block=block)
|
||||
72
tools/docs_codegen/errors.py
Normal file
72
tools/docs_codegen/errors.py
Normal file
@@ -0,0 +1,72 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
|
||||
class DocsCodegenError(RuntimeError):
|
||||
"""Raised when docs code generation fails."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
message: str,
|
||||
*,
|
||||
doc_path: Path | None = None,
|
||||
line: int | None = None,
|
||||
test_case_path: str | None = None,
|
||||
block_name: str | None = None,
|
||||
converter_tag: str | None = None,
|
||||
) -> None:
|
||||
super().__init__(message)
|
||||
self.message = message
|
||||
self.doc_path = doc_path
|
||||
self.line = line
|
||||
self.test_case_path = test_case_path
|
||||
self.block_name = block_name
|
||||
self.converter_tag = converter_tag
|
||||
|
||||
def __str__(self) -> str:
|
||||
header = "model-code generation error"
|
||||
if self.doc_path is not None:
|
||||
header = self.doc_path.as_posix()
|
||||
if self.line is not None:
|
||||
header = f"{header}:{self.line}"
|
||||
header = f"{header}: model-code generation error"
|
||||
|
||||
lines = [header]
|
||||
if self.block_name:
|
||||
lines.append(f" block_name: {self.block_name}")
|
||||
if self.test_case_path:
|
||||
lines.append(f" test_case_path: {self.test_case_path}")
|
||||
if self.converter_tag:
|
||||
lines.append(f" converter_tag: {self.converter_tag}")
|
||||
lines.append(f" detail: {self.message}")
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
def make_docs_codegen_error(
|
||||
message: str,
|
||||
*,
|
||||
block: Any | None = None,
|
||||
doc_path: Path | None = None,
|
||||
line: int | None = None,
|
||||
test_case_path: str | None = None,
|
||||
block_name: str | None = None,
|
||||
converter_tag: str | None = None,
|
||||
) -> DocsCodegenError:
|
||||
"""Build a ``DocsCodegenError``, pulling location context off ``block`` when given."""
|
||||
if block is not None:
|
||||
doc_path = getattr(block, "doc_path", doc_path)
|
||||
line = getattr(block, "directive_line", line)
|
||||
test_case_path = getattr(block, "test_case_path", test_case_path)
|
||||
block_name = getattr(block, "block_name", block_name)
|
||||
converter_tag = getattr(block, "converter_tag", converter_tag)
|
||||
|
||||
return DocsCodegenError(
|
||||
message,
|
||||
doc_path=doc_path,
|
||||
line=line,
|
||||
test_case_path=test_case_path,
|
||||
block_name=block_name,
|
||||
converter_tag=converter_tag,
|
||||
)
|
||||
140
tools/docs_codegen/generator.py
Normal file
140
tools/docs_codegen/generator.py
Normal file
@@ -0,0 +1,140 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
from tools.docs_codegen.converters import BaseConverter, GeneratedScript, build_default_converters, get_converter
|
||||
from tools.docs_codegen.errors import make_docs_codegen_error
|
||||
from tools.docs_codegen.scanner import BlockScanner, ModelCodeBlock
|
||||
from tools.docs_codegen.yaml_loader import YamlLoader
|
||||
|
||||
DEFAULT_ARTIFACT_ROOT = Path("docs/_build/doc_codegen")
|
||||
GENERATED_SCRIPT_MARKER = "{{ generated }}"
|
||||
|
||||
# One generated artifact: its repo-relative output path and the rendered script.
|
||||
GeneratedArtifact = tuple[Path, GeneratedScript]
|
||||
|
||||
|
||||
class GeneratorService:
|
||||
"""Shared generation pipeline used by both CLI and Sphinx."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
block_scanner: BlockScanner | None = None,
|
||||
yaml_loader: YamlLoader | None = None,
|
||||
converters: dict[str, BaseConverter] | None = None,
|
||||
artifact_root: str | Path = DEFAULT_ARTIFACT_ROOT,
|
||||
repo_root: str | Path | None = None,
|
||||
) -> None:
|
||||
self.repo_root = Path(repo_root) if repo_root is not None else None
|
||||
self.block_scanner = block_scanner or BlockScanner(repo_root=self.repo_root)
|
||||
self.yaml_loader = yaml_loader or YamlLoader(repo_root=self.repo_root)
|
||||
self.converters = converters or build_default_converters()
|
||||
self.artifact_root = Path(artifact_root)
|
||||
|
||||
# -- Public API ----------------------------------------------------------
|
||||
|
||||
def generate_all(self, *, dry_run: bool = False) -> list[GeneratedArtifact]:
|
||||
"""Generate artifacts for every model-code block under the documents root."""
|
||||
return self._generate_blocks(self.block_scanner.scan_default_blocks(), dry_run=dry_run)
|
||||
|
||||
def generate_document(self, doc_path: str | Path, *, dry_run: bool = False) -> list[GeneratedArtifact]:
|
||||
"""Generate artifacts for every model-code block in one document."""
|
||||
return self._generate_blocks(self.block_scanner.scan_document_blocks(doc_path), dry_run=dry_run)
|
||||
|
||||
def generate_block(
|
||||
self,
|
||||
doc_path: str | Path,
|
||||
block_name: str,
|
||||
*,
|
||||
dry_run: bool = False,
|
||||
) -> GeneratedArtifact:
|
||||
"""Generate the artifact for a single named block in a document."""
|
||||
generated_artifacts = self._generate_blocks(
|
||||
self.block_scanner.select_document_blocks(doc_path, block_name),
|
||||
dry_run=dry_run,
|
||||
)
|
||||
return generated_artifacts[0]
|
||||
|
||||
def read_generated_script(self, block: ModelCodeBlock) -> GeneratedScript:
|
||||
"""Read a previously generated artifact from disk (used by the Sphinx directive)."""
|
||||
output_path = self.output_path_for(block)
|
||||
absolute_output_path = self._absolute_output_path(output_path)
|
||||
if not absolute_output_path.exists():
|
||||
raise make_docs_codegen_error(
|
||||
f"generated artifact not found: {output_path}",
|
||||
block=block,
|
||||
)
|
||||
return GeneratedScript(content=absolute_output_path.read_text(encoding="utf-8"))
|
||||
|
||||
def output_path_for(self, block: ModelCodeBlock) -> Path:
|
||||
"""Repo-relative artifact path: ``<artifact_root>/<doc_stem>/<block_name>.sh``."""
|
||||
return self.artifact_root / block.doc_path.stem / f"{block.block_name}.sh"
|
||||
|
||||
# -- Generation pipeline -------------------------------------------------
|
||||
|
||||
def _generate_blocks(self, blocks: list[ModelCodeBlock], *, dry_run: bool) -> list[GeneratedArtifact]:
|
||||
"""Convert each block to a script, merge any raw body, and (unless dry-run) write it."""
|
||||
generated_artifacts: list[GeneratedArtifact] = []
|
||||
for block in blocks:
|
||||
converter = get_converter(self.converters, block.converter_tag, block=block)
|
||||
loaded_yaml = self.yaml_loader.load(
|
||||
test_case_path=block.test_case_path,
|
||||
block=block,
|
||||
)
|
||||
generated_script = converter.convert(loaded_yaml, block=block)
|
||||
generated_script = self._apply_block_body(generated_script, block=block)
|
||||
self._validate_generated_script(generated_script, block=block)
|
||||
output_path = self.output_path_for(block)
|
||||
|
||||
if not dry_run:
|
||||
self._write_script(self._absolute_output_path(output_path), generated_script)
|
||||
generated_artifacts.append((output_path, generated_script))
|
||||
return generated_artifacts
|
||||
|
||||
@staticmethod
|
||||
def _apply_block_body(
|
||||
generated_script: GeneratedScript,
|
||||
*,
|
||||
block: ModelCodeBlock,
|
||||
) -> GeneratedScript:
|
||||
"""Splice the converter output into the block body (at ``{{ generated }}`` or appended)."""
|
||||
if not block.raw_block_lines:
|
||||
return generated_script
|
||||
|
||||
raw_block_content = "\n".join(block.raw_block_lines)
|
||||
generated_content = generated_script.content.rstrip()
|
||||
if GENERATED_SCRIPT_MARKER in raw_block_content:
|
||||
content = raw_block_content.replace(GENERATED_SCRIPT_MARKER, generated_content)
|
||||
else:
|
||||
content = f"{raw_block_content.rstrip()}\n\n{generated_content}"
|
||||
|
||||
return GeneratedScript(content=content.rstrip() + "\n", language=generated_script.language)
|
||||
|
||||
@staticmethod
|
||||
def _validate_generated_script(generated_script: GeneratedScript, *, block: ModelCodeBlock) -> None:
|
||||
"""Guard against a converter producing an empty artifact."""
|
||||
if not generated_script.content.strip():
|
||||
raise make_docs_codegen_error("generated script content is empty", block=block)
|
||||
|
||||
@staticmethod
|
||||
def _write_script(output_path: Path, generated_script: GeneratedScript) -> None:
|
||||
"""Write an artifact to ``output_path``, creating parent directories."""
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
output_path.write_text(generated_script.content, encoding="utf-8")
|
||||
|
||||
# -- Path helpers --------------------------------------------------------
|
||||
|
||||
@property
|
||||
def _base(self) -> Path:
|
||||
"""Directory that the repo-relative artifact path resolves against for I/O."""
|
||||
return self.repo_root if self.repo_root is not None else Path.cwd()
|
||||
|
||||
def _absolute_output_path(self, output_path: Path) -> Path:
|
||||
"""Anchor a repo-relative output path to ``_base`` for filesystem reads/writes."""
|
||||
return output_path if output_path.is_absolute() else self._base / output_path
|
||||
|
||||
|
||||
def create_default_generator_service(repo_root: str | Path | None = None) -> GeneratorService:
|
||||
"""Build a ``GeneratorService`` with default scanner/loader/converters."""
|
||||
return GeneratorService(repo_root=repo_root)
|
||||
228
tools/docs_codegen/scanner.py
Normal file
228
tools/docs_codegen/scanner.py
Normal file
@@ -0,0 +1,228 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
|
||||
import regex as re
|
||||
|
||||
from tools.docs_codegen.errors import DocsCodegenError, make_docs_codegen_error
|
||||
from tools.docs_codegen.utils import trim_blank_edges
|
||||
|
||||
MODEL_CODE_DEFAULTS_PATH = Path("docs/source/tutorials/models")
|
||||
MODEL_CODE_REQUIRED_OPTION_NAMES = ("block_name", "converter_tag", "test_case_path")
|
||||
MODEL_CODE_OPTION_NAMES = (*MODEL_CODE_REQUIRED_OPTION_NAMES, "case_index", "host_index")
|
||||
MODEL_CODE_OPEN_RE = re.compile(r"^\s*```{model-code}\s*$")
|
||||
MODEL_CODE_CLOSE_RE = re.compile(r"^\s*```\s*$")
|
||||
MODEL_CODE_OPTION_RE = re.compile(r"^\s*:([A-Za-z0-9_-]+):\s*(.*?)\s*$")
|
||||
BLOCK_NAME_RE = re.compile(r"^[A-Za-z0-9._-]+$")
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ModelCodeBlock:
|
||||
"""One ``model-code`` block discovered in a documentation page."""
|
||||
|
||||
doc_path: Path
|
||||
block_name: str
|
||||
converter_tag: str
|
||||
test_case_path: str
|
||||
extra_options: tuple[tuple[str, str], ...] = ()
|
||||
directive_line: int | None = None
|
||||
raw_block_lines: tuple[str, ...] = ()
|
||||
|
||||
@property
|
||||
def key(self) -> tuple[str, str]:
|
||||
"""Identity used to detect duplicate blocks: ``(doc_path, block_name)``."""
|
||||
return (self.doc_path.as_posix(), self.block_name)
|
||||
|
||||
def get_option(self, name: str) -> str | None:
|
||||
"""Return the value of an extra (non-required) directive option, or ``None``."""
|
||||
for key, value in self.extra_options:
|
||||
if key == name:
|
||||
return value
|
||||
return None
|
||||
|
||||
|
||||
class BlockScanner:
|
||||
"""Scan markdown files for ``model-code`` directives."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
documents_root: str | Path = MODEL_CODE_DEFAULTS_PATH,
|
||||
repo_root: str | Path | None = None,
|
||||
) -> None:
|
||||
self.documents_root = Path(documents_root)
|
||||
self.repo_root = Path(repo_root) if repo_root is not None else None
|
||||
|
||||
# -- Public API ----------------------------------------------------------
|
||||
|
||||
def scan_default_blocks(self) -> list[ModelCodeBlock]:
|
||||
"""Scan every markdown file under ``documents_root`` for model-code blocks."""
|
||||
base = self._base
|
||||
models_dir = base / self.documents_root
|
||||
if not models_dir.exists():
|
||||
raise make_docs_codegen_error(
|
||||
"tutorials models directory does not exist",
|
||||
doc_path=self.documents_root,
|
||||
)
|
||||
|
||||
blocks: list[ModelCodeBlock] = []
|
||||
for absolute_doc_path in sorted(models_dir.rglob("*.md")):
|
||||
blocks.extend(self.scan_document_blocks(absolute_doc_path.relative_to(base)))
|
||||
return blocks
|
||||
|
||||
def scan_document_blocks(self, doc_path: str | Path) -> list[ModelCodeBlock]:
|
||||
"""Parse all model-code directive fences in a single markdown document."""
|
||||
repo_relative_doc_path = self._normalize_document_path(doc_path)
|
||||
absolute_doc_path = self._base / repo_relative_doc_path
|
||||
if not absolute_doc_path.exists():
|
||||
raise make_docs_codegen_error("document file does not exist", doc_path=repo_relative_doc_path)
|
||||
|
||||
lines = absolute_doc_path.read_text(encoding="utf-8").splitlines()
|
||||
blocks: list[ModelCodeBlock] = []
|
||||
line_index = 0
|
||||
|
||||
while line_index < len(lines):
|
||||
if not MODEL_CODE_OPEN_RE.match(lines[line_index]):
|
||||
line_index += 1
|
||||
continue
|
||||
|
||||
directive_line = line_index + 1
|
||||
line_index += 1
|
||||
options: dict[str, str] = {}
|
||||
body_lines: list[str] = []
|
||||
in_body = False
|
||||
|
||||
while line_index < len(lines):
|
||||
line = lines[line_index]
|
||||
if MODEL_CODE_CLOSE_RE.match(line):
|
||||
break
|
||||
|
||||
option_match = MODEL_CODE_OPTION_RE.match(line)
|
||||
if not in_body and option_match:
|
||||
options[option_match.group(1)] = option_match.group(2).strip()
|
||||
else:
|
||||
in_body = True
|
||||
body_lines.append(line)
|
||||
line_index += 1
|
||||
|
||||
if line_index >= len(lines) or not MODEL_CODE_CLOSE_RE.match(lines[line_index]):
|
||||
raise make_docs_codegen_error(
|
||||
"unclosed model-code directive fence",
|
||||
doc_path=repo_relative_doc_path,
|
||||
line=directive_line,
|
||||
)
|
||||
|
||||
blocks.append(
|
||||
self.build_block(
|
||||
options,
|
||||
doc_path=repo_relative_doc_path,
|
||||
directive_line=directive_line,
|
||||
body_lines=body_lines,
|
||||
)
|
||||
)
|
||||
line_index += 1
|
||||
|
||||
self._validate_unique_block_names(blocks)
|
||||
return blocks
|
||||
|
||||
def select_document_blocks(self, doc_path: str | Path, block_name: str | None = None) -> list[ModelCodeBlock]:
|
||||
"""Scan a document and optionally keep only the block named ``block_name``."""
|
||||
blocks = self.scan_document_blocks(doc_path)
|
||||
if block_name is None:
|
||||
return blocks
|
||||
|
||||
selected_blocks = [block for block in blocks if block.block_name == block_name]
|
||||
if not selected_blocks:
|
||||
raise make_docs_codegen_error(
|
||||
f"block_name '{block_name}' not found in document",
|
||||
doc_path=self._normalize_document_path(doc_path),
|
||||
)
|
||||
return selected_blocks
|
||||
|
||||
def build_block(
|
||||
self,
|
||||
options: Mapping[str, str],
|
||||
*,
|
||||
doc_path: str | Path,
|
||||
directive_line: int | None = None,
|
||||
body_lines: Sequence[str] = (),
|
||||
) -> ModelCodeBlock:
|
||||
"""Validate directive options and assemble a ``ModelCodeBlock``."""
|
||||
repo_relative_doc_path = self._normalize_document_path(doc_path)
|
||||
|
||||
def fail(message: str) -> DocsCodegenError:
|
||||
return make_docs_codegen_error(
|
||||
message,
|
||||
doc_path=repo_relative_doc_path,
|
||||
line=directive_line,
|
||||
test_case_path=options.get("test_case_path"),
|
||||
block_name=options.get("block_name"),
|
||||
converter_tag=options.get("converter_tag"),
|
||||
)
|
||||
|
||||
missing = [name for name in MODEL_CODE_REQUIRED_OPTION_NAMES if name not in options]
|
||||
if missing:
|
||||
raise fail(f"model-code block missing required metadata: {', '.join(missing)}")
|
||||
|
||||
extra = sorted(set(options).difference(MODEL_CODE_OPTION_NAMES))
|
||||
if extra:
|
||||
raise fail(f"model-code block contains unsupported metadata: {', '.join(extra)}")
|
||||
|
||||
normalized_options = {name: options[name].strip() for name in MODEL_CODE_OPTION_NAMES if name in options}
|
||||
empty_values = [name for name, value in normalized_options.items() if not value]
|
||||
if empty_values:
|
||||
raise fail(f"model-code block contains empty metadata: {', '.join(empty_values)}")
|
||||
|
||||
if not BLOCK_NAME_RE.fullmatch(normalized_options["block_name"]):
|
||||
raise fail("block_name may only contain letters, numbers, dots, underscores, and dashes")
|
||||
|
||||
return ModelCodeBlock(
|
||||
doc_path=repo_relative_doc_path,
|
||||
block_name=normalized_options["block_name"],
|
||||
converter_tag=normalized_options["converter_tag"],
|
||||
test_case_path=normalized_options["test_case_path"],
|
||||
extra_options=tuple(
|
||||
(name, normalized_options[name])
|
||||
for name in MODEL_CODE_OPTION_NAMES
|
||||
if name not in MODEL_CODE_REQUIRED_OPTION_NAMES and name in normalized_options
|
||||
),
|
||||
directive_line=directive_line,
|
||||
raw_block_lines=tuple(trim_blank_edges(body_lines)),
|
||||
)
|
||||
|
||||
# -- Internal helpers ----------------------------------------------------
|
||||
|
||||
@staticmethod
|
||||
def _normalize_document_path(doc_path: str | Path) -> Path:
|
||||
"""Validate that a document path is repository-relative and contained."""
|
||||
candidate = Path(doc_path)
|
||||
if candidate.is_absolute():
|
||||
raise make_docs_codegen_error("document path must be repository-relative", doc_path=candidate)
|
||||
if ".." in candidate.parts:
|
||||
raise make_docs_codegen_error("document path must stay within the repository", doc_path=candidate)
|
||||
return candidate
|
||||
|
||||
@staticmethod
|
||||
def _validate_unique_block_names(blocks: Sequence[ModelCodeBlock]) -> None:
|
||||
"""Reject documents that declare the same block_name twice."""
|
||||
seen: dict[tuple[str, str], ModelCodeBlock] = {}
|
||||
for block in blocks:
|
||||
previous = seen.get(block.key)
|
||||
if previous is None:
|
||||
seen[block.key] = block
|
||||
continue
|
||||
|
||||
raise make_docs_codegen_error(
|
||||
"duplicated block_name "
|
||||
f"'{block.block_name}' in document; previous declaration is on line {previous.directive_line}",
|
||||
block=block,
|
||||
)
|
||||
|
||||
# -- Path helpers --------------------------------------------------------
|
||||
|
||||
@property
|
||||
def _base(self) -> Path:
|
||||
"""Directory that repo-relative paths resolve against for filesystem I/O."""
|
||||
return self.repo_root if self.repo_root is not None else Path.cwd()
|
||||
93
tools/docs_codegen/sphinx_extension.py
Normal file
93
tools/docs_codegen/sphinx_extension.py
Normal file
@@ -0,0 +1,93 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Mapping
|
||||
from pathlib import Path
|
||||
|
||||
from docutils import nodes
|
||||
from docutils.parsers.rst import directives
|
||||
from sphinx.errors import SphinxError
|
||||
from sphinx.util.docutils import SphinxDirective
|
||||
|
||||
from tools.docs_codegen.errors import DocsCodegenError
|
||||
from tools.docs_codegen.generator import GeneratorService, create_default_generator_service
|
||||
from tools.docs_codegen.scanner import BlockScanner, ModelCodeBlock
|
||||
|
||||
# Anchor all repo-relative paths here instead of relying on the process CWD: the
|
||||
# docs are built from ``docs/`` (see docs/Makefile, SOURCEDIR=source), so the
|
||||
# generator must resolve paths against the repo root regardless of where
|
||||
# sphinx-build was launched.
|
||||
REPO_ROOT = Path(__file__).resolve().parents[2]
|
||||
|
||||
|
||||
def build_block_from_options(
|
||||
*,
|
||||
doc_path: Path,
|
||||
options: Mapping[str, str],
|
||||
directive_line: int | None = None,
|
||||
body_lines: list[str] | None = None,
|
||||
block_scanner: BlockScanner | None = None,
|
||||
) -> ModelCodeBlock:
|
||||
"""Build a ``ModelCodeBlock`` from directive options (no filesystem scan)."""
|
||||
scanner = block_scanner or BlockScanner(repo_root=REPO_ROOT)
|
||||
return scanner.build_block(options, doc_path=doc_path, directive_line=directive_line, body_lines=body_lines or ())
|
||||
|
||||
|
||||
def render_generated_script(
|
||||
block: ModelCodeBlock,
|
||||
*,
|
||||
service: GeneratorService | None = None,
|
||||
) -> nodes.literal_block:
|
||||
"""Read the pre-generated artifact for a block and wrap it in a docutils literal block."""
|
||||
generator_service = service or create_default_generator_service(repo_root=REPO_ROOT)
|
||||
script = generator_service.read_generated_script(block)
|
||||
literal = nodes.literal_block(script.content, script.content)
|
||||
literal["language"] = script.language
|
||||
return literal
|
||||
|
||||
|
||||
class ModelCodeDirective(SphinxDirective):
|
||||
"""Import a pre-generated shell script and render it as a code block."""
|
||||
|
||||
has_content = True
|
||||
option_spec = {
|
||||
"block_name": directives.unchanged_required,
|
||||
"converter_tag": directives.unchanged_required,
|
||||
"test_case_path": directives.unchanged_required,
|
||||
"case_index": directives.unchanged,
|
||||
"host_index": directives.unchanged,
|
||||
}
|
||||
|
||||
def run(self) -> list[nodes.Node]:
|
||||
"""Resolve the current document's block and emit its rendered code block."""
|
||||
source_relative_doc_path = Path(self.env.doc2path(self.env.docname, base=False))
|
||||
doc_path = Path("docs/source") / source_relative_doc_path
|
||||
|
||||
try:
|
||||
block = build_block_from_options(
|
||||
doc_path=doc_path,
|
||||
options=self.options,
|
||||
directive_line=self.lineno,
|
||||
body_lines=list(self.content),
|
||||
)
|
||||
return [render_generated_script(block)]
|
||||
except DocsCodegenError as exc:
|
||||
raise self.error(str(exc)) from exc
|
||||
|
||||
|
||||
def on_builder_inited(app) -> None:
|
||||
"""Sphinx ``builder-inited`` hook: regenerate all artifacts before the build reads them."""
|
||||
del app
|
||||
try:
|
||||
create_default_generator_service(repo_root=REPO_ROOT).generate_all()
|
||||
except DocsCodegenError as exc:
|
||||
raise SphinxError(str(exc)) from exc
|
||||
|
||||
|
||||
def setup(app):
|
||||
"""Sphinx extension entry point: register the directive and the build-init hook."""
|
||||
app.add_directive("model-code", ModelCodeDirective)
|
||||
app.connect("builder-inited", on_builder_inited)
|
||||
return {
|
||||
"parallel_read_safe": True,
|
||||
"parallel_write_safe": True,
|
||||
}
|
||||
230
tools/docs_codegen/utils.py
Normal file
230
tools/docs_codegen/utils.py
Normal file
@@ -0,0 +1,230 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import shlex
|
||||
from collections.abc import Mapping, Sequence
|
||||
from typing import Any, cast
|
||||
|
||||
import regex as re
|
||||
|
||||
from tools.docs_codegen.errors import make_docs_codegen_error
|
||||
|
||||
ScalarValue = str | int | float | bool | None
|
||||
|
||||
# Braced ``${VAR}`` template variables, mirroring runtime.py:TEMPLATE_VAR_RE.
|
||||
TEMPLATE_VAR_RE = re.compile(r"\$\{([A-Za-z_][A-Za-z0-9_]*)\}")
|
||||
|
||||
|
||||
def trim_blank_edges(lines: Sequence[str]) -> list[str]:
|
||||
"""Drop leading and trailing blank/whitespace-only lines."""
|
||||
start = 0
|
||||
end = len(lines)
|
||||
while start < end and not lines[start].strip():
|
||||
start += 1
|
||||
while end > start and not lines[end - 1].strip():
|
||||
end -= 1
|
||||
return list(lines[start:end])
|
||||
|
||||
|
||||
def require_mapping(value: Any, *, field_name: str, block: Any) -> dict[str, Any]:
|
||||
"""Require ``value`` to be a mapping, returning it with string-coerced keys."""
|
||||
if not isinstance(value, dict):
|
||||
raise make_docs_codegen_error(
|
||||
f"converter field '{field_name}' must be a mapping, got {type(value).__name__}",
|
||||
block=block,
|
||||
)
|
||||
return {str(key): item for key, item in value.items()}
|
||||
|
||||
|
||||
def require_mapping_list(
|
||||
yaml_root: Any,
|
||||
*,
|
||||
collection_name: str,
|
||||
block: Any,
|
||||
non_empty: bool = False,
|
||||
) -> list[dict[str, Any]]:
|
||||
"""Require ``yaml_root[collection_name]`` to be a list of mappings.
|
||||
|
||||
Validates that the YAML root is a mapping, that ``collection_name`` holds a
|
||||
list (optionally non-empty), and that every element is itself a mapping.
|
||||
Shared by ``require_indexed_mapping`` and the external-DP converters so the
|
||||
"named YAML list of mappings" pattern lives in one place.
|
||||
"""
|
||||
if not isinstance(yaml_root, dict):
|
||||
raise make_docs_codegen_error(
|
||||
f"YAML root must be a mapping, got {type(yaml_root).__name__}",
|
||||
block=block,
|
||||
)
|
||||
collection = yaml_root.get(collection_name)
|
||||
if not isinstance(collection, list) or (non_empty and not collection):
|
||||
kind = "a non-empty list" if non_empty else "a list"
|
||||
raise make_docs_codegen_error(
|
||||
f"YAML field '{collection_name}' must be {kind}",
|
||||
block=block,
|
||||
)
|
||||
return [
|
||||
require_mapping(item, field_name=f"{collection_name}[{index}]", block=block)
|
||||
for index, item in enumerate(collection)
|
||||
]
|
||||
|
||||
|
||||
def require_non_empty_string(value: Any, *, field_name: str, block: Any) -> str:
|
||||
"""Require ``value`` to be a non-blank string, returning it stripped."""
|
||||
if not isinstance(value, str) or not value.strip():
|
||||
raise make_docs_codegen_error(
|
||||
f"converter field '{field_name}' must be a non-empty string",
|
||||
block=block,
|
||||
)
|
||||
return value.strip()
|
||||
|
||||
|
||||
def require_block_index(
|
||||
*,
|
||||
block: Any,
|
||||
option_name: str,
|
||||
default: int | None = None,
|
||||
) -> int:
|
||||
"""Read a non-negative integer directive option (e.g. ``case_index``) off the block."""
|
||||
raw_index = block.get_option(option_name)
|
||||
if raw_index is None:
|
||||
if default is not None:
|
||||
return default
|
||||
raise make_docs_codegen_error(
|
||||
f"model-code block with converter_tag '{block.converter_tag}' requires {option_name}",
|
||||
block=block,
|
||||
)
|
||||
|
||||
if not raw_index.isdecimal():
|
||||
raise make_docs_codegen_error(
|
||||
f"{option_name} must be a non-negative integer, got '{raw_index}'",
|
||||
block=block,
|
||||
)
|
||||
return int(raw_index)
|
||||
|
||||
|
||||
def require_indexed_mapping(
|
||||
yaml_root: Any,
|
||||
*,
|
||||
collection_name: str,
|
||||
option_name: str,
|
||||
block: Any,
|
||||
default_index: int | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""Pick one mapping out of a YAML list, selected by a block directive option.
|
||||
|
||||
``collection_name`` is the YAML key holding the list (e.g. ``test_cases``) and
|
||||
``option_name`` is the ``model-code`` directive option carrying the index
|
||||
(e.g. ``case_index``).
|
||||
"""
|
||||
index = require_block_index(block=block, option_name=option_name, default=default_index)
|
||||
|
||||
collection = require_mapping_list(yaml_root, collection_name=collection_name, block=block)
|
||||
if index >= len(collection):
|
||||
raise make_docs_codegen_error(
|
||||
f"{option_name} {index} is out of range for '{collection_name}' with {len(collection)} items",
|
||||
block=block,
|
||||
)
|
||||
|
||||
return collection[index]
|
||||
|
||||
|
||||
def require_scalar_mapping(
|
||||
value: Any,
|
||||
*,
|
||||
field_name: str,
|
||||
block: Any,
|
||||
) -> dict[str, ScalarValue]:
|
||||
"""Require a mapping whose values are all scalars (no nested mappings/lists)."""
|
||||
mapping = require_mapping(value, field_name=field_name, block=block)
|
||||
normalized: dict[str, ScalarValue] = {}
|
||||
for key, item in mapping.items():
|
||||
if isinstance(item, (dict, list)):
|
||||
raise make_docs_codegen_error(
|
||||
f"converter field '{field_name}.{key}' must be a scalar value",
|
||||
block=block,
|
||||
)
|
||||
normalized[str(key)] = cast(ScalarValue, item)
|
||||
return normalized
|
||||
|
||||
|
||||
def require_node_field(node: Mapping[str, object], field: str, *, node_index: int, block: Any) -> object:
|
||||
"""Return a required ``config[node_index]`` field, erroring if it is missing."""
|
||||
if node.get(field) is None:
|
||||
raise make_docs_codegen_error(
|
||||
f"config[{node_index}] is missing required field '{field}'",
|
||||
block=block,
|
||||
)
|
||||
return node[field]
|
||||
|
||||
|
||||
def parse_command_tokens(value: Any, *, field_name: str, block: Any) -> list[str]:
|
||||
"""Normalize a shell string or flat token list into a list of argument tokens."""
|
||||
if isinstance(value, str):
|
||||
try:
|
||||
return shlex.split(value, posix=True)
|
||||
except ValueError as exc:
|
||||
raise make_docs_codegen_error(
|
||||
f"converter field '{field_name}' contains an invalid shell string: {exc}",
|
||||
block=block,
|
||||
) from exc
|
||||
|
||||
if isinstance(value, list) and all(not isinstance(item, (dict, list)) for item in value):
|
||||
return [str(item) for item in value]
|
||||
|
||||
raise make_docs_codegen_error(
|
||||
f"converter field '{field_name}' must be a shell string or a flat token list",
|
||||
block=block,
|
||||
)
|
||||
|
||||
|
||||
def substitute_template_positionals(
|
||||
value: str,
|
||||
*,
|
||||
positionals: Mapping[str, str],
|
||||
) -> str:
|
||||
"""Replace braced ``${VAR}`` template variables with positional shell params.
|
||||
|
||||
Only keys present in ``positionals`` are replaced; unknown braced variables
|
||||
and unbraced references like ``$SERVER_PORT`` are left untouched.
|
||||
"""
|
||||
|
||||
def repl(match: re.Match[str]) -> str:
|
||||
key = match.group(1)
|
||||
return positionals.get(key, match.group(0))
|
||||
|
||||
return TEMPLATE_VAR_RE.sub(repl, value)
|
||||
|
||||
|
||||
def render_cli_command(
|
||||
prefix: Sequence[str],
|
||||
options: Sequence[tuple[str, Sequence[str]]],
|
||||
*,
|
||||
multiline: bool,
|
||||
expand_values: bool = False,
|
||||
) -> str:
|
||||
"""Render a CLI command from a prefix and ``(flag, values)`` option groups.
|
||||
|
||||
Supports multi-value flags (e.g. ``--prefiller-hosts h1 h2``). With
|
||||
``multiline=False`` the whole command is rendered on one line. With
|
||||
``multiline=True`` each option starts on its own backslash-continued line;
|
||||
when ``expand_values`` is also set, a multi-value flag is placed on its own
|
||||
line followed by each value on its own indented line (single-value flags
|
||||
stay inline). The returned string always ends with a newline.
|
||||
"""
|
||||
prefix_str = " ".join(prefix)
|
||||
|
||||
if not multiline:
|
||||
rendered = [" ".join([flag, *[str(value) for value in values]]) for flag, values in options]
|
||||
return " ".join([prefix_str, *rendered]).rstrip() + "\n"
|
||||
|
||||
# Each entry is a logical line rendered without its trailing backslash.
|
||||
entries: list[str] = [prefix_str]
|
||||
for flag, values in options:
|
||||
str_values = [str(value) for value in values]
|
||||
if expand_values and len(str_values) > 1:
|
||||
entries.append(f" {flag}")
|
||||
entries.extend(f" {value}" for value in str_values)
|
||||
else:
|
||||
entries.append(f" {' '.join([flag, *str_values])}")
|
||||
|
||||
lines = [entry + (" \\" if index < len(entries) - 1 else "") for index, entry in enumerate(entries)]
|
||||
return "\n".join(lines) + "\n"
|
||||
90
tools/docs_codegen/yaml_loader.py
Normal file
90
tools/docs_codegen/yaml_loader.py
Normal file
@@ -0,0 +1,90 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import yaml
|
||||
|
||||
from tools.docs_codegen.errors import make_docs_codegen_error
|
||||
from tools.docs_codegen.scanner import ModelCodeBlock
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class LoadedYaml:
|
||||
"""One loaded YAML document referenced by a ``model-code`` block."""
|
||||
|
||||
yaml_path: Path
|
||||
yaml_root: Any
|
||||
|
||||
|
||||
class YamlLoader:
|
||||
"""Load and cache one repository-relative YAML file."""
|
||||
|
||||
def __init__(self, repo_root: str | Path | None = None) -> None:
|
||||
self.repo_root = Path(repo_root) if repo_root is not None else None
|
||||
self._yaml_cache: dict[Path, Any] = {}
|
||||
|
||||
# -- Public API ----------------------------------------------------------
|
||||
|
||||
def load(
|
||||
self,
|
||||
*,
|
||||
test_case_path: str,
|
||||
block: ModelCodeBlock | None = None,
|
||||
) -> LoadedYaml:
|
||||
"""Resolve, parse, and cache the YAML referenced by a model-code block."""
|
||||
yaml_path = self._resolve_test_case_path(test_case_path=test_case_path, block=block)
|
||||
yaml_root = self._load_yaml_root(yaml_path)
|
||||
return LoadedYaml(
|
||||
yaml_path=yaml_path,
|
||||
yaml_root=yaml_root,
|
||||
)
|
||||
|
||||
# -- Resolution & parsing ------------------------------------------------
|
||||
|
||||
def _resolve_test_case_path(self, *, test_case_path: str, block: ModelCodeBlock | None = None) -> Path:
|
||||
"""Resolve a repo-relative ``test_case_path`` to an absolute, contained, existing file."""
|
||||
candidate = Path(test_case_path)
|
||||
if candidate.is_absolute():
|
||||
raise make_docs_codegen_error(
|
||||
"test_case_path must be a repository-relative path",
|
||||
block=block,
|
||||
test_case_path=test_case_path,
|
||||
)
|
||||
|
||||
base = self._base.resolve()
|
||||
resolved = (base / candidate).resolve()
|
||||
if not resolved.is_relative_to(base):
|
||||
raise make_docs_codegen_error(
|
||||
"test_case_path must stay within the repository",
|
||||
block=block,
|
||||
test_case_path=test_case_path,
|
||||
)
|
||||
|
||||
if not resolved.exists():
|
||||
raise make_docs_codegen_error(
|
||||
"test_case_path file does not exist",
|
||||
block=block,
|
||||
test_case_path=test_case_path,
|
||||
)
|
||||
return resolved
|
||||
|
||||
def _load_yaml_root(self, yaml_path: Path) -> Any:
|
||||
"""Return the parsed YAML for ``yaml_path``, caching it on first load."""
|
||||
if yaml_path not in self._yaml_cache:
|
||||
self._yaml_cache[yaml_path] = self._parse_yaml_file(yaml_path)
|
||||
return self._yaml_cache[yaml_path]
|
||||
|
||||
@staticmethod
|
||||
def _parse_yaml_file(yaml_path: Path) -> Any:
|
||||
"""Read and parse one YAML file, treating an empty document as ``{}``."""
|
||||
with yaml_path.open(encoding="utf-8") as handle:
|
||||
return yaml.safe_load(handle) or {}
|
||||
|
||||
# -- Path helpers --------------------------------------------------------
|
||||
|
||||
@property
|
||||
def _base(self) -> Path:
|
||||
"""Directory that repo-relative paths resolve against for filesystem I/O."""
|
||||
return self.repo_root if self.repo_root is not None else Path.cwd()
|
||||
Reference in New Issue
Block a user