Replace v0.2.1 units validators with a declarative table

Adds UnitsValidatedModel to base.py: a model base whose subclasses declare a
UNIT_RULES table (field name -> allowed units) instead of hand-writing a
check_units method. A single inherited validator merges the rules declared
across the MRO, so each class only states the fields it introduces and shared
fields are validated once.

Migrates the six v0.2.1 unit validators (the MediaPayloadStats hierarchy plus
AggregateRequests/AggregateThroughput) to this mechanism. The two aggregate
classes multiply-inherit UnitsValidatedModel alongside their v0.2 base; the
shared validator is named distinctly from v0.2's check_units so both co-run
rather than one shadowing the other. Behavior is unchanged and covered by the
existing compat and multimodal guardrail tests; the only schema-artifact diff
is an expanded AggregateRequests description.

Signed-off-by: Brendan Slabe <slabe@google.com>

Author: Bslabe123

llmdbenchmark/analysis/benchmark_report/base.py

@@ -4,9 +4,9 @@
 
 import json
 from enum import StrEnum, auto
-from typing import Any
+from typing import Any, ClassVar
 
-from pydantic import BaseModel
+from pydantic import BaseModel, model_validator
 import yaml
 
 ###############################################################################
@@ -161,6 +161,46 @@ class Units(StrEnum):
 UNITS_MEDIA_THROUGHPUT = [Units.IMAGE_PER_S, Units.VIDEO_PER_S, Units.AUDIO_PER_S]
 UNITS_POWER = [Units.WATTS]
 
+
+###############################################################################
+# Declarative units validation
+###############################################################################
+
+
+class UnitsValidatedModel(BaseModel):
+    """Base model that validates ``Statistics`` field units declaratively.
+
+    Instead of hand-writing a ``check_units`` method per class, a subclass sets
+    ``UNIT_RULES``, mapping a field name to the list of units allowed for that
+    field's ``Statistics.units``. A single inherited validator checks every
+    rule declared anywhere in the class's MRO, so each subclass only declares
+    the fields it introduces; ``None`` (unset Optional) fields are skipped.
+
+    The validator is named distinctly from any ``check_units`` so it co-runs
+    with, rather than shadows, a hand-written validator inherited from another
+    base under multiple inheritance.
+    """
+
+    # field name -> allowed Units. Merged across the MRO at validation time.
+    UNIT_RULES: ClassVar[dict[str, list[Units]]] = {}
+
+    @model_validator(mode="after")
+    def validate_declared_units(self):
+        merged: dict[str, list[Units]] = {}
+        # Most-derived first, so a subclass rule overrides a base rule.
+        for klass in type(self).__mro__:
+            for field, allowed in klass.__dict__.get("UNIT_RULES", {}).items():
+                merged.setdefault(field, allowed)
+        for field, allowed in merged.items():
+            stat = getattr(self, field, None)
+            if stat is not None and stat.units not in allowed:
+                raise ValueError(
+                    f'Invalid units "{stat.units}" for "{field}", must be one'
+                    f" of: {' '.join(allowed)}"
+                )
+        return self
+
+
 ###############################################################################
 # Base benchmark report class
 ###############################################################################

llmdbenchmark/analysis/benchmark_report/br_v0_2_1_json_schema.json

