antirez--ds4/tests/test_server_batching.py

#!/usr/bin/env python3
"""Concurrent API correctness/load check for ds4-server session batching.

Each case is submitted twice with the same non-zero seed. The pairs run in one
cold concurrent wave and must return identical output, even though prompt sizes
and output limits differ across cases. A fresh nonce avoids accidental reuse of
an earlier run's in-memory or disk checkpoint.
"""

import argparse
import concurrent.futures
import json
import math
import statistics
import sys
import threading
import time
import urllib.error
import urllib.request


CASES = (
    ("short-greedy", 0, 80, 0.0, 101),
    ("medium-sampled", 384, 72, 0.7, 202),
    ("long-greedy", 2304, 48, 0.0, 303),
    ("short-sampled", 32, 96, 0.65, 404),
    ("very-long-sampled", 3600, 32, 0.75, 505),
    ("medium-greedy", 1024, 64, 0.0, 606),
    ("long-sampled", 2048, 40, 0.7, 707),
    ("tiny-sampled", 0, 96, 0.6, 808),
)


def percentile(values, fraction):
    ordered = sorted(values)
    return ordered[max(0, math.ceil(len(ordered) * fraction) - 1)]


def make_payload(case, case_number, nonce, stream):
    name, filler_words, max_tokens, temperature, seed = case
    filler = (" alpha beta gamma delta" * ((filler_words + 3) // 4))
    filler = " ".join(filler.split()[:filler_words])
    prompt = (
        "This is batching test %s case %d (%s). Read the filler, then write "
        "a concise explanation of why deterministic request isolation matters. "
        "Do not quote the filler.\nFILLER:\n%s"
    ) % (nonce, case_number, name, filler)
    payload = {
        "model": "deepseek-chat",
        "messages": [{"role": "user", "content": prompt}],
        "max_tokens": max_tokens,
        "temperature": temperature,
        "top_p": 0.9,
        "seed": seed + case_number * 1000,
        "stream": stream,
    }
    if stream:
        payload["stream_options"] = {"include_usage": True}
    return name, filler_words, payload


def parse_stream(raw):
    content = []
    reasoning = []
    finish = None
    usage = None
    for line in raw.decode("utf-8", errors="replace").splitlines():
        if not line.startswith("data:"):
            continue
        data = line[5:].strip()
        if not data or data == "[DONE]":
            continue
        event = json.loads(data)
        if event.get("usage"):
            usage = event["usage"]
        choices = event.get("choices") or []
        if not choices:
            continue
        choice = choices[0]
        delta = choice.get("delta") or {}
        if delta.get("content") is not None:
            content.append(delta["content"])
        if delta.get("reasoning_content") is not None:
            reasoning.append(delta["reasoning_content"])
        if choice.get("finish_reason") is not None:
            finish = choice["finish_reason"]
    return {
        "content": "".join(content),
        "reasoning": "".join(reasoning),
        "finish": finish,
        "completion_tokens": (usage or {}).get("completion_tokens"),
    }


def post_chat(url, payload, timeout, start_event):
    start_event.wait()
    body = json.dumps(payload, separators=(",", ":")).encode("utf-8")
    req = urllib.request.Request(
        url.rstrip("/") + "/v1/chat/completions",
        data=body,
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    started = time.monotonic()
    try:
        with urllib.request.urlopen(req, timeout=timeout) as response:
            raw = response.read()
    except urllib.error.HTTPError as exc:
        detail = exc.read().decode("utf-8", errors="replace")
        raise RuntimeError("HTTP %d: %s" % (exc.code, detail)) from exc
    elapsed = time.monotonic() - started

    if payload["stream"]:
        result = parse_stream(raw)
    else:
        response = json.loads(raw)
        choice = response["choices"][0]
        message = choice.get("message") or {}
        result = {
            "content": message.get("content") or "",
            "reasoning": message.get("reasoning_content") or "",
            "finish": choice.get("finish_reason"),
            "completion_tokens": (response.get("usage") or {}).get(
                "completion_tokens"
            ),
        }
    result["elapsed"] = elapsed
    return result


def comparable(result):
    return (
        result["content"],
        result["reasoning"],
        result["finish"],
        result["completion_tokens"],
    )


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--url", default="http://127.0.0.1:8000")
    parser.add_argument("--pairs", type=int, default=4)
    parser.add_argument("--workers", type=int, default=0)
    parser.add_argument("--timeout", type=float, default=1800.0)
    parser.add_argument("--stream", action="store_true")
    parser.add_argument("--nonce", default="")
    parser.add_argument(
        "--same-prompt", action="store_true",
        help="send one identical prompt and seed in every request",
    )
    parser.add_argument(
        "--max-tokens", type=int,
        help="override each case's completion limit",
    )
    parser.add_argument(
        "--case", choices=[case[0] for case in CASES],
        help="repeat one case shape instead of cycling through mixed lengths",
    )
    args = parser.parse_args()
    if args.pairs <= 0:
        parser.error("--pairs must be positive")

    nonce = args.nonce or "cold-%d" % time.time_ns()
    requests = []
    metadata = []
    for i in range(args.pairs):
        case = next((c for c in CASES if c[0] == args.case), None)
        if case is None:
            case = CASES[i % len(CASES)]
        case_number = 0 if args.same_prompt else i
        name, filler_words, payload = make_payload(
            case, case_number, nonce, args.stream
        )
        if args.max_tokens is not None:
            if args.max_tokens < 0:
                parser.error("--max-tokens must be non-negative")
            payload["max_tokens"] = args.max_tokens
        for copy in range(2):
            requests.append(payload)
            metadata.append((i, copy, name, filler_words))

    workers = args.workers or len(requests)
    workers = max(1, min(workers, len(requests)))
    start_event = threading.Event()
    wall_start = time.monotonic()
    with concurrent.futures.ThreadPoolExecutor(max_workers=workers) as executor:
        futures = [
            executor.submit(post_chat, args.url, payload, args.timeout, start_event)
            for payload in requests
        ]
        start_event.set()
        results = [future.result() for future in futures]
    wall = time.monotonic() - wall_start

    failures = 0
    for i in range(args.pairs):
        a = results[2 * i]
        b = results[2 * i + 1]
        if comparable(a) != comparable(b):
            failures += 1
            print("MISMATCH pair=%d case=%s" % (i, metadata[2 * i][2]), file=sys.stderr)
            print("  A=%s" % (json.dumps(a, ensure_ascii=True)[:1000]), file=sys.stderr)
            print("  B=%s" % (json.dumps(b, ensure_ascii=True)[:1000]), file=sys.stderr)
            for label, result in (("A", a), ("B", b)):
                matches = [
                    index for index, candidate in enumerate(results)
                    if index not in (2 * i, 2 * i + 1)
                    and comparable(candidate) == comparable(result)
                ]
                if matches:
                    print(
                        "  %s also returned by requests=%s" % (label, matches),
                        file=sys.stderr,
                    )
    if args.same_prompt:
        mismatches = [
            index for index, result in enumerate(results[1:], start=1)
            if comparable(result) != comparable(results[0])
        ]
        if mismatches:
            failures += 1
            print(
                "MISMATCH identical prompt requests=%s" % mismatches,
                file=sys.stderr,
            )

    latencies = [result["elapsed"] for result in results]
    known_tokens = [
        result["completion_tokens"]
        for result in results
        if result["completion_tokens"] is not None
    ]
    summary = {
        "status": "PASS" if failures == 0 else "FAIL",
        "pairs": args.pairs,
        "requests": len(requests),
        "workers": workers,
        "stream": args.stream,
        "wall_seconds": round(wall, 3),
        "latency_p50_seconds": round(statistics.median(latencies), 3),
        "latency_p95_seconds": round(percentile(latencies, 0.95), 3),
        "completion_tokens": sum(known_tokens) if known_tokens else None,
        "completion_tokens_per_second": (
            round(sum(known_tokens) / wall, 2) if known_tokens else None
        ),
        "nonce": nonce,
    }
    print(json.dumps(summary, sort_keys=True))
    return 0 if failures == 0 else 1


if __name__ == "__main__":
    raise SystemExit(main())