init v0.23.0

Signed-off-by: Sun Ruoxi <sunruoxi@4paradigm.com>
This commit is contained in:
2026-08-27 15:11:51 +08:00
parent b582a8e7d1
commit 7f8a1b1f7a
2849 changed files with 712887 additions and 22001 deletions

View File

@@ -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
View 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)

View 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
View 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
View 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

View 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())

View 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
View 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
View 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())

View File

@@ -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__":

View File

@@ -1,4 +1,4 @@
#!/bin/bash
#!/usr/bin/env bash
#
# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.

View 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

View File

93
tools/docs_codegen/cli.py Normal file
View 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())

View 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)

View 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,
)

View 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)

View 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()

View 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
View 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"

View 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()

View File

@@ -0,0 +1 @@
"""Sphinx helpers for documentation localization."""

View 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,
}

View 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
View 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}"

View File

@@ -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

View File

@@ -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
View 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
View 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"

View File

@@ -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

View File

@@ -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
View 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