Salvatore Sanfilippo
ab38bc6b7a
Keep exact sampled requests on ordinary batches. Clear stale verification snapshots, honor logical context limits, and reset speculative timing estimates when batch width changes. Exercise forced rejection, mixed decode, snapshot restore, disk failures and arena replacement.
332 lines
12 KiB
Python
Executable File
332 lines
12 KiB
Python
Executable File
#!/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 cancel_stream(url, payload, timeout, start_event):
|
|
"""Close a streaming response after the first generated text fragment."""
|
|
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()
|
|
with urllib.request.urlopen(req, timeout=timeout) as response:
|
|
while True:
|
|
line = response.readline()
|
|
if not line:
|
|
raise RuntimeError("stream ended before cancellation point")
|
|
if not line.startswith(b"data:"):
|
|
continue
|
|
data = line[5:].strip()
|
|
if not data or data == b"[DONE]":
|
|
continue
|
|
event = json.loads(data)
|
|
choices = event.get("choices") or []
|
|
if not choices:
|
|
continue
|
|
delta = choices[0].get("delta") or {}
|
|
if delta.get("content") or delta.get("reasoning_content"):
|
|
return time.monotonic() - started
|
|
|
|
|
|
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("--model", default="deepseek-chat")
|
|
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",
|
|
)
|
|
parser.add_argument(
|
|
"--cancel-first", type=int, default=0,
|
|
help="abort this many concurrent streams before running the pair oracle",
|
|
)
|
|
args = parser.parse_args()
|
|
if args.pairs <= 0:
|
|
parser.error("--pairs must be positive")
|
|
if args.cancel_first < 0:
|
|
parser.error("--cancel-first must not be negative")
|
|
if args.max_tokens is not None and args.max_tokens < 0:
|
|
parser.error("--max-tokens must be non-negative")
|
|
|
|
nonce = args.nonce or "cold-%d" % time.time_ns()
|
|
cancelled = []
|
|
if args.cancel_first:
|
|
cancel_start = threading.Event()
|
|
cancel_payloads = []
|
|
for i in range(args.cancel_first):
|
|
_, _, payload = make_payload(
|
|
CASES[3], i, nonce + "-cancel", True
|
|
)
|
|
payload["max_tokens"] = max(args.max_tokens or 0, 128)
|
|
payload["model"] = args.model
|
|
cancel_payloads.append(payload)
|
|
with concurrent.futures.ThreadPoolExecutor(
|
|
max_workers=args.cancel_first
|
|
) as executor:
|
|
futures = [
|
|
executor.submit(
|
|
cancel_stream, args.url, payload, args.timeout, cancel_start
|
|
)
|
|
for payload in cancel_payloads
|
|
]
|
|
cancel_start.set()
|
|
cancelled = [future.result() for future in futures]
|
|
# Give the server time to observe closed sockets and recycle every slot.
|
|
time.sleep(0.5)
|
|
|
|
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:
|
|
payload["max_tokens"] = args.max_tokens
|
|
payload["model"] = args.model
|
|
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, result in enumerate(results):
|
|
if result["finish"] not in ("stop", "length") or (
|
|
requests[i]["max_tokens"] > 0 and not (result["content"] or result["reasoning"])
|
|
):
|
|
failures += 1
|
|
print("INVALID request=%d result=%s" % (i, result), file=sys.stderr)
|
|
count = result["completion_tokens"]
|
|
if count is not None and not 0 <= count <= requests[i]["max_tokens"]:
|
|
failures += 1
|
|
print("INVALID token budget request=%d result=%s" % (i, result), file=sys.stderr)
|
|
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
|
|
),
|
|
"cancelled_streams": len(cancelled),
|
|
"cancel_latency_max_seconds": (
|
|
round(max(cancelled), 3) if cancelled else None
|
|
),
|
|
"nonce": nonce,
|
|
}
|
|
print(json.dumps(summary, sort_keys=True))
|
|
return 0 if failures == 0 else 1
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|