diff --git a/scripts/test_dspark_stream.py b/scripts/test_dspark_stream.py new file mode 100644 index 0000000..1c961dd --- /dev/null +++ b/scripts/test_dspark_stream.py @@ -0,0 +1,112 @@ +#!/usr/bin/env python3 +"""Accurate-size stream test for frank-default-large (DSpark, PARALLEL=1). + +Fires one prompt and reads back the server's OWN reported prompt_tokens so +the actual stream size is never mislabeled. Then, with parallel=1, fires a +second concurrent stream to confirm it QUEUES behind the single slot and +completes without crashing the ROCm backend. + +The prose is flowing/realistic, so a completion_tokens < max_tokens with +finish_reason "stop" is NOT a degenerate-prompt artifact. +""" +import json, sys, threading, time, urllib.request + +HOST = sys.argv[1] if len(sys.argv) > 1 else "127.0.0.1" +PORT = sys.argv[2] if len(sys.argv) > 2 else "7770" +MAX_TOKENS = int(sys.argv[3]) if len(sys.argv) > 3 else 256 +TARGET_TOKENS = int(sys.argv[4]) if len(sys.argv) > 4 else 16000 + +VAR = ( + "narrative nonfiction blends observation with reflection, ", + "and a careful writer lets each clause earn its place ", + "in the sentence's rising line, so meaning accumulates slowly. ", + "The reader is not hurried; the reader is led. Loose and periodic, ", + "the style turns on concrete nouns and verbs that carry weight, ", + "and it resists the abstract noun that hides a vague claim. ", + "A paragraph is not a list; it is a single turn of thought that ", + "opens, develops, and lands. This is the craft of prose: choose, ", + "arrange, and above all omit. Omission is the discipline that ", + "keeps the whole clear, and clarity is the courtesy we owe the ", + "willing reader." +) + +def build_prompt(topic, salt): + n = len(VAR) + parts = [] + t = 0 + while t * 76 < TARGET_TOKENS: # ~76 tokens per outer loop (measured) + for k in range(5): + parts.append(VAR[(t + k + salt) % n]) + parts.append(f" This is the {topic}, cued by note {salt} added later. ") + t += 1 + return "".join(parts) + +def post(prompt, max_tokens, timeout=3600): + body = json.dumps({ + "prompt": prompt, + "max_tokens": max_tokens, + "temperature": 0.0, + "stream": False, + }).encode() + req = urllib.request.Request( + f"http://{HOST}:{PORT}/v1/completions", + data=body, headers={"Content-Type": "application/json"}) + t0 = time.time() + with urllib.request.urlopen(req, timeout=timeout) as resp: + out = json.load(resp) + dt = time.time() - t0 + return out, dt + +# ---- Stream 1: single 16k prompt ------------------------------- +p1 = build_prompt("first subject", 1) +print(f"built prompt 1: {len(p1)} chars") +o1, t1 = post(p1, MAX_TOKENS) +u1 = o1.get("usage", {}) +print(f"STREAM1 single: prompt_tokens={u1.get('prompt_tokens')} " + f"completion_tokens={u1.get('completion_tokens')} " + f"finish={o1.get('choices',[{}])[0].get('finish_reason')} secs={t1:.1f}") + +# ---- Streams 2 & 3: concurrent, must queue (parallel=1) --------- +results = {} +def worker(key, prompt): + try: + out, dt = post(prompt, MAX_TOKENS, timeout=7200) + u = out.get("usage", {}) + results[key] = { + "prompt_tokens": u.get("prompt_tokens"), + "completion_tokens": u.get("completion_tokens"), + "finish": out.get("choices", [{}])[0].get("finish_reason"), + "secs": round(dt, 1), + } + except Exception as e: + results[key] = {"error": str(e)} + +p2 = build_prompt("second subject", 2) +p3 = build_prompt("third subject", 3) +print(f"built prompt 2 ({len(p2)} chars) and prompt 3 ({len(p3)} chars); concurrent...") +ta = threading.Thread(target=worker, args=("s2", p2)) +tb = threading.Thread(target=worker, args=("s3", p3)) +ta.start(); tb.start(); ta.join(); tb.join() +for k in sorted(results): + print(f" [{k}] {results[k]}") + +# verdict +def suspicious(r): + if "error" in r: + return True + return r.get("finish") == "stop" and r.get("completion_tokens", 0) < MAX_TOKENS + +allr = [("stream1", o1.get("choices",[{}])[0].get("finish_reason"), + u1.get("completion_tokens")), + ("s2", results.get("s2",{}).get("finish"), results.get("s2",{}).get("completion_tokens")), + ("s3", results.get("s3",{}).get("finish"), results.get("s3",{}).get("completion_tokens"))] +print("\nSUMMARY") +ok = True +for name, fin, comp in allr: + sig = fin == "stop" and comp is not None and comp < MAX_TOKENS + err = comp is None + if sig or err: + ok = False + print(f" {name}: finish={fin} completion={comp}/{MAX_TOKENS} " + f"{'SIGNATURE' if sig else ('ERROR' if err else 'ok')}") +print(" RESULT:", "PASS" if ok else "FAIL") \ No newline at end of file