training_coach/parser_llama_cpp.py

from functools import lru_cache
import logging
import os
import threading
from time import perf_counter

from training_coach.models import ParsedCheckIn
from training_coach.parser import (
    build_parser_messages,
    log_parser_messages,
    log_parser_response_text,
    parse_model_response,
)


DEFAULT_LLAMA_CPP_MODEL_REPO = "unsloth/Qwen3-1.7B-GGUF"
DEFAULT_LLAMA_CPP_MODEL_FILE = "Qwen3-1.7B-Q4_K_M.gguf"
DEFAULT_LLAMA_CPP_MAX_TOKENS = 512
DEFAULT_LLAMA_CPP_N_CTX = 2048
logger = logging.getLogger(__name__)

# llama-cpp-python is not thread-safe; serializes warmup vs. user requests.
_generate_lock = threading.Lock()


class LlamaCppRuntimeUnavailableError(RuntimeError):
    pass


# Generic JSON with no whitespace between tokens. Minified output saves roughly
# a quarter of the completion tokens versus the pretty-printed JSON the model
# prefers, and the model ignores prompt instructions to minify.
MINIFIED_JSON_GBNF = r"""
root   ::= object
value  ::= object | array | string | number | "true" | "false" | "null"
object ::= "{" ( string ":" value ("," string ":" value)* )? "}"
array  ::= "[" ( value ("," value)* )? "]"
string ::= "\"" ( [^"\\\x7F\x00-\x1F] | "\\" (["\\bfnrt] | "u" [0-9a-fA-F] [0-9a-fA-F] [0-9a-fA-F] [0-9a-fA-F]) )* "\""
number ::= ("-"? ([0-9] | [1-9] [0-9]*)) ("." [0-9]+)? ([eE] [-+]? [0-9]+)?
"""


def _load_llama_cpp():
    try:
        from llama_cpp import Llama, LlamaGrammar
    except ImportError as error:
        raise LlamaCppRuntimeUnavailableError(
            "Install llama-cpp-python to run the GGUF parser backend."
        ) from error

    return Llama, LlamaGrammar


def _optional_int_env(name: str) -> int | None:
    raw_value = os.getenv(name, "").strip()
    if not raw_value:
        return None
    return int(raw_value)


@lru_cache(maxsize=1)
def load_llama_cpp_model(
    repo_id: str = DEFAULT_LLAMA_CPP_MODEL_REPO,
    filename: str = DEFAULT_LLAMA_CPP_MODEL_FILE,
    n_ctx: int = DEFAULT_LLAMA_CPP_N_CTX,
    n_threads: int | None = None,
    n_threads_batch: int | None = None,
):
    start_time = perf_counter()
    # os_cpu_count is logged because containers report the host core count,
    # not the cgroup quota; llama.cpp defaults n_threads_batch to that host
    # count and prefill collapses under CFS throttling when oversubscribed.
    logger.info(
        "event=parser_llama_cpp_model_load_start repo=%s file=%s n_ctx=%s "
        "n_threads=%s n_threads_batch=%s os_cpu_count=%s",
        repo_id,
        filename,
        n_ctx,
        n_threads,
        n_threads_batch,
        os.cpu_count(),
    )
    Llama, _LlamaGrammar = _load_llama_cpp()
    kwargs = {
        "repo_id": repo_id,
        "filename": filename,
        "n_ctx": n_ctx,
        "verbose": False,
    }
    if n_threads is not None:
        kwargs["n_threads"] = n_threads
    if n_threads_batch is not None:
        kwargs["n_threads_batch"] = n_threads_batch

    model = Llama.from_pretrained(**kwargs)
    logger.info(
        "event=parser_llama_cpp_model_load_complete repo=%s file=%s elapsed_ms=%s",
        repo_id,
        filename,
        round((perf_counter() - start_time) * 1000),
    )
    return model


@lru_cache(maxsize=1)
def llama_cpp_json_grammar():
    _Llama, LlamaGrammar = _load_llama_cpp()
    return LlamaGrammar.from_string(MINIFIED_JSON_GBNF, verbose=False)


def build_completion_prompt(messages: list[dict[str, str]]) -> str:
    # Qwen/ChatML prompt built by hand instead of create_chat_completion: Qwen3
    # always wants to open with a <think> block, which a JSON grammar forbids,
    # pushing the model off-distribution (it returned bare "{}"). Prefilling an
    # empty think block keeps generation on-distribution and JSON-only.
    rendered = "".join(
        f"<|im_start|>{message['role']}\n{message['content']}<|im_end|>\n"
        for message in messages
    )
    return rendered + "<|im_start|>assistant\n<think>\n\n</think>\n\n"


