Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
115 changes: 115 additions & 0 deletions src/tests/test_main_router_models.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,115 @@
import json
import logging
from unittest.mock import MagicMock

import pytest

from vllm_router.routers.main_router import show_models
from vllm_router.service_discovery import EndpointInfo, ModelInfo


def _build_request():
request = MagicMock()
request.app.state = MagicMock()
request.app.state.external_provider_registry = None
return request


def _build_endpoint(model_payload):
model_id = model_payload["id"]
return EndpointInfo(
url="http://engine",
model_names=[model_id],
Id="engine-1",
added_timestamp=0.0,
model_label=model_id,
sleep=False,
model_info={model_id: ModelInfo.from_dict(model_payload)},
)


def test_model_info_preserves_extra_fields():
payload = {
"id": "test-model",
"object": "model",
"created": 123,
"owned_by": "vllm",
"root": "models/test-model",
"parent": None,
"max_model_len": 8192,
"permission": [{"id": "perm-1", "allow_sampling": True}],
}

model_info = ModelInfo.from_dict(payload)

assert model_info.extra_fields == {
"max_model_len": 8192,
"permission": [{"id": "perm-1", "allow_sampling": True}],
}
assert model_info.to_dict() == payload


@pytest.mark.asyncio
async def test_show_models_returns_full_backend_model_payload(monkeypatch):
payload = {
"id": "rich-model",
"object": "model",
"created": 1774275193,
"owned_by": "vllm",
"root": "models/rich-model",
"parent": None,
"max_model_len": 262144,
"permission": [
{
"id": "modelperm-1",
"object": "model_permission",
"allow_sampling": True,
"allow_logprobs": True,
}
],
}

monkeypatch.setattr(
"vllm_router.routers.main_router.get_service_discovery",
lambda: MagicMock(get_endpoint_info=lambda: [_build_endpoint(payload)]),
)

response = await show_models(_build_request())

assert response.status_code == 200
body = json.loads(response.body)
assert body["data"] == [payload]


@pytest.mark.asyncio
async def test_show_models_does_not_warn_for_preserved_backend_fields(
monkeypatch, caplog
):
payload = {
"id": "rich-model",
"object": "model",
"created": 1774275193,
"owned_by": "vllm",
"root": "models/rich-model",
"parent": None,
"max_model_len": 262144,
"permission": [
{
"id": "modelperm-1",
"object": "model_permission",
"allow_sampling": True,
"allow_logprobs": True,
}
],
}

monkeypatch.setattr(
"vllm_router.routers.main_router.get_service_discovery",
lambda: MagicMock(get_endpoint_info=lambda: [_build_endpoint(payload)]),
)

with caplog.at_level(logging.WARNING, logger="vllm_router.protocols"):
response = await show_models(_build_request())

assert response.status_code == 200
assert "ignored" not in caplog.text
6 changes: 4 additions & 2 deletions src/vllm_router/protocols.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
import time
from typing import List, Optional
from typing import ClassVar, List, Optional

from pydantic import BaseModel, ConfigDict, Field, model_validator

Expand All @@ -11,11 +11,12 @@
class OpenAIBaseModel(BaseModel):
# OpenAI API does allow extra fields
model_config = ConfigDict(extra="allow")
_warn_on_extra_fields: ClassVar[bool] = True

@model_validator(mode="before")
@classmethod
def __log_extra_fields__(cls, data):
if isinstance(data, dict):
if cls._warn_on_extra_fields and isinstance(data, dict):
# Get all class field names and their potential aliases
field_names = set()
for field_name, field in cls.model_fields.items():
Expand Down Expand Up @@ -43,6 +44,7 @@ class ErrorResponse(OpenAIBaseModel):


class ModelCard(OpenAIBaseModel):
_warn_on_extra_fields: ClassVar[bool] = False
id: str
object: str = "model"
created: int = Field(default_factory=lambda: int(time.time()))
Expand Down
8 changes: 1 addition & 7 deletions src/vllm_router/routers/main_router.py
Original file line number Diff line number Diff line change
Expand Up @@ -152,13 +152,7 @@ async def show_models(request: Request):
if model_id in existing_models:
continue

model_card = ModelCard(
id=model_id,
object="model",
created=model_info.created,
owned_by=model_info.owned_by,
parent=model_info.parent,
)
model_card = ModelCard(**model_info.to_dict())
Comment thread
yzhan1 marked this conversation as resolved.
model_cards.append(model_card)
existing_models.add(model_id)

Expand Down
22 changes: 19 additions & 3 deletions src/vllm_router/service_discovery.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@
import time
import uuid
from dataclasses import dataclass
from typing import Dict, List, Optional, Set
from typing import Any, Dict, List, Optional, Set

import aiohttp
import requests
Expand All @@ -32,6 +32,16 @@
logger = init_logger(__name__)

_global_service_discovery: "Optional[ServiceDiscovery]" = None
_MODEL_INFO_KNOWN_FIELDS = frozenset(
{
"id",
"object",
"created",
"owned_by",
"root",
"parent",
}
)


class ServiceDiscoveryType(enum.Enum):
Expand All @@ -51,6 +61,7 @@ class ModelInfo:
root: Optional[str] = None
parent: Optional[str] = None
is_adapter: bool = False
extra_fields: Dict[str, Any] = None

@classmethod
def from_dict(cls, data: Dict) -> "ModelInfo":
Expand All @@ -63,19 +74,24 @@ def from_dict(cls, data: Dict) -> "ModelInfo":
root=data.get("root", None),
parent=data.get("parent", None),
is_adapter=data.get("parent") is not None,
extra_fields={
k: v for k, v in data.items() if k not in _MODEL_INFO_KNOWN_FIELDS
},
)

def to_dict(self) -> Dict:
"""Convert the ModelInfo instance to a dictionary."""
return {
data = {
"id": self.id,
"object": self.object,
"created": self.created,
"owned_by": self.owned_by,
"root": self.root,
"parent": self.parent,
"is_adapter": self.is_adapter,
}
if self.extra_fields:
data.update(self.extra_fields)
return data


@dataclass
Expand Down
Loading