mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 17:30:23 +08:00
feat: 流式空闲超时、健康监控查询优化、限流桶内存上限与维护清理修复
Close #233 Co-authored-by: AAEE86 <ppk0227@hotmail.com> - cli_monitor_mixin: 引入 STREAM_IDLE_TIMEOUT_SECONDS(可通过环境变量配置), 流传输开始后若超出空闲窗口无新 chunk 则提前取消并返回 504,避免长时间挂起 - stream_context: 新增 managed_recorded_bodies 上下文管理器,确保 chunks 在 telemetry 完成后及时释放;stream_telemetry 使用该接口统一管理 response body 构建 - health endpoint: 将状态聚合改为 GROUP BY 直接统计,事件列表按 api_format 单独查询,避免单次 limit 拉取大量记录导致的遗漏与性能问题;同时过滤不活跃 provider/endpoint,与公开健康接口保持一致 - endpoint health service: 修正时间线数据按 endpoint_id 而非 key_id 聚合 - token_bucket: 引入 max_buckets/bucket_expiry 上限与定时清理,防止内存无限增长; 修复 refill_rate=0 时 get_reset_time 除零异常;新增 _is_unlimited_rate_limit 判断 - maintenance_scheduler: 调整清理顺序(先删整行再按窗口清理),新增 newer_than 边界参数,避免同一行在同一轮中被重复改写 - sync_execute: 新增 create_pending_usage 开关,允许已预创建记录的调用方跳过重复创建 - quota_reader / provider_ops balance: 小幅修复与健壮性提升 - Dockerfile: 添加 MALLOC_ARENA_MAX=2 环境变量以降低 gunicorn worker RSS - 补充相关测试覆盖
This commit is contained in:
@@ -1,11 +1,15 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import time
|
||||
from types import SimpleNamespace
|
||||
from typing import Any
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
|
||||
os.environ.setdefault("JWT_SECRET_KEY", "test-secret-key")
|
||||
|
||||
from src.api.handlers.base import stream_telemetry as stream_telemetry_module
|
||||
from src.api.handlers.base.stream_context import StreamContext
|
||||
from src.api.handlers.base.stream_telemetry import StreamTelemetryRecorder
|
||||
@@ -64,3 +68,69 @@ async def test_record_stream_stats_estimates_tokens_for_failed_partial_stream(
|
||||
assert ctx.input_tokens > 0
|
||||
assert ctx.output_tokens == max(1, len("partial output") // 4)
|
||||
recorder._dispatch_record.assert_awaited_once() # type: ignore[attr-defined]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_record_stream_stats_releases_parsed_chunks_after_dispatch(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
recorder = StreamTelemetryRecorder(
|
||||
request_id="req-release",
|
||||
user_id="1",
|
||||
api_key_id="2",
|
||||
client_ip="127.0.0.1",
|
||||
format_id="openai:chat",
|
||||
)
|
||||
recorder._get_telemetry_writer = AsyncMock( # type: ignore[method-assign]
|
||||
return_value=SimpleNamespace(include_bodies=True)
|
||||
)
|
||||
recorder._update_candidate_status = AsyncMock() # type: ignore[method-assign]
|
||||
|
||||
dispatch_payloads: list[dict[str, Any]] = []
|
||||
|
||||
async def _capture_dispatch(*args: Any, **kwargs: Any) -> None:
|
||||
dispatch_payloads.append(
|
||||
{
|
||||
"response_body": args[5],
|
||||
"client_response_body": kwargs.get("client_response_body"),
|
||||
}
|
||||
)
|
||||
|
||||
recorder._dispatch_record = _capture_dispatch # type: ignore[method-assign]
|
||||
|
||||
ctx = StreamContext(
|
||||
model="test-model",
|
||||
api_format="openai:chat",
|
||||
request_id="req-release",
|
||||
user_id=1,
|
||||
api_key_id=2,
|
||||
)
|
||||
ctx.provider_name = "test-provider"
|
||||
ctx.parsed_chunks.extend([{"type": "chunk-1"}, {"type": "chunk-2"}])
|
||||
ctx.data_count = 2
|
||||
ctx.chunk_count = 2
|
||||
|
||||
monkeypatch.setattr(stream_telemetry_module, "get_db", lambda: iter([_DummyDb()]))
|
||||
monkeypatch.setattr(
|
||||
stream_telemetry_module.SystemConfigService,
|
||||
"should_log_body",
|
||||
lambda _db: True,
|
||||
)
|
||||
monkeypatch.setattr(stream_telemetry_module.config, "stream_stats_delay", 0)
|
||||
|
||||
await recorder.record_stream_stats(
|
||||
ctx,
|
||||
original_headers={},
|
||||
original_request_body={"input": [{"content": "hello world"}]},
|
||||
start_time=time.time(),
|
||||
)
|
||||
|
||||
assert len(dispatch_payloads) == 1
|
||||
response_body = dispatch_payloads[0]["response_body"]
|
||||
assert response_body["chunks"] == [{"type": "chunk-1"}, {"type": "chunk-2"}]
|
||||
assert response_body["metadata"]["stream"] is True
|
||||
assert response_body["metadata"]["total_chunks"] == 2
|
||||
assert response_body["metadata"]["data_count"] == 2
|
||||
assert dispatch_payloads[0]["client_response_body"] is None
|
||||
assert ctx.parsed_chunks == []
|
||||
assert ctx.provider_parsed_chunks == []
|
||||
|
||||
Reference in New Issue
Block a user