Updates for pydantic export with polymorphism and general cleanup

Author: markurtz

src/guidellm/backend/__init__.py

@@ -2,7 +2,7 @@
     Backend,
     BackendType,
 )
-from .openai import OpenAIHTTPBackend
+from .openai import CHAT_COMPLETIONS_PATH, TEXT_COMPLETIONS_PATH, OpenAIHTTPBackend
 from .response import (
     RequestArgs,
     ResponseSummary,
@@ -18,4 +18,6 @@
     "Backend",
     "BackendType",
     "OpenAIHTTPBackend",
+    "TEXT_COMPLETIONS_PATH",
+    "CHAT_COMPLETIONS_PATH",
 ]

src/guidellm/backend/openai.py

@@ -16,7 +16,7 @@
 )
 from guidellm.config import settings
 
-__all__ = ["OpenAIHTTPBackend"]
+__all__ = ["OpenAIHTTPBackend", "TEXT_COMPLETIONS_PATH", "CHAT_COMPLETIONS_PATH"]
 
 
 TEXT_COMPLETIONS_PATH = "/v1/completions"

src/guidellm/backend/response.py

@@ -1,8 +1,9 @@
 from typing import Any, Dict, Literal, Optional
 
-from pydantic import BaseModel, computed_field
+from pydantic import computed_field
 
 from guidellm.config import settings
