103 lines
3.7 KiB
Python
103 lines
3.7 KiB
Python
#!/usr/bin/env python
|
||
"""E2E-проверка TTFT-watchdog обёртки воркера (e-261014163808-4da3ec).
|
||
|
||
Три сценария на фейковом стриме:
|
||
1. живой апстрим (контент сразу) → все чанки, ни одного ретрая;
|
||
2. висит до первого токена, оживёт на 2-й попытке → ретрай, ответ完整;
|
||
3. висит всегда (3 попытки) → TimeoutError с внятным текстом.
|
||
|
||
Плюс дым: сборка модели payload-ом возвращает прокси с рабочим stream().
|
||
Запуск: ./venv/bin/python tests/test_ttft_retry.py
|
||
"""
|
||
import asyncio
|
||
import os
|
||
import sys
|
||
import types
|
||
|
||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||
|
||
from harness import _llm_worker as W # noqa: E402
|
||
|
||
FAILS = []
|
||
|
||
|
||
def check(name, cond, extra=""):
|
||
print(("PASS " if cond else "FAIL ") + name + (f" | {extra}" if extra and not cond else ""))
|
||
if not cond:
|
||
FAILS.append(name)
|
||
|
||
|
||
def chunk(ct):
|
||
return {"chunk_type": ct}
|
||
|
||
|
||
def fake_stream_model(script):
|
||
"""script: список «попыток», каждая — список чанков или raise Hang."""
|
||
calls = {"n": 0}
|
||
|
||
class M:
|
||
async def stream(self, *a, **k):
|
||
i = calls["n"]
|
||
calls["n"] += 1
|
||
act = script[i]
|
||
if act == "HANG":
|
||
await asyncio.sleep(30) # дольше watchdog-бюджета
|
||
yield chunk("message_start")
|
||
return
|
||
for c in act:
|
||
yield chunk(c)
|
||
|
||
return M(), calls
|
||
|
||
|
||
async def run(model):
|
||
wrapped = W._with_ttft_retry(model, "hour", 90)
|
||
out = []
|
||
async for c in wrapped.stream():
|
||
out.append(W._chunk_type(c))
|
||
return out
|
||
|
||
|
||
async def main():
|
||
os.environ["SIMBIOT_LLM_TTFT_BUDGET"] = "0.5" # быстрый watchdog для теста
|
||
W._TTFT_FIRST_TOKEN_BUDGET = 0.5
|
||
|
||
# 1. живой апстрим
|
||
m, calls = fake_stream_model([
|
||
["message_start", "content_delta", "content_delta", "message_stop"]])
|
||
out = await run(m)
|
||
check("живой стрим проходит целиком",
|
||
out == ["message_start", "content_delta", "content_delta", "message_stop"],
|
||
str(out))
|
||
check("живой стрим — 1 попытка", calls["n"] == 1, str(calls))
|
||
|
||
# 2. виснет, оживает на 2-й
|
||
m, calls = fake_stream_model([
|
||
"HANG",
|
||
["message_start", "content_delta", "message_stop"]])
|
||
out = await run(m)
|
||
check("зависший вызов ретраится и отвечает",
|
||
out == ["message_start", "content_delta", "message_stop"], str(out))
|
||
check("ретрай сделал ровно 2 попытки", calls["n"] == 2, str(calls))
|
||
|
||
# 3. виснет всегда
|
||
m, calls = fake_stream_model(["HANG", "HANG", "HANG"])
|
||
try:
|
||
await run(m)
|
||
check("глухой апстрим даёт TimeoutError", False, "не упал")
|
||
except TimeoutError as e:
|
||
check("глухой апстрим даёт TimeoutError", "завис до первого токена" in str(e), str(e))
|
||
check("глухой апстрим — ровно 3 попытки", calls["n"] == 3, str(calls))
|
||
|
||
# 4. делегирование атрибутов прокси
|
||
m, _ = fake_stream_model([["content_delta"]])
|
||
m.foo = "bar"
|
||
wrapped = W._with_ttft_retry(m, "hour", 90)
|
||
check("прокси делегирует атрибуты", getattr(wrapped, "foo", None) == "bar")
|
||
|
||
print()
|
||
print("ИТОГ:", "ВСЕ ОК" if not FAILS else f"ПРОВАЛЫ: {FAILS}")
|
||
sys.exit(1 if FAILS else 0)
|
||
|
||
|
||
asyncio.run(main())
|