Function main

server/scripts/test_server_prefix_cache.py:82–133

()

def main():
    args = parse_args()
    base_url = args.url if args.url else spawn_server()

    # Large system prompt (~2K tokens) to make the prefill cost measurable.
    system = "You are a precise coding assistant. " * 200

    def chat(user_msg, max_tokens=8):
        payload = {
            "model": "luce-dflash",
            "messages": [
                {"role": "system", "content": system},
                {"role": "user", "content": user_msg},
            ],
            "max_tokens": max_tokens, "stream": False,
        }
        body = json.dumps(payload).encode()
        req = urllib.request.Request(
            f"{base_url}/v1/chat/completions",
            data=body, headers={"Content-Type": "application/json"})
        t0 = time.time()
        resp = urllib.request.urlopen(req, timeout=600)
        data = json.loads(resp.read())
        dt = time.time() - t0
        return dt, data["choices"][0]["message"]["content"]

    # Turn 1: cold (cache miss → snapshot taken at end)
    print("\n=== Turn 1 (cold) ===", flush=True)
    t1, r1 = chat("What is 2+2?")
    print(f"latency={t1:.2f}s  reply={r1!r}")

    # Turn 2: same system prompt → cache HIT, only suffix prefilled
    print("\n=== Turn 2 (warm) ===", flush=True)
    t2, r2 = chat("What is the capital of France?")
    print(f"latency={t2:.2f}s  reply={r2!r}")

    # Turn 3: same system prompt, third user → still warm
    print("\n=== Turn 3 (warm) ===", flush=True)
    t3, r3 = chat("Tell me about Mars.")
    print(f"latency={t3:.2f}s  reply={r3!r}")

    # Verdict
    print("\n=== Verdict ===", flush=True)
    print(f"turn_1: {t1:.2f}s")
    print(f"turn_2: {t2:.2f}s  ratio_2/1={t2/t1:.2f}")
    print(f"turn_3: {t3:.2f}s  ratio_3/1={t3/t1:.2f}")
    # Expect turn 2 and 3 prefill to be much faster (5K system prompt cached).
    # Total wall is prefill + decode; decode is ~constant (small max_tokens).
    # Conservative gate: ratio < 0.85 (turn 2 should be at least 15% faster).
    ok = (t2 / t1) < 0.85 and (t3 / t1) < 0.85
    print("\nPASS" if ok else "FAIL: prefix cache did not visibly speed up subsequent turns")
    sys.exit(0 if ok else 1)


if __name__ == "__main__":