+from guidellm.objects.pydantic import StandardBaseModel
 
 __all__ = [
     "StreamingResponseType",
@@ -15,7 +16,7 @@
 StreamingResponseType = Literal["start", "iter"]
 
 
-class StreamingTextResponse(BaseModel):
+class StreamingTextResponse(StandardBaseModel):
     """
     A model representing the response content for a streaming text request.
 
@@ -40,7 +41,7 @@ class StreamingTextResponse(BaseModel):
     request_id: Optional[str] = None
 
 
-class RequestArgs(BaseModel):
+class RequestArgs(StandardBaseModel):
     """
     A model representing the arguments for a request to a backend.
     Biases towards an HTTP request, but can be used for other types of backends.
@@ -60,7 +61,7 @@ class RequestArgs(BaseModel):
     http2: Optional[bool] = None
 
 
-class ResponseSummary(BaseModel):
+class ResponseSummary(StandardBaseModel):
     """
     A model representing a summary of a backend request.
     Always returned as the final iteration of a streaming request.

src/guidellm/benchmark/aggregator.py

@@ -13,7 +13,7 @@
     Union,
 )
 
-from pydantic import BaseModel, Field
+from pydantic import Field
 
 from guidellm.backend import ResponseSummary
 from guidellm.benchmark.benchmark import (
@@ -24,15 +24,33 @@
     GenerativeTextErrorStats,
     GenerativeTextResponseStats,
 )
-from guidellm.benchmark.profile import Profile
+from guidellm.benchmark.profile import (
+    AsyncProfile,
+    ConcurrentProfile,
+    Profile,
+    SweepProfile,
+    SynchronousProfile,
+    ThroughputProfile,
+)
 from guidellm.config import settings
-from guidellm.objects import RunningStats, Serializable, TimeRunningStats
-from guidellm.request import GenerationRequest
+from guidellm.objects import RunningStats, StandardBaseModel, TimeRunningStats
+from guidellm.request import (
+    GenerationRequest,
+    GenerativeRequestLoaderDescription,
+    RequestLoaderDescription,
+)
 from guidellm.scheduler import (
     REQ,
     RES,
+    AsyncConstantStrategy,
+    AsyncPoissonStrategy,
+    ConcurrentStrategy,
+    GenerativeRequestsWorkerDescription,
     SchedulerRequestResult,
     SchedulingStrategy,
+    SynchronousStrategy,
+    ThroughputStrategy,
+    WorkerDescription,
 )
 from guidellm.utils import check_load_processor
 
@@ -43,7 +61,7 @@
 ]
 
 
-class BenchmarkAggregator(ABC, BaseModel, Generic[BENCH, REQ, RES]):
+class BenchmarkAggregator(ABC, StandardBaseModel, Generic[BENCH, REQ, RES]):
     """
     A pydantic base class representing the base class for aggregating benchmark results.
     The purpose is to receive and process results from a Benchmarker as it iterates
@@ -55,25 +73,43 @@ class BenchmarkAggregator(ABC, BaseModel, Generic[BENCH, REQ, RES]):
     fully calculated.
     """
 
+    type_: Literal["benchmark_aggregator"] = "benchmark_aggregator"
     run_id: str = Field(
         description=(
             "The unique identifier for the encompasing benchmark run that this "
             "benchmark was a part of."
         )
     )
-    profile: Profile = Field(
+    profile: Union[
+        AsyncProfile,
+        SweepProfile,
+        ConcurrentProfile,
+        ThroughputProfile,
+        SynchronousProfile,
+        Profile,
+    ] = Field(
         description=(
             "The profile used for the entire benchamrk run that the strategy for "
             "the active benchmark was pulled from."
-        )
+        ),
+        discriminator="type_",
     )
     strategy_index: int = Field(
         description=(
             "The index of the strategy in the profile that was used for this benchmark."
         )
     )
-    strategy: SchedulingStrategy = Field(
-        description="The scheduling strategy used to run this benchmark. "
+    strategy: Union[
+        ConcurrentStrategy,
+        SchedulingStrategy,
+        ThroughputStrategy,
+        SynchronousStrategy,
+        AsyncPoissonStrategy,
+        AsyncConstantStrategy,
+        SchedulingStrategy,
+    ] = Field(
+        description="The scheduling strategy used to run this benchmark. ",
+        discriminator="type_",
     )
     max_number: Optional[int] = Field(
         description="The maximum number of requests to run for this benchmark, if any."
@@ -105,25 +141,31 @@ class BenchmarkAggregator(ABC, BaseModel, Generic[BENCH, REQ, RES]):
             "if any. These are requests that were not included in the final results."
         )
     )
-    worker_description: Optional[Serializable] = Field(
+    worker_description: Optional[
+        Union[GenerativeRequestsWorkerDescription, WorkerDescription]
+    ] = Field(
         description=(
             "The description and specifics for the worker used to resolve requests "
             "for this benchmark."
-        )
+        ),
+        discriminator="type_",
     )
-    request_loader_description: Optional[Serializable] = Field(
+    request_loader_description: Optional[
+        Union[GenerativeRequestLoaderDescription, RequestLoaderDescription]
+    ] = Field(
         description=(
             "The description and specifics for the request loader used to create "
             "requests for this benchmark."
-        )
+        ),
+        discriminator="type_",
     )
     extras: Dict[str, Any] = Field(
         description=(
             "Any additional information or metadata that was passed for this benchmark."
         )
     )
 
-    results: List[SchedulerRequestResult[GenerationRequest, ResponseSummary]] = Field(
+    results: List[SchedulerRequestResult[REQ, RES]] = Field(
         default_factory=list,
         description=(
             "The list of all results from the benchmark (complete, incomplete, error), "
@@ -423,6 +465,9 @@ def compile(self) -> BENCH:
 class GenerativeBenchmarkAggregator(
     BenchmarkAggregator[GenerativeBenchmark, GenerationRequest, ResponseSummary]
 ):
+    type_: Literal["generative_benchmark_aggregator"] = (
+        "generative_benchmark_aggregator"
+    )
     processor: Optional[Union[str, Path, Any]] = Field(
         description=(
             "The tokenizer to use for calculating token counts when none are "

src/guidellm/benchmark/benchmark.py

@@ -6,7 +6,7 @@
 
 from guidellm.benchmark.profile import Profile
 from guidellm.objects import (
-    Serializable,
+    StandardBaseModel,
     StatusDistributionSummary,
 )
 from guidellm.scheduler import SchedulerRequestInfo, SchedulingStrategy
@@ -22,7 +22,7 @@
 ]
 
 
-class BenchmarkArgs(Serializable):
+class BenchmarkArgs(StandardBaseModel):
     """
     A serializable model representing the arguments used to specify a benchmark run
     and how data was collected for it.
@@ -74,7 +74,7 @@ class BenchmarkArgs(Serializable):
     )
 
 
