mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-04 07:27:51 +00:00
b5037f28d1
* test(perf): qualify database upgrades with shared 5 RPS traffic * test(perf): share regular latency and memory limits for upgrades * fix(perf): escape index format placeholder in upgrade worker * fix(perf): handle uninitialized multiprocess failure counters * test(perf): match persisted upgrade requests by call ID --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: kerry <kerry@berri.ai>
66 lines
2.2 KiB
Python
66 lines
2.2 KiB
Python
from __future__ import annotations
|
|
|
|
import os
|
|
import sys
|
|
from pathlib import Path
|
|
from typing import Final
|
|
|
|
import pytest
|
|
from setup.run import run_until
|
|
from setup.wire import WireRecord
|
|
from test_upgrade import record
|
|
|
|
|
|
class StopAfterRecord:
|
|
def __init__(self, *, partial: Path, resume: Path, fail: bool = False) -> None:
|
|
self.partial: Final = partial
|
|
self.resume: Final = resume
|
|
self.fail: Final = fail
|
|
self.received: tuple[WireRecord, ...] = ()
|
|
|
|
def __call__(self, records: tuple[WireRecord, ...]) -> bool:
|
|
self.received += records
|
|
if self.partial.exists():
|
|
self.resume.touch()
|
|
assert not (self.fail and records), "Database failed"
|
|
return bool(records)
|
|
|
|
|
|
@pytest.mark.parametrize("fail", [False, True])
|
|
def test_controlled_process_preserves_partial_lines_and_drains_on_stop_or_failure(tmp_path: Path, fail: bool) -> None:
|
|
records: Final = tmp_path / "records.jsonl"
|
|
stop: Final = tmp_path / "stop"
|
|
partial: Final = tmp_path / "partial"
|
|
resume: Final = tmp_path / "resume"
|
|
records.touch()
|
|
encoded: Final = record().model_dump_json() + "\n"
|
|
program: Final = f"""
|
|
import time
|
|
from pathlib import Path
|
|
with Path({str(records)!r}).open("w") as f:
|
|
text = {encoded!r}
|
|
f.write(text[:10])
|
|
f.flush()
|
|
Path({str(partial)!r}).touch()
|
|
while not Path({str(resume)!r}).exists():
|
|
time.sleep(0.01)
|
|
f.write(text[10:])
|
|
f.flush()
|
|
while not Path({str(stop)!r}).exists():
|
|
time.sleep(0.01)
|
|
f.write({record(1).model_dump_json()!r} + "\\n")
|
|
f.flush()
|
|
"""
|
|
callback: Final = StopAfterRecord(partial=partial, resume=resume, fail=fail)
|
|
command: Final = [sys.executable, "-c", program]
|
|
if fail:
|
|
with pytest.raises(AssertionError, match="Database failed"):
|
|
_ = run_until(command, dict(os.environ), records, stop, callback, timeout=10)
|
|
else:
|
|
result: Final = run_until(command, dict(os.environ), records, stop, callback, timeout=10)
|
|
assert result.returncode == 0, result.stdout
|
|
assert callback.received == (record(),)
|
|
assert tuple(WireRecord.model_validate_json(line) for line in records.read_text().splitlines()) == (
|
|
record(),
|
|
record(1),
|
|
)
|