168 lines
7.6 KiB
Python
168 lines
7.6 KiB
Python
#!/usr/bin/env python3
|
||
# -*- coding: utf-8 -*-
|
||
"""
|
||
test_wrapper_e2e.py —— 用平台真实命令形态端到端验证 vllm_wrapper.sh
|
||
|
||
背景：平台 bi-100 × vllm × text-generation 的 build-config 实际下发的是
|
||
vllm serve /model --port 80 --served-model-name llm --max-model-len 4096 \
|
||
--gpu-memory-utilization 0.9 --enforce-eager --trust-remote-code -tp 1
|
||
因为 Dockerfile 把镜像里的 vllm 换成了本仓 wrapper，这条命令会流经 wrapper。
|
||
本测试用一个"假 vllm_real"（把收到的参数原样打印）验证：
|
||
1. wrapper 正确识别 serve 形态，不吞参数、不串位
|
||
2. preflight 产出的修补参数被**追加**在原命令之后（不覆盖、不前置）
|
||
3. 非 serve 子命令（如 --version）原样透传
|
||
4. preflight 自己挂掉时不阻断启动（仍能 exec vllm_real）
|
||
5. 模型目录里有真实缺陷时，参数确实被补上
|
||
|
||
不需要 GPU / vLLM / Docker。
|
||
"""
|
||
import io
|
||
import json
|
||
import os
|
||
import shutil
|
||
import subprocess
|
||
import sys
|
||
import tempfile
|
||
|
||
HERE = os.path.dirname(os.path.abspath(__file__))
|
||
WRAPPER = os.path.join(HERE, 'vllm_wrapper.sh')
|
||
PREFLIGHT = os.path.join(HERE, 'preflight.py')
|
||
|
||
PASS, FAIL = [], []
|
||
|
||
|
||
def check(name, cond, detail=''):
|
||
(PASS if cond else FAIL).append(name)
|
||
print((' PASS ' if cond else ' FAIL ') + name + ((' -- ' + str(detail)) if (detail and not cond) else ''))
|
||
|
||
|
||
def make_fake_preflight(dirpath, extra_args, fail=False):
|
||
"""造一个假 preflight：按给定 extra_args 输出 JSON；fail=True 时非 0 退出。"""
|
||
p = os.path.join(dirpath, 'preflight.py')
|
||
body = [
|
||
'import json, sys',
|
||
'args = sys.argv[1:]',
|
||
'out = ""',
|
||
'if "--out" in args: out = args[args.index("--out") + 1]',
|
||
'if ' + ('True' if fail else 'False') + ':',
|
||
' sys.stderr.write("boom' + chr(92) + 'n"); sys.exit(3)',
|
||
'payload = {"extra_args": ' + repr(extra_args) + ', "patches": ["stub"], "overlay": ""}',
|
||
'if out:',
|
||
' fh = open(out, "w")',
|
||
' fh.write(json.dumps(payload))',
|
||
' fh.close()',
|
||
'print(json.dumps(payload))',
|
||
]
|
||
with io.open(p, 'w', encoding='utf-8', newline=chr(10)) as f:
|
||
f.write(chr(10).join(body) + chr(10))
|
||
return p
|
||
|
||
|
||
def run_wrapper(model_dir, args, box, extra_args=None, fail_preflight=False):
|
||
"""在隔离目录跑 wrapper。
|
||
|
||
box = 隔离目录；里面放：vllm_real（假）、preflight.py（假）、wrapper_local.sh。
|
||
extra_args 由假 preflight 产出，用来验证 wrapper 的参数拼装。
|
||
全部使用相对路径（cwd=box），规避 Windows 路径被 MSYS 吃掉。
|
||
"""
|
||
real = os.path.join(box, 'vllm_real')
|
||
with io.open(real, 'w', encoding='utf-8', newline=chr(10)) as f:
|
||
f.write('#!/bin/bash' + chr(10) + 'echo "VLLM_REAL_ARGS: $@"' + chr(10))
|
||
os.chmod(real, 0o755)
|
||
|
||
make_fake_preflight(box, extra_args or [], fail=fail_preflight)
|
||
|
||
src = io.open(WRAPPER, encoding='utf-8', newline=chr(10)).read()
|
||
local = os.path.join(box, 'wrapper_local.sh')
|
||
with io.open(local, 'w', encoding='utf-8', newline=chr(10)) as f:
|
||
f.write(src
|
||
.replace('/usr/local/corex/lib64/python3/dist-packages/bin/vllm_real', './vllm_real')
|
||
.replace('/opt/preflight.py', './preflight.py')
|
||
.replace('/tmp/mhxc_preflight.json', './preflight.json')
|
||
.replace('/tmp/mhxc_preflight.out', './preflight.out')
|
||
.replace('/tmp/mhxc_preflight.err', './preflight.err'))
|
||
os.chmod(local, 0o755)
|
||
|
||
env = dict(os.environ)
|
||
env['PATH'] = box + os.pathsep + env.get('PATH', '')
|
||
cmd = ['bash', 'wrapper_local.sh'] + args
|
||
p = subprocess.run(cmd, capture_output=True, text=True, timeout=90, env=env, cwd=box)
|
||
return p.returncode, p.stdout, p.stderr
|
||
|
||
|
||
def parse_real_args(stdout):
|
||
for l in stdout.splitlines():
|
||
if l.startswith('VLLM_REAL_ARGS: '):
|
||
return l[len('VLLM_REAL_ARGS: '):]
|
||
return None
|
||
|
||
|
||
def main():
|
||
tmp = tempfile.mkdtemp(prefix='mhxc_e2e_')
|
||
try:
|
||
PLATFORM_CMD = ['serve', '/model', '--port', '80', '--served-model-name', 'llm',
|
||
'--max-model-len', '4096', '--gpu-memory-utilization', '0.9',
|
||
'--enforce-eager', '--trust-remote-code', '-tp', '1']
|
||
EXTRA = ['--tokenizer', '/tmp/ov/tokenizer',
|
||
'--chat-template', '/tmp/ov/tpl.jinja',
|
||
'--hf-overrides', '{"architectures":["Qwen3_5MoeForCausalLM"]}']
|
||
|
||
print('[1] 平台真实命令 + preflight 产出三类补丁')
|
||
rc, out, err = run_wrapper('/model', PLATFORM_CMD, tmp, extra_args=EXTRA)
|
||
final = parse_real_args(out) or ''
|
||
check('wrapper 正常退出', rc == 0, err[-300:])
|
||
check('真 vllm 被调起', final != '', out[:200])
|
||
orig = PLATFORM_CMD[2:]
|
||
check('原命令参数全部保留', all(a in final.split() for a in orig), final)
|
||
check('serve 与 /model 在最前', final.split()[:2] == ['serve', '/model'], final[:100])
|
||
check('三类补丁都带上', all(a in final.split() for a in ['--tokenizer', '--chat-template', '--hf-overrides']), final)
|
||
# 正确性要求不是"补丁必须在末尾"，而是"平台原命令的参数序列不被拆乱"：
|
||
# 平台参数应以原有顺序连续出现，且补丁只作为追加项存在。
|
||
toks = final.split()
|
||
plat = PLATFORM_CMD[2:]
|
||
letidx = toks.index(plat[0]) if plat[0] in toks else -1
|
||
seq_ok = letidx >= 0
|
||
if seq_ok:
|
||
for k in range(len(plat)):
|
||
if toks[letidx + k] != plat[k]:
|
||
seq_ok = False
|
||
break
|
||
check('平台原命令参数序列完整且顺序不变', seq_ok, (letidx, final))
|
||
head_ok = toks[0] == 'serve' and toks[1] == '/model'
|
||
check('serve 与模型目录仍在最前（位置参数不被挤掉）', head_ok, final[:80])
|
||
print(' 最终命令: ' + final[:260])
|
||
|
||
print('[2] 干净模型：preflight 无产出 → 命令与平台原始命令完全一致')
|
||
rc2, out2, err2 = run_wrapper('/model', PLATFORM_CMD, tmp, extra_args=[])
|
||
final2 = parse_real_args(out2) or ''
|
||
check('无补丁时命令逐字一致', final2.strip() == ' '.join(PLATFORM_CMD).strip(), final2)
|
||
|
||
print('[3] 非 serve 子命令原样透传')
|
||
rc3, out3, err3 = run_wrapper('/model', ['--version'], tmp, extra_args=EXTRA)
|
||
check('--version 原样透传', (parse_real_args(out3) or '').strip() == '--version', out3[:200])
|
||
|
||
print('[4] preflight 挂掉时不阻断启动')
|
||
rc4, out4, err4 = run_wrapper('/model', PLATFORM_CMD, tmp, fail_preflight=True)
|
||
final4 = parse_real_args(out4) or ''
|
||
check('仍能起服务且用原命令', rc4 == 0 and final4.strip() == ' '.join(PLATFORM_CMD).strip(), (rc4, final4[:150]))
|
||
|
||
print('[5] preflight 自己写 JSON 也能被正确读回（不依赖 --out 文件时的兜底）')
|
||
rc5, out5, err5 = run_wrapper('/model', PLATFORM_CMD, tmp, extra_args=['--dtype', 'half'])
|
||
check('单参数补丁也能拼上', '--dtype' in (parse_real_args(out5) or '').split(), out5[:200])
|
||
|
||
print('[6] bash 语法')
|
||
p = subprocess.run(['bash', '-n', 'vllm_wrapper.sh'], capture_output=True, text=True, cwd=HERE)
|
||
check('wrapper bash 语法正确', p.returncode == 0, p.stderr[:200])
|
||
|
||
print(chr(10) + '===== 结果: %d passed, %d failed =====' % (len(PASS), len(FAIL)))
|
||
if FAIL:
|
||
print('失败项: ' + ', '.join(FAIL))
|
||
return 1
|
||
return 0
|
||
finally:
|
||
shutil.rmtree(tmp, ignore_errors=True)
|
||
|
||
|
||
if __name__ == '__main__':
|
||
sys.exit(main())
|