-class BenchmarkRunStats(Serializable):
+class BenchmarkRunStats(StandardBaseModel):
     """
     A serializable model representing the run process statistics for the
     entire benchmark run across all requests including warmup and cooldown.
@@ -196,7 +196,7 @@ def total(self) -> int:
         return self.total_successful + self.total_incomplete + self.total_errored
 
 
-class Benchmark(Serializable):
+class Benchmark(StandardBaseModel):
     """
     The base serializable model representing a benchmark run and its results.
     Specific benchmarker implementations should extend this model to include
@@ -228,13 +228,13 @@ class Benchmark(Serializable):
             "The process statistics for the entire benchmark run across all requests."
         )
     )
-    worker: Optional[Serializable] = Field(
+    worker: Optional[StandardBaseModel] = Field(
         description=(
             "The description and specifics for the worker used to resolve requests "
             "for this benchmark."
         )
     )
-    request_loader: Optional[Serializable] = Field(
+    request_loader: Optional[StandardBaseModel] = Field(
         description=(
             "The description and specifics for the request loader used to create "
             "requests for this benchmark."
@@ -257,7 +257,7 @@ class Benchmark(Serializable):
 BENCH = TypeVar("BENCH", bound=Benchmark)
 
 
-class GenerativeTextResponseStats(Serializable):
+class GenerativeTextResponseStats(StandardBaseModel):
     """
     A serializable model representing the request values, response values, and
     statistics for a generative text response.
@@ -660,8 +660,8 @@ def from_stats(
         errored: List[GenerativeTextErrorStats],
         args: BenchmarkArgs,
         run_stats: BenchmarkRunStats,
-        worker: Optional[Serializable],
-        requests_loader: Optional[Serializable],
+        worker: Optional[StandardBaseModel],
+        requests_loader: Optional[StandardBaseModel],
         extras: Optional[Dict[str, Any]],
     ) -> "GenerativeBenchmark":
         """

src/guidellm/benchmark/benchmarker.py

@@ -20,7 +20,7 @@
 from guidellm.benchmark.aggregator import AGG, BENCH, GenerativeBenchmarkAggregator
 from guidellm.benchmark.benchmark import GenerativeBenchmark
 from guidellm.benchmark.profile import Profile
-from guidellm.objects import Serializable
+from guidellm.objects import StandardBaseModel
 from guidellm.request import GenerationRequest
 from guidellm.scheduler import (
     REQ,
@@ -35,7 +35,7 @@
 __all__ = ["Benchmarker", "BenchmarkerResult", "GenerativeBenchmarker"]
 
 
-class BenchmarkerResult(Serializable, Generic[AGG, BENCH, REQ, RES]):
+class BenchmarkerResult(StandardBaseModel, Generic[AGG, BENCH, REQ, RES]):
     type_: Literal[
         "run_start",
         "run_complete",
@@ -54,7 +54,7 @@ class BenchmarkerResult(Serializable, Generic[AGG, BENCH, REQ, RES]):
     current_result: Optional[SchedulerRequestResult[REQ, RES]] = None
 
 
-class BenchmarkerStrategyLimits(Serializable):
+class BenchmarkerStrategyLimits(StandardBaseModel):
     requests_loader_size: Optional[int] = Field(
         description="Size of the request loader.",
     )
@@ -125,7 +125,7 @@ def __init__(
         self,
         worker: RequestsWorker[REQ, RES],
         request_loader: Iterable[REQ],
-        requests_loader_description: Optional[Serializable] = None,
+        requests_loader_description: Optional[StandardBaseModel] = None,
         benchmark_save_extras: Optional[Dict[str, Any]] = None,
     ):
         self.worker = worker
@@ -291,7 +291,7 @@ def __init__(
         self,
         backend: Backend,
         request_loader: Iterable[GenerationRequest],
-        request_loader_description: Optional[Serializable] = None,
+        request_loader_description: Optional[StandardBaseModel] = None,
         benchmark_save_extras: Optional[Dict[str, Any]] = None,
         processor: Optional[Union[str, Path, PreTrainedTokenizer]] = None,
         processor_args: Optional[Dict[str, Any]] = None,