def generate_parser_response_llama_cpp(
    raw_text: str,
    *,
    repo_id: str = DEFAULT_LLAMA_CPP_MODEL_REPO,
    filename: str = DEFAULT_LLAMA_CPP_MODEL_FILE,
    max_tokens: int = DEFAULT_LLAMA_CPP_MAX_TOKENS,
    n_ctx: int = DEFAULT_LLAMA_CPP_N_CTX,
    n_threads: int | None = None,
    n_threads_batch: int | None = None,
    warn_on_truncation: bool = True,
) -> str:
    start_time = perf_counter()
    messages = build_parser_messages(raw_text)
    logger.info(
        "event=parser_llama_cpp_generate_start repo=%s file=%s text_chars=%s max_tokens=%s",
        repo_id,
        filename,
        len(raw_text),
        max_tokens,
    )
    log_parser_messages(
        backend="llama_cpp",
        model_name=f"{repo_id}/{filename}",
        messages=messages,
    )
    with _generate_lock:
        model = load_llama_cpp_model(
            repo_id=repo_id,
            filename=filename,
            n_ctx=n_ctx,
            n_threads=n_threads,
            n_threads_batch=n_threads_batch,
        )
        # The small built-in generic JSON grammar replaces the full Pydantic
        # schema grammar, which made per-token sampling unusably slow on Space
        # CPUs. Schema conformance is enforced by parse_model_response.
        response = model.create_completion(
            prompt=build_completion_prompt(messages),
            max_tokens=max_tokens,
            temperature=0,
            grammar=llama_cpp_json_grammar(),
            stop=["<|im_end|>"],
        )
    choice = response["choices"][0]
    response_text = choice["text"].strip()
    usage = response.get("usage", {})
    logger.info(
        "event=parser_llama_cpp_generate_complete repo=%s file=%s "
        "response_chars=%s finish_reason=%s prompt_tokens=%s "
        "completion_tokens=%s elapsed_ms=%s",
        repo_id,
        filename,
        len(response_text),
        choice.get("finish_reason"),
        usage.get("prompt_tokens"),
        usage.get("completion_tokens"),
        round((perf_counter() - start_time) * 1000),
    )
    if warn_on_truncation and choice.get("finish_reason") == "length":
        logger.warning(
            "event=parser_llama_cpp_truncated max_tokens=%s "
            "completion_tokens=%s",
            max_tokens,
            usage.get("completion_tokens"),
        )
    log_parser_response_text(
        backend="llama_cpp",
        model_name=f"{repo_id}/{filename}",
        response_text=response_text,
    )
    return response_text


def llama_cpp_runtime_config() -> dict:
    n_threads = _optional_int_env("LLAMA_CPP_N_THREADS")
    n_threads_batch = _optional_int_env("LLAMA_CPP_N_THREADS_BATCH")
    if n_threads_batch is None:
        # llama.cpp defaults prefill threads to the host core count, which
        # oversubscribes container cgroup quotas; match decode threads instead.
        n_threads_batch = n_threads
    return {
        "repo_id": os.getenv("LLAMA_CPP_MODEL_REPO", DEFAULT_LLAMA_CPP_MODEL_REPO),
        "filename": os.getenv("LLAMA_CPP_MODEL_FILE", DEFAULT_LLAMA_CPP_MODEL_FILE),
        "max_tokens": int(
            os.getenv("LLAMA_CPP_MAX_TOKENS", str(DEFAULT_LLAMA_CPP_MAX_TOKENS))
        ),
        "n_ctx": int(os.getenv("LLAMA_CPP_N_CTX", str(DEFAULT_LLAMA_CPP_N_CTX))),
        "n_threads": n_threads,
        "n_threads_batch": n_threads_batch,
    }


def parse_check_in_with_llama_cpp(raw_text: str) -> ParsedCheckIn:
    response_text = generate_parser_response_llama_cpp(
        raw_text,
        **llama_cpp_runtime_config(),
    )
    return parse_model_response(response_text)


def warm_up_llama_cpp_parser() -> None:
    """Load the model and prefill the constant prompt prefix at startup.

    llama.cpp reuses the KV cache for the longest common prefix between calls,
    and the parser prompt is identical up to the trailing check-in text, so a
    warmup generation makes the first real request only pay for its suffix.
    """
    start_time = perf_counter()
    logger.info("event=parser_llama_cpp_warmup_start")
    config = llama_cpp_runtime_config()
    config["max_tokens"] = 1
    try:
        generate_parser_response_llama_cpp(
            "warmup", warn_on_truncation=False, **config
        )
    except Exception:
        logger.exception("event=parser_llama_cpp_warmup_failed")
        return
    logger.info(
        "event=parser_llama_cpp_warmup_complete elapsed_ms=%s",
        round((perf_counter() - start_time) * 1000),
    )