@@ -114,7 +114,7 @@
     },
     "AggregateRequests": {
       "additionalProperties": false,
-      "description": "v0.2 request statistics, plus multi-modal payload details.",
+      "description": "v0.2 request statistics, plus multi-modal payload details.\n\nInherits the v0.2 input/output-length unit checks and adds a declarative\nrule for the new request_size field.",
       "properties": {
         "total": {
           "description": "Total number of requests sent.",

llmdbenchmark/analysis/benchmark_report/schema_v0_2_1.py

@@ -17,14 +17,18 @@
 the PR description.
 """
 
-from pydantic import BaseModel, model_validator
+from typing import ClassVar
+
+from pydantic import BaseModel
 
 from .base import (
     UNITS_MEDIA_THROUGHPUT,
     UNITS_MEMORY,
     UNITS_QUANTITY,
     UNITS_RATIO,
     UNITS_TIME,
+    Units,
+    UnitsValidatedModel,
 )
 from .schema_v0_2 import (
     MODEL_CONFIG,
@@ -66,7 +70,7 @@
 ###############################################################################
 
 
-class MediaPayloadStats(BaseModel):
+class MediaPayloadStats(UnitsValidatedModel):
     """Payload statistics shared by every media modality.
 
     All fields are distributions over the individual media instances the client
@@ -75,50 +79,32 @@ class MediaPayloadStats(BaseModel):
 
     model_config = MODEL_CONFIG.copy()
 
+    UNIT_RULES: ClassVar[dict[str, list[Units]]] = {
+        "count": UNITS_QUANTITY,
+        "bytes": UNITS_MEMORY,
+    }
+
     count: Statistics | None = None
     """Number of media instances of this modality per request."""
     bytes: Statistics | None = None
     """Encoded size per media instance."""
 
-    @model_validator(mode="after")
-    def check_media_units(self):
-        if self.count and self.count.units not in UNITS_QUANTITY:
-            raise ValueError(
-                f'Invalid units "{self.count.units}", must be one of:'
-                f" {' '.join(UNITS_QUANTITY)}"
-            )
-        if self.bytes and self.bytes.units not in UNITS_MEMORY:
-            raise ValueError(
-                f'Invalid units "{self.bytes.units}", must be one of:'
-                f" {' '.join(UNITS_MEMORY)}"
-            )
-        return self
-
 
 class VisualPayloadStats(MediaPayloadStats):
     """Payload statistics common to pixel-based modalities (image and video)."""
 
     model_config = MODEL_CONFIG.copy()
 
+    UNIT_RULES: ClassVar[dict[str, list[Units]]] = {
+        "pixels": UNITS_QUANTITY,
+        "aspect_ratio": UNITS_RATIO,
+    }
+
     pixels: Statistics | None = None
     """Pixel count per media instance (height x width, summed over frames)."""
     aspect_ratio: Statistics | None = None
     """Aspect ratio (width / height) per media instance."""
 
-    @model_validator(mode="after")
-    def check_visual_units(self):
-        if self.pixels and self.pixels.units not in UNITS_QUANTITY:
-            raise ValueError(
-                f'Invalid units "{self.pixels.units}", must be one of:'
-                f" {' '.join(UNITS_QUANTITY)}"
-            )
-        if self.aspect_ratio and self.aspect_ratio.units not in UNITS_RATIO:
-            raise ValueError(
-                f'Invalid units "{self.aspect_ratio.units}", must be one of:'
-                f" {' '.join(UNITS_RATIO)}"
-            )
-        return self
-
 
 class ImagePayloadStats(VisualPayloadStats):
     """Image payload statistics."""
@@ -131,36 +117,22 @@ class VideoPayloadStats(VisualPayloadStats):
 
     model_config = MODEL_CONFIG.copy()
 
+    UNIT_RULES: ClassVar[dict[str, list[Units]]] = {"frames": UNITS_QUANTITY}
+
     frames: Statistics | None = None
     """Number of frames per video instance."""
 
-    @model_validator(mode="after")
-    def check_video_units(self):
-        if self.frames and self.frames.units not in UNITS_QUANTITY:
-            raise ValueError(
-                f'Invalid units "{self.frames.units}", must be one of:'
-                f" {' '.join(UNITS_QUANTITY)}"
-            )
-        return self
-
 
 class AudioPayloadStats(MediaPayloadStats):
     """Audio payload statistics."""
 
     model_config = MODEL_CONFIG.copy()
 
+    UNIT_RULES: ClassVar[dict[str, list[Units]]] = {"seconds": UNITS_TIME}
+
     seconds: Statistics | None = None
     """Duration per audio instance."""
 
-    @model_validator(mode="after")
-    def check_audio_units(self):
-        if self.seconds and self.seconds.units not in UNITS_TIME:
-            raise ValueError(
-                f'Invalid units "{self.seconds.units}", must be one of:'
-                f" {' '.join(UNITS_TIME)}"
-            )
-        return self
-
 
 class MultiModalRequests(BaseModel):
     """Per-modality request payload statistics for multi-modal workloads."""
@@ -180,52 +152,41 @@ class MultiModalRequests(BaseModel):
 ###############################################################################
 
 
-class AggregateRequests(AggregateRequestsV02):
-    """v0.2 request statistics, plus multi-modal payload details."""
+class AggregateRequests(AggregateRequestsV02, UnitsValidatedModel):
+    """v0.2 request statistics, plus multi-modal payload details.
+
+    Inherits the v0.2 input/output-length unit checks and adds a declarative
+    rule for the new request_size field.
+    """
 
     model_config = MODEL_CONFIG.copy()
 
+    UNIT_RULES: ClassVar[dict[str, list[Units]]] = {"request_size": UNITS_MEMORY}
+
     request_size: Statistics | None = None
     """Total encoded request size, including all media payloads."""
     multimodal: MultiModalRequests | None = None
     """Per-modality payload statistics."""
 
-    @model_validator(mode="after")
-    def check_request_size_units(self):
-        if self.request_size and self.request_size.units not in UNITS_MEMORY:
-            raise ValueError(
-                f'Invalid units "{self.request_size.units}", must be one of:'
-                f" {' '.join(UNITS_MEMORY)}"
-            )
-        return self
 
-
-class AggregateThroughput(AggregateThroughputV02):
+class AggregateThroughput(AggregateThroughputV02, UnitsValidatedModel):
     """v0.2 throughput metrics, plus per-modality payload rates."""
 
     model_config = MODEL_CONFIG.copy()
 
+    UNIT_RULES: ClassVar[dict[str, list[Units]]] = {
+        "image_rate": UNITS_MEDIA_THROUGHPUT,
+        "video_rate": UNITS_MEDIA_THROUGHPUT,
+        "audio_rate": UNITS_MEDIA_THROUGHPUT,
+    }
+
     image_rate: Statistics | None = None
     """Image delivery rate."""
     video_rate: Statistics | None = None
     """Video delivery rate."""
     audio_rate: Statistics | None = None
     """Audio delivery rate."""
 
-    @model_validator(mode="after")
-    def check_media_rate_units(self):
-        for name, stat in (
-            ("image_rate", self.image_rate),
-            ("video_rate", self.video_rate),
-            ("audio_rate", self.audio_rate),
-        ):
-            if stat and stat.units not in UNITS_MEDIA_THROUGHPUT:
-                raise ValueError(
-                    f'Invalid units "{stat.units}" for {name}, must be one of:'
-                    f" {' '.join(UNITS_MEDIA_THROUGHPUT)}"
-                )
-        return self
-
 
 ###############################################################################
 # Containment shims: re-thread the extended aggregates up to a new report root.