llm_engine.py

# SPDX-License-Identifier: Apache-2.0
# Copyright (c) 2026 Nick Cheng

"""
LLM Engine - Separate process for running MDM inference
Handles model lifecycle, request queuing, and response delivery
"""

import multiprocessing
import time
import sys
import logging
import logging.handlers

from dataclasses import dataclass
from typing import Union, Optional

logger = logging.getLogger("LLMEngine")


@dataclass
class SamplingParams:
    model: str
    prompt: str
    max_tokens: int = 128
    block_length: int = 32
    temperature: float = 0.0
    steps: int = 128
    remasking: str = "low_confidence"


def validate_params(params: SamplingParams) -> Union[str, None]:
    """
    Validates generation parameters. Returns an error message string if invalid, or None if valid.
    """
    if params.max_tokens % params.block_length != 0:
        return f"max_tokens ({params.max_tokens}) must be divisible by block_length ({params.block_length})"

    num_blocks = params.max_tokens // params.block_length
    if params.steps % num_blocks != 0:
        return f"steps ({params.steps}) must be divisible by num_blocks ({num_blocks}) [max_tokens/block_length]"

    return None


def setup_child_logging(log_queue: multiprocessing.Queue):
    """Configures logging for child processes to send logs to the main process via queue."""
    root = logging.getLogger()
    root.setLevel(logging.INFO)
    root.handlers = []  # Clear existing handlers (avoid double logging)
    handler = logging.handlers.QueueHandler(log_queue)
    root.addHandler(handler)


def mock_engine_loop(
    request_queue: multiprocessing.Queue,
    response_queue: multiprocessing.Queue,
    log_queue: multiprocessing.Queue,
    ready_event: multiprocessing.Event,
):
    """
    Fast mock loop for testing API and infrastructure without model overhead.
    """
    setup_child_logging(log_queue)
    logger.info("Mock LLM Engine process started")

    # Simulate a short startup delay
    time.sleep(0.5)
    ready_event.set()
    logger.info("Mock Engine ready.")

    while True:
        request_id = None
        try:
            result = request_queue.get()
            request_id, params = result

            if request_id is None:  # Shutdown signal
                logger.info("Mock Engine received shutdown signal")
                break

            # Validation
            validation_error = validate_params(params)
            if validation_error:
                logger.error(
                    f"Request {request_id} failed validation: {validation_error}"
                )
                response_queue.put((request_id, False, validation_error))
                continue

            logger.info(f"Processing request {request_id} (MOCK)")

            generated_text = f"Mock response for: '{params.prompt}'"

            response_queue.put((request_id, True, generated_text))
            logger.info(f"Finished request {request_id}.")

        except Exception as e:
            logger.error(f"Error in mock engine loop: {e}")
            # If we have a request_id, try to report the error
            if request_id is not None:
                response_queue.put((request_id, False, f"Internal Error: {str(e)}"))


def engine_loop(
    request_queue: multiprocessing.Queue,
    response_queue: multiprocessing.Queue,
    log_queue: multiprocessing.Queue,
    ready_event: multiprocessing.Event,
):
    """
    Main engine loop. Handles input/output queue and performs model loading + generation.
    """
    setup_child_logging(log_queue)
    logger.info("LLM Engine process started")

    try:
        import torch
        from transformers import AutoTokenizer
        from model.modeling_llada import LLaDAModelLM
        from generate import generate_with_dual_cache

        device = "cuda" if torch.cuda.is_available() else "cpu"
        logger.info(f"Loading model on {device}...")

        checkpoint = "GSAI-ML/LLaDA-8B-Instruct"

        model = (
            LLaDAModelLM.from_pretrained(
                checkpoint, trust_remote_code=True, torch_dtype=torch.bfloat16
            )
            .to(device)
            .eval()
        )

        tokenizer = AutoTokenizer.from_pretrained(checkpoint, trust_remote_code=True)

        logger.info("Model loaded successfully.")
        ready_event.set()

    except Exception as e:
        logger.critical(f"Failed to load model environment: {e}")
        # return 503s permanently on engine bricking
        return

    while True:
        request_id = None
        try:
            result = request_queue.get()
            request_id, params = result

            if request_id is None:
                logger.info("Engine received shutdown signal")
                break

            logger.info(f"Processing request {request_id}")

            # Validation
            validation_error = validate_params(params)
            if validation_error:
                logger.error(
                    f"Request {request_id} failed validation: {validation_error}"
                )
                response_queue.put((request_id, False, validation_error))
                continue

            # Tokenize
            m = [{"role": "user", "content": params.prompt}]
            formatted_prompt = tokenizer.apply_chat_template(
                m, add_generation_prompt=True, tokenize=False
            )
            input_ids = tokenizer(formatted_prompt)["input_ids"]
            input_ids = torch.tensor(input_ids).to(device).unsqueeze(0)

            # Generate
            with torch.no_grad():
                out, nfe = generate_with_dual_cache(
                    model,
                    input_ids,
                    steps=params.steps,
                    gen_length=params.max_tokens,
                    block_length=params.block_length,
                    temperature=params.temperature,
                    remasking=params.remasking,
                )

            generated_tokens = out[0][input_ids.shape[1] :]
            answer = tokenizer.decode(generated_tokens, skip_special_tokens=True)

            response_queue.put((request_id, True, answer))
            logger.info(f"Finished request {request_id}")

        except Exception as e:
            logger.error(f"Error in engine loop: {e}")
            if request_id is not None:
                response_queue.put(
                    (request_id, False, f"Error processing request: {str(e)}")
                )


class LLMEngine:
    def __init__(self, is_mock: bool = False):
        self.request_queue = multiprocessing.Queue()
        self.response_queue = multiprocessing.Queue()
        self.log_queue = multiprocessing.Queue()
        self.ready_event = multiprocessing.Event()
        self.process = None
        self.log_listener = None
        self.is_mock = is_mock

    def start(self):
        self.log_listener = logging.handlers.QueueListener(
            self.log_queue, *logging.getLogger().handlers
        )
        self.log_listener.start()

        target_loop = mock_engine_loop if self.is_mock else engine_loop
        loop_name = "LLMEngineWorker_Mock" if self.is_mock else "LLMEngineWorker"

        self.process = multiprocessing.Process(
            target=target_loop,
            args=(
                self.request_queue,
                self.response_queue,
                self.log_queue,
                self.ready_event,
            ),
            daemon=True,
            name=loop_name,
        )
        self.process.start()

    def stop(self):
        if self.process and self.process.is_alive():
            logger.info("Sending shutdown signal to engine...")
            self.request_queue.put((None, None))

            self.process.join(timeout=1.0)

            if self.process.is_alive():
                logger.warning(
                    "Engine did not exit gracefully w/in timeout, forcing termination..."
                )
                self.process.terminate()
                self.process.join()

            logger.info("Engine stopped.")

        if self.log_listener:
            self.log_listener.stop()
            self.log_listener = None

    def submit_request(self, request_id: str, params: SamplingParams):
        self.request_queue.put((request_id, params))