475 lines
17 KiB
Python
475 lines
17 KiB
Python
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)
|