ForgeLM/experiments/bench_localllm_prefill_batch.py

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())