diff --git a/mlx_engine/model_kit/batched_model_kit.py b/mlx_engine/model_kit/batched_model_kit.py index f0f713ec..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 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/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/__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..85c92b52 --- /dev/null +++ b/mlx_engine/server/__main__.py @@ -0,0 +1,76 @@ +import argparse +import logging +import os +import signal +import threading + +from mlx_engine import load_model + +from .http import EngineRuntime, MlxEngineHttpServer + + +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.", + 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("--context-length", required=True, type=int) + parser.add_argument("--parallel-sessions", required=True, type=int) + parser.add_argument("--seed", type=int) + return parser + + +def main() -> None: + 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( + args.model, + max_kv_size=args.context_length, + max_seq_nums=args.parallel_sessions, + seed=args.seed, + trust_remote_code=False, + ) + runtime = EngineRuntime(model_kit) + server = None + + try: + server = MlxEngineHttpServer( + (args.host, args.port), + api_key=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..36533955 --- /dev/null +++ b/mlx_engine/server/chat.py @@ -0,0 +1,277 @@ +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 + + +_CHAT_TEMPLATE_CONTROL_KEYS = { + "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", +} + + +class ChatRequestError(ValueError): + """The chat request does not match the server contract.""" + + +class _ImageDataUrl(BaseModel): + url: str + + +class _TextContentPart(BaseModel): + type: Literal["text"] + text: str + + +class _InlineImageContentPart(BaseModel): + type: Literal["image_url"] + image_url: _ImageDataUrl + + +_ContentPart = Annotated[ + _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)] +_NonEmptyString = Annotated[str, Field(min_length=1)] +_TopK = Annotated[int, Field(ge=-1, le=500, strict=True)] + + +class ChatMessage(BaseModel): + model_config = ConfigDict(extra="allow") + + role: Literal["system", "user", "assistant", "tool"] + 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") + + messages: list[ChatMessage] + 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 + 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) + + +@dataclass(frozen=True) +class ChatGenerationRequest: + prompt_tokens: list[int] + generation_kwargs: dict[str, object] + + +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: + 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, SyntaxError) as error: + raise ChatRequestError("Images must contain supported image data.") from error + + +def _base64_image_data(url: str) -> tuple[str, int]: + 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.") + if data == "": + raise ChatRequestError("Images must contain valid base64 data.") + try: + image_data = base64.b64decode(data, validate=True) + except (binascii.Error, ValueError) as error: + raise ChatRequestError("Images must contain valid base64 data.") from error + return data, _validate_image_data(image_data) + + +def normalize_messages( + messages: list[ChatMessage], + *, + supports_vision: bool, +) -> 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) + 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: + 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) + ) + normalized_messages.append(normalized_message) + + return normalized_messages, images_b64 + + +def _get_chat_template(model_kit: object, *, supports_vision: bool) -> Callable: + 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( + body: object, + *, + model_kit: object, + supports_vision: bool, + tokenize: Callable[[object, str], list[int]], +) -> ChatGenerationRequest: + request = ChatCompletionRequest.model_validate(body) + if request.tools: + raise ChatRequestError("Tools are not supported yet.") + + 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), + ("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" + ) + 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, + supports_vision=supports_vision, + ) + 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 has_assistant_prefill: + 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=add_generation_prompt, + **template_kwargs, + ) + + generation_kwargs: dict[str, object] = { + "images_b64": images_b64, + "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), + ("top_p", request.top_p), + ("min_p", request.min_p), + ("repetition_penalty", request.repeat_penalty), + ): + if value is not None: + generation_kwargs[name] = value + + return ChatGenerationRequest( + prompt_tokens=tokenize(model_kit, prompt), + generation_kwargs=generation_kwargs, + ) diff --git a/mlx_engine/server/http.py b/mlx_engine/server/http.py new file mode 100644 index 00000000..5b0eb812 --- /dev/null +++ b/mlx_engine/server/http.py @@ -0,0 +1,453 @@ +from __future__ import annotations + +import hmac +from http import HTTPStatus +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +import json +import logging +import threading +from typing import Callable, Iterator +import uuid + +from pydantic import ValidationError + +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.generation_result import ( + GenerationResult, + GenerationStopCondition, +) +from mlx_engine.utils.prompt_progress_reporter import PromptProgressReporter + +from .chat import ( + ChatGenerationRequest, + ChatRequestError, + prepare_chat_generation_request, +) + + +logger = logging.getLogger(__name__) + + +_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 + + +class _ClientConnectionError(Exception): + pass + + +class _RequestBodyTooLargeError(ValueError): + pass + + +class _RequestReadTimeoutError(TimeoutError): + pass + + +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 + + 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: + return self._create_generator( + self.model_kit, + request.prompt_tokens, + request_id=request_id, + prompt_progress_reporter=prompt_progress_reporter, + **request.generation_kwargs, + ) + + def runtime_context_length(self) -> int | 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: + self._unload(self.model_kit) + + +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, + 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( + self, + is_draft: bool, + cached_tokens: int, + total_prompt_tokens: int, + prefill_tokens_processed: int, + ) -> bool: + if not is_draft and self._use_reported_prompt_tokens: + self.prompt_tokens = total_prompt_tokens + return self._report(is_draft) + + def update(self, is_draft: bool, prefill_tokens_processed: int) -> bool: + return self._report(is_draft) + + def finish( + self, + is_draft: bool, + prefill_tokens_processed: int | None = None, + ) -> bool: + return self._report(is_draft) + + 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): + 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 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) + + 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: + 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 + 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 + 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) + self.server.register_session(session) + generator = None + normal_completion = False + try: + self._start_sse_response() + reporter = _SsePromptProgressReporter( + self, + session, + len(request.prompt_tokens), + use_reported_prompt_tokens=self.server.runtime.supports_vision, + ) + 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 _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 _ClientConnectionError: + pass + finally: + 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, + generator: Iterator[GenerationResult], + reporter: _SsePromptProgressReporter, + ) -> None: + completion_tokens = 0 + terminal_sent = False + + 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 + + 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: GenerationStopCondition, + prompt_tokens: int, + completion_tokens: int, + ) -> dict: + stop_reason = 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": 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_header = self.headers.get("Content-Length") + if content_length_header is None: + raise ChatRequestError("Content-Length is required.") + try: + 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( + f"Request body exceeds the {_MAX_REQUEST_BODY_MIB} MiB limit." + ) + + 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", "").encode() + expected = f"Bearer {self.server.api_key}".encode() + return hmac.compare_digest(authorization, expected) + + 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: + 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( + status, + {"error": {"message": message}}, + close_connection=True, + ) + + 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) + + def log_message(self, format: str, *args: object) -> None: + logger.debug("HTTP %s - %s", self.address_string(), format % args) 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/mlx_engine/utils/prompt_progress_reporter.py b/mlx_engine/utils/prompt_progress_reporter.py index ac2e5f72..b85c0da9 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 + 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 @@ -246,8 +247,8 @@ 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) + prefill_tokens_processed = min(processed_tokens, prefill_tokens) if self._first_call: self._first_call = False @@ -256,17 +257,19 @@ 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 - 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 + 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 new file mode 100644 index 00000000..d5a78dd6 --- /dev/null +++ b/tests/server/test_chat.py @@ -0,0 +1,629 @@ +import json + +import pytest +from PIL import Image +from pydantic import ValidationError +from transformers.utils.chat_template_utils import render_jinja_template + +from mlx_engine.server.chat import ( + ChatMessage, + ChatRequestError, + normalize_messages, + prepare_chat_generation_request, +) + + +_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" +) +_TRUNCATED_PNG_B64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAIAAACQd1PeAAAADElEQVQ=" + + +class _FakeRenderer: + def __init__(self): + self.chat_template = "model template" + self.calls = [] + + def apply_chat_template(self, messages, **kwargs): + self.calls.append((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 + + +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"}), + 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_tokens == [1, 2, 3] + 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] + 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" + + +@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() + + 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"), + [ + ({"max_completion_tokens": 1}, "max_completion_tokens"), + ({"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}, + {"stop": [""]}, + {"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() + + 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_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() + + 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( + "control_name", + [ + "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): + 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 = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "First"}, + { + "type": "image_url", + "image_url": { + "url": f"data:image/png;base64,{_RED_PNG_B64}", + "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": f"data:image/png;base64,{_BLUE_PNG_B64}", + "detail": "auto", + }, + }, + { + "type": "image_url", + "image_url": { + "url": f"data:image/png;base64,{_RED_PNG_B64}", + "detail": "auto", + }, + }, + ], + }, + ] + + normalized, images_b64 = normalize_messages( + [ChatMessage.model_validate(message) for message in messages], + supports_vision=True, + ) + + assert images_b64 == [ + _RED_PNG_B64, + _BLUE_PNG_B64, + _RED_PNG_B64, + ] + 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": f"data:image/png;base64,{_RED_PNG_B64}", + "detail": "auto", + }, + }, + ], + } + ] + ), + model_kit=model_kit, + supports_vision=True, + tokenize=lambda _model_kit, _prompt: [7, 8], + ) + + assert request.generation_kwargs["images_b64"] == [_RED_PNG_B64] + 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_tokens == [7, 8] + assert len(tokenizer_renderer.calls) == 1 + + +def test_non_base64_image_url_is_rejected(): + with pytest.raises(ChatRequestError, match="inline base64"): + normalize_messages( + [ + ChatMessage.model_validate( + { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": {"url": "https://example.com/image.png"}, + } + ], + } + ) + ], + supports_vision=True, + ) + + +@pytest.mark.filterwarnings("ignore::PIL.Image.DecompressionBombWarning") +@pytest.mark.parametrize( + ("url", "error_message"), + [ + ("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", + ), + ], +) +def test_invalid_image_data_is_rejected_before_rendering(url, error_message): + renderer = _FakeRenderer() + + with pytest.raises(ChatRequestError, match=error_message): + 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_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"): + prepare_chat_generation_request( + _base_request( + messages=[ + { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": { + "url": f"data:image/png;base64,{_RED_PNG_B64}" + }, + } + ], + } + ] + ), + 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..925c47dc --- /dev/null +++ b/tests/server/test_http.py @@ -0,0 +1,873 @@ +from contextlib import contextmanager +import http.client +import json +import socket +import struct +import threading +import time +import weakref + +import mlx_engine.server.http as server_http +import pytest +from mlx_engine.server.http import ( + EngineRuntime, + GenerationSession, + MlxEngineHttpServer, +) +from mlx_engine.utils.generation_result import ( + GenerationResult, + GenerationStopCondition, +) +from mlx_engine.utils.prompt_progress_reporter import BatchedMlxLmReporterAdapter +from mlx_engine.utils.token import Token + + +_DECOMPRESSION_BOMB_PNG_B64 = ( + "iVBORw0KGgoAAAANSUhEUgAAJxAAACcQCAIAAAA1LPVwAAAAAElFTkSuQmCC" +) +_TRUNCATED_PNG_B64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAIAAACQd1PeAAAADElEQVQ=" + + +class _FakeRenderer: + chat_template = "model template" + + 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() + + +class _FakeVisionModelKit: + def __init__(self): + self.processor = _FakeRenderer() + + +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, *, 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: + 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 _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(), + 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_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(), + 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(), + 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(server_http._MAX_REQUEST_BODY_BYTES + 1), + ) + assert status == 413 + assert json.loads(body) == { + "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(), + 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"] + + 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["max_completion_tokens"] = 1 + 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: max_completion_tokens." + } + } + + +@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."), + (_TRUNCATED_PNG_B64, "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, + get_runtime_load_info_fn=lambda _model_kit: {}, + ) + body = _request_body() + body["messages"] = [ + { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": {"url": f"data:image/png;base64,{image_data}"}, + } + ], + } + ] + + 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": error_message}} + + +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=0, + total_prompt_tokens=3, + prefill_tokens_processed=0, + ) + assert reporter.update(is_draft=False, prefill_tokens_processed=2) + 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(), + "response_format": { + "type": "json_schema", + "json_schema": { + "schema": { + "type": "object", + "properties": {"answer": {"type": "string"}}, + } + }, + }, + }, + ) + + 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": 3, + "completion_tokens": 2, + "total_tokens": 5, + }, + "__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 + assert json.loads(generation_kwargs["json_schema"]) == { + "type": "object", + "properties": {"answer": {"type": "string"}}, + } + + +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, + ) + 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=[], + 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: + 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, + "POST", + "/v1/chat/completions", + body=_request_body(), + ) + + assert status == 200 + terminal_event = _parse_sse(response_text)[0] + assert terminal_event["usage"]["prompt_tokens"] == 9 + + +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_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_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() + stopped_request_ids = [] + large_text = "x" * (1024 * 1024) + + def create_generator(_model_kit, _prompt_tokens, **_kwargs): + 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: {}, + stop_generation_fn=stop_generation, + tokenize_fn=lambda _model_kit, _prompt: [1, 2, 3], + ) + + 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 len(stopped_request_ids) == 1 + assert stopped_request_ids[0] != "" + + +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(): + stopped_request_ids = [] + + def stop_generation(_model_kit, request_id): + stopped_request_ids.append(request_id) + + runtime = EngineRuntime( + _FakeModelKit(), + supports_vision=False, + stop_generation_fn=stop_generation, + ) + session = GenerationSession(runtime) + + session.cancel() + + assert stopped_request_ids == [session.request_id] + + +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, + ) + session = GenerationSession(runtime) + + session.cancel() + session.cancel() + + assert stop_calls == [session.request_id] + + +def test_runtime_unloads_model(): + 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() + + assert unload_calls == [model_kit] diff --git a/tests/server/test_main.py b/tests/server/test_main.py new file mode 100644 index 00000000..796dcf85 --- /dev/null +++ b/tests/server/test_main.py @@ -0,0 +1,38 @@ +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/test_batched_model_kit_scheduler.py b/tests/test_batched_model_kit_scheduler.py new file mode 100644 index 00000000..e4e40d4d --- /dev/null +++ b/tests/test_batched_model_kit_scheduler.py @@ -0,0 +1,146 @@ +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 _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") + 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 [], [] + + _install_scheduler_fakes(monkeypatch, FakeBatchGenerator) + model_kit = _make_model_kit() + model_kit._requests.put(first_request) + + 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 + + +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"] 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_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) diff --git a/tests/utils/test_prompt_progress_reporter.py b/tests/utils/test_prompt_progress_reporter.py index 67989c66..375ef749 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(1, 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)