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
3 changes: 1 addition & 2 deletions src/guidellm/backends/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@

from guidellm.extras.vllm import HAS_VLLM

from .backend import Backend, BackendArgs, BackendType
from .backend import Backend, BackendArgs
from .openai import (
AudioRequestHandler,
ChatCompletionsRequestHandler,
Expand All @@ -32,7 +32,6 @@
"AudioRequestHandler",
"Backend",
"BackendArgs",
"BackendType",
"ChatCompletionsRequestHandler",
"OpenAIHTTPBackend",
"OpenAIRequestHandler",
Expand Down
102 changes: 52 additions & 50 deletions src/guidellm/backends/backend.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,30 +8,62 @@

from __future__ import annotations

from abc import abstractmethod
from typing import Literal
from abc import ABC, abstractmethod
from typing import ClassVar

from pydantic import BaseModel, ConfigDict
from pydantic import ConfigDict, Field

from guidellm.scheduler import BackendInterface
from guidellm.schemas import GenerationRequest, GenerationResponse
from guidellm.schemas import (
GenerationRequest,
GenerationResponse,
PydanticClassRegistryMixin,
)
from guidellm.utils.registry import RegistryMixin

__all__ = [
"Backend",
"BackendArgs",
"BackendType",
]


BackendType = Literal["openai_http", "vllm_python"]
class BackendArgs(PydanticClassRegistryMixin["BackendArgs"], ABC):
"""
Base class for backend creation arguments.

This class serves as a base for defining argument models used in the creation
of backend instances. It inherits from PydanticClassRegistryMixin to enable
automatic registration of subclasses, allowing for flexible and extensible
backend configurations.

:cvar schema_discriminator: Field name for polymorphic deserialization
"""

model_config = ConfigDict(
extra="forbid",
serialize_by_alias=True,
ser_json_bytes="base64",
val_json_bytes="base64",
)

schema_discriminator: ClassVar[str] = "type"

@classmethod
def __pydantic_schema_base_type__(cls) -> type[BackendArgs]:
"""
Return base type for polymorphic validation hierarchy.

:return: Base BackendArgs class for schema validation
"""
if cls.__name__ == "BackendArgs":
return cls

class BackendArgs(BaseModel):
"""Base class for backend creation argument models."""
return BackendArgs

# Allow for extra fields until we make BackendArgs the sole source of truth
model_config = ConfigDict(extra="allow")
type_: str = Field(
Comment thread
sjmonson marked this conversation as resolved.
alias="type",
description="Type identifier for the backend configuration.",
)


class Backend(
Expand All @@ -57,18 +89,19 @@ class Backend(
::
@Backend.register("my_backend")
class MyBackend(Backend):
def __init__(self, api_key: str):
super().__init__("my_backend")
self.api_key = api_key
def __init__(self, args: MyBackendArgs):
super().__init__(args)
self.api_key = args.api_key

async def process_startup(self):
self.client = MyAPIClient(self.api_key)

backend = Backend.create("my_backend", api_key="secret")
args = MyBackendArgs(api_key="secret")
backend = Backend.create(args)
"""

@classmethod
def create(cls, type_: str, **kwargs) -> Backend:
def create(cls, args: BackendArgs) -> Backend:
"""
Create a backend instance based on the backend type.

Expand All @@ -77,6 +110,7 @@ def create(cls, type_: str, **kwargs) -> Backend:
:return: An instance of a subclass of Backend
:raises ValueError: If the backend type is not registered
"""
type_ = args.type_

backend = cls.get_registered_object(type_)

Expand All @@ -86,34 +120,15 @@ def create(cls, type_: str, **kwargs) -> Backend:
f"Available types: {list(cls.registry.keys()) if cls.registry else []}"
)

return backend(**kwargs)

@classmethod
def get_backend_args(cls, type_: str) -> type[BackendArgs]:
"""
Return the Pydantic model class for the backend's creation arguments.
return backend(args)

:param type_: The backend type identifier
:return: The backend's BackendArgs subclass
:raises ValueError: If the backend type is not registered
"""
backend_class = cls.get_registered_object(type_)

if backend_class is None:
raise ValueError(
f"Backend type '{type_}' is not registered. "
f"Available types: {list(cls.registry.keys()) if cls.registry else []}"
)

return backend_class.backend_args()

def __init__(self, type_: str):
def __init__(self, args: BackendArgs):
"""
Initialize a backend instance.

:param type_: The backend type identifier
"""
self.type_ = type_
self.type_ = args.type_

@property
def processes_limit(self) -> int | None:
Expand All @@ -130,19 +145,6 @@ def requests_limit(self) -> int | None:
"""
return None

@classmethod
@abstractmethod
def backend_args(cls) -> type[BackendArgs]:
"""
Return the Pydantic model class for this backend's creation arguments.

The model defines the parameters (e.g. target, model) that the CLI/benchmark
supply when creating the backend. Used for validation and error messages.

:return: A BackendArgs subclass whose fields are the creation params
"""
...

@abstractmethod
async def default_model(self) -> str:
"""
Expand Down
Loading
Loading