wehub-resource-sync
728 lines
26 KiB
Python
728 lines
26 KiB
Python
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()
|