zhifu-platform/tests/replica_capacity_acceptance.py

import asyncio
import json
import platform
from collections import Counter
from datetime import UTC, datetime
from pathlib import Path
from time import perf_counter

import httpx
from capacity_acceptance import accounts, percentiles
from retrieval_evaluation import publish_dataset
from test_after_sales import scenario  # noqa: F401

from zhifu.modules.knowledge.vector import connection


def summarize(samples, seconds):
    successful = [row["seconds"] for row in samples if row["status"] == "200"]
    return {
        "seconds": seconds,
        "requests": len(samples),
        "statuses": dict(Counter(row["status"] for row in samples)),
        "replicas": dict(Counter(row["replica"] for row in samples)),
        "success_rps": len(successful) / seconds,
        "success_latency": percentiles(successful),
    }


async def sustained(client, method, path, bodies, identities, seconds, concurrency=32):
    samples = []
    begin = perf_counter()

    async def worker(index):
        sequence = index
        while perf_counter() - begin < seconds:
            start = perf_counter()
            result = (
                await asyncio.gather(
                    client.request(
                        method,
                        path,
                        json=bodies[sequence % len(bodies)],
                        headers=identities[index % len(identities)],
                    ),
                    return_exceptions=True,
                )
            )[0]
            if isinstance(result, httpx.HTTPError):
                status, replica = type(result).__name__, "transport"
            elif isinstance(result, BaseException):
                raise result
            else:
                status = str(result.status_code)
                replica = result.headers.get("X-Zhifu-Replica", "direct")
                if status == "200":
                    assert result.json()["items"], result.text
            samples.append(
                {
                    "status": status,
                    "replica": replica,
                    "seconds": perf_counter() - start,
                    "window": min(int((start - begin) // 10), (seconds - 1) // 10),
                }
            )
            sequence += concurrency

    await asyncio.gather(*(worker(index) for index in range(concurrency)))
    elapsed = perf_counter() - begin
    result = summarize(samples, elapsed)
    result["windows"] = [
        {
            "start_second": index * 10,
            **summarize(
                [row for row in samples if row["window"] == index],
                10 if index < seconds // 10 - 1 else elapsed - index * 10,
            ),
        }
        for index in range(seconds // 10)
    ]
    return result


async def test_sustained_replica_capacity(scenario):  # noqa: F811
    try:
        dataset, _ = await publish_dataset(scenario)
        identities = await accounts(scenario)
        stages = []
        for name, url, replicas in (
            ("single", "http://127.0.0.1:18000", 1),
            ("replicas", "http://127.0.0.1:18080", 2),
        ):
            async with httpx.AsyncClient(
                base_url=url, timeout=35, trust_env=False, limits=httpx.Limits(max_connections=32)
            ) as client:
                for workload, method, path, bodies, seconds in (
                    (
                        "orders",
                        "GET",
                        f"/openapi/v1/orders?store_id={scenario['store']}",
                        [None],
                        30,
                    ),
                    (
                        "knowledge",
                        "POST",
                        f"/openapi/v1/stores/{scenario['store']}/knowledge/search",
                        [{"query": question["query"]} for question in dataset["queries"]],
                        60,
                    ),
                ):
                    response = await client.request(
                        method, path, json=bodies[0], headers=identities[0]
                    )
                    assert response.status_code == 200, response.text
                    result = await sustained(client, method, path, bodies, identities, seconds)
                    stages.append(
                        {
                            "deployment": name,
                            "api_replicas": replicas,
                            "inference_replicas": replicas,
                            "workload": workload,
                            "concurrency": 32,
                            **result,
                        }
                    )
        report = {
            "checked_at": datetime.now(UTC).isoformat(),
            "host": platform.platform(),
            "accounts": 16,
            "documents": 12,
            "queries": 24,
            "stages": stages,
        }
        await asyncio.to_thread(
            Path("docs/verification/replica-capacity.json").write_text,
            json.dumps(report, ensure_ascii=False, indent=2),
            encoding="utf-8",
        )
        for stage in stages:
            assert set(stage["statuses"]) <= {"200", "429"}, stage
            assert stage["statuses"].get("200", 0) > 0, stage
            if stage["deployment"] == "replicas":
                assert set(stage["replicas"]) == {"api1", "api2"}, stage
    finally:
        with connection() as (client, collection):
            client.delete(collection, filter=f'tenant == "{scenario["tenant"].hex}"')