From b46b3463bb3de112799c47a2722baf6a0be76c0e Mon Sep 17 00:00:00 2001 From: Neil Mehta Date: Wed, 22 Jul 2026 17:35:04 -0400 Subject: [PATCH 01/18] Add MLX server --- mlx_engine/server/__init__.py | 5 + mlx_engine/server/__main__.py | 82 ++++++ mlx_engine/server/chat.py | 216 ++++++++++++++++ mlx_engine/server/http.py | 463 ++++++++++++++++++++++++++++++++++ tests/server/test_chat.py | 263 +++++++++++++++++++ tests/server/test_http.py | 403 +++++++++++++++++++++++++++++ 6 files changed, 1432 insertions(+) create mode 100644 mlx_engine/server/__init__.py create mode 100644 mlx_engine/server/__main__.py create mode 100644 mlx_engine/server/chat.py create mode 100644 mlx_engine/server/http.py create mode 100644 tests/server/test_chat.py create mode 100644 tests/server/test_http.py diff --git a/mlx_engine/server/__init__.py b/mlx_engine/server/__init__.py new file mode 100644 index 00000000..fedec3e6 --- /dev/null +++ b/mlx_engine/server/__init__.py @@ -0,0 +1,5 @@ +"""Private HTTP runtime for mlx-engine.""" + +from .http import EngineRuntime, MlxEngineHttpServer + +__all__ = ["EngineRuntime", "MlxEngineHttpServer"] diff --git a/mlx_engine/server/__main__.py b/mlx_engine/server/__main__.py new file mode 100644 index 00000000..583f1dbc --- /dev/null +++ b/mlx_engine/server/__main__.py @@ -0,0 +1,82 @@ +import argparse +import logging +import signal +import threading + +from mlx_engine import load_model + +from .http import EngineRuntime, MlxEngineHttpServer + + +logger = logging.getLogger(__name__) + + +def _create_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description="Run the private mlx-engine server.") + parser.add_argument("--model", required=True) + parser.add_argument("--host", required=True) + parser.add_argument("--port", required=True, type=int) + parser.add_argument("--api-key", required=True) + parser.add_argument("--context-length", required=True, type=int) + parser.add_argument("--parallel-sessions", required=True, type=int) + parser.add_argument("--seed", type=int) + parser.add_argument("--kv-bits", type=int) + parser.add_argument("--kv-group-size", type=int) + parser.add_argument("--quantized-kv-start", type=int) + parser.add_argument( + "--log-level", + choices=("debug", "info", "warning", "error"), + default="info", + ) + return parser + + +def main() -> None: + args = _create_parser().parse_args() + log_level = getattr(logging, args.log_level.upper()) + engine_logger = logging.getLogger("mlx_engine") + engine_logger.setLevel(log_level) + for handler in engine_logger.handlers: + handler.setLevel(log_level) + + logger.info("Loading MLX model from %s", args.model) + model_kit = load_model( + args.model, + max_kv_size=args.context_length, + max_seq_nums=args.parallel_sessions, + seed=args.seed, + trust_remote_code=False, + kv_bits=args.kv_bits, + kv_group_size=args.kv_group_size, + quantized_kv_start=args.quantized_kv_start, + ) + runtime = EngineRuntime(model_kit) + server = None + + try: + server = MlxEngineHttpServer( + (args.host, args.port), + api_key=args.api_key, + runtime=runtime, + ) + + def request_shutdown(_signal_number: int, _frame: object) -> None: + logger.info("Stopping MLX server") + server.cancel_active_sessions() + threading.Thread(target=server.shutdown, daemon=True).start() + + signal.signal(signal.SIGINT, request_shutdown) + signal.signal(signal.SIGTERM, request_shutdown) + logger.info("MLX server listening on %s:%d", args.host, args.port) + server.serve_forever() + finally: + try: + if server is not None: + server.cancel_active_sessions() + server.server_close() + finally: + runtime.unload() + + +if __name__ == "__main__": + main() diff --git a/mlx_engine/server/chat.py b/mlx_engine/server/chat.py new file mode 100644 index 00000000..d89c1339 --- /dev/null +++ b/mlx_engine/server/chat.py @@ -0,0 +1,216 @@ +from dataclasses import dataclass +import re +from typing import Callable + + +_INLINE_IMAGE_PATTERN = re.compile(r"^data:image/[^;,]+;base64,(.*)$", re.DOTALL) + + +class ChatRequestError(ValueError): + """The chat request does not match the server contract.""" + + +@dataclass(frozen=True) +class ChatGenerationRequest: + prompt: str + prompt_tokens: list[int] + images_b64: list[str] + temperature: float + max_tokens: int | None + stop_strings: list[str] | None + top_p: float | None + top_k: int + min_p: float | None + repetition_penalty: float | None + + +def _normalize_content( + content: object, + *, + role: str, + images_b64: list[str], +) -> object: + if isinstance(content, str) or content is None: + return content + if not isinstance(content, list): + raise ChatRequestError( + f"Message content for role '{role}' must be text or parts." + ) + + normalized_parts: list[dict] = [] + for part in content: + if not isinstance(part, dict): + raise ChatRequestError("Message content parts must be objects.") + part_type = part.get("type") + if part_type == "text": + text = part.get("text") + if not isinstance(text, str): + raise ChatRequestError("Text content parts must contain text.") + normalized_parts.append({"type": "text", "text": text}) + continue + if part_type == "image_url": + if role not in ("user", "tool"): + raise ChatRequestError( + f"Images are not supported in '{role}' messages." + ) + image_url = part.get("image_url") + if not isinstance(image_url, dict): + raise ChatRequestError("Image content parts must contain image_url.") + url = image_url.get("url") + if not isinstance(url, str): + raise ChatRequestError("Image URLs must be strings.") + match = _INLINE_IMAGE_PATTERN.fullmatch(url) + if match is None: + raise ChatRequestError("Images must use inline base64 data URLs.") + images_b64.append(match.group(1)) + normalized_parts.append({"type": "image"}) + continue + raise ChatRequestError(f"Unsupported message content part type: {part_type!r}.") + return normalized_parts + + +def normalize_messages(messages: object) -> tuple[list[dict], list[str]]: + if not isinstance(messages, list): + raise ChatRequestError("messages must be an array.") + + normalized_messages: list[dict] = [] + images_b64: list[str] = [] + for message in messages: + if not isinstance(message, dict): + raise ChatRequestError("messages must contain objects.") + role = message.get("role") + if role not in ("system", "user", "assistant", "tool"): + raise ChatRequestError(f"Unsupported message role: {role!r}.") + + normalized_message = dict(message) + if "content" in normalized_message: + normalized_message["content"] = _normalize_content( + normalized_message["content"], + role=role, + images_b64=images_b64, + ) + normalized_messages.append(normalized_message) + + return normalized_messages, images_b64 + + +def _get_number(body: dict, field_name: str, *, required: bool) -> int | float | None: + value = body.get(field_name) + if value is None and not required: + return None + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise ChatRequestError(f"{field_name} must be a number.") + return value + + +def _get_chat_template(model_kit: object, *, supports_vision: bool) -> Callable: + tokenizer = getattr(model_kit, "tokenizer", None) + if supports_vision: + processor = getattr(model_kit, "processor", None) + candidates = [ + processor, + getattr(processor, "tokenizer", None), + tokenizer, + getattr(tokenizer, "_tokenizer", None), + ] + else: + candidates = [getattr(tokenizer, "_tokenizer", None), tokenizer] + + for renderer in candidates: + apply_chat_template = getattr(renderer, "apply_chat_template", None) + if ( + callable(apply_chat_template) + and getattr(renderer, "chat_template", None) is not None + ): + return apply_chat_template + for renderer in candidates: + apply_chat_template = getattr(renderer, "apply_chat_template", None) + if callable(apply_chat_template): + return apply_chat_template + raise ChatRequestError("The loaded model does not provide a chat template.") + + +def prepare_chat_generation_request( + body: object, + *, + model_kit: object, + supports_vision: bool, + tokenize: Callable[[object, str], list[int]], +) -> ChatGenerationRequest: + if not isinstance(body, dict): + raise ChatRequestError("The request body must be an object.") + if body.get("stream") is not True: + raise ChatRequestError("Streaming generation requires stream=true.") + + normalized_messages, images_b64 = normalize_messages(body.get("messages")) + if len(images_b64) > 0 and not supports_vision: + raise ChatRequestError("The loaded model does not support images.") + + tools = body.get("tools") + if tools is not None and not isinstance(tools, list): + raise ChatRequestError("tools must be an array.") + chat_template_kwargs = body.get("chat_template_kwargs") + if chat_template_kwargs is None: + chat_template_kwargs = {} + if not isinstance(chat_template_kwargs, dict): + raise ChatRequestError("chat_template_kwargs must be an object.") + + template_kwargs = dict(chat_template_kwargs) + if tools is not None and len(tools) > 0: + template_kwargs["tools"] = tools + apply_chat_template = _get_chat_template( + model_kit, + supports_vision=supports_vision, + ) + prompt = apply_chat_template( + normalized_messages, + tokenize=False, + add_generation_prompt=True, + **template_kwargs, + ) + if not isinstance(prompt, str): + raise ChatRequestError("The model chat template did not return text.") + + stop_strings = body.get("stop") + if stop_strings is not None and ( + not isinstance(stop_strings, list) + or any(not isinstance(stop_string, str) for stop_string in stop_strings) + ): + raise ChatRequestError("stop must be an array of strings.") + + max_tokens = _get_number(body, "max_tokens", required=False) + if max_tokens is not None and not isinstance(max_tokens, int): + raise ChatRequestError("max_tokens must be an integer.") + top_k = _get_number(body, "top_k", required=True) + if not isinstance(top_k, int): + raise ChatRequestError("top_k must be an integer.") + + return ChatGenerationRequest( + prompt=prompt, + prompt_tokens=tokenize(model_kit, prompt), + images_b64=images_b64, + temperature=float(_get_number(body, "temperature", required=True)), + max_tokens=max_tokens, + stop_strings=stop_strings, + top_p=( + None + if (top_p := _get_number(body, "top_p", required=False)) is None + else float(top_p) + ), + top_k=top_k, + min_p=( + None + if (min_p := _get_number(body, "min_p", required=False)) is None + else float(min_p) + ), + repetition_penalty=( + None + if ( + repetition_penalty := _get_number( + body, "repeat_penalty", required=False + ) + ) + is None + else float(repetition_penalty) + ), + ) diff --git a/mlx_engine/server/http.py b/mlx_engine/server/http.py new file mode 100644 index 00000000..4514353d --- /dev/null +++ b/mlx_engine/server/http.py @@ -0,0 +1,463 @@ +from __future__ import annotations + +from dataclasses import dataclass +import hmac +from http import HTTPStatus +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +import json +import logging +from queue import Queue +import threading +from typing import Callable, Iterator +import uuid + +from mlx_engine import ( + create_generator, + get_runtime_load_info, + stop_generation, + tokenize, + unload, +) +from mlx_engine.model_kit.batched_vision import BatchedVisionModelKit +from mlx_engine.utils.prompt_progress_reporter import PromptProgressReporter + +from .chat import ( + ChatGenerationRequest, + ChatRequestError, + prepare_chat_generation_request, +) + + +logger = logging.getLogger(__name__) + + +@dataclass(frozen=True) +class _PromptProgressEvent: + total_prompt_tokens: int | None = None + + +@dataclass(frozen=True) +class _GenerationResultEvent: + result: object + + +@dataclass(frozen=True) +class _GenerationErrorEvent: + error: Exception + + +class _GenerationDoneEvent: + pass + + +_GENERATION_DONE = _GenerationDoneEvent() +GenerationEvent = ( + _PromptProgressEvent + | _GenerationResultEvent + | _GenerationErrorEvent + | _GenerationDoneEvent +) + + +class EngineRuntime: + def __init__( + self, + model_kit: object, + *, + supports_vision: bool | None = None, + create_generator_fn: Callable = create_generator, + get_runtime_load_info_fn: Callable = get_runtime_load_info, + stop_generation_fn: Callable = stop_generation, + tokenize_fn: Callable = tokenize, + unload_fn: Callable = unload, + ): + self.model_kit = model_kit + self.supports_vision = ( + isinstance(model_kit, BatchedVisionModelKit) + if supports_vision is None + else supports_vision + ) + self._create_generator = create_generator_fn + self._get_runtime_load_info = get_runtime_load_info_fn + self._stop_generation = stop_generation_fn + self._tokenize = tokenize_fn + self._unload = unload_fn + self._unload_lock = threading.Lock() + self._unloaded = False + + def prepare_chat_generation(self, body: object) -> ChatGenerationRequest: + return prepare_chat_generation_request( + body, + model_kit=self.model_kit, + supports_vision=self.supports_vision, + tokenize=self._tokenize, + ) + + def create_chat_generator( + self, + request: ChatGenerationRequest, + *, + request_id: str, + prompt_progress_reporter: PromptProgressReporter, + ) -> Iterator: + generation_kwargs = { + "images_b64": request.images_b64, + "prompt_progress_reporter": prompt_progress_reporter, + "repetition_penalty": request.repetition_penalty, + "request_id": request_id, + "stop_strings": request.stop_strings, + "temp": request.temperature, + "top_k": request.top_k, + "top_p": request.top_p, + "min_p": request.min_p, + } + if request.max_tokens is not None: + generation_kwargs["max_tokens"] = request.max_tokens + return self._create_generator( + self.model_kit, + request.prompt_tokens, + **generation_kwargs, + ) + + def runtime_context_length(self) -> int | None: + runtime_info = self._get_runtime_load_info(self.model_kit) + context_length = runtime_info.get("context_length") + if isinstance(context_length, int) and context_length > 0: + return context_length + return None + + def stop(self, request_id: str) -> None: + self._stop_generation(self.model_kit, request_id) + + def unload(self) -> None: + with self._unload_lock: + if self._unloaded: + return + self._unloaded = True + self._unload(self.model_kit) + + +class _QueuePromptProgressReporter(PromptProgressReporter): + def __init__( + self, + events: Queue[GenerationEvent], + cancelled: threading.Event, + ): + self._events = events + self._cancelled = cancelled + + def begin( + self, + is_draft: bool, + cached_tokens: int, + total_prompt_tokens: int, + prefill_tokens_processed: int, + ) -> bool: + if not is_draft: + self._events.put( + _PromptProgressEvent(total_prompt_tokens=total_prompt_tokens) + ) + return not self._cancelled.is_set() + + def update(self, is_draft: bool, prefill_tokens_processed: int) -> bool: + if not is_draft: + self._events.put(_PromptProgressEvent()) + return not self._cancelled.is_set() + + def finish( + self, + is_draft: bool, + prefill_tokens_processed: int | None = None, + ) -> bool: + if not is_draft: + self._events.put(_PromptProgressEvent()) + return not self._cancelled.is_set() + + +class GenerationSession: + def __init__(self, runtime: EngineRuntime, request: ChatGenerationRequest): + self.request_id = str(uuid.uuid4()) + self.events: Queue[GenerationEvent] = Queue() + self._runtime = runtime + self._request = request + self._cancelled = threading.Event() + self._done = threading.Event() + self._thread = threading.Thread( + target=self._run, + name=f"mlx-engine-http-generation-{self.request_id}", + daemon=True, + ) + + def start(self) -> None: + self._thread.start() + + def cancel(self) -> None: + if self._cancelled.is_set() or self._done.is_set(): + return + self._cancelled.set() + try: + self._runtime.stop(self.request_id) + except Exception: + logger.exception("Failed to cancel MLX generation %s", self.request_id) + + def join(self, timeout: float | None = None) -> None: + self._thread.join(timeout) + + def _run(self) -> None: + generator = None + try: + if self._cancelled.is_set(): + return + reporter = _QueuePromptProgressReporter(self.events, self._cancelled) + generator = self._runtime.create_chat_generator( + self._request, + request_id=self.request_id, + prompt_progress_reporter=reporter, + ) + if self._cancelled.is_set(): + return + for result in generator: + self.events.put(_GenerationResultEvent(result)) + if self._cancelled.is_set(): + return + except Exception as error: + self.events.put(_GenerationErrorEvent(error)) + finally: + try: + close = getattr(generator, "close", None) + if callable(close): + close() + finally: + self._done.set() + self.events.put(_GENERATION_DONE) + + +class MlxEngineHttpServer(ThreadingHTTPServer): + daemon_threads = True + allow_reuse_address = True + + def __init__( + self, + server_address: tuple[str, int], + *, + api_key: str, + runtime: EngineRuntime, + ): + self.api_key = api_key + self.runtime = runtime + self._active_sessions: set[GenerationSession] = set() + self._active_sessions_lock = threading.Lock() + super().__init__(server_address, MlxEngineRequestHandler) + + def register_session(self, session: GenerationSession) -> None: + with self._active_sessions_lock: + self._active_sessions.add(session) + + def unregister_session(self, session: GenerationSession) -> None: + with self._active_sessions_lock: + self._active_sessions.discard(session) + + def cancel_active_sessions(self) -> None: + with self._active_sessions_lock: + sessions = list(self._active_sessions) + for session in sessions: + session.cancel() + + +class MlxEngineRequestHandler(BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + server: MlxEngineHttpServer + + def do_GET(self) -> None: + if not self._is_authorized(): + self._send_error(HTTPStatus.UNAUTHORIZED, "Unauthorized.") + return + if self.path != "/health": + self._send_error(HTTPStatus.NOT_FOUND, "Not found.") + return + + body: dict[str, object] = {"status": "ok"} + context_length = self.server.runtime.runtime_context_length() + if context_length is not None: + body["context_length"] = context_length + self._send_json(HTTPStatus.OK, body) + + def do_POST(self) -> None: + if not self._is_authorized(): + self._send_error(HTTPStatus.UNAUTHORIZED, "Unauthorized.") + return + if self.path != "/v1/chat/completions": + self._send_error(HTTPStatus.NOT_FOUND, "Not found.") + return + + try: + body = self._read_json_body() + request = self.server.runtime.prepare_chat_generation(body) + except ChatRequestError as error: + self._send_error(HTTPStatus.BAD_REQUEST, str(error)) + return + except Exception as error: + logger.exception("Failed to prepare chat request") + self._send_error(HTTPStatus.INTERNAL_SERVER_ERROR, str(error)) + return + + session = GenerationSession(self.server.runtime, request) + self.server.register_session(session) + normal_completion = False + try: + self.send_response(HTTPStatus.OK) + self.send_header("Content-Type", "text/event-stream") + self.send_header("Cache-Control", "no-cache") + self.send_header("Connection", "close") + self.end_headers() + self.close_connection = True + session.start() + self._stream_generation(session, request) + normal_completion = True + except (BrokenPipeError, ConnectionResetError, OSError): + logger.debug("Generation client disconnected: %s", session.request_id) + except Exception as error: + logger.exception("MLX generation failed") + try: + self._write_sse_json({"error": {"message": str(error)}}) + except (BrokenPipeError, ConnectionResetError, OSError): + pass + finally: + if not normal_completion: + session.cancel() + session.join(timeout=1) + self.server.unregister_session(session) + + def _stream_generation( + self, + session: GenerationSession, + request: ChatGenerationRequest, + ) -> None: + prompt_tokens = len(request.prompt_tokens) + completion_tokens = 0 + terminal_sent = False + + while True: + event = session.events.get() + if isinstance(event, _PromptProgressEvent): + if event.total_prompt_tokens is not None: + prompt_tokens = event.total_prompt_tokens + self._write_bytes(b": prompt-progress\n\n") + continue + if isinstance(event, _GenerationResultEvent): + result = event.result + tokens = getattr(result, "tokens") + completion_tokens += len(tokens) + text = getattr(result, "text") + if text != "": + self._write_sse_json( + { + "choices": [ + { + "index": 0, + "delta": {"content": text}, + "finish_reason": None, + } + ] + } + ) + stop_condition = getattr(result, "stop_condition") + if stop_condition is not None: + self._write_sse_json( + self._terminal_payload( + stop_condition=stop_condition, + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + ) + ) + terminal_sent = True + continue + if isinstance(event, _GenerationErrorEvent): + raise event.error + if isinstance(event, _GenerationDoneEvent): + break + + if not terminal_sent: + raise RuntimeError("MLX generation ended without a stop condition.") + self._write_bytes(b"data: [DONE]\n\n") + + def _terminal_payload( + self, + *, + stop_condition: object, + prompt_tokens: int, + completion_tokens: int, + ) -> dict: + stop_reason = getattr(stop_condition, "stop_reason") + if stop_reason == "eos_token": + finish_reason = "stop" + stop_metadata = {"stop_type": "eos"} + elif stop_reason == "stop_string": + finish_reason = "stop" + stop_metadata = { + "stop_type": "word", + "stopping_word": getattr(stop_condition, "stop_string"), + } + elif stop_reason == "token_limit": + finish_reason = "length" + stop_metadata = {"stop_type": "limit"} + elif stop_reason == "user_cancelled": + finish_reason = "stop" + stop_metadata = {"stop_type": "cancel"} + else: + raise RuntimeError(f"Unknown MLX stop reason: {stop_reason!r}.") + + return { + "choices": [ + { + "index": 0, + "delta": {}, + "finish_reason": finish_reason, + } + ], + "usage": { + "prompt_tokens": prompt_tokens, + "completion_tokens": completion_tokens, + "total_tokens": prompt_tokens + completion_tokens, + }, + "__lmstudio": stop_metadata, + } + + def _read_json_body(self) -> object: + content_length = self.headers.get("Content-Length") + if content_length is None: + raise ChatRequestError("Content-Length is required.") + try: + encoded_body = self.rfile.read(int(content_length)) + return json.loads(encoded_body) + except (UnicodeDecodeError, ValueError, json.JSONDecodeError) as error: + raise ChatRequestError("The request body must be valid JSON.") from error + + def _is_authorized(self) -> bool: + authorization = self.headers.get("Authorization", "") + return hmac.compare_digest( + authorization, + f"Bearer {self.server.api_key}", + ) + + def _write_sse_json(self, body: dict) -> None: + encoded_body = json.dumps(body, separators=(",", ":")).encode("utf-8") + self._write_bytes(b"data: " + encoded_body + b"\n\n") + + def _write_bytes(self, content: bytes) -> None: + self.wfile.write(content) + self.wfile.flush() + + def _send_error(self, status: HTTPStatus, message: str) -> None: + self._send_json(status, {"error": {"message": message}}) + + def _send_json(self, status: HTTPStatus, body: dict) -> None: + encoded_body = json.dumps(body, separators=(",", ":")).encode("utf-8") + self.send_response(status) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(encoded_body))) + self.end_headers() + self.wfile.write(encoded_body) + + def log_message(self, format: str, *args: object) -> None: + logger.debug("HTTP %s - %s", self.address_string(), format % args) diff --git a/tests/server/test_chat.py b/tests/server/test_chat.py new file mode 100644 index 00000000..c1ff6402 --- /dev/null +++ b/tests/server/test_chat.py @@ -0,0 +1,263 @@ +import pytest + +from mlx_engine.server.chat import ( + ChatRequestError, + normalize_messages, + prepare_chat_generation_request, +) + + +class _FakeRenderer: + def __init__(self): + self.calls = [] + + def apply_chat_template(self, messages, **kwargs): + self.calls.append((messages, kwargs)) + return "rendered prompt" + + +class _FakeTokenizerWrapper: + def __init__(self, renderer): + self._tokenizer = renderer + + +class _FakeTextModelKit: + def __init__(self, renderer): + self.tokenizer = _FakeTokenizerWrapper(renderer) + + +class _FakeVisionModelKit: + def __init__(self, renderer): + self.processor = renderer + + +def _base_request(**overrides): + request = { + "model": "ignored-single-model-id", + "messages": [{"role": "user", "content": "Hello"}], + "stream": True, + "stream_options": {"include_usage": True}, + "temperature": 0.7, + "max_tokens": 100, + "stop": ["END"], + "top_p": 0.9, + "top_k": 40, + "min_p": 0.05, + "repeat_penalty": 1.1, + } + request.update(overrides) + return request + + +def test_prepare_text_request_uses_only_supported_generation_settings(): + renderer = _FakeRenderer() + model_kit = _FakeTextModelKit(renderer) + tokenization_calls = [] + + request = prepare_chat_generation_request( + _base_request( + chat_template_kwargs={"reasoning_effort": "medium"}, + tools=[ + { + "type": "function", + "function": { + "name": "search", + "description": "Search", + "parameters": {"type": "object"}, + }, + } + ], + tool_choice="auto", + ), + model_kit=model_kit, + supports_vision=False, + tokenize=lambda received_model_kit, prompt: tokenization_calls.append( + (received_model_kit, prompt) + ) + or [1, 2, 3], + ) + + assert request.prompt == "rendered prompt" + assert request.prompt_tokens == [1, 2, 3] + assert request.images_b64 == [] + assert request.temperature == 0.7 + assert request.max_tokens == 100 + assert request.stop_strings == ["END"] + assert request.top_p == 0.9 + assert request.top_k == 40 + assert request.min_p == 0.05 + assert request.repetition_penalty == 1.1 + assert tokenization_calls == [(model_kit, "rendered prompt")] + + messages, template_kwargs = renderer.calls[0] + assert messages == [{"role": "user", "content": "Hello"}] + assert template_kwargs["tokenize"] is False + assert template_kwargs["add_generation_prompt"] is True + assert template_kwargs["reasoning_effort"] == "medium" + assert template_kwargs["tools"][0]["function"]["name"] == "search" + + +def test_normalize_images_preserves_user_and_tool_result_order(): + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "First"}, + { + "type": "image_url", + "image_url": { + "url": "data:image/jpeg;base64,first-image", + "detail": "auto", + }, + }, + {"type": "text", "text": "Second"}, + ], + }, + { + "role": "assistant", + "content": None, + "reasoning_content": "previous reasoning", + "tool_calls": [ + { + "id": "call-1", + "type": "function", + "function": {"name": "view", "arguments": "{}"}, + } + ], + }, + { + "role": "tool", + "tool_call_id": "call-1", + "content": [ + {"type": "text", "text": "Result"}, + { + "type": "image_url", + "image_url": { + "url": "data:image/png;base64,second-image", + "detail": "auto", + }, + }, + { + "type": "image_url", + "image_url": { + "url": "data:image/jpeg;base64,first-image", + "detail": "auto", + }, + }, + ], + }, + ] + + normalized, images_b64 = normalize_messages(messages) + + assert images_b64 == ["first-image", "second-image", "first-image"] + assert normalized[0]["content"] == [ + {"type": "text", "text": "First"}, + {"type": "image"}, + {"type": "text", "text": "Second"}, + ] + assert normalized[1] == messages[1] + assert normalized[2]["tool_call_id"] == "call-1" + assert normalized[2]["content"] == [ + {"type": "text", "text": "Result"}, + {"type": "image"}, + {"type": "image"}, + ] + + +def test_prepare_vision_request_forwards_base64_to_generation_boundary(): + renderer = _FakeRenderer() + model_kit = _FakeVisionModelKit(renderer) + request = prepare_chat_generation_request( + _base_request( + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "Describe this"}, + { + "type": "image_url", + "image_url": { + "url": "data:image/jpeg;base64,image-payload", + "detail": "auto", + }, + }, + ], + } + ] + ), + model_kit=model_kit, + supports_vision=True, + tokenize=lambda _model_kit, _prompt: [7, 8], + ) + + assert request.images_b64 == ["image-payload"] + assert renderer.calls[0][0][0]["content"] == [ + {"type": "text", "text": "Describe this"}, + {"type": "image"}, + ] + + +def test_vision_request_uses_processor_tokenizer_when_processor_template_is_missing(): + tokenizer_renderer = _FakeRenderer() + tokenizer_renderer.chat_template = "model template" + + class ProcessorWithoutTemplate: + chat_template = None + tokenizer = tokenizer_renderer + + def apply_chat_template(self, _messages, **_kwargs): + raise AssertionError("processor without template must not render") + + model_kit = _FakeVisionModelKit(ProcessorWithoutTemplate()) + request = prepare_chat_generation_request( + _base_request(), + model_kit=model_kit, + supports_vision=True, + tokenize=lambda _model_kit, _prompt: [7, 8], + ) + + assert request.prompt == "rendered prompt" + assert len(tokenizer_renderer.calls) == 1 + + +def test_non_base64_image_url_is_rejected(): + with pytest.raises(ChatRequestError, match="inline base64"): + normalize_messages( + [ + { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": {"url": "https://example.com/image.png"}, + } + ], + } + ] + ) + + +def test_text_model_rejects_image_request(): + renderer = _FakeRenderer() + with pytest.raises(ChatRequestError, match="does not support images"): + prepare_chat_generation_request( + _base_request( + messages=[ + { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": { + "url": "data:image/jpeg;base64,image-payload" + }, + } + ], + } + ] + ), + model_kit=_FakeTextModelKit(renderer), + supports_vision=False, + tokenize=lambda _model_kit, _prompt: [], + ) diff --git a/tests/server/test_http.py b/tests/server/test_http.py new file mode 100644 index 00000000..ec4e7060 --- /dev/null +++ b/tests/server/test_http.py @@ -0,0 +1,403 @@ +from contextlib import contextmanager +import http.client +import json +import socket +import struct +import threading +import time + +from mlx_engine.server.chat import ChatGenerationRequest +from mlx_engine.server.http import ( + EngineRuntime, + GenerationSession, + MlxEngineHttpServer, +) +from mlx_engine.utils.generation_result import ( + GenerationResult, + GenerationStopCondition, +) +from mlx_engine.utils.token import Token + + +class _FakeRenderer: + def apply_chat_template(self, messages, **kwargs): + assert messages == [{"role": "user", "content": "Hello"}] + assert kwargs["tokenize"] is False + assert kwargs["add_generation_prompt"] is True + return "rendered prompt" + + +class _FakeTokenizer: + def __init__(self): + self._tokenizer = _FakeRenderer() + + +class _FakeModelKit: + def __init__(self): + self.tokenizer = _FakeTokenizer() + + +def _request_body(): + return { + "model": "single-loaded-model", + "messages": [{"role": "user", "content": "Hello"}], + "stream": True, + "stream_options": {"include_usage": True}, + "temperature": 0.6, + "max_tokens": 32, + "stop": ["END"], + "top_p": 0.9, + "top_k": 20, + "min_p": 0.03, + "repeat_penalty": 1.05, + } + + +def _parse_sse(response_text): + events = [] + for block in response_text.split("\n\n"): + for line in block.splitlines(): + if not line.startswith("data: "): + continue + data = line.removeprefix("data: ") + if data != "[DONE]": + events.append(json.loads(data)) + return events + + +@contextmanager +def _running_server(runtime): + server = MlxEngineHttpServer( + ("127.0.0.1", 0), + api_key="secret-token", + runtime=runtime, + ) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + yield server.server_address[1] + finally: + server.cancel_active_sessions() + server.shutdown() + server.server_close() + thread.join(timeout=2) + + +def _request(port, method, path, *, body=None, authorized=True): + connection = http.client.HTTPConnection("127.0.0.1", port, timeout=5) + headers = {} + if authorized: + headers["Authorization"] = "Bearer secret-token" + if body is not None: + headers["Content-Type"] = "application/json" + encoded_body = json.dumps(body) + else: + encoded_body = None + connection.request(method, path, body=encoded_body, headers=headers) + response = connection.getresponse() + response_body = response.read().decode("utf-8") + connection.close() + return response.status, response_body + + +def test_health_requires_auth_and_reports_actualized_context_length(): + runtime = EngineRuntime( + _FakeModelKit(), + supports_vision=False, + get_runtime_load_info_fn=lambda _model_kit: {"context_length": 8192}, + ) + + with _running_server(runtime) as port: + status, body = _request(port, "GET", "/health", authorized=False) + assert status == 401 + assert json.loads(body) == {"error": {"message": "Unauthorized."}} + + status, body = _request(port, "GET", "/health") + assert status == 200 + assert json.loads(body) == {"status": "ok", "context_length": 8192} + + +def test_chat_stream_forwards_generation_settings_and_returns_usage(): + generation_calls = [] + + def create_generator(model_kit, prompt_tokens, **kwargs): + generation_calls.append((model_kit, prompt_tokens, kwargs)) + reporter = kwargs["prompt_progress_reporter"] + assert reporter.begin( + is_draft=False, + cached_tokens=3, + total_prompt_tokens=9, + prefill_tokens_processed=0, + ) + assert reporter.update(is_draft=False, prefill_tokens_processed=6) + yield GenerationResult( + text="Hello back", + tokens=[ + Token(id=10, text="Hello", logprob=-0.1), + Token(id=11, text=" back", logprob=-0.2), + ], + top_logprobs=[], + stop_condition=None, + ) + yield GenerationResult( + text="", + tokens=[], + top_logprobs=[], + stop_condition=GenerationStopCondition( + stop_reason="eos_token", + stop_string="", + stop_tokens=[2], + ), + ) + + runtime = EngineRuntime( + _FakeModelKit(), + supports_vision=False, + create_generator_fn=create_generator, + get_runtime_load_info_fn=lambda _model_kit: {}, + tokenize_fn=lambda _model_kit, prompt: [1, 2, 3] + if prompt == "rendered prompt" + else [], + ) + + with _running_server(runtime) as port: + status, response_text = _request( + port, + "POST", + "/v1/chat/completions", + body=_request_body(), + ) + + assert status == 200 + assert ": prompt-progress\n\n" in response_text + assert response_text.endswith("data: [DONE]\n\n") + events = _parse_sse(response_text) + assert events[0] == { + "choices": [ + { + "index": 0, + "delta": {"content": "Hello back"}, + "finish_reason": None, + } + ] + } + assert events[1] == { + "choices": [ + { + "index": 0, + "delta": {}, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 9, + "completion_tokens": 2, + "total_tokens": 11, + }, + "__lmstudio": {"stop_type": "eos"}, + } + + model_kit, prompt_tokens, generation_kwargs = generation_calls[0] + assert isinstance(model_kit, _FakeModelKit) + assert prompt_tokens == [1, 2, 3] + assert generation_kwargs["request_id"] != "" + assert generation_kwargs["images_b64"] == [] + assert generation_kwargs["temp"] == 0.6 + assert generation_kwargs["max_tokens"] == 32 + assert generation_kwargs["stop_strings"] == ["END"] + assert generation_kwargs["top_p"] == 0.9 + assert generation_kwargs["top_k"] == 20 + assert generation_kwargs["min_p"] == 0.03 + assert generation_kwargs["repetition_penalty"] == 1.05 + + +def test_generation_error_is_returned_inside_the_stream(): + def create_generator(_model_kit, _prompt_tokens, **_kwargs): + raise RuntimeError("generation failed") + yield + + runtime = EngineRuntime( + _FakeModelKit(), + supports_vision=False, + create_generator_fn=create_generator, + get_runtime_load_info_fn=lambda _model_kit: {}, + tokenize_fn=lambda _model_kit, _prompt: [1], + ) + + with _running_server(runtime) as port: + status, response_text = _request( + port, + "POST", + "/v1/chat/completions", + body=_request_body(), + ) + + assert status == 200 + assert _parse_sse(response_text) == [{"error": {"message": "generation failed"}}] + + +def test_client_disconnect_stops_the_active_mlx_request(): + generation_stopped = threading.Event() + stopped_request_ids = [] + + def create_generator(_model_kit, _prompt_tokens, **kwargs): + reporter = kwargs["prompt_progress_reporter"] + reporter.begin( + is_draft=False, + cached_tokens=0, + total_prompt_tokens=3, + prefill_tokens_processed=0, + ) + while not generation_stopped.is_set(): + reporter.update(is_draft=False, prefill_tokens_processed=1) + time.sleep(0.01) + yield GenerationResult( + text="", + tokens=[], + top_logprobs=[], + stop_condition=GenerationStopCondition( + stop_reason="user_cancelled", + stop_string="", + stop_tokens=[], + ), + ) + + def stop_generation(_model_kit, request_id): + stopped_request_ids.append(request_id) + generation_stopped.set() + + runtime = EngineRuntime( + _FakeModelKit(), + supports_vision=False, + create_generator_fn=create_generator, + get_runtime_load_info_fn=lambda _model_kit: {}, + stop_generation_fn=stop_generation, + tokenize_fn=lambda _model_kit, _prompt: [1, 2, 3], + ) + + with _running_server(runtime) as port: + encoded_body = json.dumps(_request_body()).encode("utf-8") + client = socket.create_connection(("127.0.0.1", port), timeout=2) + client.sendall( + b"POST /v1/chat/completions HTTP/1.1\r\n" + b"Host: 127.0.0.1\r\n" + b"Authorization: Bearer secret-token\r\n" + b"Content-Type: application/json\r\n" + + f"Content-Length: {len(encoded_body)}\r\n\r\n".encode("ascii") + + encoded_body + ) + received = b"" + while b"\r\n\r\n" not in received: + received += client.recv(4096) + assert b"200 OK" in received + client.setsockopt(socket.SOL_SOCKET, socket.SO_LINGER, struct.pack("ii", 1, 0)) + client.close() + + assert generation_stopped.wait(timeout=2) + + assert len(stopped_request_ids) == 1 + assert stopped_request_ids[0] != "" + + +def test_generation_session_cancellation_stops_the_exact_request(): + generation_stopped = threading.Event() + stopped_request_ids = [] + + def create_generator(_model_kit, _prompt_tokens, **kwargs): + reporter = kwargs["prompt_progress_reporter"] + reporter.begin( + is_draft=False, + cached_tokens=0, + total_prompt_tokens=3, + prefill_tokens_processed=0, + ) + generation_stopped.wait(timeout=2) + yield GenerationResult( + text="", + tokens=[], + top_logprobs=[], + stop_condition=GenerationStopCondition( + stop_reason="user_cancelled", + stop_string="", + stop_tokens=[], + ), + ) + + def stop_generation(_model_kit, request_id): + stopped_request_ids.append(request_id) + generation_stopped.set() + + runtime = EngineRuntime( + _FakeModelKit(), + supports_vision=False, + create_generator_fn=create_generator, + stop_generation_fn=stop_generation, + ) + request = ChatGenerationRequest( + prompt="prompt", + prompt_tokens=[1, 2, 3], + images_b64=[], + temperature=0.8, + max_tokens=None, + stop_strings=None, + top_p=0.95, + top_k=40, + min_p=0.05, + repetition_penalty=1.1, + ) + session = GenerationSession(runtime, request) + + session.start() + session.events.get(timeout=1) + session.cancel() + session.join(timeout=2) + + assert stopped_request_ids == [session.request_id] + assert generation_stopped.is_set() + + +def test_cancellation_failure_does_not_break_cleanup(): + stop_calls = [] + + def stop_generation(_model_kit, request_id): + stop_calls.append(request_id) + raise RuntimeError("backend already stopped") + + runtime = EngineRuntime( + _FakeModelKit(), + supports_vision=False, + stop_generation_fn=stop_generation, + ) + request = ChatGenerationRequest( + prompt="prompt", + prompt_tokens=[1], + images_b64=[], + temperature=0.8, + max_tokens=None, + stop_strings=None, + top_p=0.95, + top_k=40, + min_p=0.05, + repetition_penalty=1.1, + ) + session = GenerationSession(runtime, request) + + session.cancel() + session.cancel() + + assert stop_calls == [session.request_id] + + +def test_runtime_unloads_model_only_once(): + unload_calls = [] + model_kit = _FakeModelKit() + runtime = EngineRuntime( + model_kit, + supports_vision=False, + unload_fn=lambda received_model_kit: unload_calls.append(received_model_kit), + ) + + runtime.unload() + runtime.unload() + + assert unload_calls == [model_kit] From 4112e2b90f272d298180a564ee6a61f0b060ee95 Mon Sep 17 00:00:00 2001 From: Neil Mehta Date: Wed, 22 Jul 2026 18:01:40 -0400 Subject: [PATCH 02/18] Simplify MLX server request handling --- mlx_engine/server/__main__.py | 16 -- mlx_engine/server/chat.py | 266 +++++++++++++--------------------- mlx_engine/server/http.py | 63 +++----- tests/server/test_chat.py | 49 ++++--- tests/server/test_http.py | 23 +-- 5 files changed, 150 insertions(+), 267 deletions(-) diff --git a/mlx_engine/server/__main__.py b/mlx_engine/server/__main__.py index 583f1dbc..d87f2088 100644 --- a/mlx_engine/server/__main__.py +++ b/mlx_engine/server/__main__.py @@ -20,24 +20,11 @@ def _create_parser() -> argparse.ArgumentParser: parser.add_argument("--context-length", required=True, type=int) parser.add_argument("--parallel-sessions", required=True, type=int) parser.add_argument("--seed", type=int) - parser.add_argument("--kv-bits", type=int) - parser.add_argument("--kv-group-size", type=int) - parser.add_argument("--quantized-kv-start", type=int) - parser.add_argument( - "--log-level", - choices=("debug", "info", "warning", "error"), - default="info", - ) return parser def main() -> None: args = _create_parser().parse_args() - log_level = getattr(logging, args.log_level.upper()) - engine_logger = logging.getLogger("mlx_engine") - engine_logger.setLevel(log_level) - for handler in engine_logger.handlers: - handler.setLevel(log_level) logger.info("Loading MLX model from %s", args.model) model_kit = load_model( @@ -46,9 +33,6 @@ def main() -> None: max_seq_nums=args.parallel_sessions, seed=args.seed, trust_remote_code=False, - kv_bits=args.kv_bits, - kv_group_size=args.kv_group_size, - quantized_kv_start=args.quantized_kv_start, ) runtime = EngineRuntime(model_kit) server = None diff --git a/mlx_engine/server/chat.py b/mlx_engine/server/chat.py index d89c1339..f82fd4c2 100644 --- a/mlx_engine/server/chat.py +++ b/mlx_engine/server/chat.py @@ -1,133 +1,101 @@ from dataclasses import dataclass -import re -from typing import Callable +from typing import Annotated, Callable, Literal - -_INLINE_IMAGE_PATTERN = re.compile(r"^data:image/[^;,]+;base64,(.*)$", re.DOTALL) +from pydantic import BaseModel, ConfigDict, Field class ChatRequestError(ValueError): """The chat request does not match the server contract.""" +class _ImageUrl(BaseModel): + url: str + + +class _TextContentPart(BaseModel): + type: Literal["text"] + text: str + + +class _ImageContentPart(BaseModel): + type: Literal["image_url"] + image_url: _ImageUrl + + +_ContentPart = Annotated[ + _TextContentPart | _ImageContentPart, + Field(discriminator="type"), +] + + +class ChatMessage(BaseModel): + model_config = ConfigDict(extra="allow") + + role: Literal["system", "user", "assistant", "tool"] + content: str | list[_ContentPart] | None = None + + +class ChatCompletionRequest(BaseModel): + model_config = ConfigDict(extra="ignore") + + messages: list[ChatMessage] + stream: Literal[True] + temperature: float + max_tokens: int | None = None + stop: list[str] | None = None + top_p: float | None = None + top_k: int + min_p: float | None = None + repeat_penalty: float | None = None + tools: list[dict] | None = None + chat_template_kwargs: dict = Field(default_factory=dict) + + @dataclass(frozen=True) class ChatGenerationRequest: - prompt: str prompt_tokens: list[int] - images_b64: list[str] - temperature: float - max_tokens: int | None - stop_strings: list[str] | None - top_p: float | None - top_k: int - min_p: float | None - repetition_penalty: float | None + generation_kwargs: dict[str, object] -def _normalize_content( - content: object, - *, - role: str, - images_b64: list[str], -) -> object: - if isinstance(content, str) or content is None: - return content - if not isinstance(content, list): - raise ChatRequestError( - f"Message content for role '{role}' must be text or parts." - ) - - normalized_parts: list[dict] = [] - for part in content: - if not isinstance(part, dict): - raise ChatRequestError("Message content parts must be objects.") - part_type = part.get("type") - if part_type == "text": - text = part.get("text") - if not isinstance(text, str): - raise ChatRequestError("Text content parts must contain text.") - normalized_parts.append({"type": "text", "text": text}) - continue - if part_type == "image_url": - if role not in ("user", "tool"): - raise ChatRequestError( - f"Images are not supported in '{role}' messages." - ) - image_url = part.get("image_url") - if not isinstance(image_url, dict): - raise ChatRequestError("Image content parts must contain image_url.") - url = image_url.get("url") - if not isinstance(url, str): - raise ChatRequestError("Image URLs must be strings.") - match = _INLINE_IMAGE_PATTERN.fullmatch(url) - if match is None: - raise ChatRequestError("Images must use inline base64 data URLs.") - images_b64.append(match.group(1)) - normalized_parts.append({"type": "image"}) - continue - raise ChatRequestError(f"Unsupported message content part type: {part_type!r}.") - return normalized_parts - - -def normalize_messages(messages: object) -> tuple[list[dict], list[str]]: - if not isinstance(messages, list): - raise ChatRequestError("messages must be an array.") +def _base64_image_data(url: str) -> str: + header, separator, data = url.partition(",") + if ( + separator == "" + or not header.startswith("data:image/") + or not header.endswith(";base64") + ): + raise ChatRequestError("Images must use inline base64 data URLs.") + return data + +def normalize_messages(messages: list[ChatMessage]) -> tuple[list[dict], list[str]]: normalized_messages: list[dict] = [] images_b64: list[str] = [] + for message in messages: - if not isinstance(message, dict): - raise ChatRequestError("messages must contain objects.") - role = message.get("role") - if role not in ("system", "user", "assistant", "tool"): - raise ChatRequestError(f"Unsupported message role: {role!r}.") - - normalized_message = dict(message) - if "content" in normalized_message: - normalized_message["content"] = _normalize_content( - normalized_message["content"], - role=role, - images_b64=images_b64, - ) + normalized_message = message.model_dump(exclude_unset=True) + if isinstance(message.content, list): + normalized_parts: list[dict] = [] + for part in message.content: + if isinstance(part, _TextContentPart): + normalized_parts.append({"type": "text", "text": part.text}) + else: + images_b64.append(_base64_image_data(part.image_url.url)) + normalized_parts.append({"type": "image"}) + normalized_message["content"] = normalized_parts normalized_messages.append(normalized_message) return normalized_messages, images_b64 -def _get_number(body: dict, field_name: str, *, required: bool) -> int | float | None: - value = body.get(field_name) - if value is None and not required: - return None - if isinstance(value, bool) or not isinstance(value, (int, float)): - raise ChatRequestError(f"{field_name} must be a number.") - return value - - def _get_chat_template(model_kit: object, *, supports_vision: bool) -> Callable: - tokenizer = getattr(model_kit, "tokenizer", None) - if supports_vision: - processor = getattr(model_kit, "processor", None) - candidates = [ - processor, - getattr(processor, "tokenizer", None), - tokenizer, - getattr(tokenizer, "_tokenizer", None), - ] - else: - candidates = [getattr(tokenizer, "_tokenizer", None), tokenizer] - - for renderer in candidates: - apply_chat_template = getattr(renderer, "apply_chat_template", None) - if ( - callable(apply_chat_template) - and getattr(renderer, "chat_template", None) is not None - ): - return apply_chat_template - for renderer in candidates: - apply_chat_template = getattr(renderer, "apply_chat_template", None) - if callable(apply_chat_template): - return apply_chat_template - raise ChatRequestError("The loaded model does not provide a chat template.") + if not supports_vision: + return model_kit.tokenizer._tokenizer.apply_chat_template + + processor = model_kit.processor + if getattr(processor, "chat_template", None) is not None: + return processor.apply_chat_template + return processor.tokenizer.apply_chat_template def prepare_chat_generation_request( @@ -137,80 +105,40 @@ def prepare_chat_generation_request( supports_vision: bool, tokenize: Callable[[object, str], list[int]], ) -> ChatGenerationRequest: - if not isinstance(body, dict): - raise ChatRequestError("The request body must be an object.") - if body.get("stream") is not True: - raise ChatRequestError("Streaming generation requires stream=true.") - - normalized_messages, images_b64 = normalize_messages(body.get("messages")) - if len(images_b64) > 0 and not supports_vision: + request = ChatCompletionRequest.model_validate(body) + normalized_messages, images_b64 = normalize_messages(request.messages) + if images_b64 and not supports_vision: raise ChatRequestError("The loaded model does not support images.") - tools = body.get("tools") - if tools is not None and not isinstance(tools, list): - raise ChatRequestError("tools must be an array.") - chat_template_kwargs = body.get("chat_template_kwargs") - if chat_template_kwargs is None: - chat_template_kwargs = {} - if not isinstance(chat_template_kwargs, dict): - raise ChatRequestError("chat_template_kwargs must be an object.") - - template_kwargs = dict(chat_template_kwargs) - if tools is not None and len(tools) > 0: - template_kwargs["tools"] = tools - apply_chat_template = _get_chat_template( + template_kwargs = dict(request.chat_template_kwargs) + if request.tools: + template_kwargs["tools"] = request.tools + prompt = _get_chat_template( model_kit, supports_vision=supports_vision, - ) - prompt = apply_chat_template( + )( normalized_messages, tokenize=False, add_generation_prompt=True, **template_kwargs, ) - if not isinstance(prompt, str): - raise ChatRequestError("The model chat template did not return text.") - stop_strings = body.get("stop") - if stop_strings is not None and ( - not isinstance(stop_strings, list) - or any(not isinstance(stop_string, str) for stop_string in stop_strings) + generation_kwargs: dict[str, object] = { + "images_b64": images_b64, + "temp": request.temperature, + "top_k": request.top_k, + } + for name, value in ( + ("max_tokens", request.max_tokens), + ("stop_strings", request.stop), + ("top_p", request.top_p), + ("min_p", request.min_p), + ("repetition_penalty", request.repeat_penalty), ): - raise ChatRequestError("stop must be an array of strings.") - - max_tokens = _get_number(body, "max_tokens", required=False) - if max_tokens is not None and not isinstance(max_tokens, int): - raise ChatRequestError("max_tokens must be an integer.") - top_k = _get_number(body, "top_k", required=True) - if not isinstance(top_k, int): - raise ChatRequestError("top_k must be an integer.") + if value is not None: + generation_kwargs[name] = value return ChatGenerationRequest( - prompt=prompt, prompt_tokens=tokenize(model_kit, prompt), - images_b64=images_b64, - temperature=float(_get_number(body, "temperature", required=True)), - max_tokens=max_tokens, - stop_strings=stop_strings, - top_p=( - None - if (top_p := _get_number(body, "top_p", required=False)) is None - else float(top_p) - ), - top_k=top_k, - min_p=( - None - if (min_p := _get_number(body, "min_p", required=False)) is None - else float(min_p) - ), - repetition_penalty=( - None - if ( - repetition_penalty := _get_number( - body, "repeat_penalty", required=False - ) - ) - is None - else float(repetition_penalty) - ), + generation_kwargs=generation_kwargs, ) diff --git a/mlx_engine/server/http.py b/mlx_engine/server/http.py index 4514353d..745e109c 100644 --- a/mlx_engine/server/http.py +++ b/mlx_engine/server/http.py @@ -11,6 +11,8 @@ from typing import Callable, Iterator import uuid +from pydantic import ValidationError + from mlx_engine import ( create_generator, get_runtime_load_info, @@ -19,6 +21,10 @@ unload, ) from mlx_engine.model_kit.batched_vision import BatchedVisionModelKit +from mlx_engine.utils.generation_result import ( + GenerationResult, + GenerationStopCondition, +) from mlx_engine.utils.prompt_progress_reporter import PromptProgressReporter from .chat import ( @@ -38,7 +44,7 @@ class _PromptProgressEvent: @dataclass(frozen=True) class _GenerationResultEvent: - result: object + result: GenerationResult @dataclass(frozen=True) @@ -82,8 +88,6 @@ def __init__( self._stop_generation = stop_generation_fn self._tokenize = tokenize_fn self._unload = unload_fn - self._unload_lock = threading.Lock() - self._unloaded = False def prepare_chat_generation(self, body: object) -> ChatGenerationRequest: return prepare_chat_generation_request( @@ -100,40 +104,21 @@ def create_chat_generator( request_id: str, prompt_progress_reporter: PromptProgressReporter, ) -> Iterator: - generation_kwargs = { - "images_b64": request.images_b64, - "prompt_progress_reporter": prompt_progress_reporter, - "repetition_penalty": request.repetition_penalty, - "request_id": request_id, - "stop_strings": request.stop_strings, - "temp": request.temperature, - "top_k": request.top_k, - "top_p": request.top_p, - "min_p": request.min_p, - } - if request.max_tokens is not None: - generation_kwargs["max_tokens"] = request.max_tokens return self._create_generator( self.model_kit, request.prompt_tokens, - **generation_kwargs, + request_id=request_id, + prompt_progress_reporter=prompt_progress_reporter, + **request.generation_kwargs, ) def runtime_context_length(self) -> int | None: - runtime_info = self._get_runtime_load_info(self.model_kit) - context_length = runtime_info.get("context_length") - if isinstance(context_length, int) and context_length > 0: - return context_length - return None + return self._get_runtime_load_info(self.model_kit).get("context_length") def stop(self, request_id: str) -> None: self._stop_generation(self.model_kit, request_id) def unload(self) -> None: - with self._unload_lock: - if self._unloaded: - return - self._unloaded = True self._unload(self.model_kit) @@ -224,9 +209,8 @@ def _run(self) -> None: self.events.put(_GenerationErrorEvent(error)) finally: try: - close = getattr(generator, "close", None) - if callable(close): - close() + if generator is not None: + generator.close() finally: self._done.set() self.events.put(_GENERATION_DONE) @@ -293,7 +277,7 @@ def do_POST(self) -> None: try: body = self._read_json_body() request = self.server.runtime.prepare_chat_generation(body) - except ChatRequestError as error: + except (ChatRequestError, ValidationError) as error: self._send_error(HTTPStatus.BAD_REQUEST, str(error)) return except Exception as error: @@ -346,26 +330,23 @@ def _stream_generation( continue if isinstance(event, _GenerationResultEvent): result = event.result - tokens = getattr(result, "tokens") - completion_tokens += len(tokens) - text = getattr(result, "text") - if text != "": + completion_tokens += len(result.tokens) + if result.text != "": self._write_sse_json( { "choices": [ { "index": 0, - "delta": {"content": text}, + "delta": {"content": result.text}, "finish_reason": None, } ] } ) - stop_condition = getattr(result, "stop_condition") - if stop_condition is not None: + if result.stop_condition is not None: self._write_sse_json( self._terminal_payload( - stop_condition=stop_condition, + stop_condition=result.stop_condition, prompt_tokens=prompt_tokens, completion_tokens=completion_tokens, ) @@ -384,11 +365,11 @@ def _stream_generation( def _terminal_payload( self, *, - stop_condition: object, + stop_condition: GenerationStopCondition, prompt_tokens: int, completion_tokens: int, ) -> dict: - stop_reason = getattr(stop_condition, "stop_reason") + stop_reason = stop_condition.stop_reason if stop_reason == "eos_token": finish_reason = "stop" stop_metadata = {"stop_type": "eos"} @@ -396,7 +377,7 @@ def _terminal_payload( finish_reason = "stop" stop_metadata = { "stop_type": "word", - "stopping_word": getattr(stop_condition, "stop_string"), + "stopping_word": stop_condition.stop_string, } elif stop_reason == "token_limit": finish_reason = "length" diff --git a/tests/server/test_chat.py b/tests/server/test_chat.py index c1ff6402..b0340346 100644 --- a/tests/server/test_chat.py +++ b/tests/server/test_chat.py @@ -1,6 +1,7 @@ import pytest from mlx_engine.server.chat import ( + ChatMessage, ChatRequestError, normalize_messages, prepare_chat_generation_request, @@ -9,6 +10,7 @@ class _FakeRenderer: def __init__(self): + self.chat_template = "model template" self.calls = [] def apply_chat_template(self, messages, **kwargs): @@ -77,16 +79,17 @@ def test_prepare_text_request_uses_only_supported_generation_settings(): or [1, 2, 3], ) - assert request.prompt == "rendered prompt" assert request.prompt_tokens == [1, 2, 3] - assert request.images_b64 == [] - assert request.temperature == 0.7 - assert request.max_tokens == 100 - assert request.stop_strings == ["END"] - assert request.top_p == 0.9 - assert request.top_k == 40 - assert request.min_p == 0.05 - assert request.repetition_penalty == 1.1 + assert request.generation_kwargs == { + "images_b64": [], + "temp": 0.7, + "max_tokens": 100, + "stop_strings": ["END"], + "top_p": 0.9, + "top_k": 40, + "min_p": 0.05, + "repetition_penalty": 1.1, + } assert tokenization_calls == [(model_kit, "rendered prompt")] messages, template_kwargs = renderer.calls[0] @@ -148,7 +151,9 @@ def test_normalize_images_preserves_user_and_tool_result_order(): }, ] - normalized, images_b64 = normalize_messages(messages) + normalized, images_b64 = normalize_messages( + [ChatMessage.model_validate(message) for message in messages] + ) assert images_b64 == ["first-image", "second-image", "first-image"] assert normalized[0]["content"] == [ @@ -191,7 +196,7 @@ def test_prepare_vision_request_forwards_base64_to_generation_boundary(): tokenize=lambda _model_kit, _prompt: [7, 8], ) - assert request.images_b64 == ["image-payload"] + assert request.generation_kwargs["images_b64"] == ["image-payload"] assert renderer.calls[0][0][0]["content"] == [ {"type": "text", "text": "Describe this"}, {"type": "image"}, @@ -217,7 +222,7 @@ def apply_chat_template(self, _messages, **_kwargs): tokenize=lambda _model_kit, _prompt: [7, 8], ) - assert request.prompt == "rendered prompt" + assert request.prompt_tokens == [7, 8] assert len(tokenizer_renderer.calls) == 1 @@ -225,15 +230,17 @@ def test_non_base64_image_url_is_rejected(): with pytest.raises(ChatRequestError, match="inline base64"): normalize_messages( [ - { - "role": "user", - "content": [ - { - "type": "image_url", - "image_url": {"url": "https://example.com/image.png"}, - } - ], - } + ChatMessage.model_validate( + { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": {"url": "https://example.com/image.png"}, + } + ], + } + ) ] ) diff --git a/tests/server/test_http.py b/tests/server/test_http.py index ec4e7060..d6a9fcac 100644 --- a/tests/server/test_http.py +++ b/tests/server/test_http.py @@ -334,16 +334,8 @@ def stop_generation(_model_kit, request_id): stop_generation_fn=stop_generation, ) request = ChatGenerationRequest( - prompt="prompt", prompt_tokens=[1, 2, 3], - images_b64=[], - temperature=0.8, - max_tokens=None, - stop_strings=None, - top_p=0.95, - top_k=40, - min_p=0.05, - repetition_penalty=1.1, + generation_kwargs={}, ) session = GenerationSession(runtime, request) @@ -369,16 +361,8 @@ def stop_generation(_model_kit, request_id): stop_generation_fn=stop_generation, ) request = ChatGenerationRequest( - prompt="prompt", prompt_tokens=[1], - images_b64=[], - temperature=0.8, - max_tokens=None, - stop_strings=None, - top_p=0.95, - top_k=40, - min_p=0.05, - repetition_penalty=1.1, + generation_kwargs={}, ) session = GenerationSession(runtime, request) @@ -388,7 +372,7 @@ def stop_generation(_model_kit, request_id): assert stop_calls == [session.request_id] -def test_runtime_unloads_model_only_once(): +def test_runtime_unloads_model(): unload_calls = [] model_kit = _FakeModelKit() runtime = EngineRuntime( @@ -397,7 +381,6 @@ def test_runtime_unloads_model_only_once(): unload_fn=lambda received_model_kit: unload_calls.append(received_model_kit), ) - runtime.unload() runtime.unload() assert unload_calls == [model_kit] From 7457b540b379d44a0ea33d932fe8abbecaf81369 Mon Sep 17 00:00:00 2001 From: Neil Mehta Date: Thu, 23 Jul 2026 11:49:34 -0400 Subject: [PATCH 03/18] Clarify inline image schema names --- mlx_engine/server/chat.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/mlx_engine/server/chat.py b/mlx_engine/server/chat.py index f82fd4c2..9ee3f9de 100644 --- a/mlx_engine/server/chat.py +++ b/mlx_engine/server/chat.py @@ -8,7 +8,7 @@ class ChatRequestError(ValueError): """The chat request does not match the server contract.""" -class _ImageUrl(BaseModel): +class _ImageDataUrl(BaseModel): url: str @@ -17,13 +17,13 @@ class _TextContentPart(BaseModel): text: str -class _ImageContentPart(BaseModel): +class _InlineImageContentPart(BaseModel): type: Literal["image_url"] - image_url: _ImageUrl + image_url: _ImageDataUrl _ContentPart = Annotated[ - _TextContentPart | _ImageContentPart, + _TextContentPart | _InlineImageContentPart, Field(discriminator="type"), ] From 5357a16c458b32388e13bb53ac588abc5e0f4af4 Mon Sep 17 00:00:00 2001 From: Neil Mehta Date: Thu, 23 Jul 2026 13:21:30 -0400 Subject: [PATCH 04/18] Keep chat template rendering model-owned --- mlx_engine/server/chat.py | 18 ++++++++++++++++++ tests/server/test_chat.py | 24 ++++++++++++++++++++++++ 2 files changed, 42 insertions(+) diff --git a/mlx_engine/server/chat.py b/mlx_engine/server/chat.py index 9ee3f9de..6b956170 100644 --- a/mlx_engine/server/chat.py +++ b/mlx_engine/server/chat.py @@ -4,6 +4,15 @@ from pydantic import BaseModel, ConfigDict, Field +_CHAT_TEMPLATE_CONTROL_KEYS = { + "add_generation_prompt", + "chat_template", + "continue_final_message", + "tokenize", + "tools", +} + + class ChatRequestError(ValueError): """The chat request does not match the server contract.""" @@ -110,6 +119,15 @@ def prepare_chat_generation_request( if images_b64 and not supports_vision: raise ChatRequestError("The loaded model does not support images.") + overridden_controls = _CHAT_TEMPLATE_CONTROL_KEYS.intersection( + request.chat_template_kwargs + ) + if overridden_controls: + names = ", ".join(sorted(overridden_controls)) + raise ChatRequestError( + f"chat_template_kwargs cannot override server rendering controls: {names}." + ) + template_kwargs = dict(request.chat_template_kwargs) if request.tools: template_kwargs["tools"] = request.tools diff --git a/tests/server/test_chat.py b/tests/server/test_chat.py index b0340346..ae2fc005 100644 --- a/tests/server/test_chat.py +++ b/tests/server/test_chat.py @@ -100,6 +100,30 @@ def test_prepare_text_request_uses_only_supported_generation_settings(): assert template_kwargs["tools"][0]["function"]["name"] == "search" +@pytest.mark.parametrize( + "control_name", + [ + "add_generation_prompt", + "chat_template", + "continue_final_message", + "tokenize", + "tools", + ], +) +def test_chat_template_kwargs_cannot_override_server_controls(control_name): + renderer = _FakeRenderer() + + with pytest.raises(ChatRequestError, match="server rendering controls"): + prepare_chat_generation_request( + _base_request(chat_template_kwargs={control_name: "override"}), + model_kit=_FakeTextModelKit(renderer), + supports_vision=False, + tokenize=lambda _model_kit, _prompt: [], + ) + + assert renderer.calls == [] + + def test_normalize_images_preserves_user_and_tool_result_order(): messages = [ { From abbeebc28a37d04454a5df7d2a8c6cda0b0b7789 Mon Sep 17 00:00:00 2001 From: Neil Mehta Date: Thu, 23 Jul 2026 13:21:30 -0400 Subject: [PATCH 05/18] Stream HTTP generation on request threads --- mlx_engine/server/http.py | 231 ++++++++++++++------------------------ tests/server/test_http.py | 40 +------ 2 files changed, 87 insertions(+), 184 deletions(-) diff --git a/mlx_engine/server/http.py b/mlx_engine/server/http.py index 745e109c..2fb7844f 100644 --- a/mlx_engine/server/http.py +++ b/mlx_engine/server/http.py @@ -1,12 +1,10 @@ from __future__ import annotations -from dataclasses import dataclass import hmac from http import HTTPStatus from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer import json import logging -from queue import Queue import threading from typing import Callable, Iterator import uuid @@ -37,34 +35,6 @@ logger = logging.getLogger(__name__) -@dataclass(frozen=True) -class _PromptProgressEvent: - total_prompt_tokens: int | None = None - - -@dataclass(frozen=True) -class _GenerationResultEvent: - result: GenerationResult - - -@dataclass(frozen=True) -class _GenerationErrorEvent: - error: Exception - - -class _GenerationDoneEvent: - pass - - -_GENERATION_DONE = _GenerationDoneEvent() -GenerationEvent = ( - _PromptProgressEvent - | _GenerationResultEvent - | _GenerationErrorEvent - | _GenerationDoneEvent -) - - class EngineRuntime: def __init__( self, @@ -122,14 +92,40 @@ def unload(self) -> None: self._unload(self.model_kit) -class _QueuePromptProgressReporter(PromptProgressReporter): +class GenerationSession: + def __init__(self, runtime: EngineRuntime): + self.request_id = str(uuid.uuid4()) + self._runtime = runtime + self._cancelled = threading.Event() + self._done = threading.Event() + + @property + def cancelled(self) -> bool: + return self._cancelled.is_set() + + def cancel(self) -> None: + if self._cancelled.is_set() or self._done.is_set(): + return + self._cancelled.set() + try: + self._runtime.stop(self.request_id) + except Exception: + logger.exception("Failed to cancel MLX generation %s", self.request_id) + + def finish(self) -> None: + self._done.set() + + +class _SsePromptProgressReporter(PromptProgressReporter): def __init__( self, - events: Queue[GenerationEvent], - cancelled: threading.Event, + handler: MlxEngineRequestHandler, + session: GenerationSession, + prompt_tokens: int, ): - self._events = events - self._cancelled = cancelled + self._handler = handler + self._session = session + self.prompt_tokens = prompt_tokens def begin( self, @@ -139,81 +135,24 @@ def begin( prefill_tokens_processed: int, ) -> bool: if not is_draft: - self._events.put( - _PromptProgressEvent(total_prompt_tokens=total_prompt_tokens) - ) - return not self._cancelled.is_set() + self.prompt_tokens = total_prompt_tokens + return self._report(is_draft) def update(self, is_draft: bool, prefill_tokens_processed: int) -> bool: - if not is_draft: - self._events.put(_PromptProgressEvent()) - return not self._cancelled.is_set() + return self._report(is_draft) def finish( self, is_draft: bool, prefill_tokens_processed: int | None = None, ) -> bool: - if not is_draft: - self._events.put(_PromptProgressEvent()) - return not self._cancelled.is_set() - - -class GenerationSession: - def __init__(self, runtime: EngineRuntime, request: ChatGenerationRequest): - self.request_id = str(uuid.uuid4()) - self.events: Queue[GenerationEvent] = Queue() - self._runtime = runtime - self._request = request - self._cancelled = threading.Event() - self._done = threading.Event() - self._thread = threading.Thread( - target=self._run, - name=f"mlx-engine-http-generation-{self.request_id}", - daemon=True, - ) + return self._report(is_draft) - def start(self) -> None: - self._thread.start() - - def cancel(self) -> None: - if self._cancelled.is_set() or self._done.is_set(): - return - self._cancelled.set() - try: - self._runtime.stop(self.request_id) - except Exception: - logger.exception("Failed to cancel MLX generation %s", self.request_id) - - def join(self, timeout: float | None = None) -> None: - self._thread.join(timeout) - - def _run(self) -> None: - generator = None - try: - if self._cancelled.is_set(): - return - reporter = _QueuePromptProgressReporter(self.events, self._cancelled) - generator = self._runtime.create_chat_generator( - self._request, - request_id=self.request_id, - prompt_progress_reporter=reporter, - ) - if self._cancelled.is_set(): - return - for result in generator: - self.events.put(_GenerationResultEvent(result)) - if self._cancelled.is_set(): - return - except Exception as error: - self.events.put(_GenerationErrorEvent(error)) - finally: - try: - if generator is not None: - generator.close() - finally: - self._done.set() - self.events.put(_GENERATION_DONE) + def _report(self, is_draft: bool) -> bool: + if is_draft or self._session.cancelled: + return not self._session.cancelled + self._handler._write_bytes(b": prompt-progress\n\n") + return True class MlxEngineHttpServer(ThreadingHTTPServer): @@ -285,8 +224,9 @@ def do_POST(self) -> None: self._send_error(HTTPStatus.INTERNAL_SERVER_ERROR, str(error)) return - session = GenerationSession(self.server.runtime, request) + session = GenerationSession(self.server.runtime) self.server.register_session(session) + generator = None normal_completion = False try: self.send_response(HTTPStatus.OK) @@ -295,8 +235,17 @@ def do_POST(self) -> None: self.send_header("Connection", "close") self.end_headers() self.close_connection = True - session.start() - self._stream_generation(session, request) + reporter = _SsePromptProgressReporter( + self, + session, + len(request.prompt_tokens), + ) + generator = self.server.runtime.create_chat_generator( + request, + request_id=session.request_id, + prompt_progress_reporter=reporter, + ) + self._stream_generation(generator, reporter) normal_completion = True except (BrokenPipeError, ConnectionResetError, OSError): logger.debug("Generation client disconnected: %s", session.request_id) @@ -307,56 +256,46 @@ def do_POST(self) -> None: except (BrokenPipeError, ConnectionResetError, OSError): pass finally: - if not normal_completion: - session.cancel() - session.join(timeout=1) - self.server.unregister_session(session) + try: + if not normal_completion: + session.cancel() + if generator is not None: + generator.close() + finally: + session.finish() + self.server.unregister_session(session) def _stream_generation( self, - session: GenerationSession, - request: ChatGenerationRequest, + generator: Iterator[GenerationResult], + reporter: _SsePromptProgressReporter, ) -> None: - prompt_tokens = len(request.prompt_tokens) completion_tokens = 0 terminal_sent = False - while True: - event = session.events.get() - if isinstance(event, _PromptProgressEvent): - if event.total_prompt_tokens is not None: - prompt_tokens = event.total_prompt_tokens - self._write_bytes(b": prompt-progress\n\n") - continue - if isinstance(event, _GenerationResultEvent): - result = event.result - completion_tokens += len(result.tokens) - if result.text != "": - self._write_sse_json( - { - "choices": [ - { - "index": 0, - "delta": {"content": result.text}, - "finish_reason": None, - } - ] - } - ) - if result.stop_condition is not None: - self._write_sse_json( - self._terminal_payload( - stop_condition=result.stop_condition, - prompt_tokens=prompt_tokens, - completion_tokens=completion_tokens, - ) + for result in generator: + completion_tokens += len(result.tokens) + if result.text != "": + self._write_sse_json( + { + "choices": [ + { + "index": 0, + "delta": {"content": result.text}, + "finish_reason": None, + } + ] + } + ) + if result.stop_condition is not None: + self._write_sse_json( + self._terminal_payload( + stop_condition=result.stop_condition, + prompt_tokens=reporter.prompt_tokens, + completion_tokens=completion_tokens, ) - terminal_sent = True - continue - if isinstance(event, _GenerationErrorEvent): - raise event.error - if isinstance(event, _GenerationDoneEvent): - break + ) + terminal_sent = True if not terminal_sent: raise RuntimeError("MLX generation ended without a stop condition.") diff --git a/tests/server/test_http.py b/tests/server/test_http.py index d6a9fcac..f359ee8c 100644 --- a/tests/server/test_http.py +++ b/tests/server/test_http.py @@ -6,7 +6,6 @@ import threading import time -from mlx_engine.server.chat import ChatGenerationRequest from mlx_engine.server.http import ( EngineRuntime, GenerationSession, @@ -300,52 +299,21 @@ def stop_generation(_model_kit, request_id): def test_generation_session_cancellation_stops_the_exact_request(): - generation_stopped = threading.Event() stopped_request_ids = [] - def create_generator(_model_kit, _prompt_tokens, **kwargs): - reporter = kwargs["prompt_progress_reporter"] - reporter.begin( - is_draft=False, - cached_tokens=0, - total_prompt_tokens=3, - prefill_tokens_processed=0, - ) - generation_stopped.wait(timeout=2) - yield GenerationResult( - text="", - tokens=[], - top_logprobs=[], - stop_condition=GenerationStopCondition( - stop_reason="user_cancelled", - stop_string="", - stop_tokens=[], - ), - ) - def stop_generation(_model_kit, request_id): stopped_request_ids.append(request_id) - generation_stopped.set() runtime = EngineRuntime( _FakeModelKit(), supports_vision=False, - create_generator_fn=create_generator, stop_generation_fn=stop_generation, ) - request = ChatGenerationRequest( - prompt_tokens=[1, 2, 3], - generation_kwargs={}, - ) - session = GenerationSession(runtime, request) + session = GenerationSession(runtime) - session.start() - session.events.get(timeout=1) session.cancel() - session.join(timeout=2) assert stopped_request_ids == [session.request_id] - assert generation_stopped.is_set() def test_cancellation_failure_does_not_break_cleanup(): @@ -360,11 +328,7 @@ def stop_generation(_model_kit, request_id): supports_vision=False, stop_generation_fn=stop_generation, ) - request = ChatGenerationRequest( - prompt_tokens=[1], - generation_kwargs={}, - ) - session = GenerationSession(runtime, request) + session = GenerationSession(runtime) session.cancel() session.cancel() From 0e69e26eea360fe650753dbdb141148b5b95d156 Mon Sep 17 00:00:00 2001 From: Neil Mehta Date: Thu, 23 Jul 2026 14:36:52 -0400 Subject: [PATCH 06/18] Admit batched requests between generation steps --- mlx_engine/model_kit/batched_model_kit.py | 2 +- tests/test_batched_model_kit_scheduler.py | 92 +++++++++++++++++++++++ 2 files changed, 93 insertions(+), 1 deletion(-) create mode 100644 tests/test_batched_model_kit_scheduler.py diff --git a/mlx_engine/model_kit/batched_model_kit.py b/mlx_engine/model_kit/batched_model_kit.py index f0f713ec..88dcd7ac 100644 --- a/mlx_engine/model_kit/batched_model_kit.py +++ b/mlx_engine/model_kit/batched_model_kit.py @@ -385,7 +385,7 @@ def get_next_request(timeout=None): time_budget = 0.5 start = time.time() while True: - if time.time() - start > time_budget: + if not self._requests.empty() or time.time() - start > time_budget: break prompt_responses, generation_responses = batch_generator.next() diff --git a/tests/test_batched_model_kit_scheduler.py b/tests/test_batched_model_kit_scheduler.py new file mode 100644 index 00000000..fc8e3dda --- /dev/null +++ b/tests/test_batched_model_kit_scheduler.py @@ -0,0 +1,92 @@ +from contextlib import nullcontext +from queue import Queue +from types import SimpleNamespace +import threading + +import mlx_engine.model_kit.batched_model_kit as batched_model_kit_module +from mlx_engine.model_kit.batched_model_kit import BatchedModelKit +from mlx_engine.model_kit.batched_model_kit_types import GenerationRequest + + +def _request(request_id: str) -> GenerationRequest: + return GenerationRequest( + rqueue=Queue(), + prompt_tokens=[1, 2, 3], + request_id=request_id, + samplers=None, + logits_processors=[], + top_logprobs=0, + max_tokens=10, + ) + + +def test_pending_request_is_admitted_before_the_next_generation_step(monkeypatch): + first_request = _request("first") + second_request = _request("second") + inserted_request_count = 0 + generation_step_count = 0 + + class FakeBatchGenerator: + stream = object() + + def __init__(self, *_args, **_kwargs): + pass + + def insert(self, *_args, **_kwargs): + nonlocal inserted_request_count + uid = inserted_request_count + inserted_request_count += 1 + if inserted_request_count == 2: + model_kit._shutdown.set() + return (uid,) + + def next(self): + nonlocal generation_step_count + generation_step_count += 1 + if generation_step_count == 1: + model_kit._requests.put(second_request) + return [SimpleNamespace(uid=0, progress=(1, 3))], [] + model_kit._shutdown.set() + return [], [] + + monkeypatch.setattr(batched_model_kit_module, "BatchGenerator", FakeBatchGenerator) + monkeypatch.setattr( + batched_model_kit_module, + "_prepare_prompt_cache_for_generation", + lambda *_args: (None, [], [1, 2, 3]), + ) + monkeypatch.setattr( + batched_model_kit_module, + "install_mlx_compile_cache_cleanup_for_thread", + lambda: None, + ) + monkeypatch.setattr(batched_model_kit_module, "set_seed", lambda _seed: None) + monkeypatch.setattr( + batched_model_kit_module.mx, + "stream", + lambda _stream: nullcontext(), + ) + + model_kit = BatchedModelKit.__new__(BatchedModelKit) + model_kit.model = object() + model_kit.tokenizer = SimpleNamespace(detokenizer=object(), eos_token_ids=[]) + model_kit._requests = Queue() + model_kit._requests.put(first_request) + model_kit._prompt_cache = object() + model_kit._batch_results = {} + model_kit._backend_exception = None + model_kit._generation_thread = None + model_kit._shutdown = threading.Event() + model_kit._startup_complete = threading.Event() + model_kit._seed = None + model_kit._max_seq_nums = 4 + model_kit._prefill_step_size = 512 + model_kit._max_kv_size = 2048 + + scheduler_thread = threading.Thread(target=model_kit._generate) + scheduler_thread.start() + scheduler_thread.join(timeout=2) + + assert not scheduler_thread.is_alive() + assert inserted_request_count == 2 + assert generation_step_count == 1 From d281f03f78c98394e9669a9dc87e394673c2597d Mon Sep 17 00:00:00 2001 From: Neil Mehta Date: Thu, 23 Jul 2026 16:35:21 -0400 Subject: [PATCH 07/18] Handle structured output and assistant prefills --- mlx_engine/server/chat.py | 28 ++++++++++-- tests/server/test_chat.py | 95 ++++++++++++++++++++++++++++++++------- tests/server/test_http.py | 40 ++++++++++++++++- 3 files changed, 144 insertions(+), 19 deletions(-) diff --git a/mlx_engine/server/chat.py b/mlx_engine/server/chat.py index 6b956170..83582f0d 100644 --- a/mlx_engine/server/chat.py +++ b/mlx_engine/server/chat.py @@ -1,4 +1,5 @@ from dataclasses import dataclass +import json from typing import Annotated, Callable, Literal from pydantic import BaseModel, ConfigDict, Field @@ -44,6 +45,15 @@ class ChatMessage(BaseModel): content: str | list[_ContentPart] | None = None +class _JsonSchemaDefinition(BaseModel): + schema_: object = Field(alias="schema") + + +class _JsonSchemaResponseFormat(BaseModel): + type: Literal["json_schema"] + json_schema: _JsonSchemaDefinition + + class ChatCompletionRequest(BaseModel): model_config = ConfigDict(extra="ignore") @@ -57,6 +67,7 @@ class ChatCompletionRequest(BaseModel): min_p: float | None = None repeat_penalty: float | None = None tools: list[dict] | None = None + response_format: _JsonSchemaResponseFormat | None = None chat_template_kwargs: dict = Field(default_factory=dict) @@ -115,6 +126,9 @@ def prepare_chat_generation_request( tokenize: Callable[[object, str], list[int]], ) -> ChatGenerationRequest: request = ChatCompletionRequest.model_validate(body) + if request.tools: + raise ChatRequestError("Tools are not supported yet.") + normalized_messages, images_b64 = normalize_messages(request.messages) if images_b64 and not supports_vision: raise ChatRequestError("The loaded model does not support images.") @@ -129,15 +143,19 @@ def prepare_chat_generation_request( ) template_kwargs = dict(request.chat_template_kwargs) - if request.tools: - template_kwargs["tools"] = request.tools + if request.messages and request.messages[-1].role == "assistant": + template_kwargs["continue_final_message"] = True + add_generation_prompt = False + else: + add_generation_prompt = True + prompt = _get_chat_template( model_kit, supports_vision=supports_vision, )( normalized_messages, tokenize=False, - add_generation_prompt=True, + add_generation_prompt=add_generation_prompt, **template_kwargs, ) @@ -146,6 +164,10 @@ def prepare_chat_generation_request( "temp": request.temperature, "top_k": request.top_k, } + if request.response_format is not None: + generation_kwargs["json_schema"] = json.dumps( + request.response_format.json_schema.schema_ + ) for name, value in ( ("max_tokens", request.max_tokens), ("stop_strings", request.stop), diff --git a/tests/server/test_chat.py b/tests/server/test_chat.py index ae2fc005..5f53231b 100644 --- a/tests/server/test_chat.py +++ b/tests/server/test_chat.py @@ -1,4 +1,7 @@ +import json + import pytest +from pydantic import ValidationError from mlx_engine.server.chat import ( ChatMessage, @@ -57,20 +60,7 @@ def test_prepare_text_request_uses_only_supported_generation_settings(): tokenization_calls = [] request = prepare_chat_generation_request( - _base_request( - chat_template_kwargs={"reasoning_effort": "medium"}, - tools=[ - { - "type": "function", - "function": { - "name": "search", - "description": "Search", - "parameters": {"type": "object"}, - }, - } - ], - tool_choice="auto", - ), + _base_request(chat_template_kwargs={"reasoning_effort": "medium"}), model_kit=model_kit, supports_vision=False, tokenize=lambda received_model_kit, prompt: tokenization_calls.append( @@ -96,8 +86,83 @@ def test_prepare_text_request_uses_only_supported_generation_settings(): assert messages == [{"role": "user", "content": "Hello"}] assert template_kwargs["tokenize"] is False assert template_kwargs["add_generation_prompt"] is True + assert "continue_final_message" not in template_kwargs assert template_kwargs["reasoning_effort"] == "medium" - assert template_kwargs["tools"][0]["function"]["name"] == "search" + + +def test_tools_are_rejected(): + renderer = _FakeRenderer() + + with pytest.raises(ChatRequestError, match="Tools are not supported yet"): + prepare_chat_generation_request( + _base_request( + tools=[ + { + "type": "function", + "function": {"name": "search"}, + } + ] + ), + model_kit=_FakeTextModelKit(renderer), + supports_vision=False, + tokenize=lambda _model_kit, _prompt: [], + ) + + assert renderer.calls == [] + + +def test_json_schema_is_forwarded_to_generation(): + renderer = _FakeRenderer() + schema = { + "type": "object", + "properties": {"answer": {"type": "string"}}, + "required": ["answer"], + } + + request = prepare_chat_generation_request( + _base_request( + response_format={ + "type": "json_schema", + "json_schema": {"name": "answer", "schema": schema}, + } + ), + model_kit=_FakeTextModelKit(renderer), + supports_vision=False, + tokenize=lambda _model_kit, _prompt: [], + ) + + assert json.loads(request.generation_kwargs["json_schema"]) == schema + + +def test_unsupported_response_format_is_rejected(): + with pytest.raises(ValidationError, match="json_schema"): + prepare_chat_generation_request( + _base_request(response_format={"type": "json_object"}), + model_kit=_FakeTextModelKit(_FakeRenderer()), + supports_vision=False, + tokenize=lambda _model_kit, _prompt: [], + ) + + +def test_final_assistant_message_is_rendered_as_a_prefill(): + renderer = _FakeRenderer() + + prepare_chat_generation_request( + _base_request( + messages=[ + {"role": "user", "content": "Respond with JSON"}, + {"role": "assistant", "content": '{"answer":'}, + ] + ), + model_kit=_FakeTextModelKit(renderer), + supports_vision=False, + tokenize=lambda _model_kit, _prompt: [], + ) + + messages, template_kwargs = renderer.calls[0] + assert messages[-1] == {"role": "assistant", "content": '{"answer":'} + assert template_kwargs["add_generation_prompt"] is False + assert template_kwargs["continue_final_message"] is True @pytest.mark.parametrize( diff --git a/tests/server/test_http.py b/tests/server/test_http.py index f359ee8c..fc18beef 100644 --- a/tests/server/test_http.py +++ b/tests/server/test_http.py @@ -164,7 +164,18 @@ def create_generator(model_kit, prompt_tokens, **kwargs): port, "POST", "/v1/chat/completions", - body=_request_body(), + body={ + **_request_body(), + "response_format": { + "type": "json_schema", + "json_schema": { + "schema": { + "type": "object", + "properties": {"answer": {"type": "string"}}, + } + }, + }, + }, ) assert status == 200 @@ -208,6 +219,33 @@ def create_generator(model_kit, prompt_tokens, **kwargs): assert generation_kwargs["top_k"] == 20 assert generation_kwargs["min_p"] == 0.03 assert generation_kwargs["repetition_penalty"] == 1.05 + assert json.loads(generation_kwargs["json_schema"]) == { + "type": "object", + "properties": {"answer": {"type": "string"}}, + } + + +def test_tools_are_rejected_before_streaming(): + runtime = EngineRuntime( + _FakeModelKit(), + supports_vision=False, + get_runtime_load_info_fn=lambda _model_kit: {}, + ) + body = _request_body() + body["tools"] = [{"type": "function", "function": {"name": "search"}}] + + with _running_server(runtime) as port: + status, response_body = _request( + port, + "POST", + "/v1/chat/completions", + body=body, + ) + + assert status == 400 + assert json.loads(response_body) == { + "error": {"message": "Tools are not supported yet."} + } def test_generation_error_is_returned_inside_the_stream(): From 4bd8b0e8fd9eb15e5fffe35620d259b0394e05f2 Mon Sep 17 00:00:00 2001 From: Neil Mehta Date: Thu, 23 Jul 2026 17:15:15 -0400 Subject: [PATCH 08/18] Validate server request boundaries --- mlx_engine/server/chat.py | 10 +++++++++- mlx_engine/server/http.py | 27 ++++++++++++++++++++++--- tests/server/test_chat.py | 26 ++++++++++++++++++++++++ tests/server/test_http.py | 42 +++++++++++++++++++++++++++++++++++++++ 4 files changed, 101 insertions(+), 4 deletions(-) diff --git a/mlx_engine/server/chat.py b/mlx_engine/server/chat.py index 83582f0d..b651e983 100644 --- a/mlx_engine/server/chat.py +++ b/mlx_engine/server/chat.py @@ -129,6 +129,14 @@ def prepare_chat_generation_request( if request.tools: raise ChatRequestError("Tools are not supported yet.") + has_assistant_prefill = bool( + request.messages and request.messages[-1].role == "assistant" + ) + if request.response_format is not None and has_assistant_prefill: + raise ChatRequestError( + "Structured output is not supported with assistant prefills." + ) + normalized_messages, images_b64 = normalize_messages(request.messages) if images_b64 and not supports_vision: raise ChatRequestError("The loaded model does not support images.") @@ -143,7 +151,7 @@ def prepare_chat_generation_request( ) template_kwargs = dict(request.chat_template_kwargs) - if request.messages and request.messages[-1].role == "assistant": + if has_assistant_prefill: template_kwargs["continue_final_message"] = True add_generation_prompt = False else: diff --git a/mlx_engine/server/http.py b/mlx_engine/server/http.py index 2fb7844f..366bca08 100644 --- a/mlx_engine/server/http.py +++ b/mlx_engine/server/http.py @@ -35,6 +35,13 @@ logger = logging.getLogger(__name__) +_MAX_REQUEST_BODY_BYTES = 500 * 1024 * 1024 + + +class _RequestBodyTooLargeError(ValueError): + pass + + class EngineRuntime: def __init__( self, @@ -216,6 +223,9 @@ def do_POST(self) -> None: try: body = self._read_json_body() request = self.server.runtime.prepare_chat_generation(body) + except _RequestBodyTooLargeError as error: + self._send_error(HTTPStatus.REQUEST_ENTITY_TOO_LARGE, str(error)) + return except (ChatRequestError, ValidationError) as error: self._send_error(HTTPStatus.BAD_REQUEST, str(error)) return @@ -344,11 +354,22 @@ def _terminal_payload( } def _read_json_body(self) -> object: - content_length = self.headers.get("Content-Length") - if content_length is None: + content_length_header = self.headers.get("Content-Length") + if content_length_header is None: raise ChatRequestError("Content-Length is required.") try: - encoded_body = self.rfile.read(int(content_length)) + content_length = int(content_length_header) + except ValueError as error: + raise ChatRequestError( + "Content-Length must be a positive integer." + ) from error + if content_length <= 0: + raise ChatRequestError("Content-Length must be a positive integer.") + if content_length > _MAX_REQUEST_BODY_BYTES: + raise _RequestBodyTooLargeError("Request body exceeds the 500 MiB limit.") + + try: + encoded_body = self.rfile.read(content_length) return json.loads(encoded_body) except (UnicodeDecodeError, ValueError, json.JSONDecodeError) as error: raise ChatRequestError("The request body must be valid JSON.") from error diff --git a/tests/server/test_chat.py b/tests/server/test_chat.py index 5f53231b..b57d4cba 100644 --- a/tests/server/test_chat.py +++ b/tests/server/test_chat.py @@ -144,6 +144,32 @@ def test_unsupported_response_format_is_rejected(): ) +def test_structured_output_with_assistant_prefill_is_rejected(): + renderer = _FakeRenderer() + + with pytest.raises( + ChatRequestError, + match="Structured output is not supported with assistant prefills", + ): + prepare_chat_generation_request( + _base_request( + messages=[ + {"role": "user", "content": "Respond with JSON"}, + {"role": "assistant", "content": '{"answer":'}, + ], + response_format={ + "type": "json_schema", + "json_schema": {"schema": {"type": "object"}}, + }, + ), + model_kit=_FakeTextModelKit(renderer), + supports_vision=False, + tokenize=lambda _model_kit, _prompt: [], + ) + + assert renderer.calls == [] + + def test_final_assistant_message_is_rendered_as_a_prefill(): renderer = _FakeRenderer() diff --git a/tests/server/test_http.py b/tests/server/test_http.py index fc18beef..7eb53908 100644 --- a/tests/server/test_http.py +++ b/tests/server/test_http.py @@ -99,6 +99,23 @@ def _request(port, method, path, *, body=None, authorized=True): return response.status, response_body +def _request_with_content_length(port, content_length): + connection = http.client.HTTPConnection("127.0.0.1", port, timeout=5) + connection.request( + "POST", + "/v1/chat/completions", + headers={ + "Authorization": "Bearer secret-token", + "Content-Type": "application/json", + "Content-Length": content_length, + }, + ) + response = connection.getresponse() + response_body = response.read().decode("utf-8") + connection.close() + return response.status, response_body + + def test_health_requires_auth_and_reports_actualized_context_length(): runtime = EngineRuntime( _FakeModelKit(), @@ -116,6 +133,31 @@ def test_health_requires_auth_and_reports_actualized_context_length(): assert json.loads(body) == {"status": "ok", "context_length": 8192} +def test_invalid_and_oversized_content_lengths_are_rejected(): + runtime = EngineRuntime( + _FakeModelKit(), + supports_vision=False, + get_runtime_load_info_fn=lambda _model_kit: {}, + ) + + with _running_server(runtime) as port: + for content_length in ("invalid", "-1", "0"): + status, body = _request_with_content_length(port, content_length) + assert status == 400 + assert json.loads(body) == { + "error": {"message": "Content-Length must be a positive integer."} + } + + status, body = _request_with_content_length( + port, + str(500 * 1024 * 1024 + 1), + ) + assert status == 413 + assert json.loads(body) == { + "error": {"message": "Request body exceeds the 500 MiB limit."} + } + + def test_chat_stream_forwards_generation_settings_and_returns_usage(): generation_calls = [] From 6083c4fc69cf16340e139c6ac872f6f744dbe343 Mon Sep 17 00:00:00 2001 From: Neil Mehta Date: Mon, 27 Jul 2026 15:04:29 -0400 Subject: [PATCH 09/18] Validate generation requests and prompt usage --- .../model_kit/batched_vision/model_kit.py | 6 +- mlx_engine/server/chat.py | 38 ++++++-- mlx_engine/utils/prompt_progress_reporter.py | 16 ++-- tests/server/test_chat.py | 92 +++++++++++++++++++ tests/server/test_http.py | 74 +++++++++++++++ tests/test_batched_vision_model_kit.py | 68 ++++++++++++++ tests/test_vision_models.py | 2 +- tests/utils/test_prompt_progress_reporter.py | 25 +++++ 8 files changed, 303 insertions(+), 18 deletions(-) diff --git a/mlx_engine/model_kit/batched_vision/model_kit.py b/mlx_engine/model_kit/batched_vision/model_kit.py index 14bc53ed..aadebe55 100644 --- a/mlx_engine/model_kit/batched_vision/model_kit.py +++ b/mlx_engine/model_kit/batched_vision/model_kit.py @@ -512,12 +512,12 @@ def _insert_prepared_request( ) detokenizer = self._new_detokenizer() - total_prompt_tokens = max(0, prompt_token_count - 1) - cached_tokens = min(cached_prefix_len, total_prompt_tokens) + prefill_tokens = max(0, prompt_token_count - 1) + cached_tokens = min(cached_prefix_len, prefill_tokens) request.rqueue.put( PromptProgressBeginEvent( cached_tokens=cached_tokens, - total_prompt_tokens=total_prompt_tokens, + total_prompt_tokens=prompt_token_count, prefill_tokens_processed=0, ) ) diff --git a/mlx_engine/server/chat.py b/mlx_engine/server/chat.py index b651e983..3bd76911 100644 --- a/mlx_engine/server/chat.py +++ b/mlx_engine/server/chat.py @@ -36,6 +36,10 @@ class _InlineImageContentPart(BaseModel): _TextContentPart | _InlineImageContentPart, Field(discriminator="type"), ] +_NonNegativeFloat = Annotated[float, Field(ge=0, allow_inf_nan=False)] +_Probability = Annotated[float, Field(ge=0, le=1, allow_inf_nan=False)] +_PositiveInt = Annotated[int, Field(gt=0, strict=True)] +_TopK = Annotated[int, Field(ge=-1, le=500, strict=True)] class ChatMessage(BaseModel): @@ -59,13 +63,19 @@ class ChatCompletionRequest(BaseModel): messages: list[ChatMessage] stream: Literal[True] - temperature: float - max_tokens: int | None = None + temperature: _NonNegativeFloat + max_tokens: _PositiveInt | None = None stop: list[str] | None = None - top_p: float | None = None - top_k: int - min_p: float | None = None - repeat_penalty: float | None = None + top_p: _Probability | None = None + top_k: _TopK + min_p: _Probability | None = None + repeat_penalty: _NonNegativeFloat | None = None + seed: int | None = None + logprobs: bool | None = None + top_logprobs: int | None = None + logit_bias: dict | None = None + presence_penalty: float | None = None + frequency_penalty: float | None = None tools: list[dict] | None = None response_format: _JsonSchemaResponseFormat | None = None chat_template_kwargs: dict = Field(default_factory=dict) @@ -129,6 +139,22 @@ def prepare_chat_generation_request( if request.tools: raise ChatRequestError("Tools are not supported yet.") + unsupported_controls = [ + name + for name, requested in ( + ("seed", request.seed is not None), + ("logprobs", request.logprobs is True), + ("top_logprobs", request.top_logprobs is not None), + ("logit_bias", request.logit_bias is not None), + ("presence_penalty", request.presence_penalty is not None), + ("frequency_penalty", request.frequency_penalty is not None), + ) + if requested + ] + if unsupported_controls: + names = ", ".join(unsupported_controls) + raise ChatRequestError(f"Unsupported generation controls: {names}.") + has_assistant_prefill = bool( request.messages and request.messages[-1].role == "assistant" ) diff --git a/mlx_engine/utils/prompt_progress_reporter.py b/mlx_engine/utils/prompt_progress_reporter.py index ac2e5f72..26b51209 100644 --- a/mlx_engine/utils/prompt_progress_reporter.py +++ b/mlx_engine/utils/prompt_progress_reporter.py @@ -32,8 +32,10 @@ def begin( Args: is_draft: True if this is for the draft model, False for main model. cached_tokens: Number of tokens already in the KV cache. - total_prompt_tokens: Total number of tokens in the prompt. - prefill_tokens_processed: Number of tokens processed so far (usually 0 at begin). + total_prompt_tokens: Full prompt token count, including the final token used + to seed autoregressive decoding. + prefill_tokens_processed: Number of tokens prefilled so far. This may finish + one token short of total_prompt_tokens because the final token seeds decoding. Returns: True to continue processing, False to cancel. @@ -226,9 +228,8 @@ class BatchedMlxLmReporterAdapter: Adapts a PromptProgressReporter to the BatchedModelKit.generate callback. Converts (processed_tokens, total_tokens) -> None to reporter method calls. - Automatically calls finish() when processed_tokens - 1 >= total_tokens. - We need the off-by-one since mlx-lm prefills every token except for the last one, - since that token is needed to start the auto-regressive decoding + Reports the full prompt length while calling finish() after all but the final + decode seed token have been prefilled. Unlike MlxLmReporterAdapter, do not throw when we receive a stop request. Return False so batched schedulers can cooperatively cancel at chunk @@ -246,8 +247,7 @@ def __call__(self, processed_tokens: int, total_tokens: int) -> bool: if self._finished: return True - # mlx-lm tells us how many total prompt tokens there are. It leaves one unprocessed to seed the decode. Make that adjustment here - total_tokens = max(0, total_tokens - 1) + prefill_tokens = max(0, total_tokens - 1) if self._first_call: self._first_call = False @@ -261,7 +261,7 @@ def __call__(self, processed_tokens: int, total_tokens: int) -> bool: if not should_continue: return False - if processed_tokens >= total_tokens: + if processed_tokens >= prefill_tokens: self._finished = True return self._reporter.finish( is_draft=False, prefill_tokens_processed=processed_tokens diff --git a/tests/server/test_chat.py b/tests/server/test_chat.py index b57d4cba..e099475a 100644 --- a/tests/server/test_chat.py +++ b/tests/server/test_chat.py @@ -90,6 +90,98 @@ def test_prepare_text_request_uses_only_supported_generation_settings(): assert template_kwargs["reasoning_effort"] == "medium" +def test_supported_generation_boundaries_and_unknown_future_fields_are_accepted(): + renderer = _FakeRenderer() + + request = prepare_chat_generation_request( + _base_request( + temperature=0, + max_tokens=1, + top_p=1, + top_k=-1, + min_p=0, + repeat_penalty=0, + logprobs=False, + future_sampling_control={"enabled": True}, + ), + model_kit=_FakeTextModelKit(renderer), + supports_vision=False, + tokenize=lambda _model_kit, _prompt: [], + ) + + assert request.generation_kwargs == { + "images_b64": [], + "temp": 0, + "max_tokens": 1, + "stop_strings": ["END"], + "top_p": 1, + "top_k": -1, + "min_p": 0, + "repetition_penalty": 0, + } + + +@pytest.mark.parametrize( + ("overrides", "control_name"), + [ + ({"seed": 0}, "seed"), + ({"logprobs": True}, "logprobs"), + ({"top_logprobs": 5}, "top_logprobs"), + ({"logit_bias": {"1": 1}}, "logit_bias"), + ({"presence_penalty": 0}, "presence_penalty"), + ({"frequency_penalty": 0}, "frequency_penalty"), + ], +) +def test_unsupported_generation_controls_are_rejected(overrides, control_name): + renderer = _FakeRenderer() + + with pytest.raises(ChatRequestError, match=control_name): + prepare_chat_generation_request( + _base_request(**overrides), + model_kit=_FakeTextModelKit(renderer), + supports_vision=False, + tokenize=lambda _model_kit, _prompt: [], + ) + + assert renderer.calls == [] + + +@pytest.mark.parametrize( + "overrides", + [ + {"temperature": -0.1}, + {"temperature": float("nan")}, + {"temperature": float("inf")}, + {"max_tokens": 0}, + {"max_tokens": -1}, + {"max_tokens": 1.5}, + {"top_p": -0.1}, + {"top_p": 1.1}, + {"top_p": float("nan")}, + {"top_k": -2}, + {"top_k": 501}, + {"top_k": 1.5}, + {"min_p": -0.1}, + {"min_p": 1.1}, + {"min_p": float("nan")}, + {"repeat_penalty": -0.1}, + {"repeat_penalty": float("nan")}, + ], +) +def test_invalid_generation_settings_are_rejected_before_rendering(overrides): + renderer = _FakeRenderer() + + with pytest.raises(ValidationError): + prepare_chat_generation_request( + _base_request(**overrides), + model_kit=_FakeTextModelKit(renderer), + supports_vision=False, + tokenize=lambda _model_kit, _prompt: [], + ) + + assert renderer.calls == [] + + def test_tools_are_rejected(): renderer = _FakeRenderer() diff --git a/tests/server/test_http.py b/tests/server/test_http.py index 7eb53908..5e00a18e 100644 --- a/tests/server/test_http.py +++ b/tests/server/test_http.py @@ -15,6 +15,7 @@ GenerationResult, GenerationStopCondition, ) +from mlx_engine.utils.prompt_progress_reporter import BatchedMlxLmReporterAdapter from mlx_engine.utils.token import Token @@ -158,6 +159,39 @@ def test_invalid_and_oversized_content_lengths_are_rejected(): } +def test_invalid_generation_settings_are_rejected_before_streaming(): + runtime = EngineRuntime( + _FakeModelKit(), + supports_vision=False, + get_runtime_load_info_fn=lambda _model_kit: {}, + ) + + with _running_server(runtime) as port: + invalid_body = _request_body() + invalid_body["temperature"] = -0.1 + status, response_body = _request( + port, + "POST", + "/v1/chat/completions", + body=invalid_body, + ) + assert status == 400 + assert "temperature" in json.loads(response_body)["error"]["message"] + + unsupported_body = _request_body() + unsupported_body["seed"] = 0 + status, response_body = _request( + port, + "POST", + "/v1/chat/completions", + body=unsupported_body, + ) + assert status == 400 + assert json.loads(response_body) == { + "error": {"message": "Unsupported generation controls: seed."} + } + + def test_chat_stream_forwards_generation_settings_and_returns_usage(): generation_calls = [] @@ -267,6 +301,46 @@ def create_generator(model_kit, prompt_tokens, **kwargs): } +def test_batched_prompt_usage_includes_the_decode_seed_token(): + def create_generator(_model_kit, prompt_tokens, **kwargs): + reporter = BatchedMlxLmReporterAdapter( + kwargs["prompt_progress_reporter"], + emit_begin=True, + ) + assert reporter(0, len(prompt_tokens)) + assert reporter(len(prompt_tokens) - 1, len(prompt_tokens)) + yield GenerationResult( + text="", + tokens=[], + top_logprobs=[], + stop_condition=GenerationStopCondition( + stop_reason="eos_token", + stop_string="", + stop_tokens=[2], + ), + ) + + runtime = EngineRuntime( + _FakeModelKit(), + supports_vision=False, + create_generator_fn=create_generator, + get_runtime_load_info_fn=lambda _model_kit: {}, + tokenize_fn=lambda _model_kit, _prompt: [1, 2, 3], + ) + + with _running_server(runtime) as port: + status, response_text = _request( + port, + "POST", + "/v1/chat/completions", + body=_request_body(), + ) + + assert status == 200 + terminal_event = _parse_sse(response_text)[0] + assert terminal_event["usage"]["prompt_tokens"] == 3 + + def test_tools_are_rejected_before_streaming(): runtime = EngineRuntime( _FakeModelKit(), diff --git a/tests/test_batched_vision_model_kit.py b/tests/test_batched_vision_model_kit.py index 37c36a65..6ec69028 100644 --- a/tests/test_batched_vision_model_kit.py +++ b/tests/test_batched_vision_model_kit.py @@ -12,6 +12,12 @@ from mlx_engine.model_kit.batched_vision.prompt_cache.types import ( PromptImageSpan, ) +from mlx_engine.model_kit.batched_vision.prompt_inputs import PreparedPrompt +from mlx_engine.model_kit.batched_vision.request_lifecycle import ( + GenerationRequest, + PreparedInsert, +) +from mlx_engine.utils.prompt_progress_events import PromptProgressBeginEvent def test_global_no_chunked_prefill_exempts_gemma4_visual_policy(): @@ -108,6 +114,68 @@ def fake_load_model(model_path, **kwargs): assert call["patched_model"] is loaded_model +def test_insert_reports_full_prepared_prompt_length(monkeypatch): + class FakeBatchGenerator: + def insert(self, *_args, **_kwargs): + return 7 + + monkeypatch.setattr( + model_kit_module, + "build_prompt_kwargs", + lambda *_args, **_kwargs: {"inputs_embeds": "embeddings"}, + ) + monkeypatch.setattr( + model_kit_module, + "build_prefix_cache_chunks", + lambda *_args, **_kwargs: [], + ) + monkeypatch.setattr( + model_kit_module, + "first_unsaved_prefix_cache_chunk_index", + lambda *_args, **_kwargs: 0, + ) + + request = GenerationRequest( + rqueue=Queue(), + prompt_tokens=[1], + request_id="request", + images_b64=None, + sampler=lambda logits: logits, + logits_processors=[], + top_logprobs=0, + max_tokens=1, + ) + prepared = PreparedInsert( + request=request, + prepared_prompt=PreparedPrompt( + prompt_input_ids=[1, 2, 3], + raw_inputs=None, + image_spans=[], + ), + restored=None, + ) + kit = object.__new__(BatchedVisionModelKit) + kit._shutdown = SimpleNamespace(is_set=lambda: True) + kit.model = SimpleNamespace(no_chunked_prefill=False) + kit.model_type = "other_vlm" + kit._uses_gemma4_bidirectional_visual_attention = False + kit._vision_feature_memoizer = None + kit._prompt_cache_coordinator = SimpleNamespace( + save_prompt_cache_snapshot=lambda *_args, **_kwargs: None + ) + kit._prompt_cache_store = SimpleNamespace(can_store_records=lambda: False) + kit._new_detokenizer = lambda: object() + active = {} + + kit._insert_prepared_request(FakeBatchGenerator(), prepared, active) + + progress = request.rqueue.get_nowait() + assert isinstance(progress, PromptProgressBeginEvent) + assert progress.total_prompt_tokens == 3 + assert progress.cached_tokens == 0 + assert active[7].request_id == "request" + + def test_generate_fatal_error_preserves_state_for_exception_propagation(monkeypatch): class FakeBatchGenerator: def __init__(self): diff --git a/tests/test_vision_models.py b/tests/test_vision_models.py index cdf9525b..892e58ea 100644 --- a/tests/test_vision_models.py +++ b/tests/test_vision_models.py @@ -72,7 +72,7 @@ def _assert_cached_follow_up_prefill_is_small( ) -> None: assert finish_event["type"] == "finish" assert finish_event["prefill_tokens_processed"] == ( - begin_event["total_prompt_tokens"] - begin_event["cached_tokens"] + begin_event["total_prompt_tokens"] - begin_event["cached_tokens"] - 1 ) assert finish_event["prefill_tokens_processed"] <= CACHING_TEST_PREFILL_STEP_SIZE diff --git a/tests/utils/test_prompt_progress_reporter.py b/tests/utils/test_prompt_progress_reporter.py index 67989c66..c1011b3e 100644 --- a/tests/utils/test_prompt_progress_reporter.py +++ b/tests/utils/test_prompt_progress_reporter.py @@ -1,6 +1,7 @@ import unittest from typing import Optional from mlx_engine.utils.prompt_progress_reporter import ( + BatchedMlxLmReporterAdapter, ForwardingReporter, MlxLmReporterAdapter, ) @@ -205,3 +206,27 @@ def update(self, *args, **kwargs): adapter(0, 100) # Begin - succeeds with self.assertRaises(StopPromptProcessing): adapter(50, 100) # Update - cancels + + +class TestBatchedMlxLmReporterAdapter(unittest.TestCase): + def test_reports_full_prompt_length_and_seed_excluded_prefill(self): + inner = MockReporter(return_value=True) + adapter = BatchedMlxLmReporterAdapter(inner, emit_begin=True) + + adapter(0, 3) + adapter(2, 3) + + self.assertEqual(inner.events[0]["type"], "begin") + self.assertEqual(inner.events[0]["total_prompt_tokens"], 3) + self.assertEqual(inner.events[-1]["type"], "finish") + self.assertEqual(inner.events[-1]["prefill_tokens_processed"], 2) + + def test_one_token_prompt_reports_usage_without_prefill_work(self): + inner = MockReporter(return_value=True) + adapter = BatchedMlxLmReporterAdapter(inner, emit_begin=True) + + adapter(0, 1) + + self.assertEqual(inner.events[0]["total_prompt_tokens"], 1) + self.assertEqual(inner.events[-1]["type"], "finish") + self.assertEqual(inner.events[-1]["prefill_tokens_processed"], 0) From ef0524f1dcc7b31d00dca6cd95e476e9af32c5fb Mon Sep 17 00:00:00 2001 From: Neil Mehta Date: Mon, 27 Jul 2026 16:08:23 -0400 Subject: [PATCH 10/18] Harden server request and usage accounting --- mlx_engine/server/__main__.py | 18 +++- mlx_engine/server/chat.py | 3 +- mlx_engine/server/http.py | 5 +- mlx_engine/utils/prompt_progress_reporter.py | 13 +-- tests/server/test_chat.py | 1 + tests/server/test_http.py | 89 ++++++++++++++++++-- tests/server/test_main.py | 36 ++++++++ tests/utils/test_prompt_progress_reporter.py | 2 +- 8 files changed, 146 insertions(+), 21 deletions(-) create mode 100644 tests/server/test_main.py diff --git a/mlx_engine/server/__main__.py b/mlx_engine/server/__main__.py index d87f2088..85c92b52 100644 --- a/mlx_engine/server/__main__.py +++ b/mlx_engine/server/__main__.py @@ -1,5 +1,6 @@ import argparse import logging +import os import signal import threading @@ -11,12 +12,17 @@ logger = logging.getLogger(__name__) +_API_KEY_ENV_VAR = "MLX_ENGINE_API_KEY" + + def _create_parser() -> argparse.ArgumentParser: - parser = argparse.ArgumentParser(description="Run the private mlx-engine server.") + parser = argparse.ArgumentParser( + description="Run the private mlx-engine server.", + epilog=f"Authentication is configured through {_API_KEY_ENV_VAR}.", + ) parser.add_argument("--model", required=True) parser.add_argument("--host", required=True) parser.add_argument("--port", required=True, type=int) - parser.add_argument("--api-key", required=True) parser.add_argument("--context-length", required=True, type=int) parser.add_argument("--parallel-sessions", required=True, type=int) parser.add_argument("--seed", type=int) @@ -24,7 +30,11 @@ def _create_parser() -> argparse.ArgumentParser: def main() -> None: - args = _create_parser().parse_args() + parser = _create_parser() + args = parser.parse_args() + api_key = os.environ.get(_API_KEY_ENV_VAR) + if not api_key: + parser.error(f"{_API_KEY_ENV_VAR} environment variable is required") logger.info("Loading MLX model from %s", args.model) model_kit = load_model( @@ -40,7 +50,7 @@ def main() -> None: try: server = MlxEngineHttpServer( (args.host, args.port), - api_key=args.api_key, + api_key=api_key, runtime=runtime, ) diff --git a/mlx_engine/server/chat.py b/mlx_engine/server/chat.py index 3bd76911..0c2822ff 100644 --- a/mlx_engine/server/chat.py +++ b/mlx_engine/server/chat.py @@ -39,6 +39,7 @@ class _InlineImageContentPart(BaseModel): _NonNegativeFloat = Annotated[float, Field(ge=0, allow_inf_nan=False)] _Probability = Annotated[float, Field(ge=0, le=1, allow_inf_nan=False)] _PositiveInt = Annotated[int, Field(gt=0, strict=True)] +_NonEmptyString = Annotated[str, Field(min_length=1)] _TopK = Annotated[int, Field(ge=-1, le=500, strict=True)] @@ -65,7 +66,7 @@ class ChatCompletionRequest(BaseModel): stream: Literal[True] temperature: _NonNegativeFloat max_tokens: _PositiveInt | None = None - stop: list[str] | None = None + stop: list[_NonEmptyString] | None = None top_p: _Probability | None = None top_k: _TopK min_p: _Probability | None = None diff --git a/mlx_engine/server/http.py b/mlx_engine/server/http.py index 366bca08..81eaa420 100644 --- a/mlx_engine/server/http.py +++ b/mlx_engine/server/http.py @@ -129,9 +129,11 @@ def __init__( handler: MlxEngineRequestHandler, session: GenerationSession, prompt_tokens: int, + use_reported_prompt_tokens: bool, ): self._handler = handler self._session = session + self._use_reported_prompt_tokens = use_reported_prompt_tokens self.prompt_tokens = prompt_tokens def begin( @@ -141,7 +143,7 @@ def begin( total_prompt_tokens: int, prefill_tokens_processed: int, ) -> bool: - if not is_draft: + if not is_draft and self._use_reported_prompt_tokens: self.prompt_tokens = total_prompt_tokens return self._report(is_draft) @@ -249,6 +251,7 @@ def do_POST(self) -> None: self, session, len(request.prompt_tokens), + use_reported_prompt_tokens=self.server.runtime.supports_vision, ) generator = self.server.runtime.create_chat_generator( request, diff --git a/mlx_engine/utils/prompt_progress_reporter.py b/mlx_engine/utils/prompt_progress_reporter.py index 26b51209..b85c0da9 100644 --- a/mlx_engine/utils/prompt_progress_reporter.py +++ b/mlx_engine/utils/prompt_progress_reporter.py @@ -228,8 +228,8 @@ class BatchedMlxLmReporterAdapter: Adapts a PromptProgressReporter to the BatchedModelKit.generate callback. Converts (processed_tokens, total_tokens) -> None to reporter method calls. - Reports the full prompt length while calling finish() after all but the final - decode seed token have been prefilled. + Includes the decode seed in the token total but excludes it from prefill + progress. Unlike MlxLmReporterAdapter, do not throw when we receive a stop request. Return False so batched schedulers can cooperatively cancel at chunk @@ -248,6 +248,7 @@ def __call__(self, processed_tokens: int, total_tokens: int) -> bool: return True prefill_tokens = max(0, total_tokens - 1) + prefill_tokens_processed = min(processed_tokens, prefill_tokens) if self._first_call: self._first_call = False @@ -256,7 +257,7 @@ def __call__(self, processed_tokens: int, total_tokens: int) -> bool: is_draft=False, cached_tokens=0, total_prompt_tokens=total_tokens, - prefill_tokens_processed=processed_tokens, + prefill_tokens_processed=prefill_tokens_processed, ) if not should_continue: return False @@ -264,9 +265,11 @@ def __call__(self, processed_tokens: int, total_tokens: int) -> bool: if processed_tokens >= prefill_tokens: self._finished = True return self._reporter.finish( - is_draft=False, prefill_tokens_processed=processed_tokens + is_draft=False, + prefill_tokens_processed=prefill_tokens_processed, ) return self._reporter.update( - is_draft=False, prefill_tokens_processed=processed_tokens + is_draft=False, + prefill_tokens_processed=prefill_tokens_processed, ) diff --git a/tests/server/test_chat.py b/tests/server/test_chat.py index e099475a..9053447b 100644 --- a/tests/server/test_chat.py +++ b/tests/server/test_chat.py @@ -155,6 +155,7 @@ def test_unsupported_generation_controls_are_rejected(overrides, control_name): {"max_tokens": 0}, {"max_tokens": -1}, {"max_tokens": 1.5}, + {"stop": [""]}, {"top_p": -0.1}, {"top_p": 1.1}, {"top_p": float("nan")}, diff --git a/tests/server/test_http.py b/tests/server/test_http.py index 5e00a18e..71874dd5 100644 --- a/tests/server/test_http.py +++ b/tests/server/test_http.py @@ -20,6 +20,8 @@ class _FakeRenderer: + chat_template = "model template" + def apply_chat_template(self, messages, **kwargs): assert messages == [{"role": "user", "content": "Hello"}] assert kwargs["tokenize"] is False @@ -37,6 +39,11 @@ def __init__(self): self.tokenizer = _FakeTokenizer() +class _FakeVisionModelKit: + def __init__(self): + self.processor = _FakeRenderer() + + def _request_body(): return { "model": "single-loaded-model", @@ -178,6 +185,17 @@ def test_invalid_generation_settings_are_rejected_before_streaming(): assert status == 400 assert "temperature" in json.loads(response_body)["error"]["message"] + empty_stop_body = _request_body() + empty_stop_body["stop"] = [""] + status, response_body = _request( + port, + "POST", + "/v1/chat/completions", + body=empty_stop_body, + ) + assert status == 400 + assert "stop" in json.loads(response_body)["error"]["message"] + unsupported_body = _request_body() unsupported_body["seed"] = 0 status, response_body = _request( @@ -200,11 +218,11 @@ def create_generator(model_kit, prompt_tokens, **kwargs): reporter = kwargs["prompt_progress_reporter"] assert reporter.begin( is_draft=False, - cached_tokens=3, - total_prompt_tokens=9, + cached_tokens=0, + total_prompt_tokens=3, prefill_tokens_processed=0, ) - assert reporter.update(is_draft=False, prefill_tokens_processed=6) + assert reporter.update(is_draft=False, prefill_tokens_processed=2) yield GenerationResult( text="Hello back", tokens=[ @@ -276,9 +294,9 @@ def create_generator(model_kit, prompt_tokens, **kwargs): } ], "usage": { - "prompt_tokens": 9, + "prompt_tokens": 3, "completion_tokens": 2, - "total_tokens": 11, + "total_tokens": 5, }, "__lmstudio": {"stop_type": "eos"}, } @@ -301,14 +319,21 @@ def create_generator(model_kit, prompt_tokens, **kwargs): } -def test_batched_prompt_usage_includes_the_decode_seed_token(): +def test_batched_text_cache_hit_preserves_full_prompt_usage(): + request_count = 0 + def create_generator(_model_kit, prompt_tokens, **kwargs): + nonlocal request_count reporter = BatchedMlxLmReporterAdapter( kwargs["prompt_progress_reporter"], emit_begin=True, ) - assert reporter(0, len(prompt_tokens)) - assert reporter(len(prompt_tokens) - 1, len(prompt_tokens)) + if request_count == 0: + assert reporter(0, len(prompt_tokens)) + assert reporter(len(prompt_tokens) - 1, len(prompt_tokens)) + else: + assert reporter(1, 1) + request_count += 1 yield GenerationResult( text="", tokens=[], @@ -328,6 +353,52 @@ def create_generator(_model_kit, prompt_tokens, **kwargs): tokenize_fn=lambda _model_kit, _prompt: [1, 2, 3], ) + with _running_server(runtime) as port: + responses = [ + _request( + port, + "POST", + "/v1/chat/completions", + body=_request_body(), + ) + for _ in range(2) + ] + + assert request_count == 2 + for status, response_text in responses: + assert status == 200 + terminal_event = _parse_sse(response_text)[0] + assert terminal_event["usage"]["prompt_tokens"] == 3 + + +def test_vision_usage_uses_the_prepared_prompt_length(): + def create_generator(_model_kit, _prompt_tokens, **kwargs): + reporter = kwargs["prompt_progress_reporter"] + assert reporter.begin( + is_draft=False, + cached_tokens=0, + total_prompt_tokens=9, + prefill_tokens_processed=0, + ) + yield GenerationResult( + text="", + tokens=[], + top_logprobs=[], + stop_condition=GenerationStopCondition( + stop_reason="eos_token", + stop_string="", + stop_tokens=[2], + ), + ) + + runtime = EngineRuntime( + _FakeVisionModelKit(), + supports_vision=True, + create_generator_fn=create_generator, + get_runtime_load_info_fn=lambda _model_kit: {}, + tokenize_fn=lambda _model_kit, _prompt: [1, 2, 3], + ) + with _running_server(runtime) as port: status, response_text = _request( port, @@ -338,7 +409,7 @@ def create_generator(_model_kit, prompt_tokens, **kwargs): assert status == 200 terminal_event = _parse_sse(response_text)[0] - assert terminal_event["usage"]["prompt_tokens"] == 3 + assert terminal_event["usage"]["prompt_tokens"] == 9 def test_tools_are_rejected_before_streaming(): diff --git a/tests/server/test_main.py b/tests/server/test_main.py new file mode 100644 index 00000000..e7ae08dc --- /dev/null +++ b/tests/server/test_main.py @@ -0,0 +1,36 @@ +import sys + +import pytest + +from mlx_engine.server import __main__ as server_main + + +_REQUIRED_ARGS = [ + "--model", + "model-path", + "--host", + "127.0.0.1", + "--port", + "1234", + "--context-length", + "4096", + "--parallel-sessions", + "2", +] + + +def test_api_key_is_not_a_command_line_argument(): + args = server_main._create_parser().parse_args(_REQUIRED_ARGS) + + assert not hasattr(args, "api_key") + + +def test_api_key_environment_variable_is_required(monkeypatch, capsys): + monkeypatch.delenv("MLX_ENGINE_API_KEY", raising=False) + monkeypatch.setattr(sys, "argv", ["mlx-engine-server", *_REQUIRED_ARGS]) + + with pytest.raises(SystemExit) as error: + server_main.main() + + assert error.value.code == 2 + assert "MLX_ENGINE_API_KEY environment variable is required" in capsys.readouterr().err diff --git a/tests/utils/test_prompt_progress_reporter.py b/tests/utils/test_prompt_progress_reporter.py index c1011b3e..375ef749 100644 --- a/tests/utils/test_prompt_progress_reporter.py +++ b/tests/utils/test_prompt_progress_reporter.py @@ -225,7 +225,7 @@ def test_one_token_prompt_reports_usage_without_prefill_work(self): inner = MockReporter(return_value=True) adapter = BatchedMlxLmReporterAdapter(inner, emit_begin=True) - adapter(0, 1) + adapter(1, 1) self.assertEqual(inner.events[0]["total_prompt_tokens"], 1) self.assertEqual(inner.events[-1]["type"], "finish") From a17e12f3041d8b0369c5d0d393a9619bf2cbe360 Mon Sep 17 00:00:00 2001 From: Neil Mehta Date: Mon, 27 Jul 2026 16:11:26 -0400 Subject: [PATCH 11/18] Format server startup test --- tests/server/test_main.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/server/test_main.py b/tests/server/test_main.py index e7ae08dc..796dcf85 100644 --- a/tests/server/test_main.py +++ b/tests/server/test_main.py @@ -33,4 +33,6 @@ def test_api_key_environment_variable_is_required(monkeypatch, capsys): server_main.main() assert error.value.code == 2 - assert "MLX_ENGINE_API_KEY environment variable is required" in capsys.readouterr().err + assert ( + "MLX_ENGINE_API_KEY environment variable is required" in capsys.readouterr().err + ) From 191760450249da568fd65c52af99cf6fa697dcd6 Mon Sep 17 00:00:00 2001 From: Neil Mehta Date: Tue, 28 Jul 2026 11:23:58 -0400 Subject: [PATCH 12/18] Normalize text messages and close error responses --- mlx_engine/server/chat.py | 17 ++++++++++--- mlx_engine/server/http.py | 17 +++++++++++-- tests/server/test_chat.py | 53 +++++++++++++++++++++++++++++++++++++-- tests/server/test_http.py | 29 +++++++++++++++++++++ 4 files changed, 109 insertions(+), 7 deletions(-) diff --git a/mlx_engine/server/chat.py b/mlx_engine/server/chat.py index 0c2822ff..5e33fe7b 100644 --- a/mlx_engine/server/chat.py +++ b/mlx_engine/server/chat.py @@ -99,7 +99,11 @@ def _base64_image_data(url: str) -> str: return data -def normalize_messages(messages: list[ChatMessage]) -> tuple[list[dict], list[str]]: +def normalize_messages( + messages: list[ChatMessage], + *, + supports_vision: bool, +) -> tuple[list[dict], list[str]]: normalized_messages: list[dict] = [] images_b64: list[str] = [] @@ -107,13 +111,17 @@ def normalize_messages(messages: list[ChatMessage]) -> tuple[list[dict], list[st normalized_message = message.model_dump(exclude_unset=True) if isinstance(message.content, list): normalized_parts: list[dict] = [] + text_parts: list[str] = [] for part in message.content: if isinstance(part, _TextContentPart): + text_parts.append(part.text) normalized_parts.append({"type": "text", "text": part.text}) else: images_b64.append(_base64_image_data(part.image_url.url)) normalized_parts.append({"type": "image"}) - normalized_message["content"] = normalized_parts + normalized_message["content"] = ( + normalized_parts if supports_vision else "".join(text_parts) + ) normalized_messages.append(normalized_message) return normalized_messages, images_b64 @@ -164,7 +172,10 @@ def prepare_chat_generation_request( "Structured output is not supported with assistant prefills." ) - normalized_messages, images_b64 = normalize_messages(request.messages) + normalized_messages, images_b64 = normalize_messages( + request.messages, + supports_vision=supports_vision, + ) if images_b64 and not supports_vision: raise ChatRequestError("The loaded model does not support images.") diff --git a/mlx_engine/server/http.py b/mlx_engine/server/http.py index 81eaa420..c67cbef3 100644 --- a/mlx_engine/server/http.py +++ b/mlx_engine/server/http.py @@ -393,13 +393,26 @@ def _write_bytes(self, content: bytes) -> None: self.wfile.flush() def _send_error(self, status: HTTPStatus, message: str) -> None: - self._send_json(status, {"error": {"message": message}}) + self._send_json( + status, + {"error": {"message": message}}, + close_connection=True, + ) - def _send_json(self, status: HTTPStatus, body: dict) -> None: + def _send_json( + self, + status: HTTPStatus, + body: dict, + *, + close_connection: bool = False, + ) -> None: encoded_body = json.dumps(body, separators=(",", ":")).encode("utf-8") self.send_response(status) self.send_header("Content-Type", "application/json") self.send_header("Content-Length", str(len(encoded_body))) + if close_connection: + self.send_header("Connection", "close") + self.close_connection = True self.end_headers() self.wfile.write(encoded_body) diff --git a/tests/server/test_chat.py b/tests/server/test_chat.py index 9053447b..23fb2424 100644 --- a/tests/server/test_chat.py +++ b/tests/server/test_chat.py @@ -2,6 +2,7 @@ import pytest from pydantic import ValidationError +from transformers.utils.chat_template_utils import render_jinja_template from mlx_engine.server.chat import ( ChatMessage, @@ -21,6 +22,23 @@ def apply_chat_template(self, messages, **kwargs): return "rendered prompt" +class _TransformersTextRenderer: + chat_template = ( + "{% for message in messages %}" + "{{ message['content'] | trim }}" + "{% endfor %}" + "{% if add_generation_prompt %}assistant{% endif %}" + ) + + def apply_chat_template(self, messages, **kwargs): + rendered, _generation_indices = render_jinja_template( + conversations=[messages], + chat_template=self.chat_template, + **kwargs, + ) + return rendered[0] + + class _FakeTokenizerWrapper: def __init__(self, renderer): self._tokenizer = renderer @@ -90,6 +108,35 @@ def test_prepare_text_request_uses_only_supported_generation_settings(): assert template_kwargs["reasoning_effort"] == "medium" +@pytest.mark.parametrize( + ("content", "expected_prompt"), + [ + ( + [ + {"type": "text", "text": "First"}, + {"type": "text", "text": "Second"}, + ], + "FirstSecondassistant", + ), + ([], "assistant"), + ], +) +def test_text_content_parts_are_strings_before_transformers_template( + content, + expected_prompt, +): + rendered_prompts = [] + + prepare_chat_generation_request( + _base_request(messages=[{"role": "user", "content": content}]), + model_kit=_FakeTextModelKit(_TransformersTextRenderer()), + supports_vision=False, + tokenize=lambda _model_kit, prompt: rendered_prompts.append(prompt) or [], + ) + + assert rendered_prompts == [expected_prompt] + + def test_supported_generation_boundaries_and_unknown_future_fields_are_accepted(): renderer = _FakeRenderer() @@ -360,7 +407,8 @@ def test_normalize_images_preserves_user_and_tool_result_order(): ] normalized, images_b64 = normalize_messages( - [ChatMessage.model_validate(message) for message in messages] + [ChatMessage.model_validate(message) for message in messages], + supports_vision=True, ) assert images_b64 == ["first-image", "second-image", "first-image"] @@ -449,7 +497,8 @@ def test_non_base64_image_url_is_rejected(): ], } ) - ] + ], + supports_vision=True, ) diff --git a/tests/server/test_http.py b/tests/server/test_http.py index 71874dd5..5e942363 100644 --- a/tests/server/test_http.py +++ b/tests/server/test_http.py @@ -141,6 +141,35 @@ def test_health_requires_auth_and_reports_actualized_context_length(): assert json.loads(body) == {"status": "ok", "context_length": 8192} +def test_rejected_post_closes_connection_before_unread_body_can_be_reused(): + runtime = EngineRuntime( + _FakeModelKit(), + supports_vision=False, + get_runtime_load_info_fn=lambda _model_kit: {}, + ) + body = b'{"messages":[]}' + + with _running_server(runtime) as port: + with socket.create_connection(("127.0.0.1", port), timeout=2) as client: + client.sendall( + b"POST /v1/chat/completions HTTP/1.1\r\n" + b"Host: 127.0.0.1\r\n" + b"Content-Type: application/json\r\n" + + f"Content-Length: {len(body)}\r\n\r\n".encode("ascii") + + body + + b"GET /health HTTP/1.1\r\n" + b"Host: 127.0.0.1\r\n" + b"Authorization: Bearer secret-token\r\n\r\n" + ) + response = b"" + while chunk := client.recv(4096): + response += chunk + + assert response.count(b"HTTP/1.1") == 1 + assert b"HTTP/1.1 401 Unauthorized" in response + assert b"\r\nConnection: close\r\n" in response + + def test_invalid_and_oversized_content_lengths_are_rejected(): runtime = EngineRuntime( _FakeModelKit(), From 104fd6c0f2307a12911025474b43797651756cb0 Mon Sep 17 00:00:00 2001 From: Neil Mehta Date: Wed, 29 Jul 2026 10:51:16 -0400 Subject: [PATCH 13/18] Bound stream writes and correct prefill progress --- mlx_engine/server/http.py | 35 ++++++--- mlx_engine/utils/prompt_progress_events.py | 20 +++-- tests/server/test_http.py | 86 ++++++++++++++++++---- tests/utils/test_prompt_progress_events.py | 58 ++++++++++----- 4 files changed, 152 insertions(+), 47 deletions(-) diff --git a/mlx_engine/server/http.py b/mlx_engine/server/http.py index c67cbef3..78269c27 100644 --- a/mlx_engine/server/http.py +++ b/mlx_engine/server/http.py @@ -36,6 +36,11 @@ _MAX_REQUEST_BODY_BYTES = 500 * 1024 * 1024 +_SSE_WRITE_TIMEOUT_SECONDS = 30.0 + + +class _ClientConnectionError(Exception): + pass class _RequestBodyTooLargeError(ValueError): @@ -241,12 +246,7 @@ def do_POST(self) -> None: generator = None normal_completion = False try: - self.send_response(HTTPStatus.OK) - self.send_header("Content-Type", "text/event-stream") - self.send_header("Cache-Control", "no-cache") - self.send_header("Connection", "close") - self.end_headers() - self.close_connection = True + self._start_sse_response() reporter = _SsePromptProgressReporter( self, session, @@ -260,13 +260,13 @@ def do_POST(self) -> None: ) self._stream_generation(generator, reporter) normal_completion = True - except (BrokenPipeError, ConnectionResetError, OSError): + except _ClientConnectionError: logger.debug("Generation client disconnected: %s", session.request_id) except Exception as error: logger.exception("MLX generation failed") try: self._write_sse_json({"error": {"message": str(error)}}) - except (BrokenPipeError, ConnectionResetError, OSError): + except _ClientConnectionError: pass finally: try: @@ -384,13 +384,28 @@ def _is_authorized(self) -> bool: f"Bearer {self.server.api_key}", ) + def _start_sse_response(self) -> None: + try: + self.connection.settimeout(_SSE_WRITE_TIMEOUT_SECONDS) + self.close_connection = True + self.send_response(HTTPStatus.OK) + self.send_header("Content-Type", "text/event-stream") + self.send_header("Cache-Control", "no-cache") + self.send_header("Connection", "close") + self.end_headers() + except OSError as error: + raise _ClientConnectionError from error + def _write_sse_json(self, body: dict) -> None: encoded_body = json.dumps(body, separators=(",", ":")).encode("utf-8") self._write_bytes(b"data: " + encoded_body + b"\n\n") def _write_bytes(self, content: bytes) -> None: - self.wfile.write(content) - self.wfile.flush() + try: + self.wfile.write(content) + self.wfile.flush() + except OSError as error: + raise _ClientConnectionError from error def _send_error(self, status: HTTPStatus, message: str) -> None: self._send_json( diff --git a/mlx_engine/utils/prompt_progress_events.py b/mlx_engine/utils/prompt_progress_events.py index 6239e798..1c3ef88b 100644 --- a/mlx_engine/utils/prompt_progress_events.py +++ b/mlx_engine/utils/prompt_progress_events.py @@ -52,13 +52,21 @@ def __init__( self._percent_callback = percent_callback self._context: Optional[ProgressContext] = None - def _emit_percent(self, prefill_tokens_processed: int) -> None: + def _emit_percent( + self, + prefill_tokens_processed: int, + *, + is_final: bool = False, + ) -> None: if self._percent_callback is None or self._context is None: return - tokens_to_prefill = ( - self._context.total_prompt_tokens - self._context.cached_tokens - ) - if tokens_to_prefill <= 0: + if is_final: + self._percent_callback(100.0) + return + + prefill_tokens = max(0, self._context.total_prompt_tokens - 1) + tokens_to_prefill = max(0, prefill_tokens - self._context.cached_tokens) + if tokens_to_prefill == 0: self._percent_callback(100.0) else: percent = (prefill_tokens_processed / tokens_to_prefill) * 100.0 @@ -104,5 +112,5 @@ def finish( ) should_continue = self._progress_callback(event, is_draft) if not is_draft and prefill_tokens_processed is not None: - self._emit_percent(prefill_tokens_processed) + self._emit_percent(prefill_tokens_processed, is_final=True) return should_continue diff --git a/tests/server/test_http.py b/tests/server/test_http.py index 5e942363..a9c1b6ef 100644 --- a/tests/server/test_http.py +++ b/tests/server/test_http.py @@ -6,6 +6,7 @@ import threading import time +import mlx_engine.server.http as server_http from mlx_engine.server.http import ( EngineRuntime, GenerationSession, @@ -73,12 +74,14 @@ def _parse_sse(response_text): @contextmanager -def _running_server(runtime): +def _running_server(runtime, *, send_buffer_size=None): server = MlxEngineHttpServer( ("127.0.0.1", 0), api_key="secret-token", runtime=runtime, ) + if send_buffer_size is not None: + server.socket.setsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF, send_buffer_size) thread = threading.Thread(target=server.serve_forever, daemon=True) thread.start() try: @@ -464,29 +467,84 @@ def test_tools_are_rejected_before_streaming(): } -def test_generation_error_is_returned_inside_the_stream(): +def test_generation_errors_are_returned_inside_the_stream(): + for generation_error in ( + RuntimeError("generation failed"), + OSError("backend I/O failed"), + ): + + def create_generator(_model_kit, _prompt_tokens, **_kwargs): + raise generation_error + yield + + runtime = EngineRuntime( + _FakeModelKit(), + supports_vision=False, + create_generator_fn=create_generator, + get_runtime_load_info_fn=lambda _model_kit: {}, + tokenize_fn=lambda _model_kit, _prompt: [1], + ) + + with _running_server(runtime) as port: + status, response_text = _request( + port, + "POST", + "/v1/chat/completions", + body=_request_body(), + ) + + assert status == 200 + assert _parse_sse(response_text) == [ + {"error": {"message": str(generation_error)}} + ] + + +def test_stalled_sse_write_cancels_the_active_mlx_request(monkeypatch): + monkeypatch.setattr(server_http, "_SSE_WRITE_TIMEOUT_SECONDS", 0.05) + generation_stopped = threading.Event() + stopped_request_ids = [] + large_text = "x" * (1024 * 1024) + def create_generator(_model_kit, _prompt_tokens, **_kwargs): - raise RuntimeError("generation failed") - yield + while not generation_stopped.is_set(): + yield GenerationResult( + text=large_text, + tokens=[], + top_logprobs=[], + stop_condition=None, + ) + + def stop_generation(_model_kit, request_id): + stopped_request_ids.append(request_id) + generation_stopped.set() runtime = EngineRuntime( _FakeModelKit(), supports_vision=False, create_generator_fn=create_generator, get_runtime_load_info_fn=lambda _model_kit: {}, - tokenize_fn=lambda _model_kit, _prompt: [1], + stop_generation_fn=stop_generation, + tokenize_fn=lambda _model_kit, _prompt: [1, 2, 3], ) - with _running_server(runtime) as port: - status, response_text = _request( - port, - "POST", - "/v1/chat/completions", - body=_request_body(), - ) + with _running_server(runtime, send_buffer_size=4096) as port: + encoded_body = json.dumps(_request_body()).encode("utf-8") + with socket.socket() as client: + client.setsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF, 4096) + client.settimeout(2) + client.connect(("127.0.0.1", port)) + client.sendall( + b"POST /v1/chat/completions HTTP/1.1\r\n" + b"Host: 127.0.0.1\r\n" + b"Authorization: Bearer secret-token\r\n" + b"Content-Type: application/json\r\n" + + f"Content-Length: {len(encoded_body)}\r\n\r\n".encode("ascii") + + encoded_body + ) + assert generation_stopped.wait(timeout=2) - assert status == 200 - assert _parse_sse(response_text) == [{"error": {"message": "generation failed"}}] + assert len(stopped_request_ids) == 1 + assert stopped_request_ids[0] != "" def test_client_disconnect_stops_the_active_mlx_request(): diff --git a/tests/utils/test_prompt_progress_events.py b/tests/utils/test_prompt_progress_events.py index 38cb380e..f228350e 100644 --- a/tests/utils/test_prompt_progress_events.py +++ b/tests/utils/test_prompt_progress_events.py @@ -4,6 +4,7 @@ PromptProgressBeginEvent, PromptProgressEvent, ) +from mlx_engine.utils.prompt_progress_reporter import BatchedMlxLmReporterAdapter class TestPromptProgressCallbackReporter(unittest.TestCase): @@ -85,12 +86,12 @@ def test_percent_callback_emitted_on_begin(self): reporter.begin( is_draft=False, cached_tokens=50, - total_prompt_tokens=100, + total_prompt_tokens=101, prefill_tokens_processed=10, ) self.assertEqual(len(self.percents), 1) - # percent = prefill / (total - cached) * 100 = 10 / 50 * 100 = 20 + # percent = prefill / (total - seed - cached) * 100 = 10 / 50 * 100 self.assertEqual(self.percents[0], 20.0) def test_percent_callback_emitted_on_update(self): @@ -100,7 +101,7 @@ def test_percent_callback_emitted_on_update(self): reporter.begin( is_draft=False, cached_tokens=20, - total_prompt_tokens=100, + total_prompt_tokens=101, prefill_tokens_processed=0, ) self.percents.clear() @@ -108,7 +109,7 @@ def test_percent_callback_emitted_on_update(self): reporter.update(is_draft=False, prefill_tokens_processed=40) self.assertEqual(len(self.percents), 1) - # percent = prefill / (total - cached) * 100 = 40 / 80 * 100 = 50 + # percent = prefill / (total - seed - cached) * 100 = 40 / 80 * 100 self.assertEqual(self.percents[0], 50.0) def test_percent_callback_emitted_on_finish(self): @@ -118,7 +119,7 @@ def test_percent_callback_emitted_on_finish(self): reporter.begin( is_draft=False, cached_tokens=20, - total_prompt_tokens=100, + total_prompt_tokens=101, prefill_tokens_processed=0, ) self.percents.clear() @@ -126,17 +127,16 @@ def test_percent_callback_emitted_on_finish(self): reporter.finish(is_draft=False, prefill_tokens_processed=80) self.assertEqual(len(self.percents), 1) - # percent = prefill / (total - cached) * 100 = 80 / 80 * 100 = 100 self.assertEqual(self.percents[0], 100.0) - def test_percent_emitted_on_finish_using_last_value(self): + def test_finish_uses_last_processed_value_and_emits_100_percent(self): reporter = PromptProgressCallbackReporter( self.progress_callback, percent_callback=self.percent_callback ) reporter.begin( is_draft=False, cached_tokens=20, - total_prompt_tokens=100, + total_prompt_tokens=101, prefill_tokens_processed=0, ) reporter.update(is_draft=False, prefill_tokens_processed=60) @@ -144,9 +144,33 @@ def test_percent_emitted_on_finish_using_last_value(self): reporter.finish(is_draft=False, prefill_tokens_processed=None) - self.assertEqual(len(self.percents), 1) - # percent = prefill / (total - cached) * 100 = 60 / 80 * 100 = 75 - self.assertEqual(self.percents[0], 75.0) + self.assertEqual(self.events[-1]["event"].prefill_tokens_processed, 60) + self.assertEqual(self.percents, [100.0]) + + def test_batched_adapter_reaches_100_without_processing_decode_seed(self): + reporter = PromptProgressCallbackReporter( + self.progress_callback, percent_callback=self.percent_callback + ) + adapter = BatchedMlxLmReporterAdapter(reporter, emit_begin=True) + + adapter(0, 3) + adapter(2, 3) + + self.assertEqual(self.percents[-1], 100.0) + final_event = self.events[-1]["event"] + self.assertTrue(final_event.is_final) + self.assertEqual(final_event.prefill_tokens_processed, 2) + + def test_one_token_batched_prompt_has_no_prefill_work(self): + reporter = PromptProgressCallbackReporter( + self.progress_callback, percent_callback=self.percent_callback + ) + adapter = BatchedMlxLmReporterAdapter(reporter, emit_begin=True) + + adapter(1, 1) + + self.assertEqual(self.percents, [100.0, 100.0]) + self.assertEqual(self.events[-1]["event"].prefill_tokens_processed, 0) def test_draft_events_do_not_emit_percent(self): reporter = PromptProgressCallbackReporter( @@ -205,7 +229,7 @@ def test_percent_clamped_to_100(self): reporter.begin( is_draft=False, cached_tokens=80, - total_prompt_tokens=100, + total_prompt_tokens=101, prefill_tokens_processed=0, ) self.percents.clear() @@ -222,7 +246,7 @@ def test_percent_clamped_to_0(self): reporter.begin( is_draft=False, cached_tokens=-10, # edge case - total_prompt_tokens=100, + total_prompt_tokens=101, prefill_tokens_processed=0, ) self.percents.clear() @@ -239,7 +263,7 @@ def test_full_sequence(self): reporter.begin( is_draft=False, cached_tokens=20, - total_prompt_tokens=100, + total_prompt_tokens=101, prefill_tokens_processed=0, ) reporter.update(is_draft=False, prefill_tokens_processed=40) @@ -252,7 +276,7 @@ def test_full_sequence(self): begin_event = self.events[0]["event"] self.assertIsInstance(begin_event, PromptProgressBeginEvent) self.assertEqual(begin_event.cached_tokens, 20) - self.assertEqual(begin_event.total_prompt_tokens, 100) + self.assertEqual(begin_event.total_prompt_tokens, 101) self.assertEqual(begin_event.prefill_tokens_processed, 0) self.assertFalse(self.events[0]["is_draft"]) @@ -325,11 +349,11 @@ def test_all_cached_returns_100_percent(self): ) reporter.begin( is_draft=False, - cached_tokens=100, + cached_tokens=99, total_prompt_tokens=100, prefill_tokens_processed=0, ) - # tokens_to_prefill = 0, so returns 100% (everything cached) + # The 99 prefill tokens are cached; the remaining token seeds decode. self.assertEqual(len(self.percents), 1) self.assertEqual(self.percents[0], 100.0) From 2f4ef0bff70f92471e11f2b567337d1db03dd06b Mon Sep 17 00:00:00 2001 From: Neil Mehta Date: Wed, 29 Jul 2026 14:53:52 -0400 Subject: [PATCH 14/18] Harden chat request validation --- mlx_engine/server/chat.py | 20 +++++++++ mlx_engine/server/http.py | 12 +++-- tests/server/test_chat.py | 63 ++++++++++++++++++++++++--- tests/server/test_http.py | 92 ++++++++++++++++++++++++++++++++++++++- 4 files changed, 174 insertions(+), 13 deletions(-) diff --git a/mlx_engine/server/chat.py b/mlx_engine/server/chat.py index 5e33fe7b..1fb3f246 100644 --- a/mlx_engine/server/chat.py +++ b/mlx_engine/server/chat.py @@ -1,3 +1,5 @@ +import base64 +import binascii from dataclasses import dataclass import json from typing import Annotated, Callable, Literal @@ -9,8 +11,20 @@ "add_generation_prompt", "chat_template", "continue_final_message", + "conversation", + "documents", + "load_audio_from_video", + "max_length", + "messages", + "padding", + "processor_kwargs", + "return_assistant_tokens_mask", + "return_dict", + "return_tensors", "tokenize", + "tokenizer_kwargs", "tools", + "truncation", } @@ -96,6 +110,12 @@ def _base64_image_data(url: str) -> str: or not header.endswith(";base64") ): raise ChatRequestError("Images must use inline base64 data URLs.") + if data == "": + raise ChatRequestError("Images must contain valid base64 data.") + try: + base64.b64decode(data, validate=True) + except (binascii.Error, ValueError) as error: + raise ChatRequestError("Images must contain valid base64 data.") from error return data diff --git a/mlx_engine/server/http.py b/mlx_engine/server/http.py index 78269c27..2c061a61 100644 --- a/mlx_engine/server/http.py +++ b/mlx_engine/server/http.py @@ -35,7 +35,8 @@ logger = logging.getLogger(__name__) -_MAX_REQUEST_BODY_BYTES = 500 * 1024 * 1024 +_MAX_REQUEST_BODY_MIB = 64 +_MAX_REQUEST_BODY_BYTES = _MAX_REQUEST_BODY_MIB * 1024 * 1024 _SSE_WRITE_TIMEOUT_SECONDS = 30.0 @@ -228,8 +229,9 @@ def do_POST(self) -> None: return try: - body = self._read_json_body() - request = self.server.runtime.prepare_chat_generation(body) + request = self.server.runtime.prepare_chat_generation( + self._read_json_body() + ) except _RequestBodyTooLargeError as error: self._send_error(HTTPStatus.REQUEST_ENTITY_TOO_LARGE, str(error)) return @@ -369,7 +371,9 @@ def _read_json_body(self) -> object: if content_length <= 0: raise ChatRequestError("Content-Length must be a positive integer.") if content_length > _MAX_REQUEST_BODY_BYTES: - raise _RequestBodyTooLargeError("Request body exceeds the 500 MiB limit.") + raise _RequestBodyTooLargeError( + f"Request body exceeds the {_MAX_REQUEST_BODY_MIB} MiB limit." + ) try: encoded_body = self.rfile.read(content_length) diff --git a/tests/server/test_chat.py b/tests/server/test_chat.py index 23fb2424..a918062e 100644 --- a/tests/server/test_chat.py +++ b/tests/server/test_chat.py @@ -337,8 +337,20 @@ def test_final_assistant_message_is_rendered_as_a_prefill(): "add_generation_prompt", "chat_template", "continue_final_message", + "conversation", + "documents", + "load_audio_from_video", + "max_length", + "messages", + "padding", + "processor_kwargs", + "return_assistant_tokens_mask", + "return_dict", + "return_tensors", "tokenize", + "tokenizer_kwargs", "tools", + "truncation", ], ) def test_chat_template_kwargs_cannot_override_server_controls(control_name): @@ -364,7 +376,7 @@ def test_normalize_images_preserves_user_and_tool_result_order(): { "type": "image_url", "image_url": { - "url": "data:image/jpeg;base64,first-image", + "url": "data:image/jpeg;base64,Zmlyc3QtaW1hZ2U=", "detail": "auto", }, }, @@ -391,14 +403,14 @@ def test_normalize_images_preserves_user_and_tool_result_order(): { "type": "image_url", "image_url": { - "url": "data:image/png;base64,second-image", + "url": "data:image/png;base64,c2Vjb25kLWltYWdl", "detail": "auto", }, }, { "type": "image_url", "image_url": { - "url": "data:image/jpeg;base64,first-image", + "url": "data:image/jpeg;base64,Zmlyc3QtaW1hZ2U=", "detail": "auto", }, }, @@ -411,7 +423,11 @@ def test_normalize_images_preserves_user_and_tool_result_order(): supports_vision=True, ) - assert images_b64 == ["first-image", "second-image", "first-image"] + assert images_b64 == [ + "Zmlyc3QtaW1hZ2U=", + "c2Vjb25kLWltYWdl", + "Zmlyc3QtaW1hZ2U=", + ] assert normalized[0]["content"] == [ {"type": "text", "text": "First"}, {"type": "image"}, @@ -439,7 +455,7 @@ def test_prepare_vision_request_forwards_base64_to_generation_boundary(): { "type": "image_url", "image_url": { - "url": "data:image/jpeg;base64,image-payload", + "url": "data:image/jpeg;base64,aW1hZ2UtcGF5bG9hZA==", "detail": "auto", }, }, @@ -452,7 +468,7 @@ def test_prepare_vision_request_forwards_base64_to_generation_boundary(): tokenize=lambda _model_kit, _prompt: [7, 8], ) - assert request.generation_kwargs["images_b64"] == ["image-payload"] + assert request.generation_kwargs["images_b64"] == ["aW1hZ2UtcGF5bG9hZA=="] assert renderer.calls[0][0][0]["content"] == [ {"type": "text", "text": "Describe this"}, {"type": "image"}, @@ -502,6 +518,39 @@ def test_non_base64_image_url_is_rejected(): ) +@pytest.mark.parametrize( + "url", + [ + "data:image/jpeg;base64,", + "data:image/jpeg;base64,not-valid-base64!", + ], +) +def test_invalid_base64_image_data_is_rejected_before_rendering(url): + renderer = _FakeRenderer() + + with pytest.raises(ChatRequestError, match="valid base64 data"): + prepare_chat_generation_request( + _base_request( + messages=[ + { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": {"url": url}, + } + ], + } + ] + ), + model_kit=_FakeVisionModelKit(renderer), + supports_vision=True, + tokenize=lambda _model_kit, _prompt: [], + ) + + assert renderer.calls == [] + + def test_text_model_rejects_image_request(): renderer = _FakeRenderer() with pytest.raises(ChatRequestError, match="does not support images"): @@ -514,7 +563,7 @@ def test_text_model_rejects_image_request(): { "type": "image_url", "image_url": { - "url": "data:image/jpeg;base64,image-payload" + "url": "data:image/jpeg;base64,aW1hZ2UtcGF5bG9hZA==" }, } ], diff --git a/tests/server/test_http.py b/tests/server/test_http.py index a9c1b6ef..dffbb638 100644 --- a/tests/server/test_http.py +++ b/tests/server/test_http.py @@ -5,6 +5,7 @@ import struct import threading import time +import weakref import mlx_engine.server.http as server_http from mlx_engine.server.http import ( @@ -190,14 +191,68 @@ def test_invalid_and_oversized_content_lengths_are_rejected(): status, body = _request_with_content_length( port, - str(500 * 1024 * 1024 + 1), + str(server_http._MAX_REQUEST_BODY_BYTES + 1), ) assert status == 413 assert json.loads(body) == { - "error": {"message": "Request body exceeds the 500 MiB limit."} + "error": { + "message": ( + "Request body exceeds the " + f"{server_http._MAX_REQUEST_BODY_MIB} MiB limit." + ) + } } +def test_parsed_request_body_is_released_before_generation(monkeypatch): + class WeakReferenceableDict(dict): + pass + + pending_bodies = [WeakReferenceableDict(_request_body())] + body_reference = weakref.ref(pending_bodies[0]) + body_released_before_generation = [] + original_json_loads = json.loads + + def parse_request_body(value): + if isinstance(value, bytes): + return pending_bodies.pop() + return original_json_loads(value) + + monkeypatch.setattr(server_http.json, "loads", parse_request_body) + + def create_generator(_model_kit, _prompt_tokens, **_kwargs): + body_released_before_generation.append(body_reference() is None) + yield GenerationResult( + text="", + tokens=[], + top_logprobs=[], + stop_condition=GenerationStopCondition( + stop_reason="eos_token", + stop_string="", + stop_tokens=[2], + ), + ) + + runtime = EngineRuntime( + _FakeModelKit(), + supports_vision=False, + create_generator_fn=create_generator, + get_runtime_load_info_fn=lambda _model_kit: {}, + tokenize_fn=lambda _model_kit, _prompt: [1], + ) + + with _running_server(runtime) as port: + status, _response_body = _request( + port, + "POST", + "/v1/chat/completions", + body=_request_body(), + ) + + assert status == 200 + assert body_released_before_generation == [True] + + def test_invalid_generation_settings_are_rejected_before_streaming(): runtime = EngineRuntime( _FakeModelKit(), @@ -242,6 +297,39 @@ def test_invalid_generation_settings_are_rejected_before_streaming(): } +def test_invalid_base64_image_is_rejected_before_streaming(): + runtime = EngineRuntime( + _FakeVisionModelKit(), + supports_vision=True, + get_runtime_load_info_fn=lambda _model_kit: {}, + ) + body = _request_body() + body["messages"] = [ + { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": {"url": "data:image/jpeg;base64,not-valid-base64!"}, + } + ], + } + ] + + with _running_server(runtime) as port: + status, response_body = _request( + port, + "POST", + "/v1/chat/completions", + body=body, + ) + + assert status == 400 + assert json.loads(response_body) == { + "error": {"message": "Images must contain valid base64 data."} + } + + def test_chat_stream_forwards_generation_settings_and_returns_usage(): generation_calls = [] From 1af112a3df6ee77eb7b1c2bf5d5d5775403568f7 Mon Sep 17 00:00:00 2001 From: Neil Mehta Date: Wed, 29 Jul 2026 16:14:18 -0400 Subject: [PATCH 15/18] Harden server request handling --- mlx_engine/server/chat.py | 2 ++ mlx_engine/server/http.py | 24 +++++++++++--- tests/server/test_chat.py | 1 + tests/server/test_http.py | 68 +++++++++++++++++++++++++++++++++++++-- 4 files changed, 88 insertions(+), 7 deletions(-) diff --git a/mlx_engine/server/chat.py b/mlx_engine/server/chat.py index 1fb3f246..69c386a1 100644 --- a/mlx_engine/server/chat.py +++ b/mlx_engine/server/chat.py @@ -80,6 +80,7 @@ class ChatCompletionRequest(BaseModel): stream: Literal[True] temperature: _NonNegativeFloat max_tokens: _PositiveInt | None = None + max_completion_tokens: _PositiveInt | None = None stop: list[_NonEmptyString] | None = None top_p: _Probability | None = None top_k: _TopK @@ -171,6 +172,7 @@ def prepare_chat_generation_request( unsupported_controls = [ name for name, requested in ( + ("max_completion_tokens", request.max_completion_tokens is not None), ("seed", request.seed is not None), ("logprobs", request.logprobs is True), ("top_logprobs", request.top_logprobs is not None), diff --git a/mlx_engine/server/http.py b/mlx_engine/server/http.py index 2c061a61..5b0eb812 100644 --- a/mlx_engine/server/http.py +++ b/mlx_engine/server/http.py @@ -37,6 +37,7 @@ _MAX_REQUEST_BODY_MIB = 64 _MAX_REQUEST_BODY_BYTES = _MAX_REQUEST_BODY_MIB * 1024 * 1024 +_REQUEST_READ_TIMEOUT_SECONDS = 30.0 _SSE_WRITE_TIMEOUT_SECONDS = 30.0 @@ -48,6 +49,10 @@ class _RequestBodyTooLargeError(ValueError): pass +class _RequestReadTimeoutError(TimeoutError): + pass + + class EngineRuntime: def __init__( self, @@ -187,6 +192,11 @@ def __init__( self._active_sessions_lock = threading.Lock() super().__init__(server_address, MlxEngineRequestHandler) + def get_request(self): + connection, client_address = super().get_request() + connection.settimeout(_REQUEST_READ_TIMEOUT_SECONDS) + return connection, client_address + def register_session(self, session: GenerationSession) -> None: with self._active_sessions_lock: self._active_sessions.add(session) @@ -235,6 +245,9 @@ def do_POST(self) -> None: except _RequestBodyTooLargeError as error: self._send_error(HTTPStatus.REQUEST_ENTITY_TOO_LARGE, str(error)) return + except _RequestReadTimeoutError as error: + self._send_error(HTTPStatus.REQUEST_TIMEOUT, str(error)) + return except (ChatRequestError, ValidationError) as error: self._send_error(HTTPStatus.BAD_REQUEST, str(error)) return @@ -377,16 +390,17 @@ def _read_json_body(self) -> object: try: encoded_body = self.rfile.read(content_length) + except TimeoutError as error: + raise _RequestReadTimeoutError("Request body read timed out.") from error + try: return json.loads(encoded_body) except (UnicodeDecodeError, ValueError, json.JSONDecodeError) as error: raise ChatRequestError("The request body must be valid JSON.") from error def _is_authorized(self) -> bool: - authorization = self.headers.get("Authorization", "") - return hmac.compare_digest( - authorization, - f"Bearer {self.server.api_key}", - ) + authorization = self.headers.get("Authorization", "").encode() + expected = f"Bearer {self.server.api_key}".encode() + return hmac.compare_digest(authorization, expected) def _start_sse_response(self) -> None: try: diff --git a/tests/server/test_chat.py b/tests/server/test_chat.py index a918062e..7678ea16 100644 --- a/tests/server/test_chat.py +++ b/tests/server/test_chat.py @@ -171,6 +171,7 @@ def test_supported_generation_boundaries_and_unknown_future_fields_are_accepted( @pytest.mark.parametrize( ("overrides", "control_name"), [ + ({"max_completion_tokens": 1}, "max_completion_tokens"), ({"seed": 0}, "seed"), ({"logprobs": True}, "logprobs"), ({"top_logprobs": 5}, "top_logprobs"), diff --git a/tests/server/test_http.py b/tests/server/test_http.py index dffbb638..d6693910 100644 --- a/tests/server/test_http.py +++ b/tests/server/test_http.py @@ -145,6 +145,68 @@ def test_health_requires_auth_and_reports_actualized_context_length(): assert json.loads(body) == {"status": "ok", "context_length": 8192} +def test_non_ascii_authorization_is_rejected(): + runtime = EngineRuntime( + _FakeModelKit(), + supports_vision=False, + get_runtime_load_info_fn=lambda _model_kit: {}, + ) + + with _running_server(runtime) as port: + connection = http.client.HTTPConnection("127.0.0.1", port, timeout=2) + connection.request( + "GET", + "/health", + headers={"Authorization": "Bearer \xff"}, + ) + response = connection.getresponse() + response_body = response.read().decode("utf-8") + connection.close() + + assert response.status == 401 + assert json.loads(response_body) == {"error": {"message": "Unauthorized."}} + + +def test_partial_headers_time_out(monkeypatch): + monkeypatch.setattr(server_http, "_REQUEST_READ_TIMEOUT_SECONDS", 0.05) + runtime = EngineRuntime( + _FakeModelKit(), + supports_vision=False, + get_runtime_load_info_fn=lambda _model_kit: {}, + ) + + with _running_server(runtime) as port: + with socket.create_connection(("127.0.0.1", port), timeout=2) as client: + client.sendall(b"GET /health HTTP/1.1\r\n") + assert client.recv(1) == b"" + + +def test_partial_request_body_returns_request_timeout(monkeypatch): + monkeypatch.setattr(server_http, "_REQUEST_READ_TIMEOUT_SECONDS", 0.05) + runtime = EngineRuntime( + _FakeModelKit(), + supports_vision=False, + get_runtime_load_info_fn=lambda _model_kit: {}, + ) + + with _running_server(runtime) as port: + with socket.create_connection(("127.0.0.1", port), timeout=2) as client: + client.sendall( + b"POST /v1/chat/completions HTTP/1.1\r\n" + b"Host: 127.0.0.1\r\n" + b"Authorization: Bearer secret-token\r\n" + b"Content-Type: application/json\r\n" + b"Content-Length: 2\r\n\r\n" + b"{" + ) + response = b"" + while chunk := client.recv(4096): + response += chunk + + assert b"HTTP/1.1 408 Request Timeout" in response + assert b"Request body read timed out." in response + + def test_rejected_post_closes_connection_before_unread_body_can_be_reused(): runtime = EngineRuntime( _FakeModelKit(), @@ -284,7 +346,7 @@ def test_invalid_generation_settings_are_rejected_before_streaming(): assert "stop" in json.loads(response_body)["error"]["message"] unsupported_body = _request_body() - unsupported_body["seed"] = 0 + unsupported_body["max_completion_tokens"] = 1 status, response_body = _request( port, "POST", @@ -293,7 +355,9 @@ def test_invalid_generation_settings_are_rejected_before_streaming(): ) assert status == 400 assert json.loads(response_body) == { - "error": {"message": "Unsupported generation controls: seed."} + "error": { + "message": "Unsupported generation controls: max_completion_tokens." + } } From b1b8868e1f400e0c407885ad437f4c09314212da Mon Sep 17 00:00:00 2001 From: Neil Mehta Date: Wed, 29 Jul 2026 17:08:03 -0400 Subject: [PATCH 16/18] Reject decompression bomb images --- mlx_engine/server/chat.py | 20 +++++++++++++++- tests/server/test_chat.py | 48 +++++++++++++++++++++++++++------------ tests/server/test_http.py | 23 +++++++++++++++---- 3 files changed, 71 insertions(+), 20 deletions(-) diff --git a/mlx_engine/server/chat.py b/mlx_engine/server/chat.py index 69c386a1..c3c767d7 100644 --- a/mlx_engine/server/chat.py +++ b/mlx_engine/server/chat.py @@ -1,9 +1,11 @@ import base64 import binascii from dataclasses import dataclass +from io import BytesIO import json from typing import Annotated, Callable, Literal +from PIL import Image from pydantic import BaseModel, ConfigDict, Field @@ -103,6 +105,21 @@ class ChatGenerationRequest: generation_kwargs: dict[str, object] +def _validate_image_data(data: bytes) -> None: + try: + with Image.open(BytesIO(data)) as image: + max_image_pixels = Image.MAX_IMAGE_PIXELS + if ( + max_image_pixels is not None + and image.width * image.height > max_image_pixels + ): + raise ChatRequestError("Image dimensions are too large.") + except Image.DecompressionBombError as error: + raise ChatRequestError("Image dimensions are too large.") from error + except OSError as error: + raise ChatRequestError("Images must contain supported image data.") from error + + def _base64_image_data(url: str) -> str: header, separator, data = url.partition(",") if ( @@ -114,9 +131,10 @@ def _base64_image_data(url: str) -> str: if data == "": raise ChatRequestError("Images must contain valid base64 data.") try: - base64.b64decode(data, validate=True) + image_data = base64.b64decode(data, validate=True) except (binascii.Error, ValueError) as error: raise ChatRequestError("Images must contain valid base64 data.") from error + _validate_image_data(image_data) return data diff --git a/tests/server/test_chat.py b/tests/server/test_chat.py index 7678ea16..8ff33499 100644 --- a/tests/server/test_chat.py +++ b/tests/server/test_chat.py @@ -12,6 +12,20 @@ ) +_RED_PNG_B64 = ( + "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAIAAACQd1PeAAAADElEQVR4nGP4z8AAAAMBAQDJ" + "/pLvAAAAAElFTkSuQmCC" +) +_BLUE_PNG_B64 = ( + "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAIAAACQd1PeAAAADElEQVR4nGNgYPgPAAEDAQAI" + "icLsAAAAAElFTkSuQmCC" +) +# A 10,000 x 10,000 PNG header with no allocated pixel data. +_DECOMPRESSION_BOMB_PNG_B64 = ( + "iVBORw0KGgoAAAANSUhEUgAAJxAAACcQCAIAAAA1LPVwAAAAAElFTkSuQmCC" +) + + class _FakeRenderer: def __init__(self): self.chat_template = "model template" @@ -377,7 +391,7 @@ def test_normalize_images_preserves_user_and_tool_result_order(): { "type": "image_url", "image_url": { - "url": "data:image/jpeg;base64,Zmlyc3QtaW1hZ2U=", + "url": f"data:image/png;base64,{_RED_PNG_B64}", "detail": "auto", }, }, @@ -404,14 +418,14 @@ def test_normalize_images_preserves_user_and_tool_result_order(): { "type": "image_url", "image_url": { - "url": "data:image/png;base64,c2Vjb25kLWltYWdl", + "url": f"data:image/png;base64,{_BLUE_PNG_B64}", "detail": "auto", }, }, { "type": "image_url", "image_url": { - "url": "data:image/jpeg;base64,Zmlyc3QtaW1hZ2U=", + "url": f"data:image/png;base64,{_RED_PNG_B64}", "detail": "auto", }, }, @@ -425,9 +439,9 @@ def test_normalize_images_preserves_user_and_tool_result_order(): ) assert images_b64 == [ - "Zmlyc3QtaW1hZ2U=", - "c2Vjb25kLWltYWdl", - "Zmlyc3QtaW1hZ2U=", + _RED_PNG_B64, + _BLUE_PNG_B64, + _RED_PNG_B64, ] assert normalized[0]["content"] == [ {"type": "text", "text": "First"}, @@ -456,7 +470,7 @@ def test_prepare_vision_request_forwards_base64_to_generation_boundary(): { "type": "image_url", "image_url": { - "url": "data:image/jpeg;base64,aW1hZ2UtcGF5bG9hZA==", + "url": f"data:image/png;base64,{_RED_PNG_B64}", "detail": "auto", }, }, @@ -469,7 +483,7 @@ def test_prepare_vision_request_forwards_base64_to_generation_boundary(): tokenize=lambda _model_kit, _prompt: [7, 8], ) - assert request.generation_kwargs["images_b64"] == ["aW1hZ2UtcGF5bG9hZA=="] + assert request.generation_kwargs["images_b64"] == [_RED_PNG_B64] assert renderer.calls[0][0][0]["content"] == [ {"type": "text", "text": "Describe this"}, {"type": "image"}, @@ -519,17 +533,23 @@ def test_non_base64_image_url_is_rejected(): ) +@pytest.mark.filterwarnings("ignore::PIL.Image.DecompressionBombWarning") @pytest.mark.parametrize( - "url", + ("url", "error_message"), [ - "data:image/jpeg;base64,", - "data:image/jpeg;base64,not-valid-base64!", + ("data:image/jpeg;base64,", "valid base64 data"), + ("data:image/jpeg;base64,not-valid-base64!", "valid base64 data"), + ("data:image/png;base64,bm90IGFuIGltYWdl", "supported image data"), + ( + f"data:image/png;base64,{_DECOMPRESSION_BOMB_PNG_B64}", + "Image dimensions are too large", + ), ], ) -def test_invalid_base64_image_data_is_rejected_before_rendering(url): +def test_invalid_image_data_is_rejected_before_rendering(url, error_message): renderer = _FakeRenderer() - with pytest.raises(ChatRequestError, match="valid base64 data"): + with pytest.raises(ChatRequestError, match=error_message): prepare_chat_generation_request( _base_request( messages=[ @@ -564,7 +584,7 @@ def test_text_model_rejects_image_request(): { "type": "image_url", "image_url": { - "url": "data:image/jpeg;base64,aW1hZ2UtcGF5bG9hZA==" + "url": f"data:image/png;base64,{_RED_PNG_B64}" }, } ], diff --git a/tests/server/test_http.py b/tests/server/test_http.py index d6693910..4b536496 100644 --- a/tests/server/test_http.py +++ b/tests/server/test_http.py @@ -8,6 +8,7 @@ import weakref import mlx_engine.server.http as server_http +import pytest from mlx_engine.server.http import ( EngineRuntime, GenerationSession, @@ -21,6 +22,11 @@ from mlx_engine.utils.token import Token +_DECOMPRESSION_BOMB_PNG_B64 = ( + "iVBORw0KGgoAAAANSUhEUgAAJxAAACcQCAIAAAA1LPVwAAAAAElFTkSuQmCC" +) + + class _FakeRenderer: chat_template = "model template" @@ -361,7 +367,16 @@ def test_invalid_generation_settings_are_rejected_before_streaming(): } -def test_invalid_base64_image_is_rejected_before_streaming(): +@pytest.mark.filterwarnings("ignore::PIL.Image.DecompressionBombWarning") +@pytest.mark.parametrize( + ("image_data", "error_message"), + [ + ("not-valid-base64!", "Images must contain valid base64 data."), + ("bm90IGFuIGltYWdl", "Images must contain supported image data."), + (_DECOMPRESSION_BOMB_PNG_B64, "Image dimensions are too large."), + ], +) +def test_invalid_image_is_rejected_before_streaming(image_data, error_message): runtime = EngineRuntime( _FakeVisionModelKit(), supports_vision=True, @@ -374,7 +389,7 @@ def test_invalid_base64_image_is_rejected_before_streaming(): "content": [ { "type": "image_url", - "image_url": {"url": "data:image/jpeg;base64,not-valid-base64!"}, + "image_url": {"url": f"data:image/png;base64,{image_data}"}, } ], } @@ -389,9 +404,7 @@ def test_invalid_base64_image_is_rejected_before_streaming(): ) assert status == 400 - assert json.loads(response_body) == { - "error": {"message": "Images must contain valid base64 data."} - } + assert json.loads(response_body) == {"error": {"message": error_message}} def test_chat_stream_forwards_generation_settings_and_returns_usage(): From 46240a28efc956108ffebeca2ef169f33238200b Mon Sep 17 00:00:00 2001 From: Neil Mehta Date: Wed, 29 Jul 2026 17:41:42 -0400 Subject: [PATCH 17/18] Validate complete image requests --- mlx_engine/server/chat.py | 31 ++++++++++++++++---------- tests/server/test_chat.py | 32 +++++++++++++++++++++++++++ tests/server/test_http.py | 46 +++++++++++++++++++++++++++++++++++++++ 3 files changed, 97 insertions(+), 12 deletions(-) diff --git a/mlx_engine/server/chat.py b/mlx_engine/server/chat.py index c3c767d7..36533955 100644 --- a/mlx_engine/server/chat.py +++ b/mlx_engine/server/chat.py @@ -105,22 +105,26 @@ class ChatGenerationRequest: generation_kwargs: dict[str, object] -def _validate_image_data(data: bytes) -> None: +def _validate_image_pixel_count(pixel_count: int) -> None: + max_image_pixels = Image.MAX_IMAGE_PIXELS + if max_image_pixels is not None and pixel_count > max_image_pixels: + raise ChatRequestError("Image dimensions are too large.") + + +def _validate_image_data(data: bytes) -> int: try: with Image.open(BytesIO(data)) as image: - max_image_pixels = Image.MAX_IMAGE_PIXELS - if ( - max_image_pixels is not None - and image.width * image.height > max_image_pixels - ): - raise ChatRequestError("Image dimensions are too large.") + pixel_count = image.width * image.height + _validate_image_pixel_count(pixel_count) + image.verify() + return pixel_count except Image.DecompressionBombError as error: raise ChatRequestError("Image dimensions are too large.") from error - except OSError as error: + except (OSError, SyntaxError) as error: raise ChatRequestError("Images must contain supported image data.") from error -def _base64_image_data(url: str) -> str: +def _base64_image_data(url: str) -> tuple[str, int]: header, separator, data = url.partition(",") if ( separator == "" @@ -134,8 +138,7 @@ def _base64_image_data(url: str) -> str: image_data = base64.b64decode(data, validate=True) except (binascii.Error, ValueError) as error: raise ChatRequestError("Images must contain valid base64 data.") from error - _validate_image_data(image_data) - return data + return data, _validate_image_data(image_data) def normalize_messages( @@ -145,6 +148,7 @@ def normalize_messages( ) -> tuple[list[dict], list[str]]: normalized_messages: list[dict] = [] images_b64: list[str] = [] + total_image_pixels = 0 for message in messages: normalized_message = message.model_dump(exclude_unset=True) @@ -156,7 +160,10 @@ def normalize_messages( text_parts.append(part.text) normalized_parts.append({"type": "text", "text": part.text}) else: - images_b64.append(_base64_image_data(part.image_url.url)) + image_b64, image_pixels = _base64_image_data(part.image_url.url) + total_image_pixels += image_pixels + _validate_image_pixel_count(total_image_pixels) + images_b64.append(image_b64) normalized_parts.append({"type": "image"}) normalized_message["content"] = ( normalized_parts if supports_vision else "".join(text_parts) diff --git a/tests/server/test_chat.py b/tests/server/test_chat.py index 8ff33499..d5a78dd6 100644 --- a/tests/server/test_chat.py +++ b/tests/server/test_chat.py @@ -1,6 +1,7 @@ import json import pytest +from PIL import Image from pydantic import ValidationError from transformers.utils.chat_template_utils import render_jinja_template @@ -24,6 +25,7 @@ _DECOMPRESSION_BOMB_PNG_B64 = ( "iVBORw0KGgoAAAANSUhEUgAAJxAAACcQCAIAAAA1LPVwAAAAAElFTkSuQmCC" ) +_TRUNCATED_PNG_B64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAIAAACQd1PeAAAADElEQVQ=" class _FakeRenderer: @@ -540,6 +542,10 @@ def test_non_base64_image_url_is_rejected(): ("data:image/jpeg;base64,", "valid base64 data"), ("data:image/jpeg;base64,not-valid-base64!", "valid base64 data"), ("data:image/png;base64,bm90IGFuIGltYWdl", "supported image data"), + ( + f"data:image/png;base64,{_TRUNCATED_PNG_B64}", + "supported image data", + ), ( f"data:image/png;base64,{_DECOMPRESSION_BOMB_PNG_B64}", "Image dimensions are too large", @@ -572,6 +578,32 @@ def test_invalid_image_data_is_rejected_before_rendering(url, error_message): assert renderer.calls == [] +def test_aggregate_image_pixels_are_bounded_before_rendering(monkeypatch): + monkeypatch.setattr(Image, "MAX_IMAGE_PIXELS", 2) + renderer = _FakeRenderer() + image_part = { + "type": "image_url", + "image_url": {"url": f"data:image/png;base64,{_RED_PNG_B64}"}, + } + + with pytest.raises(ChatRequestError, match="Image dimensions are too large"): + prepare_chat_generation_request( + _base_request( + messages=[ + { + "role": "user", + "content": [image_part, image_part, image_part], + } + ] + ), + model_kit=_FakeVisionModelKit(renderer), + supports_vision=True, + tokenize=lambda _model_kit, _prompt: [], + ) + + assert renderer.calls == [] + + def test_text_model_rejects_image_request(): renderer = _FakeRenderer() with pytest.raises(ChatRequestError, match="does not support images"): diff --git a/tests/server/test_http.py b/tests/server/test_http.py index 4b536496..925c47dc 100644 --- a/tests/server/test_http.py +++ b/tests/server/test_http.py @@ -25,6 +25,7 @@ _DECOMPRESSION_BOMB_PNG_B64 = ( "iVBORw0KGgoAAAANSUhEUgAAJxAAACcQCAIAAAA1LPVwAAAAAElFTkSuQmCC" ) +_TRUNCATED_PNG_B64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAIAAACQd1PeAAAADElEQVQ=" class _FakeRenderer: @@ -373,6 +374,7 @@ def test_invalid_generation_settings_are_rejected_before_streaming(): [ ("not-valid-base64!", "Images must contain valid base64 data."), ("bm90IGFuIGltYWdl", "Images must contain supported image data."), + (_TRUNCATED_PNG_B64, "Images must contain supported image data."), (_DECOMPRESSION_BOMB_PNG_B64, "Image dimensions are too large."), ], ) @@ -664,6 +666,50 @@ def create_generator(_model_kit, _prompt_tokens, **_kwargs): ] +def test_mid_stream_generation_errors_emit_a_recognized_error_frame(): + generation_error = RuntimeError("generation failed after output") + + def create_generator(_model_kit, _prompt_tokens, **_kwargs): + yield GenerationResult( + text="partial output", + tokens=[Token(id=10, text="partial output", logprob=-0.1)], + top_logprobs=[], + stop_condition=None, + ) + raise generation_error + + runtime = EngineRuntime( + _FakeModelKit(), + supports_vision=False, + create_generator_fn=create_generator, + get_runtime_load_info_fn=lambda _model_kit: {}, + tokenize_fn=lambda _model_kit, _prompt: [1], + ) + + with _running_server(runtime) as port: + status, response_text = _request( + port, + "POST", + "/v1/chat/completions", + body=_request_body(), + ) + + assert status == 200 + assert _parse_sse(response_text) == [ + { + "choices": [ + { + "index": 0, + "delta": {"content": "partial output"}, + "finish_reason": None, + } + ] + }, + {"error": {"message": str(generation_error)}}, + ] + assert "data: [DONE]" not in response_text + + def test_stalled_sse_write_cancels_the_active_mlx_request(monkeypatch): monkeypatch.setattr(server_http, "_SSE_WRITE_TIMEOUT_SECONDS", 0.05) generation_stopped = threading.Event() From 92833844d3ad801ff7b99d8477cd67c73a91e991 Mon Sep 17 00:00:00 2001 From: Neil Mehta Date: Thu, 30 Jul 2026 13:43:22 -0400 Subject: [PATCH 18/18] Prevent batched request starvation --- mlx_engine/model_kit/batched_model_kit.py | 63 ++++++------ tests/test_batched_model_kit_scheduler.py | 118 ++++++++++++++++------ 2 files changed, 116 insertions(+), 65 deletions(-) diff --git a/mlx_engine/model_kit/batched_model_kit.py b/mlx_engine/model_kit/batched_model_kit.py index 88dcd7ac..678e2076 100644 --- a/mlx_engine/model_kit/batched_model_kit.py +++ b/mlx_engine/model_kit/batched_model_kit.py @@ -334,10 +334,9 @@ def get_next_request(timeout=None): timeout: None | float = None if (len(self._batch_results) > 0) else 0.1 request = get_next_request(timeout=timeout) - # We got a request + # Handle at most one request before advancing the current batch. if request is not None: if isinstance(request, CancelGenerationRequest): - # Handle cancel request found_request_id = False request_id = request.request_id for uid, entry in self._batch_results.items(): @@ -349,45 +348,40 @@ def get_next_request(timeout=None): break if not found_request_id: logger.warning(f"Could not cancel {request_id=} (id not found)") - continue - - with mx.stream(batch_generator.stream): - cache, cached_prefix, rest = _prepare_prompt_cache_for_generation( - self._prompt_cache, current_model_key, request.prompt_tokens - ) - - # Keep cache allocation on the same MLX stream as generation. - (uid,) = batch_generator.insert( - [rest], - [request.max_tokens], - caches=[cache], - all_tokens=[cached_prefix], - samplers=[request.samplers], - logits_processors=[request.logits_processors], - ) + else: + with mx.stream(batch_generator.stream): + cache, cached_prefix, rest = ( + _prepare_prompt_cache_for_generation( + self._prompt_cache, + current_model_key, + request.prompt_tokens, + ) + ) - # Track this request - self._batch_results[uid] = { - "cache_key": request.prompt_tokens[:], - "rqueue": request.rqueue, - "detokenizer": self.tokenizer.detokenizer, - "top_logprobs": request.top_logprobs, - "request_id": request.request_id, - } + # Keep cache allocation on the same MLX stream as generation. + (uid,) = batch_generator.insert( + [rest], + [request.max_tokens], + caches=[cache], + all_tokens=[cached_prefix], + samplers=[request.samplers], + logits_processors=[request.logits_processors], + ) - # Check for new requests - continue + self._batch_results[uid] = { + "cache_key": request.prompt_tokens[:], + "rqueue": request.rqueue, + "detokenizer": self.tokenizer.detokenizer, + "top_logprobs": request.top_logprobs, + "request_id": request.request_id, + } - # No request so serve from the current batch - if len(self._batch_results) == 0: + if self._shutdown.is_set() or len(self._batch_results) == 0: continue time_budget = 0.5 start = time.time() while True: - if not self._requests.empty() or time.time() - start > time_budget: - break - prompt_responses, generation_responses = batch_generator.next() if not prompt_responses and not generation_responses: break @@ -449,6 +443,9 @@ def get_next_request(timeout=None): ) del self._batch_results[r.uid] + if not self._requests.empty() or time.time() - start > time_budget: + break + for entry in self._batch_results.values(): entry["rqueue"].put(RequestCancelled("Model shutdown requested")) diff --git a/tests/test_batched_model_kit_scheduler.py b/tests/test_batched_model_kit_scheduler.py index fc8e3dda..e4e40d4d 100644 --- a/tests/test_batched_model_kit_scheduler.py +++ b/tests/test_batched_model_kit_scheduler.py @@ -20,6 +20,48 @@ def _request(request_id: str) -> GenerationRequest: ) +def _make_model_kit() -> BatchedModelKit: + model_kit = BatchedModelKit.__new__(BatchedModelKit) + model_kit.model = object() + model_kit.tokenizer = SimpleNamespace(detokenizer=object(), eos_token_ids=[]) + model_kit._requests = Queue() + model_kit._prompt_cache = object() + model_kit._batch_results = {} + model_kit._backend_exception = None + model_kit._generation_thread = None + model_kit._shutdown = threading.Event() + model_kit._startup_complete = threading.Event() + model_kit._seed = None + model_kit._max_seq_nums = 4 + model_kit._prefill_step_size = 512 + model_kit._max_kv_size = 2048 + return model_kit + + +def _install_scheduler_fakes(monkeypatch, batch_generator_type) -> None: + monkeypatch.setattr( + batched_model_kit_module, + "BatchGenerator", + batch_generator_type, + ) + monkeypatch.setattr( + batched_model_kit_module, + "_prepare_prompt_cache_for_generation", + lambda *_args: (None, [], [1, 2, 3]), + ) + monkeypatch.setattr( + batched_model_kit_module, + "install_mlx_compile_cache_cleanup_for_thread", + lambda: None, + ) + monkeypatch.setattr(batched_model_kit_module, "set_seed", lambda _seed: None) + monkeypatch.setattr( + batched_model_kit_module.mx, + "stream", + lambda _stream: nullcontext(), + ) + + def test_pending_request_is_admitted_before_the_next_generation_step(monkeypatch): first_request = _request("first") second_request = _request("second") @@ -49,39 +91,9 @@ def next(self): model_kit._shutdown.set() return [], [] - monkeypatch.setattr(batched_model_kit_module, "BatchGenerator", FakeBatchGenerator) - monkeypatch.setattr( - batched_model_kit_module, - "_prepare_prompt_cache_for_generation", - lambda *_args: (None, [], [1, 2, 3]), - ) - monkeypatch.setattr( - batched_model_kit_module, - "install_mlx_compile_cache_cleanup_for_thread", - lambda: None, - ) - monkeypatch.setattr(batched_model_kit_module, "set_seed", lambda _seed: None) - monkeypatch.setattr( - batched_model_kit_module.mx, - "stream", - lambda _stream: nullcontext(), - ) - - model_kit = BatchedModelKit.__new__(BatchedModelKit) - model_kit.model = object() - model_kit.tokenizer = SimpleNamespace(detokenizer=object(), eos_token_ids=[]) - model_kit._requests = Queue() + _install_scheduler_fakes(monkeypatch, FakeBatchGenerator) + model_kit = _make_model_kit() model_kit._requests.put(first_request) - model_kit._prompt_cache = object() - model_kit._batch_results = {} - model_kit._backend_exception = None - model_kit._generation_thread = None - model_kit._shutdown = threading.Event() - model_kit._startup_complete = threading.Event() - model_kit._seed = None - model_kit._max_seq_nums = 4 - model_kit._prefill_step_size = 512 - model_kit._max_kv_size = 2048 scheduler_thread = threading.Thread(target=model_kit._generate) scheduler_thread.start() @@ -90,3 +102,45 @@ def next(self): assert not scheduler_thread.is_alive() assert inserted_request_count == 2 assert generation_step_count == 1 + + +def test_active_generation_progresses_while_request_queue_stays_populated(monkeypatch): + events = [] + inserted_request_count = 0 + generation_step_count = 0 + + class FakeBatchGenerator: + stream = object() + + def __init__(self, *_args, **_kwargs): + pass + + def insert(self, *_args, **_kwargs): + nonlocal inserted_request_count + uid = inserted_request_count + inserted_request_count += 1 + events.append("insert") + if inserted_request_count < 3: + model_kit._requests.put(_request(f"queued-{inserted_request_count}")) + return (uid,) + + def next(self): + nonlocal generation_step_count + generation_step_count += 1 + events.append("next") + if generation_step_count < 3: + assert not model_kit._requests.empty() + if generation_step_count == 3: + model_kit._shutdown.set() + return [], [] + + _install_scheduler_fakes(monkeypatch, FakeBatchGenerator) + model_kit = _make_model_kit() + model_kit._requests.put(_request("first")) + + scheduler_thread = threading.Thread(target=model_kit._generate) + scheduler_thread.start() + scheduler_thread.join(timeout=2) + + assert not scheduler_thread.is_alive() + assert events == ["insert", "next", "insert", "next", "insert", "next"]