Fixes for e2e enablement

Author: markurtz

src/guidellm/__main__.py

@@ -3,6 +3,7 @@
     CompilableAggregator,
     GenerativeRequestsAggregator,
     GenerativeStatsProgressAggregator,
+    InjectExtrasAggregator,
     SchedulerStatsAggregator,
     SerializableAggregator,
 )
@@ -62,6 +63,7 @@
     "GenerativeRequestStats",
     "GenerativeRequestsAggregator",
     "GenerativeStatsProgressAggregator",
+    "InjectExtrasAggregator",
     "Profile",
     "ProfileType",
     "SchedulerStatsAggregator",

src/guidellm/benchmark/__init__.py

@@ -33,7 +33,7 @@
     runtime_checkable,
 )
 
-import numpy
+import numpy as np
 from pydantic import Field, PrivateAttr
 
 from guidellm.backend import (
@@ -70,6 +70,7 @@
     "CompilableAggregator",
     "GenerativeRequestsAggregator",
     "GenerativeStatsProgressAggregator",
+    "InjectExtrasAggregator",
     "SchedulerStatsAggregator",
     "SerializableAggregator",
 ]
@@ -284,6 +285,47 @@ def compile(
         """
 
 
+@SerializableAggregator.register("inject_extras")
+class InjectExtrasAggregator(
+    SerializableAggregator[ResponseT, RequestT, MeasuredRequestTimingsT], InfoMixin
+):
+    """
+    Aggregator for injecting extra metadata into the output.
+    """
+
+    @classmethod
+    def validated_kwargs(cls, extras: dict[str, Any], **kwargs) -> dict[str, Any]:
+        return {"extras": extras}
+
+    type_: Literal["inject_extras"] = Field(default="inject_extras")
+    extras: dict[str, Any] | None = Field(default_factory=None)
+
+    def __call__(
+        self,
+        agg_state: dict[str, Any],
+        response: ResponseT | None,
+        request: RequestT,
+        request_info: ScheduledRequestInfo[MeasuredRequestTimingsT],
+        scheduler_state: SchedulerState,
+    ) -> dict[str, Any] | None:
+        """
+        Inject extra metadata into the aggregation state.
+
+        :param agg_state: Current aggregation state to update.
+        :param response: Response generated for the request, if successful.
+        :param request: The processed request object.
+        :param request_info: Scheduling metadata and timing information.
+        :param scheduler_state: Current scheduler execution state.
+        :return: Updated aggregation state with injected extras.
+        """
+        return None
+
+    def compile(
+        self, agg_state: dict[str, Any], scheduler_state: SchedulerState
+    ) -> dict[str, Any]:
+        return {"extras": self.extras} if self.extras else {}
+
+
 @SerializableAggregator.register("scheduler_stats")
 class SchedulerStatsAggregator(
     SerializableAggregator[ResponseT, RequestT, MeasuredRequestTimingsT], InfoMixin
@@ -600,7 +642,7 @@ def __call__(
                 self.add_aggregate_metric_rate(f"{prefix}prompt_tokens", agg_state)
             )
             self.add_aggregate_metric(
-                f"{prefix}prompt_tokens", agg_state, response.prompt_tokens
+                "prompt_tokens", agg_state, response.prompt_tokens
             )
             agg_state["prompt_tokens_per_request"] = self.add_aggregate_metric_rate(
                 "prompt_tokens", agg_state
@@ -842,26 +884,32 @@ def compile(
             "requests": StatusBreakdown(
                 successful=(
                     list(
-                        numpy.random.choice(
-                            successful, size=self.request_samples, replace=False
+                        np.random.choice(
+                            successful,
+                            size=min(self.request_samples, len(successful)),
+                            replace=False,
                         )
                     )
                     if self.request_samples
                     else successful
                 ),
                 incomplete=(
                     list(
-                        numpy.random.choice(
-                            incomplete, size=self.request_samples, replace=False
+                        np.random.choice(
+                            incomplete,
+                            size=min(self.request_samples, len(incomplete)),
+                            replace=False,
                         )
                     )
                     if self.request_samples
                     else incomplete
                 ),
                 errored=(
                     list(
-                        numpy.random.choice(
-                            errored, size=self.request_samples, replace=False
+                        np.random.choice(
+                            errored,
+                            size=min(self.request_samples, len(errored)),
+                            replace=False,
                         )
                     )
                     if self.request_samples