62 lines
1.9 KiB
Python
62 lines
1.9 KiB
Python
#!/usr/bin/env python3
|
|
"""Issue one non-streaming completion request and record wall/usage."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import time
|
|
import urllib.request
|
|
from pathlib import Path
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--port", type=int, required=True)
|
|
parser.add_argument("--request", type=Path, required=True)
|
|
parser.add_argument("--output", type=Path, required=True)
|
|
parser.add_argument("--timeout-seconds", type=float, default=1800)
|
|
return parser.parse_args()
|
|
|
|
|
|
def main() -> None:
|
|
args = parse_args()
|
|
body = json.loads(args.request.read_text())
|
|
expected_input = len(body["prompt"])
|
|
expected_output = int(body["max_tokens"])
|
|
request = urllib.request.Request(
|
|
f"http://127.0.0.1:{args.port}/v1/completions",
|
|
data=json.dumps(body, separators=(",", ":")).encode(),
|
|
headers={"Content-Type": "application/json"},
|
|
method="POST",
|
|
)
|
|
started = time.monotonic()
|
|
with urllib.request.urlopen(
|
|
request, timeout=args.timeout_seconds
|
|
) as response:
|
|
payload = json.load(response)
|
|
wall_ms = (time.monotonic() - started) * 1000
|
|
usage = payload.get("usage") or {}
|
|
observed = (
|
|
int(usage.get("prompt_tokens", -1)),
|
|
int(usage.get("completion_tokens", -1)),
|
|
)
|
|
if observed != (expected_input, expected_output):
|
|
raise ValueError(
|
|
f"usage mismatch: expected {expected_input}+{expected_output}, "
|
|
f"observed {observed[0]}+{observed[1]}"
|
|
)
|
|
result = {
|
|
"wall_ms": wall_ms,
|
|
"input_tokens": observed[0],
|
|
"output_tokens": observed[1],
|
|
"usage": usage,
|
|
}
|
|
args.output.parent.mkdir(parents=True, exist_ok=True)
|
|
args.output.write_text(json.dumps(result, indent=2, sort_keys=True) + "\n")
|
|
print(json.dumps(result, sort_keys=True), flush=True)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|