modify: 流式输出

This commit is contained in:
2026-08-21 14:40:36 +08:00
parent 94c50353f4
commit cc8713d5e7
8 changed files with 1132 additions and 23 deletions
+92
View File
@@ -264,6 +264,93 @@ async def upload(file: UploadFile = File(...)):
> **约定**:插件 `__init__.py` 中导入 `router` 实例(`from .routes import router`),网关通过 `getattr(module, "router", None)` 自动发现。 > **约定**:插件 `__init__.py` 中导入 `router` 实例(`from .routes import router`),网关通过 `getattr(module, "router", None)` 自动发现。
### 流式端点(SSE)
基于 Zenoh Query/Reply 原生"多回复"实现流式输出:Worker 的 `@rpc_method` 返回生成器时,
框架逐块 `query.reply()` 推送(见 `utils/worker` README「流式端点」),网关以 SSE
(`text/event-stream`)逐块转发给 HTTP 客户端。**无额外端口、无额外进程、无新依赖**;
未声明 `stream` 的端点行为与之前完全一致。
#### TOML 配置
在方法节中声明 `stream = true`(示例 `config/workers/example/stream_chat.toml`):
```toml
[POST]
summary = "模拟 AI 流式对话(SSE,逐 token 输出)"
auth_required = false
stream = true
timeout = 300.0 # 流总寿命上限(秒,zenoh query deadline 硬上限)
stream_idle_timeout = 10.0 # 网关侧连续无数据块的最大间隔(秒)
params = ["prompt: str", "tokens: int"]
```
| 字段 | 默认值 | 说明 |
| --- | --- | --- |
| `stream` | `false` | 流式端点:以 SSE(`text/event-stream`)输出 |
| `timeout` | `3.0` | 流总寿命上限(秒)。**流式端点必须配大**(如 300s),到期后 zenoh 丢弃后续回复 |
| `stream_idle_timeout` | `10.0` | 空闲超时(秒):超过该间隔无新数据块则网关终止流 |
> `timeout`(总寿命硬上限)与 `stream_idle_timeout`(空闲超时)是两回事,流式场景两个都要配。
#### 请求与响应
- 请求与普通端点一致(GET 走 query/path 参数,POST/PUT/PATCH 走 body 模型),鉴权依赖同样生效
- 响应为 `text/event-stream`,每块一个 `data: {json}\n\n` 事件(SSE 标准分帧)
- 响应头:`Cache-Control: no-cache`、`X-Accel-Buffering: no`(配合反代关闭缓冲)
- OpenAPI 不声明流式端点的响应 schema(逐块 JSON dict,无法用单一模型描述)
```bash
curl -N -X POST localhost:8000/api/v1/calc/stream_chat \
-H 'Content-Type: application/json' -d '{"prompt":"你好","tokens":3}'
# data: {"token":"你好·1","index":0}
# data: {"token":"你好·2","index":1}
# data: {"token":"你好·3","index":2}
# data: {"done":true}
```
#### 错误语义
| 场景 | 表现 |
| --- | --- |
| 无可用 worker(首个数据块前) | HTTP 503(与普通端点一致) |
| Worker 中途异常 / 流式端点初始化失败 | SSE error 事件 `data: {"error":"..."}` 收尾(HTTP 200 已提交,无法再改状态码) |
| 空闲超时 | SSE error 事件 `data: {"error":"流式空闲超时 ..."}` 收尾 |
| 客户端断开 | 网关终止迭代,Worker 侧 reply 失败自动停流,不空转 |
#### 部署注意
生产环境经 nginx 反代时需关闭缓冲,否则 SSE 会被整段缓冲、失去逐块即时性:
```nginx
location /api/v1/ {
proxy_buffering off;
}
```
#### 本地联调前注意
流式能力在 `chongming-worker` 框架(`utils/worker/`,独立仓库)。**本地未 push 的框架改动
不会自动生效**——venv 里安装的是 `uv sync` 从 git 源拉取的旧版副本,旧版 `@rpc_method`
不认 `stream` 参数,会报 `TypeError: rpc_method() got an unexpected keyword argument 'stream'`
并连锁导致契约校验失败(启动终止)。本地联调时先以可编辑方式装入 venv:
```bash
uv pip install -e utils/worker
```
正式发布仍走「开发与发布检查清单」流程(工具包 push → `uv lock --upgrade-package
chongming-worker` → `uv sync`,可编辑安装会被锁文件版本替换)。
#### 测试
```
uv run python test_stream_worker.py # Worker 框架流式回复单测
uv run python test_stream_gateway.py # 网关 call_stream 单测
uv run python test_stream_e2e_zenoh.py # 真实 Zenoh 会话端到端(peer 模式,无需 router)
uv run python test_stream_route.py # 路由层 SSE 输出与兼容性回归(非流式端点不变)
```
--- ---
## 进程插件与混合部署 ## 进程插件与混合部署
@@ -450,6 +537,9 @@ python publisher.py example add '{"a":3,"b":5}' # 调用 add
5. 重启服务 5. 重启服务
``` ```
流式端点(SSE)在此基础上额外两步:TOML 方法节声明 `stream = true` 并配大 `timeout`;
`@rpc_method(..., stream=True)` 函数返回生成器(yield 逐块),详见上文「流式端点(SSE)」。
### 常用命令 ### 常用命令
```bash ```bash
@@ -566,6 +656,8 @@ uv pip install -e utils/<name> # 安装/更新工具包子包
| `internal` | `<操作>.toml` | `true` 则仅限内部 RPC,不注册 HTTP 路由 | | `internal` | `<操作>.toml` | `true` 则仅限内部 RPC,不注册 HTTP 路由 |
| `cacheable` | `<操作>.toml` | 端点是否可缓存 | | `cacheable` | `<操作>.toml` | 端点是否可缓存 |
| `summary` | `<操作>.toml` | 端点描述(用于 OpenAPI / 日志) | | `summary` | `<操作>.toml` | 端点描述(用于 OpenAPI / 日志) |
| `stream` | `<操作>.toml` | `true` 则流式输出(SSE),需同时配大 `timeout` |
| `stream_idle_timeout`| `<操作>.toml` | 流式空闲超时(秒,默认 10),超时无新数据块则终止流 |
--- ---
+148 -7
View File
@@ -5,15 +5,17 @@
为每个端点生成带有正确 Pydantic 类型注解的处理函数,确保 OpenAPI 文档 为每个端点生成带有正确 Pydantic 类型注解的处理函数,确保 OpenAPI 文档
正确反映参数类型、响应结构,并注册到 components/schemas。 正确反映参数类型、响应结构,并注册到 components/schemas。
""" """
from typing import Any, Dict, List, Optional from typing import Any, AsyncGenerator, Dict, List, Optional
from inspect import Parameter, Signature from inspect import Parameter, Signature
import orjson
from fastapi import FastAPI, HTTPException, APIRouter, Query, Path as FastAPIPath, Request, Depends from fastapi import FastAPI, HTTPException, APIRouter, Query, Path as FastAPIPath, Request, Depends
from fastapi.responses import StreamingResponse
from pydantic import create_model, Field, BaseModel from pydantic import create_model, Field, BaseModel
from chongming_worker.route_config import WORKER_CONFIG_DIR, load_worker_routes_from_dir from chongming_worker.route_config import WORKER_CONFIG_DIR, load_worker_routes_from_dir
from .worker_manager import WorkerManager from .worker_manager import WorkerManager, StreamIdleTimeoutError
from .logger import get_logger from .logger import get_logger
from .auth import store_user_to_request from .auth import store_user_to_request
@@ -92,6 +94,10 @@ class RouteConfig:
self.cacheable: bool = data.get("cacheable", False) self.cacheable: bool = data.get("cacheable", False)
self.auth_required: bool = data.get("auth_required", True) self.auth_required: bool = data.get("auth_required", True)
self.internal: bool = data.get("internal", False) self.internal: bool = data.get("internal", False)
# 流式端点:以 SSE(text/event-stream)逐块输出
self.stream: bool = data.get("stream", False)
# 流式空闲超时(秒):连续未收到数据块的最大间隔(timeout 为流总寿命上限)
self.stream_idle_timeout: float = data.get("stream_idle_timeout", 10.0)
def parse_params(self) -> Dict[str, str]: def parse_params(self) -> Dict[str, str]:
"""解析 params 列表为 {name: type} 字典""" """解析 params 列表为 {name: type} 字典"""
@@ -353,6 +359,116 @@ def _make_query_endpoint(worker_manager: WorkerManager, rc: RouteConfig, worker_
return endpoint return endpoint
# 流式端点响应头:禁止缓存;通知反向代理(nginx)关闭缓冲保证 SSE 即时性
_STREAM_HEADERS = {
"Cache-Control": "no-cache",
"X-Accel-Buffering": "no",
}
def _make_stream_endpoint(
worker_manager: WorkerManager,
rc: RouteConfig,
worker_type: str,
path_params: List[str] = (),
query_params: Optional[Dict[str, str]] = None,
RequestModel: Optional[type[BaseModel]] = None,
):
"""
创建流式端点(TOML 中 stream = true):以 SSE(text/event-stream)逐块输出。
请求参数形态与普通端点一致(GET 走 query/path 参数,POST/PUT/PATCH 走 body 模型)。
错误语义:
- 初始错误(无可用 worker 等)-> HTTP 503
- 流中错误(Worker 返回 {"error": ...} 收尾块 / 空闲超时)-> SSE error 事件
(此时 HTTP 200 已提交,无法改状态码)
不校验响应模型:流式输出是逐块 JSON dict。
"""
path = rc.path
subject = rc.subject
http_method = rc.method
timeout = rc.timeout
idle_timeout = rc.stream_idle_timeout
param_names = list((query_params or {}).keys())
def build_payload(request: Request, body=None, kwargs=None) -> Dict[str, Any]:
"""构造信封协议 payload(与普通端点一致)"""
meta = {"method": http_method}
user_info = getattr(request.state, "user", None)
if user_info:
meta["user_info"] = user_info
if body is not None:
data = body.model_dump()
else:
data = {k: kwargs[k] for k in param_names}
return {"meta": meta, "data": data}
def _sse_frame(data: Dict[str, Any]) -> str:
return f"data: {orjson.dumps(data).decode('utf-8')}\n\n"
async def _sse_body(gen) -> AsyncGenerator[str, None]:
"""SSE 分帧;流中异常转为 SSE error 事件(HTTP 200 已提交)"""
try:
async for chunk in gen:
yield f"data: {chunk.decode('utf-8')}\n\n"
except StreamIdleTimeoutError as e:
logger.warning("路由 %s 流式空闲超时: %s", path, e)
yield _sse_frame({"error": str(e)})
except Exception as e:
logger.error("路由 %s 流式中断: %s", path, e)
yield _sse_frame({"error": "内部服务器错误"})
async def _start_stream(payload: Dict[str, Any]) -> StreamingResponse:
try:
gen = await worker_manager.call_stream(
worker_type=worker_type,
subject=subject,
payload=payload,
timeout=timeout,
idle_timeout=idle_timeout,
)
except RuntimeError as e:
logger.warning("路由 %s [%s] 流式调用失败: %s", path, worker_type, e)
raise HTTPException(status_code=503, detail=str(e))
return StreamingResponse(
_sse_body(gen),
media_type="text/event-stream",
headers=_STREAM_HEADERS,
)
if RequestModel is not None:
async def endpoint(request: Request, body: RequestModel) -> StreamingResponse: # type: ignore
return await _start_stream(build_payload(request, body=body))
return endpoint
# GET/DELETE 形态
if not param_names:
async def _simple_endpoint(request: Request) -> StreamingResponse:
return await _start_stream(build_payload(request, kwargs=None))
return _simple_endpoint
async def endpoint(request: Request, **kwargs: Any) -> StreamingResponse:
return await _start_stream(build_payload(request, kwargs=kwargs))
# 附加类型注解与签名(与 _make_query_endpoint 一致),保证 OpenAPI 参数文档正确
params: List[Parameter] = [
Parameter("request", Parameter.POSITIONAL_OR_KEYWORD, annotation=Request)
]
annotations: Dict[str, Any] = {}
for name, type_str in (query_params or {}).items():
py_type = _parse_python_type(type_str)
annotations[name] = py_type
if name in path_params:
params.append(Parameter(name, Parameter.KEYWORD_ONLY, default=FastAPIPath(...), annotation=py_type))
else:
params.append(Parameter(name, Parameter.KEYWORD_ONLY, default=Query(...), annotation=py_type))
endpoint.__annotations__ = annotations
endpoint.__signature__ = Signature(params, return_annotation=StreamingResponse) # type: ignore
return endpoint
def _register_single_route( def _register_single_route(
router: APIRouter, router: APIRouter,
worker_manager: WorkerManager, worker_manager: WorkerManager,
@@ -391,16 +507,41 @@ def _register_single_route(
request_model_name = f"{worker_type}_{clean_path}_{method_tag}_Request" request_model_name = f"{worker_type}_{clean_path}_{method_tag}_Request"
# 生成响应模型 # 生成响应模型
ResponseModel = generate_response_model( if rc.stream:
rc.response_model, response_model_name # 流式端点:逐块 SSE 输出,不校验响应模型(OpenAPI 不声明响应 schema)
) ResponseModel = None
if rc.method in ("POST", "PUT", "PATCH"):
if rc.method in ("POST", "PUT", "PATCH"): # body 形态:RequestModel 传给端点做请求体校验
RequestModel = generate_request_model(
param_spec, request_model_name
)
endpoint = _make_stream_endpoint(
worker_manager, rc, worker_type, path_params, query_params, RequestModel
)
else:
# GET/DELETE 形态:query/path 参数走签名注解;
# RequestModel 仅注册到 extra_schemas 供 OpenAPI 展示(与普通 GET 一致)
if param_spec:
RequestModel = generate_request_model(
param_spec, request_model_name
)
if extra_schemas is not None and request_model_name not in extra_schemas:
extra_schemas[request_model_name] = RequestModel
endpoint = _make_stream_endpoint(
worker_manager, rc, worker_type, path_params, query_params, None
)
elif rc.method in ("POST", "PUT", "PATCH"):
ResponseModel = generate_response_model(
rc.response_model, response_model_name
)
RequestModel = generate_request_model( RequestModel = generate_request_model(
param_spec, request_model_name param_spec, request_model_name
) )
endpoint = _make_body_endpoint(worker_manager, rc, worker_type, RequestModel) endpoint = _make_body_endpoint(worker_manager, rc, worker_type, RequestModel)
else: else:
ResponseModel = generate_response_model(
rc.response_model, response_model_name
)
# 对 GET 端点也生成 Request 模型并注册到 extra_schemas # 对 GET 端点也生成 Request 模型并注册到 extra_schemas
if param_spec: if param_spec:
RequestModel = generate_request_model( RequestModel = generate_request_model(
+126 -15
View File
@@ -2,7 +2,7 @@ import asyncio
import random import random
import json import json
import orjson import orjson
from typing import Any, Dict, List, Optional from typing import Any, AsyncGenerator, Dict, List, Optional
import zenoh import zenoh
from chongming_cache import Cache from chongming_cache import Cache
from .logger import get_logger from .logger import get_logger
@@ -12,6 +12,22 @@ logger = get_logger("worker_manager")
_SENTINEL = object() # 哨兵值,用于标记"No reply received" _SENTINEL = object() # 哨兵值,用于标记"No reply received"
class StreamIdleTimeoutError(TimeoutError):
"""流式调用空闲超时:在 idle_timeout 秒内未收到任何数据块"""
def _next_reply(iterator):
"""线程安全地获取迭代器的下一个元素。
用 for 循环 + 哨兵代替直接 next(),避免 StopIteration 在
asyncio.to_thread 中传播的问题(与 call() 的 _get_first 同模式)。
迭代耗尽时返回 _SENTINEL。
"""
for item in iterator:
return item
return _SENTINEL
class WorkerManager: class WorkerManager:
""" """
统一 Worker 管理器 统一 Worker 管理器
@@ -130,6 +146,114 @@ class WorkerManager:
return False return False
return True return True
async def _raise_no_worker(self, worker_type: str) -> None:
"""
无可用 Worker 时的统一诊断:列出当前已知的所有 worker 类型并抛 RuntimeError。
供 call() / call_stream() 共用,保证两类调用错误信息一致。
"""
async with self._lock:
all_known = list(self._workers_cache.keys())
local_known = list(self._local_workers.keys())
all_types = sorted(set(all_known + local_known))
msg = (
f"类型 '{worker_type}' 没有可用 Worker。"
f"当前已知的 Worker 类型: {all_types if all_types else '(无)'}。"
f"请检查: 1) config/workers/config.toml 是否配置了该类型 "
f"2) 插件初始化是否成功(查看启动日志中的 '加载 Worker 插件失败' 错误)"
f"3) Worker 进程是否正常运行"
)
logger.warning(msg)
raise RuntimeError(msg)
async def call_stream(
self,
worker_type: str,
subject: str,
payload: Dict[str, Any],
timeout: float = 300.0,
idle_timeout: float = 10.0,
) -> AsyncGenerator[bytes, None]:
"""
流式调用:消费 Worker 的多回复流,逐块产出 JSON 字节。
与 call() 的区别:
- 不合并回复(ConsolidationMode.NONE),逐块 yield,实现流式传输
- 不重试(流一旦开始输出无法安全重放,中途失败只能重建)
- 超时模型: timeout = 流总寿命上限(zenoh query deadline,硬上限,
流式端点须配大如 300s);idle_timeout = 连续无数据块的最大间隔
- 错误约定: 块内含 "error" 键(Worker 收尾错误块)或 reply.err
(zenoh 层错误)均为流的终止信号,转发后结束
Args:
worker_type: Worker 类型标识(如 "example")
subject: 调用的主题(如 "stream_chat")
payload: 请求参数(信封格式,由路由层构造)
timeout: 流总寿命上限(秒),需大于业务最长生成时长
idle_timeout: 空闲超时(秒),超时抛 StreamIdleTimeoutError
Yields:
bytes: 每个回复块的 JSON 字节
Raises:
RuntimeError: 无可用 worker
StreamIdleTimeoutError: 空闲超时
"""
# 急切解析 worker:无可用时立即抛错,让路由层映射 HTTP 503
# (异步生成器函数体在首次迭代才执行,故这里用普通 async 函数返回内部生成器)
worker_expr = await self.get_worker(worker_type)
if not worker_expr:
await self._raise_no_worker(worker_type)
full_key = f"{worker_expr}/{subject}"
reply_iter = self.session.get(
full_key,
payload=orjson.dumps(payload),
consolidation=zenoh.ConsolidationMode.NONE,
timeout=timeout,
)
async def _stream():
# 逐块消费;zenoh reply 迭代器是同步阻塞的,必须放线程执行,
# 避免阻塞事件循环(与 call() 的 _get_first 同模式)
while True:
try:
reply = await asyncio.wait_for(
asyncio.to_thread(_next_reply, reply_iter),
timeout=idle_timeout,
)
except asyncio.TimeoutError:
raise StreamIdleTimeoutError(
f"Worker {worker_expr} 流式空闲超时 ({idle_timeout}s 无数据)"
) from None
if reply is _SENTINEL:
break
if reply.ok: # type: ignore
data = bytes(reply.ok.payload) if reply.ok.payload else b"" # type: ignore
if not data:
continue
decoded = orjson.loads(data)
if isinstance(decoded, dict) and decoded.get("error"):
# Worker 错误收尾块:转发后终止流
yield data
return
yield data
else:
# zenoh 层错误回复(如 queryable 异常):转发后终止流
error_payload = reply.err.payload # type: ignore
error_msg = str(error_payload) if error_payload else "unknown error"
logger.warning(
f"Worker {worker_expr} 流式返回错误: {error_msg}"
)
yield orjson.dumps(
{"error": f"Worker {worker_expr} 返回错误: {error_msg}"}
)
return
return _stream()
async def call( async def call(
self, self,
worker_type: str, worker_type: str,
@@ -166,20 +290,7 @@ class WorkerManager:
# 2. 获取可用 worker # 2. 获取可用 worker
worker_expr = await self.get_worker(worker_type) worker_expr = await self.get_worker(worker_type)
if not worker_expr: if not worker_expr:
# 给出诊断信息:列出当前已知的所有 worker 类型 await self._raise_no_worker(worker_type)
async with self._lock:
all_known = list(self._workers_cache.keys())
local_known = list(self._local_workers.keys())
all_types = sorted(set(all_known + local_known))
msg = (
f"类型 '{worker_type}' 没有可用 Worker。"
f"当前已知的 Worker 类型: {all_types if all_types else '(无)'}。"
f"请检查: 1) config/workers/config.toml 是否配置了该类型 "
f"2) 插件初始化是否成功(查看启动日志中的 '加载 Worker 插件失败' 错误)"
f"3) Worker 进程是否正常运行"
)
logger.warning(msg)
raise RuntimeError(msg)
full_key = f"{worker_expr}/{subject}" full_key = f"{worker_expr}/{subject}"
payload_bytes = orjson.dumps(payload) payload_bytes = orjson.dumps(payload)
+180
View File
@@ -0,0 +1,180 @@
"""
真实 Zenoh 会话下的流式端到端测试(peer 模式,无需 router)
单进程内起两个 Zenoh 会话(listen + connect),跑通完整生产链路:
@rpc_method(stream=True) 注册 -> create_worker 声明 Queryable
-> WorkerManager.call_stream 消费多回复 -> 逐块到达
验证:
1. 流式方法:3+1 块按序到达(真实 Zenoh 多回复 + ConsolidationMode.NONE)
2. 流式中途异常:{"error": ...} 收尾块转发后终止
3. 非流式方法:仍是单回复(兼容性回归)
4. generate_route_configs 输出 stream 标记、非流式不含
5. decorators 对生成器函数的调用不消费生成器
运行: uv run python test_stream_e2e_zenoh.py
"""
import asyncio
import importlib.util
import logging
import sys
import types
from pathlib import Path
import orjson
import zenoh
REPO_ROOT = Path(__file__).resolve().parent
sys.path.insert(0, str(REPO_ROOT / "utils" / "worker" / "src"))
sys.path.insert(0, str(REPO_ROOT / "src"))
# ---- 网关依赖隔离(同 test_stream_gateway):桩掉 chongming 包链与 chongming_cache ----
_pkg = types.ModuleType("chongming")
_pkg.__path__ = []
_app = types.ModuleType("chongming.app")
_app.__path__ = []
_core = types.ModuleType("chongming.app.core")
_core.__path__ = []
_logger_stub = types.ModuleType("chongming.app.core.logger")
_logger_stub.get_logger = lambda name: logging.getLogger(name)
_cache_stub = types.ModuleType("chongming_cache")
_cache_stub.Cache = object
for mod in (_pkg, _app, _core, _logger_stub, _cache_stub):
sys.modules[mod.__name__] = mod
_wm_path = REPO_ROOT / "src" / "chongming" / "app" / "core" / "worker_manager.py"
_spec = importlib.util.spec_from_file_location(
"chongming.app.core.worker_manager", _wm_path
)
_wm_mod = importlib.util.module_from_spec(_spec)
sys.modules["chongming.app.core.worker_manager"] = _wm_mod
_spec.loader.exec_module(_wm_mod) # type: ignore[union-attr]
WorkerManager = _wm_mod.WorkerManager
StreamIdleTimeoutError = _wm_mod.StreamIdleTimeoutError
from chongming_worker import clear_registry, rpc_method # noqa: E402
from chongming_worker.decorators import generate_route_configs # noqa: E402
from chongming_worker.runtime import create_worker # noqa: E402
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)
LISTEN_PORT = 7447
LISTEN = f"tcp/127.0.0.1:{LISTEN_PORT}"
def open_session(listen=None):
"""peer 模式会话:worker 侧 listen,网关侧 connect(显式端点,禁用组播)"""
c = zenoh.Config()
c.insert_json5("scouting/multicast/enabled", "false")
if listen:
c.insert_json5("listen/endpoints", f'["{listen}"]')
else:
c.insert_json5("connect/endpoints", f'["{LISTEN}"]')
return zenoh.open(c)
async def run_checks():
clear_registry()
# ---- 流式端点(返回异步生成器)----
@rpc_method(subject="streamtest.echo", method="POST", stream=True)
async def echo_stream(prompt: str = "hi", tokens: int = 3):
for i in range(tokens):
await asyncio.sleep(0.01)
yield {"token": f"{prompt}-{i}"}
yield {"done": True}
# ---- 流式端点:中途异常(框架应发 {"error": ...} 收尾块)----
@rpc_method(subject="streamtest.fail", method="POST", stream=True)
async def fail_stream():
yield {"token": "a"}
raise ValueError("mid-boom")
# ---- 非流式端点(兼容性对照)----
@rpc_method(subject="streamtest.add", method="GET")
async def add(a: float, b: float):
return a + b
# 1. 配置生成:流式标记
configs = generate_route_configs("streamtest")
stream_cfg = next(c for c in configs if c["subject"] == "streamtest.echo")
check("generate_route_configs 输出 stream 标记", stream_cfg.get("stream") is True)
add_cfg = next(c for c in configs if c["subject"] == "streamtest.add")
check("非流式配置不含 stream 标记", "stream" not in add_cfg)
# 2. 真实 Zenoh 会话(peer 模式,无 router)
session_worker = open_session(listen=LISTEN)
session_gw = open_session()
# 3. 启动真实 Worker(create_worker 声明 Queryable)
worker_task = asyncio.create_task(
create_worker(session_worker, "streamtest", "w1")
)
await asyncio.sleep(0.5) # 等待 queryable 在 peer 间可见
try:
# 4. 网关 WorkerManager.call_stream 真实消费流式端点
wm = WorkerManager(session_gw)
wm.register_worker("streamtest", "w1")
async def _consume(subject, data, idle_timeout=2.0):
gen = await wm.call_stream(
"streamtest", subject,
{"meta": {"method": "POST"}, "data": data},
timeout=10.0, idle_timeout=idle_timeout,
)
return [orjson.loads(b) async for b in gen]
chunks = await _consume("streamtest.echo", {"prompt": "hi", "tokens": 3})
check(
"流式 3+1 块按序到达(真实多回复)",
chunks == [
{"token": "hi-0"}, {"token": "hi-1"}, {"token": "hi-2"}, {"done": True},
],
f"chunks={chunks}",
)
# 5. 流式中途异常:错误块收尾
chunks = await _consume("streamtest.fail", {})
check(
"流式中途异常 error 块收尾",
chunks == [{"token": "a"}, {"error": "mid-boom"}],
f"chunks={chunks}",
)
# 6. 非流式兼容:call() 仍单回复
result = await wm.call(
"streamtest", "streamtest.add",
{"meta": {"method": "GET"}, "data": {"a": 2, "b": 3}},
timeout=10.0,
)
check("非流式 call() 正常单回复", result == 5.0, f"result={result}")
# 7. 负载均衡路径兼容:liveliness 发现的 worker 也能流式(同 call)
worker_expr = await wm.get_worker("streamtest")
check("get_worker 仍可解析本地 worker", worker_expr is not None)
finally:
worker_task.cancel()
await asyncio.gather(worker_task, return_exceptions=True)
session_gw.close()
session_worker.close()
def run():
asyncio.run(run_checks())
print(f"\n结果: {len(failures)} 个失败")
return 1 if failures else 0
if __name__ == "__main__":
sys.exit(run())
+246
View File
@@ -0,0 +1,246 @@
"""
网关 call_stream 单测(依赖隔离:桩掉 chongming 包链与 chongming_cache)
直接加载 src/chongming/app/core/worker_manager.py 模块(避开 chongming.app
整链导入与 config.toml 依赖),验证:
1. 多块顺序产出(ConsolidationMode.NONE 逐块消费)
2. Worker 错误收尾块 {"error": ...}:转发后终止流
3. zenoh 错误回复(reply.err):转发错误块后终止流
4. 空闲超时:抛 StreamIdleTimeoutError
5. 无可用 worker:抛 RuntimeError
6. 兼容性回归:call()(非流式路径)仍只取第一个回复
运行: uv run python test_stream_gateway.py
"""
import asyncio
import importlib.util
import logging
import sys
import time
import types
from pathlib import Path
from types import SimpleNamespace
import orjson
import zenoh
REPO_ROOT = Path(__file__).resolve().parent
# ---- 依赖隔离:桩掉网关包链(chongming.app 顶层导入 fastapi/uvicorn/config)----
_pkg = types.ModuleType("chongming")
_pkg.__path__ = []
_app = types.ModuleType("chongming.app")
_app.__path__ = []
_core = types.ModuleType("chongming.app.core")
_core.__path__ = []
_logger_stub = types.ModuleType("chongming.app.core.logger")
_logger_stub.get_logger = lambda name: logging.getLogger(name)
_cache_stub = types.ModuleType("chongming_cache")
_cache_stub.Cache = object # WorkerManager 仅作类型标注使用
for mod in (_pkg, _app, _core, _logger_stub, _cache_stub):
sys.modules[mod.__name__] = mod
_wm_path = REPO_ROOT / "src" / "chongming" / "app" / "core" / "worker_manager.py"
_spec = importlib.util.spec_from_file_location(
"chongming.app.core.worker_manager", _wm_path
)
_wm_mod = importlib.util.module_from_spec(_spec)
sys.modules["chongming.app.core.worker_manager"] = _wm_mod
_spec.loader.exec_module(_wm_mod) # type: ignore[union-attr]
WorkerManager = _wm_mod.WorkerManager
StreamIdleTimeoutError = _wm_mod.StreamIdleTimeoutError
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)
class FakeReply:
"""模拟 zenoh.Reply:ok / err 二选一"""
def __init__(self, payload=None, error=None):
if error is None:
self.ok = SimpleNamespace(payload=payload)
self.err = None
else:
self.ok = None
self.err = SimpleNamespace(payload=error)
class FakeSession:
"""模拟 zenoh.Session.get():返回 reply 迭代器(可迭代对象)"""
def __init__(self, reply_source):
self._reply_source = reply_source
self.get_calls = []
def get(self, key_expr, payload, consolidation=None, timeout=None):
self.get_calls.append((str(key_expr), payload, consolidation, timeout))
return iter(self._reply_source)
def make_manager(replies, worker_id="worker-0"):
wm = WorkerManager(FakeSession(replies))
wm.register_worker("example", worker_id)
return wm
def test_idle_timeout():
# 模拟阻塞式 reply 迭代器:每个 __next__ 阻塞 0.3s(首个块 0.3s 后才到达),
# 而空闲超时仅 0.05s -> 应抛 StreamIdleTimeoutError
class BlockingReplyIterator:
def __init__(self):
self._done = False
def __iter__(self):
return self
def __next__(self):
time.sleep(0.3)
if self._done:
raise StopIteration
self._done = True
return FakeReply(payload=orjson.dumps({"token": "slow"}))
wm = WorkerManager(FakeSession(BlockingReplyIterator()))
def collect_stream(wm, idle_timeout=1.0):
"""消费 call_stream 全部块并返回解码后的列表"""
async def _run():
gen = await wm.call_stream(
"example", "stream_chat", {"meta": {"method": "POST"}, "data": {}},
timeout=5.0, idle_timeout=idle_timeout,
)
return [orjson.loads(b) async for b in gen]
return asyncio.run(_run())
def test_multiple_chunks():
replies = [
FakeReply(payload=orjson.dumps({"token": "a"})),
FakeReply(payload=orjson.dumps({"token": "b"})),
]
wm = make_manager(replies)
chunks = collect_stream(wm)
check("多块顺序产出", chunks == [{"token": "a"}, {"token": "b"}])
check(
"使用 ConsolidationMode.NONE",
wm.session.get_calls[0][2] is zenoh.ConsolidationMode.NONE,
)
check("key_expr 正确", wm.session.get_calls[0][0] == "worker/example/worker-0/stream_chat")
def test_error_chunk_terminates():
replies = [
FakeReply(payload=orjson.dumps({"token": "a"})),
FakeReply(payload=orjson.dumps({"error": "boom"})),
FakeReply(payload=orjson.dumps({"token": "c"})), # 不应被消费
]
wm = make_manager(replies)
chunks = collect_stream(wm)
check("错误块转发后终止", chunks == [{"token": "a"}, {"error": "boom"}])
def test_err_reply_terminates():
replies = [
FakeReply(payload=orjson.dumps({"token": "a"})),
FakeReply(error="oops"),
]
wm = make_manager(replies)
chunks = collect_stream(wm)
check(
"reply.err 转发错误块后终止",
chunks == [{"token": "a"}, {"error": "Worker worker/example/worker-0 返回错误: oops"}],
)
def test_idle_timeout():
# 模拟阻塞式 reply 迭代器:每个 __next__ 阻塞 0.3s(首个块 0.3s 后才到达),
# 而空闲超时仅 0.05s -> 应抛 StreamIdleTimeoutError
class BlockingReplyIterator:
def __init__(self):
self._done = False
def __iter__(self):
return self
def __next__(self):
time.sleep(0.3)
if self._done:
raise StopIteration
self._done = True
return FakeReply(payload=orjson.dumps({"token": "slow"}))
wm = WorkerManager(FakeSession(BlockingReplyIterator()))
wm.register_worker("example", "worker-0")
async def _run():
gen = await wm.call_stream(
"example", "stream_chat", {}, timeout=5.0, idle_timeout=0.05
)
async for _ in gen:
pass
try:
asyncio.run(_run())
check("空闲超时抛 StreamIdleTimeoutError", False)
except StreamIdleTimeoutError as e:
check("空闲超时抛 StreamIdleTimeoutError", True, str(e))
def test_no_worker():
wm = WorkerManager(FakeSession([])) # 未注册任何 worker
async def _run():
gen = await wm.call_stream("missing", "x", {})
async for _ in gen:
pass
try:
asyncio.run(_run())
check("无 worker 抛 RuntimeError", False)
except RuntimeError as e:
check("无 worker 抛 RuntimeError", True, str(e)[:30])
def test_call_compat_first_reply_only():
"""兼容性回归:非流式 call() 仍只取第一个回复,行为不变"""
replies = [
FakeReply(payload=orjson.dumps({"result": 1})),
FakeReply(payload=orjson.dumps({"result": 2})),
]
wm = make_manager(replies)
async def _run():
return await wm.call(
"example", "stream_chat", {"meta": {"method": "POST"}, "data": {}},
timeout=5.0,
)
result = asyncio.run(_run())
check("call() 仍只取第一个回复(非流式兼容)", result == {"result": 1})
def run():
test_multiple_chunks()
test_error_chunk_terminates()
test_err_reply_terminates()
test_idle_timeout()
test_no_worker()
test_call_compat_first_reply_only()
print(f"\n结果: {len(failures)} 个失败")
return 1 if failures else 0
if __name__ == "__main__":
sys.exit(run())
+221
View File
@@ -0,0 +1,221 @@
"""
网关路由层流式端点单测(依赖隔离:桩掉 chongming 包链、auth 与 chongming_cache)
直接加载 src/chongming/app/core/route_generator.py,用 FastAPI TestClient 验证:
1. 流式 POST 端点:SSE 输出(media_type / 逐块 data: 帧 / 响应头)
2. 流式 GET 端点:query 参数正确透传到 payload
3. 初始错误(无可用 worker)-> HTTP 503
4. 流中空闲超时 -> SSE error 事件(HTTP 200 已提交)
5. 兼容性回归:非流式端点输出与改动前一致(JSON 响应)
6. 非流式端点仍走原 call() 逻辑(不误入流式分支)
运行: uv run python test_stream_route.py
"""
import importlib.util
import logging
import sys
import types
from pathlib import Path
import orjson
from fastapi import FastAPI
from fastapi.testclient import TestClient
REPO_ROOT = Path(__file__).resolve().parent
sys.path.insert(0, str(REPO_ROOT / "src"))
sys.path.insert(0, str(REPO_ROOT / "utils" / "worker" / "src"))
# ---- 依赖隔离:桩掉 chongming 包链 / auth / chongming_cache ----
_pkg = types.ModuleType("chongming")
_pkg.__path__ = []
_app = types.ModuleType("chongming.app")
_app.__path__ = []
_core = types.ModuleType("chongming.app.core")
_core.__path__ = []
_logger_stub = types.ModuleType("chongming.app.core.logger")
_logger_stub.get_logger = lambda name: logging.getLogger(name)
_auth_stub = types.ModuleType("chongming.app.core.auth")
_auth_stub.store_user_to_request = lambda x: x # 仅用于 Depends,桩掉鉴权
_cache_stub = types.ModuleType("chongming_cache")
_cache_stub.Cache = object
for mod in (_pkg, _app, _core, _logger_stub, _auth_stub, _cache_stub):
sys.modules[mod.__name__] = mod
# 加载真实 worker_manager(route_generator 依赖其 StreamIdleTimeoutError)
_wm_path = REPO_ROOT / "src" / "chongming" / "app" / "core" / "worker_manager.py"
_wm_spec = importlib.util.spec_from_file_location(
"chongming.app.core.worker_manager", _wm_path
)
_wm_mod = importlib.util.module_from_spec(_wm_spec)
sys.modules["chongming.app.core.worker_manager"] = _wm_mod
_wm_spec.loader.exec_module(_wm_mod) # type: ignore[union-attr]
# 加载真实 route_generator
_rg_path = REPO_ROOT / "src" / "chongming" / "app" / "core" / "route_generator.py"
_rg_spec = importlib.util.spec_from_file_location(
"chongming.app.core.route_generator", _rg_path
)
_rg_mod = importlib.util.module_from_spec(_rg_spec)
sys.modules["chongming.app.core.route_generator"] = _rg_mod
_rg_spec.loader.exec_module(_rg_mod) # type: ignore[union-attr]
register_route_group = _rg_mod.register_route_group
StreamIdleTimeoutError = _wm_mod.StreamIdleTimeoutError
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)
class FakeWorkerManager:
"""模拟 WorkerManager:call() 同步返回;call_stream() 逐块产出"""
def __init__(self, chunks=None, stream_error=None):
self.chunks = chunks if chunks is not None else [{"token": "a"}, {"token": "b"}, {"done": True}]
self.stream_error = stream_error # "no_worker" | "idle" | None
self.call_payloads = []
self.stream_payloads = []
async def call(self, worker_type, subject, payload, timeout=3.0, retries=1, cacheable=False):
self.call_payloads.append((worker_type, subject, payload))
data = payload["data"]
return {"result": data.get("a", 0) + data.get("b", 0)}
async def call_stream(self, worker_type, subject, payload, timeout=300.0, idle_timeout=10.0):
if self.stream_error == "no_worker":
raise RuntimeError(f"类型 '{worker_type}' 没有可用 Worker。")
self.stream_payloads.append((worker_type, subject, payload))
async def _gen():
if self.stream_error == "idle":
yield orjson.dumps({"token": "a"})
raise StreamIdleTimeoutError("流式空闲超时 (0.1s 无数据)")
for c in self.chunks:
yield orjson.dumps(c)
return _gen()
def build_app(fwm):
app = FastAPI()
routes = [
# 流式 POST(body 参数)
{
"subject": "streamtest.chat", "path": "/stream/chat", "method": "POST",
"params": ["prompt: str", "tokens: int"], "summary": "sse chat",
"stream": True, "timeout": 300.0, "stream_idle_timeout": 1.0,
"auth_required": False,
},
# 流式 GET(query 参数)
{
"subject": "streamtest.progress", "path": "/stream/progress", "method": "GET",
"params": ["job_id: str", "interval: float"], "summary": "sse progress",
"stream": True, "timeout": 60.0, "auth_required": False,
},
# 非流式对照(兼容性回归)
{
"subject": "streamtest.add", "path": "/calc/add", "method": "GET",
"params": ["a: float", "b: float"], "summary": "add",
"timeout": 3.0, "auth_required": False,
},
{
"subject": "streamtest.echo", "path": "/calc/echo", "method": "POST",
"params": ["msg: str"], "summary": "echo",
"timeout": 3.0, "auth_required": False,
},
]
register_route_group(app, fwm, "streamtest", routes, prefix="/api/v1")
return TestClient(app)
def test_stream_post_sse():
fwm = FakeWorkerManager()
client = build_app(fwm)
r = client.post("/api/v1/stream/chat", json={"prompt": "hi", "tokens": 3})
check("流式 POST 状态码 200", r.status_code == 200, str(r.status_code))
check(
"SSE media_type",
r.headers["content-type"].startswith("text/event-stream"),
r.headers.get("content-type", ""),
)
check(
"SSE 逐块 data: 帧",
r.text
== 'data: {"token":"a"}\n\ndata: {"token":"b"}\n\ndata: {"done":true}\n\n',
repr(r.text),
)
check("防缓冲响应头", r.headers.get("cache-control") == "no-cache")
check(
"body 参数透传",
fwm.stream_payloads[0][2]["data"] == {"prompt": "hi", "tokens": 3},
)
def test_stream_get_query_params():
fwm = FakeWorkerManager()
client = build_app(fwm)
r = client.get("/api/v1/stream/progress?job_id=job-1&interval=0.5")
check("流式 GET 状态码 200", r.status_code == 200)
check(
"GET query 参数透传",
fwm.stream_payloads[0][2]["data"] == {"job_id": "job-1", "interval": 0.5},
str(fwm.stream_payloads[0][2]["data"]),
)
check("GET 流式输出 SSE", r.text.startswith("data: "), repr(r.text[:40]))
def test_stream_no_worker_503():
fwm = FakeWorkerManager(stream_error="no_worker")
client = build_app(fwm)
r = client.post("/api/v1/stream/chat", json={"prompt": "hi", "tokens": 1})
check("初始错误(无 worker)-> 503", r.status_code == 503, str(r.status_code))
def test_stream_idle_timeout_sse_event():
fwm = FakeWorkerManager(stream_error="idle")
client = build_app(fwm)
r = client.post("/api/v1/stream/chat", json={"prompt": "hi", "tokens": 1})
check("空闲超时保持 200(已提交)", r.status_code == 200)
check(
"空闲超时转为 SSE error 事件",
r.text == 'data: {"token":"a"}\n\ndata: {"error":"流式空闲超时 (0.1s 无数据)"}\n\n',
repr(r.text),
)
def test_non_stream_compat():
"""兼容性回归:非流式端点走原 call() 逻辑,输出不变"""
fwm = FakeWorkerManager()
client = build_app(fwm)
r = client.get("/api/v1/calc/add?a=2&b=3")
check("非流式 GET 仍为 JSON 响应", r.status_code == 200 and r.json() == {"result": 5.0},
f"status={r.status_code} body={r.text}")
r = client.post("/api/v1/calc/echo", json={"msg": "hello"})
check("非流式 POST 仍为 JSON 响应",
r.status_code == 200 and r.json() == {"result": 0.0},
f"status={r.status_code} body={r.text}")
check("非流式调用未进入 call_stream", len(fwm.stream_payloads) == 0)
check("非流式调用仍走 call()", len(fwm.call_payloads) == 2)
def run():
test_stream_post_sse()
test_stream_get_query_params()
test_stream_no_worker_503()
test_stream_idle_timeout_sse_event()
test_non_stream_compat()
print(f"\n结果: {len(failures)} 个失败")
return 1 if failures else 0
if __name__ == "__main__":
sys.exit(run())
+118
View File
@@ -0,0 +1,118 @@
"""
Worker 框架流式回复单测(不依赖真实 zenoh 会话)
直接测 runtime.stream_replies_to_query(fake query 仅需 key_expr 与 reply):
1. 异步生成器:多块顺序 reply,payload 为 orjson JSON
2. 同步生成器:同样逐块 reply
3. 生成器中途异常:{"error": ...} 收尾块后停流
4. reply 失败(查询已终止):立即停流且不抛异常
运行: uv run python test_stream_worker.py
"""
import asyncio
import sys
from pathlib import Path
import orjson
REPO_ROOT = Path(__file__).resolve().parent
# 优先使用本地框架 checkout(即使 chongming_worker 已通过 uv 安装也以本地为准)
sys.path.insert(0, str(REPO_ROOT / "utils" / "worker" / "src"))
from chongming_worker.runtime import stream_replies_to_query # noqa: E402
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)
class FakeQuery:
"""仅实现 stream_replies_to_query 用到的两个接口(key_expr / reply)"""
def __init__(self, key_expr="worker/example/worker-0/example.stream_chat"):
self.key_expr = key_expr
self.replies = []
self.fail_reply = False
def reply(self, key_expr, payload):
if self.fail_reply:
raise RuntimeError("query terminated")
self.replies.append((key_expr, payload))
def test_async_generator():
q = FakeQuery()
async def gen():
yield {"token": "a"}
yield {"token": "b"}
asyncio.run(stream_replies_to_query(q, gen()))
check("异步生成器逐块回复", len(q.replies) == 2)
check("回复 key_expr 正确", all(k == q.key_expr for k, _ in q.replies))
check(
"payload 为 JSON dict",
[orjson.loads(p) for _, p in q.replies] == [{"token": "a"}, {"token": "b"}],
)
def test_sync_generator():
q = FakeQuery()
def gen():
yield {"n": 1}
yield {"n": 2}
asyncio.run(stream_replies_to_query(q, gen()))
check(
"同步生成器逐块回复",
[orjson.loads(p) for _, p in q.replies] == [{"n": 1}, {"n": 2}],
)
def test_mid_stream_exception():
q = FakeQuery()
async def gen():
yield {"token": "a"}
raise ValueError("boom")
asyncio.run(stream_replies_to_query(q, gen()))
check("中途异常收尾 error 块", len(q.replies) == 2)
check("error 块内容", orjson.loads(q.replies[1][1]) == {"error": "boom"})
def test_reply_failure_stops():
q = FakeQuery()
q.fail_reply = True
yielded = []
async def gen():
yielded.append("yield-1")
yield {"token": "a"}
yielded.append("yield-2") # 首个块 reply 失败后不应被执行到
yield {"token": "b"}
# 不应抛异常,且首个块 reply 失败后立即停流
asyncio.run(stream_replies_to_query(q, gen()))
check("reply 失败不抛异常", True)
check("reply 失败后立即停流", yielded == ["yield-1"])
def run():
test_async_generator()
test_sync_generator()
test_mid_stream_exception()
test_reply_failure_stops()
print(f"\n结果: {len(failures)} 个失败")
return 1 if failures else 0
if __name__ == "__main__":
sys.exit(run())
Generated
+1 -1
View File
@@ -486,7 +486,7 @@ source = { git = "http://124.71.68.6:3000/chongming/chongming-stream-manager.git
[[package]] [[package]]
name = "chongming-worker" name = "chongming-worker"
version = "0.2.0" version = "0.2.0"
source = { git = "http://124.71.68.6:3000/chongming/chongming-utils-worker.git#d8230bc0ca913a1253d9db9d9af231da9fa7c779" } source = { git = "http://124.71.68.6:3000/chongming/chongming-utils-worker.git#6dbdce177b2e44f2d0fd84df51d49bd690a8926b" }
dependencies = [ dependencies = [
{ name = "eclipse-zenoh" }, { name = "eclipse-zenoh" },
{ name = "pydantic" }, { name = "pydantic" },