Download tests/test_server_batching.py from julianmb/ds4fa: direct link, hf CLI and curl.
- Browser
- Download file 8.87 kB
-
https://huggingface.co/julianmb/ds4fa/resolve/main/tests/test_server_batching.py
- Command line
-
hf download hf://julianmb/ds4fa/tests/test_server_batching.py
-
curl -L -o test_server_batching.py https://huggingface.co/julianmb/ds4fa/resolve/main/tests/test_server_batching.py
8.87 kB
| #!/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()) | |