Changes to support input validation to match OpenAI behavior.

serve/mlc_serve/engine/async_connector.py

@@ -1,6 +1,6 @@
 import asyncio
 import logging
-from typing import AsyncIterator
+from typing import AsyncIterator, Any
 
 from .base import (
     InferenceEngine,
@@ -15,7 +15,9 @@
 
 
 class TextGenerationError(Exception):
-    pass
+    def __init__(self, error: Any) -> None:
+        self.error = error
+        super(self).__init__(error)
 
 
 class AsyncEngineConnector:
@@ -98,6 +100,7 @@ async def _get_queue_item_until_stopped(self, queue: ResultQueue) -> RequestOutp
 
         if wait_shutdown_task.done():
             if self.engine_loop_exception is not None:
+                import pdb; pdb.set_trace()
                 raise RuntimeError(
                     f"InferenceEngine raised exception: {self.engine_loop_exception}"
                 )

serve/mlc_serve/engine/base.py

@@ -1,6 +1,8 @@
+from __future__ import annotations
+
 from dataclasses import dataclass, field
 from enum import Enum
-from typing import Optional, List
+from typing import List, Callable, Any, Optional
 import json
 import inspect
 from .sampling_params import SamplingParams, SamplingType
@@ -41,10 +43,10 @@ def get_engine_config(dict_config, enable_check = True):
         # assert engine_config.max_num_sequences > 0
         # assert engine_config.max_num_sequences * engine_config.max_input_len == engine_config.max_num_batched_tokens
 
-        assert (engine_config.min_decode_steps > 0) and (engine_config.max_decode_steps > 0) 
+        assert (engine_config.min_decode_steps > 0) and (engine_config.max_decode_steps > 0)
         assert engine_config.max_decode_steps > engine_config.min_decode_steps
         assert engine_config.prompt_allocate_ratio > 0
-    
+
     return engine_config
 
 @dataclass
@@ -75,12 +77,23 @@ class FinishReason(Enum):
     Length = "length"
     Cancelled = "cancelled"
 
+# A single token.
+Token = List[int]
+
+@dataclass
+class ValidationError:
+    msg: str
+
+# The type signature of the token validation callback.
+ValidateTokensCallback = Callable[["Request", List[Token]], ValidationError]
 
 @dataclass
 class Request:
     request_id: RequestId
     messages: list[ChatMessage]
 
+    # Perform request validation post-tokenization, used by the HTTP layer to control validation.
+    validate_tokens: Optional[ValidateTokensCallback] = None
     # Number of sequences to generate
     num_sequences: int = 1
     # TODO: should `best_of` be handled in the serving layer?
@@ -209,6 +222,7 @@ class RequestState:
     stopping_criteria: StoppingCriteria
     debug_options: DebugOptions
     is_ended: bool = False
+    validation_err: Optional[ValidationError] = None
 
 def check_stopping_sequences(stopping_criteria, output_text, delta, is_ended):
     if stopping_criteria.stop_sequences:
@@ -226,4 +240,4 @@ def check_stopping_sequences(stopping_criteria, output_text, delta, is_ended):
                     output_text = output_text[:output_text.find(t) + len(t)]
                 is_ended = True
                 break
-    return output_text, delta, is_ended
+    return output_text, delta, is_ended

serve/mlc_serve/engine/staging_engine.py

@@ -5,7 +5,7 @@
 import multiprocessing
 import queue
 from threading import Lock
-from typing import Callable
+from typing import Callable, Optional
 
 from .base import (
     InferenceStepResult,
@@ -84,6 +84,8 @@ def add(self, requests: list[Request]):
             # TODO: verify that request id is unique
             if req.num_sequences > 1:
                 raise RuntimeError("num_sequences > 1 is not supported for now")
+
+            # If the request violates the tokenization, this returns None, so skip.
             state = self._get_new_request_state(req)
             new_request_states.append(state)
 
@@ -192,14 +194,18 @@ def step(self) -> InferenceStepResult:
     def _is_ready_to_serve(self) -> bool:
         return self.worker_process is not None and self.worker_process.is_alive()
 
-    def _get_new_request_state(self, request: Request) -> RequestState:
+    def _get_new_request_state(self, request: Request) -> Optional[RequestState]:
         if request.debug_options.prompt is not None:
             prompt = request.debug_options.prompt
         else:
             prompt = self.conversation_template.apply(request.messages)
 
         prompt_tokens = self.tokenizer.encode(prompt)
 
+        validation_err = None
+        if request.validate_tokens is not None:
+            validation_err = request.validate_tokens(request, prompt_tokens)
+
         return RequestState(
             request_id=request.request_id,
             token_ids=prompt_tokens,
@@ -209,6 +215,7 @@ def _get_new_request_state(self, request: Request) -> RequestState:
             stopping_criteria=request.stopping_criteria,
             debug_options=request.debug_options,
             output_text="",
+            validation_err=validation_err,
         )
 
     def _decode_last_output(self, state: RequestState) -> str:

serve/mlc_serve/engine/staging_engine_worker.py

@@ -7,13 +7,15 @@
 from collections import deque
 from dataclasses import dataclass
 from threading import Condition, Lock, Thread
-from typing import Callable, Optional, Union
+from typing import Callable, Optional, Union, Any
 
 from .base import FinishReason, RequestId, RequestState
 from .model_module import DecodeRequest, ModelModule, PrefillRequest, SequenceId
+import structlog
 
 logger = logging.getLogger(__name__)
 
+LOG = structlog.stdlib.get_logger(__name__)
 
 @dataclass
 class ShutdownCommand:
@@ -77,7 +79,16 @@ def __init__(
 
     def add(self, request_states: list[RequestState]):
         with self.queue_lock:
-            self.queue.extend(request_states)
+            # States which have been invalidated should never be added, directly
+            # cancel them instead.
+            valid_states = []
+            for request_state in request_states:
+                if request_state.validation_err is not None:
+                    self.cancelled_requests.append(request_state)
+                else:
+                    valid_states.append(request_state)
+            LOG.info("valid_states", valid_states=valid_states)
+            self.queue.extend(valid_states)
             self.has_new_requests.notify_all()
 
     def cancel(self, request_id: RequestId):
@@ -102,7 +113,7 @@ def wait_for_request(self, timeout_seconds=None) -> bool:
             )
 
     def has_pending_requests(self) -> bool:
-        return self.queue or self.current_batch
+        return self.queue or self.current_batch or self.cancelled_requests
 
     def step(self) -> GenerationLoopWorkerOutput:
         logger.debug("Starting new inference step.")
@@ -130,12 +141,17 @@ def step(self) -> GenerationLoopWorkerOutput:
                 self._remove_request_from_batch(state.request_id)
 
         for state in self.cancelled_requests:
+            err = None
+            if state.validation_err:
+                err = state.validation_err
+
             outputs.append(
                 SequenceGenerationOutput(
                     # TODO: support multi-sequence
                     id=SequenceId(state.request_id, 0),
                     new_tokens=[],
                     finish_reason=FinishReason.Cancelled,
+                    error = err
                 )
             )
             if state.request_id in self.current_batch: