@@ -1,4 +1,4 @@
|
||||
#!/bin/bash
|
||||
#!/usr/bin/env bash
|
||||
|
||||
#
|
||||
# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
|
||||
|
||||
474
tools/ai_qos.py
Normal file
474
tools/ai_qos.py
Normal file
@@ -0,0 +1,474 @@
|
||||
import argparse
|
||||
import importlib
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
from vllm_ascend import ai_qos
|
||||
|
||||
VISIBLE_DEVICE_ENV = "ASCEND_RT_VISIBLE_DEVICES"
|
||||
|
||||
MASTER_ID_AIV_DATA = 11
|
||||
MASTER_ID_AIV_INS = 12
|
||||
MASTER_ID_SDMA = 13
|
||||
MASTER_ID_PCIEDMA = 7
|
||||
FUSE_SELECT_MAX = 1
|
||||
SDMA_MATA_BW_LOW = 0
|
||||
SDMA_MATA_BW_HIGH = 1
|
||||
SDMA_MATA_HARDLIMIT = 0
|
||||
D2D_VL_INIT = 0
|
||||
H2D_VL_INIT = 1
|
||||
STATE_SDMA_MATA_LEN = 4
|
||||
STATE_QOS_TUPLE_LEN = 5
|
||||
STATE_FUSE_GBL_LEN = 3
|
||||
FUSE_APPLY_ENABLE = 1
|
||||
FUSE_APPLY_AUTOQOS_FUSE_EN = 1
|
||||
|
||||
DEFAULT_STATE_PATH = Path(__file__).resolve().parent / "ai_qos_state.json"
|
||||
|
||||
# Shown when unset cannot parse the state file; state file is removed after printing.
|
||||
UNSET_STATE_PARSE_FAILED_MSG = (
|
||||
"Failed to parse the state file. Please reboot the server to restore endpoint-side QoS settings. "
|
||||
"On the switch side, log in and run `sys-view` to enter system view, then run "
|
||||
"`display current-configuration` to show the current configuration. Find the previously applied "
|
||||
"switch-side QoS commands, re-enter each command with the `undo` prefix, and finally run `commit` "
|
||||
"to complete configuration rollback."
|
||||
)
|
||||
|
||||
|
||||
def _remove_state_file(path: Path) -> None:
|
||||
try:
|
||||
path.unlink()
|
||||
except OSError as e:
|
||||
print(
|
||||
f"Warning: could not remove state file {path}: {e}. Please remove this file manually.",
|
||||
file=sys.stderr,
|
||||
)
|
||||
|
||||
|
||||
def _unset_state_parse_failed(state_path: Path) -> None:
|
||||
print(UNSET_STATE_PARSE_FAILED_MSG, file=sys.stderr)
|
||||
_remove_state_file(state_path)
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
def _parse_and_validate_unset_state(data: object) -> tuple[dict, list[str], dict[str, list[int]], dict[str, list[int]]]:
|
||||
"""Validate unset JSON shape and types; raise ValueError on any failure."""
|
||||
if not isinstance(data, dict):
|
||||
raise ValueError("root must be object")
|
||||
oq = data.get("original_qos")
|
||||
if not isinstance(oq, dict):
|
||||
raise ValueError("original_qos")
|
||||
pc = data.get("printed_commands")
|
||||
if not isinstance(pc, list) or not all(isinstance(x, str) for x in pc):
|
||||
raise ValueError("printed_commands")
|
||||
osm_raw = data.get("original_sdma_mata", {})
|
||||
if not isinstance(osm_raw, dict):
|
||||
raise ValueError("original_sdma_mata")
|
||||
ofu_raw = data.get("original_fuse", {})
|
||||
if not isinstance(ofu_raw, dict):
|
||||
raise ValueError("original_fuse")
|
||||
|
||||
validated_oq: dict[str, dict[str, list[int]]] = {}
|
||||
for dev_s, masters in oq.items():
|
||||
try:
|
||||
_dev = int(dev_s)
|
||||
except (TypeError, ValueError) as e:
|
||||
raise ValueError("original_qos device key") from e
|
||||
if not isinstance(masters, dict):
|
||||
raise ValueError("original_qos masters")
|
||||
vm: dict[str, list[int]] = {}
|
||||
for m_s, tup in masters.items():
|
||||
try:
|
||||
_ = int(m_s)
|
||||
except (TypeError, ValueError) as e:
|
||||
raise ValueError("original_qos master key") from e
|
||||
if not isinstance(tup, list) or len(tup) != STATE_QOS_TUPLE_LEN:
|
||||
raise ValueError("original_qos tuple")
|
||||
try:
|
||||
vm[str(m_s)] = [int(x) for x in tup]
|
||||
except (TypeError, ValueError) as e:
|
||||
raise ValueError("original_qos tuple values") from e
|
||||
validated_oq[str(_dev)] = vm
|
||||
|
||||
validated_osm: dict[str, list[int]] = {}
|
||||
for dev_s, mata in osm_raw.items():
|
||||
try:
|
||||
_dev = int(dev_s)
|
||||
except (TypeError, ValueError) as e:
|
||||
raise ValueError("original_sdma_mata device key") from e
|
||||
if not isinstance(mata, list) or len(mata) != STATE_SDMA_MATA_LEN:
|
||||
raise ValueError("original_sdma_mata tuple")
|
||||
try:
|
||||
validated_osm[str(_dev)] = [int(x) for x in mata]
|
||||
except (TypeError, ValueError) as e:
|
||||
raise ValueError("original_sdma_mata values") from e
|
||||
|
||||
validated_ofu: dict[str, list[int]] = {}
|
||||
for dev_s, gbl in ofu_raw.items():
|
||||
try:
|
||||
_dev = int(dev_s)
|
||||
except (TypeError, ValueError) as e:
|
||||
raise ValueError("original_fuse device key") from e
|
||||
if not isinstance(gbl, list) or len(gbl) != STATE_FUSE_GBL_LEN:
|
||||
raise ValueError("original_fuse tuple")
|
||||
try:
|
||||
validated_ofu[str(_dev)] = [int(x) for x in gbl]
|
||||
except (TypeError, ValueError) as e:
|
||||
raise ValueError("original_fuse values") from e
|
||||
|
||||
return validated_oq, pc, validated_osm, validated_ofu
|
||||
|
||||
|
||||
def _print_config_block(lines: list[str]) -> None:
|
||||
print("system-view")
|
||||
for line in lines:
|
||||
print(line)
|
||||
print("commit")
|
||||
|
||||
|
||||
def _device_list() -> list[int]:
|
||||
device_str = os.getenv(VISIBLE_DEVICE_ENV, "").strip()
|
||||
if not device_str:
|
||||
try:
|
||||
torch = importlib.import_module("torch")
|
||||
count = int(torch.npu.device_count())
|
||||
except Exception as e:
|
||||
print(
|
||||
f"Error: {VISIBLE_DEVICE_ENV} is unset and failed to run torch.npu.device_count().",
|
||||
file=sys.stderr,
|
||||
)
|
||||
print(f"Details: {e}", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
if count <= 0:
|
||||
print("Error: no visible NPU devices found.", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
return list(range(count))
|
||||
out: list[int] = []
|
||||
for dev in device_str.split(","):
|
||||
part = dev.strip()
|
||||
if not part:
|
||||
print(
|
||||
f"Error: invalid {VISIBLE_DEVICE_ENV} value (empty segment): {device_str!r}",
|
||||
file=sys.stderr,
|
||||
)
|
||||
sys.exit(1)
|
||||
try:
|
||||
d = int(part, 10)
|
||||
except ValueError:
|
||||
print(
|
||||
f"Error: {VISIBLE_DEVICE_ENV} must be comma-separated integers; got {device_str!r}",
|
||||
file=sys.stderr,
|
||||
)
|
||||
sys.exit(1)
|
||||
if d < 0:
|
||||
print(
|
||||
f"Error: {VISIBLE_DEVICE_ENV} device id must be non-negative; got {d}",
|
||||
file=sys.stderr,
|
||||
)
|
||||
sys.exit(1)
|
||||
out.append(d)
|
||||
if not out:
|
||||
print(f"Error: {VISIBLE_DEVICE_ENV} must list at least one device.", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
return out
|
||||
|
||||
|
||||
def _capture_original_qos(device_list: list[int], masterid_table: dict[str, int]) -> dict[str, dict[str, list]]:
|
||||
original: dict[str, dict[str, list]] = {}
|
||||
for device_id in device_list:
|
||||
key_d = str(device_id)
|
||||
original[key_d] = {}
|
||||
for _accu, master_id in masterid_table.items():
|
||||
ret, m, mpamid, q, pmg, mode = ai_qos.get_qos(device_id, master_id)
|
||||
if ret != 0:
|
||||
print(
|
||||
f"Warning: get_qos failed (dev={device_id} master={master_id} ret={ret}); not saved for restore.",
|
||||
file=sys.stderr,
|
||||
)
|
||||
continue
|
||||
original[key_d][str(master_id)] = [m, mpamid, q, pmg, mode]
|
||||
return original
|
||||
|
||||
|
||||
def _capture_original_sdma_mata(
|
||||
device_list: list[int],
|
||||
) -> dict[str, list]:
|
||||
out: dict[str, list] = {}
|
||||
for device_id in device_list:
|
||||
ret_q, _m, mpamid, _q, _pmg, _mode = ai_qos.get_qos(device_id, MASTER_ID_SDMA)
|
||||
if ret_q != 0:
|
||||
print(
|
||||
f"Warning: get_qos(SDMA) failed (dev={device_id} ret={ret_q}); SDMA mata not saved for restore.",
|
||||
file=sys.stderr,
|
||||
)
|
||||
continue
|
||||
ret, bw_lo, bw_hi, hard = ai_qos.get_bw(device_id, mpamid)
|
||||
if ret == 0:
|
||||
out[str(device_id)] = [int(mpamid), int(bw_lo), int(bw_hi), int(hard)]
|
||||
return out
|
||||
|
||||
|
||||
def _capture_original_fuse(
|
||||
device_list: list[int],
|
||||
) -> dict[str, list]:
|
||||
out: dict[str, list] = {}
|
||||
for device_id in device_list:
|
||||
ret, en, aut, fuse = ai_qos.get_fuse_mode(device_id)
|
||||
if ret == 0:
|
||||
out[str(device_id)] = [int(en), int(aut), int(fuse)]
|
||||
return out
|
||||
|
||||
|
||||
def _merge_baseline_for_new_devices(
|
||||
device_list: list[int],
|
||||
original_qos: dict[str, dict[str, list]],
|
||||
original_sdma_mata: dict[str, list],
|
||||
original_fuse: dict[str, list],
|
||||
masterid_table: dict[str, int],
|
||||
) -> None:
|
||||
missing: list[int] = []
|
||||
for d in device_list:
|
||||
ds = str(d)
|
||||
if ds not in original_qos or ds not in original_sdma_mata or ds not in original_fuse:
|
||||
missing.append(d)
|
||||
if not missing:
|
||||
return
|
||||
oq = _capture_original_qos(missing, masterid_table)
|
||||
for k, v in oq.items():
|
||||
original_qos[k] = v
|
||||
osm = _capture_original_sdma_mata(missing)
|
||||
for k, v in osm.items():
|
||||
original_sdma_mata[k] = v
|
||||
ofu = _capture_original_fuse(missing)
|
||||
for k, v in ofu.items():
|
||||
original_fuse[k] = v
|
||||
|
||||
|
||||
def _load_first_apply_baseline(
|
||||
state_path: Path,
|
||||
) -> tuple[dict, dict, dict] | None:
|
||||
if not state_path.is_file():
|
||||
return None
|
||||
try:
|
||||
data = json.loads(state_path.read_text(encoding="utf-8"))
|
||||
except (OSError, json.JSONDecodeError):
|
||||
return None
|
||||
oq = data.get("original_qos")
|
||||
if not isinstance(oq, dict):
|
||||
return None
|
||||
osm = data.get("original_sdma_mata", {})
|
||||
ofu = data.get("original_fuse", {})
|
||||
if not isinstance(osm, dict):
|
||||
osm = {}
|
||||
if not isinstance(ofu, dict):
|
||||
ofu = {}
|
||||
return (oq, osm, ofu)
|
||||
|
||||
|
||||
def run_unset(state_path: Path) -> None:
|
||||
if not state_path.is_file():
|
||||
print(f"No state file at {state_path}; nothing to undo.", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
try:
|
||||
data = json.loads(state_path.read_text(encoding="utf-8"))
|
||||
except (OSError, json.JSONDecodeError):
|
||||
_unset_state_parse_failed(state_path)
|
||||
|
||||
try:
|
||||
original_qos, printed, original_sdma_mata, original_fuse = _parse_and_validate_unset_state(data)
|
||||
except ValueError:
|
||||
_unset_state_parse_failed(state_path)
|
||||
|
||||
for dev_s, mata in original_sdma_mata.items():
|
||||
device_id = int(dev_s)
|
||||
mid, bw_lo, bw_hi, hard = mata
|
||||
r = ai_qos.set_bw(device_id, mid, bw_lo, bw_hi, hard)
|
||||
if r != 0:
|
||||
print(
|
||||
f"Warning: restore SDMA mata (dev={device_id} mpamid={mid}) failed, ret = {r}",
|
||||
file=sys.stderr,
|
||||
)
|
||||
|
||||
for dev_s, masters in original_qos.items():
|
||||
device_id = int(dev_s)
|
||||
for m_s, tup in masters.items():
|
||||
master_id = int(m_s)
|
||||
_m, mpamid, qos, pmg, mode = tup
|
||||
ai_qos.set_qos(device_id, master_id, mpamid, qos, pmg, mode)
|
||||
|
||||
for dev_s, gbl in original_fuse.items():
|
||||
device_id = int(dev_s)
|
||||
en, aut, fmode = gbl
|
||||
r = ai_qos.set_fuse_gbl_config(device_id, en, aut, fmode)
|
||||
if r != 0:
|
||||
print(
|
||||
f"Warning: restore fuse gbl (dev={device_id}) failed, ret = {r}",
|
||||
file=sys.stderr,
|
||||
)
|
||||
|
||||
_print_config_block([f"undo {line}" for line in printed])
|
||||
|
||||
try:
|
||||
state_path.unlink()
|
||||
except OSError as e:
|
||||
print(
|
||||
f"Warning: could not remove state file {state_path}: {e}. Please remove this file manually.",
|
||||
file=sys.stderr,
|
||||
)
|
||||
|
||||
|
||||
class AiqosConfig:
|
||||
def __init__(self, aiqos_config: dict):
|
||||
self.mode = aiqos_config.get("mode")
|
||||
self.aiqos_priority = aiqos_config.get("aiqos_priority")
|
||||
self.aiqos_table = {
|
||||
"AIV_D2D": {"low": (1, 0, 0, 1), "middle": (3, 4, 1, 2), "high": (5, 5, 2, 3)},
|
||||
"AIV_H2D": {"low": (1, 0, 3, 1), "middle": (3, 4, 4, 2), "high": (5, 5, 5, 3)},
|
||||
"SDMA_D2D": {"low": (2, 0, 0, 1), "middle": (4, 4, 1, 2), "high": (6, 5, 2, 3)},
|
||||
"SDMA_H2D": {"low": (2, 0, 3, 1), "middle": (4, 4, 4, 2), "high": (6, 5, 5, 3)},
|
||||
"PCIEDMA_H2D": {"low": (0, 0, 3, 1), "high": (7, 5, 5, 3)},
|
||||
}
|
||||
self.masterid_table = {
|
||||
"AIV_DATA": MASTER_ID_AIV_DATA,
|
||||
"AIV_INS": MASTER_ID_AIV_INS,
|
||||
"SDMA": MASTER_ID_SDMA,
|
||||
"PCIEDMA": MASTER_ID_PCIEDMA,
|
||||
}
|
||||
|
||||
def set_qos(self, state_path: Path) -> None:
|
||||
device_list = _device_list()
|
||||
baseline = _load_first_apply_baseline(state_path)
|
||||
if baseline is not None:
|
||||
original_qos, original_sdma_mata, original_fuse = baseline
|
||||
_merge_baseline_for_new_devices(
|
||||
device_list,
|
||||
original_qos,
|
||||
original_sdma_mata,
|
||||
original_fuse,
|
||||
self.masterid_table,
|
||||
)
|
||||
else:
|
||||
original_qos = _capture_original_qos(device_list, self.masterid_table)
|
||||
original_sdma_mata = _capture_original_sdma_mata(
|
||||
device_list,
|
||||
)
|
||||
original_fuse = _capture_original_fuse(device_list)
|
||||
|
||||
attributes = ["sqos", "dqos", "vl", "pri"]
|
||||
for op_type in self.aiqos_table:
|
||||
level = self.aiqos_priority.get(op_type)
|
||||
config_tuple = self.aiqos_table.get(op_type).get(level)
|
||||
for idx, attr in enumerate(attributes):
|
||||
var_name = f"{op_type.lower()}_{attr}"
|
||||
setattr(self, var_name, config_tuple[idx])
|
||||
aiv_qos = min(self.aiv_d2d_sqos, self.aiv_h2d_sqos)
|
||||
sdma_qos = min(self.sdma_d2d_sqos, self.sdma_h2d_sqos)
|
||||
pcie_qos = self.pciedma_h2d_sqos
|
||||
fuse_mode = FUSE_SELECT_MAX
|
||||
|
||||
command_types = {
|
||||
"aiv_d2d": aiv_qos,
|
||||
"aiv_h2d": aiv_qos,
|
||||
"sdma_d2d": sdma_qos,
|
||||
"sdma_h2d": sdma_qos,
|
||||
"pciedma_h2d": pcie_qos,
|
||||
}
|
||||
|
||||
def generate_command(qos_value: int, dqos: int, vl: int, pri: int, vl_init: int) -> str:
|
||||
return (
|
||||
f"hccs qos remap {qos_value} {vl_init} {dqos}\n"
|
||||
f"hccs vl remap peer-type cpu {dqos} {vl_init} {vl}\n"
|
||||
f"hccs vl remap peer-type npu {dqos} {vl_init} {vl}\n"
|
||||
f"hccs vl remap peer-type sw {dqos} {vl_init} {vl}\n"
|
||||
f"hccs sp peer-type cpu {vl} {pri}\n"
|
||||
f"hccs sp peer-type npu {vl} {pri}\n"
|
||||
f"hccs sp peer-type sw {vl} {pri}\n"
|
||||
)
|
||||
|
||||
cmd_set: set[str] = set()
|
||||
for cmd_type, qos_value in command_types.items():
|
||||
dqos = getattr(self, f"{cmd_type}_dqos")
|
||||
vl = getattr(self, f"{cmd_type}_vl")
|
||||
pri = getattr(self, f"{cmd_type}_pri")
|
||||
vl_init = D2D_VL_INIT
|
||||
if "h2d" in cmd_type:
|
||||
vl_init = H2D_VL_INIT
|
||||
cmd_str = generate_command(qos_value, dqos, vl, pri, vl_init)
|
||||
for sub_str in cmd_str.split("\n"):
|
||||
if sub_str.strip():
|
||||
cmd_set.add(sub_str)
|
||||
printed_commands = sorted(cmd_set)
|
||||
|
||||
for device_id in device_list:
|
||||
ai_qos.set_fuse_gbl_config(device_id, FUSE_APPLY_ENABLE, FUSE_APPLY_AUTOQOS_FUSE_EN, fuse_mode)
|
||||
for accu, master_id in self.masterid_table.items():
|
||||
ret, _master, mpamid, qos, pmg, mode = ai_qos.get_qos(device_id, master_id)
|
||||
if ret != 0:
|
||||
print(
|
||||
f"get_qos failed (dev={device_id} master={master_id} ret={ret}).",
|
||||
file=sys.stderr,
|
||||
)
|
||||
continue
|
||||
if accu.startswith("AIV"):
|
||||
ai_qos.set_qos(device_id, master_id, mpamid, aiv_qos, pmg, mode)
|
||||
elif accu.startswith("SDMA"):
|
||||
ai_qos.set_bw(device_id, mpamid, SDMA_MATA_BW_LOW, SDMA_MATA_BW_HIGH, SDMA_MATA_HARDLIMIT)
|
||||
ai_qos.set_qos(device_id, master_id, mpamid, sdma_qos, pmg, mode)
|
||||
else:
|
||||
ai_qos.set_qos(device_id, master_id, mpamid, pcie_qos, pmg, mode)
|
||||
|
||||
_print_config_block(printed_commands)
|
||||
|
||||
state_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
state_path.write_text(
|
||||
json.dumps(
|
||||
{
|
||||
"original_qos": original_qos,
|
||||
"original_sdma_mata": original_sdma_mata,
|
||||
"original_fuse": original_fuse,
|
||||
"printed_commands": printed_commands,
|
||||
},
|
||||
indent=2,
|
||||
),
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(
|
||||
description="AI QoS tuning for Ascend NPU. "
|
||||
" Multiple apply reuses the "
|
||||
"first snapshot in the state file; unset restores that baseline and removes the file."
|
||||
)
|
||||
parser.add_argument(
|
||||
"command",
|
||||
nargs="?",
|
||||
default="apply",
|
||||
choices=["apply", "unset"],
|
||||
help='Run "unset" to restore the first-apply snapshot and delete the state file.',
|
||||
)
|
||||
parser.add_argument("--mode", type=str, default="auto", choices=["auto", "manual"])
|
||||
parser.add_argument("--AIV_D2D", type=str, default="high", choices=["low", "middle", "high"])
|
||||
parser.add_argument("--AIV_H2D", type=str, default="high", choices=["low", "middle", "high"])
|
||||
parser.add_argument("--SDMA_D2D", type=str, default="high", choices=["low", "middle", "high"])
|
||||
parser.add_argument("--SDMA_H2D", type=str, default="low", choices=["low", "middle", "high"])
|
||||
parser.add_argument("--PCIEDMA_H2D", type=str, default="high", choices=["low", "high"])
|
||||
args = parser.parse_args()
|
||||
state_path = DEFAULT_STATE_PATH
|
||||
|
||||
if args.command == "unset":
|
||||
run_unset(state_path)
|
||||
else:
|
||||
aiqos_config = {
|
||||
"mode": args.mode,
|
||||
"aiqos_priority": {
|
||||
"AIV_D2D": args.AIV_D2D,
|
||||
"AIV_H2D": args.AIV_H2D,
|
||||
"SDMA_D2D": args.SDMA_D2D,
|
||||
"SDMA_H2D": args.SDMA_H2D,
|
||||
"PCIEDMA_H2D": args.PCIEDMA_H2D,
|
||||
},
|
||||
}
|
||||
AiqosConfig(aiqos_config).set_qos(state_path)
|
||||
61
tools/ai_qos/CMakeLists.txt
Normal file
61
tools/ai_qos/CMakeLists.txt
Normal file
@@ -0,0 +1,61 @@
|
||||
cmake_minimum_required(VERSION 3.16)
|
||||
project(ai_qos LANGUAGES CXX)
|
||||
find_package(Python3 COMPONENTS Interpreter Development.Module REQUIRED)
|
||||
execute_process(
|
||||
COMMAND python3 -c "import pybind11, os; print(os.path.dirname(os.path.dirname(pybind11.get_include())))"
|
||||
OUTPUT_VARIABLE PYBIND11_ROOT
|
||||
OUTPUT_STRIP_TRAILING_WHITESPACE
|
||||
)
|
||||
list(APPEND CMAKE_PREFIX_PATH "${PYBIND11_ROOT}")
|
||||
find_package(pybind11 CONFIG REQUIRED)
|
||||
pybind11_add_module(ai_qos MODULE ai_qos.cpp)
|
||||
target_compile_features(ai_qos PRIVATE cxx_std_11)
|
||||
# ---- DSMI lookup (driver SDK, not CANN-only) ----
|
||||
set(_dsmi_inc_paths
|
||||
/usr/local/Ascend/driver/include
|
||||
/usr/local/Ascend/ascend-toolkit/latest/driver/include
|
||||
/usr/local/Ascend/driver/kernel/inc/driver
|
||||
)
|
||||
set(_dsmi_lib_paths
|
||||
/usr/local/Ascend/driver/lib64/driver
|
||||
/usr/local/Ascend/ascend-toolkit/latest/driver/lib64
|
||||
)
|
||||
if(DEFINED ENV{ASCEND_HOME_PATH} AND NOT "$ENV{ASCEND_HOME_PATH}" STREQUAL "")
|
||||
get_filename_component(_ascend_parent "$ENV{ASCEND_HOME_PATH}" DIRECTORY)
|
||||
list(INSERT _dsmi_inc_paths 0
|
||||
"${_ascend_parent}/driver/include"
|
||||
"$ENV{ASCEND_HOME_PATH}/driver/include"
|
||||
)
|
||||
list(INSERT _dsmi_lib_paths 0
|
||||
"${_ascend_parent}/driver/lib64/driver"
|
||||
"$ENV{ASCEND_HOME_PATH}/driver/lib64"
|
||||
)
|
||||
endif()
|
||||
find_path(DSMI_INCLUDE_DIR
|
||||
NAMES dsmi_common_interface.h
|
||||
PATHS ${_dsmi_inc_paths}
|
||||
NO_DEFAULT_PATH
|
||||
)
|
||||
find_library(DSMI_LIBRARY
|
||||
NAMES drvdsmi_host
|
||||
PATHS ${_dsmi_lib_paths}
|
||||
NO_DEFAULT_PATH
|
||||
)
|
||||
if(NOT DSMI_INCLUDE_DIR)
|
||||
message(FATAL_ERROR
|
||||
"ai_qos: dsmi_common_interface.h not found. "
|
||||
"Pass -DDSMI_INCLUDE_DIR=/path/to/include"
|
||||
)
|
||||
endif()
|
||||
if(NOT DSMI_LIBRARY)
|
||||
message(FATAL_ERROR
|
||||
"ai_qos: drvdsmi_host not found. "
|
||||
"Pass -DDSMI_LIBRARY=/path/to/libdrvdsmi_host.so"
|
||||
)
|
||||
endif()
|
||||
target_include_directories(ai_qos PRIVATE
|
||||
${pybind11_INCLUDE_DIRS}
|
||||
${DSMI_INCLUDE_DIR}
|
||||
)
|
||||
target_link_libraries(ai_qos PRIVATE ${DSMI_LIBRARY})
|
||||
install(TARGETS ai_qos LIBRARY DESTINATION .)
|
||||
136
tools/ai_qos/ai_qos.cpp
Normal file
136
tools/ai_qos/ai_qos.cpp
Normal file
@@ -0,0 +1,136 @@
|
||||
#include <pybind11/pybind11.h>
|
||||
#include <stdio.h>
|
||||
#include <stdlib.h>
|
||||
#include <string.h>
|
||||
#include <tuple>
|
||||
#include "dsmi_common_interface.h"
|
||||
|
||||
#define DSMI_QOS_INDEX_OFFSET 8
|
||||
#define DSMI_QOS_INDEX_LEN 8U
|
||||
#define DSMI_QOS_MAIN_INDEX_OFFSET 8U
|
||||
#define DSMI_QOS_SUB_INDEX_OFFSET 16U
|
||||
#define DSMI_QOS_THIRD_INDEX_OFFSET 24U
|
||||
|
||||
#define PCIEDMA_MASTER 7
|
||||
#define SDMA_MASTER 13
|
||||
|
||||
#define DSMI_QOS_SUB_CMD_MAKE(qos_index, qos_sub_cmd) (((qos_index) << DSMI_QOS_INDEX_OFFSET) | (qos_sub_cmd))
|
||||
#define DSMI_QOS_SUB_CMD_MAKE_V2(qos_main_index, qos_sub_index, qos_third_index, qos_sub_cmd) \
|
||||
((((qos_main_index) & ((1U << DSMI_QOS_INDEX_LEN) - 1U)) << DSMI_QOS_MAIN_INDEX_OFFSET) | \
|
||||
(((qos_sub_index) & ((1U << DSMI_QOS_INDEX_LEN) - 1U)) << DSMI_QOS_SUB_INDEX_OFFSET) | \
|
||||
(((qos_third_index) & ((1U << DSMI_QOS_INDEX_LEN) - 1U)) << DSMI_QOS_THIRD_INDEX_OFFSET) | (qos_sub_cmd))
|
||||
|
||||
int set_fuse_gbl_config(unsigned int device_id, uint32_t enable, uint32_t autoqos_fuse_en, int mpamqos_fuse_mode) {
|
||||
struct qos_gbl_config gblCfg = {0};
|
||||
gblCfg.enable = enable;
|
||||
gblCfg.autoqos_fuse_en = autoqos_fuse_en;
|
||||
gblCfg.mpamqos_fuse_mode = mpamqos_fuse_mode;
|
||||
int ret = dsmi_set_device_info(device_id, DSMI_MAIN_CMD_QOS, static_cast<uint32_t>(DSMI_QOS_SUB_GLOBAL_CONFIG),
|
||||
static_cast<void*>(&gblCfg), sizeof(struct qos_gbl_config));
|
||||
if (ret != 0) {
|
||||
printf("[dev:%d] set fuse gbl (en=%u auto=%u mode=%d) failed, ret = %d\n", device_id, enable, autoqos_fuse_en,
|
||||
mpamqos_fuse_mode, ret);
|
||||
return ret;
|
||||
}
|
||||
return ret;
|
||||
}
|
||||
|
||||
int set_qos(unsigned int device_id, int master, int mpamid, int qos, int pmg, int mode) {
|
||||
struct qos_master_config masterCfg = {0};
|
||||
masterCfg.master = master;
|
||||
masterCfg.mpamid = mpamid;
|
||||
masterCfg.qos = qos;
|
||||
masterCfg.pmg = pmg;
|
||||
if (master == PCIEDMA_MASTER) {
|
||||
masterCfg.bitmap[0] = 0x1;
|
||||
} else if (master == SDMA_MASTER) {
|
||||
masterCfg.bitmap[0] = 0xffffffffffffffff;
|
||||
}
|
||||
masterCfg.mode = mode;
|
||||
int ret = dsmi_set_device_info(device_id, DSMI_MAIN_CMD_QOS, static_cast<uint32_t>(DSMI_QOS_SUB_MASTER_CONFIG),
|
||||
static_cast<void*>(&masterCfg), sizeof(struct qos_master_config));
|
||||
if (ret != 0) {
|
||||
printf("[dev:%d] set qos = %d failed, ret = %d\n", device_id, qos, ret);
|
||||
return ret;
|
||||
}
|
||||
return ret;
|
||||
}
|
||||
|
||||
int set_bw(unsigned int device_id, int mpamid, int bw_low, int bw_high, int hardlimit) {
|
||||
struct qos_mata_config mataCfg = {0};
|
||||
mataCfg.mpamid = mpamid;
|
||||
mataCfg.bw_low = bw_low;
|
||||
mataCfg.bw_high = bw_high;
|
||||
mataCfg.hardlimit = hardlimit;
|
||||
int ret = dsmi_set_device_info(device_id, DSMI_MAIN_CMD_QOS, static_cast<uint32_t>(DSMI_QOS_SUB_MATA_CONFIG),
|
||||
static_cast<void*>(&mataCfg), sizeof(struct qos_mata_config));
|
||||
if (ret != 0) {
|
||||
printf("[dev:%d] mpamid: %d set bw: %d-%d failed, ret = %d\n", device_id, mpamid, bw_low, bw_high, ret);
|
||||
return ret;
|
||||
}
|
||||
return ret;
|
||||
}
|
||||
|
||||
std::tuple<int, unsigned int, unsigned int, int> get_bw(unsigned int device_id, int mpamid) {
|
||||
struct qos_mata_config mataCfg = {0};
|
||||
mataCfg.mpamid = mpamid;
|
||||
uint32_t size = sizeof(struct qos_mata_config);
|
||||
uint32_t subCmd = static_cast<uint32_t>(DSMI_QOS_SUB_CMD_MAKE(mataCfg.mpamid, DSMI_QOS_SUB_MATA_CONFIG));
|
||||
int ret = dsmi_get_device_info(device_id, DSMI_MAIN_CMD_QOS, subCmd, static_cast<void*>(&mataCfg), &size);
|
||||
if (ret != 0 || size != sizeof(struct qos_mata_config)) {
|
||||
printf("[dev:%d] mpamid: %d get bw failed, ret = %d, size = %u, main cmd = %#x, sub cmd = %#x\n", device_id, mpamid,
|
||||
ret, size, DSMI_MAIN_CMD_QOS, subCmd);
|
||||
int err = (ret != 0) ? ret : -1;
|
||||
return std::make_tuple(err, 0, 0, 0);
|
||||
}
|
||||
return std::make_tuple(ret, mataCfg.bw_low, mataCfg.bw_high, mataCfg.hardlimit);
|
||||
}
|
||||
|
||||
std::tuple<int, int, int, int, int, unsigned int> get_qos(unsigned int device_id, int master) {
|
||||
struct qos_master_config masterCfg = {0};
|
||||
masterCfg.master = master;
|
||||
uint32_t size = sizeof(struct qos_master_config);
|
||||
int coreid = 0;
|
||||
uint32_t subCmd = static_cast<uint32_t>(
|
||||
DSMI_QOS_SUB_CMD_MAKE_V2(static_cast<uint32_t>(masterCfg.master), coreid, 0, DSMI_QOS_SUB_MASTER_CONFIG));
|
||||
int ret = dsmi_get_device_info(device_id, DSMI_MAIN_CMD_QOS, subCmd, static_cast<void*>(&masterCfg), &size);
|
||||
if (ret != 0 || size != sizeof(struct qos_master_config)) {
|
||||
printf("[dev:%d] get qos failed, ret = %d, size = %u, main cmd = %#x, sub cmd = %#x\n", device_id, ret, size,
|
||||
DSMI_MAIN_CMD_QOS, subCmd);
|
||||
int err = (ret != 0) ? ret : -1;
|
||||
return std::make_tuple(err, 0, 0, 0, 0, 0U);
|
||||
}
|
||||
return std::make_tuple(0, static_cast<int>(masterCfg.master), static_cast<int>(masterCfg.mpamid),
|
||||
static_cast<int>(masterCfg.qos), static_cast<int>(masterCfg.pmg), masterCfg.mode);
|
||||
}
|
||||
|
||||
std::tuple<int, unsigned int, unsigned int, unsigned int> get_fuse_mode(unsigned int device_id) {
|
||||
struct qos_gbl_config gblCfg = {0};
|
||||
uint32_t size = sizeof(struct qos_gbl_config);
|
||||
int ret = dsmi_get_device_info(device_id, DSMI_MAIN_CMD_QOS, static_cast<uint32_t>(DSMI_QOS_SUB_GLOBAL_CONFIG),
|
||||
static_cast<void*>(&gblCfg), &size);
|
||||
if (ret != 0 || size != sizeof(struct qos_gbl_config)) {
|
||||
printf("[dev:%d] get fuse mode failed, ret = %d, size = %u, main cmd = %#x, sub cmd = %#x\n", device_id, ret, size,
|
||||
DSMI_MAIN_CMD_QOS, DSMI_QOS_SUB_GLOBAL_CONFIG);
|
||||
int err = (ret != 0) ? ret : -1;
|
||||
return std::make_tuple(err, 0U, 0U, 0U);
|
||||
}
|
||||
return std::make_tuple(0, gblCfg.enable, gblCfg.autoqos_fuse_en, gblCfg.mpamqos_fuse_mode);
|
||||
}
|
||||
|
||||
namespace py = pybind11;
|
||||
PYBIND11_MODULE(ai_qos, m) {
|
||||
m.doc() = "AI QoS(Quality of Service) control module for hardware resource management";
|
||||
m.def("set_qos", &set_qos, py::arg("device_id"), py::arg("master"), py::arg("mpamid"), py::arg("qos"), py::arg("pmg"),
|
||||
py::arg("mode"));
|
||||
m.def("set_fuse_gbl_config", &set_fuse_gbl_config, py::arg("device_id"), py::arg("enable"),
|
||||
py::arg("autoqos_fuse_en"), py::arg("mpamqos_fuse_mode"));
|
||||
m.def("get_qos", &get_qos,
|
||||
"Returns (ret, master, mpamid, qos, pmg, mode). ret==0 on success; on failure ret is DSMI error (or -1 if size "
|
||||
"mismatch).",
|
||||
py::arg("device_id"), py::arg("master"));
|
||||
m.def("set_bw", &set_bw, py::arg("device_id"), py::arg("mpamid"), py::arg("bw_low"), py::arg("bw_high"),
|
||||
py::arg("hardlimit"));
|
||||
m.def("get_bw", &get_bw, py::arg("device_id"), py::arg("mpamid"));
|
||||
m.def("get_fuse_mode", &get_fuse_mode, py::arg("device_id"));
|
||||
}
|
||||
334
tools/aisbench.py
Normal file
334
tools/aisbench.py
Normal file
@@ -0,0 +1,334 @@
|
||||
# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
|
||||
# Copyright 2023 The vLLM team.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
# This file is a part of the vllm-ascend project.
|
||||
#
|
||||
import hashlib
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import subprocess
|
||||
import tempfile
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import filelock
|
||||
import huggingface_hub
|
||||
import pandas as pd
|
||||
import regex as re
|
||||
from modelscope import snapshot_download # type: ignore
|
||||
|
||||
BENCHMARK_HOME = os.getenv("BENCHMARK_HOME", os.path.abspath("./benchmark"))
|
||||
DATASET_CONF_DIR = os.path.join(BENCHMARK_HOME, "ais_bench", "benchmark", "configs", "datasets")
|
||||
REQUEST_CONF_DIR = os.path.join(BENCHMARK_HOME, "ais_bench", "benchmark", "configs", "models", "vllm_api")
|
||||
DATASET_DIR = os.path.join(BENCHMARK_HOME, "ais_bench", "datasets")
|
||||
|
||||
|
||||
class AisbenchRunner:
|
||||
RESULT_MSG = {"performance": "Performance Result files located in ", "accuracy": "write csv to "}
|
||||
DATASET_RENAME = {"aime2024": "aime", "gsm8k-lite": "gsm8k", "textvqa-lite": "textvqa"}
|
||||
|
||||
def _run_aisbench_task(self):
|
||||
dataset_conf = self.dataset_conf.split("/")[-1]
|
||||
if self.task_type == "accuracy":
|
||||
aisbench_cmd = ["ais_bench", "--models", f"{self.request_conf}_custom", "--datasets", f"{dataset_conf}"]
|
||||
if self.task_type == "performance":
|
||||
aisbench_cmd = [
|
||||
"ais_bench",
|
||||
"--models",
|
||||
f"{self.request_conf}_custom",
|
||||
"--datasets",
|
||||
f"{dataset_conf}_custom",
|
||||
"--mode",
|
||||
"perf",
|
||||
]
|
||||
if self.num_prompts:
|
||||
aisbench_cmd.extend(["--num-prompts", str(self.num_prompts)])
|
||||
self.stdout_file = f"output_{self.task_type}.txt"
|
||||
aisbench_cmd = " ".join(aisbench_cmd) + f" --debug > {self.stdout_file} 2>&1 &"
|
||||
print(f"running aisbench cmd: {aisbench_cmd}")
|
||||
self.proc: subprocess.Popen = subprocess.Popen(aisbench_cmd, shell=True)
|
||||
|
||||
def __init__(self, model: str, port: int, aisbench_config: dict, host_ip: str = "localhost", verify=True):
|
||||
self.model = model
|
||||
self.dataset_path = aisbench_config.get("dataset_path_local")
|
||||
if not self.dataset_path:
|
||||
self.dataset_path = maybe_download_from_modelscope(aisbench_config["dataset_path"], repo_type="dataset")
|
||||
self.model_path = aisbench_config.get("model_path")
|
||||
if not self.model_path:
|
||||
self.model_path = maybe_download_from_modelscope(model)
|
||||
assert self.dataset_path is not None and self.model_path is not None, (
|
||||
f"Failed to download dataset or model: dataset={self.dataset_path}, model={self.model_path}"
|
||||
)
|
||||
self.port = port
|
||||
self.host_ip = host_ip
|
||||
self.task_type = aisbench_config["case_type"]
|
||||
self.request_conf = aisbench_config["request_conf"]
|
||||
self.dataset_conf = aisbench_config.get("dataset_conf")
|
||||
self.num_prompts = aisbench_config.get("num_prompts")
|
||||
self.max_out_len = aisbench_config["max_out_len"]
|
||||
self.batch_size = aisbench_config["batch_size"]
|
||||
self.request_rate = aisbench_config.get("request_rate", 0)
|
||||
self.trust_remote_code = aisbench_config.get("trust_remote_code", True)
|
||||
self.temperature = aisbench_config.get("temperature")
|
||||
self.top_k = aisbench_config.get("top_k")
|
||||
self.top_p = aisbench_config.get("top_p")
|
||||
self.seed = aisbench_config.get("seed")
|
||||
self.min_p = aisbench_config.get("min_p")
|
||||
self.presence_penalty = aisbench_config.get("presence_penalty")
|
||||
self.repetition_penalty = aisbench_config.get("repetition_penalty")
|
||||
self.no_pred = aisbench_config.get("no_pred")
|
||||
self.thinking = aisbench_config.get("thinking")
|
||||
self.exp_folder = None
|
||||
self.result_line = None
|
||||
self._init_dataset_conf()
|
||||
self._init_request_conf()
|
||||
self._run_aisbench_task()
|
||||
self._wait_for_task()
|
||||
if verify:
|
||||
self.baseline = aisbench_config.get("baseline", 1)
|
||||
if self.task_type == "accuracy":
|
||||
self.threshold = aisbench_config.get("threshold", 1)
|
||||
self._accuracy_verify()
|
||||
if self.task_type == "performance":
|
||||
self.threshold = aisbench_config.get("threshold", 0.97)
|
||||
self._performance_verify()
|
||||
|
||||
def _init_dataset_conf(self):
|
||||
if self.task_type == "accuracy":
|
||||
dataset_name = os.path.basename(self.dataset_path)
|
||||
dataset_rename = self.DATASET_RENAME.get(dataset_name, dataset_name)
|
||||
dst_dir = os.path.join(DATASET_DIR, dataset_rename)
|
||||
command = ["cp", "-r", self.dataset_path, dst_dir]
|
||||
subprocess.call(command)
|
||||
if self.task_type == "performance":
|
||||
conf_path = os.path.join(DATASET_CONF_DIR, f"{self.dataset_conf}.py")
|
||||
if self.dataset_conf.startswith("textvqa"):
|
||||
self.dataset_path = os.path.join(self.dataset_path, "textvqa_val.jsonl")
|
||||
with open(conf_path, encoding="utf-8") as f:
|
||||
content = f.read()
|
||||
content = re.sub(r"path=.*", f'path="{self.dataset_path}",', content)
|
||||
conf_path_new = os.path.join(DATASET_CONF_DIR, f"{self.dataset_conf}_custom.py")
|
||||
with open(conf_path_new, "w", encoding="utf-8") as f:
|
||||
f.write(content)
|
||||
|
||||
def _init_request_conf(self):
|
||||
conf_path = os.path.join(REQUEST_CONF_DIR, f"{self.request_conf}.py")
|
||||
with open(conf_path, encoding="utf-8") as f:
|
||||
content = f.read()
|
||||
content = re.sub(r"model=.*", f'model="{self.model}",', content)
|
||||
content = re.sub(r"host_port.*", f"host_port={self.port},", content)
|
||||
content = re.sub(r"host_ip.*", f'host_ip="{self.host_ip}",', content)
|
||||
content = re.sub(r"max_out_len.*", f"max_out_len={self.max_out_len},", content)
|
||||
content = re.sub(r"batch_size.*", f"batch_size={self.batch_size},", content)
|
||||
content = re.sub(r"trust_remote_code=.*", f"trust_remote_code={self.trust_remote_code},", content)
|
||||
if self.top_p:
|
||||
content = re.sub(r"ignore_eos.*", f"ignore_eos=False,\n top_p={self.top_p},", content)
|
||||
if self.top_k:
|
||||
content = re.sub(r"ignore_eos.*", f"ignore_eos=False,\n top_k={self.top_k},", content)
|
||||
if self.seed:
|
||||
content = re.sub(r"ignore_eos.*", f"ignore_eos=False,\n seed={self.seed},", content)
|
||||
if self.min_p:
|
||||
content = re.sub(r"ignore_eos.*", f"ignore_eos=False,\n min_p={self.min_p},", content)
|
||||
if self.presence_penalty:
|
||||
content = re.sub(
|
||||
r"ignore_eos.*", f"ignore_eos=False,\n presence_penalty={self.presence_penalty},", content
|
||||
)
|
||||
if self.repetition_penalty:
|
||||
content = re.sub(
|
||||
r"ignore_eos.*",
|
||||
f"ignore_eos=False,\n repetition_penalty={self.repetition_penalty},",
|
||||
content,
|
||||
)
|
||||
if self.thinking:
|
||||
field_thinking = 'chat_template_kwargs={"thinking": True}'
|
||||
content = re.sub(r"ignore_eos.*", f"ignore_eos=False,\n {field_thinking},", content)
|
||||
if self.task_type == "performance":
|
||||
content = re.sub(r"path=.*", f'path="{self.model_path}",', content)
|
||||
content = re.sub(r"request_rate.*", f"request_rate={self.request_rate},", content)
|
||||
content = re.sub(r"temperature.*", "temperature=0,", content)
|
||||
content = re.sub(r"ignore_eos.*", "ignore_eos=True,", content)
|
||||
if self.task_type == "accuracy":
|
||||
content = re.sub(r"temperature.*", "temperature=0.6,", content)
|
||||
if self.temperature is not None:
|
||||
content = re.sub(r"temperature.*", f"temperature={self.temperature},", content)
|
||||
if self.no_pred:
|
||||
content = re.sub(r"pred_postprocessor.*", "#pred_postprocessor", content)
|
||||
conf_path_new = os.path.join(REQUEST_CONF_DIR, f"{self.request_conf}_custom.py")
|
||||
with open(conf_path_new, "w", encoding="utf-8") as f:
|
||||
f.write(content)
|
||||
print(f"The request config is\n {content}")
|
||||
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc_value, traceback):
|
||||
self.proc.terminate()
|
||||
try:
|
||||
self.proc.wait(8)
|
||||
except subprocess.TimeoutExpired:
|
||||
# force kill if needed
|
||||
self.proc.kill()
|
||||
|
||||
def _wait_for_exp_folder(self):
|
||||
while True:
|
||||
line = self._check_runtime_stdout()
|
||||
if "Current exp folder: " in line:
|
||||
self.exp_folder = re.search(r"Current exp folder: (.*)", line).group(1)
|
||||
print(f"Current exp folder: {self.exp_folder}")
|
||||
return
|
||||
self._check_runtime_error(line)
|
||||
|
||||
def _wait_for_task(self):
|
||||
self._wait_for_exp_folder()
|
||||
result_msg = self.RESULT_MSG[self.task_type]
|
||||
while True:
|
||||
line = self._check_runtime_stdout()
|
||||
if result_msg in line:
|
||||
print(line)
|
||||
self.result_line = line
|
||||
return
|
||||
self._check_runtime_error(line)
|
||||
|
||||
def _check_runtime_stdout(self):
|
||||
time.sleep(5)
|
||||
cmd = f"tail -100 {self.stdout_file}"
|
||||
return subprocess.check_output(cmd, shell=True).decode()
|
||||
|
||||
@staticmethod
|
||||
def _check_runtime_error(line):
|
||||
if "ERROR" in line:
|
||||
print(line)
|
||||
error_msg = "Some errors happened to Aisbench runtime"
|
||||
raise RuntimeError(error_msg) from None
|
||||
|
||||
def _get_result_performance(self):
|
||||
result_dir = re.search(r"Performance Result files located in (.*)", self.result_line).group(1)[:-1]
|
||||
dataset_type = self.dataset_conf.split("/")[0]
|
||||
result_csv_file = os.path.join(result_dir, f"{dataset_type}.csv")
|
||||
result_json_file = os.path.join(result_dir, f"{dataset_type}.json")
|
||||
self.result_csv = pd.read_csv(result_csv_file, index_col=0)
|
||||
print("Getting performance results from file: ", result_json_file)
|
||||
with open(result_json_file, encoding="utf-8") as f:
|
||||
self.result_json = json.load(f)
|
||||
self.result = [self.result_csv, self.result_json]
|
||||
|
||||
def _get_result_accuracy(self):
|
||||
acc_file = re.search(r"write csv to (.*)", self.result_line).group(1)
|
||||
df = pd.read_csv(acc_file)
|
||||
self.result = float(df.iloc[0, -1])
|
||||
|
||||
def _performance_verify(self):
|
||||
self._get_result_performance()
|
||||
output_throughput = self.result_json["Output Token Throughput"]["total"].replace("token/s", "")
|
||||
assert float(output_throughput) >= self.threshold * self.baseline, (
|
||||
"Performance verification failed. "
|
||||
f"The current Output Token Throughput is {output_throughput} token/s, "
|
||||
f"which is not greater than or equal to {self.threshold} * baseline {self.baseline}."
|
||||
)
|
||||
|
||||
def _accuracy_verify(self):
|
||||
self._get_result_accuracy()
|
||||
acc_value = self.result
|
||||
assert self.baseline - self.threshold <= acc_value <= self.baseline + self.threshold, (
|
||||
"Accuracy verification failed. "
|
||||
f"The accuracy of {self.dataset_path} is {acc_value}, "
|
||||
f"which is not within {self.threshold} relative to baseline {self.baseline}."
|
||||
)
|
||||
|
||||
|
||||
def run_aisbench_cases(model, port, aisbench_cases, server_args="", host_ip="localhost"):
|
||||
aisbench_results = []
|
||||
aisbench_errors = []
|
||||
total = sum(1 for c in aisbench_cases if c)
|
||||
idx = 0
|
||||
for aisbench_case in aisbench_cases:
|
||||
if not aisbench_case:
|
||||
continue
|
||||
idx += 1
|
||||
case_name = aisbench_case.get("case_name", "unknown")
|
||||
case_type = aisbench_case.get("case_type", "unknown")
|
||||
logging.info("=" * 60)
|
||||
logging.info("[%d/%d] Starting benchmark: %s (type=%s)", idx, total, case_name, case_type)
|
||||
logging.info("=" * 60)
|
||||
try:
|
||||
with AisbenchRunner(model=model, port=port, host_ip=host_ip, aisbench_config=aisbench_case) as aisbench:
|
||||
aisbench_results.append(aisbench.result)
|
||||
logging.info("[%d/%d] Finished benchmark: %s", idx, total, case_name)
|
||||
except Exception as e:
|
||||
aisbench_results.append("")
|
||||
aisbench_errors.append([aisbench_case, e])
|
||||
logging.error("[%d/%d] Benchmark failed: %s, reason: %s", idx, total, case_name, e)
|
||||
print(e)
|
||||
for failed_case, error_info in aisbench_errors:
|
||||
error_msg = f"The following aisbench case failed: {failed_case}, reason is {error_info}"
|
||||
if server_args:
|
||||
error_msg += f"\nserver_args are {server_args}"
|
||||
logging.error(error_msg)
|
||||
assert not aisbench_errors, "some aisbench cases failed, info were shown above."
|
||||
return aisbench_results
|
||||
|
||||
|
||||
def get_TTFT(results):
|
||||
TTFT = []
|
||||
for i in range(len(results)):
|
||||
TTFT.append(float(results[i][0].loc["TTFT", "Average"][:-3]))
|
||||
return TTFT
|
||||
|
||||
|
||||
temp_dir = tempfile.gettempdir()
|
||||
|
||||
|
||||
def get_lock(model_name_or_path: str | Path, cache_dir: str | None = None):
|
||||
lock_dir = cache_dir or temp_dir
|
||||
model_name_or_path = str(model_name_or_path)
|
||||
os.makedirs(os.path.dirname(lock_dir), exist_ok=True)
|
||||
model_name = model_name_or_path.replace("/", "-")
|
||||
hash_name = hashlib.sha256(model_name.encode()).hexdigest()
|
||||
# add hash to avoid conflict with old users' lock files
|
||||
lock_file_name = hash_name + model_name + ".lock"
|
||||
# mode 0o666 is required for the filelock to be shared across users
|
||||
lock = filelock.FileLock(os.path.join(lock_dir, lock_file_name), mode=0o666)
|
||||
return lock
|
||||
|
||||
|
||||
def maybe_download_from_modelscope(
|
||||
model: str,
|
||||
repo_type: str = "model",
|
||||
revision: str | None = None,
|
||||
download_dir: str | None = None,
|
||||
ignore_patterns: str | list[str] | None = None,
|
||||
allow_patterns: list[str] | str | None = None,
|
||||
) -> str:
|
||||
"""
|
||||
Download model/dataset from ModelScope hub.
|
||||
Returns the path to the downloaded model, or None if the model is not
|
||||
downloaded from ModelScope.
|
||||
"""
|
||||
# Use file lock to prevent multiple processes from
|
||||
# downloading the same model weights at the same time.
|
||||
with get_lock(model, download_dir):
|
||||
if not os.path.exists(model):
|
||||
model_path = snapshot_download(
|
||||
model_id=model,
|
||||
repo_type=repo_type,
|
||||
cache_dir=download_dir,
|
||||
local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE,
|
||||
revision=revision,
|
||||
ignore_file_pattern=ignore_patterns,
|
||||
allow_patterns=allow_patterns,
|
||||
)
|
||||
else:
|
||||
model_path = model
|
||||
return model_path
|
||||
87
tools/check_boolean_context_manager.py
Normal file
87
tools/check_boolean_context_manager.py
Normal file
@@ -0,0 +1,87 @@
|
||||
#!/usr/bin/env python3
|
||||
#
|
||||
# Copyright (c) 2026 Huawei Technologies Co., Ltd. All Rights Reserved.
|
||||
# Copyright 2023 The vLLM team.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
# This file is a part of the vllm-ascend project.
|
||||
# Adapted from https://github.com/vllm-project/vllm/tree/main/tools
|
||||
#
|
||||
"""Lint: detect `with a() and b():` (boolean op in with-statement context).
|
||||
|
||||
Using `and`/`or` to combine context managers is almost always a bug:
|
||||
|
||||
with ctx_a() and ctx_b(): # BUG: only ctx_b is entered
|
||||
with ctx_a() or ctx_b(): # BUG: only ctx_a is entered
|
||||
|
||||
The correct way to combine context managers is:
|
||||
|
||||
with ctx_a(), ctx_b(): # comma-separated
|
||||
with (ctx_a(), ctx_b()): # parenthesized (Python 3.10+)
|
||||
with contextlib.ExitStack() ... # ExitStack
|
||||
"""
|
||||
|
||||
import ast
|
||||
import sys
|
||||
|
||||
|
||||
def check_file(filepath: str) -> list[str]:
|
||||
try:
|
||||
with open(filepath, encoding="utf-8") as f:
|
||||
source = f.read()
|
||||
except (OSError, UnicodeDecodeError):
|
||||
return []
|
||||
|
||||
try:
|
||||
tree = ast.parse(source, filename=filepath)
|
||||
except SyntaxError:
|
||||
return []
|
||||
|
||||
violations = []
|
||||
for node in ast.walk(tree):
|
||||
if isinstance(node, (ast.With, ast.AsyncWith)):
|
||||
for item in node.items:
|
||||
if isinstance(item.context_expr, ast.BoolOp):
|
||||
op = "and" if isinstance(item.context_expr.op, ast.And) else "or"
|
||||
violations.append(
|
||||
f"{filepath}:{item.context_expr.lineno}: "
|
||||
f"boolean `{op}` used to combine context managers "
|
||||
"in `with` statement; use a comma instead"
|
||||
)
|
||||
return violations
|
||||
|
||||
|
||||
def main() -> int:
|
||||
if len(sys.argv) < 2:
|
||||
print("Usage: check_boolean_context_manager.py <file> ...", file=sys.stderr)
|
||||
return 1
|
||||
|
||||
all_violations = []
|
||||
for filepath in sys.argv[1:]:
|
||||
all_violations.extend(check_file(filepath))
|
||||
|
||||
if all_violations:
|
||||
print(
|
||||
"Boolean operator used to combine context managers in a `with` "
|
||||
"statement.\n"
|
||||
"Use `with a(), b():` or `with (a(), b()):` instead.\n"
|
||||
)
|
||||
for violation in all_violations:
|
||||
print(f" {violation}")
|
||||
return 1
|
||||
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
98
tools/check_forbidden_imports.py
Normal file
98
tools/check_forbidden_imports.py
Normal file
@@ -0,0 +1,98 @@
|
||||
#!/usr/bin/env python3
|
||||
#
|
||||
# Copyright (c) 2026 Huawei Technologies Co., Ltd. All Rights Reserved.
|
||||
# Copyright 2023 The vLLM team.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
# This file is a part of the vllm-ascend project.
|
||||
# Adapted from https://github.com/vllm-project/vllm/tree/main/tools
|
||||
#
|
||||
|
||||
import sys
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
import regex as re
|
||||
|
||||
|
||||
@dataclass
|
||||
class ForbiddenImport:
|
||||
pattern: str
|
||||
tip: str
|
||||
allowed_pattern: re.Pattern = re.compile(r"^$")
|
||||
allowed_files: set[str] = field(default_factory=set)
|
||||
|
||||
|
||||
CHECK_IMPORTS = {
|
||||
"pickle/cloudpickle": ForbiddenImport(
|
||||
pattern=(
|
||||
r"^\s*(import\s+(pickle|cloudpickle)(\s|$|\sas)"
|
||||
r"|from\s+(pickle|cloudpickle)\s+import\b)"
|
||||
),
|
||||
tip=("Avoid using pickle or cloudpickle or add this file to tools/check_forbidden_imports.py."),
|
||||
allowed_files={
|
||||
"vllm_ascend/distributed/kv_transfer/kv_pool/cpu_offload/metadata.py",
|
||||
"vllm_ascend/distributed/weight_transfer/npu_ipc_engine.py",
|
||||
"tests/ut/distributed/test_hccl_weight_transfer.py",
|
||||
},
|
||||
),
|
||||
"re": ForbiddenImport(
|
||||
pattern=r"^\s*(?:import\s+re(?:$|\s|,)|from\s+re\s+import)",
|
||||
tip="Replace 'import re' with 'import regex as re' or 'import regex'.",
|
||||
allowed_pattern=re.compile(r"^\s*import\s+regex(\s*|\s+as\s+re\s*)$"),
|
||||
),
|
||||
"triton": ForbiddenImport(
|
||||
pattern=r"^(from|import)\s+triton(\s|\.|$)",
|
||||
tip=("Use 'from vllm.triton_utils import triton'/'tl'."),
|
||||
allowed_pattern=re.compile(
|
||||
r"^\s*import\s+triton\.language\.extra\.cann\.extension\s+as\s+_extension_module(\s+#.*)?$"
|
||||
),
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
def check_file(path: str) -> int:
|
||||
try:
|
||||
with open(path, encoding="utf-8") as f:
|
||||
content = f.read()
|
||||
except (OSError, UnicodeDecodeError):
|
||||
return []
|
||||
|
||||
return_code = 0
|
||||
for import_name, forbidden_import in CHECK_IMPORTS.items():
|
||||
if path in forbidden_import.allowed_files:
|
||||
continue
|
||||
|
||||
for match in re.finditer(forbidden_import.pattern, content, re.MULTILINE):
|
||||
if forbidden_import.allowed_pattern.match(match.group()):
|
||||
continue
|
||||
|
||||
line_num = content[: match.start() + 1].count("\n") + 1
|
||||
print(
|
||||
f"{path}:{line_num}: "
|
||||
"\033[91merror:\033[0m "
|
||||
f"Found forbidden import: {import_name}. {forbidden_import.tip}"
|
||||
)
|
||||
return_code = 1
|
||||
|
||||
return return_code
|
||||
|
||||
|
||||
def main() -> int:
|
||||
return_code = 0
|
||||
for path in sys.argv[1:]:
|
||||
return_code |= check_file(path)
|
||||
return return_code
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
69
tools/check_logger.sh
Executable file
69
tools/check_logger.sh
Executable file
@@ -0,0 +1,69 @@
|
||||
#!/bin/bash
|
||||
#
|
||||
# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
# This file is a part of the vllm-ascend project.
|
||||
#
|
||||
# Check that vllm_ascend modules do not use init_logger(__name__).
|
||||
#
|
||||
# vllm's logging config registers a handler only for the "vllm" logger
|
||||
# namespace. Any logger created via init_logger(__name__) inside a
|
||||
# vllm_ascend module ends up in the "vllm_ascend.*" namespace, which has
|
||||
# no handler, so every log call is silently dropped.
|
||||
#
|
||||
# The correct pattern is:
|
||||
# from vllm.logger import logger
|
||||
#
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
REPO_ROOT="$(cd "$SCRIPT_DIR/.." && pwd)"
|
||||
PATCH_DIR="$REPO_ROOT/vllm_ascend/"
|
||||
|
||||
VIOLATIONS=0
|
||||
|
||||
for FILE in $(find "$PATCH_DIR" -type f -name "*.py" 2>/dev/null); do
|
||||
[[ -f "$FILE" ]] || continue
|
||||
|
||||
# Find lines that call init_logger(__name__)
|
||||
while IFS= read -r MATCH; do
|
||||
LINENUM=$(echo "$MATCH" | cut -d: -f1)
|
||||
LINE=$(echo "$MATCH" | cut -d: -f2-)
|
||||
if [[ $VIOLATIONS -eq 0 ]]; then
|
||||
echo ""
|
||||
fi
|
||||
echo " $FILE:$LINENUM: $LINE"
|
||||
VIOLATIONS=$(( VIOLATIONS + 1 ))
|
||||
done < <(grep -n 'init_logger[[:space:]]*([[:space:]]*__name__[[:space:]]*)' "$FILE" 2>/dev/null || true)
|
||||
done
|
||||
|
||||
if [[ $VIOLATIONS -gt 0 ]]; then
|
||||
echo ""
|
||||
echo "Found $VIOLATIONS violation(s): init_logger(__name__) must not be used in vllm_ascend modules."
|
||||
echo ""
|
||||
echo "vllm's logging handler is registered only for the 'vllm' namespace."
|
||||
echo "Loggers created with init_logger(__name__) inside vllm_ascend end up"
|
||||
echo "in the 'vllm_ascend.*' namespace, which has no handler — all log"
|
||||
echo "messages are silently dropped."
|
||||
echo ""
|
||||
echo "Fix: replace"
|
||||
echo " from vllm.logger import init_logger"
|
||||
echo " logger = init_logger(__name__)"
|
||||
echo "with"
|
||||
echo " from vllm.logger import logger"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
exit 0
|
||||
203
tools/check_long_functions.py
Executable file
203
tools/check_long_functions.py
Executable file
@@ -0,0 +1,203 @@
|
||||
#!/usr/bin/env python3
|
||||
#
|
||||
# Copyright (c) 2026 Huawei Technologies Co., Ltd. All Rights Reserved.
|
||||
# Copyright 2023 The vLLM team.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
# This file is a part of the vllm-ascend project.
|
||||
#
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import ast
|
||||
import subprocess
|
||||
import sys
|
||||
|
||||
MAX_FUNCTION_LINES = 100
|
||||
|
||||
|
||||
def _get_changed_lines(filepath: str) -> set[int]:
|
||||
"""Return added line numbers from staged git diff.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
filepath : str
|
||||
File path to inspect.
|
||||
|
||||
Returns
|
||||
-------
|
||||
set[int]
|
||||
1-indexed added line numbers from staged changes.
|
||||
|
||||
Notes
|
||||
-----
|
||||
If no staged diff exists (e.g. CI --all-files mode),
|
||||
an empty set is returned so that existing functions
|
||||
are not incorrectly flagged.
|
||||
"""
|
||||
try:
|
||||
result = subprocess.run(
|
||||
["git", "diff", "--cached", "--", filepath],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
check=True,
|
||||
)
|
||||
except (subprocess.CalledProcessError, FileNotFoundError):
|
||||
return set()
|
||||
|
||||
if not result.stdout.strip():
|
||||
return set()
|
||||
|
||||
changed: set[int] = set()
|
||||
current_line = 0
|
||||
|
||||
for line in result.stdout.splitlines():
|
||||
if line.startswith("@@"):
|
||||
# Example:
|
||||
# @@ -10,3 +20,8 @@
|
||||
parts = line.split()
|
||||
new_info = parts[2].lstrip("+")
|
||||
current_line = int(new_info.split(",")[0])
|
||||
|
||||
elif line.startswith("+") and not line.startswith("+++"):
|
||||
changed.add(current_line)
|
||||
current_line += 1
|
||||
|
||||
elif not line.startswith("-"):
|
||||
current_line += 1
|
||||
|
||||
return changed
|
||||
|
||||
|
||||
def _has_comment(source_lines: list[str], start_line_1: int, end_line_1: int) -> bool:
|
||||
"""Check whether a function contains Python comments.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
source_lines : list[str]
|
||||
Source code lines.
|
||||
start_line_1 : int
|
||||
1-indexed function start line.
|
||||
end_line_1 : int
|
||||
1-indexed function end line.
|
||||
|
||||
Returns
|
||||
-------
|
||||
bool
|
||||
True if any valid '#' comment exists.
|
||||
"""
|
||||
start_idx = max(start_line_1 - 1, 0)
|
||||
end_idx = min(end_line_1, len(source_lines))
|
||||
|
||||
for i in range(start_idx, end_idx):
|
||||
stripped = source_lines[i].strip()
|
||||
|
||||
if not stripped:
|
||||
continue
|
||||
# Full-line comment
|
||||
if stripped.startswith("#"):
|
||||
return True
|
||||
# Inline comment
|
||||
if "#" in stripped:
|
||||
in_string = False
|
||||
quote = ""
|
||||
for ch in stripped:
|
||||
if ch in ('"', "'"):
|
||||
if not in_string:
|
||||
in_string = True
|
||||
quote = ch
|
||||
elif ch == quote:
|
||||
in_string = False
|
||||
elif ch == "#" and not in_string:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def check_file(filepath: str, changed_lines: set[int]) -> list[str]:
|
||||
"""Check one Python file for undocumented long functions."""
|
||||
try:
|
||||
with open(filepath, encoding="utf-8") as f:
|
||||
source = f.read()
|
||||
except (OSError, UnicodeDecodeError):
|
||||
return []
|
||||
|
||||
try:
|
||||
tree = ast.parse(source, filename=filepath)
|
||||
except SyntaxError:
|
||||
return []
|
||||
|
||||
source_lines = source.splitlines()
|
||||
violations: list[str] = []
|
||||
|
||||
for node in ast.walk(tree):
|
||||
if not isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
|
||||
continue
|
||||
|
||||
# Only check newly added functions.
|
||||
if node.lineno not in changed_lines:
|
||||
continue
|
||||
if node.end_lineno is None:
|
||||
continue
|
||||
func_lines = node.end_lineno - node.lineno + 1
|
||||
if func_lines <= MAX_FUNCTION_LINES:
|
||||
continue
|
||||
# Has docstring
|
||||
if ast.get_docstring(node) is not None:
|
||||
continue
|
||||
# Has inline comments
|
||||
if _has_comment(source_lines, node.lineno, node.end_lineno):
|
||||
continue
|
||||
|
||||
violations.append(
|
||||
f"{filepath}:{node.lineno}: "
|
||||
f"function '{node.name}' "
|
||||
f"is {func_lines} lines "
|
||||
f"(>{MAX_FUNCTION_LINES}) "
|
||||
f"without comments or docstring"
|
||||
)
|
||||
return violations
|
||||
|
||||
|
||||
def main() -> int:
|
||||
if len(sys.argv) < 2:
|
||||
return 0
|
||||
|
||||
all_violations: list[str] = []
|
||||
for filepath in sys.argv[1:]:
|
||||
changed_lines = _get_changed_lines(filepath)
|
||||
|
||||
# No newly added lines in this file.
|
||||
if not changed_lines:
|
||||
continue
|
||||
|
||||
all_violations.extend(check_file(filepath, changed_lines))
|
||||
|
||||
if all_violations:
|
||||
print(
|
||||
"Functions longer than "
|
||||
f"{MAX_FUNCTION_LINES} lines "
|
||||
"must include comments or a docstring.\n"
|
||||
"Add a docstring or inline comments "
|
||||
"to explain the function logic.\n"
|
||||
)
|
||||
|
||||
for violation in all_violations:
|
||||
print(f" {violation}")
|
||||
|
||||
return 1
|
||||
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
@@ -17,6 +17,7 @@
|
||||
# Adapted from https://github.com/vllm-project/vllm/tree/main/tools
|
||||
#
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
|
||||
VLLM_ASCEND_SRC = "vllm_ascend"
|
||||
@@ -36,8 +37,8 @@ def check_init_file_in_package(directory):
|
||||
return False
|
||||
|
||||
# If any .py file exists, we expect an __init__.py
|
||||
if any(f.endswith('.py') for f in files):
|
||||
init_file = os.path.join(directory, '__init__.py')
|
||||
if any(f.endswith(".py") for f in files):
|
||||
init_file = os.path.join(directory, "__init__.py")
|
||||
if not os.path.isfile(init_file):
|
||||
return False
|
||||
return True
|
||||
@@ -45,12 +46,36 @@ def check_init_file_in_package(directory):
|
||||
|
||||
def find_missing_init_dirs(src_dir):
|
||||
"""
|
||||
Walk through the src_dir and return subdirectories missing __init__.py.
|
||||
Return tracked package directories in src_dir missing __init__.py.
|
||||
|
||||
Prefer git-tracked files so ignored/untracked generated directories are not
|
||||
treated as package violations.
|
||||
"""
|
||||
missing_init = set()
|
||||
for dirpath, _, _ in os.walk(src_dir):
|
||||
if not check_init_file_in_package(dirpath):
|
||||
missing_init.add(dirpath)
|
||||
|
||||
try:
|
||||
result = subprocess.run(
|
||||
["git", "ls-files", "-z", src_dir],
|
||||
check=True,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
)
|
||||
tracked_files = {f for f in result.stdout.split("\0") if f}
|
||||
except (subprocess.CalledProcessError, FileNotFoundError):
|
||||
tracked_files = None
|
||||
|
||||
if tracked_files is not None:
|
||||
candidate_dirs = {os.path.dirname(f) for f in tracked_files if f.endswith(".py")}
|
||||
for dirpath in candidate_dirs:
|
||||
init_path = f"{dirpath}/__init__.py" if dirpath else "__init__.py"
|
||||
if init_path not in tracked_files:
|
||||
missing_init.add(dirpath)
|
||||
|
||||
else:
|
||||
for dirpath, _, _ in os.walk(src_dir):
|
||||
if not check_init_file_in_package(dirpath):
|
||||
missing_init.add(dirpath)
|
||||
|
||||
return missing_init
|
||||
|
||||
|
||||
@@ -62,14 +87,12 @@ def main():
|
||||
all_missing.update(missing)
|
||||
|
||||
if all_missing:
|
||||
print(
|
||||
"❌ Missing '__init__.py' files in the following Python package directories:"
|
||||
)
|
||||
print("❌ Missing '__init__.py' files in the following Python package directories:")
|
||||
for pkg in sorted(all_missing):
|
||||
print(f" - {pkg}")
|
||||
sys.exit(1)
|
||||
else:
|
||||
print("✅ All Python packages have __init__.py files.")
|
||||
print("All Python packages have __init__.py files.")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
#!/bin/bash
|
||||
#!/usr/bin/env bash
|
||||
|
||||
#
|
||||
# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
|
||||
|
||||
425
tools/collect_user_first_contribution.sh
Executable file
425
tools/collect_user_first_contribution.sh
Executable file
@@ -0,0 +1,425 @@
|
||||
#!/usr/bin/env bash
|
||||
|
||||
#
|
||||
# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
|
||||
# Copyright 2023 The vLLM team.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
# This file is a part of the vllm-ascend project.
|
||||
# Adapted from https://github.com/vllm-project/vllm/tree/main/tools
|
||||
#
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
# Default configuration
|
||||
DEFAULT_REPO="vllm-project/vllm-ascend"
|
||||
DEFAULT_CONTRIBUTORS_FILE="docs/source/community/contributors.md"
|
||||
|
||||
function usage() {
|
||||
echo "This script collects contributors' first contributions and updates the contributors.md file."
|
||||
echo "Supports incremental updates by tracking the last commit hash."
|
||||
echo ""
|
||||
echo "Please set the environment variable GITHUB_TOKEN with repo read permission."
|
||||
echo "Refer to https://docs.github.com/en/rest/authentication/authenticating-to-the-rest-api?apiVersion=2022-11-28"
|
||||
echo ""
|
||||
echo "Usage: $0 [options]"
|
||||
echo " $0 --full # Force full refresh (ignore last commit hash)"
|
||||
echo " $0 --help"
|
||||
echo ""
|
||||
echo "Options:"
|
||||
echo " --full Force full refresh, recalculate all contributors"
|
||||
echo " --repo=OWNER/REPO Specify GitHub repository (default: ${DEFAULT_REPO})"
|
||||
echo " --file=PATH Specify contributors.md path (default: ${DEFAULT_CONTRIBUTORS_FILE})"
|
||||
echo ""
|
||||
echo "Examples:"
|
||||
echo " $0 # Incremental update from last commit"
|
||||
echo " $0 --full # Full refresh"
|
||||
}
|
||||
|
||||
# Parse arguments
|
||||
REPO="${DEFAULT_REPO}"
|
||||
CONTRIBUTORS_FILE="${DEFAULT_CONTRIBUTORS_FILE}"
|
||||
FORCE_FULL=false
|
||||
|
||||
for arg in "$@"; do
|
||||
case $arg in
|
||||
--help)
|
||||
usage
|
||||
exit 0
|
||||
;;
|
||||
--full)
|
||||
FORCE_FULL=true
|
||||
shift
|
||||
;;
|
||||
--repo=*)
|
||||
REPO="${arg#*=}"
|
||||
shift
|
||||
;;
|
||||
--file=*)
|
||||
CONTRIBUTORS_FILE="${arg#*=}"
|
||||
shift
|
||||
;;
|
||||
*)
|
||||
echo "Unknown argument: $arg"
|
||||
usage
|
||||
exit 1
|
||||
;;
|
||||
esac
|
||||
done
|
||||
|
||||
GITHUB_TOKEN="${GITHUB_TOKEN:-}"
|
||||
|
||||
if [ -z "$GITHUB_TOKEN" ]; then
|
||||
echo "Error: Please set the environment variable GITHUB_TOKEN with repo read permission."
|
||||
echo "Refer to https://docs.github.com/en/rest/authentication/authenticating-to-the-rest-api?apiVersion=2022-11-28"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
# Get the script directory to find the project root
|
||||
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
PROJECT_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd)"
|
||||
|
||||
# Resolve contributors file path
|
||||
if [[ "$CONTRIBUTORS_FILE" != /* ]]; then
|
||||
CONTRIBUTORS_FILE="${PROJECT_ROOT}/${CONTRIBUTORS_FILE}"
|
||||
fi
|
||||
|
||||
if [ ! -f "$CONTRIBUTORS_FILE" ]; then
|
||||
echo "Error: Contributors file not found: ${CONTRIBUTORS_FILE}"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
# Change to project root for git operations
|
||||
cd "$PROJECT_ROOT"
|
||||
|
||||
# Get current HEAD commit hash
|
||||
CURRENT_HEAD=$(git rev-parse HEAD)
|
||||
CURRENT_HEAD_SHORT="${CURRENT_HEAD:0:7}"
|
||||
|
||||
echo "Repository: ${REPO}"
|
||||
echo "Contributors file: ${CONTRIBUTORS_FILE}"
|
||||
echo "Current HEAD: ${CURRENT_HEAD_SHORT}"
|
||||
echo ""
|
||||
|
||||
# Function to extract last commit hash from contributors file
|
||||
get_last_commit_hash() {
|
||||
local file="$1"
|
||||
# Look for comment line with last commit hash: <!-- last_commit: abc1234 -->
|
||||
grep -o '<!-- last_commit: [a-f0-9]* -->' "$file" 2>/dev/null | sed 's/<!-- last_commit: \([a-f0-9]*\) -->/\1/' || echo ""
|
||||
}
|
||||
|
||||
# Function to extract current contributor count from file
|
||||
get_current_contributor_count() {
|
||||
local file="$1"
|
||||
# Find the first row number in the table (most recent contributor)
|
||||
grep -o '| [0-9]* |' "$file" 2>/dev/null | head -1 | grep -o '[0-9]*' || echo "0"
|
||||
}
|
||||
|
||||
# Function to extract GitHub login from noreply email
|
||||
# Format: ID+username@users.noreply.github.com or username@users.noreply.github.com
|
||||
extract_login_from_noreply_email() {
|
||||
local email="$1"
|
||||
if [[ "$email" == *@users.noreply.github.com ]]; then
|
||||
# Remove the domain part
|
||||
local local_part="${email%@users.noreply.github.com}"
|
||||
# Check if it's in format "ID+username" or just "username"
|
||||
if [[ "$local_part" == *+* ]]; then
|
||||
# Format: ID+username -> extract username
|
||||
echo "${local_part#*+}"
|
||||
else
|
||||
# Format: username
|
||||
echo "$local_part"
|
||||
fi
|
||||
else
|
||||
echo ""
|
||||
fi
|
||||
}
|
||||
|
||||
# Function to get GitHub login for a commit
|
||||
get_github_login() {
|
||||
local sha="$1"
|
||||
local email="$2"
|
||||
local api_url="https://api.github.com/repos/${REPO}/commits/${sha}"
|
||||
local resp
|
||||
resp=$(curl -s -H "Authorization: token ${GITHUB_TOKEN}" -H "Accept: application/vnd.github.v3+json" "$api_url")
|
||||
local login
|
||||
login=$(echo "$resp" | jq -r '.author.login // empty' 2>/dev/null || echo "")
|
||||
|
||||
# If no login from API, try to extract from noreply email
|
||||
if [ -z "$login" ]; then
|
||||
login=$(extract_login_from_noreply_email "$email")
|
||||
fi
|
||||
|
||||
echo "$login"
|
||||
}
|
||||
|
||||
# Check if we should do incremental update
|
||||
LAST_COMMIT=""
|
||||
INCREMENTAL=false
|
||||
|
||||
if [ "$FORCE_FULL" = false ]; then
|
||||
LAST_COMMIT=$(get_last_commit_hash "$CONTRIBUTORS_FILE")
|
||||
if [ -n "$LAST_COMMIT" ] && [ "$LAST_COMMIT" != "$CURRENT_HEAD" ]; then
|
||||
# Check if LAST_COMMIT is an ancestor of CURRENT_HEAD
|
||||
if git merge-base --is-ancestor "$LAST_COMMIT" "$CURRENT_HEAD" 2>/dev/null; then
|
||||
INCREMENTAL=true
|
||||
echo "Incremental update from commit: ${LAST_COMMIT:0:7}"
|
||||
else
|
||||
echo "Warning: Last commit ${LAST_COMMIT:0:7} is not an ancestor of current HEAD, doing full refresh."
|
||||
fi
|
||||
elif [ "$LAST_COMMIT" = "$CURRENT_HEAD" ]; then
|
||||
echo "Already up to date (HEAD matches last recorded commit)."
|
||||
echo "Use --full to force a full refresh."
|
||||
exit 0
|
||||
fi
|
||||
fi
|
||||
|
||||
if [ "$INCREMENTAL" = true ]; then
|
||||
# Incremental update: get new commits since last commit
|
||||
echo ""
|
||||
echo "Fetching new commits..."
|
||||
|
||||
# Get all commits in time order, format: sha|email|name|date
|
||||
ALLCOMMITS=$(mktemp)
|
||||
git log --pretty=format:'%H|%aE|%aN|%cI' --reverse "${LAST_COMMIT}..${CURRENT_HEAD}" > "$ALLCOMMITS"
|
||||
|
||||
# Get the first commit for each author email (from all history, but we'll filter to new ones)
|
||||
ALL_HISTORY=$(mktemp)
|
||||
git log --pretty=format:'%H|%aE|%aN|%cI' --reverse --all > "$ALL_HISTORY"
|
||||
|
||||
# First commit by email (from all history)
|
||||
FIRST_BY_EMAIL=$(mktemp)
|
||||
awk -F'|' '!seen[$2]++ {print $2 "|" $1 "|" $4 "|" $3}' "$ALL_HISTORY" > "$FIRST_BY_EMAIL"
|
||||
|
||||
# New SHAs in this range
|
||||
NEW_SHAS=$(mktemp)
|
||||
git rev-list "${LAST_COMMIT}..${CURRENT_HEAD}" > "$NEW_SHAS"
|
||||
|
||||
# Extract existing contributor logins from the file for deduplication
|
||||
EXISTING_LOGINS=$(mktemp)
|
||||
grep -oE '\[@[^]]+\]' "$CONTRIBUTORS_FILE" 2>/dev/null | sed 's/\[@//;s/\]//' | sort -u > "$EXISTING_LOGINS" || true
|
||||
|
||||
# Collect new contributors (first commit is in the new range)
|
||||
NEW_CONTRIBUTORS=$(mktemp)
|
||||
count=0
|
||||
skipped=0
|
||||
|
||||
while IFS='|' read -r email sha date name; do
|
||||
if grep -Fxq "$sha" "$NEW_SHAS"; then
|
||||
# Query GitHub API
|
||||
login=$(get_github_login "$sha" "$email")
|
||||
|
||||
# Skip if no GitHub login
|
||||
if [ -z "$login" ]; then
|
||||
continue
|
||||
fi
|
||||
|
||||
# Check if contributor already exists (deduplication)
|
||||
if grep -Fxq "$login" "$EXISTING_LOGINS"; then
|
||||
echo "Skipping duplicate contributor: $login"
|
||||
((skipped++)) || true
|
||||
continue
|
||||
fi
|
||||
|
||||
# Format date
|
||||
formatted_date=$(echo "$date" | cut -d'T' -f1 | tr '-' '/')
|
||||
short_sha="${sha:0:7}"
|
||||
|
||||
echo "${login}|${sha}|${short_sha}|${formatted_date}" >> "$NEW_CONTRIBUTORS"
|
||||
((count++)) || true
|
||||
fi
|
||||
done < "$FIRST_BY_EMAIL"
|
||||
|
||||
NEW_COUNT=$(wc -l < "$NEW_CONTRIBUTORS" | tr -d ' ')
|
||||
echo "Found ${NEW_COUNT} new contributors"
|
||||
if [ "$skipped" -gt 0 ]; then
|
||||
echo "Skipped ${skipped} duplicate contributors"
|
||||
fi
|
||||
|
||||
if [ "$NEW_COUNT" -eq 0 ]; then
|
||||
echo "No new contributors found."
|
||||
rm -f "$ALLCOMMITS" "$ALL_HISTORY" "$FIRST_BY_EMAIL" "$NEW_SHAS" "$NEW_CONTRIBUTORS"
|
||||
exit 0
|
||||
fi
|
||||
|
||||
# Get current contributor count
|
||||
CURRENT_COUNT=$(get_current_contributor_count "$CONTRIBUTORS_FILE")
|
||||
echo "Current contributor count: ${CURRENT_COUNT}"
|
||||
|
||||
# Generate new rows (sorted by date descending, newest gets highest number)
|
||||
NEW_ROWS=$(mktemp)
|
||||
sort -t'|' -k4 -r "$NEW_CONTRIBUTORS" | awk -F'|' -v start="$CURRENT_COUNT" -v new_count="$NEW_COUNT" -v repo="$REPO" '
|
||||
BEGIN { nr = start + new_count }
|
||||
{
|
||||
login = $1
|
||||
sha = $2
|
||||
short_sha = $3
|
||||
date = $4
|
||||
|
||||
# All contributors now have GitHub login
|
||||
printf "| %d | [@%s](https://github.com/%s) | %s | [%s](https://github.com/%s/commit/%s) |\n", nr, login, login, date, short_sha, repo, sha
|
||||
nr--
|
||||
}' > "$NEW_ROWS"
|
||||
|
||||
# Update the file
|
||||
TEMP_FILE=$(mktemp)
|
||||
CURRENT_DATE=$(date +%Y-%m-%d)
|
||||
NEW_TOTAL=$((CURRENT_COUNT + NEW_COUNT))
|
||||
|
||||
# Track if we just wrote the table header (to insert new rows after separator)
|
||||
WROTE_HEADER=false
|
||||
# Track if we are skipping old header lines (between <!-- last_commit and | Number |)
|
||||
SKIP_OLD_HEADER=false
|
||||
|
||||
while IFS= read -r line || [ -n "$line" ]; do
|
||||
if [[ "$line" == "<!-- last_commit:"* ]]; then
|
||||
# Start skipping old header lines
|
||||
SKIP_OLD_HEADER=true
|
||||
continue
|
||||
elif [[ "$SKIP_OLD_HEADER" == true && "$line" != "| Number | Contributor | Date | Commit ID |" ]]; then
|
||||
# Skip all old header lines (Updated on, Every release, empty lines)
|
||||
continue
|
||||
elif [[ "$line" == "| Number | Contributor | Date | Commit ID |" ]]; then
|
||||
SKIP_OLD_HEADER=false
|
||||
# Insert new content before the table header
|
||||
echo "<!-- last_commit: ${CURRENT_HEAD} -->" >> "$TEMP_FILE"
|
||||
echo "" >> "$TEMP_FILE"
|
||||
echo "Every release of vLLM Ascend would not have been possible without the following contributors:" >> "$TEMP_FILE"
|
||||
echo "" >> "$TEMP_FILE"
|
||||
echo "Updated on ${CURRENT_DATE}:" >> "$TEMP_FILE"
|
||||
echo "" >> "$TEMP_FILE"
|
||||
echo "$line" >> "$TEMP_FILE"
|
||||
WROTE_HEADER=true
|
||||
elif [[ "$WROTE_HEADER" == true && "$line" == "|:"* ]]; then
|
||||
# This is the separator line after header - write it, then insert new rows
|
||||
echo "$line" >> "$TEMP_FILE"
|
||||
cat "$NEW_ROWS" >> "$TEMP_FILE"
|
||||
WROTE_HEADER=false
|
||||
else
|
||||
# Existing rows keep their original numbers (new rows are inserted above)
|
||||
echo "$line" >> "$TEMP_FILE"
|
||||
fi
|
||||
done < "$CONTRIBUTORS_FILE"
|
||||
|
||||
mv "$TEMP_FILE" "$CONTRIBUTORS_FILE"
|
||||
|
||||
echo ""
|
||||
echo "Done! Added ${NEW_COUNT} new contributors. Total: ${NEW_TOTAL}"
|
||||
|
||||
# Cleanup
|
||||
rm -f "$ALLCOMMITS" "$ALL_HISTORY" "$FIRST_BY_EMAIL" "$NEW_SHAS" "$NEW_CONTRIBUTORS" "$NEW_ROWS" "$EXISTING_LOGINS"
|
||||
|
||||
else
|
||||
# Full refresh
|
||||
echo "Performing full refresh..."
|
||||
echo ""
|
||||
|
||||
# All commits in time order
|
||||
ALLCOMMITS=$(mktemp)
|
||||
git log --pretty=format:'%H|%aE|%aN|%cI' --reverse --all > "$ALLCOMMITS"
|
||||
|
||||
# First commit by email
|
||||
FIRST_BY_EMAIL=$(mktemp)
|
||||
awk -F'|' '!seen[$2]++ {print $2 "|" $1 "|" $4 "|" $3}' "$ALLCOMMITS" > "$FIRST_BY_EMAIL"
|
||||
|
||||
# Collect all contributors
|
||||
CONTRIBUTORS_DATA=$(mktemp)
|
||||
TOTAL=$(wc -l < "$FIRST_BY_EMAIL" | tr -d ' ')
|
||||
CURRENT=0
|
||||
|
||||
echo "Processing ${TOTAL} contributors..."
|
||||
|
||||
while IFS='|' read -r email sha date name; do
|
||||
CURRENT=$((CURRENT + 1))
|
||||
printf "\rProcessing: %d/%d" "$CURRENT" "$TOTAL"
|
||||
|
||||
login=$(get_github_login "$sha" "$email")
|
||||
formatted_date=$(echo "$date" | cut -d'T' -f1 | tr '-' '/')
|
||||
short_sha="${sha:0:7}"
|
||||
|
||||
if [ -n "$login" ]; then
|
||||
echo "${login}|${sha}|${short_sha}|${formatted_date}" >> "$CONTRIBUTORS_DATA"
|
||||
fi
|
||||
# Skip contributors without GitHub login (cannot be linked to GitHub ID)
|
||||
done < "$FIRST_BY_EMAIL"
|
||||
|
||||
echo ""
|
||||
echo ""
|
||||
|
||||
# Deduplicate by GitHub login (same user may have multiple emails)
|
||||
# Keep the earliest commit (first occurrence) for each login
|
||||
DEDUPED_DATA=$(mktemp)
|
||||
awk -F'|' '!seen[$1]++' "$CONTRIBUTORS_DATA" > "$DEDUPED_DATA"
|
||||
mv "$DEDUPED_DATA" "$CONTRIBUTORS_DATA"
|
||||
|
||||
CONTRIBUTOR_COUNT=$(wc -l < "$CONTRIBUTORS_DATA" | tr -d ' ')
|
||||
echo "Found ${CONTRIBUTOR_COUNT} unique contributors"
|
||||
|
||||
# Generate new content
|
||||
NEW_SECTION=$(mktemp)
|
||||
CURRENT_DATE=$(date +%Y-%m-%d)
|
||||
|
||||
{
|
||||
echo "<!-- last_commit: ${CURRENT_HEAD} -->"
|
||||
echo ""
|
||||
echo "Every release of vLLM Ascend would not have been possible without the following contributors:"
|
||||
echo ""
|
||||
echo "Updated on ${CURRENT_DATE}:"
|
||||
echo ""
|
||||
echo "| Number | Contributor | Date | Commit ID |"
|
||||
echo "|:------:|:-----------:|:-----:|:---------:|"
|
||||
|
||||
sort -t'|' -k4 -r "$CONTRIBUTORS_DATA" | awk -F'|' -v total="$CONTRIBUTOR_COUNT" -v repo="$REPO" '
|
||||
BEGIN { nr = total }
|
||||
{
|
||||
login = $1
|
||||
sha = $2
|
||||
short_sha = $3
|
||||
date = $4
|
||||
|
||||
# All contributors now have GitHub login
|
||||
printf "| %d | [@%s](https://github.com/%s) | %s | [%s](https://github.com/%s/commit/%s) |\n", nr, login, login, date, short_sha, repo, sha
|
||||
nr--
|
||||
}'
|
||||
} > "$NEW_SECTION"
|
||||
|
||||
# Update the file
|
||||
TEMP_FILE=$(mktemp)
|
||||
FOUND_CONTRIBUTORS=false
|
||||
|
||||
while IFS= read -r line || [ -n "$line" ]; do
|
||||
if [[ "$line" == "## Contributors" ]]; then
|
||||
FOUND_CONTRIBUTORS=true
|
||||
echo "$line" >> "$TEMP_FILE"
|
||||
cat "$NEW_SECTION" >> "$TEMP_FILE"
|
||||
break
|
||||
else
|
||||
echo "$line" >> "$TEMP_FILE"
|
||||
fi
|
||||
done < "$CONTRIBUTORS_FILE"
|
||||
|
||||
if ! $FOUND_CONTRIBUTORS; then
|
||||
echo "" >> "$TEMP_FILE"
|
||||
echo "## Contributors" >> "$TEMP_FILE"
|
||||
cat "$NEW_SECTION" >> "$TEMP_FILE"
|
||||
echo ""
|
||||
echo "Warning: '## Contributors' section not found, appended at the end."
|
||||
fi
|
||||
|
||||
mv "$TEMP_FILE" "$CONTRIBUTORS_FILE"
|
||||
|
||||
echo "Done! Contributors list has been updated in: ${CONTRIBUTORS_FILE}"
|
||||
|
||||
# Cleanup
|
||||
rm -f "$ALLCOMMITS" "$FIRST_BY_EMAIL" "$CONTRIBUTORS_DATA" "$NEW_SECTION"
|
||||
fi
|
||||
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()
|
||||
1
tools/docs_i18n/__init__.py
Normal file
1
tools/docs_i18n/__init__.py
Normal file
@@ -0,0 +1 @@
|
||||
"""Sphinx helpers for documentation localization."""
|
||||
71
tools/docs_i18n/sphinx_extension.py
Normal file
71
tools/docs_i18n/sphinx_extension.py
Normal file
@@ -0,0 +1,71 @@
|
||||
from docutils import nodes
|
||||
from sphinx.errors import SphinxError
|
||||
from sphinx.transforms import SphinxTransform
|
||||
|
||||
|
||||
def _is_inside_tab_content(node: nodes.Node) -> bool:
|
||||
"""Return whether a node is inside a sphinx-design tab content node."""
|
||||
parent = node.parent
|
||||
while parent is not None:
|
||||
if isinstance(parent, nodes.Element) and parent.get("design_component") == "tab-content":
|
||||
return True
|
||||
parent = parent.parent
|
||||
return False
|
||||
|
||||
|
||||
class RestoreTabTableCellSource(SphinxTransform):
|
||||
"""Restore source metadata for tables nested in sphinx-design tabs."""
|
||||
|
||||
# Run before Sphinx's PreserveTranslatableMessages (10) and Locale (20)
|
||||
# transforms so both gettext extraction and localized HTML see the cells.
|
||||
default_priority = 5
|
||||
|
||||
def apply(self, **kwargs) -> None:
|
||||
for table in self.document.findall(nodes.table):
|
||||
if not _is_inside_tab_content(table):
|
||||
continue
|
||||
|
||||
source = table.source or self.document.get("source")
|
||||
line = table.line or 0
|
||||
for paragraph in table.findall(nodes.paragraph):
|
||||
if not paragraph.source:
|
||||
paragraph.source = source
|
||||
if paragraph.line is None:
|
||||
paragraph.line = line
|
||||
|
||||
|
||||
class ValidateTabTableCellTranslations(SphinxTransform):
|
||||
"""Fail a localized build when a table in a tab was not translated."""
|
||||
|
||||
# Run after Sphinx's Locale transform (20), which sets the ``translated``
|
||||
# attribute after applying the message catalog to each translatable node.
|
||||
default_priority = 30
|
||||
|
||||
def apply(self, **kwargs) -> None:
|
||||
if not self.config.validate_tab_table_translations:
|
||||
return
|
||||
|
||||
untranslated = []
|
||||
for table in self.document.findall(nodes.table):
|
||||
if not _is_inside_tab_content(table):
|
||||
continue
|
||||
|
||||
for paragraph in table.findall(nodes.paragraph):
|
||||
# Complex tables may contain structural paragraphs without
|
||||
# visible text. Sphinx does not extract or translate them.
|
||||
if paragraph.astext().strip() and not paragraph.get("translated", False):
|
||||
untranslated.append(f"{paragraph.source}:{paragraph.line}: {paragraph.astext()!r}")
|
||||
|
||||
if untranslated:
|
||||
details = "\n".join(f"- {item}" for item in untranslated)
|
||||
raise SphinxError("localized table cells nested in tabs were not translated:\n" + details)
|
||||
|
||||
|
||||
def setup(app):
|
||||
app.add_config_value("validate_tab_table_translations", False, "env")
|
||||
app.add_transform(RestoreTabTableCellSource)
|
||||
app.add_transform(ValidateTabTableCellTranslations)
|
||||
return {
|
||||
"parallel_read_safe": True,
|
||||
"parallel_write_safe": True,
|
||||
}
|
||||
87
tools/format_contributors.py
Normal file
87
tools/format_contributors.py
Normal file
@@ -0,0 +1,87 @@
|
||||
# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
|
||||
# Copyright 2023 The vLLM team.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
# This file is a part of the vllm-ascend project.
|
||||
# Adapted from https://github.com/vllm-project/vllm/tree/main/tools
|
||||
#
|
||||
import argparse
|
||||
import sys
|
||||
from datetime import datetime
|
||||
|
||||
import regex as re
|
||||
|
||||
p = re.compile(r"@(?P<user>[A-Za-z0-9-_]+)[^\`]*\`(?P<sha>[0-9a-fA-F]+)\`\s*[-–—]\s*(?P<date>.+)$")
|
||||
|
||||
|
||||
def parse_lines(lines):
|
||||
items = []
|
||||
for ln in lines:
|
||||
ln = ln.strip()
|
||||
if not ln:
|
||||
continue
|
||||
m = p.search(ln)
|
||||
if not m:
|
||||
continue
|
||||
user = m.group("user")
|
||||
sha = m.group("sha")
|
||||
datestr = m.group("date").strip()
|
||||
try:
|
||||
dt = datetime.fromisoformat(datestr)
|
||||
except Exception:
|
||||
# fallback: try to parse common formats
|
||||
try:
|
||||
dt = datetime.strptime(datestr, "%Y/%m/%d")
|
||||
except Exception:
|
||||
continue
|
||||
items.append((dt, user, sha, datestr))
|
||||
return items
|
||||
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser(
|
||||
description="Format and sort contributor lines by date (newest first). Outputs markdown table by default."
|
||||
)
|
||||
ap.add_argument(
|
||||
"file", nargs="?", help="input file (default stdin), output from collect_user_first_contribution.sh"
|
||||
)
|
||||
ap.add_argument("--start", type=int, default=1, help="minimum number for table (oldest row will have this number)")
|
||||
ap.add_argument("--repo", default="vllm-project/vllm-ascend", help="repo used for commit links")
|
||||
args = ap.parse_args()
|
||||
|
||||
if args.file:
|
||||
with open(args.file, encoding="utf-8") as f:
|
||||
lines = f.readlines()
|
||||
else:
|
||||
lines = sys.stdin.readlines()
|
||||
|
||||
items = parse_lines(lines)
|
||||
# sort newest first
|
||||
items.sort(key=lambda x: x[0], reverse=True)
|
||||
|
||||
# Outputs markdown table (sorted by date), the minimum number is args.start
|
||||
count = len(items)
|
||||
if count == 0:
|
||||
return
|
||||
n = args.start + count - 1
|
||||
for dt, user, sha, datestr in items:
|
||||
short = sha[:7]
|
||||
date_short = dt.strftime("%Y/%m/%d")
|
||||
user_url = f"https://github.com/{user}"
|
||||
commit_url = f"https://github.com/{args.repo}/commit/{sha}"
|
||||
print(f"| {n} | [@{user}]({user_url}) | {date_short} | [{short}]({commit_url}) |")
|
||||
n -= 1
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
340
tools/mooncake_installer.sh
Normal file
340
tools/mooncake_installer.sh
Normal file
@@ -0,0 +1,340 @@
|
||||
#!/usr/bin/env bash
|
||||
|
||||
#
|
||||
# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
|
||||
# Copyright 2023 The vLLM team.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
# This file is a part of the vllm-ascend project.
|
||||
# This is script is inspired from https://github.com/kvcache-ai/Mooncake/blob/main/dependencies.sh
|
||||
#
|
||||
|
||||
# Color definitions
|
||||
GREEN="\033[0;32m"
|
||||
BLUE="\033[0;34m"
|
||||
YELLOW="\033[0;33m"
|
||||
RED="\033[0;31m"
|
||||
NC="\033[0m" # No Color
|
||||
|
||||
# Configuration
|
||||
REPO_ROOT=`pwd`
|
||||
GITHUB_PROXY=${GITHUB_PROXY:-"https://github.com"}
|
||||
GOVER=1.23.8
|
||||
YALANTINGLIBS_VERSION=0.5.6
|
||||
|
||||
# Function to print section headers
|
||||
print_section() {
|
||||
echo -e "\n${BLUE}=== $1 ===${NC}"
|
||||
}
|
||||
|
||||
# Function to print success messages
|
||||
print_success() {
|
||||
echo -e "${GREEN}✓ $1${NC}"
|
||||
}
|
||||
|
||||
# Function to print error messages and exit
|
||||
print_error() {
|
||||
echo -e "${RED}✗ ERROR: $1${NC}"
|
||||
exit 1
|
||||
}
|
||||
|
||||
# Function to check command success
|
||||
check_success() {
|
||||
if [ $? -ne 0 ]; then
|
||||
print_error "$1"
|
||||
fi
|
||||
}
|
||||
|
||||
|
||||
# Parse command line arguments
|
||||
SKIP_CONFIRM=false
|
||||
for arg in "$@"; do
|
||||
case $arg in
|
||||
-y|--yes)
|
||||
SKIP_CONFIRM=true
|
||||
;;
|
||||
-h|--help)
|
||||
echo -e "${YELLOW}Mooncake Dependencies Installer${NC}"
|
||||
echo -e "Usage: ./dependencies.sh [OPTIONS]"
|
||||
echo -e "\nOptions:"
|
||||
echo -e " -y, --yes Skip confirmation and install all dependencies"
|
||||
echo -e " -h, --help Show this help message and exit"
|
||||
exit 0
|
||||
;;
|
||||
esac
|
||||
done
|
||||
|
||||
# Print welcome message
|
||||
echo -e "${YELLOW}Mooncake Dependencies Installer${NC}"
|
||||
echo -e "This script will install all required dependencies for Mooncake."
|
||||
echo -e "The following components will be installed:"
|
||||
echo -e " - System packages (build tools, libraries)"
|
||||
echo -e " - yalantinglibs"
|
||||
echo -e " - Git submodules"
|
||||
echo -e " - Go $GOVER"
|
||||
echo
|
||||
|
||||
# Ask for confirmation unless -y flag is used
|
||||
if [ "$SKIP_CONFIRM" = false ]; then
|
||||
read -p "Do you want to continue? [Y/n] " -n 1 -r
|
||||
echo
|
||||
if [[ ! $REPLY =~ ^[Yy]$ ]] && [[ ! $REPLY = "" ]]; then
|
||||
echo -e "${YELLOW}Installation cancelled.${NC}"
|
||||
exit 0
|
||||
fi
|
||||
fi
|
||||
|
||||
# Define a function to handle the git clone operation
|
||||
clone_repo_if_not_exists() {
|
||||
local repo_dir=$1
|
||||
local repo_url=$2
|
||||
|
||||
if [ ! -d "$repo_dir" ]; then
|
||||
git clone --depth 1 "$repo_url"
|
||||
else
|
||||
echo "Directory $repo_dir already exists, skipping clone."
|
||||
fi
|
||||
}
|
||||
|
||||
|
||||
# Update package lists
|
||||
print_section "Updating package lists"
|
||||
check_success "Failed to update package lists"
|
||||
|
||||
# Install system packages
|
||||
print_section "Installing system packages"
|
||||
echo -e "${YELLOW}This may take a few minutes...${NC}"
|
||||
|
||||
# System detection and dependency installation
|
||||
if command -v apt-get &> /dev/null; then
|
||||
echo "Detected apt-get. Using Debian-based package manager."
|
||||
apt-get update
|
||||
apt-get install -y build-essential \
|
||||
cmake \
|
||||
git \
|
||||
wget \
|
||||
unzip \
|
||||
libibverbs-dev \
|
||||
libgoogle-glog-dev \
|
||||
libgtest-dev \
|
||||
libjsoncpp-dev \
|
||||
libunwind-dev \
|
||||
libnuma-dev \
|
||||
libpython3-dev \
|
||||
libboost-all-dev \
|
||||
libssl-dev \
|
||||
libgrpc-dev \
|
||||
libgrpc++-dev \
|
||||
libprotobuf-dev \
|
||||
libyaml-cpp-dev \
|
||||
protobuf-compiler-grpc \
|
||||
libcurl4-openssl-dev \
|
||||
libhiredis-dev \
|
||||
pkg-config \
|
||||
patchelf \
|
||||
mpich \
|
||||
libmpich-dev
|
||||
apt purge -y openmpi-bin libopenmpi-dev || true
|
||||
elif command -v yum &> /dev/null; then
|
||||
echo "Detected yum. Using Red Hat-based package manager."
|
||||
yum makecache
|
||||
yum install -y \
|
||||
gcc \
|
||||
gcc-c++ \
|
||||
make \
|
||||
cmake \
|
||||
unzip \
|
||||
git \
|
||||
wget \
|
||||
libibverbs-devel \
|
||||
numactl-devel \
|
||||
gflags-devel \
|
||||
glog-devel \
|
||||
gtest \
|
||||
gtest-devel \
|
||||
jsoncpp-devel \
|
||||
mpich \
|
||||
mpich-devel \
|
||||
boost-devel \
|
||||
openssl-devel \
|
||||
hiredis-devel \
|
||||
python3-devel \
|
||||
curl-devel \
|
||||
patchelf
|
||||
|
||||
# install yaml-cpp
|
||||
cd "${REPO_ROOT}/thirdparties"
|
||||
clone_repo_if_not_exists "yaml-cpp" https://github.com/jbeder/yaml-cpp.git
|
||||
cd yaml-cpp || exit
|
||||
rm -rf build
|
||||
mkdir -p build && cd build
|
||||
cmake ..
|
||||
make -j$(nproc)
|
||||
make install
|
||||
cd "${REPO_ROOT}"
|
||||
else
|
||||
echo "Unsupported package manager. Please install the dependencies manually."
|
||||
exit 1
|
||||
fi
|
||||
|
||||
check_success "Failed to install system packages"
|
||||
print_success "System packages installed successfully"
|
||||
|
||||
# Install yalantinglibs
|
||||
print_section "Installing yalantinglibs"
|
||||
|
||||
# Check if thirdparties directory exists
|
||||
if [ ! -d "${REPO_ROOT}/thirdparties" ]; then
|
||||
mkdir -p "${REPO_ROOT}/thirdparties"
|
||||
check_success "Failed to create thirdparties directory"
|
||||
fi
|
||||
|
||||
# Change to thirdparties directory
|
||||
cd "${REPO_ROOT}/thirdparties"
|
||||
check_success "Failed to change to thirdparties directory"
|
||||
|
||||
# Check if yalantinglibs is already installed
|
||||
if [ -d "yalantinglibs-${YALANTINGLIBS_VERSION}" ]; then
|
||||
echo -e "${YELLOW}yalantinglibs-${YALANTINGLIBS_VERSION} directory already exists. Removing for fresh install...${NC}"
|
||||
rm -rf yalantinglibs-${YALANTINGLIBS_VERSION}
|
||||
check_success "Failed to remove existing yalantinglibs directory"
|
||||
fi
|
||||
|
||||
# Download yalantinglibs
|
||||
YALANTINGLIBS_ZIPFILE="yalantinglibs-${YALANTINGLIBS_VERSION}.zip"
|
||||
echo "Downloading yalantinglibs ${YALANTINGLIBS_VERSION} from ${GITHUB_PROXY}/alibaba/yalantinglibs/archive/refs/tags/${YALANTINGLIBS_VERSION}.zip"
|
||||
wget -q --show-progress -O ${YALANTINGLIBS_ZIPFILE} ${GITHUB_PROXY}/alibaba/yalantinglibs/archive/refs/tags/${YALANTINGLIBS_VERSION}.zip
|
||||
check_success "Failed to download yalantinglibs"
|
||||
|
||||
# Extract yalantinglibs
|
||||
echo "Extracting yalantinglibs..."
|
||||
unzip -q ${YALANTINGLIBS_ZIPFILE}
|
||||
check_success "Failed to extract yalantinglibs"
|
||||
|
||||
# Clean up downloaded ZIP file
|
||||
rm -f ${YALANTINGLIBS_ZIPFILE}
|
||||
check_success "Failed to clean up downloaded ZIP file"
|
||||
|
||||
# Build and install yalantinglibs
|
||||
cd yalantinglibs-${YALANTINGLIBS_VERSION}
|
||||
check_success "Failed to change to yalantinglibs directory"
|
||||
|
||||
mkdir -p build
|
||||
check_success "Failed to create build directory"
|
||||
|
||||
cd build
|
||||
check_success "Failed to change to build directory"
|
||||
|
||||
echo "Configuring yalantinglibs..."
|
||||
cmake .. -DBUILD_EXAMPLES=OFF -DBUILD_BENCHMARK=OFF -DBUILD_UNIT_TESTS=OFF
|
||||
check_success "Failed to configure yalantinglibs"
|
||||
|
||||
echo "Building yalantinglibs (using $(nproc) cores)..."
|
||||
cmake --build . -j$(nproc)
|
||||
check_success "Failed to build yalantinglibs"
|
||||
|
||||
echo "Installing yalantinglibs..."
|
||||
cmake --install .
|
||||
check_success "Failed to install yalantinglibs"
|
||||
|
||||
print_success "yalantinglibs installed successfully"
|
||||
|
||||
# Initialize and update git submodules
|
||||
print_section "Initializing Git Submodules"
|
||||
|
||||
# Check if .gitmodules exists
|
||||
if [ -f "${REPO_ROOT}/.gitmodules" ]; then
|
||||
# Check if submodules are already initialized by looking for the .git directory in the first submodule
|
||||
FIRST_SUBMODULE=$(grep "path" ${REPO_ROOT}/.gitmodules | head -1 | awk '{print $3}')
|
||||
|
||||
echo "Enter repository root: ${REPO_ROOT}"
|
||||
cd "${REPO_ROOT}"
|
||||
check_success "Failed to change to repository root directory"
|
||||
|
||||
if [ -d "${REPO_ROOT}/${FIRST_SUBMODULE}/.git" ] || [ -f "${REPO_ROOT}/${FIRST_SUBMODULE}/.git" ]; then
|
||||
echo -e "${YELLOW}Git submodules already initialized. Skipping...${NC}"
|
||||
else
|
||||
echo "Initializing git submodules..."
|
||||
git submodule update --init
|
||||
check_success "Failed to initialize git submodules"
|
||||
|
||||
print_success "Git submodules initialized and updated successfully"
|
||||
fi
|
||||
else
|
||||
echo -e "${YELLOW}No .gitmodules file found. Skipping...${NC}"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
print_section "Installing Go $GOVER"
|
||||
|
||||
install_go() {
|
||||
ARCH=$(uname -m)
|
||||
if [ "$ARCH" = "aarch64" ]; then
|
||||
ARCH="arm64"
|
||||
elif [ "$ARCH" = "x86_64" ]; then
|
||||
ARCH="amd64"
|
||||
else
|
||||
echo "Unsupported architecture: $ARCH"
|
||||
exit 1
|
||||
fi
|
||||
# Download Go
|
||||
echo "Downloading Go $GOVER..."
|
||||
wget -q --show-progress https://golang.google.cn/dl/go$GOVER.linux-$ARCH.tar.gz
|
||||
check_success "Failed to download Go $GOVER"
|
||||
|
||||
# Install Go
|
||||
echo "Installing Go $GOVER..."
|
||||
tar -C /usr/local -xzf go$GOVER.linux-$ARCH.tar.gz
|
||||
check_success "Failed to install Go $GOVER"
|
||||
|
||||
# Clean up downloaded file
|
||||
rm -f go$GOVER.linux-$ARCH.tar.gz
|
||||
check_success "Failed to clean up Go installation file"
|
||||
|
||||
print_success "Go $GOVER installed successfully"
|
||||
}
|
||||
|
||||
# Check if Go is already installed
|
||||
if command -v go &> /dev/null; then
|
||||
GO_VERSION=$(go version | awk '{print $3}')
|
||||
if [[ "$GO_VERSION" == "go$GOVER" ]]; then
|
||||
echo -e "${YELLOW}Go $GOVER is already installed. Skipping...${NC}"
|
||||
else
|
||||
echo -e "${YELLOW}Found Go $GO_VERSION. Will install Go $GOVER...${NC}"
|
||||
install_go
|
||||
fi
|
||||
else
|
||||
install_go
|
||||
fi
|
||||
|
||||
# Add Go to PATH if not already there
|
||||
if ! grep -q "export PATH=\$PATH:/usr/local/go/bin" ~/.bashrc; then
|
||||
echo -e "${YELLOW}Adding Go to your PATH in ~/.bashrc${NC}"
|
||||
echo 'export PATH=$PATH:/usr/local/go/bin' >> ~/.bashrc
|
||||
echo -e "${YELLOW}Please run 'source ~/.bashrc' or start a new terminal to use Go${NC}"
|
||||
fi
|
||||
|
||||
# Return to the repository root
|
||||
cd "${REPO_ROOT}"
|
||||
|
||||
# Print summary
|
||||
print_section "Installation Complete"
|
||||
echo -e "${GREEN}All dependencies have been successfully installed!${NC}"
|
||||
echo -e "The following components were installed:"
|
||||
echo -e " ${GREEN}✓${NC} System packages"
|
||||
echo -e " ${GREEN}✓${NC} yalantinglibs"
|
||||
echo -e " ${GREEN}✓${NC} Git submodules"
|
||||
echo -e " ${GREEN}✓${NC} Go $GOVER"
|
||||
echo
|
||||
echo -e "You can now build and run Mooncake."
|
||||
echo -e "${YELLOW}Note: You may need to restart your terminal or run 'source ~/.bashrc' to use Go.${NC}"
|
||||
@@ -1,4 +1,4 @@
|
||||
#!/bin/bash
|
||||
#!/usr/bin/env bash
|
||||
|
||||
#
|
||||
# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
|
||||
@@ -30,9 +30,13 @@ if [ $PYTHON_VERSION == "local" ]; then
|
||||
PYTHON_VERSION=$(python -c 'import sys; print(f"{sys.version_info.major}.{sys.version_info.minor}")')
|
||||
fi
|
||||
|
||||
# Define colors
|
||||
GREEN='\033[0;32m'
|
||||
NC='\033[0m' # No Color
|
||||
|
||||
run_mypy() {
|
||||
echo "Running mypy on $1"
|
||||
mypy --check-untyped-defs --follow-imports skip --python-version "${PYTHON_VERSION}" "$@"
|
||||
echo -e "${GREEN}Running mypy for $1 on python version: ${PYTHON_VERSION}${NC}"
|
||||
mypy --follow-imports skip --check-untyped-defs --python-version "${PYTHON_VERSION}" "$@"
|
||||
}
|
||||
|
||||
run_mypy vllm_ascend
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
#!/bin/bash
|
||||
#!/usr/bin/env bash
|
||||
|
||||
#
|
||||
# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
|
||||
|
||||
45
tools/send_mm_request.py
Normal file
45
tools/send_mm_request.py
Normal file
@@ -0,0 +1,45 @@
|
||||
import base64
|
||||
import os
|
||||
|
||||
import huggingface_hub
|
||||
import requests
|
||||
from modelscope import snapshot_download # type: ignore
|
||||
|
||||
mm_dir = snapshot_download(
|
||||
"vllm-ascend/mm_request",
|
||||
repo_type="dataset",
|
||||
local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE,
|
||||
)
|
||||
image_path = os.path.join(mm_dir, "test_mm2.jpg")
|
||||
with open(image_path, "rb") as image_file:
|
||||
image_data = base64.b64encode(image_file.read()).decode("utf-8")
|
||||
|
||||
data = {
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "What is the content of this image?"},
|
||||
{"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{image_data}"}},
|
||||
],
|
||||
}
|
||||
],
|
||||
"eos_token_id": [1, 106],
|
||||
"pad_token_id": 0,
|
||||
"top_k": 64,
|
||||
"top_p": 0.95,
|
||||
"max_tokens": 8192,
|
||||
"stream": False,
|
||||
}
|
||||
|
||||
headers = {"Accept": "application/json", "Content-Type": "application/json"}
|
||||
|
||||
|
||||
def send_image_request(model, server):
|
||||
data["model"] = model
|
||||
url = server.url_for("v1", "chat", "completions")
|
||||
response = requests.post(url, headers=headers, json=data)
|
||||
print("Status Code:", response.status_code)
|
||||
response_json = response.json()
|
||||
print("Response:", response_json)
|
||||
assert response_json["choices"][0]["message"]["content"], "empty response"
|
||||
39
tools/send_request.py
Normal file
39
tools/send_request.py
Normal file
@@ -0,0 +1,39 @@
|
||||
from typing import Any
|
||||
|
||||
import requests
|
||||
|
||||
|
||||
def send_v1_completions(prompt, model, server, request_args=None):
|
||||
data: dict[str, Any] = {"model": model, "prompt": prompt}
|
||||
if request_args:
|
||||
data.update(request_args)
|
||||
url = server.url_for("v1", "completions")
|
||||
response = requests.post(url, json=data)
|
||||
print(f"Status Code: {response.status_code}")
|
||||
response_json = response.json()
|
||||
print(f"Response json: {response_json}")
|
||||
response_text = response_json["choices"][0]["text"]
|
||||
print(f"Response: {response_text}")
|
||||
assert response_text, "empty response"
|
||||
|
||||
|
||||
def send_v1_chat_completions(prompt, model, server, request_args=None):
|
||||
data: dict[str, Any] = {
|
||||
"model": model,
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": prompt,
|
||||
}
|
||||
],
|
||||
}
|
||||
if request_args:
|
||||
data.update(request_args)
|
||||
url = server.url_for("v1", "chat", "completions")
|
||||
response = requests.post(url, json=data)
|
||||
print(f"Status Code: {response.status_code}")
|
||||
response_json = response.json()
|
||||
print(f"Response json: {response_json}")
|
||||
response_text = response_json["choices"][0]["message"]["content"]
|
||||
print(f"Response: {response_text}")
|
||||
assert response_text, "empty response"
|
||||
@@ -1,4 +1,4 @@
|
||||
#!/bin/bash
|
||||
#!/usr/bin/env bash
|
||||
|
||||
#
|
||||
# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
|
||||
@@ -19,13 +19,13 @@
|
||||
# Adapted from https://github.com/vllm-project/vllm/tree/main/tools
|
||||
#
|
||||
|
||||
set -e
|
||||
set -euo pipefail
|
||||
|
||||
scversion="stable"
|
||||
shellcheck_args=(-S error -s bash)
|
||||
|
||||
if [ -d "shellcheck-${scversion}" ]; then
|
||||
PATH="$PATH:$(pwd)/shellcheck-${scversion}"
|
||||
export PATH
|
||||
export PATH="$PATH:$(pwd)/shellcheck-${scversion}"
|
||||
fi
|
||||
|
||||
if ! [ -x "$(command -v shellcheck)" ]; then
|
||||
@@ -34,12 +34,35 @@ if ! [ -x "$(command -v shellcheck)" ]; then
|
||||
exit 1
|
||||
fi
|
||||
|
||||
# automatic local install if linux x86_64
|
||||
wget -qO- "https://github.com/koalaman/shellcheck/releases/download/${scversion?}/shellcheck-${scversion?}.linux.x86_64.tar.xz" | tar -xJv
|
||||
PATH="$PATH:$(pwd)/shellcheck-${scversion}"
|
||||
export PATH
|
||||
export PATH="$PATH:$(pwd)/shellcheck-${scversion}"
|
||||
fi
|
||||
|
||||
# should enable this
|
||||
# find . -path ./.git -prune -o -name "*.sh" -print0 \
|
||||
# | xargs -0 -I {} sh -c 'git check-ignore -q "{}" || shellcheck -s bash "{}"'
|
||||
if [ -n "${SHELLCHECK_OPTS:-}" ]; then
|
||||
# Split caller-provided options the same way shell would.
|
||||
# shellcheck disable=SC2206
|
||||
extra_shellcheck_args=(${SHELLCHECK_OPTS})
|
||||
shellcheck_args+=("${extra_shellcheck_args[@]}")
|
||||
fi
|
||||
|
||||
if [ "$#" -eq 0 ]; then
|
||||
while IFS= read -r tracked_file; do
|
||||
shellcheck "${shellcheck_args[@]}" "$tracked_file"
|
||||
done < <(git ls-files "*.sh")
|
||||
exit 0
|
||||
fi
|
||||
|
||||
for file in "$@"; do
|
||||
if git check-ignore -q "$file"; then
|
||||
continue
|
||||
fi
|
||||
|
||||
case "$file" in
|
||||
*.csh|*.tcsh)
|
||||
# Skip C shell scripts because this checker only supports sh-like shells.
|
||||
continue
|
||||
;;
|
||||
esac
|
||||
|
||||
shellcheck "${shellcheck_args[@]}" "$file"
|
||||
done
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
#!/bin/bash
|
||||
#!/usr/bin/env bash
|
||||
|
||||
#
|
||||
# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
|
||||
|
||||
157
tools/vllm_bench.py
Normal file
157
tools/vllm_bench.py
Normal file
@@ -0,0 +1,157 @@
|
||||
# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
|
||||
# Copyright 2023 The vLLM team.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
# This file is a part of the vllm-ascend project.
|
||||
#
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import subprocess
|
||||
from datetime import datetime
|
||||
|
||||
from .aisbench import maybe_download_from_modelscope
|
||||
|
||||
|
||||
class VllmbenchRunner:
|
||||
def _run_vllm_bench_task(self):
|
||||
vllm_bench_cmd = [
|
||||
"vllm",
|
||||
"bench",
|
||||
"serve",
|
||||
"--backend",
|
||||
"openai-chat",
|
||||
"--trust-remote-code",
|
||||
"--served-model-name",
|
||||
str(self.model_name),
|
||||
"--model",
|
||||
self.model_path,
|
||||
"--tokenizer",
|
||||
self.model_path,
|
||||
"--metric-percentiles",
|
||||
"50,90,99",
|
||||
"--host",
|
||||
self.host_ip,
|
||||
"--port",
|
||||
str(self.port),
|
||||
"--save-result",
|
||||
"--result-filename",
|
||||
self.result_filename,
|
||||
"--endpoint",
|
||||
"/v1/chat/completions",
|
||||
"--ready-check-timeout-sec",
|
||||
"0",
|
||||
]
|
||||
self._concat_config_args(vllm_bench_cmd)
|
||||
print(f"running vllm_bench cmd: {' '.join(vllm_bench_cmd)}")
|
||||
self.proc: subprocess.Popen = subprocess.Popen(
|
||||
vllm_bench_cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True
|
||||
)
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model_name: str,
|
||||
port: int,
|
||||
config: dict,
|
||||
baseline: float,
|
||||
threshold: float = 0.97,
|
||||
model_path: str = "",
|
||||
host_ip: str = "localhost",
|
||||
):
|
||||
self.model_name = model_name
|
||||
self.model_path = model_path
|
||||
if not self.model_path:
|
||||
self.model_path = maybe_download_from_modelscope(model_name)
|
||||
assert self.model_path is not None, f"Failed to download model: model={self.model_path}"
|
||||
self.port = port
|
||||
self.host_ip = host_ip
|
||||
curr_time = datetime.now().strftime("%Y%m%d%H%M%S")
|
||||
self.result_filename = f"result_vllm_bench_{curr_time}.json"
|
||||
self.config = config
|
||||
self.baseline = baseline
|
||||
self.threshold = threshold
|
||||
|
||||
self._run_vllm_bench_task()
|
||||
self._wait_for_task()
|
||||
self._performance_verify()
|
||||
|
||||
def _concat_config_args(self, vllm_bench_cmd):
|
||||
if "ignore_eos" in self.config:
|
||||
if self.config["ignore_eos"]:
|
||||
self.config["ignore_eos"] = ""
|
||||
else:
|
||||
self.config.pop("ignore_eos")
|
||||
for key, value in self.config.items():
|
||||
key = "--" + key.replace("_", "-")
|
||||
vllm_bench_cmd += [key, str(value)]
|
||||
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc_value, traceback):
|
||||
self.proc.terminate()
|
||||
try:
|
||||
self.proc.wait(8)
|
||||
except subprocess.TimeoutExpired:
|
||||
# force kill if needed
|
||||
self.proc.kill()
|
||||
|
||||
def _wait_for_task(self):
|
||||
"""Wait for the vllm bench command to complete and check the execution result"""
|
||||
|
||||
stdout, stderr = self.proc.communicate()
|
||||
|
||||
if self.proc.returncode != 0:
|
||||
logging.error("vllm bench command failed, return code: %s", self.proc.returncode)
|
||||
logging.error("Standard output: %s", stdout)
|
||||
logging.error("Standard error: %s", stderr)
|
||||
raise RuntimeError(f"vllm bench command execution failed: {stderr}")
|
||||
|
||||
logging.info("vllm bench command completed, return code: %s", self.proc.returncode)
|
||||
if stdout:
|
||||
lines = stdout.split("\n")
|
||||
last_lines = lines[-100:] if len(lines) > 100 else lines
|
||||
logging.info("Last %s lines of standard output:", len(last_lines))
|
||||
for line in last_lines:
|
||||
logging.info(line)
|
||||
else:
|
||||
logging.info("Standard output is empty")
|
||||
|
||||
def _get_result(self):
|
||||
result_file = os.path.join(os.getcwd(), self.result_filename)
|
||||
print("Getting performance results from file: ", result_file)
|
||||
with open(result_file, encoding="utf-8") as f:
|
||||
self.result = json.load(f)
|
||||
|
||||
def _performance_verify(self):
|
||||
self._get_result()
|
||||
output_throughput = self.result["output_throughput"]
|
||||
assert float(output_throughput) >= self.baseline * self.threshold, (
|
||||
"Performance verification failed. "
|
||||
f"The current Output Token Throughput is {output_throughput} token/s, "
|
||||
f"which is not greater than or equal to {self.threshold} * baseline {self.baseline}."
|
||||
)
|
||||
|
||||
|
||||
def run_vllm_bench_case(model_name, port, config, baseline, threshold=0.97, model_path="", host_ip="localhost"):
|
||||
try:
|
||||
with VllmbenchRunner(
|
||||
model_name, port, config, baseline, threshold, model_path=model_path, host_ip=host_ip
|
||||
) as vllm_bench:
|
||||
vllm_bench_result = vllm_bench.result
|
||||
except Exception as e:
|
||||
print(e)
|
||||
error_msg = f"vllm_bench run failed, reason is {e}"
|
||||
logging.error(error_msg)
|
||||
raise RuntimeError(error_msg) from e
|
||||
return vllm_bench_result
|
||||
Reference in New Issue
Block a user