585 lines
18 KiB
Python
585 lines
18 KiB
Python
from __future__ import annotations
|
|
|
|
from io import StringIO
|
|
from pathlib import Path
|
|
from textwrap import dedent
|
|
|
|
import pytest
|
|
|
|
from tools.docs_codegen.cli import main
|
|
from tools.docs_codegen.converters import RUN_DP_TEMPLATE_POSITIONALS
|
|
from tools.docs_codegen.errors import DocsCodegenError
|
|
from tools.docs_codegen.generator import GeneratorService
|
|
from tools.docs_codegen.scanner import BlockScanner
|
|
from tools.docs_codegen.utils import substitute_template_positionals
|
|
|
|
|
|
def _write_text(path: Path, content: str) -> None:
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
path.write_text(dedent(content).lstrip(), encoding="utf-8")
|
|
|
|
|
|
def _write_single_node_yaml(tmp_path: Path) -> Path:
|
|
case_path = tmp_path / "cases" / "single_node.yaml"
|
|
_write_text(
|
|
case_path,
|
|
"""
|
|
test_cases:
|
|
- name: default-case
|
|
model: default/model
|
|
envs:
|
|
SERVER_PORT: 8123
|
|
server_cmd: []
|
|
- name: selected-case
|
|
model: "Qwen/Test Model"
|
|
envs:
|
|
HCCL_BUFFSIZE: 1024
|
|
PROMPT: "hello world"
|
|
SERVER_PORT: DEFAULT_PORT
|
|
server_cmd:
|
|
- "--tensor-parallel-size"
|
|
- 2
|
|
- "--kv-transfer-config"
|
|
- '{"foo": "bar", "enabled": true}'
|
|
server_cmd_extra: "--trust-remote-code --enable-expert-parallel"
|
|
""",
|
|
)
|
|
return case_path.relative_to(tmp_path)
|
|
|
|
|
|
def _write_multi_node_yaml(tmp_path: Path, *, invalid_command: bool = False) -> Path:
|
|
case_path = tmp_path / "cases" / "multi_node.yaml"
|
|
second_command = (
|
|
'"vllm serve multi-node/model extra-positional"'
|
|
if invalid_command
|
|
else """
|
|
- vllm
|
|
- serve
|
|
- multi-node/model
|
|
- "--headless"
|
|
- "--port"
|
|
- "$SERVER_PORT"
|
|
"""
|
|
)
|
|
_write_text(
|
|
case_path,
|
|
f"""
|
|
deployment:
|
|
- envs:
|
|
LOCAL_IP: 127.0.0.1
|
|
server_cmd: "vllm serve first-host --port 8000"
|
|
- envs:
|
|
MASTER_IP: 10.0.0.1
|
|
SERVER_PORT: 9000
|
|
server_cmd: {second_command.rstrip()}
|
|
""",
|
|
)
|
|
return case_path.relative_to(tmp_path)
|
|
|
|
|
|
def _write_model_code_doc(tmp_path: Path, content: str, *, name: str = "Demo.md") -> Path:
|
|
doc_path = tmp_path / "docs" / "models" / name
|
|
_write_text(doc_path, content)
|
|
return doc_path.relative_to(tmp_path)
|
|
|
|
|
|
def _generate_block(tmp_path: Path, doc_path: Path, block_name: str) -> str:
|
|
service = GeneratorService(artifact_root="artifacts")
|
|
return service.generate_block(doc_path, block_name, dry_run=True)[1].content
|
|
|
|
|
|
# Explicit CPU smart-UT routing keeps this guard out of the --run-all-cpu bucket.
|
|
def test_block_scanner_parses_metadata_and_trims_raw_block(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
|
|
single_yaml = _write_single_node_yaml(tmp_path)
|
|
doc_path = _write_model_code_doc(
|
|
tmp_path,
|
|
f"""
|
|
# Demo
|
|
|
|
```{{model-code}}
|
|
:block_name: single
|
|
:converter_tag: single_node
|
|
:test_case_path: {single_yaml}
|
|
:case_index: 1
|
|
|
|
set -eux
|
|
{{{{ generated }}}}
|
|
|
|
```
|
|
""",
|
|
)
|
|
monkeypatch.chdir(tmp_path)
|
|
|
|
blocks = BlockScanner().scan_document_blocks(doc_path)
|
|
|
|
assert len(blocks) == 1
|
|
block = blocks[0]
|
|
assert block.doc_path == doc_path
|
|
assert block.block_name == "single"
|
|
assert block.converter_tag == "single_node"
|
|
assert block.test_case_path == single_yaml.as_posix()
|
|
assert block.extra_options == (("case_index", "1"),)
|
|
assert block.directive_line == 3
|
|
assert block.raw_block_lines == ("set -eux", "{{ generated }}")
|
|
|
|
|
|
def test_block_scanner_rejects_duplicate_block_names(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
|
|
single_yaml = _write_single_node_yaml(tmp_path)
|
|
doc_path = _write_model_code_doc(
|
|
tmp_path,
|
|
f"""
|
|
```{{model-code}}
|
|
:block_name: serve
|
|
:converter_tag: single_node
|
|
:test_case_path: {single_yaml}
|
|
```
|
|
|
|
```{{model-code}}
|
|
:block_name: serve
|
|
:converter_tag: single_node
|
|
:test_case_path: {single_yaml}
|
|
```
|
|
""",
|
|
name="Duplicate.md",
|
|
)
|
|
monkeypatch.chdir(tmp_path)
|
|
|
|
with pytest.raises(DocsCodegenError) as exc_info:
|
|
BlockScanner().scan_document_blocks(doc_path)
|
|
|
|
error_message = str(exc_info.value)
|
|
assert "docs/models/Duplicate.md:7: model-code generation error" in error_message
|
|
assert "block_name: serve" in error_message
|
|
assert "duplicated block_name 'serve'" in error_message
|
|
assert "previous declaration is on line 1" in error_message
|
|
|
|
|
|
def test_block_scanner_rejects_unsupported_metadata(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
|
|
single_yaml = _write_single_node_yaml(tmp_path)
|
|
doc_path = _write_model_code_doc(
|
|
tmp_path,
|
|
f"""
|
|
```{{model-code}}
|
|
:block_name: serve
|
|
:converter_tag: single_node
|
|
:test_case_path: {single_yaml}
|
|
:unknown_option: value
|
|
```
|
|
""",
|
|
)
|
|
monkeypatch.chdir(tmp_path)
|
|
|
|
with pytest.raises(DocsCodegenError) as exc_info:
|
|
BlockScanner().scan_document_blocks(doc_path)
|
|
|
|
assert "unsupported metadata: unknown_option" in str(exc_info.value)
|
|
|
|
|
|
def test_single_node_converter_uses_case_index_defaults_and_extra_args(
|
|
tmp_path: Path,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
):
|
|
single_yaml = _write_single_node_yaml(tmp_path)
|
|
doc_path = _write_model_code_doc(
|
|
tmp_path,
|
|
f"""
|
|
```{{model-code}}
|
|
:block_name: single
|
|
:converter_tag: single_node
|
|
:test_case_path: {single_yaml}
|
|
:case_index: 1
|
|
|
|
set -eux
|
|
{{{{ generated }}}}
|
|
```
|
|
""",
|
|
)
|
|
monkeypatch.chdir(tmp_path)
|
|
|
|
script = _generate_block(tmp_path, doc_path, "single")
|
|
|
|
assert script.startswith("set -eux\nexport HCCL_BUFFSIZE=1024")
|
|
assert 'export PROMPT="hello world"' in script
|
|
assert "export SERVER_PORT=8000" in script
|
|
assert "ignored/model" not in script
|
|
assert "vllm serve 'Qwen/Test Model' \\" in script
|
|
assert "--tensor-parallel-size 2 \\" in script
|
|
assert '"foo": "bar"' in script
|
|
assert '"enabled": true' in script
|
|
assert "--trust-remote-code \\" in script
|
|
assert "--enable-expert-parallel" in script
|
|
|
|
|
|
def test_single_node_converter_defaults_to_first_case_and_preserves_port(
|
|
tmp_path: Path,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
):
|
|
single_yaml = _write_single_node_yaml(tmp_path)
|
|
doc_path = _write_model_code_doc(
|
|
tmp_path,
|
|
f"""
|
|
```{{model-code}}
|
|
:block_name: default_case
|
|
:converter_tag: single_node
|
|
:test_case_path: {single_yaml}
|
|
```
|
|
""",
|
|
)
|
|
monkeypatch.chdir(tmp_path)
|
|
|
|
script = _generate_block(tmp_path, doc_path, "default_case")
|
|
|
|
assert script == "export SERVER_PORT=8123\n\nvllm serve default/model\n"
|
|
|
|
|
|
def test_single_node_converter_reports_invalid_case_index(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
|
|
single_yaml = _write_single_node_yaml(tmp_path)
|
|
doc_path = _write_model_code_doc(
|
|
tmp_path,
|
|
f"""
|
|
```{{model-code}}
|
|
:block_name: missing_case
|
|
:converter_tag: single_node
|
|
:test_case_path: {single_yaml}
|
|
:case_index: 3
|
|
```
|
|
""",
|
|
)
|
|
monkeypatch.chdir(tmp_path)
|
|
|
|
with pytest.raises(DocsCodegenError) as exc_info:
|
|
_generate_block(tmp_path, doc_path, "missing_case")
|
|
|
|
assert "case_index 3 is out of range for 'test_cases' with 2 items" in str(exc_info.value)
|
|
|
|
|
|
def test_multi_node_converter_uses_host_index_and_token_list(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
|
|
multi_yaml = _write_multi_node_yaml(tmp_path)
|
|
doc_path = _write_model_code_doc(
|
|
tmp_path,
|
|
f"""
|
|
```{{model-code}}
|
|
:block_name: worker
|
|
:converter_tag: multi_node
|
|
:test_case_path: {multi_yaml}
|
|
:host_index: 1
|
|
```
|
|
""",
|
|
)
|
|
monkeypatch.chdir(tmp_path)
|
|
|
|
script = _generate_block(tmp_path, doc_path, "worker")
|
|
|
|
assert script.startswith("export MASTER_IP=10.0.0.1\nexport SERVER_PORT=9000")
|
|
assert "vllm serve multi-node/model \\" in script
|
|
assert "--headless \\" in script
|
|
assert "--port $SERVER_PORT" in script
|
|
assert "first-host" not in script
|
|
|
|
|
|
def test_multi_node_converter_requires_host_index(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
|
|
multi_yaml = _write_multi_node_yaml(tmp_path)
|
|
doc_path = _write_model_code_doc(
|
|
tmp_path,
|
|
f"""
|
|
```{{model-code}}
|
|
:block_name: worker
|
|
:converter_tag: multi_node
|
|
:test_case_path: {multi_yaml}
|
|
```
|
|
""",
|
|
)
|
|
monkeypatch.chdir(tmp_path)
|
|
|
|
with pytest.raises(DocsCodegenError) as exc_info:
|
|
_generate_block(tmp_path, doc_path, "worker")
|
|
|
|
assert "converter_tag 'multi_node' requires host_index" in str(exc_info.value)
|
|
|
|
|
|
def test_multi_node_converter_rejects_extra_positional_args(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
|
|
multi_yaml = _write_multi_node_yaml(tmp_path, invalid_command=True)
|
|
doc_path = _write_model_code_doc(
|
|
tmp_path,
|
|
f"""
|
|
```{{model-code}}
|
|
:block_name: worker
|
|
:converter_tag: multi_node
|
|
:test_case_path: {multi_yaml}
|
|
:host_index: 1
|
|
```
|
|
""",
|
|
)
|
|
monkeypatch.chdir(tmp_path)
|
|
|
|
with pytest.raises(DocsCodegenError) as exc_info:
|
|
_generate_block(tmp_path, doc_path, "worker")
|
|
|
|
assert "unsupported positional argument 'extra-positional'" in str(exc_info.value)
|
|
|
|
|
|
def test_generator_service_writes_selected_artifact(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
|
|
single_yaml = _write_single_node_yaml(tmp_path)
|
|
doc_path = _write_model_code_doc(
|
|
tmp_path,
|
|
f"""
|
|
```{{model-code}}
|
|
:block_name: default_case
|
|
:converter_tag: single_node
|
|
:test_case_path: {single_yaml}
|
|
```
|
|
""",
|
|
)
|
|
monkeypatch.chdir(tmp_path)
|
|
|
|
service = GeneratorService(artifact_root="artifacts")
|
|
output_path, generated_script = service.generate_block(doc_path, "default_case", dry_run=False)
|
|
|
|
assert output_path == Path("artifacts/Demo/default_case.sh")
|
|
assert output_path.read_text(encoding="utf-8") == generated_script.content
|
|
assert generated_script.content == "export SERVER_PORT=8123\n\nvllm serve default/model\n"
|
|
|
|
|
|
def test_cli_generates_block_to_stdout(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
|
|
single_yaml = _write_single_node_yaml(tmp_path)
|
|
doc_path = _write_model_code_doc(
|
|
tmp_path,
|
|
f"""
|
|
```{{model-code}}
|
|
:block_name: default_case
|
|
:converter_tag: single_node
|
|
:test_case_path: {single_yaml}
|
|
```
|
|
""",
|
|
)
|
|
monkeypatch.chdir(tmp_path)
|
|
stdout = StringIO()
|
|
stderr = StringIO()
|
|
|
|
exit_code = main(["--block", f"{doc_path}::default_case", "--dry-run", "--stdout"], stdout=stdout, stderr=stderr)
|
|
|
|
assert exit_code == 0
|
|
assert stderr.getvalue() == ""
|
|
assert stdout.getvalue() == (
|
|
"docs/_build/doc_codegen/Demo/default_case.sh\nexport SERVER_PORT=8123\n\nvllm serve default/model\n"
|
|
)
|
|
|
|
|
|
def test_cli_rejects_invalid_block_reference():
|
|
stdout = StringIO()
|
|
stderr = StringIO()
|
|
|
|
exit_code = main(["--block", "docs/model.md"], stdout=stdout, stderr=stderr)
|
|
|
|
assert exit_code == 1
|
|
assert stdout.getvalue() == ""
|
|
assert "block reference must use '<doc_path>::<block_name>'" in stderr.getvalue()
|
|
|
|
|
|
def _write_external_dp_yaml(tmp_path: Path, *, routing_type: str = "disaggregated_prefill") -> Path:
|
|
case_path = tmp_path / "cases" / "external_dp.yaml"
|
|
_write_text(
|
|
case_path,
|
|
f"""
|
|
model: "Eco-Tech/GLM-Test"
|
|
num_nodes: 2
|
|
|
|
routing:
|
|
type: "{routing_type}"
|
|
groups:
|
|
prefiller: [0]
|
|
decoder: [1]
|
|
|
|
config:
|
|
- node_index: 0
|
|
port_start: 7100
|
|
dp_rpc_port: 12321
|
|
dp_size: 4
|
|
dp_size_local: 2
|
|
dp_rank_start: 0
|
|
tp_size: 8
|
|
dp_address: "${{NODE_0_IP}}"
|
|
- node_index: 1
|
|
port_start: 7200
|
|
dp_rpc_port: 12321
|
|
dp_size: 8
|
|
dp_size_local: 4
|
|
dp_rank_start: 0
|
|
tp_size: 4
|
|
dp_address: "${{NODE_1_IP}}"
|
|
|
|
env_common: &env_common
|
|
HCCL_BUFFSIZE: "1024"
|
|
OMP_PROC_BIND: "false"
|
|
|
|
templates:
|
|
- node_index: 0
|
|
envs:
|
|
<<: *env_common
|
|
ASCEND_RT_VISIBLE_DEVICES: "${{VISIBLE_DEVICES}}"
|
|
server_cmd_template:
|
|
- --host
|
|
- "0.0.0.0"
|
|
- --port
|
|
- ${{PORT}}
|
|
- --data-parallel-size
|
|
- ${{DP_SIZE}}
|
|
- --data-parallel-rank
|
|
- ${{DP_RANK}}
|
|
- --tensor-parallel-size
|
|
- ${{TP_SIZE}}
|
|
- --profiler-config
|
|
- '{{"profiler":"torch","with_stack":false}}'
|
|
- --kv-transfer-config
|
|
- '{{"kv_connector": "MooncakeConnectorV1", "kv_role": "kv_producer", "kv_port": "30000"}}'
|
|
- node_index: 1
|
|
envs:
|
|
<<: *env_common
|
|
ASCEND_RT_VISIBLE_DEVICES: "${{VISIBLE_DEVICES}}"
|
|
server_cmd_template:
|
|
- --host
|
|
- "0.0.0.0"
|
|
- --port
|
|
- ${{PORT}}
|
|
- --data-parallel-size
|
|
- ${{DP_SIZE}}
|
|
- --tensor-parallel-size
|
|
- ${{TP_SIZE}}
|
|
""",
|
|
)
|
|
return case_path.relative_to(tmp_path)
|
|
|
|
|
|
def test_substitute_template_positionals():
|
|
positionals = RUN_DP_TEMPLATE_POSITIONALS
|
|
assert substitute_template_positionals("${DP_SIZE}", positionals=positionals) == "$3"
|
|
assert substitute_template_positionals("--port ${PORT}", positionals=positionals) == "--port $2"
|
|
# Unknown braced variables and unbraced refs are left untouched.
|
|
assert substitute_template_positionals("${UNKNOWN}", positionals=positionals) == "${UNKNOWN}"
|
|
assert substitute_template_positionals("$SERVER_PORT", positionals=positionals) == "$SERVER_PORT"
|
|
|
|
|
|
def test_external_dp_template_converter_maps_positionals(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
|
|
external_yaml = _write_external_dp_yaml(tmp_path)
|
|
doc_path = _write_model_code_doc(
|
|
tmp_path,
|
|
f"""
|
|
```{{model-code}}
|
|
:block_name: prefill_n0
|
|
:converter_tag: external_dp_template
|
|
:test_case_path: {external_yaml}
|
|
:host_index: 0
|
|
```
|
|
""",
|
|
)
|
|
monkeypatch.chdir(tmp_path)
|
|
|
|
script = _generate_block(tmp_path, doc_path, "prefill_n0")
|
|
|
|
assert "export HCCL_BUFFSIZE=1024" in script
|
|
assert "export ASCEND_RT_VISIBLE_DEVICES=$1" in script
|
|
assert "SERVER_PORT" not in script
|
|
assert "vllm serve Eco-Tech/GLM-Test \\" in script
|
|
assert "--port $2 \\" in script
|
|
assert "--data-parallel-size $3 \\" in script
|
|
assert "--data-parallel-rank $4 \\" in script
|
|
assert "--tensor-parallel-size $7 \\" in script
|
|
# Space-free JSON values are quoted (not just whitespace-containing ones).
|
|
assert '--profiler-config \'{"profiler":"torch","with_stack":false}\' \\' in script
|
|
assert '"kv_role": "kv_producer"' in script
|
|
assert "${DP_SIZE}" not in script
|
|
|
|
|
|
def test_external_dp_template_converter_requires_host_index(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
|
|
external_yaml = _write_external_dp_yaml(tmp_path)
|
|
doc_path = _write_model_code_doc(
|
|
tmp_path,
|
|
f"""
|
|
```{{model-code}}
|
|
:block_name: prefill_n0
|
|
:converter_tag: external_dp_template
|
|
:test_case_path: {external_yaml}
|
|
```
|
|
""",
|
|
)
|
|
monkeypatch.chdir(tmp_path)
|
|
|
|
with pytest.raises(DocsCodegenError) as exc_info:
|
|
_generate_block(tmp_path, doc_path, "prefill_n0")
|
|
|
|
assert "converter_tag 'external_dp_template' requires host_index" in str(exc_info.value)
|
|
|
|
|
|
def test_external_dp_launch_converter_combines_all_nodes(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
|
|
external_yaml = _write_external_dp_yaml(tmp_path)
|
|
doc_path = _write_model_code_doc(
|
|
tmp_path,
|
|
f"""
|
|
```{{model-code}}
|
|
:block_name: launch
|
|
:converter_tag: external_dp_launch
|
|
:test_case_path: {external_yaml}
|
|
```
|
|
""",
|
|
)
|
|
monkeypatch.chdir(tmp_path)
|
|
|
|
script = _generate_block(tmp_path, doc_path, "launch")
|
|
|
|
assert script == (
|
|
"python launch_online_dp.py --dp-size 4 --tp-size 8 --dp-size-local 2 --dp-rank-start 0 "
|
|
"--dp-address ${NODE_0_IP} --dp-rpc-port 12321 --vllm-start-port 7100\n\n"
|
|
"python launch_online_dp.py --dp-size 8 --tp-size 4 --dp-size-local 4 --dp-rank-start 0 "
|
|
"--dp-address ${NODE_1_IP} --dp-rpc-port 12321 --vllm-start-port 7200\n"
|
|
)
|
|
|
|
|
|
def test_external_dp_proxy_converter_expands_groups(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
|
|
external_yaml = _write_external_dp_yaml(tmp_path)
|
|
doc_path = _write_model_code_doc(
|
|
tmp_path,
|
|
f"""
|
|
```{{model-code}}
|
|
:block_name: proxy
|
|
:converter_tag: external_dp_proxy
|
|
:test_case_path: {external_yaml}
|
|
```
|
|
""",
|
|
)
|
|
monkeypatch.chdir(tmp_path)
|
|
|
|
script = _generate_block(tmp_path, doc_path, "proxy")
|
|
|
|
assert script.startswith("python load_balance_proxy_server_example.py \\")
|
|
# Single-value flags stay inline; multi-value flags expand one value per line.
|
|
assert " --host ${NODE_0_IP} \\" in script
|
|
assert " --port 1999 \\" in script
|
|
assert " --prefiller-hosts \\\n ${NODE_0_IP} \\\n ${NODE_0_IP} \\" in script
|
|
assert " --prefiller-ports \\\n 7100 \\\n 7101 \\" in script
|
|
assert (
|
|
" --decoder-hosts \\\n ${NODE_1_IP} \\\n ${NODE_1_IP} \\\n ${NODE_1_IP} \\\n ${NODE_1_IP} \\"
|
|
in script
|
|
)
|
|
assert " --decoder-ports \\\n 7200 \\\n 7201 \\\n 7202 \\\n 7203" in script
|
|
assert script.rstrip().endswith(" 7203")
|
|
|
|
|
|
def test_external_dp_proxy_converter_rejects_unsupported_routing(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
|
|
external_yaml = _write_external_dp_yaml(tmp_path, routing_type="generic_dp")
|
|
doc_path = _write_model_code_doc(
|
|
tmp_path,
|
|
f"""
|
|
```{{model-code}}
|
|
:block_name: proxy
|
|
:converter_tag: external_dp_proxy
|
|
:test_case_path: {external_yaml}
|
|
```
|
|
""",
|
|
)
|
|
monkeypatch.chdir(tmp_path)
|
|
|
|
with pytest.raises(DocsCodegenError) as exc_info:
|
|
_generate_block(tmp_path, doc_path, "proxy")
|
|
|
|
assert "only supports routing.type 'disaggregated_prefill'" in str(exc_info.value)
|