Initial ForgeLM import

experiments/bench_localllm_prefill_batch.py

@@ -0,0 +1,175 @@
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import subprocess
+import sys
+import time
+import urllib.error
+import urllib.request
+from pathlib import Path
+
+
+def run(command: list[str], cwd: Path) -> bool:
+    print(f"$ {' '.join(command)}")
+    result = subprocess.run(command, cwd=cwd)
+    return result.returncode == 0
+
+
+def discover_model(root: Path) -> Path | None:
+    explicit = os.environ.get("LOCALLLM_TEST_MODEL")
+    if explicit:
+        path = Path(explicit)
+        return path if path.exists() else None
+
+    models_dir = root / "models"
+    candidates = sorted(models_dir.glob("*.gguf"))
+    return candidates[0] if candidates else None
+
+
+def wait_json(url: str, timeout_secs: float) -> dict | None:
+    deadline = time.time() + timeout_secs
+    while time.time() < deadline:
+        try:
+            with urllib.request.urlopen(url, timeout=5) as response:
+                return json.load(response)
+        except Exception:
+            time.sleep(1)
+    return None
+
+
+def post_json(url: str, payload: dict, timeout_secs: float) -> dict | None:
+    data = json.dumps(payload).encode("utf-8")
+    request = urllib.request.Request(
+        url,
+        data=data,
+        headers={"Content-Type": "application/json"},
+        method="POST",
+    )
+    with urllib.request.urlopen(request, timeout=timeout_secs) as response:
+        return json.load(response)
+
+
+def run_live_completion(root: Path) -> bool:
+    model = discover_model(root)
+    if model is None:
+        print("SKIP live completion: no .gguf model found")
+        return False
+
+    stdout_path = root / "experiments" / "localllm-live-stdout.log"
+    stderr_path = root / "experiments" / "localllm-live-stderr.log"
+    stdout_path.parent.mkdir(parents=True, exist_ok=True)
+
+    env = os.environ.copy()
+    env.update(
+        {
+            "LOCALLLM_MODEL": str(model),
+            "LOCALLLM_ADDR": "127.0.0.1:18083",
+            "LOCALLLM_BATCH": "128",
+            "LOCALLLM_UBATCH": "128",
+            "LOCALLLM_CTX": "16384",
+        }
+    )
+
+    stdout_file = stdout_path.open("w", encoding="utf-8")
+    stderr_file = stderr_path.open("w", encoding="utf-8")
+    process: subprocess.Popen[str] | None = None
+
+    try:
+        process = subprocess.Popen(
+            [str(root / "target" / "debug" / "localllm.exe")],
+            cwd=root,
+            env=env,
+            stdout=stdout_file,
+            stderr=stderr_file,
+            text=True,
+        )
+
+        health = wait_json("http://127.0.0.1:18083/health", timeout_secs=90)
+        if not health or health.get("status") != "ok":
+            return False
+
+        models = wait_json("http://127.0.0.1:18083/v1/models", timeout_secs=15)
+        if not models or not models.get("data"):
+            return False
+
+        model_id = models["data"][0]["id"]
+        prompt = (
+            "token " * 220
+        ) + "Explain in one sentence why chunked prompt prefill matters."
+        response = post_json(
+            "http://127.0.0.1:18083/v1/chat/completions",
+            {
+                "model": model_id,
+                "messages": [{"role": "user", "content": prompt}],
+                "max_tokens": 128,
+            },
+            timeout_secs=180,
+        )
+
+        return bool(response and response.get("choices"))
+    except (
+        OSError,
+        urllib.error.URLError,
+        urllib.error.HTTPError,
+        TimeoutError,
+        json.JSONDecodeError,
+    ):
+        return False
+    finally:
+        if process is not None:
+            process.terminate()
+            try:
+                process.wait(timeout=10)
+            except subprocess.TimeoutExpired:
+                process.kill()
+                process.wait(timeout=10)
+        stdout_file.close()
+        stderr_file.close()
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--root", required=True)
+    args = parser.parse_args()
+
+    root = Path(args.root)
+    checks = {
+        "prefill_plan_passed": run(
+            [
+                "cargo",
+                "test",
+                "-p",
+                "inference-core",
+                "prefill_plan_",
+            ],
+            root,
+        ),
+        "cargo_check_passed": run(
+            ["cargo", "check", "-p", "inference-core"],
+            root,
+        ),
+        "server_build_passed": run(
+            ["cargo", "build", "-p", "inference-server"],
+            root,
+        ),
+    }
+
+    checks["live_completion_passed"] = checks[
+        "server_build_passed"
+    ] and run_live_completion(root)
+
+    passed = sum(1 for ok in checks.values() if ok)
+    total = len(checks)
+    score = (passed / total) * 100.0
+
+    print(f"METRIC prefill_batch_assert_score={score:.1f}")
+    for name, ok in checks.items():
+        print(f"METRIC {name}={1 if ok else 0}")
+
+    return 0 if passed == total else 1
+
+
+if __name__ == "__main__":
+    sys.exit(main())