import asyncio import json import random from random import randint from typing import Any, AsyncGenerator, Dict, Optional, Union from fastapi import FastAPI, HTTPException, Request from starlette.responses import JSONResponse, StreamingResponse from ray.llm._internal.common.utils.cloud_utils import LoraMirrorConfig from ray.llm._internal.serve.core.configs.llm_config import ( DiskMultiplexConfig, LLMConfig, ) from ray.llm._internal.serve.core.configs.openai_api_models import ( ChatCompletionRequest, ChatCompletionResponse, CompletionRequest, CompletionResponse, DetokenizeRequest, DetokenizeResponse, EmbeddingRequest, EmbeddingResponse, ErrorResponse, ScoreRequest, ScoreResponse, TokenizeRequest, TokenizeResponse, TranscriptionRequest, TranscriptionResponse, ) from ray.llm._internal.serve.core.engine.protocol import LLMEngine from ray.llm._internal.serve.core.protocol import RawRequestInfo from ray.llm._internal.serve.engines.vllm.kv_transfer.base import ( DefaultConnectorBackend, ) from ray.llm._internal.serve.utils.lora_serve_utils import LoraModelLoader from ray.serve.context import ( _get_internal_replica_context, _get_serve_request_context, ) class MockVLLMEngine(LLMEngine): """Mock vLLM Engine that generates fake text responses. - In case of LoRA it generates a prefix with the model name in the text part of the response. """ def __init__(self, llm_config: LLMConfig): """Create a mock vLLM Engine. Args: llm_config: The llm configuration for this engine """ self.llm_config = llm_config # The mock skips engine init, where setup_engine_backend attaches this. if llm_config.engine_kwargs.get("kv_transfer_config"): llm_config._kv_connector_backend = DefaultConnectorBackend(llm_config) self.started = False self._current_lora_model: Dict[str, DiskMultiplexConfig] = {} self._is_sleeping = False self._is_paused = False async def start(self): """Start the mock engine.""" self.started = True def routing_stats(self) -> Dict[str, Any]: """Mock engine advertises no routing stats (no KV-cache events).""" return {} async def resolve_lora(self, lora_model: DiskMultiplexConfig): """Resolve/load a LoRA model.""" self._current_lora_model[lora_model.model_id] = lora_model async def check_health(self) -> None: """Check the health of the mock engine.""" if not self.started: raise RuntimeError("Engine not started") async def reset_prefix_cache(self) -> None: """Reset the prefix cache of the mock engine.""" if not self.started: raise RuntimeError("Engine not started") async def start_profile(self) -> None: """Start profiling of the mock engine.""" if not self.started: raise RuntimeError("Engine not started") async def stop_profile(self) -> None: """Stop profiling of the mock engine.""" if not self.started: raise RuntimeError("Engine not started") async def sleep(self, **kwargs: Any) -> None: """Put the mock engine to sleep. This mimics vLLM's behavior: resets prefix cache and sets sleeping state. Args: **kwargs: Engine-specific options. """ if not self.started: raise RuntimeError("Engine not started") # vLLM resets prefix cache on sleep await self.reset_prefix_cache() self._is_sleeping = True async def wakeup(self, **kwargs: Any) -> None: """Wake up the mock engine from sleep. Args: **kwargs: Engine-specific options. """ if not self.started: raise RuntimeError("Engine not started") self._is_sleeping = False async def is_sleeping(self) -> bool: """Check if the mock engine is sleeping. Returns: True if the engine is sleeping, False otherwise. """ return self._is_sleeping async def pause(self, **kwargs: Any) -> None: """Pause generation on the mock engine. This mimics vLLM's behavior: halts generation while keeping weights in GPU. Args: **kwargs: Engine-specific options (mode, clear_cache). """ if not self.started: raise RuntimeError("Engine not started") # vLLM optionally clears cache on pause if kwargs.get("clear_cache", True): await self.reset_prefix_cache() self._is_paused = True async def resume(self, **kwargs: Any) -> None: """Resume generation on the mock engine after pause. Args: **kwargs: Engine-specific options. """ if not self.started: raise RuntimeError("Engine not started") self._is_paused = False async def is_paused(self) -> bool: """Check if the mock engine is paused. Returns: True if the engine is paused, False otherwise. """ return self._is_paused async def build_asgi_app(self): """Build a minimal ASGI app for direct-streaming tests.""" app = FastAPI() @app.middleware("http") async def _tag_serving_replica(request: Request, call_next): # Tag each response with the serving replica and the session id it # saw, so direct-streaming tests can assert affinity over HAProxy. response = await call_next(request) ctx = _get_internal_replica_context() if ctx is not None: response.headers["x-replica-id"] = ctx.replica_id.unique_id response.headers[ "x-serve-session-id" ] = _get_serve_request_context().session_id return response def check_model(model: Optional[str]) -> None: if model is not None and model != self.llm_config.model_id: raise HTTPException( status_code=404, detail=f"Could not find model {model}", ) async def to_response(gen): try: first = await gen.__anext__() except StopAsyncIteration: return JSONResponse(content={}) if isinstance(first, ErrorResponse): raise HTTPException( status_code=first.error.code, detail=first.error.message, ) if isinstance(first, str): async def stream(): yield first async for item in gen: if isinstance(item, str): yield item else: yield f"data: {item.model_dump_json()}\n\n" return StreamingResponse(stream(), media_type="text/event-stream") return JSONResponse(content=first.model_dump()) @app.get("/v1/models") async def models(): return { "object": "list", "data": [ { "id": self.llm_config.model_id, "object": "model", "created": 0, "owned_by": "mock", "metadata": {"input_modality": "text"}, } ], } @app.post("/v1/chat/completions") async def chat_completions(request: Request): body = ChatCompletionRequest.model_validate(await request.json()) check_model(body.model) return await to_response(self.chat(body)) @app.post("/v1/completions") async def completions(request: Request): body = CompletionRequest.model_validate(await request.json()) check_model(body.model) return await to_response(self.completions(body)) return app async def chat( self, request: ChatCompletionRequest, raw_request_info: Optional[RawRequestInfo] = None, ) -> AsyncGenerator[Union[str, ChatCompletionResponse, ErrorResponse], None]: """Mock chat completion.""" if not self.started: raise RuntimeError("Engine not started") # Extract prompt text from messages prompt_text = "" if request.messages: for message in request.messages: if hasattr(message, "content") and message.content: prompt_text += str(message.content) + " " max_tokens = getattr(request, "max_tokens", None) or randint(1, 10) # Generate streaming response async for response in self._generate_chat_response( request=request, prompt_text=prompt_text.strip(), max_tokens=max_tokens ): yield response async def completions( self, request: CompletionRequest, raw_request_info: Optional[RawRequestInfo] = None, ) -> AsyncGenerator[Union[str, CompletionResponse, ErrorResponse], None]: """Mock text completion.""" if not self.started: raise RuntimeError("Engine not started") prompt_text = str(request.prompt) if request.prompt else "" max_tokens = getattr(request, "max_tokens", None) or randint(5, 20) # Generate streaming response async for response in self._generate_completion_response( request=request, prompt_text=prompt_text, max_tokens=max_tokens ): yield response async def embeddings( self, request: EmbeddingRequest, raw_request_info: Optional[RawRequestInfo] = None, ) -> AsyncGenerator[Union[str, EmbeddingResponse, ErrorResponse], None]: """Mock embeddings generation.""" if not self.started: raise RuntimeError("Engine not started") # Generate a mock embedding response embedding_data = [] inputs = request.input if isinstance(request.input, list) else [request.input] for i, text in enumerate(inputs): # Generate random embedding vector dimensions = getattr(request, "dimensions", None) or 1536 embedding = [random.uniform(-1, 1) for _ in range(dimensions)] embedding_data.append( {"object": "embedding", "embedding": embedding, "index": i} ) response = EmbeddingResponse( object="list", data=embedding_data, model=request.model or self.llm_config.model_id, usage={ "prompt_tokens": len(str(request.input).split()), "total_tokens": len(str(request.input).split()), }, ) yield response async def transcriptions( self, request: TranscriptionRequest, raw_request_info: Optional[RawRequestInfo] = None, ) -> AsyncGenerator[Union[str, TranscriptionResponse, ErrorResponse], None]: """Mock transcription generation.""" if not self.started: raise RuntimeError("Engine not started") # Extract audio file info language = getattr(request, "language", "en") temperature = getattr(request, "temperature", 0.0) # Generate transcription response async for response in self._generate_transcription_response( request=request, language=language, temperature=temperature ): yield response async def score( self, request: ScoreRequest, raw_request_info: Optional[RawRequestInfo] = None, ) -> AsyncGenerator[Union[str, ScoreResponse, ErrorResponse], None]: """Mock score generation for text pairs.""" if not self.started: raise RuntimeError("Engine not started") # Extract text_1 and text_2 from the request text_1 = getattr(request, "text_1", "") text_2 = getattr(request, "text_2", "") # Convert to lists if they aren't already text_1_list = text_1 if isinstance(text_1, list) else [text_1] text_2_list = text_2 if isinstance(text_2, list) else [text_2] # Generate mock scores for each pair score_data = [] for i, (t1, t2) in enumerate(zip(text_1_list, text_2_list)): # Generate a random score (can be any float value) score = random.uniform(-10.0, 10.0) score_data.append({"object": "score", "score": score, "index": i}) # Create the response response = ScoreResponse( object="list", data=score_data, model=request.model or self.llm_config.model_id, usage={ "prompt_tokens": len(str(text_1).split()) + len(str(text_2).split()), "total_tokens": len(str(text_1).split()) + len(str(text_2).split()), }, ) yield response async def tokenize( self, request: TokenizeRequest, raw_request_info: Optional[RawRequestInfo] = None, ) -> AsyncGenerator[Union[TokenizeResponse, ErrorResponse], None]: """Mock tokenize generation.""" if not self.started: raise RuntimeError("Engine not started") # Get prompt text from the request prompt = getattr(request, "prompt", None) if prompt is None: # For TokenizeChatRequest, messages would be used messages = getattr(request, "messages", []) prompt = " ".join(str(getattr(m, "content", "")) for m in messages if m) # Generate mock token IDs (simple: use character codes) prompt_str = str(prompt) if prompt else "" tokens = [ord(c) for c in prompt_str] # Optionally generate token strings return_token_strs = getattr(request, "return_token_strs", False) token_strs = list(prompt_str) if return_token_strs else None response = TokenizeResponse( count=len(tokens), max_model_len=4096, # Mock max model length tokens=tokens, token_strs=token_strs, ) yield response async def detokenize( self, request: DetokenizeRequest, raw_request_info: Optional[RawRequestInfo] = None, ) -> AsyncGenerator[Union[DetokenizeResponse, ErrorResponse], None]: """Mock detokenize generation.""" if not self.started: raise RuntimeError("Engine not started") # Get tokens from the request tokens = getattr(request, "tokens", []) # Convert token IDs back to characters (inverse of our mock tokenize) prompt = "".join(chr(t) if 0 <= t < 0x110000 else "?" for t in tokens) response = DetokenizeResponse(prompt=prompt) yield response def _maybe_attach_kv_transfer_params(self, request, response) -> None: """Stamp the serving replica id into ``kv_transfer_params`` for P/D tests. The orchestrator sends the prefill request with ``remote_engine_id`` unset; fill it with this replica's id so the response reports the prefill replica. On the decode request the id is already set and passes through. Lets tests observe that the session id pinned the prefill replica, not just the decode ingress. """ params = getattr(request, "kv_transfer_params", None) if not params: return params = dict(params) if params.get("remote_engine_id") is None: ctx = _get_internal_replica_context() if ctx is not None: params["remote_engine_id"] = ctx.replica_id.unique_id response.kv_transfer_params = params async def _generate_chat_response( self, request: ChatCompletionRequest, prompt_text: str, max_tokens: int ) -> AsyncGenerator[Union[str, ChatCompletionResponse], None]: """Generate mock chat completion response.""" request_id = request.request_id or f"chatcmpl-{random.randint(1000, 9999)}" # # Use request.model if provided, otherwise fall back to llm_config.model_id model_name = request.model or self.llm_config.model_id lora_prefix = ( "" if request.model not in self._current_lora_model else f"[lora_model] {request.model}: " ) if request.stream: # Streaming response - return SSE formatted strings created_time = int(asyncio.get_event_loop().time()) for i in range(max_tokens): if i == 0: token = f"{lora_prefix}test_{i} " else: token = f"test_{i} " if i == max_tokens - 1: # no space for the last token token = f"test_{i}" # Create streaming chunk choice = { "index": 0, "delta": { "content": token, "role": "assistant" if i == 0 else None, }, "finish_reason": "stop" if i == max_tokens - 1 else None, } chunk_data = { "id": request_id, "object": "chat.completion.chunk", "created": created_time, "model": model_name, "choices": [choice], } # Format as SSE yield f"data: {json.dumps(chunk_data)}\n\n" await asyncio.sleep(0.01) # Simulate processing time # Send final [DONE] message yield "data: [DONE]\n\n" else: # Non-streaming response - return response object generated_text = " ".join([f"test_{i}" for i in range(max_tokens)]) generated_text = f"{lora_prefix}{generated_text}" choice = { "index": 0, "message": {"role": "assistant", "content": generated_text}, "finish_reason": "stop", } response = ChatCompletionResponse( id=request_id, object="chat.completion", created=int(asyncio.get_event_loop().time()), model=model_name, choices=[choice], usage={ "prompt_tokens": len(prompt_text.split()), "completion_tokens": max_tokens, "total_tokens": len(prompt_text.split()) + max_tokens, }, ) self._maybe_attach_kv_transfer_params(request, response) yield response async def _generate_completion_response( self, request: CompletionRequest, prompt_text: str, max_tokens: int ) -> AsyncGenerator[Union[str, CompletionResponse], None]: """Generate mock completion response.""" request_id = request.request_id or f"cmpl-{random.randint(1000, 9999)}" model_name = request.model or self.llm_config.model_id lora_prefix = ( "" if request.model not in self._current_lora_model else f"[lora_model] {request.model}: " ) if request.stream: # Streaming response - return SSE formatted strings created_time = int(asyncio.get_event_loop().time()) for i in range(max_tokens): if i == 0: token = f"{lora_prefix}test_{i} " else: token = f"test_{i} " if i == max_tokens - 1: # no space for the last token token = f"test_{i}" choice = { "index": 0, "text": token, "finish_reason": "stop" if i == max_tokens - 1 else None, } chunk_data = { "id": request_id, "object": "text_completion", "created": created_time, "model": model_name, "choices": [choice], } # Format as SSE yield f"data: {json.dumps(chunk_data)}\n\n" await asyncio.sleep(0.01) # Send final [DONE] message yield "data: [DONE]\n\n" else: # Non-streaming response - return response object generated_text = " ".join([f"test_{i}" for i in range(max_tokens)]) generated_text = f"{lora_prefix}{generated_text}" choice = {"index": 0, "text": generated_text, "finish_reason": "stop"} response = CompletionResponse( id=request_id, object="text_completion", created=int(asyncio.get_event_loop().time()), model=model_name, choices=[choice], usage={ "prompt_tokens": len(prompt_text.split()), "completion_tokens": max_tokens, "total_tokens": len(prompt_text.split()) + max_tokens, }, ) self._maybe_attach_kv_transfer_params(request, response) yield response async def _generate_transcription_response( self, request: TranscriptionRequest, language: str, temperature: float, ) -> AsyncGenerator[Union[str, TranscriptionResponse], None]: """Generate mock transcription response.""" request_id = request.request_id or f"transcribe-{random.randint(1000, 9999)}" lora_prefix = ( "" if request.model not in self._current_lora_model else f"[lora_model] {request.model}: " ) # Generate mock transcription text with LoRA prefix mock_transcription_text = ( f"Mock transcription in {language} language with temperature {temperature}" ) if lora_prefix: mock_transcription_text = f"{lora_prefix}{mock_transcription_text}" model_name = request.model or self.llm_config.model_id if request.stream: # Streaming response - return SSE formatted strings created_time = int(asyncio.get_event_loop().time()) # Split transcription into words for streaming words = mock_transcription_text.split() for i, word in enumerate(words): # Create streaming chunk choice = { "delta": { "content": word + (" " if i < len(words) - 1 else ""), }, } chunk_data = { "delta": None, "type": None, "logprobs": None, "id": request_id, "object": "transcription.chunk", "created": created_time, "model": model_name, "choices": [choice], } # Format as SSE yield f"data: {json.dumps(chunk_data)}\n\n" await asyncio.sleep(0.01) # Simulate processing time # Send final chunk with finish_reason final_choice = { "delta": { "content": "", "finish_reason": "stop", "stop_reason": None, }, } final_chunk_data = { "delta": None, "type": None, "logprobs": None, "id": request_id, "object": "transcription.chunk", "created": created_time, "model": model_name, "choices": [final_choice], } yield f"data: {json.dumps(final_chunk_data)}\n\n" # Send final [DONE] message yield "data: [DONE]\n\n" else: # Non-streaming response - return response object response = TranscriptionResponse( text=mock_transcription_text, logprobs=None, usage={ "seconds": 5.0, "type": "duration", }, ) yield response class MockAsyncLLM: """Mock vLLM's AsyncLLM: ``generate`` replays a fixed list of ``RequestOutput``s, with ``error_after`` raising mid-stream.""" def __init__(self, script, error_after=None): self.script = script self.error_after = error_after async def generate(self, prompt, sampling_params, request_id, **kwargs): for i, output in enumerate(self.script): if self.error_after is not None and i == self.error_after: raise RuntimeError("engine failure") yield output class FakeLoraModelLoader(LoraModelLoader): """Fake LoRA model loader for testing that bypasses S3 entirely.""" async def load_model_from_config( self, lora_model_id: str, llm_config ) -> DiskMultiplexConfig: """Load a fake LoRA model without any S3 access.""" return DiskMultiplexConfig( model_id=lora_model_id, max_total_tokens=llm_config.max_request_context_length, local_path="/fake/local/path", lora_assigned_int_id=random.randint(1, 100), ) async def load_model( self, lora_model_id: str, lora_mirror_config: LoraMirrorConfig ) -> DiskMultiplexConfig: """Load a fake LoRA model.""" return DiskMultiplexConfig( model_id=lora_model_id, max_total_tokens=lora_mirror_config.max_total_tokens, local_path="/fake/local/path", lora_assigned_int_id=random.randint(1, 100), ) class PGCreationMockEngine(MockVLLMEngine): """ A wrapper around the mock engine that forces it to create the placement group on startup, simulating the real vLLM initialization sequence. """ def __init__(self, llm_config, *args, **kwargs): super().__init__(llm_config, *args, **kwargs) self.engine_config = llm_config.get_engine_config() self.engine_config.get_or_create_pg()