conc_long.py 1.9 KB

1234567891011121314151617181920212223
  1. #!/usr/bin/env python3
  2. """N concurrent distinct needle prompts of ~<tokens> each. Prints per-request result and wall time."""
  3. import json, urllib.request, time, sys, random, concurrent.futures as cf
  4. URL = "http://10.0.0.1:8888/v1/chat/completions"
  5. target = int(sys.argv[1]); n = int(sys.argv[2]); base = int(sys.argv[3]) if len(sys.argv) > 3 else 100
  6. words = ["harbour","lantern","meadow","copper","violin","orchard","compass","thistle","ember","granite","willow","saffron","anchor","quartz","ledger","falcon"]
  7. def build(seed, reps):
  8. rnd = random.Random(seed)
  9. para = lambda i: f"Entry {seed}-{i}: the {rnd.choice(words)} keeper noted {rnd.randint(1,999)} {rnd.choice(words)} crates by the {rnd.choice(words)} gate before the {rnd.choice(words)} bell rang at dusk. "
  10. return "".join(para(i) for i in range(reps)) + f"\nThe secret code word is PELICAN-{seed}.\n" + "".join(para(reps + i) for i in range(20)) + "\nWhat is the secret code word? Answer with the code word only."
  11. def ask(p, n=20):
  12. body = dict(model="deepseek-v4.1-flash", temperature=0, max_tokens=n, stream=False, chat_template_kwargs={"thinking": False}, messages=[{"role": "user", "content": p}])
  13. t0 = time.time()
  14. r = json.load(urllib.request.urlopen(urllib.request.Request(URL, data=json.dumps(body).encode(), headers={"Content-Type": "application/json"}), timeout=3000))
  15. return r["usage"]["prompt_tokens"], time.time() - t0, r["choices"][0]["message"]["content"].strip()
  16. n50, _, _ = ask(build(1, 50), 1); per = (n50 - 20) / 50.0; reps = int((target - 200) / per)
  17. t0 = time.time()
  18. with cf.ThreadPoolExecutor(n) as ex:
  19. rs = list(ex.map(lambda s: ask(build(s, reps)), range(base, base + n)))
  20. wall = time.time() - t0
  21. for i, (pt, dt, txt) in enumerate(rs):
  22. print(f" [{i}] prompt_tokens={pt} {dt:.1f}s answer={txt[:20]!r} {'OK' if txt == f'PELICAN-{base+i}' else 'WRONG'}")
  23. print(f"CONC {n} x ~{target}: wall {wall:.1f}s, total prompt tokens {sum(r[0] for r in rs)}")