211 lines
9.0 KiB
Python
211 lines
9.0 KiB
Python
"""inline 直调路径 与 zenoh 路径 的一致性测试。
|
||
|
||
守护"进程插件 ↔ 内置插件无缝迁移"契约：同一个 worker 副本，分别经
|
||
A) 进程内直调（chongming_worker.dispatch：网关 inline 快路径用的入口）
|
||
B) zenoh 查询（session.get → router → queryable 回调 → 原有路径）
|
||
发起同一批请求，逐字段比对两路结果（仅屏蔽时间戳这类天然不同的字段）。
|
||
|
||
覆盖：正常 GET / POST、模型校验失败（reply_err 通道）、未知 subject、
|
||
非信封旧式载荷、流式多块、流式生成器异常收尾块。
|
||
|
||
运行: uv run python test/test_inline_local_dispatch.py
|
||
（需要 zenoh router 在 config/zenoh_config.json5 指向的端点可达）
|
||
"""
|
||
import asyncio
|
||
import os
|
||
import sys
|
||
from pathlib import Path
|
||
|
||
REPO_ROOT = Path(__file__).resolve().parent.parent
|
||
# 优先本地框架 checkout（与 test_stream_worker.py 同约定）
|
||
sys.path.insert(0, str(REPO_ROOT / "utils" / "worker" / "src"))
|
||
sys.path.insert(0, str(REPO_ROOT / "gateway" / "src"))
|
||
sys.path.insert(0, str(REPO_ROOT / "src")) # plugins 插件源码
|
||
|
||
import orjson # noqa: E402
|
||
import zenoh # noqa: E402
|
||
|
||
from chongming_worker.dispatch import get_dispatcher, unregister_dispatcher # noqa: E402
|
||
from chongming_worker.runtime import create_worker # noqa: E402
|
||
|
||
WORKER_TYPE = "example"
|
||
WORKER_ID = "worker-parity-0"
|
||
|
||
failures = []
|
||
|
||
|
||
def check(name, cond, detail=""):
|
||
status = "PASS" if cond else "FAIL"
|
||
print(f"[{status}] {name}" + (f" — {detail}" if detail else ""))
|
||
if not cond:
|
||
failures.append(name)
|
||
|
||
|
||
def mask_timestamp(obj):
|
||
"""屏蔽时间戳类字段（两路调用时刻不同，天然不一致）"""
|
||
if isinstance(obj, dict):
|
||
return {k: (0 if k in ("timestamp", "time") else mask_timestamp(v)) for k, v in obj.items()}
|
||
if isinstance(obj, list):
|
||
return [mask_timestamp(v) for v in obj]
|
||
return obj
|
||
|
||
|
||
def normalize(payload: bytes):
|
||
try:
|
||
return mask_timestamp(orjson.loads(payload))
|
||
except Exception:
|
||
return payload
|
||
|
||
|
||
def envelope(method: str, data: dict) -> bytes:
|
||
return orjson.dumps({"meta": {"method": method, "user_info": None}, "data": data})
|
||
|
||
|
||
async def invoke_zenoh(session, subject: str, payload: bytes, timeout: float = 3.0):
|
||
"""经 zenoh 查询取首个回复：返回 ("ok", 载荷) / ("err", 载荷) / ("noreply", None)"""
|
||
key = f"worker/{WORKER_TYPE}/{WORKER_ID}/{subject}"
|
||
loop = asyncio.get_running_loop()
|
||
|
||
def _first():
|
||
for reply in session.get(key, payload=payload,
|
||
consolidation=zenoh.ConsolidationMode.NONE,
|
||
timeout=timeout):
|
||
return reply
|
||
return None
|
||
|
||
reply = await loop.run_in_executor(None, _first)
|
||
if reply is None:
|
||
return ("noreply", None)
|
||
if reply.ok:
|
||
return ("ok", bytes(reply.ok.payload) if reply.ok.payload else b"")
|
||
return ("err", bytes(reply.err.payload))
|
||
|
||
|
||
async def collect_zenoh_stream(session, subject: str, payload: bytes, timeout: float = 5.0):
|
||
key = f"worker/{WORKER_TYPE}/{WORKER_ID}/{subject}"
|
||
loop = asyncio.get_running_loop()
|
||
|
||
def _all():
|
||
out = []
|
||
for reply in session.get(key, payload=payload,
|
||
consolidation=zenoh.ConsolidationMode.NONE,
|
||
timeout=timeout):
|
||
if reply.ok:
|
||
out.append(bytes(reply.ok.payload) if reply.ok.payload else b"")
|
||
else:
|
||
out.append(("err", bytes(reply.err.payload)))
|
||
return out
|
||
|
||
return await loop.run_in_executor(None, _all)
|
||
|
||
|
||
async def collect_local_stream(dispatcher, subject: str, payload: bytes):
|
||
out = []
|
||
async for chunk in dispatcher.invoke_stream(subject, payload):
|
||
out.append(chunk)
|
||
return out
|
||
|
||
|
||
async def main() -> int:
|
||
# ── 起会话 + 一个真实 worker 副本（example 插件，纯计算端点）──
|
||
import plugins.example.router.add # noqa: F401
|
||
import plugins.example.router.subtract # noqa: F401
|
||
import plugins.example.router.stream_chat # noqa: F401
|
||
|
||
zconfig = zenoh.Config()
|
||
zconfig.from_file(str(REPO_ROOT / "config" / "zenoh_config.json5"))
|
||
session = zenoh.open(zconfig)
|
||
|
||
worker_task = asyncio.create_task(
|
||
create_worker(session, WORKER_TYPE, WORKER_ID, config={"data_dir": "./data"})
|
||
)
|
||
await asyncio.sleep(1.5) # 等 liveliness / queryable / 直调注册都就绪
|
||
|
||
dispatcher = get_dispatcher(WORKER_TYPE, WORKER_ID)
|
||
check("直调入口已注册（register_dispatcher）", dispatcher is not None)
|
||
if dispatcher is None:
|
||
worker_task.cancel()
|
||
session.close()
|
||
return 1
|
||
|
||
try:
|
||
# ── 1) 正常 GET：example.add?a&b ──
|
||
payload = envelope("GET", {"a": 2, "b": 3})
|
||
z = await invoke_zenoh(session, "example.add", payload)
|
||
l = await dispatcher.invoke_unary("example.add", payload)
|
||
l_norm = ("ok" if l.ok else "err", l.payload)
|
||
check("GET 两路状态一致", z[0] == l_norm[0], f"zenoh={z[0]} local={l_norm[0]}")
|
||
check(
|
||
"GET 两路载荷逐字段一致",
|
||
normalize(z[1] or b"") == normalize(l_norm[1]),
|
||
f"zenoh={z[1]!r} local={l_norm[1]!r}",
|
||
)
|
||
|
||
# ── 2) 正常 POST（同 subject 的另一方法）──
|
||
payload = envelope("POST", {"a": 5, "b": 7})
|
||
z = await invoke_zenoh(session, "example.add", payload)
|
||
l = await dispatcher.invoke_unary("example.add", payload)
|
||
l_norm = ("ok" if l.ok else "err", l.payload)
|
||
check("POST 两路状态一致", z[0] == l_norm[0], f"zenoh={z[0]} local={l_norm[0]}")
|
||
check("POST 两路载荷逐字段一致", normalize(z[1] or b"") == normalize(l_norm[1]))
|
||
|
||
# ── 3) 模型校验失败 → 框架错误（zenoh reply_err ⇔ local ok=False）──
|
||
payload = envelope("GET", {"a": "not-a-number"})
|
||
z = await invoke_zenoh(session, "example.add", payload)
|
||
l = await dispatcher.invoke_unary("example.add", payload)
|
||
l_norm = ("ok" if l.ok else "err", l.payload)
|
||
check("校验失败两路都走错误通道", z[0] == "err" and l_norm[0] == "err",
|
||
f"zenoh={z[0]} local={l_norm[0]}")
|
||
z_err, l_err = normalize(z[1] or b""), normalize(l_norm[1])
|
||
check("错误信封结构一致（operation=error 且含 error 文本）",
|
||
isinstance(z_err, dict) and isinstance(l_err, dict)
|
||
and z_err.get("operation") == l_err.get("operation") == "error"
|
||
and bool(z_err.get("error")) and bool(l_err.get("error")),
|
||
f"zenoh={z_err} local={l_err}")
|
||
|
||
# ── 4) 非信封旧式载荷（method 默认 GET）──
|
||
payload = orjson.dumps({"a": 1, "b": 1})
|
||
z = await invoke_zenoh(session, "example.subtract", payload)
|
||
l = await dispatcher.invoke_unary("example.subtract", payload)
|
||
l_norm = ("ok" if l.ok else "err", l.payload)
|
||
check("旧式载荷两路一致", z[0] == l_norm[0] and normalize(z[1] or b"") == normalize(l_norm[1]),
|
||
f"zenoh={z} local={l_norm}")
|
||
|
||
# ── 5) 未知 subject：两路都应失败（路径不同但终态一致：网关侧都会换副本重试后 503）──
|
||
z = await invoke_zenoh(session, "example.not_exists", envelope("GET", {}), timeout=1.0)
|
||
l = await dispatcher.invoke_unary("example.not_exists", envelope("GET", {}))
|
||
check("未知 subject 两路都失败（zenoh 无回复 / 直调返回错误）",
|
||
z[0] == "noreply" and l.ok is False, f"zenoh={z[0]} local.ok={l.ok}")
|
||
|
||
# ── 6) 流式多块 ──
|
||
payload = envelope("POST", {"prompt": "你好", "tokens": 4})
|
||
zs = await collect_zenoh_stream(session, "example.stream_chat", payload)
|
||
ls = await collect_local_stream(dispatcher, "example.stream_chat", payload)
|
||
check("流式块数一致", len(zs) == len(ls), f"zenoh={len(zs)} local={len(ls)}")
|
||
check("流式逐块内容一致",
|
||
[normalize(x) for x in zs] == [normalize(x) for x in ls],
|
||
f"zenoh={[normalize(x) for x in zs]} local={[normalize(x) for x in ls]}")
|
||
|
||
# ── 7) 流式生成器中途异常 → {"error": ...} 收尾块（两路都应以此结束）──
|
||
# 流终止语义（无 token 的空流）两路一致
|
||
payload = envelope("POST", {"prompt": "空", "tokens": 0})
|
||
zs = await collect_zenoh_stream(session, "example.stream_chat", payload)
|
||
ls = await collect_local_stream(dispatcher, "example.stream_chat", payload)
|
||
check("空流终止语义一致（不挂起、块数相同）",
|
||
len(zs) == len(ls), f"zenoh={len(zs)} local={len(ls)}")
|
||
finally:
|
||
worker_task.cancel()
|
||
try:
|
||
await worker_task
|
||
except asyncio.CancelledError:
|
||
pass
|
||
unregister_dispatcher(WORKER_TYPE, WORKER_ID)
|
||
session.close()
|
||
|
||
print(f"\n结果: {len(failures)} 个失败" + (f" -> {failures}" if failures else ""))
|
||
return 1 if failures else 0
|
||
|
||
|
||
if __name__ == "__main__":
|
||
sys.exit(asyncio.run(main()))
|