zhifu-platform/tests/inference_capacity_acceptance.py

import asyncio
import json
from collections import Counter
from pathlib import Path
from time import perf_counter

import httpx
from prometheus_client.parser import text_string_to_metric_families


async def test_http_overload_and_recovery():
    async with httpx.AsyncClient(
        base_url="http://127.0.0.1:8200",
        timeout=30,
        trust_env=False,
        limits=httpx.Limits(max_connections=320, max_keepalive_connections=32),
    ) as client:
        start = perf_counter()
        responses = await asyncio.gather(
            *(
                client.post("/embeddings", json={"input": [f"查询第 {index} 个包裹的配送信息"]})
                for index in range(320)
            )
        )
        counts = Counter(response.status_code for response in responses)
        duration = perf_counter() - start
        assert set(counts) <= {200, 429}, counts
        assert counts[200] > 0 and counts[429] > 0, counts
        for response in responses:
            if response.status_code == 200:
                assert len(response.json()["vectors"][0]) == 1024
        recovery = await client.post("/embeddings", json={"input": ["查询物流"]})
        assert recovery.status_code == 200, recovery.text
        metrics = await client.get("/metrics")
        samples = [
            sample
            for family in text_string_to_metric_families(metrics.text)
            for sample in family.samples
        ]
        assert any(
            sample.name == "zhifu_embedding_queue_texts" and sample.value == 0 for sample in samples
        )
        await asyncio.to_thread(
            Path("docs/verification/inference-overload.json").write_text,
            json.dumps(
                {
                    "requests": 320,
                    "statuses": dict(counts),
                    "seconds": duration,
                    "recovery_status": recovery.status_code,
                    "queue_after": 0,
                },
                indent=2,
            ),
            encoding="utf-8",
        )