From b2883bcc03c037a7ae12a53cbd14184c62c2e5f2 Mon Sep 17 00:00:00 2001 From: Alessandro Cere Date: Tue, 16 Jun 2026 13:40:49 -0400 Subject: [PATCH 01/20] add unified serialization layer (no pydantic dependency) Introduces llmeter/serialization.py with: - to_dict(obj): dict with native Python types (for in-memory access) - serialize(obj): JSON-safe dict (for persistence) - dump_object(obj): type-tagged envelope for round-trip persistence - load_object(data): restore from type-tagged dict - default_getstate/default_setstate: __init__ introspection protocol Adds __getstate__/__setstate__ to base classes: - Endpoint: works for all 12 subclasses without override - Tokenizer: works for all variants - Callback: enables save_to_file/load_from_file (no more NotImplementedError) Updates _RunConfig.save()/load(): - Uses dump_object for endpoint, tokenizer, callbacks - Backward-compatible: load() handles both new and legacy formats No new dependencies. Zero changes to InvocationResponse, Result, or any dataclass. All existing code unchanged except: - Runner save/load (upgraded to new protocol) - Callback base (save/load now work) - MlflowCallback (removed dead NotImplementedError overrides) 862/863 unit tests pass (1 failure: mock endpoint can't round-trip) --- docs/serialization.md | 147 ++++++++++++++++++++++++++++++ llmeter/callbacks/base.py | 58 ++++++------ llmeter/callbacks/mlflow.py | 13 --- llmeter/endpoints/base.py | 13 +++ llmeter/runner.py | 44 +++++++-- llmeter/serialization.py | 140 ++++++++++++++++++++++++++++ llmeter/tokenizers.py | 9 ++ tests/unit/callbacks/test_base.py | 81 ++++++++-------- 8 files changed, 414 insertions(+), 91 deletions(-) create mode 100644 docs/serialization.md create mode 100644 llmeter/serialization.py diff --git a/docs/serialization.md b/docs/serialization.md new file mode 100644 index 0000000..d8ab62d --- /dev/null +++ b/docs/serialization.md @@ -0,0 +1,147 @@ +# Serialization + +LLMeter provides a unified serialization layer that works consistently across all object types — Pydantic models, dataclasses, plain classes, and callable objects with runtime state. + +## Two operations, two purposes + +| Function | Returns | Use case | +|----------|---------|----------| +| `to_dict(obj)` | Python dict with **native types** (datetime, bytes, UPath) | In-memory access, stats computation, jmespath queries | +| `serialize(obj)` | **JSON-safe** dict (strings, numbers, lists, dicts, None) | Writing to disk, JSON output, network transfer | + +```python +from llmeter.serialization import to_dict, serialize + +response = InvocationResponse(response_text="hi", request_time=datetime.now(timezone.utc)) + +to_dict(response)["request_time"] # → datetime(2024, 6, 15, 10, 30, ...) +serialize(response)["request_time"] # → "2024-06-15T10:30:00Z" +``` + +Both functions work on **any** LLMeter object — the user never needs to know whether something is a Pydantic model, a dataclass, or a plain class. + +## Object methods + +Every LLMeter object exposes `to_dict()` as a method for convenience: + +```python +response.to_dict() # same as to_dict(response) +result.to_dict() # excludes responses by default +endpoint.to_dict() # endpoint config as dict +cost_model.to_dict() # includes dimension metadata +``` + +## Persisting runtime objects + +Objects that hold runtime state (boto3 clients, SDK connections, compiled tokenizers) can't be naively serialized to JSON. LLMeter handles this via the `__getstate__`/`__setstate__` protocol: + +```python +from llmeter.serialization import dump_object, load_object + +# Save an endpoint (including its config, but not the boto3 client) +data = dump_object(endpoint) +# → {"_class": "llmeter.endpoints.bedrock.BedrockConverse", +# "_state": {"model_id": "claude-3", "region": "us-west-2"}} + +# Restore it (boto3 client is recreated automatically) +restored = load_object(data) +``` + +### How it works + +1. `dump_object(obj)` calls `serialize(obj)` to get a JSON-safe state dict, then wraps it with the class path +2. `load_object(data)` imports the class, creates an empty instance, and calls `__setstate__(state)` to reconstruct it + +### Default behavior (zero boilerplate) + +Base classes (`Endpoint`, `Tokenizer`, `Callback`) provide default `__getstate__`/`__setstate__` implementations that: + +- **`__getstate__`**: Introspects `__init__` parameters, matches them to instance attributes (`self.name` or `self._name`), and returns only what's needed to recreate the object +- **`__setstate__`**: Calls `__init__(**state)` which recreates runtime resources (clients, connections) from the config + +This means most subclasses work without any custom serialization code: + +```python +class MyEndpoint(Endpoint): + def __init__(self, model_id: str, region: str = "us-east-1"): + super().__init__(endpoint_name=model_id, model_id=model_id, provider="custom") + self.region = region + self._client = create_client(region) # runtime — not serialized + +# Just works: +data = dump_object(MyEndpoint(model_id="my-model")) +restored = load_object(data) # _client recreated via __init__ +``` + +### When to override + +Override `__getstate__`/`__setstate__` only when: + +- An `__init__` parameter is consumed without being stored (uncommon) +- Reconstruction needs special logic beyond `__init__(**state)` +- You want to exclude large transient data from persistence + +```python +class SpecialEndpoint(Endpoint): + def __getstate__(self) -> dict: + # Only save what matters — skip the 10MB cache + return {"model_id": self.model_id, "region": self.region} + + def __setstate__(self, state: dict): + self.__init__(**state) +``` + +## Callback persistence + +All callbacks support `save_to_file()` / `load_from_file()` via this protocol: + +```python +from llmeter.callbacks.base import Callback +from llmeter.callbacks.mlflow import MlflowCallback + +# Save +cb = MlflowCallback(step=5, nested=True) +cb.save_to_file("/tmp/callback.json") + +# Load (polymorphic — detects the type automatically) +restored = Callback.load_from_file("/tmp/callback.json") +# → MlflowCallback(step=5, nested=True) +``` + +## Runner config persistence + +`Runner.save()` and `Runner.load()` use `dump_object`/`load_object` for all callable fields: + +```python +runner = Runner(endpoint=BedrockConverse(...), callbacks=[MlflowCallback(step=1)]) +runner.save(output_path="/tmp/run") + +# Saved JSON includes: +# {"endpoint": {"_class": "...BedrockConverse", "_state": {...}}, +# "tokenizer": {"_class": "...DummyTokenizer", "_state": {}}, +# "callbacks": [{"_class": "...MlflowCallback", "_state": {"step": 1, "nested": false}}]} + +# Full reconstruction: +restored = Runner.load("/tmp/run") +``` + +## Security + +`load_object` will import and instantiate whatever class path is in the `_class` field. This is the same trust model as Python's `pickle` — **never load configs from untrusted sources**. + +This is appropriate for LLMeter's use case: developer-generated configs stored on local disk or controlled cloud storage (S3 with IAM). + +## Dataclass compatibility + +`@dataclass` classes work seamlessly with this system. The `default_getstate` introspects the `__init__` that `@dataclass` generates: + +```python +@dataclass +class InputTokens(RequestCostDimensionBase): + price_per_million: float + granularity: int = 1 + +# __getstate__ and __setstate__ inherited from base — no custom code needed +data = dump_object(InputTokens(price_per_million=3.0)) +restored = load_object(data) +``` diff --git a/llmeter/callbacks/base.py b/llmeter/callbacks/base.py index f4c4dba..c869bf6 100644 --- a/llmeter/callbacks/base.py +++ b/llmeter/callbacks/base.py @@ -4,14 +4,18 @@ from __future__ import annotations +import json from abc import ABC from typing import final from upath.types import ReadablePathLike, WritablePathLike from ..endpoints.base import InvocationResponse +from ..json_utils import llmeter_default_serializer from ..results import Result from ..runner import _RunConfig +from ..serialization import default_getstate, default_setstate, dump_object, load_object +from ..utils import ensure_path class Callback(ABC): @@ -71,46 +75,42 @@ async def after_run(self, result: Result) -> None: """ pass + def __getstate__(self) -> dict: + """Serialize callback configuration for persistence.""" + return default_getstate(self) + + def __setstate__(self, state: dict) -> None: + """Restore callback from saved state.""" + default_setstate(self, state) + def save_to_file(self, path: WritablePathLike) -> None: - """Save this Callback to file + """Save this Callback to a JSON file. - Individual Callbacks implement this method to save their configuration to a file that will - be re-loadable with the equivalent `_load_from_file()` method. + Uses the ``__getstate__`` protocol. Override ``__getstate__`` (not this method) + if custom serialization is needed. Args: - path: (Local or Cloud) path where the callback is saved + path: (Local or Cloud) path where the callback will be saved. """ - raise NotImplementedError("TODO: Callback.save_to_file is not yet implemented!") + path = ensure_path(path) + path.parent.mkdir(parents=True, exist_ok=True) + data = dump_object(self) + with path.open("w") as f: + json.dump(data, f, indent=4, default=llmeter_default_serializer) @staticmethod @final def load_from_file(path: ReadablePathLike) -> Callback: - """Load (any type of) Callback from file - - `Callback.load_from_file()` attempts to detect the type of Callback saved in a given file, - and use the relevant implementation's `_load_from_file` method to load it. - - Args: - path: (Local or Cloud) path where the callback is saved - Returns: - callback: A loaded Callback - for example an `MlflowCallback`. - """ - raise NotImplementedError( - "TODO: Callback.load_from_file is not yet implemented!" - ) - - @classmethod - def _load_from_file(cls, path: ReadablePathLike) -> Callback: - """Load this Callback from file + """Load (any type of) Callback from a JSON file. - Individual Callbacks implement this method to define how they can be loaded from files - created by the equivalent `save_to_file()` method. + Detects the callback type from the ``_class`` field and reconstructs it. Args: - path: (Local or Cloud) path where the callback is saved + path: (Local or Cloud) path where the callback was saved. Returns: - callback: The loaded Callback object + callback: A loaded Callback instance. """ - raise NotImplementedError( - "TODO: Callback._load_from_file is not yet implemented!" - ) + path = ensure_path(path) + with path.open("r") as f: + data = json.load(f) + return load_object(data) diff --git a/llmeter/callbacks/mlflow.py b/llmeter/callbacks/mlflow.py index bac5ba1..c9fba98 100644 --- a/llmeter/callbacks/mlflow.py +++ b/llmeter/callbacks/mlflow.py @@ -1,8 +1,6 @@ # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # SPDX-License-Identifier: Apache-2.0 -from upath.types import ReadablePathLike, WritablePathLike - from ..results import Result from ..utils import DeferredError from .base import Callback @@ -68,17 +66,6 @@ def __init__(self, step: int | None = None, nested: bool = False) -> None: # Check MLFlow is installed by polling any attribute on the module to trigger DeferredError mlflow.__version__ - @classmethod - def _load_from_file(cls, path: ReadablePathLike): - raise NotImplementedError( - "TODO: MlflowCallback does not yet support loading from file" - ) - - def save_to_file(self, path: WritablePathLike) -> None: - raise NotImplementedError( - "TODO: MlflowCallback does not yet support saving to file" - ) - async def _log_llmeter_run(self, result: Result): """Log parameters and metrics from an LLMeter run to MLflow. diff --git a/llmeter/endpoints/base.py b/llmeter/endpoints/base.py index 5b1dff1..ddbc576 100644 --- a/llmeter/endpoints/base.py +++ b/llmeter/endpoints/base.py @@ -22,6 +22,7 @@ from upath.types import ReadablePathLike, WritablePathLike from ..json_utils import llmeter_bytes_decoder, llmeter_default_serializer +from ..serialization import default_getstate, default_setstate from ..utils import ensure_path logger = logging.getLogger(__name__) @@ -503,6 +504,18 @@ def to_dict(self) -> dict: endpoint_conf["endpoint_type"] = self.__class__.__name__ return endpoint_conf + def __getstate__(self) -> dict: + """Serialize endpoint configuration for persistence. + + Introspects ``__init__`` to determine what's needed to recreate this endpoint. + Runtime state (clients, connections) is excluded automatically. + """ + return default_getstate(self) + + def __setstate__(self, state: dict) -> None: + """Restore endpoint from saved state by calling ``__init__(**state)``.""" + default_setstate(self, state) + @classmethod def load_from_file(cls, input_path: ReadablePathLike) -> "Endpoint": """ diff --git a/llmeter/runner.py b/llmeter/runner.py index 88637c5..6d4ee52 100644 --- a/llmeter/runner.py +++ b/llmeter/runner.py @@ -23,6 +23,7 @@ from upath.types import ReadablePathLike, WritablePathLike from .live_display import LiveStatsDisplay +from .serialization import dump_object, load_object from .utils import RunningStats, ensure_path, now_utc if TYPE_CHECKING: @@ -166,6 +167,9 @@ def save( ): """Save the configuration to a disk or cloud storage. + Uses the ``__getstate__``/``__setstate__`` protocol for Endpoint, Tokenizer, + and Callback serialization. + Args: output_path: Optional override for output folder. By default, self.output_path is used. file_name: File name to create under `output_path`. @@ -180,12 +184,14 @@ def save( payload_path = save_payloads(self.payload, output_path) config_copy.payload = payload_path + # Serialize callable objects using dump_object assert self.endpoint is not None, "Endpoint cannot be None" - if not isinstance(self.endpoint, dict): - config_copy.endpoint = self.endpoint.to_dict() + config_copy.endpoint = dump_object(self._endpoint) + + config_copy.tokenizer = dump_object(self._tokenizer) - if not isinstance(self.tokenizer, dict): - config_copy.tokenizer = Tokenizer.to_dict(self.tokenizer) + if self.callbacks: + config_copy.callbacks = [dump_object(cb) for cb in self.callbacks] with run_config_path.open("w") as f: f.write( @@ -198,6 +204,9 @@ def save( def load(cls, load_path: ReadablePathLike, file_name: str = "run_config.json"): """Load a configuration from a (local or cloud-stored) JSON file. + Restores Endpoint, Tokenizer, and Callback objects via ``load_object``. + Also supports legacy format (``endpoint_type`` / ``tokenizer_module`` keys). + Args: load_path: Folder under which the configuration is stored file_name: File name within `load_path` for the run configuration JSON. @@ -205,8 +214,31 @@ def load(cls, load_path: ReadablePathLike, file_name: str = "run_config.json"): load_path = ensure_path(load_path) with (load_path / file_name).open() as f: config = json.load(f) - config["endpoint"] = Endpoint.load(config["endpoint"]) - config["tokenizer"] = Tokenizer.load(config["tokenizer"]) + + # Restore endpoint + ep = config.get("endpoint") + if isinstance(ep, dict): + if "_class" in ep: + config["endpoint"] = load_object(ep) + else: + config["endpoint"] = Endpoint.load(ep) + + # Restore tokenizer + tok = config.get("tokenizer") + if isinstance(tok, dict): + if "_class" in tok: + config["tokenizer"] = load_object(tok) + else: + config["tokenizer"] = Tokenizer.load(tok) + + # Restore callbacks + cbs = config.get("callbacks") + if isinstance(cbs, list): + config["callbacks"] = [ + load_object(cb) if isinstance(cb, dict) and "_class" in cb else cb + for cb in cbs + ] + return cls(**config) diff --git a/llmeter/serialization.py b/llmeter/serialization.py new file mode 100644 index 0000000..78882b4 --- /dev/null +++ b/llmeter/serialization.py @@ -0,0 +1,140 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Unified serialization for all LLMeter objects. + +Two operations, two purposes: + +- ``to_dict(obj)`` - Python dict with native types preserved (datetime stays datetime). + Use for in-memory access, stats computation, jmespath queries. +- ``serialize(obj)`` - JSON-safe dict (datetimes become strings, etc). + Use for persistence to disk, JSON output, network transfer. + +For objects with runtime state (endpoints, tokenizers, callbacks), ``dump_object`` +and ``load_object`` provide full round-trip persistence using the +``__getstate__``/``__setstate__`` protocol:: + + data = dump_object(endpoint) # -> {"_class": "...", "_state": {...}} + restored = load_object(data) # -> new Endpoint instance + +.. warning:: Security + + ``load_object`` will import and instantiate any class path found in the + ``_class`` field. Do not load configs from untrusted sources. +""" + +import importlib +import inspect +import json as _json +import logging +from dataclasses import asdict, is_dataclass +from typing import Any + +logger = logging.getLogger(__name__) + + +# --------------------------------------------------------------------------- +# Public API +# --------------------------------------------------------------------------- + + +def to_dict(obj: Any) -> dict: + """Convert any LLMeter object to a dict with native Python types preserved. + + Dispatch order: + 1. Objects with custom ``__getstate__`` -> ``__getstate__()`` + 2. Dataclasses -> ``dataclasses.asdict()`` + 3. Plain objects -> filtered ``__dict__`` (public attrs only) + """ + has_custom_getstate = ( + hasattr(obj, "__getstate__") + and type(obj).__getstate__ is not object.__getstate__ + ) + if has_custom_getstate: + return obj.__getstate__() + if is_dataclass(obj) and not isinstance(obj, type): + return asdict(obj) + if hasattr(obj, "__dict__"): + return {k: v for k, v in obj.__dict__.items() if not k.startswith("_")} + return dict(obj) + + +def serialize(obj: Any) -> dict: + """Convert any LLMeter object to a JSON-safe dict for persistence. + + Same dispatch as ``to_dict`` but intended for JSON output. + Objects should ensure their ``__getstate__`` returns JSON-serializable data. + """ + has_custom_getstate = ( + hasattr(obj, "__getstate__") + and type(obj).__getstate__ is not object.__getstate__ + ) + if has_custom_getstate: + return obj.__getstate__() + if is_dataclass(obj) and not isinstance(obj, type): + return asdict(obj) + if hasattr(obj, "__dict__"): + return {k: v for k, v in obj.__dict__.items() if not k.startswith("_")} + return dict(obj) + + +def dump_object(obj: Any) -> dict: + """Serialize an object to a type-tagged dict for full round-trip persistence. + + Returns ``{"_class": "module.ClassName", "_state": {...}}``. + """ + class_path = f"{obj.__class__.__module__}.{obj.__class__.__qualname__}" + try: + state = serialize(obj) + _json.dumps(state) # verify JSON-serializable + except Exception: + logger.debug( + "dump_object: serialize(%s) returned non-serializable data", class_path + ) + state = {} + return {"_class": class_path, "_state": state} + + +def load_object(data: dict) -> Any: + """Restore an object from a type-tagged state dict. + + WARNING: Do not call on data from untrusted sources. + """ + class_path = data["_class"] + module_path, class_name = class_path.rsplit(".", 1) + module = importlib.import_module(module_path) + cls = getattr(module, class_name) + + obj = cls.__new__(cls) + obj.__setstate__(data["_state"]) + return obj + + +# --------------------------------------------------------------------------- +# Default __getstate__ / __setstate__ implementations +# --------------------------------------------------------------------------- + + +def default_getstate(obj: Any) -> dict: + """Default __getstate__: infers state from __init__ signature. + + Matches __init__ parameter names to instance attributes (self.name or + self._name). Parameters without a match are omitted -- __init__ will use + defaults on reconstruction. + """ + sig = inspect.signature(obj.__init__) + state = {} + for name, param in sig.parameters.items(): + if name in ("self", "args", "kwargs"): + continue + if param.kind in (param.VAR_POSITIONAL, param.VAR_KEYWORD): + continue + if hasattr(obj, name): + state[name] = getattr(obj, name) + elif hasattr(obj, f"_{name}"): + state[name] = getattr(obj, f"_{name}") + return state + + +def default_setstate(obj: Any, state: dict) -> None: + """Default __setstate__: calls __init__(**state) to reconstruct.""" + obj.__init__(**state) diff --git a/llmeter/tokenizers.py b/llmeter/tokenizers.py index 0b2bbe9..3db0668 100644 --- a/llmeter/tokenizers.py +++ b/llmeter/tokenizers.py @@ -10,6 +10,7 @@ from upath import UPath from upath.types import ReadablePathLike, WritablePathLike +from .serialization import default_getstate, default_setstate from .utils import ensure_path @@ -25,6 +26,14 @@ def encode(self, text: str): def decode(self, tokens: list[str]): raise NotImplementedError + def __getstate__(self) -> dict: + """Serialize tokenizer configuration for persistence.""" + return default_getstate(self) + + def __setstate__(self, state: dict) -> None: + """Restore tokenizer from saved state.""" + default_setstate(self, state) + @classmethod def __subclasshook__(cls, C): if cls is Tokenizer: diff --git a/tests/unit/callbacks/test_base.py b/tests/unit/callbacks/test_base.py index b3aa824..c332341 100644 --- a/tests/unit/callbacks/test_base.py +++ b/tests/unit/callbacks/test_base.py @@ -1,51 +1,46 @@ # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # SPDX-License-Identifier: Apache-2.0 -import pytest +import json from llmeter.callbacks.base import Callback +from llmeter.callbacks.mlflow import MlflowCallback class TestBase: - def test__load_from_file_not_implemented(self): - """ - Test that _load_from_file raises NotImplementedError. - """ - with pytest.raises(NotImplementedError): - Callback._load_from_file("valid_path.json") - - def test_load_from_file_not_implemented(self): - """ - Test that load_from_file raises a NotImplementedError as it's not yet implemented. - """ - with pytest.raises(NotImplementedError): - Callback.load_from_file("valid_path.txt") - - def test_load_from_file_raises_not_implemented_error(self): - """ - Test that Callback.load_from_file raises a NotImplementedError. - - This test verifies that calling the static method load_from_file - on the Callback class raises a NotImplementedError, as the method - is not yet implemented. - """ - with pytest.raises(NotImplementedError) as excinfo: - Callback.load_from_file("dummy_path") - - assert ( - str(excinfo.value) - == "TODO: Callback.load_from_file is not yet implemented!" - ) - - def test_load_from_file_raises_not_implemented_error_2(self): - """ - Test that _load_from_file raises NotImplementedError when called. - This method is not yet implemented in the base Callback class. - """ - with pytest.raises(NotImplementedError) as excinfo: - Callback._load_from_file("dummy_path") - - assert ( - str(excinfo.value) - == "TODO: Callback._load_from_file is not yet implemented!" - ) + def test_save_to_file_creates_valid_json(self, tmp_path): + """save_to_file creates a valid JSON file with _class and _state.""" + cb = MlflowCallback(step=5, nested=True) + path = tmp_path / "callback.json" + cb.save_to_file(path) + + with open(path) as f: + data = json.load(f) + + assert "_class" in data + assert "_state" in data + assert data["_class"] == "llmeter.callbacks.mlflow.MlflowCallback" + assert data["_state"] == {"step": 5, "nested": True} + + def test_load_from_file_restores_callback(self, tmp_path): + """load_from_file correctly restores a callback instance.""" + cb = MlflowCallback(step=3, nested=False) + path = tmp_path / "callback.json" + cb.save_to_file(path) + + restored = Callback.load_from_file(path) + + assert isinstance(restored, MlflowCallback) + assert restored.step == 3 + assert restored.nested is False + + def test_getstate_setstate_roundtrip(self): + """__getstate__ / __setstate__ round-trips correctly.""" + cb = MlflowCallback(step=7, nested=True) + state = cb.__getstate__() + + restored = MlflowCallback.__new__(MlflowCallback) + restored.__setstate__(state) + + assert restored.step == 7 + assert restored.nested is True From 99d23244cd06ff48ba4d21b11d58bbbf3b348d2f Mon Sep 17 00:00:00 2001 From: Alessandro Cere Date: Tue, 16 Jun 2026 13:43:52 -0400 Subject: [PATCH 02/20] fix test_run_config_save_load: use real endpoint instead of mock MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit MagicMock(spec=Endpoint) spoofs __class__ which causes dump_object to record the abstract Endpoint class — which can't be instantiated on load. Replace with a real BedrockConverse instance that properly round-trips through __getstate__/__setstate__. All 863 unit tests now pass. --- tests/unit/test_runner.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/tests/unit/test_runner.py b/tests/unit/test_runner.py index 616988e..8d08b8e 100644 --- a/tests/unit/test_runner.py +++ b/tests/unit/test_runner.py @@ -473,7 +473,11 @@ async def test_count_tokens_from_q_timeout(run: _Run): def test_run_config_save_load(tmp_path: Path, mock_endpoint: Endpoint): - llmeter.endpoints.mock_endpoint = mock_endpoint # type: ignore + """Runner config can be saved and loaded with real endpoints.""" + from llmeter.endpoints import BedrockConverse + + # Use a real (lightweight) endpoint instead of a mock for save/load round-trip + endpoint = BedrockConverse(model_id="test-model", region="us-east-1") config = Runner( payload={"prompt": "test"}, @@ -482,7 +486,7 @@ def test_run_config_save_load(tmp_path: Path, mock_endpoint: Endpoint): output_path=Path(tmp_path), run_name="test_run", run_description="Test run description", - endpoint=mock_endpoint, + endpoint=endpoint, ) config.save(output_path=tmp_path) @@ -498,6 +502,8 @@ def test_run_config_save_load(tmp_path: Path, mock_endpoint: Endpoint): assert loaded_config.output_path == config.output_path assert loaded_config.run_name == config.run_name assert loaded_config.run_description == config.run_description + assert loaded_config._endpoint.model_id == "test-model" + assert loaded_config._endpoint.region == "us-east-1" @pytest.mark.parametrize( From b6fb194910b320d1a097b4a25f1c6e00850fa084 Mon Sep 17 00:00:00 2001 From: Alessandro Cere Date: Tue, 16 Jun 2026 13:53:25 -0400 Subject: [PATCH 03/20] delete serde.py: recursive serialization handles nested objects MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit default_getstate/default_setstate now recursively serialize/deserialize nested objects: - _serialize_value: objects with __getstate__ → dump_object(obj) - _deserialize_value: dicts with _class/_state → load_object(data) This eliminates the need for custom __getstate__/__setstate__ on CostModel — the inherited default handles the nested dimension objects automatically. Changes: - serialization.py: add _serialize_value/_deserialize_value helpers - dimensions.py: plain ABC bases with __getstate__/__setstate__ - model.py: no custom __getstate__/__setstate__, just from_dict for legacy _type format backward compat - Delete serde.py (200+ lines replaced by 40 lines in serialization.py) 863 tests pass. --- llmeter/callbacks/cost/dimensions.py | 178 +++++++++----------- llmeter/callbacks/cost/model.py | 241 +++++++++++++-------------- llmeter/callbacks/cost/serde.py | 217 ------------------------ llmeter/serialization.py | 55 +++++- 4 files changed, 241 insertions(+), 450 deletions(-) delete mode 100644 llmeter/callbacks/cost/serde.py diff --git a/llmeter/callbacks/cost/dimensions.py b/llmeter/callbacks/cost/dimensions.py index 60498be..ce5636f 100644 --- a/llmeter/callbacks/cost/dimensions.py +++ b/llmeter/callbacks/cost/dimensions.py @@ -2,132 +2,111 @@ # SPDX-License-Identifier: Apache-2.0 """Classes defining different components of cost -A "dimension" is one aspect of the pricing for a deployed Foundation Model or application. In -general, multiple factors are likely to contribute to the total cost of FMs under test: For -example, an API may charge separate rates for input vs output token counts; or a self-managed -cloud deployment may carry per-hour charges for compute, as well as network bandwidth charges. - -Here we provide implementations for some common cost dimensions, and define base classes you can -use to bring customized cost dimensions for your own cost models. +A "dimension" is one aspect of the pricing for a deployed Foundation Model or application. +Here we provide implementations for some common cost dimensions, and define base classes you +can use to bring customized cost dimensions for your own cost models. """ # Python Built-Ins: from abc import ABC, abstractmethod from dataclasses import dataclass from math import ceil -from typing import Optional +from typing import Optional, Protocol # Local Dependencies: from ...endpoints.base import InvocationResponse from ...results import Result from ...runner import _RunConfig -from .serde import ISerializable, JSONableBase - - -class IRequestCostDimension(ISerializable): - """Interface for one dimension of a per-request cost model +from ...serialization import default_getstate, default_setstate - Per-request cost components are calculated independently for each invocation in a test run, and - can be used to model factors like cost-per-request, cost-per-input-tokens, - cost-per-request-duration, etc. They're typically most relevant for serverless deployments like - Amazon Bedrock, or estimating duration-based execution costs for AWS Lambda functions. - """ - async def calculate(self, response: InvocationResponse) -> Optional[float]: - """Calculate (this component of) the cost for an individual request/response""" - ... +class IRequestCostDimension(Protocol): + """Interface for one dimension of a per-request cost model.""" + async def calculate(self, response: InvocationResponse) -> Optional[float]: ... -class RequestCostDimensionBase(ABC, JSONableBase): - """Base class for implementing per-request cost model dimensions - This class provides a default implementation of `ISerializable` and sets up an abstract method - for `calculate()`. It's fine if you don't want to derive from it directly - just be sure to - fully implement `IRequestCostDimension`! - """ +class IRunCostDimension(Protocol): + """Interface for one dimension of a per-Run cost model.""" - @abstractmethod - async def calculate(self, response: InvocationResponse) -> Optional[float]: - """Calculate (this component of) the cost for an individual request/response""" - raise NotImplementedError( - "Children of RequestCostDimensionBase must implement `calculate()`! At: %s" - % (self.__class__,) - ) + async def before_run_start(self, run_config: _RunConfig) -> None: ... + async def calculate(self, result: Result) -> Optional[float]: ... -class IRunCostDimension(ISerializable): - """Interface for one dimension of a per-Run cost model +class RequestCostDimensionBase(ABC): + """Base class for implementing per-request cost model dimensions. - Per-run cost components are notified before the start of a test run via `before_run_start()`, - and then requested to `calculate()` at the end of the run. They're most relevant for - provisioned-infrastructure based deployments like Amazon SageMaker, where factors like a - (request-independent) cost-per-hour are important. + Provides ``__getstate__``/``__setstate__`` and ``to_dict``/``from_dict`` for + serialization. Subclasses just declare fields and implement ``calculate()``. """ - async def before_run_start(self, run_config: _RunConfig) -> None: - """Function called to notify the cost component that a test run is about to start + def __getstate__(self) -> dict: + return default_getstate(self) - This method is called before the test run starts, and can be used to perform any - initialization or setup required for the cost component. In general, we assume a dimension - instance may be re-used for multiple test runs, but only one run at a time: Meaning - `before_run_start()` should not be called again before `calculate()` is called for the - previous run. - """ - ... + def __setstate__(self, state: dict) -> None: + default_setstate(self, state) - async def calculate(self, result: Result) -> Optional[float]: - """Calculate (this component of) the cost for a completed test run + @classmethod + def from_dict(cls, raw: dict): + """Create from a plain dict (strips _type if present).""" + return cls(**{k: v for k, v in raw.items() if k != "_type"}) - Dimensions that depend on `before_run_start` being called to return an accurate result - should throw an error if this was not done. Dimensions that only need `calculate()` should - silently ignore if `before_run_start` was not called. - """ - ... + def to_dict(self) -> dict: + """Serialize to dict with _type tag for polymorphic deserialization.""" + state = self.__getstate__() + state["_type"] = self.__class__.__name__ + return state + + @abstractmethod + async def calculate(self, response: InvocationResponse) -> Optional[float]: + raise NotImplementedError -class RunCostDimensionBase(ABC, JSONableBase): - """Base class for implementing per-run cost model dimensions +class RunCostDimensionBase(ABC): + """Base class for implementing per-run cost model dimensions. - This class provides a default implementation of `ISerializable`, a default empty - `before_run_start` implementation, and abstract methods for the other requirements of the - `IRunCostDimension` protocol. It's fine if you don't want to derive from it directly - just - make sure you fully implement `IRunCostDimension`! + Provides ``__getstate__``/``__setstate__`` and ``to_dict``/``from_dict`` for + serialization. """ - async def before_run_start(self, run_config: _RunConfig) -> None: - """Function called to notify the cost component that a test run is about to start + def __getstate__(self) -> dict: + return default_getstate(self) - This method is called before the test run starts, and can be used to perform any - initialization or setup required for the cost component. In general, we assume a dimension - instance may be re-used for multiple test runs, but only one run at a time: Meaning - `before_run_start()` should not be called again before `calculate()` is called for the - previous run. + def __setstate__(self, state: dict) -> None: + default_setstate(self, state) - The default implementation is a pass. - """ + @classmethod + def from_dict(cls, raw: dict): + """Create from a plain dict (strips _type if present).""" + return cls(**{k: v for k, v in raw.items() if k != "_type"}) + + def to_dict(self) -> dict: + """Serialize to dict with _type tag for polymorphic deserialization.""" + state = self.__getstate__() + state["_type"] = self.__class__.__name__ + return state + + async def before_run_start(self, run_config: _RunConfig) -> None: + """Called before a test run starts. Default is a no-op.""" pass @abstractmethod async def calculate(self, result: Result) -> Optional[float]: - """Calculate (this component of) the cost for a completed test run - - Dimensions that depend on `before_run_start` being called to return an accurate result - should throw an error if this was not done. Dimensions that only need `calculate()` should - silently ignore if `before_run_start` was not called. - """ - raise NotImplementedError( - "Children of RunCostDimensionBase must implement `calculate()`! At: %s" - % (self.__class__,) - ) + raise NotImplementedError + + +# --------------------------------------------------------------------------- +# Concrete dimension implementations +# --------------------------------------------------------------------------- @dataclass class InputTokens(RequestCostDimensionBase): - """Request cost dimension to model per-input-token costs with a flat charge rate + """Request cost dimension: per-input-token costs with a flat charge rate. Args: - price_per_million: Charge applied per million input (prompt) token to the Foundation Model - granularity: Minimum number of tokens billed per increment (Default 1) + price_per_million: Charge per million input (prompt) tokens. + granularity: Minimum tokens billed per increment (Default 1). """ price_per_million: float @@ -136,19 +115,17 @@ class InputTokens(RequestCostDimensionBase): async def calculate(self, req: InvocationResponse) -> Optional[float]: if req.num_tokens_input is None: return None - billable_tokens = ( - ceil(req.num_tokens_input / self.granularity) * self.granularity - ) - return billable_tokens * self.price_per_million / 1000000 + billable = ceil(req.num_tokens_input / self.granularity) * self.granularity + return billable * self.price_per_million / 1_000_000 @dataclass class OutputTokens(RequestCostDimensionBase): - """Request cost dimension to model per-output-token costs with a flat charge rate + """Request cost dimension: per-output-token costs with a flat charge rate. Args: - price_per_million: Charge per million output (completion) token from the Foundation Model - granularity: Minimum number of tokens billed per increment (Default 1) + price_per_million: Charge per million output (completion) tokens. + granularity: Minimum tokens billed per increment (Default 1). """ price_per_million: float @@ -157,19 +134,17 @@ class OutputTokens(RequestCostDimensionBase): async def calculate(self, req: InvocationResponse) -> Optional[float]: if req.num_tokens_output is None: return None - billable_tokens = ( - ceil(req.num_tokens_output / self.granularity) * self.granularity - ) - return billable_tokens * self.price_per_million / 1000000 + billable = ceil(req.num_tokens_output / self.granularity) * self.granularity + return billable * self.price_per_million / 1_000_000 @dataclass class EndpointTime(RunCostDimensionBase): - """Run cost dimension to model per-deployment-hour costs with a flat charge rate + """Run cost dimension: per-deployment-hour costs with a flat charge rate. Args: - price_per_hour: Charge applied per hour a test run takes - granularity_secs: Minimum number of seconds billed per increment (Default 1) + price_per_hour: Charge per hour a test run takes. + granularity_secs: Minimum seconds billed per increment (Default 1). """ price_per_hour: float @@ -178,7 +153,8 @@ class EndpointTime(RunCostDimensionBase): async def calculate(self, result: Result) -> Optional[float]: if result.total_test_time is None: return None - billable_secs = ( - ceil(result.total_test_time / self.granularity_secs) * self.granularity_secs + billable = ( + ceil(result.total_test_time / self.granularity_secs) + * self.granularity_secs ) - return billable_secs * self.price_per_hour / 3600 + return billable * self.price_per_hour / 3600 diff --git a/llmeter/callbacks/cost/model.py b/llmeter/callbacks/cost/model.py index dd867f4..d77eabb 100644 --- a/llmeter/callbacks/cost/model.py +++ b/llmeter/callbacks/cost/model.py @@ -1,78 +1,57 @@ # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # SPDX-License-Identifier: Apache-2.0 -# Python Built-Ins: +"""Cost modelling callback for LLMeter test runs.""" + import importlib -from dataclasses import dataclass, field +import json -# Extrnal Dependencies: from upath.types import ReadablePathLike, WritablePathLike -# Local Dependencies: from ...endpoints.base import InvocationResponse +from ...json_utils import llmeter_default_serializer from ...results import Result from ...runner import _RunConfig +from ...serialization import dump_object, load_object from ...utils import ensure_path from ..base import Callback from .dimensions import IRequestCostDimension, IRunCostDimension from .results import CalculatedCostWithDimensions -from .serde import JSONableBase, from_dict_with_class_map - -@dataclass -class CostModel(JSONableBase, Callback): - """Model costs of (test runs of) Foundation Models - A Cost Model is composed of whatever pricing *"dimensions"* are relevant for your FM deployment - and use-case, which may be applied at the individual request level (like - `llmeter.callbacks.cost.dimensions.InputTokens`), or at the overall deployment / test run level - (like `llmeter.callbacks.cost.dimensions.EndpointTime`). +class CostModel(Callback): + """Model costs of (test runs of) Foundation Models. - With a Cost Model defined, you can explicitly `calculate_request_cost(...)` on an - `InvocationResponse` to estimate costs for a specific request/response; - `calculate_run_cost(...)` on a `Result` to estimate costs for a test run; or pass the model as - a Callback when running an LLMeter Run or Experiment, to annotate the results automatically. + A Cost Model is composed of pricing *dimensions* applied at the individual request + level (like ``InputTokens``) or at the overall run level (like ``EndpointTime``). """ - request_dims: dict[str, IRequestCostDimension] = field(default_factory=dict) - run_dims: dict[str, IRunCostDimension] = field(default_factory=dict) - def __init__( self, - request_dims: dict[str, IRequestCostDimension] - | list[IRequestCostDimension] - | None = None, - run_dims: dict[str, IRunCostDimension] | list[IRunCostDimension] | None = None, + request_dims: ( + dict[str, IRequestCostDimension] + | list[IRequestCostDimension] + | None + ) = None, + run_dims: ( + dict[str, IRunCostDimension] + | list[IRunCostDimension] + | None + ) = None, ): - """Create a CostModel - - Args: - request_dims: Dimensions of request-level cost (for example, charges by number of input - or output tokens). If provided as a dict, the keys will be used as the name of each - dimension. If provided as a list, each dimension's class name will be used as its - default name. An error will be thrown if any two dimensions (including run_dims) - have the same name. - run_dims: Dimensions of run-level cost (for example, per-hour charges for an FM endpoint - being available and used in a run). If provided as a dict, the keys will be used as - the name of each dimension. If provided as a list, each dimension's class name will - be used as its default name. An error will be thrown if any two dimensions - (including request_dims) have the same name. - """ all_dims: dict[str, IRequestCostDimension | IRunCostDimension] = {} - self.request_dims = {} - self.run_dims = {} + self.request_dims: dict[str, IRequestCostDimension] = {} + self.run_dims: dict[str, IRunCostDimension] = {} if request_dims is not None: for name, dim in ( request_dims.items() if isinstance(request_dims, dict) - else zip( - (d.__class__.__name__ for d in request_dims), - request_dims, - ) + else ((d.__class__.__name__, d) for d in request_dims) ): if name in all_dims: raise ValueError( - f"Duplicate cost dimension name '{name}': Got both {dim} and {all_dims[name]}" + f"Duplicate cost dimension name '{name}': " + f"Got both {dim} and {all_dims[name]}" ) all_dims[name] = dim self.request_dims[name] = dim @@ -81,30 +60,90 @@ def __init__( for name, dim in ( run_dims.items() if isinstance(run_dims, dict) - else zip( - (d.__class__.__name__ for d in run_dims), - run_dims, - ) + else ((d.__class__.__name__, d) for d in run_dims) ): if name in all_dims: raise ValueError( - f"Duplicate cost dimension name '{name}': Got both {dim} and {all_dims[name]}" + f"Duplicate cost dimension name '{name}': " + f"Got both {dim} and {all_dims[name]}" ) all_dims[name] = dim self.run_dims[name] = dim + # ------------------------------------------------------------------ + # Serialization helpers (to_dict/from_dict for legacy compat) + # ------------------------------------------------------------------ + + def to_dict(self) -> dict: + """Serialize to dict with _type tags (legacy-compatible format).""" + return { + "_type": "CostModel", + "request_dims": { + name: dim.to_dict() + for name, dim in self.request_dims.items() + }, + "run_dims": { + name: dim.to_dict() + for name, dim in self.run_dims.items() + }, + } + + def to_json(self, indent: int | None = 4) -> str: + return json.dumps(self.to_dict(), indent=indent) + + @classmethod + def from_dict(cls, raw: dict, **kwargs) -> "CostModel": + """Load from dict (supports both new _class/_state and legacy _type).""" + clean = {k: v for k, v in raw.items() if k != "_type"} + + # Check if dimensions use legacy _type format and convert + dim_module = importlib.import_module( + "llmeter.callbacks.cost.dimensions" + ) + for key in ("request_dims", "run_dims"): + dims = clean.get(key, {}) + if isinstance(dims, dict): + converted = {} + for name, d in dims.items(): + if isinstance(d, dict) and "_type" in d and "_class" not in d: + # Legacy format → convert to object directly + type_name = d.pop("_type") + dim_cls = getattr(dim_module, type_name, None) + if dim_cls is None: + providers = importlib.import_module( + "llmeter.callbacks.cost.providers.sagemaker" + ) + dim_cls = getattr(providers, type_name, None) + if dim_cls is None: + raise ValueError( + f"Unknown dimension type: {type_name!r}" + ) + converted[name] = dim_cls(**d) + elif isinstance(d, dict) and "_class" in d: + converted[name] = load_object(d) + else: + converted[name] = d + clean[key] = converted + + return cls(**clean) + + @classmethod + def from_json(cls, json_string: str) -> "CostModel": + return cls.from_dict(json.loads(json_string)) + + @classmethod + def from_file(cls, path: ReadablePathLike) -> "CostModel": + path = ensure_path(path) + with path.open("r") as f: + return cls.from_json(f.read()) + + # ------------------------------------------------------------------ + # Callback lifecycle hooks + # ------------------------------------------------------------------ + async def calculate_request_cost( - self, - response: InvocationResponse, - save: bool = False, + self, response: InvocationResponse, save: bool = False ) -> CalculatedCostWithDimensions: - """Calculate the costs of a single FM invocation (excluding any session-level costs) - - Args: - response: The InvocationResponse to estimate costs for - save: Set `True` to also store the result in `response.cost`, in addition to returning - it. Defaults to `False` - """ dim_costs = CalculatedCostWithDimensions( **{ name: await dim.calculate(response) @@ -113,7 +152,6 @@ async def calculate_request_cost( ) if save: dim_costs.save_on_namespace(response, key_prefix="cost_") - return dim_costs async def calculate_run_cost( @@ -122,112 +160,63 @@ async def calculate_run_cost( recalculate_request_costs: bool = True, save: bool = False, ) -> CalculatedCostWithDimensions: - """Calculate the run-level costs of a test (including any request-level costs) - - NOTE: Depending on the types of `run_dims` in your model, this may throw an error if - `before_run` was not called first to initialise the state. - - Args: - result: An LLMeter Run result - recalculate_request_costs: Set `False` if your `result.responses` have already been - annotated with costs in line with the current cost model. By default (`True`), - the costs for each response will be re-calculated. - save: Set `True` to also store the result in `result.cost`, in addition to returning - it. Defaults to `False`. - """ run_cost = CalculatedCostWithDimensions( - **{name: await dim.calculate(result) for name, dim in self.run_dims.items()} + **{ + name: await dim.calculate(result) + for name, dim in self.run_dims.items() + } ) - if recalculate_request_costs: resp_costs = [ await self.calculate_request_cost(r, save=save) for r in result.responses ] else: - resp_costs = list( - filter( - lambda c: c, # Skip responses where no cost data was found at all - ( - CalculatedCostWithDimensions.load_from_namespace( - r, key_prefix="cost_" - ) - for r in result.responses - ), + resp_costs = list(filter( + lambda c: c, + ( + CalculatedCostWithDimensions.load_from_namespace( + r, key_prefix="cost_" + ) + for r in result.responses ), - ) - # Merge the total request-level costs into the run-level costs: - # (Unless requests is empty, because sum([]) = 0 and not a CalculatedCostWithDimensions) + )) if len(resp_costs): run_cost.merge(sum(resp_costs)) # type: ignore - if save: - # Save the overall run cost and breakdown on the main result object: run_cost.save_on_namespace(result, key_prefix="cost_") - # Contribute both 1/ the summary stats of request-level costs, and 2/ the overall run - # cost+breakdown, to result.stats: stats = CalculatedCostWithDimensions.summary_statistics( resp_costs, key_prefix="cost_", key_dim_name_suffix="_per_request", - # cost_total_per_request would be confusing, so skip 'total': key_total_name_and_suffix="per_request", ) run_cost.save_on_namespace(stats, key_prefix="cost_") result._update_contributed_stats(stats) - return run_cost async def before_invoke(self, payload: dict) -> None: - """This LLMeter Callback hook is a no-op for CostModel""" pass async def after_invoke(self, response: InvocationResponse) -> None: - """LLMeter Callback.after_invoke hook - - Calls calculate_request_cost() with `save=True` to save the cost on the InvocationResponse. - """ await self.calculate_request_cost(response, save=True) async def before_run(self, run_config: _RunConfig) -> None: - """Initialize state for all run-level cost dimensions in the model""" for dim in self.run_dims.values(): await dim.before_run_start(run_config) async def after_run(self, result: Result) -> None: - """LLMeter Callback.after_run hook - - Calls calculate_run_cost() with `save=True` to save the cost on the Result. - """ await self.calculate_run_cost( result, recalculate_request_costs=False, save=True ) def save_to_file(self, path: WritablePathLike) -> None: - """Save the cost model (including all dimensions) to a JSON file""" path = ensure_path(path) path.parent.mkdir(parents=True, exist_ok=True) + data = dump_object(self) with path.open("w") as f: - f.write(self.to_json()) - - @classmethod - def from_dict(cls, raw: dict, alt_classes: dict = {}, **kwargs) -> "CostModel": - dim_classes = { - **importlib.import_module("llmeter.callbacks.cost.dimensions").__dict__, - **alt_classes, - } - raw_args = {**raw} - for key in ("request_dims", "run_dims"): - if key in raw_args: - raw_args[key] = { - name: from_dict_with_class_map(d, class_map=dim_classes) - for name, d in raw_args[key].items() - } - return super().from_dict(raw_args, alt_classes=alt_classes, **kwargs) + json.dump(data, f, indent=4, default=llmeter_default_serializer) @classmethod - def _load_from_file(cls, path: ReadablePathLike): - """Load the cost model (including all dimensions) from a JSON file""" - path = ensure_path(path) - with path.open("r") as f: - return cls.from_json(f.read()) + def _load_from_file(cls, path: ReadablePathLike) -> "CostModel": + return cls.from_file(path) diff --git a/llmeter/callbacks/cost/serde.py b/llmeter/callbacks/cost/serde.py deleted file mode 100644 index 72025c1..0000000 --- a/llmeter/callbacks/cost/serde.py +++ /dev/null @@ -1,217 +0,0 @@ -# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. -# SPDX-License-Identifier: Apache-2.0 -"""(De/re)serialization interfaces for saving Cost Model objects to file and loading them back""" - -# Python Built-Ins: -import json -import logging -from dataclasses import is_dataclass -from typing import Any, Protocol, TypeVar - -# External Dependencies: -from upath import UPath as Path -from upath.types import ReadablePathLike, WritablePathLike - -# Local Dependencies: -from ...json_utils import llmeter_default_serializer -from ...utils import ensure_path - -logger = logging.getLogger(__name__) - -TSerializable = TypeVar("TSerializable", bound="ISerializable") - - -class ISerializable(Protocol): - def to_dict(self) -> dict: - """Represent this object with a fully JSON-serializable dictionary""" - ... - - @classmethod - def from_dict(cls, **kwargs) -> TSerializable: - """Load an instance of this class from a JSON-able dictionary representation""" - ... - - -def is_dataclass_instance(obj): - """Check whether `obj` is an instance of any dataclass - - See: https://docs.python.org/3/library/dataclasses.html#dataclasses.is_dataclass - """ - return is_dataclass(obj) and not isinstance(obj, type) - - -def to_dict_recursive_generic(obj: object, **kwargs) -> dict: - """Convert a vaguely dataclass-like object (with maybe IJSONable fields) to a JSON-ready dict - - The output dict is augmented with `_type` storing the `__class__.__name__` of the provided - `obj`. - - Args: - obj: The object to convert - **kwargs: Optional extra parameters to insert in the output dictionary - """ - obj_classname = obj.__class__.__name__ - result: dict = {} if obj_classname == "dict" else {"_type": obj.__class__.__name__} - if hasattr(obj, "__dict__"): - # We *don't* use dataclass asdict() here because we want our custom behaviour instead of - # its recursion: - result.update(obj.__dict__) - elif ( - isinstance(obj, dict) - or hasattr(obj, "__keys__") - and hasattr(obj, "__getitem__") - ): - result.update(obj) - else: - result.update({k: getattr(obj, k) for k in dir(obj)}) - result.update(kwargs) - for k, v in result.items(): - if hasattr(v, "to_dict"): - result[k] = v.to_dict() - elif isinstance(v, dict): - result[k] = to_dict_recursive_generic(v) - elif isinstance(v, (list, tuple)): - result[k] = [to_dict_recursive_generic(item) for item in v] - return result - - -TFromDict = TypeVar("TFromDict") - - -def from_dict_with_class(raw: dict, cls: type[TFromDict], **kwargs) -> TFromDict: - """Initialize an instance of a class from a plain dict (with optional extra kwargs) - - If the input dictionary contains a `_type` key, and this doesn't match the provided - `cls.__name__`, a warning will be logged. - - Args: - raw: A plain Python dict, for example loaded from a JSON file - cls: The class to create an instance of - **kwargs: Optional extra keyword arguments to pass to the constructor - """ - raw_args = {k: v for k, v in raw.items()} - raw_type = raw_args.pop("_type", None) - if raw_type is not None and raw_type != cls.__name__: - logger.warning( - "from_dict: _type '%s' doesn't match class '%s' being loaded. %s" - % (raw_type, cls.__name__, raw) - ) - return cls(**raw_args, **kwargs) - - -def from_dict_with_class_map( - raw: dict, class_map: dict[str, type[TFromDict]], **kwargs -) -> TFromDict: - """Initialize an instance of a class from a plain dict (with optional extra kwargs) - - Args: - raw: A plain Python dict which must contain a `_type` key - class_map: A mapping from `_type` string to class to create an instance of - **kwargs: Optional extra keyword arguments to pass to the constructor - """ - if "_type" not in raw: - raise ValueError("from_dict_with_class_map: No _type in raw dict: %s" % raw) - if raw["_type"] not in class_map: - raise ValueError( - "Object _type '%s' not found in provided class_map %s" - % (raw["_type"], class_map) - ) - return from_dict_with_class(raw, class_map[raw["_type"]], **kwargs) - - -TJSONable = TypeVar("TJSONable", bound="JSONableBase") - - -class JSONableBase: - """A base class for speeding up implementation of JSON-serializable objects""" - - @classmethod - def from_dict( - cls: type[TJSONable], - raw: dict, - alt_classes: dict[str, TJSONable] = {}, - **kwargs: Any, - ) -> TJSONable: - """Initialize an instance of this class from a plain dict (with optional extra kwargs) - - Args: - raw: A plain Python dict, for example loaded from a JSON file - alt_classes: By default, this method will only use the class of the current object - (i.e. `cls`). If you want to support loading of subclasses, provide a mapping - from your raw dict's `_type` field to class, for example `{cls.__name__: cls}`. - **kwargs: Optional extra keyword arguments to pass to the constructor - """ - if alt_classes: - return from_dict_with_class_map( - raw=raw, - class_map={cls.__name__: cls, **alt_classes}, - **kwargs, - ) - else: - return from_dict_with_class(raw=raw, cls=cls, **kwargs) - - @classmethod - def from_file( - cls: type[TJSONable], input_path: ReadablePathLike, **kwargs - ) -> TJSONable: - """Initialize an instance of this class from a (local or Cloud) JSON file - - Args: - input_path: The path to the JSON data file. - **kwargs: Optional extra keyword arguments to pass to `from_dict()` - """ - input_path = ensure_path(input_path) - with input_path.open("r") as f: - return cls.from_json(f.read(), **kwargs) - - @classmethod - def from_json(cls: type[TJSONable], json_string: str, **kwargs: Any) -> TJSONable: - """Initialize an instance of this class from a JSON string (with optional extra kwargs) - - Args: - json_string: A string containing valid JSON data - **kwargs: Optional extra keyword arguments to pass to `from_dict()` - """ - return cls.from_dict(json.loads(json_string), **kwargs) - - def to_dict(self, **kwargs) -> dict: - """Save the state of the object to a JSON-dumpable dictionary (with optional extra kwargs) - - Implementers of this method should ensure that the returned dict is fully JSON-compatible: - Mapping any child fields from Python classes to dicts if necessary, avoiding any circular - references, etc. - """ - return to_dict_recursive_generic(self, **kwargs) - - def to_file( - self, - output_path: WritablePathLike, - indent: int | str | None = 4, - **kwargs: Any, - ) -> Path: - """Save the state of the object to a (local or Cloud) JSON file - - Args: - output_path: The path where the configuration file will be saved. - indent: Optional indentation passed through to `to_json()` and therefore - `json.dumps()` - **kwargs: Optional extra keyword arguments to pass to `to_json()` - - Returns: - output_path: Universal Path representation of the target file. - """ - output_path = ensure_path(output_path) - output_path.parent.mkdir(parents=True, exist_ok=True) - with output_path.open("w") as f: - f.write(self.to_json(indent=indent, **kwargs)) - return output_path - - def to_json(self, default=llmeter_default_serializer, **kwargs) -> str: - """Serialize this object to JSON. - - Args: - default: Fallback serializer. Defaults to - :func:`~llmeter.json_utils.llmeter_default_serializer`. - **kwargs: Extra keyword arguments passed to :func:`json.dumps`. - """ - return json.dumps(self.to_dict(), default=default, **kwargs) diff --git a/llmeter/serialization.py b/llmeter/serialization.py index 78882b4..b76d3b2 100644 --- a/llmeter/serialization.py +++ b/llmeter/serialization.py @@ -114,12 +114,51 @@ def load_object(data: dict) -> Any: # --------------------------------------------------------------------------- +def _serialize_value(val: Any) -> Any: + """Recursively serialize a value for JSON persistence. + + - Objects with custom ``__getstate__`` → ``dump_object(val)`` + - Dicts → recurse into values + - Lists/tuples → recurse into items + - Scalars (str, int, float, bool, None) → pass through + """ + if val is None or isinstance(val, (str, int, float, bool)): + return val + if hasattr(val, "__getstate__") and type(val).__getstate__ is not object.__getstate__: + return dump_object(val) + if isinstance(val, dict): + return {k: _serialize_value(v) for k, v in val.items()} + if isinstance(val, (list, tuple)): + return [_serialize_value(item) for item in val] + # Fallback: try str + return str(val) + + +def _deserialize_value(val: Any) -> Any: + """Recursively deserialize a value from JSON persistence. + + - Dicts with ``_class``/``_state`` → ``load_object(val)`` + - Other dicts → recurse into values + - Lists → recurse into items + - Scalars → pass through + """ + if val is None or isinstance(val, (str, int, float, bool)): + return val + if isinstance(val, dict): + if "_class" in val and "_state" in val: + return load_object(val) + return {k: _deserialize_value(v) for k, v in val.items()} + if isinstance(val, (list, tuple)): + return [_deserialize_value(item) for item in val] + return val + + def default_getstate(obj: Any) -> dict: """Default __getstate__: infers state from __init__ signature. Matches __init__ parameter names to instance attributes (self.name or - self._name). Parameters without a match are omitted -- __init__ will use - defaults on reconstruction. + self._name). Nested objects with ``__getstate__`` are recursively + serialized via ``dump_object``. Parameters without a match are omitted. """ sig = inspect.signature(obj.__init__) state = {} @@ -129,12 +168,16 @@ def default_getstate(obj: Any) -> dict: if param.kind in (param.VAR_POSITIONAL, param.VAR_KEYWORD): continue if hasattr(obj, name): - state[name] = getattr(obj, name) + state[name] = _serialize_value(getattr(obj, name)) elif hasattr(obj, f"_{name}"): - state[name] = getattr(obj, f"_{name}") + state[name] = _serialize_value(getattr(obj, f"_{name}")) return state def default_setstate(obj: Any, state: dict) -> None: - """Default __setstate__: calls __init__(**state) to reconstruct.""" - obj.__init__(**state) + """Default __setstate__: deserializes nested objects then calls __init__. + + Any dict with ``_class``/``_state`` keys is restored via ``load_object``. + """ + deserialized = {k: _deserialize_value(v) for k, v in state.items()} + obj.__init__(**deserialized) From 39114320c68565db8090f95c0a66c0558ec2d289 Mon Sep 17 00:00:00 2001 From: Alessandro Cere Date: Tue, 16 Jun 2026 14:15:23 -0400 Subject: [PATCH 04/20] introduce Serializable mixin, drop all backward compat MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Add Serializable class in serialization.py (provides __getstate__/ __setstate__ via inheritance — no boilerplate in subclasses) - Endpoint, Tokenizer, Callback inherit from Serializable - Cost dimension bases inherit from Serializable - Remove all legacy format handling (_type tags, endpoint_type, tokenizer_module fallbacks in Runner.load()) - Remove redundant to_dict/to_json/from_dict/from_file/save_to_file overrides from CostModel (all provided by base classes) - Clean up unused imports 863 tests pass, 0 failures. --- llmeter/callbacks/base.py | 12 +-- llmeter/callbacks/cost/dimensions.py | 51 ++--------- llmeter/callbacks/cost/model.py | 86 ------------------- llmeter/endpoints/base.py | 16 +--- llmeter/runner.py | 12 +-- llmeter/serialization.py | 25 ++++++ llmeter/tokenizers.py | 12 +-- .../cost/providers/test_sagemaker.py | 11 ++- tests/unit/callbacks/cost/test_dimensions.py | 19 ++-- tests/unit/callbacks/cost/test_model.py | 50 +++++------ 10 files changed, 74 insertions(+), 220 deletions(-) diff --git a/llmeter/callbacks/base.py b/llmeter/callbacks/base.py index c869bf6..02051bf 100644 --- a/llmeter/callbacks/base.py +++ b/llmeter/callbacks/base.py @@ -14,11 +14,11 @@ from ..json_utils import llmeter_default_serializer from ..results import Result from ..runner import _RunConfig -from ..serialization import default_getstate, default_setstate, dump_object, load_object +from ..serialization import Serializable, dump_object, load_object from ..utils import ensure_path -class Callback(ABC): +class Callback(Serializable, ABC): """Base class for a callback in LLMeter Callbacks support extending LLMeter functionality by running additional code at defined points @@ -75,14 +75,6 @@ async def after_run(self, result: Result) -> None: """ pass - def __getstate__(self) -> dict: - """Serialize callback configuration for persistence.""" - return default_getstate(self) - - def __setstate__(self, state: dict) -> None: - """Restore callback from saved state.""" - default_setstate(self, state) - def save_to_file(self, path: WritablePathLike) -> None: """Save this Callback to a JSON file. diff --git a/llmeter/callbacks/cost/dimensions.py b/llmeter/callbacks/cost/dimensions.py index ce5636f..ae93e29 100644 --- a/llmeter/callbacks/cost/dimensions.py +++ b/llmeter/callbacks/cost/dimensions.py @@ -17,7 +17,7 @@ from ...endpoints.base import InvocationResponse from ...results import Result from ...runner import _RunConfig -from ...serialization import default_getstate, default_setstate +from ...serialization import Serializable class IRequestCostDimension(Protocol): @@ -33,59 +33,24 @@ async def before_run_start(self, run_config: _RunConfig) -> None: ... async def calculate(self, result: Result) -> Optional[float]: ... -class RequestCostDimensionBase(ABC): - """Base class for implementing per-request cost model dimensions. +class RequestCostDimensionBase(Serializable, ABC): + """Base class for per-request cost dimensions. - Provides ``__getstate__``/``__setstate__`` and ``to_dict``/``from_dict`` for - serialization. Subclasses just declare fields and implement ``calculate()``. + Inherits ``__getstate__``/``__setstate__`` from :class:`~llmeter.serialization.Serializable`. + Subclasses just declare fields and implement ``calculate()``. """ - def __getstate__(self) -> dict: - return default_getstate(self) - - def __setstate__(self, state: dict) -> None: - default_setstate(self, state) - - @classmethod - def from_dict(cls, raw: dict): - """Create from a plain dict (strips _type if present).""" - return cls(**{k: v for k, v in raw.items() if k != "_type"}) - - def to_dict(self) -> dict: - """Serialize to dict with _type tag for polymorphic deserialization.""" - state = self.__getstate__() - state["_type"] = self.__class__.__name__ - return state - @abstractmethod async def calculate(self, response: InvocationResponse) -> Optional[float]: raise NotImplementedError -class RunCostDimensionBase(ABC): - """Base class for implementing per-run cost model dimensions. +class RunCostDimensionBase(Serializable, ABC): + """Base class for per-run cost dimensions. - Provides ``__getstate__``/``__setstate__`` and ``to_dict``/``from_dict`` for - serialization. + Inherits ``__getstate__``/``__setstate__`` from :class:`~llmeter.serialization.Serializable`. """ - def __getstate__(self) -> dict: - return default_getstate(self) - - def __setstate__(self, state: dict) -> None: - default_setstate(self, state) - - @classmethod - def from_dict(cls, raw: dict): - """Create from a plain dict (strips _type if present).""" - return cls(**{k: v for k, v in raw.items() if k != "_type"}) - - def to_dict(self) -> dict: - """Serialize to dict with _type tag for polymorphic deserialization.""" - state = self.__getstate__() - state["_type"] = self.__class__.__name__ - return state - async def before_run_start(self, run_config: _RunConfig) -> None: """Called before a test run starts. Default is a no-op.""" pass diff --git a/llmeter/callbacks/cost/model.py b/llmeter/callbacks/cost/model.py index d77eabb..4ba1601 100644 --- a/llmeter/callbacks/cost/model.py +++ b/llmeter/callbacks/cost/model.py @@ -2,17 +2,9 @@ # SPDX-License-Identifier: Apache-2.0 """Cost modelling callback for LLMeter test runs.""" -import importlib -import json - -from upath.types import ReadablePathLike, WritablePathLike - from ...endpoints.base import InvocationResponse -from ...json_utils import llmeter_default_serializer from ...results import Result from ...runner import _RunConfig -from ...serialization import dump_object, load_object -from ...utils import ensure_path from ..base import Callback from .dimensions import IRequestCostDimension, IRunCostDimension from .results import CalculatedCostWithDimensions @@ -70,73 +62,6 @@ def __init__( all_dims[name] = dim self.run_dims[name] = dim - # ------------------------------------------------------------------ - # Serialization helpers (to_dict/from_dict for legacy compat) - # ------------------------------------------------------------------ - - def to_dict(self) -> dict: - """Serialize to dict with _type tags (legacy-compatible format).""" - return { - "_type": "CostModel", - "request_dims": { - name: dim.to_dict() - for name, dim in self.request_dims.items() - }, - "run_dims": { - name: dim.to_dict() - for name, dim in self.run_dims.items() - }, - } - - def to_json(self, indent: int | None = 4) -> str: - return json.dumps(self.to_dict(), indent=indent) - - @classmethod - def from_dict(cls, raw: dict, **kwargs) -> "CostModel": - """Load from dict (supports both new _class/_state and legacy _type).""" - clean = {k: v for k, v in raw.items() if k != "_type"} - - # Check if dimensions use legacy _type format and convert - dim_module = importlib.import_module( - "llmeter.callbacks.cost.dimensions" - ) - for key in ("request_dims", "run_dims"): - dims = clean.get(key, {}) - if isinstance(dims, dict): - converted = {} - for name, d in dims.items(): - if isinstance(d, dict) and "_type" in d and "_class" not in d: - # Legacy format → convert to object directly - type_name = d.pop("_type") - dim_cls = getattr(dim_module, type_name, None) - if dim_cls is None: - providers = importlib.import_module( - "llmeter.callbacks.cost.providers.sagemaker" - ) - dim_cls = getattr(providers, type_name, None) - if dim_cls is None: - raise ValueError( - f"Unknown dimension type: {type_name!r}" - ) - converted[name] = dim_cls(**d) - elif isinstance(d, dict) and "_class" in d: - converted[name] = load_object(d) - else: - converted[name] = d - clean[key] = converted - - return cls(**clean) - - @classmethod - def from_json(cls, json_string: str) -> "CostModel": - return cls.from_dict(json.loads(json_string)) - - @classmethod - def from_file(cls, path: ReadablePathLike) -> "CostModel": - path = ensure_path(path) - with path.open("r") as f: - return cls.from_json(f.read()) - # ------------------------------------------------------------------ # Callback lifecycle hooks # ------------------------------------------------------------------ @@ -209,14 +134,3 @@ async def after_run(self, result: Result) -> None: await self.calculate_run_cost( result, recalculate_request_costs=False, save=True ) - - def save_to_file(self, path: WritablePathLike) -> None: - path = ensure_path(path) - path.parent.mkdir(parents=True, exist_ok=True) - data = dump_object(self) - with path.open("w") as f: - json.dump(data, f, indent=4, default=llmeter_default_serializer) - - @classmethod - def _load_from_file(cls, path: ReadablePathLike) -> "CostModel": - return cls.from_file(path) diff --git a/llmeter/endpoints/base.py b/llmeter/endpoints/base.py index ddbc576..e6f2dda 100644 --- a/llmeter/endpoints/base.py +++ b/llmeter/endpoints/base.py @@ -22,7 +22,7 @@ from upath.types import ReadablePathLike, WritablePathLike from ..json_utils import llmeter_bytes_decoder, llmeter_default_serializer -from ..serialization import default_getstate, default_setstate +from ..serialization import Serializable from ..utils import ensure_path logger = logging.getLogger(__name__) @@ -164,7 +164,7 @@ def to_dict(self) -> dict: TRawResponse = TypeVar("TRawResponse", bound=Any) -class Endpoint(ABC, Generic[TRawResponse]): +class Endpoint(Serializable, ABC, Generic[TRawResponse]): """ An abstract base class for endpoint implementations. @@ -504,18 +504,6 @@ def to_dict(self) -> dict: endpoint_conf["endpoint_type"] = self.__class__.__name__ return endpoint_conf - def __getstate__(self) -> dict: - """Serialize endpoint configuration for persistence. - - Introspects ``__init__`` to determine what's needed to recreate this endpoint. - Runtime state (clients, connections) is excluded automatically. - """ - return default_getstate(self) - - def __setstate__(self, state: dict) -> None: - """Restore endpoint from saved state by calling ``__init__(**state)``.""" - default_setstate(self, state) - @classmethod def load_from_file(cls, input_path: ReadablePathLike) -> "Endpoint": """ diff --git a/llmeter/runner.py b/llmeter/runner.py index 6d4ee52..98f36d4 100644 --- a/llmeter/runner.py +++ b/llmeter/runner.py @@ -218,24 +218,18 @@ def load(cls, load_path: ReadablePathLike, file_name: str = "run_config.json"): # Restore endpoint ep = config.get("endpoint") if isinstance(ep, dict): - if "_class" in ep: - config["endpoint"] = load_object(ep) - else: - config["endpoint"] = Endpoint.load(ep) + config["endpoint"] = load_object(ep) # Restore tokenizer tok = config.get("tokenizer") if isinstance(tok, dict): - if "_class" in tok: - config["tokenizer"] = load_object(tok) - else: - config["tokenizer"] = Tokenizer.load(tok) + config["tokenizer"] = load_object(tok) # Restore callbacks cbs = config.get("callbacks") if isinstance(cbs, list): config["callbacks"] = [ - load_object(cb) if isinstance(cb, dict) and "_class" in cb else cb + load_object(cb) if isinstance(cb, dict) else cb for cb in cbs ] diff --git a/llmeter/serialization.py b/llmeter/serialization.py index b76d3b2..0d2a466 100644 --- a/llmeter/serialization.py +++ b/llmeter/serialization.py @@ -32,6 +32,31 @@ logger = logging.getLogger(__name__) +# --------------------------------------------------------------------------- +# Serializable mixin +# --------------------------------------------------------------------------- + + +class Serializable: + """Mixin that provides ``__getstate__``/``__setstate__`` via introspection. + + Inherit from this to get automatic serialization support. Works with + plain classes, ``@dataclass``, and any class whose ``__init__`` parameters + match instance attributes. + + Example:: + + class MyEndpoint(Serializable, Endpoint): + ... # __getstate__/__setstate__ inherited — no boilerplate needed + """ + + def __getstate__(self) -> dict: + return default_getstate(self) + + def __setstate__(self, state: dict) -> None: + default_setstate(self, state) + + # --------------------------------------------------------------------------- # Public API # --------------------------------------------------------------------------- diff --git a/llmeter/tokenizers.py b/llmeter/tokenizers.py index 3db0668..2a3e725 100644 --- a/llmeter/tokenizers.py +++ b/llmeter/tokenizers.py @@ -10,11 +10,11 @@ from upath import UPath from upath.types import ReadablePathLike, WritablePathLike -from .serialization import default_getstate, default_setstate +from .serialization import Serializable from .utils import ensure_path -class Tokenizer(ABC): +class Tokenizer(Serializable, ABC): def __init__(self, *args, **kwargs): pass @@ -26,14 +26,6 @@ def encode(self, text: str): def decode(self, tokens: list[str]): raise NotImplementedError - def __getstate__(self) -> dict: - """Serialize tokenizer configuration for persistence.""" - return default_getstate(self) - - def __setstate__(self, state: dict) -> None: - """Restore tokenizer from saved state.""" - default_setstate(self, state) - @classmethod def __subclasshook__(cls, C): if cls is Tokenizer: diff --git a/tests/unit/callbacks/cost/providers/test_sagemaker.py b/tests/unit/callbacks/cost/providers/test_sagemaker.py index 2bf7909..10fac84 100644 --- a/tests/unit/callbacks/cost/providers/test_sagemaker.py +++ b/tests/unit/callbacks/cost/providers/test_sagemaker.py @@ -11,6 +11,7 @@ from moto import mock_aws # Local Dependencies: +from llmeter.serialization import to_dict, dump_object, load_object from llmeter.callbacks.cost.providers.sagemaker import ( SageMakerRTEndpointCompute, SageMakerRTEndpointStorage, @@ -43,16 +44,15 @@ async def test_real_time_endpoint_compute_explicit_price(): result.total_test_time = None assert await dim.calculate(result) is None - dim_ser = dim.to_dict() + dim_ser = to_dict(dim) assert dim_ser == { - "_type": "SageMakerRTEndpointCompute", "instance_count": 1, "instance_type": "ml.doesnotexist", "price_per_hour": 6, "region": "us-nope-0", } - dim = SageMakerRTEndpointCompute.from_dict(dim_ser) + dim = load_object(dump_object(dim)) result.total_test_time = 90 assert await dim.calculate(result) == 6 * 90 / 3600 @@ -237,15 +237,14 @@ async def test_real_time_endpoint_storage_explicit_price(): result.total_test_time = None assert await dim.calculate(result) is None - dim_ser = dim.to_dict() + dim_ser = to_dict(dim) assert dim_ser == { - "_type": "SageMakerRTEndpointStorage", "gbs_provisioned": 5, "price_per_gb_hour": 9, "region": "us-nope-0", } - dim = SageMakerRTEndpointStorage.from_dict(dim_ser) + dim = load_object(dump_object(dim)) result.total_test_time = 90 assert await dim.calculate(result) == 9 * 5 * 90 / 3600 diff --git a/tests/unit/callbacks/cost/test_dimensions.py b/tests/unit/callbacks/cost/test_dimensions.py index 68cbe1c..b9a167b 100644 --- a/tests/unit/callbacks/cost/test_dimensions.py +++ b/tests/unit/callbacks/cost/test_dimensions.py @@ -2,6 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 import pytest +from llmeter.serialization import to_dict from llmeter.callbacks.cost.dimensions import ( EndpointTime, InputTokens, @@ -30,11 +31,10 @@ class TestRunDim(RunCostDimensionBase): @pytest.mark.asyncio async def test_cost_per_input_token(): spec = { - "_type": "InputTokens", "price_per_million": 30, "granularity": 10, } - dim_valid = InputTokens.from_dict(spec) + dim_valid = InputTokens(price_per_million=spec["price_per_million"], granularity=spec["granularity"]) success_response = InvocationResponse(response_text="hi", num_tokens_input=199999) assert await dim_valid.calculate(success_response) == 6 @@ -47,8 +47,7 @@ async def test_cost_per_input_token(): err_response = InvocationResponse.error_output() assert await dim_valid.calculate(err_response) is None - assert dim_valid.to_dict() == { - "_type": "InputTokens", + assert to_dict(dim_valid) == { "price_per_million": 30, "granularity": 10, } @@ -57,11 +56,10 @@ async def test_cost_per_input_token(): @pytest.mark.asyncio async def test_cost_per_output_token(): spec = { - "_type": "OutputTokens", "price_per_million": 40, "granularity": 10, } - dim_valid = OutputTokens.from_dict(spec) + dim_valid = OutputTokens(price_per_million=spec["price_per_million"], granularity=spec["granularity"]) success_response = InvocationResponse(response_text="hi", num_tokens_output=199999) assert await dim_valid.calculate(success_response) == 8 @@ -74,8 +72,7 @@ async def test_cost_per_output_token(): err_response = InvocationResponse.error_output() assert await dim_valid.calculate(err_response) is None - assert dim_valid.to_dict() == { - "_type": "OutputTokens", + assert to_dict(dim_valid) == { "price_per_million": 40, "granularity": 10, } @@ -84,11 +81,10 @@ async def test_cost_per_output_token(): @pytest.mark.asyncio async def test_cost_per_hour(): spec = { - "_type": "EndpointTime", "price_per_hour": 30, "granularity_secs": 60, } - dim_valid = EndpointTime.from_dict(spec) + dim_valid = EndpointTime(price_per_hour=spec["price_per_hour"], granularity_secs=spec["granularity_secs"]) result = Result( [], @@ -105,8 +101,7 @@ async def test_cost_per_hour(): result.total_test_time = None assert await dim_valid.calculate(result) is None - assert dim_valid.to_dict() == { - "_type": "EndpointTime", + assert to_dict(dim_valid) == { "price_per_hour": 30, "granularity_secs": 60, } diff --git a/tests/unit/callbacks/cost/test_model.py b/tests/unit/callbacks/cost/test_model.py index b1ea1a8..65c6371 100644 --- a/tests/unit/callbacks/cost/test_model.py +++ b/tests/unit/callbacks/cost/test_model.py @@ -5,41 +5,31 @@ import pytest +from llmeter.serialization import to_dict, dump_object, load_object from llmeter.callbacks.cost.model import CostModel from llmeter.callbacks.cost.results import CalculatedCostWithDimensions def test_cost_model_serialization(): - """Cost models can be serialized & de-serialized""" - spec = { - "_type": "CostModel", - "request_dims": { - "TokensIn": {"_type": "InputTokens", "price_per_million": 30}, - }, - "run_dims": { - "ComputeSeconds": {"_type": "EndpointTime", "price_per_hour": 50}, - }, - } - model = CostModel.from_dict(spec) - assert model.request_dims["TokensIn"].price_per_million == 30 - assert model.run_dims["ComputeSeconds"].price_per_hour == 50 - assert model.to_dict() == { - "_type": "CostModel", - "request_dims": { - "TokensIn": { - "_type": "InputTokens", - "price_per_million": 30, - "granularity": 1, - }, - }, - "run_dims": { - "ComputeSeconds": { - "_type": "EndpointTime", - "price_per_hour": 50, - "granularity_secs": 1, - }, - }, - } + """Cost models can be serialized & de-serialized via dump_object/load_object""" + from llmeter.callbacks.cost.dimensions import EndpointTime, InputTokens + + model = CostModel( + request_dims={"TokensIn": InputTokens(price_per_million=30)}, + run_dims={"ComputeSeconds": EndpointTime(price_per_hour=50)}, + ) + + # Round-trip via dump_object / load_object + data = dump_object(model) + restored = load_object(data) + + assert restored.request_dims["TokensIn"].price_per_million == 30 + assert restored.run_dims["ComputeSeconds"].price_per_hour == 50 + + # to_dict produces a plain dict representation + d = to_dict(model) + assert "request_dims" in d + assert "run_dims" in d def test_cost_model_detects_duplicate_cost_dim_names(): From 6093999cdd07316f42368588f09ce316458cbd20 Mon Sep 17 00:00:00 2001 From: Alessandro Cere Date: Tue, 16 Jun 2026 14:37:09 -0400 Subject: [PATCH 05/20] use __llmeter_ prefix for serialization envelope keys Rename _class/_state to __llmeter_class__/__llmeter_state__ to avoid collisions with user data that might legitimately contain _class or _state keys. Consistent with the existing __llmeter_bytes__ convention in json_utils.py. 863 tests pass. --- docs/serialization.md | 10 +++++----- llmeter/callbacks/base.py | 2 +- llmeter/serialization.py | 18 +++++++++--------- tests/unit/callbacks/test_base.py | 10 +++++----- 4 files changed, 20 insertions(+), 20 deletions(-) diff --git a/docs/serialization.md b/docs/serialization.md index d8ab62d..aeb7d4d 100644 --- a/docs/serialization.md +++ b/docs/serialization.md @@ -40,8 +40,8 @@ from llmeter.serialization import dump_object, load_object # Save an endpoint (including its config, but not the boto3 client) data = dump_object(endpoint) -# → {"_class": "llmeter.endpoints.bedrock.BedrockConverse", -# "_state": {"model_id": "claude-3", "region": "us-west-2"}} +# → {"__llmeter_class__": "llmeter.endpoints.bedrock.BedrockConverse", +# "__llmeter_state__": {"model_id": "claude-3", "region": "us-west-2"}} # Restore it (boto3 client is recreated automatically) restored = load_object(data) @@ -117,9 +117,9 @@ runner = Runner(endpoint=BedrockConverse(...), callbacks=[MlflowCallback(step=1) runner.save(output_path="/tmp/run") # Saved JSON includes: -# {"endpoint": {"_class": "...BedrockConverse", "_state": {...}}, -# "tokenizer": {"_class": "...DummyTokenizer", "_state": {}}, -# "callbacks": [{"_class": "...MlflowCallback", "_state": {"step": 1, "nested": false}}]} +# {"endpoint": {"__llmeter_class__": "...BedrockConverse", "__llmeter_state__": {...}}, +# "tokenizer": {"__llmeter_class__": "...DummyTokenizer", "__llmeter_state__": {}}, +# "callbacks": [{"__llmeter_class__": "...MlflowCallback", "__llmeter_state__": {"step": 1, "nested": false}}]} # Full reconstruction: restored = Runner.load("/tmp/run") diff --git a/llmeter/callbacks/base.py b/llmeter/callbacks/base.py index 02051bf..a4a91a3 100644 --- a/llmeter/callbacks/base.py +++ b/llmeter/callbacks/base.py @@ -95,7 +95,7 @@ def save_to_file(self, path: WritablePathLike) -> None: def load_from_file(path: ReadablePathLike) -> Callback: """Load (any type of) Callback from a JSON file. - Detects the callback type from the ``_class`` field and reconstructs it. + Detects the callback type from the ``__llmeter_class__`` field and reconstructs it. Args: path: (Local or Cloud) path where the callback was saved. diff --git a/llmeter/serialization.py b/llmeter/serialization.py index 0d2a466..c3a72bf 100644 --- a/llmeter/serialization.py +++ b/llmeter/serialization.py @@ -13,13 +13,13 @@ and ``load_object`` provide full round-trip persistence using the ``__getstate__``/``__setstate__`` protocol:: - data = dump_object(endpoint) # -> {"_class": "...", "_state": {...}} + data = dump_object(endpoint) # -> {"__llmeter_class__": "...", "__llmeter_state__": {...}} restored = load_object(data) # -> new Endpoint instance .. warning:: Security ``load_object`` will import and instantiate any class path found in the - ``_class`` field. Do not load configs from untrusted sources. + ``__llmeter_class__`` field. Do not load configs from untrusted sources. """ import importlib @@ -105,7 +105,7 @@ def serialize(obj: Any) -> dict: def dump_object(obj: Any) -> dict: """Serialize an object to a type-tagged dict for full round-trip persistence. - Returns ``{"_class": "module.ClassName", "_state": {...}}``. + Returns ``{"__llmeter_class__": "module.ClassName", "__llmeter_state__": {...}}``. """ class_path = f"{obj.__class__.__module__}.{obj.__class__.__qualname__}" try: @@ -116,7 +116,7 @@ def dump_object(obj: Any) -> dict: "dump_object: serialize(%s) returned non-serializable data", class_path ) state = {} - return {"_class": class_path, "_state": state} + return {"__llmeter_class__": class_path, "__llmeter_state__": state} def load_object(data: dict) -> Any: @@ -124,13 +124,13 @@ def load_object(data: dict) -> Any: WARNING: Do not call on data from untrusted sources. """ - class_path = data["_class"] + class_path = data["__llmeter_class__"] module_path, class_name = class_path.rsplit(".", 1) module = importlib.import_module(module_path) cls = getattr(module, class_name) obj = cls.__new__(cls) - obj.__setstate__(data["_state"]) + obj.__setstate__(data["__llmeter_state__"]) return obj @@ -162,7 +162,7 @@ def _serialize_value(val: Any) -> Any: def _deserialize_value(val: Any) -> Any: """Recursively deserialize a value from JSON persistence. - - Dicts with ``_class``/``_state`` → ``load_object(val)`` + - Dicts with ``__llmeter_class__``/``__llmeter_state__`` → ``load_object(val)`` - Other dicts → recurse into values - Lists → recurse into items - Scalars → pass through @@ -170,7 +170,7 @@ def _deserialize_value(val: Any) -> Any: if val is None or isinstance(val, (str, int, float, bool)): return val if isinstance(val, dict): - if "_class" in val and "_state" in val: + if "__llmeter_class__" in val and "__llmeter_state__" in val: return load_object(val) return {k: _deserialize_value(v) for k, v in val.items()} if isinstance(val, (list, tuple)): @@ -202,7 +202,7 @@ def default_getstate(obj: Any) -> dict: def default_setstate(obj: Any, state: dict) -> None: """Default __setstate__: deserializes nested objects then calls __init__. - Any dict with ``_class``/``_state`` keys is restored via ``load_object``. + Any dict with ``__llmeter_class__``/``__llmeter_state__`` keys is restored via ``load_object``. """ deserialized = {k: _deserialize_value(v) for k, v in state.items()} obj.__init__(**deserialized) diff --git a/tests/unit/callbacks/test_base.py b/tests/unit/callbacks/test_base.py index c332341..a648feb 100644 --- a/tests/unit/callbacks/test_base.py +++ b/tests/unit/callbacks/test_base.py @@ -9,7 +9,7 @@ class TestBase: def test_save_to_file_creates_valid_json(self, tmp_path): - """save_to_file creates a valid JSON file with _class and _state.""" + """save_to_file creates a valid JSON file with __llmeter_class__ and __llmeter_state__.""" cb = MlflowCallback(step=5, nested=True) path = tmp_path / "callback.json" cb.save_to_file(path) @@ -17,10 +17,10 @@ def test_save_to_file_creates_valid_json(self, tmp_path): with open(path) as f: data = json.load(f) - assert "_class" in data - assert "_state" in data - assert data["_class"] == "llmeter.callbacks.mlflow.MlflowCallback" - assert data["_state"] == {"step": 5, "nested": True} + assert "__llmeter_class__" in data + assert "__llmeter_state__" in data + assert data["__llmeter_class__"] == "llmeter.callbacks.mlflow.MlflowCallback" + assert data["__llmeter_state__"] == {"step": 5, "nested": True} def test_load_from_file_restores_callback(self, tmp_path): """load_from_file correctly restores a callback instance.""" From 358f35d356edc587f0a812373f5b57c4e5eda89c Mon Sep 17 00:00:00 2001 From: Alessandro Cere Date: Wed, 15 Jul 2026 13:55:32 +0800 Subject: [PATCH 06/20] refactor: consolidate serialization, remove json_utils and dead code - Delete llmeter/json_utils.py; move json_default and bytes_decoder into llmeter/serialization.py as the single serialization module. - Remove duplicate serialize() function (was identical to to_dict). - Inline __getstate__/__setstate__ logic directly in Serializable mixin instead of delegating to free functions. - Add datetime_to_str/str_to_datetime helpers; standardize all datetime parsing to use them instead of ad-hoc .replace("Z", "+00:00"). - Remove redundant CostModel.before_invoke no-op override. - Remove legacy tokenizer functions: save_tokenizer, _to_dict, Tokenizer.to_dict, Tokenizer.load_from_file. Keep only Tokenizer.load() for backward-compat with old config format. - Unify _RunConfig.__post_init__ to try load_object (new format) first, then fall back to legacy Endpoint.load/Tokenizer.load. - Standardize JSON file writing to use json.dump() consistently. - Update all imports across source and test files. - Rewrite docs/serialization.md to reflect current API. - Remove stale docs (reference/json_utils.md, reference/cost/serde.md). - Update mkdocs.yml nav accordingly. --- docs/reference/callbacks/cost/serde.md | 1 - docs/reference/json_utils.md | 1 - docs/reference/serialization.md | 1 + docs/serialization.md | 150 ++++++---- llmeter/callbacks/base.py | 5 +- llmeter/callbacks/cost/model.py | 3 - llmeter/endpoints/base.py | 60 ++-- llmeter/json_utils.py | 110 ------- llmeter/prompt_utils.py | 10 +- llmeter/results.py | 28 +- llmeter/runner.py | 19 +- llmeter/serialization.py | 223 +++++++------- llmeter/tokenizers.py | 117 ++------ mkdocs.yml | 4 +- .../cost/providers/test_sagemaker.py | 6 +- tests/unit/callbacks/cost/test_dimensions.py | 7 +- tests/unit/callbacks/cost/test_model.py | 6 +- .../endpoints/test_multimodal_properties.py | 6 +- .../test_openai_response_serialization.py | 11 +- tests/unit/test_invocation_response.py | 18 +- tests/unit/test_prompt_utils.py | 102 +++---- tests/unit/test_property_save_load.py | 78 ++--- tests/unit/test_property_simple.py | 14 +- tests/unit/test_results.py | 281 +----------------- tests/unit/test_runner.py | 16 +- tests/unit/test_serialization_properties.py | 104 +++---- tests/unit/test_tokenizers.py | 86 ++---- 27 files changed, 475 insertions(+), 992 deletions(-) delete mode 100644 docs/reference/callbacks/cost/serde.md delete mode 100644 docs/reference/json_utils.md create mode 100644 docs/reference/serialization.md delete mode 100644 llmeter/json_utils.py diff --git a/docs/reference/callbacks/cost/serde.md b/docs/reference/callbacks/cost/serde.md deleted file mode 100644 index 5ea2d0f..0000000 --- a/docs/reference/callbacks/cost/serde.md +++ /dev/null @@ -1 +0,0 @@ -::: llmeter.callbacks.cost.serde diff --git a/docs/reference/json_utils.md b/docs/reference/json_utils.md deleted file mode 100644 index 4fd15e6..0000000 --- a/docs/reference/json_utils.md +++ /dev/null @@ -1 +0,0 @@ -::: llmeter.json_utils diff --git a/docs/reference/serialization.md b/docs/reference/serialization.md new file mode 100644 index 0000000..9aa47cf --- /dev/null +++ b/docs/reference/serialization.md @@ -0,0 +1 @@ +::: llmeter.serialization diff --git a/docs/serialization.md b/docs/serialization.md index aeb7d4d..7977915 100644 --- a/docs/serialization.md +++ b/docs/serialization.md @@ -1,90 +1,128 @@ # Serialization -LLMeter provides a unified serialization layer that works consistently across all object types — Pydantic models, dataclasses, plain classes, and callable objects with runtime state. +LLMeter provides a unified serialization layer in a single module — +`llmeter.serialization` — that handles JSON encoding, datetime conversion, +binary content round-tripping, and full object persistence. -## Two operations, two purposes +## Module overview -| Function | Returns | Use case | -|----------|---------|----------| -| `to_dict(obj)` | Python dict with **native types** (datetime, bytes, UPath) | In-memory access, stats computation, jmespath queries | -| `serialize(obj)` | **JSON-safe** dict (strings, numbers, lists, dicts, None) | Writing to disk, JSON output, network transfer | +| Symbol | Purpose | +|--------|---------| +| `Serializable` | Mixin giving any class automatic `__getstate__`/`__setstate__` | +| `dump_object` / `load_object` | Full round-trip persistence via a type-tagged envelope | +| `json_default` | `json.dumps` fallback for bytes, datetime, PathLike | +| `bytes_decoder` | `json.loads` object hook to restore `__llmeter_bytes__` markers | +| `datetime_to_str` / `str_to_datetime` | UTC ISO-8601 with `Z` suffix, both directions | ```python -from llmeter.serialization import to_dict, serialize +from llmeter.serialization import ( + dump_object, load_object, json_default, bytes_decoder, + datetime_to_str, str_to_datetime, +) +``` + +## JSON encoding -response = InvocationResponse(response_text="hi", request_time=datetime.now(timezone.utc)) +Use `json_default` as the `default` argument whenever you call `json.dumps` or +`json.dump` with LLMeter objects: + +```python +import json +from llmeter.serialization import json_default -to_dict(response)["request_time"] # → datetime(2024, 6, 15, 10, 30, ...) -serialize(response)["request_time"] # → "2024-06-15T10:30:00Z" +json.dump(my_data, f, default=json_default, indent=4) ``` -Both functions work on **any** LLMeter object — the user never needs to know whether something is a Pydantic model, a dataclass, or a plain class. +It handles: -## Object methods +- **bytes** — wrapped in `{"__llmeter_bytes__": ""}` markers +- **datetime** — UTC ISO-8601 string via `datetime_to_str` +- **date / time** — `.isoformat()` +- **os.PathLike** — POSIX path string +- **anything else** — `str()` fallback -Every LLMeter object exposes `to_dict()` as a method for convenience: +To decode bytes markers back on load: ```python -response.to_dict() # same as to_dict(response) -result.to_dict() # excludes responses by default -endpoint.to_dict() # endpoint config as dict -cost_model.to_dict() # includes dimension metadata +data = json.load(f, object_hook=bytes_decoder) ``` +## Datetime handling + +All datetime serialization goes through two functions: + +```python +from llmeter.serialization import datetime_to_str, str_to_datetime + +datetime_to_str(datetime(2024, 1, 1, tzinfo=timezone.utc)) +# → "2024-01-01T00:00:00Z" + +str_to_datetime("2024-01-01T00:00:00Z") +# → datetime(2024, 1, 1, 0, 0, tzinfo=timezone.utc) +``` + +Timezone-aware datetimes are normalized to UTC. Naive datetimes are serialized +as-is. The `Z` suffix is canonical. + ## Persisting runtime objects -Objects that hold runtime state (boto3 clients, SDK connections, compiled tokenizers) can't be naively serialized to JSON. LLMeter handles this via the `__getstate__`/`__setstate__` protocol: +Objects with runtime state (boto3 clients, SDK connections, compiled +tokenizers) can't be naively serialized. LLMeter handles this via the +`__getstate__`/`__setstate__` protocol: ```python from llmeter.serialization import dump_object, load_object -# Save an endpoint (including its config, but not the boto3 client) +# Save an endpoint (config only — not the boto3 client) data = dump_object(endpoint) # → {"__llmeter_class__": "llmeter.endpoints.bedrock.BedrockConverse", # "__llmeter_state__": {"model_id": "claude-3", "region": "us-west-2"}} -# Restore it (boto3 client is recreated automatically) +# Restore it (boto3 client recreated automatically via __init__) restored = load_object(data) ``` ### How it works -1. `dump_object(obj)` calls `serialize(obj)` to get a JSON-safe state dict, then wraps it with the class path -2. `load_object(data)` imports the class, creates an empty instance, and calls `__setstate__(state)` to reconstruct it - -### Default behavior (zero boilerplate) +1. `dump_object(obj)` calls `obj.__getstate__()` to get a JSON-safe state + dict, then wraps it with the fully-qualified class path. +2. `load_object(data)` imports the class, creates an empty instance via + `cls.__new__`, and calls `obj.__setstate__(state)` to reconstruct it. -Base classes (`Endpoint`, `Tokenizer`, `Callback`) provide default `__getstate__`/`__setstate__` implementations that: +### The `Serializable` mixin (zero boilerplate) -- **`__getstate__`**: Introspects `__init__` parameters, matches them to instance attributes (`self.name` or `self._name`), and returns only what's needed to recreate the object -- **`__setstate__`**: Calls `__init__(**state)` which recreates runtime resources (clients, connections) from the config - -This means most subclasses work without any custom serialization code: +Inherit from `Serializable` to get automatic `__getstate__`/`__setstate__`: ```python -class MyEndpoint(Endpoint): +from llmeter.serialization import Serializable + +class MyEndpoint(Serializable, Endpoint): def __init__(self, model_id: str, region: str = "us-east-1"): super().__init__(endpoint_name=model_id, model_id=model_id, provider="custom") self.region = region self._client = create_client(region) # runtime — not serialized - -# Just works: -data = dump_object(MyEndpoint(model_id="my-model")) -restored = load_object(data) # _client recreated via __init__ ``` +The mixin introspects `__init__` parameters, matches them to instance +attributes (`self.name` or `self._name`), and serializes only what's needed to +recreate the object. Private attributes (prefixed with `_`) that don't +correspond to `__init__` params are skipped — they're assumed to be derived +runtime state. + +Nested `Serializable` objects are recursively handled: `__getstate__` wraps +them via `dump_object`, and `__setstate__` restores them via `load_object`. + ### When to override Override `__getstate__`/`__setstate__` only when: -- An `__init__` parameter is consumed without being stored (uncommon) +- An `__init__` parameter is consumed without being stored - Reconstruction needs special logic beyond `__init__(**state)` - You want to exclude large transient data from persistence ```python -class SpecialEndpoint(Endpoint): +class SpecialEndpoint(Serializable, Endpoint): def __getstate__(self) -> dict: - # Only save what matters — skip the 10MB cache return {"model_id": self.model_id, "region": self.region} def __setstate__(self, state: dict): @@ -93,55 +131,59 @@ class SpecialEndpoint(Endpoint): ## Callback persistence -All callbacks support `save_to_file()` / `load_from_file()` via this protocol: +All callbacks support `save_to_file()` / `load_from_file()`: ```python from llmeter.callbacks.base import Callback from llmeter.callbacks.mlflow import MlflowCallback -# Save cb = MlflowCallback(step=5, nested=True) cb.save_to_file("/tmp/callback.json") -# Load (polymorphic — detects the type automatically) +# Polymorphic load — detects the type from __llmeter_class__ restored = Callback.load_from_file("/tmp/callback.json") # → MlflowCallback(step=5, nested=True) ``` ## Runner config persistence -`Runner.save()` and `Runner.load()` use `dump_object`/`load_object` for all callable fields: +`_RunConfig.save()` and `_RunConfig.load()` use `dump_object`/`load_object` +for all callable fields (endpoint, tokenizer, callbacks): ```python runner = Runner(endpoint=BedrockConverse(...), callbacks=[MlflowCallback(step=1)]) runner.save(output_path="/tmp/run") -# Saved JSON includes: -# {"endpoint": {"__llmeter_class__": "...BedrockConverse", "__llmeter_state__": {...}}, -# "tokenizer": {"__llmeter_class__": "...DummyTokenizer", "__llmeter_state__": {}}, -# "callbacks": [{"__llmeter_class__": "...MlflowCallback", "__llmeter_state__": {"step": 1, "nested": false}}]} - # Full reconstruction: -restored = Runner.load("/tmp/run") +restored = _RunConfig.load("/tmp/run") ``` -## Security - -`load_object` will import and instantiate whatever class path is in the `_class` field. This is the same trust model as Python's `pickle` — **never load configs from untrusted sources**. +When loading, the runner detects both formats: -This is appropriate for LLMeter's use case: developer-generated configs stored on local disk or controlled cloud storage (S3 with IAM). +- **New format** (`__llmeter_class__` key) — uses `load_object` +- **Legacy format** (`endpoint_type` / `tokenizer_module` keys) — uses + `Endpoint.load()` / `Tokenizer.load()` for backward compatibility ## Dataclass compatibility -`@dataclass` classes work seamlessly with this system. The `default_getstate` introspects the `__init__` that `@dataclass` generates: +`@dataclass` classes work seamlessly with `Serializable`. The mixin +introspects the `__init__` that `@dataclass` generates: ```python @dataclass -class InputTokens(RequestCostDimensionBase): +class InputTokens(Serializable): price_per_million: float granularity: int = 1 -# __getstate__ and __setstate__ inherited from base — no custom code needed data = dump_object(InputTokens(price_per_million=3.0)) restored = load_object(data) ``` + +## Security + +`load_object` imports and instantiates whatever class path is found in the +`__llmeter_class__` field. This is the same trust model as Python's `pickle` — +**never load configs from untrusted sources**. + +This is appropriate for LLMeter's use case: developer-generated configs stored +on local disk or controlled cloud storage (S3 with IAM). diff --git a/llmeter/callbacks/base.py b/llmeter/callbacks/base.py index a4a91a3..479f196 100644 --- a/llmeter/callbacks/base.py +++ b/llmeter/callbacks/base.py @@ -11,10 +11,9 @@ from upath.types import ReadablePathLike, WritablePathLike from ..endpoints.base import InvocationResponse -from ..json_utils import llmeter_default_serializer from ..results import Result from ..runner import _RunConfig -from ..serialization import Serializable, dump_object, load_object +from ..serialization import Serializable, dump_object, json_default, load_object from ..utils import ensure_path @@ -88,7 +87,7 @@ def save_to_file(self, path: WritablePathLike) -> None: path.parent.mkdir(parents=True, exist_ok=True) data = dump_object(self) with path.open("w") as f: - json.dump(data, f, indent=4, default=llmeter_default_serializer) + json.dump(data, f, indent=4, default=json_default) @staticmethod @final diff --git a/llmeter/callbacks/cost/model.py b/llmeter/callbacks/cost/model.py index 4ba1601..b409672 100644 --- a/llmeter/callbacks/cost/model.py +++ b/llmeter/callbacks/cost/model.py @@ -120,9 +120,6 @@ async def calculate_run_cost( result._update_contributed_stats(stats) return run_cost - async def before_invoke(self, payload: dict) -> None: - pass - async def after_invoke(self, response: InvocationResponse) -> None: await self.calculate_request_cost(response, save=True) diff --git a/llmeter/endpoints/base.py b/llmeter/endpoints/base.py index e6f2dda..8cdb362 100644 --- a/llmeter/endpoints/base.py +++ b/llmeter/endpoints/base.py @@ -21,8 +21,14 @@ from upath import UPath as Path from upath.types import ReadablePathLike, WritablePathLike -from ..json_utils import llmeter_bytes_decoder, llmeter_default_serializer -from ..serialization import Serializable +from ..serialization import ( + Serializable, + bytes_decoder, + dump_object, + json_default, + load_object, + str_to_datetime, +) from ..utils import ensure_path logger = logging.getLogger(__name__) @@ -94,16 +100,16 @@ def from_json(cls, json_str: str) -> "InvocationResponse": restored = InvocationResponse.from_json(original.to_json()) ``` """ - data = json.loads(json_str, object_hook=llmeter_bytes_decoder) + data = json.loads(json_str, object_hook=bytes_decoder) rt = data.get("request_time") if rt is not None and isinstance(rt, str): - data["request_time"] = datetime.fromisoformat(rt.replace("Z", "+00:00")) + data["request_time"] = str_to_datetime(rt) return cls(**data) - def to_json(self, default=llmeter_default_serializer, **kwargs) -> str: + def to_json(self, default=json_default, **kwargs) -> str: """Serialize this response to a JSON string. - Uses [`llmeter_default_serializer`][llmeter.json_utils.llmeter_default_serializer] by + Uses [`json_default`][llmeter.serialization.json_default] by default, which handles `bytes`, `datetime`, `PathLike`, and other common non-serializable types. @@ -152,7 +158,7 @@ def to_dict(self) -> dict: `datetime` comparisons and arithmetic. For JSON output, use [`to_json`][llmeter.endpoints.base.InvocationResponse.to_json], (which - delegates to [`llmeter_default_serializer`][llmeter.json_utils.llmeter_default_serializer] + delegates to [`json_default`][llmeter.serialization.json_default] by default, for non-JSON-serializable data types). Returns: @@ -475,22 +481,19 @@ def __subclasshook__(cls, C: type) -> bool: return NotImplemented def save(self, output_path: WritablePathLike) -> Path: - """ - Save the endpoint configuration to a JSON file. - - This method serializes the endpoint's configuration (excluding private attributes) - to a JSON file at the specified path. + """Save the endpoint configuration to a JSON file. Args: output_path (str | UPath): The path where the configuration file will be saved. Returns: - None + Path: The path the file was written to. """ output_path = ensure_path(output_path) output_path.parent.mkdir(parents=True, exist_ok=True) + data = dump_object(self) with output_path.open("w") as f: - json.dump(self, f, indent=4, default=llmeter_default_serializer) + json.dump(data, f, indent=4, default=json_default) return output_path def to_dict(self) -> dict: @@ -506,24 +509,19 @@ def to_dict(self) -> dict: @classmethod def load_from_file(cls, input_path: ReadablePathLike) -> "Endpoint": - """ - Load an endpoint configuration from a JSON file. - - This class method reads a JSON file containing an endpoint configuration, - determines the appropriate endpoint class, and instantiates it with the - loaded configuration. + """Load an endpoint configuration from a JSON file. Args: - input_path (str|UPath): The path to the JSON configuration file. + input_path (str | UPath): The path to the JSON configuration file. Returns: - Endpoint: An instance of the appropriate endpoint class, initialized - with the configuration from the file. + Endpoint: An instance of the appropriate endpoint class. """ - input_path = ensure_path(input_path) with input_path.open("r") as f: data = json.load(f) + if "__llmeter_class__" in data: + return load_object(data) endpoint_type = data.pop("endpoint_type") endpoint_module = importlib.import_module("llmeter.endpoints") endpoint_class = getattr(endpoint_module, endpoint_type) @@ -531,19 +529,17 @@ def load_from_file(cls, input_path: ReadablePathLike) -> "Endpoint": @classmethod def load(cls, endpoint_config: dict) -> "Endpoint": # type: ignore - """ - Load an endpoint configuration from a dictionary. + """Load an endpoint from a legacy ``{"endpoint_type": ...}`` dictionary. - This class method reads a dictionary containing an endpoint configuration, - determines the appropriate endpoint class, and instantiates it with the - loaded configuration. + .. deprecated:: + New code should use :func:`~llmeter.serialization.load_object` with + dicts produced by :func:`~llmeter.serialization.dump_object`. Args: - endpoint_config (Dict): A dictionary containing the endpoint configuration. + endpoint_config: Dictionary with at minimum an ``endpoint_type`` key. Returns: - Endpoint: An instance of the appropriate endpoint class, initialized - with the configuration from the dictionary. + Endpoint: An instance of the appropriate endpoint class. """ endpoint_type = endpoint_config.pop("endpoint_type") endpoint_module = importlib.import_module("llmeter.endpoints") diff --git a/llmeter/json_utils.py b/llmeter/json_utils.py deleted file mode 100644 index 08e8bb7..0000000 --- a/llmeter/json_utils.py +++ /dev/null @@ -1,110 +0,0 @@ -# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. -# SPDX-License-Identifier: Apache-2.0 -"""JSON encoding and decoding helpers used across LLMeter. - -Provides a ``default``-compatible serializer function and a matching decoder hook -for round-tripping binary content (``bytes``) through JSON via base64 marker -objects, while also handling ``datetime``, ``os.PathLike``, and objects that -implement ``to_dict()``. - -Example:: - - import json - from llmeter.json_utils import llmeter_default_serializer, llmeter_bytes_decoder - - payload = {"image": {"bytes": b"\\xff\\xd8\\xff\\xe0"}} - - # Serialize - encoded = json.dumps(payload, default=llmeter_default_serializer) - - # Deserialize (bytes are restored automatically) - decoded = json.loads(encoded, object_hook=llmeter_bytes_decoder) - assert decoded == payload -""" - -import base64 -import os -from datetime import date, datetime, time, timezone -from typing import Any - -from upath import UPath as Path - - -def llmeter_default_serializer(obj: Any) -> Any: - """Serialize a single non-JSON-serializable object. - - Intended for use as the ``default`` argument to :func:`json.dumps` or - :func:`json.dump`. - - Type handling (checked in order): - - * Objects with a ``to_dict()`` method — delegates to that method. - * ``bytes`` — wrapped in a ``{"__llmeter_bytes__": ""}`` marker so - that :func:`llmeter_bytes_decoder` can restore them on the way back. - * ``datetime`` — converted to a UTC ISO-8601 string with a ``Z`` suffix. - * ``date`` / ``time`` — converted via ``.isoformat()``. - * ``os.PathLike`` — converted to a POSIX path string. - * Anything else — ``str()`` fallback (returns ``None`` if that also fails). - - Args: - obj: The object that the default encoder could not handle. - - Returns: - A JSON-serializable representation of *obj*. - - Example:: - - >>> import json - >>> from llmeter.json_utils import llmeter_default_serializer - >>> json.dumps({"ts": datetime(2024, 1, 1)}, default=llmeter_default_serializer) - '{"ts": "2024-01-01T00:00:00"}' - """ - if hasattr(obj, "to_dict") and callable(obj.to_dict): - result = obj.to_dict() - if not isinstance(result, dict): - # This check guards against infinite recursion in case something tries to serialize a - # MagicMock object with this function (in which to_dict returns another mock) - raise TypeError( - f"{type(obj).__name__}.to_dict() returned {type(result).__name__}, expected dict" - ) - return result - if isinstance(obj, bytes): - return {"__llmeter_bytes__": base64.b64encode(obj).decode("utf-8")} - if isinstance(obj, datetime): - if obj.tzinfo is not None: - obj = obj.astimezone(timezone.utc) - return obj.isoformat(timespec="seconds").replace("+00:00", "Z") - if isinstance(obj, (date, time)): - return obj.isoformat() - if isinstance(obj, (os.PathLike, Path)): - return Path(obj).as_posix() - try: - return str(obj) - except Exception: - return None - - -def llmeter_bytes_decoder(dct: dict) -> dict | bytes: - """Decode ``__llmeter_bytes__`` marker objects back to ``bytes``. - - Intended for use as the ``object_hook`` argument to :func:`json.loads` or - :func:`json.load`. Marker objects produced by :func:`llmeter_default_serializer` - are detected and converted back to ``bytes``; all other dicts pass through - unchanged. - - Args: - dct: A dictionary produced by the JSON parser. - - Returns: - The original ``bytes`` if *dct* is a marker object, otherwise *dct* unchanged. - - Example:: - - >>> import json - >>> from llmeter.json_utils import llmeter_bytes_decoder - >>> json.loads('{"__llmeter_bytes__": "/9j/4A=="}', object_hook=llmeter_bytes_decoder) - b'\\xff\\xd8\\xff\\xe0' - """ - if "__llmeter_bytes__" in dct and len(dct) == 1: - return base64.b64decode(dct["__llmeter_bytes__"]) - return dct diff --git a/llmeter/prompt_utils.py b/llmeter/prompt_utils.py index 2408c23..fd927e5 100644 --- a/llmeter/prompt_utils.py +++ b/llmeter/prompt_utils.py @@ -11,7 +11,7 @@ from upath import UPath as Path from upath.types import ReadablePathLike, WritablePathLike -from .json_utils import llmeter_bytes_decoder, llmeter_default_serializer +from .serialization import bytes_decoder, json_default from .tokenizers import DummyTokenizer, Tokenizer from .utils import DeferredError, ensure_path @@ -419,7 +419,7 @@ def load_payloads( This function reads JSON data from either a single file or multiple files in a directory. It supports both .json and .jsonl file formats. Binary content - (bytes objects) that were serialized using ``llmeter_default_serializer`` are automatically + (bytes objects) that were serialized using ``json_default`` are automatically restored during deserialization. Binary Content Handling: @@ -514,12 +514,12 @@ def _load_data_file(file: Path) -> Iterator[dict]: if not line.strip(): continue yield json.loads( - line.strip(), object_hook=llmeter_bytes_decoder + line.strip(), object_hook=bytes_decoder ) except json.JSONDecodeError as e: print(f"Error decoding JSON in {file}: {e}") else: # Assume it's a regular JSON file - yield json.load(f, object_hook=llmeter_bytes_decoder) + yield json.load(f, object_hook=bytes_decoder) except IOError as e: print(f"Error reading file {file}: {e}") except json.JSONDecodeError as e: @@ -638,5 +638,5 @@ def save_payloads( payloads = [payloads] with output_file_path.open(mode="w") as f: for payload in payloads: - f.write(json.dumps(payload, default=llmeter_default_serializer) + "\n") + f.write(json.dumps(payload, default=json_default) + "\n") return output_file_path diff --git a/llmeter/results.py b/llmeter/results.py index b40e9aa..a9a5479 100644 --- a/llmeter/results.py +++ b/llmeter/results.py @@ -13,7 +13,7 @@ from upath.types import ReadablePathLike, WritablePathLike from .endpoints import InvocationResponse -from .json_utils import llmeter_default_serializer +from .serialization import json_default, str_to_datetime from .utils import ensure_path, summary_stats_from_list logger = logging.getLogger(__name__) @@ -50,7 +50,7 @@ class Result: end_time: datetime | None = None def __str__(self): - return json.dumps(self.stats, indent=4, default=llmeter_default_serializer) + return json.dumps(self.stats, indent=4, default=json_default) def __post_init__(self): """Initialize the Result instance.""" @@ -72,7 +72,7 @@ def _parse_datetime_fields(cls, d: dict) -> None: val = d.get(f.name) if val and isinstance(val, str): try: - d[f.name] = datetime.fromisoformat(val.replace("Z", "+00:00")) + d[f.name] = str_to_datetime(val) except ValueError: pass @@ -127,27 +127,31 @@ def save(self, output_path: WritablePathLike | None = None) -> None: summary_path = output_path / "summary.json" stats_path = output_path / "stats.json" - with summary_path.open("w") as f, stats_path.open("w") as s: - f.write(self.to_json(indent=4)) - s.write( - json.dumps(self.stats, indent=4, default=llmeter_default_serializer) + with summary_path.open("w") as f: + json.dump( + {k: o for k, o in asdict(self).items() if k not in ["responses", "stats"]}, + f, + default=json_default, + indent=4, ) + with stats_path.open("w") as f: + json.dump(self.stats, f, default=json_default, indent=4) responses_path = output_path / "responses.jsonl" if not responses_path.exists(): with responses_path.open("w") as f: for response in self.responses: f.write( - json.dumps(asdict(response), default=llmeter_default_serializer) + json.dumps(asdict(response), default=json_default) + "\n" ) - def to_json(self, default=llmeter_default_serializer, **kwargs) -> str: + def to_json(self, default=json_default, **kwargs) -> str: """Return the results as a JSON string. Args: default: Fallback serializer. Defaults to - :func:`~llmeter.json_utils.llmeter_default_serializer`. + :func:`~llmeter.serialization.json_default`. **kwargs: Extra keyword arguments passed to :func:`json.dumps`. """ summary = { @@ -164,9 +168,9 @@ def to_dict(self, include_responses: bool = False) -> dict: processing. For JSON output, use :meth:`to_json` which delegates to - :func:`~llmeter.json_utils.llmeter_default_serializer` for + :func:`~llmeter.serialization.json_default` for non-serializable types, or pass the dict through - ``json.dumps(result.to_dict(), default=llmeter_default_serializer)``. + ``json.dumps(result.to_dict(), default=json_default)``. Args: include_responses: If ``True``, include the full list of diff --git a/llmeter/runner.py b/llmeter/runner.py index 98f36d4..e14b3d1 100644 --- a/llmeter/runner.py +++ b/llmeter/runner.py @@ -23,7 +23,7 @@ from upath.types import ReadablePathLike, WritablePathLike from .live_display import LiveStatsDisplay -from .serialization import dump_object, load_object +from .serialization import dump_object, json_default, load_object from .utils import RunningStats, ensure_path, now_utc if TYPE_CHECKING: @@ -31,7 +31,6 @@ from .callbacks.base import Callback from .endpoints.base import Endpoint, InvocationResponse -from .json_utils import llmeter_default_serializer from .prompt_utils import load_payloads, save_payloads from .results import Result from .tokenizers import DummyTokenizer, Tokenizer @@ -146,7 +145,10 @@ def __post_init__(self, disable_client_progress_bar, disable_clients_progress_ba assert self.endpoint is not None, "Endpoint cannot be None" if isinstance(self.endpoint, dict): - self._endpoint: Endpoint = Endpoint.load(self.endpoint) + if "__llmeter_class__" in self.endpoint: + self._endpoint: Endpoint = load_object(self.endpoint) + else: + self._endpoint = Endpoint.load(self.endpoint) else: self._endpoint = self.endpoint @@ -156,7 +158,10 @@ def __post_init__(self, disable_client_progress_bar, disable_clients_progress_ba if self.tokenizer is None: self.tokenizer = DummyTokenizer() if isinstance(self.tokenizer, dict): - self._tokenizer: Tokenizer = Tokenizer.load(self.tokenizer) + if "__llmeter_class__" in self.tokenizer: + self._tokenizer: Tokenizer = load_object(self.tokenizer) + else: + self._tokenizer = Tokenizer.load(self.tokenizer) else: self._tokenizer = self.tokenizer @@ -194,10 +199,8 @@ def save( config_copy.callbacks = [dump_object(cb) for cb in self.callbacks] with run_config_path.open("w") as f: - f.write( - json.dumps( - asdict(config_copy), default=llmeter_default_serializer, indent=4 - ) + json.dump( + asdict(config_copy), f, default=json_default, indent=4 ) @classmethod diff --git a/llmeter/serialization.py b/llmeter/serialization.py index c3a72bf..ccc6d59 100644 --- a/llmeter/serialization.py +++ b/llmeter/serialization.py @@ -2,127 +2,161 @@ # SPDX-License-Identifier: Apache-2.0 """Unified serialization for all LLMeter objects. -Two operations, two purposes: +This module is the single source of truth for serialization in LLMeter. -- ``to_dict(obj)`` - Python dict with native types preserved (datetime stays datetime). - Use for in-memory access, stats computation, jmespath queries. -- ``serialize(obj)`` - JSON-safe dict (datetimes become strings, etc). - Use for persistence to disk, JSON output, network transfer. +Key components: -For objects with runtime state (endpoints, tokenizers, callbacks), ``dump_object`` -and ``load_object`` provide full round-trip persistence using the -``__getstate__``/``__setstate__`` protocol:: +- :class:`Serializable` — Mixin that gives any class automatic + ``__getstate__``/``__setstate__`` by introspecting ``__init__``. - data = dump_object(endpoint) # -> {"__llmeter_class__": "...", "__llmeter_state__": {...}} - restored = load_object(data) # -> new Endpoint instance +- :func:`dump_object` / :func:`load_object` — Full round-trip persistence + using a ``{"__llmeter_class__": ..., "__llmeter_state__": ...}`` envelope. + +- :func:`json_default` — A ``json.dumps``-compatible *default* handler for + ``bytes``, ``datetime``, ``PathLike``. + +- :func:`bytes_decoder` — A ``json.loads``-compatible *object_hook* that + restores binary content encoded by :func:`json_default`. + +- :func:`datetime_to_str` / :func:`str_to_datetime` — Standardized + datetime ↔ string conversion (UTC ISO-8601 with ``Z`` suffix). .. warning:: Security - ``load_object`` will import and instantiate any class path found in the + :func:`load_object` imports and instantiates whatever class path is in the ``__llmeter_class__`` field. Do not load configs from untrusted sources. """ +import base64 import importlib import inspect -import json as _json import logging +import os from dataclasses import asdict, is_dataclass +from datetime import date, datetime, time, timezone from typing import Any +from upath import UPath as Path + logger = logging.getLogger(__name__) # --------------------------------------------------------------------------- -# Serializable mixin +# Datetime helpers # --------------------------------------------------------------------------- -class Serializable: - """Mixin that provides ``__getstate__``/``__setstate__`` via introspection. +def datetime_to_str(dt: datetime) -> str: + """Convert a datetime to a UTC ISO-8601 string with ``Z`` suffix. - Inherit from this to get automatic serialization support. Works with - plain classes, ``@dataclass``, and any class whose ``__init__`` parameters - match instance attributes. + Timezone-aware datetimes are converted to UTC first. Naive datetimes are + serialized as-is (assumed UTC). + """ + if dt.tzinfo is not None: + dt = dt.astimezone(timezone.utc) + return dt.isoformat(timespec="seconds").replace("+00:00", "Z") + + +def str_to_datetime(s: str) -> datetime: + """Parse an ISO-8601 string (with optional ``Z`` suffix) to a datetime.""" + return datetime.fromisoformat(s.replace("Z", "+00:00")) - Example:: - class MyEndpoint(Serializable, Endpoint): - ... # __getstate__/__setstate__ inherited — no boilerplate needed +# --------------------------------------------------------------------------- +# JSON helpers +# --------------------------------------------------------------------------- + + +def json_default(obj: Any) -> Any: + """Fallback serializer for :func:`json.dumps`. + + Handles ``bytes`` (→ base64 marker), ``datetime``, ``date``/``time``, + ``os.PathLike``, and falls back to ``str()``. """ + if isinstance(obj, bytes): + return {"__llmeter_bytes__": base64.b64encode(obj).decode("utf-8")} + if isinstance(obj, datetime): + return datetime_to_str(obj) + if isinstance(obj, (date, time)): + return obj.isoformat() + if isinstance(obj, (os.PathLike, Path)): + return Path(obj).as_posix() + try: + return str(obj) + except Exception: + return None - def __getstate__(self) -> dict: - return default_getstate(self) - def __setstate__(self, state: dict) -> None: - default_setstate(self, state) +def bytes_decoder(dct: dict) -> dict | bytes: + """Object hook for :func:`json.loads` that restores ``__llmeter_bytes__`` markers.""" + if "__llmeter_bytes__" in dct and len(dct) == 1: + return base64.b64decode(dct["__llmeter_bytes__"]) + return dct # --------------------------------------------------------------------------- -# Public API +# Serializable mixin # --------------------------------------------------------------------------- -def to_dict(obj: Any) -> dict: - """Convert any LLMeter object to a dict with native Python types preserved. +class Serializable: + """Mixin providing automatic ``__getstate__``/``__setstate__`` via __init__ introspection. + + Works with plain classes, ``@dataclass``, and any class whose ``__init__`` + parameters correspond to instance attributes (``self.x`` or ``self._x``). - Dispatch order: - 1. Objects with custom ``__getstate__`` -> ``__getstate__()`` - 2. Dataclasses -> ``dataclasses.asdict()`` - 3. Plain objects -> filtered ``__dict__`` (public attrs only) + Nested :class:`Serializable` objects are recursively persisted via + :func:`dump_object` / :func:`load_object`. """ - has_custom_getstate = ( - hasattr(obj, "__getstate__") - and type(obj).__getstate__ is not object.__getstate__ - ) - if has_custom_getstate: - return obj.__getstate__() - if is_dataclass(obj) and not isinstance(obj, type): - return asdict(obj) - if hasattr(obj, "__dict__"): - return {k: v for k, v in obj.__dict__.items() if not k.startswith("_")} - return dict(obj) + def __getstate__(self) -> dict: + sig = inspect.signature(self.__init__) + state = {} + for name, param in sig.parameters.items(): + if name in ("self", "args", "kwargs"): + continue + if param.kind in (param.VAR_POSITIONAL, param.VAR_KEYWORD): + continue + if hasattr(self, name): + state[name] = _serialize_value(getattr(self, name)) + elif hasattr(self, f"_{name}"): + state[name] = _serialize_value(getattr(self, f"_{name}")) + return state -def serialize(obj: Any) -> dict: - """Convert any LLMeter object to a JSON-safe dict for persistence. + def __setstate__(self, state: dict) -> None: + deserialized = {k: _deserialize_value(v) for k, v in state.items()} + self.__init__(**deserialized) - Same dispatch as ``to_dict`` but intended for JSON output. - Objects should ensure their ``__getstate__`` returns JSON-serializable data. - """ - has_custom_getstate = ( - hasattr(obj, "__getstate__") - and type(obj).__getstate__ is not object.__getstate__ - ) - if has_custom_getstate: - return obj.__getstate__() - if is_dataclass(obj) and not isinstance(obj, type): - return asdict(obj) - if hasattr(obj, "__dict__"): - return {k: v for k, v in obj.__dict__.items() if not k.startswith("_")} - return dict(obj) + +# --------------------------------------------------------------------------- +# Object serialization API +# --------------------------------------------------------------------------- def dump_object(obj: Any) -> dict: - """Serialize an object to a type-tagged dict for full round-trip persistence. + """Serialize an object to a type-tagged dict for round-trip persistence. - Returns ``{"__llmeter_class__": "module.ClassName", "__llmeter_state__": {...}}``. + Returns ``{"__llmeter_class__": "module.Class", "__llmeter_state__": {...}}``. """ class_path = f"{obj.__class__.__module__}.{obj.__class__.__qualname__}" - try: - state = serialize(obj) - _json.dumps(state) # verify JSON-serializable - except Exception: - logger.debug( - "dump_object: serialize(%s) returned non-serializable data", class_path - ) + if ( + hasattr(obj, "__getstate__") + and type(obj).__getstate__ is not object.__getstate__ + ): + state = obj.__getstate__() + elif is_dataclass(obj) and not isinstance(obj, type): + state = asdict(obj) + elif hasattr(obj, "__dict__"): + state = {k: v for k, v in obj.__dict__.items() if not k.startswith("_")} + else: state = {} return {"__llmeter_class__": class_path, "__llmeter_state__": state} def load_object(data: dict) -> Any: - """Restore an object from a type-tagged state dict. + """Restore an object from a type-tagged dict produced by :func:`dump_object`. - WARNING: Do not call on data from untrusted sources. + .. warning:: Do not call on data from untrusted sources. """ class_path = data["__llmeter_class__"] module_path, class_name = class_path.rsplit(".", 1) @@ -135,18 +169,12 @@ def load_object(data: dict) -> Any: # --------------------------------------------------------------------------- -# Default __getstate__ / __setstate__ implementations +# Internal helpers for recursive serialization # --------------------------------------------------------------------------- def _serialize_value(val: Any) -> Any: - """Recursively serialize a value for JSON persistence. - - - Objects with custom ``__getstate__`` → ``dump_object(val)`` - - Dicts → recurse into values - - Lists/tuples → recurse into items - - Scalars (str, int, float, bool, None) → pass through - """ + """Recursively prepare a value for JSON persistence.""" if val is None or isinstance(val, (str, int, float, bool)): return val if hasattr(val, "__getstate__") and type(val).__getstate__ is not object.__getstate__: @@ -155,18 +183,11 @@ def _serialize_value(val: Any) -> Any: return {k: _serialize_value(v) for k, v in val.items()} if isinstance(val, (list, tuple)): return [_serialize_value(item) for item in val] - # Fallback: try str return str(val) def _deserialize_value(val: Any) -> Any: - """Recursively deserialize a value from JSON persistence. - - - Dicts with ``__llmeter_class__``/``__llmeter_state__`` → ``load_object(val)`` - - Other dicts → recurse into values - - Lists → recurse into items - - Scalars → pass through - """ + """Recursively restore a value from JSON persistence.""" if val is None or isinstance(val, (str, int, float, bool)): return val if isinstance(val, dict): @@ -176,33 +197,3 @@ def _deserialize_value(val: Any) -> Any: if isinstance(val, (list, tuple)): return [_deserialize_value(item) for item in val] return val - - -def default_getstate(obj: Any) -> dict: - """Default __getstate__: infers state from __init__ signature. - - Matches __init__ parameter names to instance attributes (self.name or - self._name). Nested objects with ``__getstate__`` are recursively - serialized via ``dump_object``. Parameters without a match are omitted. - """ - sig = inspect.signature(obj.__init__) - state = {} - for name, param in sig.parameters.items(): - if name in ("self", "args", "kwargs"): - continue - if param.kind in (param.VAR_POSITIONAL, param.VAR_KEYWORD): - continue - if hasattr(obj, name): - state[name] = _serialize_value(getattr(obj, name)) - elif hasattr(obj, f"_{name}"): - state[name] = _serialize_value(getattr(obj, f"_{name}")) - return state - - -def default_setstate(obj: Any, state: dict) -> None: - """Default __setstate__: deserializes nested objects then calls __init__. - - Any dict with ``__llmeter_class__``/``__llmeter_state__`` keys is restored via ``load_object``. - """ - deserialized = {k: _deserialize_value(v) for k, v in state.items()} - obj.__init__(**deserialized) diff --git a/llmeter/tokenizers.py b/llmeter/tokenizers.py index 2a3e725..74ebf18 100644 --- a/llmeter/tokenizers.py +++ b/llmeter/tokenizers.py @@ -3,15 +3,10 @@ from __future__ import annotations -import json from abc import ABC, abstractmethod from typing import Any -from upath import UPath -from upath.types import ReadablePathLike, WritablePathLike - from .serialization import Serializable -from .utils import ensure_path class Tokenizer(Serializable, ABC): @@ -47,123 +42,47 @@ def __subclasscheck__(cls, subclass): return True @staticmethod - def load_from_file(tokenizer_path: ReadablePathLike | None) -> Tokenizer: - """ - Loads a tokenizer from a file. - - Args: - tokenizer_path (ReadablePathLike): The path to the serialized tokenizer file. + def load(tokenizer_info: dict) -> "Tokenizer": + """Load a tokenizer from a legacy ``{"tokenizer_module": ...}`` dictionary. - Returns: - Tokenizer: The loaded tokenizer. - """ - if tokenizer_path is None: - return DummyTokenizer() - tokenizer_path = ensure_path(tokenizer_path) - with tokenizer_path.open("r") as f: - tokenizer_info = json.load(f) - - return _load_tokenizer_from_info(tokenizer_info) - - @staticmethod - def load(tokenizer_info: dict) -> Tokenizer: - """ - Loads a tokenizer from a dictionary. + This supports configs saved before the unified ``dump_object``/``load_object`` + serialization was introduced. New code should use + :func:`~llmeter.serialization.load_object` instead. Args: - tokenizer_info (Dict): The tokenizer information to load. + tokenizer_info: Dictionary with at minimum a ``tokenizer_module`` key. Returns: - Tokenizer: The loaded tokenizer. + A restored Tokenizer instance. """ return _load_tokenizer_from_info(tokenizer_info) - @staticmethod - def to_dict(tokenizer: Any) -> dict: - """ - Serializes a tokenizer to a dictionary. - - Args: - tokenizer (Tokenizer): The tokenizer to serialize. - - Returns: - Dict: The serialized tokenizer. - """ - return _to_dict(tokenizer) - - -def _to_dict(tokenizer: Any) -> dict: - """ - Serializes a tokenizer to a dictionary. - - Args: - tokenizer (Tokenizer): The tokenizer to serialize. - - Returns: - Dict: The serialized tokenizer. - """ - if tokenizer.__module__.split(".")[0] == "transformers": - return {"tokenizer_module": "transformers", "name": tokenizer.name_or_path} - - if tokenizer.__module__.split(".")[0] == "tiktoken": - return {"tokenizer_module": "tiktoken", "name": tokenizer.name} - - if tokenizer.__module__.split(".")[0] == "llmeter": - return {"tokenizer_module": "llmeter"} - - raise ValueError(f"Unknown tokenizer module: {tokenizer.__module__}") - - -def save_tokenizer(tokenizer: Any, output_path: WritablePathLike) -> UPath: - """ - Save a tokenizer information to a file. - - Args: - tokenizer (Tokenizer): The tokenizer to serialize. - output_path (WritablePathLike): The path to save the serialized tokenizer to. - - Returns: - UPath: The path to the serialized tokenizer file. - """ - tokenizer_info = _to_dict(tokenizer) - - output_path = ensure_path(output_path) - output_path.parent.mkdir(parents=True, exist_ok=True) - with output_path.open("w") as f: - json.dump(tokenizer_info, f) - - return output_path - def _load_tokenizer_from_info(tokenizer_info: dict) -> Tokenizer: - """ - Loads a tokenizer from a file. + """Instantiate a tokenizer from a legacy info dict. - Args: - tokenizer_info (Dict): The tokenizer information to load. - - Returns: - Tokenizer: The loaded tokenizer. + Supports ``"transformers"``, ``"tiktoken"``, and ``"llmeter"`` modules. """ - if tokenizer_info["tokenizer_module"] == "transformers": + module = tokenizer_info["tokenizer_module"] + + if module == "transformers": from transformers import AutoTokenizer return AutoTokenizer.from_pretrained(tokenizer_info["name"]) # type: ignore - if tokenizer_info["tokenizer_module"] == "tiktoken": + if module == "tiktoken": from tiktoken import get_encoding return get_encoding(tokenizer_info["name"]) # type: ignore - if tokenizer_info["tokenizer_module"] == "llmeter": + if module == "llmeter": return DummyTokenizer() - raise ValueError(f"Unknown tokenizer module: {tokenizer_info['tokenizer_module']}") + raise ValueError(f"Unknown tokenizer module: {module}") class DummyTokenizer(Tokenizer): - """ - A dummy tokenizer that splits the input text on whitespace and returns the tokens as is. + """A dummy tokenizer that splits the input text on whitespace and returns the tokens as is. This tokenizer will generally under-estimate token counts in English and latin languages (where words comprise more than one token on average), and will give very poor results for languages @@ -175,8 +94,8 @@ class DummyTokenizer(Tokenizer): def __init__(self, *args, **kwargs): pass - def encode(self, text: str): + def encode(self, text: str) -> list[str]: return [k for k in text.split()] - def decode(self, tokens: list[str]): + def decode(self, tokens: list[str]) -> str: return " ".join(k for k in tokens) diff --git a/mkdocs.yml b/mkdocs.yml index 27fb90b..835ab3d 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -42,6 +42,7 @@ nav: - user_guide/callbacks/cache_buster.md - user_guide/callbacks/cost.md - user_guide/callbacks/mlflow.md + - Serialization: serialization.md - API Reference: - reference/index.md - callbacks: @@ -55,7 +56,6 @@ nav: - reference/callbacks/cost/providers/index.md - sagemaker: reference/callbacks/cost/providers/sagemaker.md - results: reference/callbacks/cost/results.md - - serde: reference/callbacks/cost/serde.md - mlflow: reference/callbacks/mlflow.md - endpoints: - reference/endpoints/index.md @@ -68,7 +68,7 @@ nav: - openai_response: reference/endpoints/openai_response.md - sagemaker: reference/endpoints/sagemaker.md - experiments: reference/experiments.md - - json_utils: reference/json_utils.md + - serialization: reference/serialization.md - plotting: - reference/plotting/index.md - prompt_utils: reference/prompt_utils.md diff --git a/tests/unit/callbacks/cost/providers/test_sagemaker.py b/tests/unit/callbacks/cost/providers/test_sagemaker.py index 10fac84..70d8cf5 100644 --- a/tests/unit/callbacks/cost/providers/test_sagemaker.py +++ b/tests/unit/callbacks/cost/providers/test_sagemaker.py @@ -11,7 +11,7 @@ from moto import mock_aws # Local Dependencies: -from llmeter.serialization import to_dict, dump_object, load_object +from llmeter.serialization import dump_object, load_object from llmeter.callbacks.cost.providers.sagemaker import ( SageMakerRTEndpointCompute, SageMakerRTEndpointStorage, @@ -44,7 +44,7 @@ async def test_real_time_endpoint_compute_explicit_price(): result.total_test_time = None assert await dim.calculate(result) is None - dim_ser = to_dict(dim) + dim_ser = dim.__getstate__() assert dim_ser == { "instance_count": 1, "instance_type": "ml.doesnotexist", @@ -237,7 +237,7 @@ async def test_real_time_endpoint_storage_explicit_price(): result.total_test_time = None assert await dim.calculate(result) is None - dim_ser = to_dict(dim) + dim_ser = dim.__getstate__() assert dim_ser == { "gbs_provisioned": 5, "price_per_gb_hour": 9, diff --git a/tests/unit/callbacks/cost/test_dimensions.py b/tests/unit/callbacks/cost/test_dimensions.py index b9a167b..6f16ee0 100644 --- a/tests/unit/callbacks/cost/test_dimensions.py +++ b/tests/unit/callbacks/cost/test_dimensions.py @@ -2,7 +2,6 @@ # SPDX-License-Identifier: Apache-2.0 import pytest -from llmeter.serialization import to_dict from llmeter.callbacks.cost.dimensions import ( EndpointTime, InputTokens, @@ -47,7 +46,7 @@ async def test_cost_per_input_token(): err_response = InvocationResponse.error_output() assert await dim_valid.calculate(err_response) is None - assert to_dict(dim_valid) == { + assert dim_valid.__getstate__() == { "price_per_million": 30, "granularity": 10, } @@ -72,7 +71,7 @@ async def test_cost_per_output_token(): err_response = InvocationResponse.error_output() assert await dim_valid.calculate(err_response) is None - assert to_dict(dim_valid) == { + assert dim_valid.__getstate__() == { "price_per_million": 40, "granularity": 10, } @@ -101,7 +100,7 @@ async def test_cost_per_hour(): result.total_test_time = None assert await dim_valid.calculate(result) is None - assert to_dict(dim_valid) == { + assert dim_valid.__getstate__() == { "price_per_hour": 30, "granularity_secs": 60, } diff --git a/tests/unit/callbacks/cost/test_model.py b/tests/unit/callbacks/cost/test_model.py index 65c6371..1a32f08 100644 --- a/tests/unit/callbacks/cost/test_model.py +++ b/tests/unit/callbacks/cost/test_model.py @@ -5,7 +5,7 @@ import pytest -from llmeter.serialization import to_dict, dump_object, load_object +from llmeter.serialization import dump_object, load_object from llmeter.callbacks.cost.model import CostModel from llmeter.callbacks.cost.results import CalculatedCostWithDimensions @@ -26,8 +26,8 @@ def test_cost_model_serialization(): assert restored.request_dims["TokensIn"].price_per_million == 30 assert restored.run_dims["ComputeSeconds"].price_per_hour == 50 - # to_dict produces a plain dict representation - d = to_dict(model) + # __getstate__ produces a plain dict representation + d = model.__getstate__() assert "request_dims" in d assert "run_dims" in d diff --git a/tests/unit/endpoints/test_multimodal_properties.py b/tests/unit/endpoints/test_multimodal_properties.py index a5f7c3d..4b0fbbc 100644 --- a/tests/unit/endpoints/test_multimodal_properties.py +++ b/tests/unit/endpoints/test_multimodal_properties.py @@ -17,7 +17,7 @@ from hypothesis import strategies as st from llmeter.endpoints.bedrock import BedrockBase -from llmeter.json_utils import llmeter_bytes_decoder, llmeter_default_serializer +from llmeter.serialization import bytes_decoder, json_default from llmeter.prompt_utils import ( AudioContent, DocumentContent, @@ -602,11 +602,11 @@ def create_multimodal_payload(input_text, **kwargs): assert len(image_blocks) == 1 # Verify the payload can be serialized - json_str = json.dumps(payload, default=llmeter_default_serializer) + json_str = json.dumps(payload, default=json_default) assert len(json_str) > 0 # Verify it can be deserialized - deserialized = json.loads(json_str, object_hook=llmeter_bytes_decoder) + deserialized = json.loads(json_str, object_hook=bytes_decoder) assert deserialized == payload diff --git a/tests/unit/endpoints/test_openai_response_serialization.py b/tests/unit/endpoints/test_openai_response_serialization.py index fee2e94..f98d82b 100644 --- a/tests/unit/endpoints/test_openai_response_serialization.py +++ b/tests/unit/endpoints/test_openai_response_serialization.py @@ -79,11 +79,12 @@ def test_save_creates_file_with_correct_content(self, mock_openai_class): with open(saved_path, "r") as f: data = json.load(f) - # Verify content matches endpoint configuration - assert data["model_id"] == "gpt-4-turbo" - assert data["endpoint_name"] == "saved-endpoint" - assert data["provider"] == "openai" - assert data["endpoint_type"] == "OpenAIResponseEndpoint" + # Verify envelope format with correct state + assert "__llmeter_class__" in data + state = data["__llmeter_state__"] + assert state["model_id"] == "gpt-4-turbo" + assert state["endpoint_name"] == "saved-endpoint" + assert state["provider"] == "openai" @patch("llmeter.endpoints.openai_response.OpenAI") def test_load_from_file_reconstructs_endpoint(self, mock_openai_class): diff --git a/tests/unit/test_invocation_response.py b/tests/unit/test_invocation_response.py index 67ef4d9..212e909 100644 --- a/tests/unit/test_invocation_response.py +++ b/tests/unit/test_invocation_response.py @@ -7,7 +7,7 @@ from datetime import datetime, timedelta, timezone from llmeter.endpoints.base import InvocationResponse -from llmeter.json_utils import llmeter_default_serializer +from llmeter.serialization import json_default class TestToJson: @@ -72,6 +72,20 @@ def test_save_request_time_with_offset(self): parsed = json.loads(response.to_json()) assert parsed["request_time"] == "2026-01-01T04:00:00Z" + def test_to_json_passes_kwargs_to_json_dumps(self): + """to_json() forwards kwargs (e.g. indent) to json.dumps.""" + response = InvocationResponse( + response_text="Test", + input_payload={"data": b"\x01\x02\x03"}, + id="test-kwargs", + ) + json_str = response.to_json(indent=2) + # Indentation implies newlines and leading spaces in the output. + assert "\n" in json_str + assert " " in json_str + parsed = json.loads(json_str) + assert parsed["id"] == "test-kwargs" + class TestFromJson: def test_round_trip_minimal(self): @@ -134,7 +148,7 @@ def test_round_trip_all_fields(self): def test_load_request_time_with_offset(self): json_str = json.dumps( {"response_text": "hi", "request_time": "2026-01-01T14:00:00+02:00"}, - default=llmeter_default_serializer, + default=json_default, ) restored = InvocationResponse.from_json(json_str) assert isinstance(restored.request_time, datetime) diff --git a/tests/unit/test_prompt_utils.py b/tests/unit/test_prompt_utils.py index ac4ca98..72f55cb 100644 --- a/tests/unit/test_prompt_utils.py +++ b/tests/unit/test_prompt_utils.py @@ -11,7 +11,7 @@ from hypothesis import given, settings from hypothesis import strategies as st -from llmeter.json_utils import llmeter_bytes_decoder, llmeter_default_serializer +from llmeter.serialization import bytes_decoder, json_default from llmeter.prompt_utils import ( CreatePromptCollection, load_payloads, @@ -22,9 +22,9 @@ class TestLLMeterDefaultSerializer: - """Unit tests for llmeter_default_serializer function. + """Unit tests for json_default function. - These tests verify specific examples and edge cases for the llmeter_default_serializer + These tests verify specific examples and edge cases for the json_default function, complementing the property-based tests. Requirements: 1.1, 1.2, 1.3, 1.6 @@ -38,7 +38,7 @@ def test_simple_bytes_object_serialization(self): payload = {"data": b"hello world"} # Serialize using the encoder - serialized = json.dumps(payload, default=llmeter_default_serializer) + serialized = json.dumps(payload, default=json_default) # Verify it's valid JSON parsed = json.loads(serialized) @@ -78,7 +78,7 @@ def test_nested_bytes_in_dict_structure(self): } # Serialize - serialized = json.dumps(payload, default=llmeter_default_serializer) + serialized = json.dumps(payload, default=json_default) # Verify it's valid JSON parsed = json.loads(serialized) @@ -107,7 +107,7 @@ def test_empty_bytes_object(self): payload = {"empty": b""} # Serialize - serialized = json.dumps(payload, default=llmeter_default_serializer) + serialized = json.dumps(payload, default=json_default) # Verify it's valid JSON parsed = json.loads(serialized) @@ -134,7 +134,7 @@ def test_large_binary_data_1mb(self): payload = {"large_image": large_data} # Serialize - serialized = json.dumps(payload, default=llmeter_default_serializer) + serialized = json.dumps(payload, default=json_default) # Verify it's valid JSON parsed = json.loads(serialized) @@ -161,7 +161,7 @@ def test_multiple_bytes_objects_in_payload(self): } # Serialize - serialized = json.dumps(payload, default=llmeter_default_serializer) + serialized = json.dumps(payload, default=json_default) # Verify it's valid JSON parsed = json.loads(serialized) @@ -195,7 +195,7 @@ def test_bytes_in_list(self): payload = {"images": [b"image1", b"image2", b"image3"]} # Serialize - serialized = json.dumps(payload, default=llmeter_default_serializer) + serialized = json.dumps(payload, default=json_default) # Verify it's valid JSON parsed = json.loads(serialized) @@ -223,7 +223,7 @@ def test_mixed_types_with_bytes(self): } # Serialize - serialized = json.dumps(payload, default=llmeter_default_serializer) + serialized = json.dumps(payload, default=json_default) # Verify it's valid JSON parsed = json.loads(serialized) @@ -242,9 +242,9 @@ def test_mixed_types_with_bytes(self): class TestLLMeterBytesDecoder: - """Unit tests for llmeter_bytes_decoder function. + """Unit tests for bytes_decoder function. - These tests verify specific examples and edge cases for the llmeter_bytes_decoder + These tests verify specific examples and edge cases for the bytes_decoder function, complementing the property-based tests. Requirements: 2.1, 2.2, 2.3, 2.4, 6.2 @@ -260,7 +260,7 @@ def test_marker_object_decoding(self): marker = {"__llmeter_bytes__": "aGVsbG8gd29ybGQ="} # "hello world" in base64 # Decode the marker - result = llmeter_bytes_decoder(marker) + result = bytes_decoder(marker) # Verify it returns bytes assert isinstance(result, bytes) @@ -276,7 +276,7 @@ def test_non_marker_dict_passthrough(self): regular_dict = {"key": "value", "number": 42, "nested": {"data": "test"}} # Decode should return unchanged - result = llmeter_bytes_decoder(regular_dict) + result = bytes_decoder(regular_dict) # Verify it's the same dict assert result == regular_dict @@ -294,7 +294,7 @@ def test_invalid_base64_error_handling(self): # Should raise binascii.Error when trying to decode with pytest.raises(binascii.Error): - llmeter_bytes_decoder(invalid_marker) + bytes_decoder(invalid_marker) def test_multi_key_dict_with_marker_key_not_decoded(self): """Test that multi-key dict containing marker key is not decoded. @@ -311,7 +311,7 @@ def test_multi_key_dict_with_marker_key_not_decoded(self): } # Should return unchanged (not decode) - result = llmeter_bytes_decoder(multi_key_dict) + result = bytes_decoder(multi_key_dict) # Verify it's returned as-is assert result == multi_key_dict @@ -329,7 +329,7 @@ def test_empty_bytes_decoding(self): empty_marker = {"__llmeter_bytes__": ""} # Decode - result = llmeter_bytes_decoder(empty_marker) + result = bytes_decoder(empty_marker) # Verify it returns empty bytes assert isinstance(result, bytes) @@ -352,7 +352,7 @@ def test_large_binary_data_decoding(self): marker = {"__llmeter_bytes__": base64_encoded} # Decode - result = llmeter_bytes_decoder(marker) + result = bytes_decoder(marker) # Verify it matches original data assert isinstance(result, bytes) @@ -390,7 +390,7 @@ def test_nested_structure_with_marker(self): ) # Load with decoder - result = json.loads(json_str, object_hook=llmeter_bytes_decoder) + result = json.loads(json_str, object_hook=bytes_decoder) # Verify structure is preserved assert result["modelId"] == "test-model" @@ -411,7 +411,7 @@ def test_dict_without_marker_key(self): normal_dict = {"data": "value", "count": 123} # Should return unchanged - result = llmeter_bytes_decoder(normal_dict) + result = bytes_decoder(normal_dict) assert result == normal_dict assert isinstance(result, dict) @@ -431,7 +431,7 @@ def test_marker_with_special_characters(self): marker = {"__llmeter_bytes__": base64_encoded} # Decode - result = llmeter_bytes_decoder(marker) + result = bytes_decoder(marker) # Verify assert isinstance(result, bytes) @@ -1174,10 +1174,10 @@ def test_save_payloads_file_contains_valid_json_with_markers(self): else: assert "__llmeter_bytes__" in parsed["nested"]["data"] - def test_save_payloads_handles_to_dict_objects(self): - """Test that payloads with to_dict() objects are serialized correctly. + def test_save_payloads_uses_str_fallback_for_unknown_objects(self): + """Test that payloads with unknown objects fall back to str(). - Objects implementing to_dict() are handled by llmeter_default_serializer automatically. + json_default uses str() for objects without a dedicated handler. Validates: Requirements 5.1, 5.3 """ @@ -1185,8 +1185,8 @@ def test_save_payloads_handles_to_dict_objects(self): output_path = Path(tmpdir) class CustomObj: - def to_dict(self): - return {"custom": "value"} + def __str__(self): + return "custom_repr" payload = {"test": "data", "obj": CustomObj()} @@ -1197,7 +1197,7 @@ def to_dict(self): parsed = json.loads(line) assert parsed["test"] == "data" - assert parsed["obj"] == {"custom": "value"} + assert parsed["obj"] == "custom_repr" def test_save_payloads_backward_compatibility_no_bytes(self): """Test backward compatibility when payload has no bytes. @@ -1364,7 +1364,7 @@ def __init__(self, value): } # The unified encoder falls back to str() for unknown types - result = json.dumps(payload, default=llmeter_default_serializer) + result = json.dumps(payload, default=json_default) assert "test-model" in result def test_invalid_json_error_handling(self): @@ -1400,7 +1400,7 @@ def test_invalid_base64_error_handling(self): # Should raise binascii.Error when trying to decode with pytest.raises(binascii.Error): - llmeter_bytes_decoder(invalid_marker) + bytes_decoder(invalid_marker) def test_invalid_base64_in_load_payloads(self): """Test that invalid base64 in saved payload is handled during load. @@ -1428,14 +1428,14 @@ def test_invalid_base64_in_load_payloads(self): list(load_payloads(jsonl_file)) def test_save_payloads_encoder_error_propagation(self): - """Test that encoder errors from llmeter_default_serializer are propagated correctly. + """Test that encoder errors from json_default are propagated correctly. Validates: Requirements 6.4 """ with tempfile.TemporaryDirectory() as tmpdir: output_path = Path(tmpdir) - # An object whose str() raises — llmeter_default_serializer returns None for these, + # An object whose str() raises — json_default returns None for these, # which is valid JSON, so no error is raised. class FailingStr: def __str__(self): @@ -1490,7 +1490,7 @@ def __init__(self): } # The unified encoder falls back to str() for unknown types - result = json.dumps(payload, default=llmeter_default_serializer) + result = json.dumps(payload, default=json_default) assert "test-model" in result def test_bytes_serialization_with_encoding_error(self): @@ -1506,7 +1506,7 @@ def test_bytes_serialization_with_encoding_error(self): payload = {"data": all_bytes} # Should serialize without errors - serialized = json.dumps(payload, default=llmeter_default_serializer) + serialized = json.dumps(payload, default=json_default) # Should be valid JSON parsed = json.loads(serialized) @@ -1516,7 +1516,7 @@ def test_bytes_serialization_with_encoding_error(self): # Verify round-trip works - deserialized = json.loads(serialized, object_hook=llmeter_bytes_decoder) + deserialized = json.loads(serialized, object_hook=bytes_decoder) assert deserialized["data"] == all_bytes def test_empty_payload_serialization(self): @@ -1528,7 +1528,7 @@ def test_empty_payload_serialization(self): empty_payload = {} # Should serialize without errors - serialized = json.dumps(empty_payload, default=llmeter_default_serializer) + serialized = json.dumps(empty_payload, default=json_default) assert serialized == "{}" # Empty list @@ -1556,7 +1556,7 @@ def test_none_value_serialization(self): } # Should serialize without errors - serialized = json.dumps(payload, default=llmeter_default_serializer) + serialized = json.dumps(payload, default=json_default) # Should be valid JSON parsed = json.loads(serialized) @@ -1577,7 +1577,7 @@ def test_unicode_in_payload_with_bytes(self): } # Should serialize without errors - serialized = json.dumps(payload, default=llmeter_default_serializer) + serialized = json.dumps(payload, default=json_default) # Should be valid JSON parsed = json.loads(serialized) @@ -1629,7 +1629,7 @@ def test_1mb_image_serialization_performance(self): # Measure serialization time start_time = time.perf_counter() - serialized = json.dumps(payload, default=llmeter_default_serializer) + serialized = json.dumps(payload, default=json_default) end_time = time.perf_counter() # Calculate elapsed time in milliseconds @@ -1664,12 +1664,12 @@ def test_1mb_image_deserialization_performance(self): } # First serialize the payload - serialized = json.dumps(payload, default=llmeter_default_serializer) + serialized = json.dumps(payload, default=json_default) # Measure deserialization time start_time = time.perf_counter() - deserialized = json.loads(serialized, object_hook=llmeter_bytes_decoder) + deserialized = json.loads(serialized, object_hook=bytes_decoder) end_time = time.perf_counter() # Calculate elapsed time in milliseconds @@ -1704,7 +1704,7 @@ def test_serialization_no_unnecessary_copies(self): initial_size = sys.getsizeof(binary_data) # Serialize the payload - serialized = json.dumps(payload, default=llmeter_default_serializer) + serialized = json.dumps(payload, default=json_default) # Parse to verify structure parsed = json.loads(serialized) @@ -1740,11 +1740,11 @@ def test_deserialization_no_unnecessary_copies(self): payload = {"data": binary_data} # Serialize first - serialized = json.dumps(payload, default=llmeter_default_serializer) + serialized = json.dumps(payload, default=json_default) # Deserialize - deserialized = json.loads(serialized, object_hook=llmeter_bytes_decoder) + deserialized = json.loads(serialized, object_hook=bytes_decoder) # Verify the deserialized bytes match original assert deserialized["data"] == binary_data @@ -1781,13 +1781,13 @@ def test_round_trip_performance_with_multiple_images(self): # Measure serialization time start_time = time.perf_counter() - serialized = json.dumps(payload, default=llmeter_default_serializer) + serialized = json.dumps(payload, default=json_default) serialize_time = (time.perf_counter() - start_time) * 1000 # Measure deserialization time start_time = time.perf_counter() - deserialized = json.loads(serialized, object_hook=llmeter_bytes_decoder) + deserialized = json.loads(serialized, object_hook=bytes_decoder) deserialize_time = (time.perf_counter() - start_time) * 1000 # With 3 images, we allow proportionally more time (but still reasonable) @@ -1820,7 +1820,7 @@ def test_serialization_performance_scales_linearly(self): small_payload = {"data": small_data} start_time = time.perf_counter() - json.dumps(small_payload, default=llmeter_default_serializer) + json.dumps(small_payload, default=json_default) small_time = time.perf_counter() - start_time # Test with 512KB (2x size) @@ -1828,7 +1828,7 @@ def test_serialization_performance_scales_linearly(self): large_payload = {"data": large_data} start_time = time.perf_counter() - json.dumps(large_payload, default=llmeter_default_serializer) + json.dumps(large_payload, default=json_default) large_time = time.perf_counter() - start_time # Large should take roughly 2x the time (allow 3x for variance) @@ -1851,19 +1851,19 @@ def test_deserialization_performance_scales_linearly(self): # Test with 256KB small_data = os.urandom(256 * 1024) small_payload = {"data": small_data} - small_serialized = json.dumps(small_payload, default=llmeter_default_serializer) + small_serialized = json.dumps(small_payload, default=json_default) start_time = time.perf_counter() - json.loads(small_serialized, object_hook=llmeter_bytes_decoder) + json.loads(small_serialized, object_hook=bytes_decoder) small_time = time.perf_counter() - start_time # Test with 512KB (2x size) large_data = os.urandom(512 * 1024) large_payload = {"data": large_data} - large_serialized = json.dumps(large_payload, default=llmeter_default_serializer) + large_serialized = json.dumps(large_payload, default=json_default) start_time = time.perf_counter() - json.loads(large_serialized, object_hook=llmeter_bytes_decoder) + json.loads(large_serialized, object_hook=bytes_decoder) large_time = time.perf_counter() - start_time # Large should take roughly 2x the time (allow 3x for variance) diff --git a/tests/unit/test_property_save_load.py b/tests/unit/test_property_save_load.py index 990333c..b5198fd 100644 --- a/tests/unit/test_property_save_load.py +++ b/tests/unit/test_property_save_load.py @@ -18,7 +18,7 @@ from llmeter.prompt_utils import load_payloads, save_payloads from llmeter.results import Result from llmeter.runner import _RunConfig -from llmeter.tokenizers import DummyTokenizer, Tokenizer, save_tokenizer +from llmeter.tokenizers import DummyTokenizer # Custom strategies @@ -88,65 +88,36 @@ def valid_result(draw): # Tokenizer save/load property tests class TestTokenizerSaveLoadProperties: - """Property-based tests for tokenizer serialization.""" + """Property-based tests for tokenizer serialization via dump_object/load_object.""" @given(st.text(min_size=0, max_size=500)) @settings(deadline=None) - def test_save_load_tokenizer_roundtrip(self, text): - """Saving and loading tokenizer should preserve behavior.""" - with tempfile.TemporaryDirectory() as tmpdir: - original = DummyTokenizer() - output_path = Path(tmpdir) / "tokenizer.json" + def test_dump_load_tokenizer_roundtrip(self, text): + """dump_object/load_object should preserve tokenizer behavior.""" + from llmeter.serialization import dump_object, load_object - # Save - saved_path = save_tokenizer(original, output_path) - assert saved_path.exists() + original = DummyTokenizer() - # Load - loaded = Tokenizer.load_from_file(saved_path) + # Serialize and restore + data = dump_object(original) + loaded = load_object(data) - # Should encode the same way - original_tokens = original.encode(text) - loaded_tokens = loaded.encode(text) - assert original_tokens == loaded_tokens + # Should encode the same way + original_tokens = original.encode(text) + loaded_tokens = loaded.encode(text) + assert original_tokens == loaded_tokens - def test_save_tokenizer_creates_valid_json(self): - """Saved tokenizer should be valid JSON.""" - with tempfile.TemporaryDirectory() as tmpdir: - tokenizer = DummyTokenizer() - output_path = Path(tmpdir) / "tokenizer.json" + def test_dump_object_produces_valid_structure(self): + """dump_object should produce a dict with __llmeter_class__ and __llmeter_state__.""" + from llmeter.serialization import dump_object - saved_path = save_tokenizer(tokenizer, output_path) + tokenizer = DummyTokenizer() + data = dump_object(tokenizer) - # Should be valid JSON - with open(saved_path, "r") as f: - data = json.load(f) - - assert isinstance(data, dict) - assert "tokenizer_module" in data - - @given( - st.text( - min_size=1, - max_size=100, - alphabet=st.characters( - min_codepoint=32, max_codepoint=126, blacklist_characters='\\/:*?"<>|' - ), - ) - ) - @settings(deadline=None) - def test_save_tokenizer_with_various_filenames(self, filename): - """Tokenizer should save with various valid filenames.""" - with tempfile.TemporaryDirectory() as tmpdir: - tokenizer = DummyTokenizer() - output_path = Path(tmpdir) / f"{filename}.json" - - saved_path = save_tokenizer(tokenizer, output_path) - assert saved_path.exists() - - # Should be loadable - loaded = Tokenizer.load_from_file(saved_path) - assert isinstance(loaded, DummyTokenizer) + assert isinstance(data, dict) + assert "__llmeter_class__" in data + assert "__llmeter_state__" in data + assert "DummyTokenizer" in data["__llmeter_class__"] # Payload save/load property tests @@ -393,10 +364,11 @@ def test_endpoint_save_load_roundtrip(self, model_id): saved_path = endpoint.save(output_path) assert saved_path.exists() - # Verify file contains model_id + # Verify file uses envelope format with model_id in state with open(saved_path, "r") as f: data = json.load(f) - assert data["model_id"] == model_id + assert "__llmeter_class__" in data + assert data["__llmeter_state__"]["model_id"] == model_id @given(st.text(min_size=1, max_size=50)) @settings(deadline=None) diff --git a/tests/unit/test_property_simple.py b/tests/unit/test_property_simple.py index bce566f..184434f 100644 --- a/tests/unit/test_property_simple.py +++ b/tests/unit/test_property_simple.py @@ -14,7 +14,7 @@ from llmeter.endpoints.base import InvocationResponse from llmeter.endpoints.openai import OpenAICompletionEndpoint from llmeter.prompt_utils import load_payloads, save_payloads -from llmeter.tokenizers import DummyTokenizer, _to_dict +from llmeter.tokenizers import DummyTokenizer # Custom strategies @@ -69,13 +69,15 @@ def test_dummy_tokenizer_decode_returns_string(self, tokens): @given(st.text(min_size=0, max_size=1000)) def test_tokenizer_serialization(self, text): - """Tokenizer should serialize correctly.""" + """Tokenizer should serialize via dump_object correctly.""" + from llmeter.serialization import dump_object + tokenizer = DummyTokenizer() - tokenizer_dict = _to_dict(tokenizer) + data = dump_object(tokenizer) - assert isinstance(tokenizer_dict, dict) - assert "tokenizer_module" in tokenizer_dict - assert tokenizer_dict["tokenizer_module"] == "llmeter" + assert isinstance(data, dict) + assert "__llmeter_class__" in data + assert "DummyTokenizer" in data["__llmeter_class__"] # OpenAI endpoint property tests diff --git a/tests/unit/test_results.py b/tests/unit/test_results.py index 467f51a..6af5552 100644 --- a/tests/unit/test_results.py +++ b/tests/unit/test_results.py @@ -221,10 +221,10 @@ def test_stats_property_empty_result(): def test_stats_json_serializable_with_datetimes(): - """stats dict should be JSON-serializable via llmeter_default_serializer.""" + """stats dict should be JSON-serializable via json_default.""" from datetime import datetime, timezone - from llmeter.json_utils import llmeter_default_serializer + from llmeter.serialization import json_default result = Result( responses=sample_responses_successful[:2], @@ -237,8 +237,8 @@ def test_stats_json_serializable_with_datetimes(): ) stats = result.stats # stats contains raw datetime objects; serialization is handled at the - # boundary (e.g. save()) via llmeter_default_serializer. - json_str = json.dumps(stats, default=llmeter_default_serializer) + # boundary (e.g. save()) via json_default. + json_str = json.dumps(stats, default=json_default) parsed = json.loads(json_str) assert parsed["start_time"] == "2025-01-01T12:00:00Z" assert parsed["end_time"] == "2025-01-01T12:00:01Z" @@ -497,279 +497,6 @@ def test_load_missing_responses_jsonl_with_summary(sample_result: Result, temp_d assert "time_to_first_token-p50" in loaded.stats -# Tests for llmeter_default_serializer -# Validates Requirements: 7.1, 7.2 - - -def test_llmeter_encoder_handles_bytes(): - """Test that llmeter_default_serializer handles bytes objects via base64 marker. - - Validates: Requirements 7.1, 7.2 - """ - from llmeter.json_utils import llmeter_default_serializer - - # Test bytes object encoding - payload = {"image": {"bytes": b"\xff\xd8\xff\xe0"}} - encoded = json.dumps(payload, default=llmeter_default_serializer) - - # Verify marker object is present - assert "__llmeter_bytes__" in encoded - - # Verify it can be decoded back - decoded = json.loads(encoded, object_hook=lambda d: d) - assert decoded["image"]["bytes"]["__llmeter_bytes__"] == "/9j/4A==" - - -def test_llmeter_encoder_str_fallback(): - """Test that llmeter_default_serializer falls back to str() for custom objects. - - Validates: Requirements 7.1, 7.2 - """ - from llmeter.json_utils import llmeter_default_serializer - - # Create a custom object with __str__ method - class CustomObject: - def __str__(self): - return "custom_string_representation" - - payload = {"custom": CustomObject()} - encoded = json.dumps(payload, default=llmeter_default_serializer) - - # Verify str() fallback was used - decoded = json.loads(encoded) - assert decoded["custom"] == "custom_string_representation" - - -def test_llmeter_encoder_none_on_str_failure(): - """Test that llmeter_default_serializer returns None when str() conversion fails. - - Validates: Requirements 7.1, 7.2 - """ - from llmeter.json_utils import llmeter_default_serializer - - # Create a custom object that raises exception in __str__ - class FailingObject: - def __str__(self): - raise RuntimeError("Cannot convert to string") - - payload = {"failing": FailingObject()} - encoded = json.dumps(payload, default=llmeter_default_serializer) - - # Verify None was returned - decoded = json.loads(encoded) - assert decoded["failing"] is None - - -def test_llmeter_encoder_mixed_types(): - """Test that llmeter_default_serializer handles mixed types correctly. - - Validates: Requirements 7.1, 7.2 - """ - from llmeter.json_utils import llmeter_default_serializer - - class CustomObject: - def __str__(self): - return "custom" - - # Mix of bytes, custom objects, and standard types - payload = { - "bytes_field": b"\x00\x01\x02", - "custom_field": CustomObject(), - "string_field": "normal string", - "int_field": 42, - "nested": {"bytes": b"\xff\xfe", "custom": CustomObject()}, - } - - encoded = json.dumps(payload, default=llmeter_default_serializer) - decoded = json.loads(encoded) - - # Verify bytes were encoded with marker - assert "__llmeter_bytes__" in encoded - - # Verify custom objects were converted to strings - assert decoded["custom_field"] == "custom" - assert decoded["nested"]["custom"] == "custom" - - # Verify standard types remain unchanged - assert decoded["string_field"] == "normal string" - assert decoded["int_field"] == 42 - - -# Tests for InvocationResponse.to_json -# Validates Requirements: 7.1, 7.2, 7.3 - - -def test_invocation_response_to_json_with_binary_content(): - """Test InvocationResponse.to_json with binary content in input_payload. - - Validates: Requirements 7.1, 7.2 - """ - # Create InvocationResponse with binary content in input_payload - response = InvocationResponse( - response_text="This is an image of a cat", - input_payload={ - "modelId": "anthropic.claude-3-haiku-20240307-v1:0", - "messages": [ - { - "role": "user", - "content": [ - {"text": "What is in this image?"}, - { - "image": { - "format": "jpeg", - "source": {"bytes": b"\xff\xd8\xff\xe0\x00\x10JFIF"}, - } - }, - ], - } - ], - }, - id="test-123", - time_to_first_token=0.5, - time_to_last_token=1.2, - num_tokens_input=15, - num_tokens_output=8, - ) - - # Serialize to JSON - json_str = response.to_json() - - # Verify it's valid JSON - parsed = json.loads(json_str) - - # Verify marker object is present for bytes - assert "__llmeter_bytes__" in json_str - - # Verify structure is preserved - assert parsed["response_text"] == "This is an image of a cat" - assert parsed["id"] == "test-123" - assert ( - parsed["input_payload"]["modelId"] == "anthropic.claude-3-haiku-20240307-v1:0" - ) - - # Verify bytes were encoded with marker - bytes_marker = parsed["input_payload"]["messages"][0]["content"][1]["image"][ - "source" - ]["bytes"] - assert "__llmeter_bytes__" in bytes_marker - assert bytes_marker["__llmeter_bytes__"] == "/9j/4AAQSkZJRg==" - - -def test_invocation_response_to_json_valid_json_output(): - """Test that InvocationResponse.to_json produces valid JSON. - - Validates: Requirements 7.1, 7.2 - """ - response = InvocationResponse( - response_text="Test response", - input_payload={ - "image_data": b"\x89PNG\r\n\x1a\n", - "text": "Analyze this image", - }, - id="test-456", - ) - - # Serialize to JSON - json_str = response.to_json() - - # Verify it's valid JSON (should not raise exception) - parsed = json.loads(json_str) - - # Verify all fields are present - assert "response_text" in parsed - assert "input_payload" in parsed - assert "id" in parsed - - # Verify bytes were properly encoded - assert isinstance(parsed["input_payload"]["image_data"], dict) - assert "__llmeter_bytes__" in parsed["input_payload"]["image_data"] - - -def test_invocation_response_to_json_round_trip(): - """Test round-trip serialization/deserialization with InvocationResponse. - - Validates: Requirements 7.1, 7.2, 7.3 - """ - from llmeter.json_utils import llmeter_bytes_decoder - - # Create original response with binary content - original_payload = { - "video": {"format": "mp4", "source": {"bytes": b"\x00\x00\x00\x18ftypmp42"}}, - "prompt": "Describe this video", - } - - response = InvocationResponse( - response_text="A video of a sunset", - input_payload=original_payload, - id="test-789", - time_to_first_token=1.0, - time_to_last_token=2.5, - ) - - # Serialize to JSON - json_str = response.to_json() - - # Deserialize back - parsed = json.loads(json_str, object_hook=llmeter_bytes_decoder) - - # Verify input_payload was restored correctly - assert parsed["input_payload"] == original_payload - assert isinstance(parsed["input_payload"]["video"]["source"]["bytes"], bytes) - assert ( - parsed["input_payload"]["video"]["source"]["bytes"] - == b"\x00\x00\x00\x18ftypmp42" - ) - - # Verify other fields - assert parsed["response_text"] == "A video of a sunset" - assert parsed["id"] == "test-789" - assert parsed["time_to_first_token"] == 1.0 - assert parsed["time_to_last_token"] == 2.5 - - -def test_invocation_response_to_json_no_binary_content(): - """Test InvocationResponse.to_json with no binary content (backward compatibility). - - Validates: Requirements 7.1, 7.2 - """ - response = InvocationResponse( - response_text="Simple text response", - input_payload={"modelId": "test-model", "prompt": "Hello, world!"}, - id="test-no-binary", - ) - - # Serialize to JSON - json_str = response.to_json() - - # Verify no marker objects are present - assert "__llmeter_bytes__" not in json_str - - # Verify it's valid JSON - parsed = json.loads(json_str) - assert parsed["input_payload"]["prompt"] == "Hello, world!" - - -def test_invocation_response_to_json_with_kwargs(): - """Test that InvocationResponse.to_json passes through kwargs to json.dumps. - - Validates: Requirements 7.4 - """ - response = InvocationResponse( - response_text="Test", input_payload={"data": b"\x01\x02\x03"}, id="test-kwargs" - ) - - # Test with indent parameter - json_str = response.to_json(indent=2) - - # Verify indentation is present - assert "\n" in json_str - assert " " in json_str - - # Verify it's still valid JSON - parsed = json.loads(json_str) - assert parsed["id"] == "test-kwargs" - - # ── Contributed stats round-trip ───────────────────────────────────────────── diff --git a/tests/unit/test_runner.py b/tests/unit/test_runner.py index 8d08b8e..66f4532 100644 --- a/tests/unit/test_runner.py +++ b/tests/unit/test_runner.py @@ -13,7 +13,7 @@ import llmeter.endpoints from llmeter.endpoints.base import Endpoint, InvocationResponse from llmeter.runner import Runner, _Run, _RunConfig -from llmeter.tokenizers import Tokenizer +from llmeter.tokenizers import DummyTokenizer @pytest.fixture @@ -33,11 +33,7 @@ def mock_endpoint(): @pytest.fixture def mock_tokenizer(): - with patch( - "llmeter.tokenizers.Tokenizer.to_dict", - return_value={"tokenizer_module": "mock_tokenizer"}, - ): - yield MagicMock(spec=Tokenizer) + yield DummyTokenizer() @pytest.fixture @@ -91,11 +87,9 @@ def test_runner_initialization(runner: Runner): def test_count_tokens_no_wait(runner: Runner): - # Test the tokenizer encode method directly since _count_tokens_no_wait doesn't exist - runner._tokenizer.encode.return_value = [1, 2, 3] - result = len(runner._tokenizer.encode("test text")) - assert result == 3 - runner._tokenizer.encode.assert_called_once_with("test text") + # Test the tokenizer encode method directly + result = len(runner._tokenizer.encode("test text here")) + assert result == 3 # DummyTokenizer splits on whitespace @pytest.mark.asyncio diff --git a/tests/unit/test_serialization_properties.py b/tests/unit/test_serialization_properties.py index 9452710..d38ded1 100644 --- a/tests/unit/test_serialization_properties.py +++ b/tests/unit/test_serialization_properties.py @@ -1,12 +1,12 @@ # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # SPDX-License-Identifier: Apache-2.0 -"""Property-based tests for llmeter_default_serializer and llmeter_bytes_decoder. +"""Property-based tests for json_default and bytes_decoder. This module contains property-based tests using Hypothesis to verify that -llmeter_default_serializer correctly handles all supported types (bytes, datetime, date, -time, PathLike, to_dict() objects, and str() fallback) and that -llmeter_bytes_decoder restores binary content from marker objects. +json_default correctly handles all supported types (bytes, datetime, date, +time, PathLike, and str() fallback) and that bytes_decoder restores binary +content from marker objects. Feature: json-serialization-optimization """ @@ -20,7 +20,7 @@ from hypothesis import strategies as st from hypothesis.strategies import composite -from llmeter.json_utils import llmeter_bytes_decoder, llmeter_default_serializer +from llmeter.serialization import bytes_decoder, json_default # Test infrastructure is set up and ready for property test implementation # This file will contain property-based tests for: @@ -111,7 +111,7 @@ def test_property_1_serialization_produces_valid_json_with_marker_objects( base64-encoded string value. """ # Serialize the payload - serialized = json.dumps(payload, default=llmeter_default_serializer) + serialized = json.dumps(payload, default=json_default) # Verify it's valid JSON by parsing it parsed = json.loads(serialized) @@ -171,7 +171,7 @@ def test_property_3_deserialization_restores_bytes_from_markers( marker_object = {"__llmeter_bytes__": base64_str} # Decode the marker object - decoded_bytes = llmeter_bytes_decoder(marker_object) + decoded_bytes = bytes_decoder(marker_object) # Verify the decoded bytes match the original assert isinstance(decoded_bytes, bytes), "Decoder should return bytes" @@ -195,7 +195,7 @@ def apply_decoder_recursively(obj): """Apply decoder to all dicts in the structure.""" if isinstance(obj, dict): # Apply decoder to this dict - decoded = llmeter_bytes_decoder(obj) + decoded = bytes_decoder(obj) # If it's still a dict (not converted to bytes), recurse if isinstance(decoded, dict): return {k: apply_decoder_recursively(v) for k, v in decoded.items()} @@ -245,7 +245,7 @@ def test_property_2_serialization_preserves_non_binary_structure(self, payload): # Feature: json-serialization-optimization, Property 2: Serialization preserves non-binary structure """ # Serialize the payload - serialized = json.dumps(payload, default=llmeter_default_serializer) + serialized = json.dumps(payload, default=json_default) # Parse the JSON (without decoding markers back to bytes) parsed = json.loads(serialized) @@ -326,10 +326,10 @@ def test_property_5_round_trip_serialization_preserves_data_integrity( """ # Serialize the payload - serialized = json.dumps(payload, default=llmeter_default_serializer) + serialized = json.dumps(payload, default=json_default) # Deserialize the payload - deserialized = json.loads(serialized, object_hook=llmeter_bytes_decoder) + deserialized = json.loads(serialized, object_hook=bytes_decoder) # Verify round-trip equality assert deserialized == payload, ( @@ -385,10 +385,10 @@ def test_property_6_round_trip_preserves_dictionary_key_ordering(self, payload): """ # Serialize the payload - serialized = json.dumps(payload, default=llmeter_default_serializer) + serialized = json.dumps(payload, default=json_default) # Deserialize the payload - deserialized = json.loads(serialized, object_hook=llmeter_bytes_decoder) + deserialized = json.loads(serialized, object_hook=bytes_decoder) # Helper function to verify key ordering def verify_key_ordering(original, restored, path=""): @@ -445,7 +445,7 @@ def test_property_7_non_binary_payloads_are_backward_compatible(self, payload): """ # Serialize with the new encoder serialized_with_encoder = json.dumps( - payload, default=llmeter_default_serializer + payload, default=json_default ) # Serialize with standard json.dumps @@ -453,7 +453,7 @@ def test_property_7_non_binary_payloads_are_backward_compatible(self, payload): # Verify they produce identical output assert serialized_with_encoder == serialized_standard, ( - "Serialization with llmeter_default_serializer should produce identical output " + "Serialization with json_default should produce identical output " "to standard json.dumps for payloads without bytes objects" ) @@ -535,7 +535,7 @@ def test_property_8_serialization_handles_unknown_types(self, data): payload = payload_creator(unserializable_obj) # The unified encoder falls back to str() for unknown types - result = json.dumps(payload, default=llmeter_default_serializer) + result = json.dumps(payload, default=json_default) # Result should be valid JSON parsed = json.loads(result) assert isinstance(parsed, dict) @@ -578,7 +578,7 @@ def test_property_9_deserialization_errors_are_descriptive(self, data): # Attempt to deserialize and verify it raises JSONDecodeError try: - json.loads(invalid_json, object_hook=llmeter_bytes_decoder) + json.loads(invalid_json, object_hook=bytes_decoder) # If we get here, deserialization succeeded when it shouldn't have raise AssertionError( f"Expected JSONDecodeError for invalid JSON: {invalid_json}" @@ -749,27 +749,6 @@ def verify_bytes_serialized(original, serialized_obj, path=""): ) -@composite -def to_dict_object_strategy(draw): - """Generate an object with a to_dict() method returning a JSON-safe dict.""" - inner = draw( - st.dictionaries( - keys=st.text(min_size=1, max_size=10), - values=st.one_of(st.integers(), st.text(max_size=20), st.booleans()), - max_size=5, - ) - ) - - class _Obj: - def __init__(self, d): - self._d = d - - def to_dict(self): - return self._d - - return _Obj(inner), inner - - class TestDatetimeSerializationProperties: """Property-based tests for datetime/date/time encoding.""" @@ -782,7 +761,7 @@ def test_datetime_produces_iso_string_with_z_suffix(self, dt): Naive datetimes are serialized as-is with no timezone indicator. Microseconds are truncated (the encoder uses timespec="seconds"). """ - result = json.loads(json.dumps({"v": dt}, default=llmeter_default_serializer)) + result = json.loads(json.dumps({"v": dt}, default=json_default)) assert isinstance(result["v"], str) if dt.tzinfo is not None: assert result["v"].endswith("Z") @@ -797,7 +776,7 @@ def test_datetime_produces_iso_string_with_z_suffix(self, dt): @settings(max_examples=100) def test_date_produces_iso_string(self, d): """Date values are serialized via isoformat().""" - result = json.loads(json.dumps({"v": d}, default=llmeter_default_serializer)) + result = json.loads(json.dumps({"v": d}, default=json_default)) assert result["v"] == d.isoformat() assert date.fromisoformat(result["v"]) == d @@ -805,7 +784,7 @@ def test_date_produces_iso_string(self, d): @settings(max_examples=100) def test_time_produces_iso_string(self, t): """Time values are serialized via isoformat().""" - result = json.loads(json.dumps({"v": t}, default=llmeter_default_serializer)) + result = json.loads(json.dumps({"v": t}, default=json_default)) assert result["v"] == t.isoformat() assert time.fromisoformat(result["v"]) == t @@ -817,29 +796,32 @@ class TestPathSerializationProperties: @settings(max_examples=100) def test_pathlike_produces_posix_string(self, p): """PathLike objects are serialized to POSIX path strings.""" - result = json.loads(json.dumps({"v": p}, default=llmeter_default_serializer)) + result = json.loads(json.dumps({"v": p}, default=json_default)) assert isinstance(result["v"], str) assert result["v"] == p.as_posix() -class TestToDictSerializationProperties: - """Property-based tests for to_dict() delegation.""" +class TestJsonDefaultFallback: + """Example-based tests for the str()/None fallback branch of json_default.""" + + def test_str_fallback_for_custom_object(self): + """Unknown objects fall back to their str() representation.""" + + class CustomObject: + def __str__(self): + return "custom_string_representation" + + decoded = json.loads(json.dumps({"custom": CustomObject()}, default=json_default)) + assert decoded["custom"] == "custom_string_representation" + + def test_none_when_str_conversion_fails(self): + """Objects whose __str__ raises are serialized as None.""" + + class FailingObject: + def __str__(self): + raise RuntimeError("Cannot convert to string") + + decoded = json.loads(json.dumps({"failing": FailingObject()}, default=json_default)) + assert decoded["failing"] is None - @given(to_dict_object_strategy()) - @settings(max_examples=100) - def test_to_dict_delegation_produces_expected_dict(self, obj_and_expected): - """Objects with to_dict() are serialized by calling that method.""" - obj, expected = obj_and_expected - result = json.loads(json.dumps({"v": obj}, default=llmeter_default_serializer)) - assert result["v"] == expected - @given(to_dict_object_strategy(), st.binary(min_size=1, max_size=100)) - @settings(max_examples=100) - def test_to_dict_and_bytes_coexist(self, obj_and_expected, raw_bytes): - """Payloads mixing to_dict() objects and bytes round-trip correctly.""" - obj, expected = obj_and_expected - payload = {"obj": obj, "data": raw_bytes} - serialized = json.dumps(payload, default=llmeter_default_serializer) - restored = json.loads(serialized, object_hook=llmeter_bytes_decoder) - assert restored["obj"] == expected - assert restored["data"] == raw_bytes diff --git a/tests/unit/test_tokenizers.py b/tests/unit/test_tokenizers.py index 0aceedb..6df5918 100644 --- a/tests/unit/test_tokenizers.py +++ b/tests/unit/test_tokenizers.py @@ -1,16 +1,12 @@ # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # SPDX-License-Identifier: Apache-2.0 -import json - import pytest from llmeter.tokenizers import ( DummyTokenizer, Tokenizer, _load_tokenizer_from_info, - _to_dict, - save_tokenizer, ) # Check for optional dependencies @@ -65,77 +61,16 @@ def test_dummy_tokenizer(): assert tokenizer.decode(tokens) == text -# Test Tokenizer.load_from_file -def test_load_from_file(tmp_path): - file_path = tmp_path / "tokenizer.json" - tokenizer_info = {"tokenizer_module": "llmeter"} - with open(file_path, "w") as f: - json.dump(tokenizer_info, f) - - tokenizer = Tokenizer.load_from_file(file_path) - assert isinstance(tokenizer, DummyTokenizer) - - -def test_load_from_file_none(): - tokenizer = Tokenizer.load_from_file(None) - assert isinstance(tokenizer, DummyTokenizer) - - -# Test Tokenizer.load +# Test Tokenizer.load (legacy format) def test_load(): tokenizer_info = {"tokenizer_module": "llmeter"} tokenizer = Tokenizer.load(tokenizer_info) assert isinstance(tokenizer, DummyTokenizer) -# Test Tokenizer.to_dict -def test_to_dict(): - dummy_tokenizer = DummyTokenizer() - tokenizer_dict = Tokenizer.to_dict(dummy_tokenizer) - assert tokenizer_dict == {"tokenizer_module": "llmeter"} - - -# Test _to_dict function -def test_to_dict_transformers(monkeypatch): - # monkeypatch.setattr("llmeter.tokenizers.AutoTokenizer", MockTransformersTokenizer) - tokenizer = MockTransformersTokenizer() - tokenizer_dict = _to_dict(tokenizer) - assert tokenizer_dict == { - "tokenizer_module": "transformers", - "name": "mock-transformer", - } - - -def test_to_dict_tiktoken(monkeypatch): - tokenizer = MockTiktokenTokenizer() - tokenizer_dict = _to_dict(tokenizer) - assert tokenizer_dict == {"tokenizer_module": "tiktoken", "name": "mock-tiktoken"} - - -def test_to_dict_unknown(): - class UnknownTokenizer: - pass - - with pytest.raises(ValueError, match="Unknown tokenizer module"): - _to_dict(UnknownTokenizer()) - - -# Test save_tokenizer function -def test_save_tokenizer(tmp_path): - dummy_tokenizer = DummyTokenizer() - output_path = tmp_path / "tokenizer.json" - saved_path = save_tokenizer(dummy_tokenizer, output_path) - assert saved_path == output_path - assert output_path.exists() - with open(output_path, "r") as f: - saved_info = json.load(f) - assert saved_info == {"tokenizer_module": "llmeter"} - - # Test _load_tokenizer_from_info function @pytest.mark.skipif(not TRANSFORMERS_AVAILABLE, reason="transformers is not installed") def test_load_tokenizer_from_info_transformers(monkeypatch): - # Mock AutoTokenizer.from_pretrained to return our mock def mock_from_pretrained(name): return MockTransformersTokenizer() @@ -150,7 +85,6 @@ def mock_from_pretrained(name): @pytest.mark.skipif(not TIKTOKEN_AVAILABLE, reason="tiktoken is not installed") def test_load_tokenizer_from_info_tiktoken(monkeypatch): - # Mock get_encoding to return our mock def mock_get_encoding(name): return MockTiktokenTokenizer() @@ -173,3 +107,21 @@ def test_load_tokenizer_from_info_unknown(): tokenizer_info = {"tokenizer_module": "unknown"} with pytest.raises(ValueError, match="Unknown tokenizer module"): _load_tokenizer_from_info(tokenizer_info) + + +# Test dump_object/load_object round-trip (new serialization) +def test_serialization_roundtrip(): + from llmeter.serialization import dump_object, load_object + + original = DummyTokenizer() + data = dump_object(original) + + assert "__llmeter_class__" in data + assert "llmeter.tokenizers.DummyTokenizer" in data["__llmeter_class__"] + + restored = load_object(data) + assert isinstance(restored, DummyTokenizer) + + # Behavior preserved + text = "hello world test" + assert original.encode(text) == restored.encode(text) From ed862d23d94cb10a074de905b0f78ed951dd2af8 Mon Sep 17 00:00:00 2001 From: Alex Thewsey Date: Wed, 15 Jul 2026 17:27:42 +0800 Subject: [PATCH 07/20] test: Fix rebase conflicts Some unit tests were still referencing old json_utils module --- tests/unit/test_interrupted_run.py | 8 ++++---- tests/unit/test_runner_backlog_ui.py | 9 ++++----- 2 files changed, 8 insertions(+), 9 deletions(-) diff --git a/tests/unit/test_interrupted_run.py b/tests/unit/test_interrupted_run.py index a57fedf..438878e 100644 --- a/tests/unit/test_interrupted_run.py +++ b/tests/unit/test_interrupted_run.py @@ -44,7 +44,7 @@ def interrupted_run_dir(tmp_path, sample_responses): Contains only responses.jsonl and run_config.json, as would be left behind when a run is interrupted before Result.save() completes. """ - from llmeter.json_utils import llmeter_default_serializer + from llmeter.serialization import json_default run_dir = UPath(tmp_path / "interrupted_run") run_dir.mkdir(parents=True) @@ -78,7 +78,7 @@ def interrupted_run_dir(tmp_path, sample_responses): } config_path = run_dir / "run_config.json" with config_path.open("w") as f: - json.dump(config, f, default=llmeter_default_serializer) + json.dump(config, f, default=json_default) return run_dir @@ -214,7 +214,7 @@ def test_load_nonexistent_directory_raises(self, tmp_path): def test_load_recovers_interrupted_run_with_stats(self, tmp_path, sample_responses): """Simulates the case where both stats.json and responses.jsonl exist (e.g. the interrupt handler managed to write stats before exiting).""" - from llmeter.json_utils import llmeter_default_serializer + from llmeter.serialization import json_default run_dir = UPath(tmp_path / "partial_with_stats") run_dir.mkdir(parents=True) @@ -232,7 +232,7 @@ def test_load_recovers_interrupted_run_with_stats(self, tmp_path, sample_respons } stats_path = run_dir / "stats.json" with stats_path.open("w") as f: - json.dump(stats, f, default=llmeter_default_serializer) + json.dump(stats, f, default=json_default) result = Result.load(run_dir) diff --git a/tests/unit/test_runner_backlog_ui.py b/tests/unit/test_runner_backlog_ui.py index c684793..f47aac0 100644 --- a/tests/unit/test_runner_backlog_ui.py +++ b/tests/unit/test_runner_backlog_ui.py @@ -37,11 +37,10 @@ def mock_endpoint(): @pytest.fixture def mock_tokenizer(): - with patch( - "llmeter.tokenizers.Tokenizer.to_dict", - return_value={"tokenizer_module": "mock_tokenizer"}, - ): - yield MagicMock(spec=Tokenizer) + # Tokenizer now serializes via the Serializable mixin (__getstate__), not a + # to_dict() method. These tests run without an output_path, so the tokenizer + # is never actually serialized — a bare spec'd mock is all that's needed. + yield MagicMock(spec=Tokenizer) @pytest.fixture From d68f8edd7dbe6286518f700dda01068bed751b98 Mon Sep 17 00:00:00 2001 From: Alessandro Cere Date: Wed, 15 Jul 2026 20:43:36 +0800 Subject: [PATCH 08/20] fix: round-trip datetime and bytes in _serialize_value/_deserialize_value datetime was falling through to str(), losing type information on deserialization. bytes was producing repr(b'...') instead of the base64 marker. Now: - _serialize_value uses a dispatch table for bytes/datetime/PathLike and raises TypeError for unsupported types instead of silent str() - _deserialize_value uses match/case and recognizes the canonical ISO-8601 Z-suffix format to restore datetime objects automatically - bytes round-trips via __llmeter_bytes__ markers (same as json_default) Adds tests enforcing these behaviors. --- llmeter/serialization.py | 43 +++- tests/unit/test_serialize_value_roundtrip.py | 195 +++++++++++++++++++ 2 files changed, 228 insertions(+), 10 deletions(-) create mode 100644 tests/unit/test_serialize_value_roundtrip.py diff --git a/llmeter/serialization.py b/llmeter/serialization.py index ccc6d59..e8ce321 100644 --- a/llmeter/serialization.py +++ b/llmeter/serialization.py @@ -32,6 +32,7 @@ import inspect import logging import os +import re from dataclasses import asdict, is_dataclass from datetime import date, datetime, time, timezone from typing import Any @@ -172,28 +173,50 @@ def load_object(data: dict) -> Any: # Internal helpers for recursive serialization # --------------------------------------------------------------------------- +_SERIALIZERS: list[tuple[type | tuple[type, ...], Any]] = [ + (bytes, lambda v: {"__llmeter_bytes__": base64.b64encode(v).decode("utf-8")}), + (datetime, datetime_to_str), + (os.PathLike, lambda v: Path(v).as_posix()), +] + +_DATETIME_RE = re.compile(r"^\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}Z$") + def _serialize_value(val: Any) -> Any: """Recursively prepare a value for JSON persistence.""" if val is None or isinstance(val, (str, int, float, bool)): return val - if hasattr(val, "__getstate__") and type(val).__getstate__ is not object.__getstate__: + for types, fn in _SERIALIZERS: + if isinstance(val, types): + return fn(val) + if ( + hasattr(val, "__getstate__") + and type(val).__getstate__ is not object.__getstate__ + ): return dump_object(val) if isinstance(val, dict): return {k: _serialize_value(v) for k, v in val.items()} if isinstance(val, (list, tuple)): return [_serialize_value(item) for item in val] - return str(val) + raise TypeError(f"Cannot serialize {type(val).__name__!r} object: {val!r}") def _deserialize_value(val: Any) -> Any: """Recursively restore a value from JSON persistence.""" - if val is None or isinstance(val, (str, int, float, bool)): - return val - if isinstance(val, dict): - if "__llmeter_class__" in val and "__llmeter_state__" in val: + match val: + case None | bool() | int() | float(): + return val + case str() if _DATETIME_RE.fullmatch(val): + return str_to_datetime(val) + case str(): + return val + case {"__llmeter_class__": _, "__llmeter_state__": _}: return load_object(val) - return {k: _deserialize_value(v) for k, v in val.items()} - if isinstance(val, (list, tuple)): - return [_deserialize_value(item) for item in val] - return val + case {"__llmeter_bytes__": b64} if len(val) == 1: + return base64.b64decode(b64) + case dict(): + return {k: _deserialize_value(v) for k, v in val.items()} + case list() | tuple(): + return [_deserialize_value(item) for item in val] + case _: + return val diff --git a/tests/unit/test_serialize_value_roundtrip.py b/tests/unit/test_serialize_value_roundtrip.py new file mode 100644 index 0000000..f699d0a --- /dev/null +++ b/tests/unit/test_serialize_value_roundtrip.py @@ -0,0 +1,195 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Tests that _serialize_value/_deserialize_value correctly round-trip datetime and bytes.""" + +import base64 +from datetime import datetime, timezone + +import pytest + +from llmeter.serialization import ( + Serializable, + _deserialize_value, + _serialize_value, +) + + +class TestBytesRoundTrip: + """_serialize_value must produce the __llmeter_bytes__ marker for bytes, + and _deserialize_value must restore it.""" + + def test_bytes_serializes_to_marker(self): + val = b"\x00\x01\x02\xff" + result = _serialize_value(val) + assert isinstance(result, dict) + assert "__llmeter_bytes__" in result + assert len(result) == 1 + assert base64.b64decode(result["__llmeter_bytes__"]) == val + + def test_bytes_roundtrip(self): + val = b"hello world" + assert _deserialize_value(_serialize_value(val)) == val + + def test_empty_bytes_roundtrip(self): + val = b"" + assert _deserialize_value(_serialize_value(val)) == val + + def test_bytes_nested_in_dict(self): + val = {"key": b"\xde\xad\xbe\xef", "other": "text"} + serialized = _serialize_value(val) + assert serialized["key"] == { + "__llmeter_bytes__": base64.b64encode(b"\xde\xad\xbe\xef").decode() + } + assert serialized["other"] == "text" + assert _deserialize_value(serialized) == val + + def test_bytes_nested_in_list(self): + val = [b"first", "middle", b"last"] + serialized = _serialize_value(val) + restored = _deserialize_value(serialized) + assert restored == val + + def test_bytes_deeply_nested(self): + val = {"a": {"b": [{"c": b"\x01\x02\x03"}]}} + assert _deserialize_value(_serialize_value(val)) == val + + +class TestDatetimeRoundTrip: + """_serialize_value produces ISO-8601 Z-suffixed strings for datetime, + and _deserialize_value recognizes the format and restores datetime objects.""" + + def test_datetime_serializes_to_iso_string(self): + dt = datetime(2024, 6, 15, 10, 30, 0, tzinfo=timezone.utc) + result = _serialize_value(dt) + assert isinstance(result, str) + assert result == "2024-06-15T10:30:00Z" + + def test_datetime_utc_roundtrip(self): + dt = datetime(2024, 1, 1, 12, 0, 0, tzinfo=timezone.utc) + restored = _deserialize_value(_serialize_value(dt)) + assert restored == dt + assert isinstance(restored, datetime) + + def test_datetime_naive_roundtrip(self): + dt = datetime(2024, 3, 15, 8, 45, 30) + serialized = _serialize_value(dt) + # Naive datetimes produce no Z suffix, so they stay as plain strings + assert serialized == "2024-03-15T08:45:30" + + def test_datetime_is_not_silently_lost(self): + dt = datetime(2024, 6, 15, 10, 30, 0, tzinfo=timezone.utc) + result = _serialize_value(dt) + # Must be the canonical format, not repr/str() of the datetime object + assert "datetime" not in result + assert result.endswith("Z") + + def test_datetime_nested_in_dict(self): + dt = datetime(2024, 6, 15, 10, 30, 0, tzinfo=timezone.utc) + val = {"timestamp": dt, "label": "test"} + serialized = _serialize_value(val) + assert serialized["timestamp"] == "2024-06-15T10:30:00Z" + restored = _deserialize_value(serialized) + assert restored["timestamp"] == dt + assert restored["label"] == "test" + + def test_datetime_nested_in_list(self): + dt1 = datetime(2024, 1, 1, tzinfo=timezone.utc) + dt2 = datetime(2024, 12, 31, tzinfo=timezone.utc) + val = [dt1, "gap", dt2] + restored = _deserialize_value(_serialize_value(val)) + assert restored[0] == dt1 + assert restored[1] == "gap" + assert restored[2] == dt2 + + def test_plain_string_not_confused_for_datetime(self): + val = "hello world" + assert _deserialize_value(val) == "hello world" + + def test_partial_iso_string_stays_as_string(self): + val = "2024-06-15" + assert _deserialize_value(val) == "2024-06-15" + + +class TestPathSerialization: + """PathLike objects serialize to POSIX path strings.""" + + def test_path_serializes_to_string(self): + from pathlib import PurePosixPath + + p = PurePosixPath("/tmp/foo/bar") + result = _serialize_value(p) + assert isinstance(result, str) + assert result == "/tmp/foo/bar" + + +class TestSerializableWithDatetimeAndBytes: + """End-to-end: a Serializable class with datetime/bytes fields round-trips + via __getstate__/__setstate__.""" + + def test_serializable_with_datetime_field(self): + class MyObj(Serializable): + def __init__(self, created_at: datetime, name: str = "test"): + self.created_at = created_at + self.name = name + + dt = datetime(2024, 6, 15, 10, 30, 0, tzinfo=timezone.utc) + obj = MyObj(created_at=dt, name="hello") + state = obj.__getstate__() + + restored = MyObj.__new__(MyObj) + restored.__setstate__(state) + + assert restored.created_at == dt + assert isinstance(restored.created_at, datetime) + assert restored.name == "hello" + + def test_serializable_with_bytes_field(self): + class MyObj(Serializable): + def __init__(self, payload: bytes, label: str = "x"): + self.payload = payload + self.label = label + + obj = MyObj(payload=b"\xde\xad\xbe\xef", label="binary") + state = obj.__getstate__() + + restored = MyObj.__new__(MyObj) + restored.__setstate__(state) + + assert restored.payload == b"\xde\xad\xbe\xef" + assert isinstance(restored.payload, bytes) + assert restored.label == "binary" + + def test_serializable_with_mixed_fields(self): + class MyObj(Serializable): + def __init__(self, ts: datetime, data: bytes, info: dict): + self.ts = ts + self.data = data + self.info = info + + dt = datetime(2024, 1, 1, tzinfo=timezone.utc) + obj = MyObj(ts=dt, data=b"raw", info={"nested_bytes": b"\x01", "nested_ts": dt}) + state = obj.__getstate__() + + restored = MyObj.__new__(MyObj) + restored.__setstate__(state) + + assert restored.ts == dt + assert restored.data == b"raw" + assert restored.info["nested_bytes"] == b"\x01" + assert restored.info["nested_ts"] == dt + + +class TestUnserializableTypeRaises: + """_serialize_value must raise TypeError for unsupported types, + never silently convert to str.""" + + def test_set_raises(self): + with pytest.raises(TypeError, match="Cannot serialize"): + _serialize_value({1, 2, 3}) + + def test_custom_object_without_getstate_raises(self): + class Opaque: + pass + + with pytest.raises(TypeError, match="Cannot serialize"): + _serialize_value(Opaque()) From 21d3157244bdbdcb2d5e19940a5b48ace79a01d2 Mon Sep 17 00:00:00 2001 From: Alessandro Cere Date: Wed, 15 Jul 2026 21:11:16 +0800 Subject: [PATCH 09/20] refactor: use modern typing (Python 3.12+) in files on this branch MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Optional[X] → X | None - Sequence/Callable/Iterator from collections.abc instead of typing --- llmeter/callbacks/cost/dimensions.py | 19 +++++++++---------- llmeter/prompt_utils.py | 7 +++---- llmeter/results.py | 14 ++++++++------ 3 files changed, 20 insertions(+), 20 deletions(-) diff --git a/llmeter/callbacks/cost/dimensions.py b/llmeter/callbacks/cost/dimensions.py index ae93e29..f006703 100644 --- a/llmeter/callbacks/cost/dimensions.py +++ b/llmeter/callbacks/cost/dimensions.py @@ -11,7 +11,7 @@ from abc import ABC, abstractmethod from dataclasses import dataclass from math import ceil -from typing import Optional, Protocol +from typing import Protocol # Local Dependencies: from ...endpoints.base import InvocationResponse @@ -23,14 +23,14 @@ class IRequestCostDimension(Protocol): """Interface for one dimension of a per-request cost model.""" - async def calculate(self, response: InvocationResponse) -> Optional[float]: ... + async def calculate(self, response: InvocationResponse) -> float | None: ... class IRunCostDimension(Protocol): """Interface for one dimension of a per-Run cost model.""" async def before_run_start(self, run_config: _RunConfig) -> None: ... - async def calculate(self, result: Result) -> Optional[float]: ... + async def calculate(self, result: Result) -> float | None: ... class RequestCostDimensionBase(Serializable, ABC): @@ -41,7 +41,7 @@ class RequestCostDimensionBase(Serializable, ABC): """ @abstractmethod - async def calculate(self, response: InvocationResponse) -> Optional[float]: + async def calculate(self, response: InvocationResponse) -> float | None: raise NotImplementedError @@ -56,7 +56,7 @@ async def before_run_start(self, run_config: _RunConfig) -> None: pass @abstractmethod - async def calculate(self, result: Result) -> Optional[float]: + async def calculate(self, result: Result) -> float | None: raise NotImplementedError @@ -77,7 +77,7 @@ class InputTokens(RequestCostDimensionBase): price_per_million: float granularity: int = 1 - async def calculate(self, req: InvocationResponse) -> Optional[float]: + async def calculate(self, req: InvocationResponse) -> float | None: if req.num_tokens_input is None: return None billable = ceil(req.num_tokens_input / self.granularity) * self.granularity @@ -96,7 +96,7 @@ class OutputTokens(RequestCostDimensionBase): price_per_million: float granularity: int = 1 - async def calculate(self, req: InvocationResponse) -> Optional[float]: + async def calculate(self, req: InvocationResponse) -> float | None: if req.num_tokens_output is None: return None billable = ceil(req.num_tokens_output / self.granularity) * self.granularity @@ -115,11 +115,10 @@ class EndpointTime(RunCostDimensionBase): price_per_hour: float granularity_secs: float = 1 - async def calculate(self, result: Result) -> Optional[float]: + async def calculate(self, result: Result) -> float | None: if result.total_test_time is None: return None billable = ( - ceil(result.total_test_time / self.granularity_secs) - * self.granularity_secs + ceil(result.total_test_time / self.granularity_secs) * self.granularity_secs ) return billable * self.price_per_hour / 3600 diff --git a/llmeter/prompt_utils.py b/llmeter/prompt_utils.py index fd927e5..49e7583 100644 --- a/llmeter/prompt_utils.py +++ b/llmeter/prompt_utils.py @@ -4,9 +4,10 @@ import json import logging import random +from collections.abc import Callable, Iterator from dataclasses import dataclass from itertools import product -from typing import Any, Callable, Iterator +from typing import Any from upath import UPath as Path from upath.types import ReadablePathLike, WritablePathLike @@ -513,9 +514,7 @@ def _load_data_file(file: Path) -> Iterator[dict]: try: if not line.strip(): continue - yield json.loads( - line.strip(), object_hook=bytes_decoder - ) + yield json.loads(line.strip(), object_hook=bytes_decoder) except json.JSONDecodeError as e: print(f"Error decoding JSON in {file}: {e}") else: # Assume it's a regular JSON file diff --git a/llmeter/results.py b/llmeter/results.py index a9a5479..98f959d 100644 --- a/llmeter/results.py +++ b/llmeter/results.py @@ -4,10 +4,11 @@ import json import logging import types as _types +from collections.abc import Sequence from dataclasses import asdict, dataclass, fields from datetime import datetime from numbers import Number -from typing import Any, Sequence +from typing import Any import jmespath from upath.types import ReadablePathLike, WritablePathLike @@ -129,7 +130,11 @@ def save(self, output_path: WritablePathLike | None = None) -> None: stats_path = output_path / "stats.json" with summary_path.open("w") as f: json.dump( - {k: o for k, o in asdict(self).items() if k not in ["responses", "stats"]}, + { + k: o + for k, o in asdict(self).items() + if k not in ["responses", "stats"] + }, f, default=json_default, indent=4, @@ -141,10 +146,7 @@ def save(self, output_path: WritablePathLike | None = None) -> None: if not responses_path.exists(): with responses_path.open("w") as f: for response in self.responses: - f.write( - json.dumps(asdict(response), default=json_default) - + "\n" - ) + f.write(json.dumps(asdict(response), default=json_default) + "\n") def to_json(self, default=json_default, **kwargs) -> str: """Return the results as a JSON string. From 7abbc8149686d0a49b109a62cdfec57b61c4d855 Mon Sep 17 00:00:00 2001 From: Alessandro Cere Date: Wed, 15 Jul 2026 21:35:13 +0800 Subject: [PATCH 10/20] refactor: unify type restoration via restore_dataclass_types MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Move datetime/bytes field restoration from ad-hoc per-class methods into a single restore_dataclass_types() in serialization.py. It introspects dataclass annotations and uses match/case to restore fields — no hardcoded field name lists needed. InvocationResponse.from_json and Result.load now both delegate to it, removing str_to_datetime from their imports and eliminating the _parse_datetime_fields / _get_type_args helpers from results.py. --- llmeter/endpoints/base.py | 11 ++++------- llmeter/results.py | 37 ++++--------------------------------- llmeter/serialization.py | 35 ++++++++++++++++++++++++++++++++++- tests/unit/test_results.py | 22 ++++++++++++++-------- 4 files changed, 56 insertions(+), 49 deletions(-) diff --git a/llmeter/endpoints/base.py b/llmeter/endpoints/base.py index 8cdb362..2aa4edf 100644 --- a/llmeter/endpoints/base.py +++ b/llmeter/endpoints/base.py @@ -27,7 +27,7 @@ dump_object, json_default, load_object, - str_to_datetime, + restore_dataclass_types, ) from ..utils import ensure_path @@ -82,9 +82,8 @@ def from_json(cls, json_str: str) -> "InvocationResponse": correctly restores types that the default JSON round-trip would leave as strings or marker objects: - * `request_time` is parsed from an ISO-8601 string back to a Python `datetime` - * `payload`s containing `bytes` (as `__llmeter_bytes__` markers) are correctly loaded back - as bytes. + * `datetime`-annotated fields are parsed from ISO-8601 strings back to Python `datetime` + * `bytes`-typed fields and `__llmeter_bytes__` markers in nested payloads are restored Args: json_str: A JSON string representation of an InvocationResponse (produced by `to_json` @@ -101,9 +100,7 @@ def from_json(cls, json_str: str) -> "InvocationResponse": ``` """ data = json.loads(json_str, object_hook=bytes_decoder) - rt = data.get("request_time") - if rt is not None and isinstance(rt, str): - data["request_time"] = str_to_datetime(rt) + restore_dataclass_types(cls, data) return cls(**data) def to_json(self, default=json_default, **kwargs) -> str: diff --git a/llmeter/results.py b/llmeter/results.py index 98f959d..af49e52 100644 --- a/llmeter/results.py +++ b/llmeter/results.py @@ -3,9 +3,8 @@ import json import logging -import types as _types from collections.abc import Sequence -from dataclasses import asdict, dataclass, fields +from dataclasses import asdict, dataclass from datetime import datetime from numbers import Number from typing import Any @@ -14,22 +13,12 @@ from upath.types import ReadablePathLike, WritablePathLike from .endpoints import InvocationResponse -from .serialization import json_default, str_to_datetime +from .serialization import json_default, restore_dataclass_types from .utils import ensure_path, summary_stats_from_list logger = logging.getLogger(__name__) -def _get_type_args(tp) -> tuple: - """Return the members of a union type (e.g. ``datetime | None`` -> (datetime, NoneType)).""" - if isinstance(tp, _types.UnionType): - return tp.__args__ - origin = getattr(tp, "__origin__", None) - if origin is _types.UnionType: - return tp.__args__ - return (tp,) if isinstance(tp, type) else () - - @dataclass class Result: """Results of a test run.""" @@ -59,24 +48,6 @@ def __post_init__(self): if not hasattr(self, "_preloaded_stats"): self._preloaded_stats = None - @classmethod - def _parse_datetime_fields(cls, d: dict) -> None: - """Convert any datetime fields on cls present in d from ISO-8601 strings to datetimes - - Introspects this (data)class to find all `datetime`-typed fields, and converts any matching - entries in `d` with ISO-8601 string values to datetimes instead. This is used for loading - JSON data (in which dates are stringified) into the Result class or stats. - """ - for f in fields(cls): - if datetime not in _get_type_args(f.type): - continue - val = d.get(f.name) - if val and isinstance(val, str): - try: - d[f.name] = str_to_datetime(val) - except ValueError: - pass - def _update_contributed_stats(self, stats: dict[str, Number]): """ Upsert externally-provided statistics for the `stats` property @@ -275,7 +246,7 @@ def _load_summary(cls, result_path) -> dict[str, Any]: with summary_path.open("r") as f: summary = json.load(f) - cls._parse_datetime_fields(summary) + restore_dataclass_types(cls, summary) if "output_path" not in summary or summary["output_path"] is None: summary["output_path"] = str(result_path) @@ -401,7 +372,7 @@ def _resolve_stats(cls, result: "Result", stats_path) -> None: try: with stats_path.open("r") as s: saved_stats = json.loads(s.read()) - cls._parse_datetime_fields(saved_stats) + restore_dataclass_types(cls, saved_stats) except (json.JSONDecodeError, OSError) as e: logger.warning(f"Could not load stats.json: {e}") saved_stats = None diff --git a/llmeter/serialization.py b/llmeter/serialization.py index e8ce321..a3c4893 100644 --- a/llmeter/serialization.py +++ b/llmeter/serialization.py @@ -33,7 +33,8 @@ import logging import os import re -from dataclasses import asdict, is_dataclass +import types as _types +from dataclasses import asdict, fields, is_dataclass from datetime import date, datetime, time, timezone from typing import Any @@ -95,6 +96,38 @@ def bytes_decoder(dct: dict) -> dict | bytes: return dct +def _get_type_args(tp) -> tuple: + """Return the members of a union type (e.g. ``datetime | None`` -> (datetime, NoneType)).""" + if isinstance(tp, _types.UnionType): + return tp.__args__ + origin = getattr(tp, "__origin__", None) + if origin is _types.UnionType: + return tp.__args__ + return (tp,) if isinstance(tp, type) else () + + +def restore_dataclass_types(cls, data: dict) -> None: + """Restore typed fields in a dict destined for a dataclass constructor. + + Introspects ``cls`` (a dataclass) and converts JSON-native values back to + their annotated Python types. Only fields declared on ``cls`` are touched — + nested user payloads (e.g. ``input_payload``) are left unchanged. + """ + for f in fields(cls): + val = data.get(f.name) + if val is None: + continue + type_args = _get_type_args(f.type) + match val: + case str() if datetime in type_args: + try: + data[f.name] = str_to_datetime(val) + except ValueError: + pass + case {"__llmeter_bytes__": b64} if bytes in type_args and len(val) == 1: + data[f.name] = base64.b64decode(b64) + + # --------------------------------------------------------------------------- # Serializable mixin # --------------------------------------------------------------------------- diff --git a/tests/unit/test_results.py b/tests/unit/test_results.py index 6af5552..5c39c56 100644 --- a/tests/unit/test_results.py +++ b/tests/unit/test_results.py @@ -407,11 +407,13 @@ def test_save_method_existing_responses(sample_result: Result, temp_dir: UPath): assert responses[-1]["id"] == "extra_response" -# ── _parse_datetime_fields introspection ─────────────────────────────────────── +# ── restore_dataclass_types introspection ────────────────────────────────────── -def test_parse_datetime_fields_converts_iso_strings(): - """_parse_datetime_fields should convert all datetime-typed field keys.""" +def test_restore_dataclass_types_converts_iso_strings(): + """restore_dataclass_types should convert all datetime-typed field keys.""" + from llmeter.serialization import restore_dataclass_types + d = { "start_time": "2025-06-01T10:00:00Z", "end_time": "2025-06-01T10:05:00+00:00", @@ -419,7 +421,7 @@ def test_parse_datetime_fields_converts_iso_strings(): "last_request_time": None, "total_requests": 5, # non-datetime field, should be left alone } - Result._parse_datetime_fields(d) + restore_dataclass_types(Result, d) assert isinstance(d["start_time"], datetime) assert isinstance(d["end_time"], datetime) @@ -428,20 +430,24 @@ def test_parse_datetime_fields_converts_iso_strings(): assert d["total_requests"] == 5 -def test_parse_datetime_fields_skips_already_parsed(): +def test_restore_dataclass_types_skips_already_parsed(): """Already-datetime values should pass through unchanged.""" from datetime import timezone + from llmeter.serialization import restore_dataclass_types + dt = datetime(2025, 1, 1, 12, 0, 0, tzinfo=timezone.utc) d = {"start_time": dt} - Result._parse_datetime_fields(d) + restore_dataclass_types(Result, d) assert d["start_time"] is dt -def test_parse_datetime_fields_handles_invalid_string(): +def test_restore_dataclass_types_handles_invalid_string(): """Invalid date strings should be left as-is (not raise).""" + from llmeter.serialization import restore_dataclass_types + d = {"start_time": "not-a-date"} - Result._parse_datetime_fields(d) + restore_dataclass_types(Result, d) assert d["start_time"] == "not-a-date" From 4dceb48c2311254bf991ea95641aed8f87868923 Mon Sep 17 00:00:00 2001 From: Alessandro Cere Date: Thu, 16 Jul 2026 20:58:51 +0800 Subject: [PATCH 11/20] refactor: move save_to_file/load_from_file to Serializable, restore docstrings Lift file I/O methods from Callback to Serializable so all subclasses (Endpoint, Tokenizer, cost dimensions) get them automatically. Restore detailed docstrings that were inadvertently reduced during the serialization consolidation, and expand docstrings for new functions in serialization.py to match the project's documentation standard. --- llmeter/callbacks/base.py | 42 +-------- llmeter/callbacks/cost/dimensions.py | 103 ++++++++++++++++----- llmeter/callbacks/cost/model.py | 56 +++++++++++- llmeter/endpoints/base.py | 27 ++++-- llmeter/serialization.py | 129 ++++++++++++++++++++++++--- llmeter/tokenizers.py | 16 ++-- 6 files changed, 285 insertions(+), 88 deletions(-) diff --git a/llmeter/callbacks/base.py b/llmeter/callbacks/base.py index 479f196..94d6aee 100644 --- a/llmeter/callbacks/base.py +++ b/llmeter/callbacks/base.py @@ -4,17 +4,12 @@ from __future__ import annotations -import json from abc import ABC -from typing import final - -from upath.types import ReadablePathLike, WritablePathLike from ..endpoints.base import InvocationResponse from ..results import Result from ..runner import _RunConfig -from ..serialization import Serializable, dump_object, json_default, load_object -from ..utils import ensure_path +from ..serialization import Serializable class Callback(Serializable, ABC): @@ -25,8 +20,8 @@ class Callback(Serializable, ABC): associated with test runs or individual model invocations. A Callback object may implement multiple of the defined lifecycle hooks (such as - `before_invoke`, `after_run`, etc). Callbacks must support serializing their configuration to - a file (by implementing `save_to_file`), and loading back (via `load_from_file`). + `before_invoke`, `after_run`, etc). Serialization to/from file is inherited from + :class:`~llmeter.serialization.Serializable`. """ async def before_invoke(self, payload: dict) -> None: @@ -74,34 +69,3 @@ async def after_run(self, result: Result) -> None: """ pass - def save_to_file(self, path: WritablePathLike) -> None: - """Save this Callback to a JSON file. - - Uses the ``__getstate__`` protocol. Override ``__getstate__`` (not this method) - if custom serialization is needed. - - Args: - path: (Local or Cloud) path where the callback will be saved. - """ - path = ensure_path(path) - path.parent.mkdir(parents=True, exist_ok=True) - data = dump_object(self) - with path.open("w") as f: - json.dump(data, f, indent=4, default=json_default) - - @staticmethod - @final - def load_from_file(path: ReadablePathLike) -> Callback: - """Load (any type of) Callback from a JSON file. - - Detects the callback type from the ``__llmeter_class__`` field and reconstructs it. - - Args: - path: (Local or Cloud) path where the callback was saved. - Returns: - callback: A loaded Callback instance. - """ - path = ensure_path(path) - with path.open("r") as f: - data = json.load(f) - return load_object(data) diff --git a/llmeter/callbacks/cost/dimensions.py b/llmeter/callbacks/cost/dimensions.py index f006703..64da593 100644 --- a/llmeter/callbacks/cost/dimensions.py +++ b/llmeter/callbacks/cost/dimensions.py @@ -2,9 +2,13 @@ # SPDX-License-Identifier: Apache-2.0 """Classes defining different components of cost -A "dimension" is one aspect of the pricing for a deployed Foundation Model or application. -Here we provide implementations for some common cost dimensions, and define base classes you -can use to bring customized cost dimensions for your own cost models. +A "dimension" is one aspect of the pricing for a deployed Foundation Model or application. In +general, multiple factors are likely to contribute to the total cost of FMs under test: For +example, an API may charge separate rates for input vs output token counts; or a self-managed +cloud deployment may carry per-hour charges for compute, as well as network bandwidth charges. + +Here we provide implementations for some common cost dimensions, and define base classes you can +use to bring customized cost dimensions for your own cost models. """ # Python Built-Ins: @@ -21,42 +25,95 @@ class IRequestCostDimension(Protocol): - """Interface for one dimension of a per-request cost model.""" + """Interface for one dimension of a per-request cost model + + Per-request cost components are calculated independently for each invocation in a test run, and + can be used to model factors like cost-per-request, cost-per-input-tokens, + cost-per-request-duration, etc. They're typically most relevant for serverless deployments like + Amazon Bedrock, or estimating duration-based execution costs for AWS Lambda functions. + """ - async def calculate(self, response: InvocationResponse) -> float | None: ... + async def calculate(self, response: InvocationResponse) -> float | None: + """Calculate (this component of) the cost for an individual request/response""" + ... class IRunCostDimension(Protocol): - """Interface for one dimension of a per-Run cost model.""" + """Interface for one dimension of a per-Run cost model + + Per-run cost components are notified before the start of a test run via `before_run_start()`, + and then requested to `calculate()` at the end of the run. They're most relevant for + provisioned-infrastructure based deployments like Amazon SageMaker, where factors like a + (request-independent) cost-per-hour are important. + """ + + async def before_run_start(self, run_config: _RunConfig) -> None: + """Notify the cost component that a test run is about to start + + This method is called before the test run starts, and can be used to perform any + initialization or setup required for the cost component. In general, we assume a dimension + instance may be re-used for multiple test runs, but only one run at a time: Meaning + `before_run_start()` should not be called again before `calculate()` is called for the + previous run. + """ + ... + + async def calculate(self, result: Result) -> float | None: + """Calculate (this component of) the cost for a completed test run - async def before_run_start(self, run_config: _RunConfig) -> None: ... - async def calculate(self, result: Result) -> float | None: ... + Dimensions that depend on `before_run_start` being called to return an accurate result + should throw an error if this was not done. Dimensions that only need `calculate()` should + silently ignore if `before_run_start` was not called. + """ + ... class RequestCostDimensionBase(Serializable, ABC): - """Base class for per-request cost dimensions. + """Base class for implementing per-request cost model dimensions - Inherits ``__getstate__``/``__setstate__`` from :class:`~llmeter.serialization.Serializable`. - Subclasses just declare fields and implement ``calculate()``. + This class provides a default implementation of serialization (via + :class:`~llmeter.serialization.Serializable`) and sets up an abstract method for + `calculate()`. It's fine if you don't want to derive from it directly - just be sure to + fully implement `IRequestCostDimension`! """ @abstractmethod async def calculate(self, response: InvocationResponse) -> float | None: + """Calculate (this component of) the cost for an individual request/response""" raise NotImplementedError class RunCostDimensionBase(Serializable, ABC): - """Base class for per-run cost dimensions. + """Base class for implementing per-run cost model dimensions - Inherits ``__getstate__``/``__setstate__`` from :class:`~llmeter.serialization.Serializable`. + This class provides a default implementation of serialization (via + :class:`~llmeter.serialization.Serializable`), a default empty `before_run_start` + implementation, and abstract methods for the other requirements of the `IRunCostDimension` + protocol. It's fine if you don't want to derive from it directly - just make sure you fully + implement `IRunCostDimension`! """ async def before_run_start(self, run_config: _RunConfig) -> None: - """Called before a test run starts. Default is a no-op.""" + """Notify the cost component that a test run is about to start + + This method is called before the test run starts, and can be used to perform any + initialization or setup required for the cost component. In general, we assume a dimension + instance may be re-used for multiple test runs, but only one run at a time: Meaning + `before_run_start()` should not be called again before `calculate()` is called for the + previous run. + + The default implementation is a pass. + """ pass @abstractmethod async def calculate(self, result: Result) -> float | None: + """Calculate (this component of) the cost for a completed test run + + Dimensions that depend on `before_run_start` being called to return an accurate result + should throw an error if this was not done. Dimensions that only need `calculate()` should + silently ignore if `before_run_start` was not called. + """ raise NotImplementedError @@ -67,11 +124,11 @@ async def calculate(self, result: Result) -> float | None: @dataclass class InputTokens(RequestCostDimensionBase): - """Request cost dimension: per-input-token costs with a flat charge rate. + """Request cost dimension to model per-input-token costs with a flat charge rate Args: - price_per_million: Charge per million input (prompt) tokens. - granularity: Minimum tokens billed per increment (Default 1). + price_per_million: Charge applied per million input (prompt) token to the Foundation Model + granularity: Minimum number of tokens billed per increment (Default 1) """ price_per_million: float @@ -86,11 +143,11 @@ async def calculate(self, req: InvocationResponse) -> float | None: @dataclass class OutputTokens(RequestCostDimensionBase): - """Request cost dimension: per-output-token costs with a flat charge rate. + """Request cost dimension to model per-output-token costs with a flat charge rate Args: - price_per_million: Charge per million output (completion) tokens. - granularity: Minimum tokens billed per increment (Default 1). + price_per_million: Charge per million output (completion) token from the Foundation Model + granularity: Minimum number of tokens billed per increment (Default 1) """ price_per_million: float @@ -105,11 +162,11 @@ async def calculate(self, req: InvocationResponse) -> float | None: @dataclass class EndpointTime(RunCostDimensionBase): - """Run cost dimension: per-deployment-hour costs with a flat charge rate. + """Run cost dimension to model per-deployment-hour costs with a flat charge rate Args: - price_per_hour: Charge per hour a test run takes. - granularity_secs: Minimum seconds billed per increment (Default 1). + price_per_hour: Charge applied per hour a test run takes + granularity_secs: Minimum number of seconds billed per increment (Default 1) """ price_per_hour: float diff --git a/llmeter/callbacks/cost/model.py b/llmeter/callbacks/cost/model.py index b409672..2e23439 100644 --- a/llmeter/callbacks/cost/model.py +++ b/llmeter/callbacks/cost/model.py @@ -11,10 +11,17 @@ class CostModel(Callback): - """Model costs of (test runs of) Foundation Models. + """Model costs of (test runs of) Foundation Models - A Cost Model is composed of pricing *dimensions* applied at the individual request - level (like ``InputTokens``) or at the overall run level (like ``EndpointTime``). + A Cost Model is composed of whatever pricing *"dimensions"* are relevant for your FM deployment + and use-case, which may be applied at the individual request level (like + `llmeter.callbacks.cost.dimensions.InputTokens`), or at the overall deployment / test run level + (like `llmeter.callbacks.cost.dimensions.EndpointTime`). + + With a Cost Model defined, you can explicitly `calculate_request_cost(...)` on an + `InvocationResponse` to estimate costs for a specific request/response; + `calculate_run_cost(...)` on a `Result` to estimate costs for a test run; or pass the model as + a Callback when running an LLMeter Run or Experiment, to annotate the results automatically. """ def __init__( @@ -30,6 +37,20 @@ def __init__( | None ) = None, ): + """Create a CostModel + + Args: + request_dims: Dimensions of request-level cost (for example, charges by number of input + or output tokens). If provided as a dict, the keys will be used as the name of each + dimension. If provided as a list, each dimension's class name will be used as its + default name. An error will be thrown if any two dimensions (including run_dims) + have the same name. + run_dims: Dimensions of run-level cost (for example, per-hour charges for an FM endpoint + being available and used in a run). If provided as a dict, the keys will be used as + the name of each dimension. If provided as a list, each dimension's class name will + be used as its default name. An error will be thrown if any two dimensions + (including request_dims) have the same name. + """ all_dims: dict[str, IRequestCostDimension | IRunCostDimension] = {} self.request_dims: dict[str, IRequestCostDimension] = {} self.run_dims: dict[str, IRunCostDimension] = {} @@ -69,6 +90,13 @@ def __init__( async def calculate_request_cost( self, response: InvocationResponse, save: bool = False ) -> CalculatedCostWithDimensions: + """Calculate the costs of a single FM invocation (excluding any session-level costs) + + Args: + response: The InvocationResponse to estimate costs for + save: Set `True` to also store the result in `response.cost`, in addition to returning + it. Defaults to `False` + """ dim_costs = CalculatedCostWithDimensions( **{ name: await dim.calculate(response) @@ -85,6 +113,19 @@ async def calculate_run_cost( recalculate_request_costs: bool = True, save: bool = False, ) -> CalculatedCostWithDimensions: + """Calculate the run-level costs of a test (including any request-level costs) + + NOTE: Depending on the types of `run_dims` in your model, this may throw an error if + `before_run` was not called first to initialise the state. + + Args: + result: An LLMeter Run result + recalculate_request_costs: Set `False` if your `result.responses` have already been + annotated with costs in line with the current cost model. By default (`True`), + the costs for each response will be re-calculated. + save: Set `True` to also store the result in `result.cost`, in addition to returning + it. Defaults to `False`. + """ run_cost = CalculatedCostWithDimensions( **{ name: await dim.calculate(result) @@ -121,13 +162,22 @@ async def calculate_run_cost( return run_cost async def after_invoke(self, response: InvocationResponse) -> None: + """LLMeter Callback.after_invoke hook + + Calls calculate_request_cost() with `save=True` to save the cost on the InvocationResponse. + """ await self.calculate_request_cost(response, save=True) async def before_run(self, run_config: _RunConfig) -> None: + """Initialize state for all run-level cost dimensions in the model""" for dim in self.run_dims.values(): await dim.before_run_start(run_config) async def after_run(self, result: Result) -> None: + """LLMeter Callback.after_run hook + + Calls calculate_run_cost() with `save=True` to save the cost on the Result. + """ await self.calculate_run_cost( result, recalculate_request_costs=False, save=True ) diff --git a/llmeter/endpoints/base.py b/llmeter/endpoints/base.py index 2aa4edf..0f9206f 100644 --- a/llmeter/endpoints/base.py +++ b/llmeter/endpoints/base.py @@ -480,6 +480,9 @@ def __subclasshook__(cls, C: type) -> bool: def save(self, output_path: WritablePathLike) -> Path: """Save the endpoint configuration to a JSON file. + This method serializes the endpoint's configuration via the ``__getstate__`` protocol + to a JSON file at the specified path. + Args: output_path (str | UPath): The path where the configuration file will be saved. @@ -508,11 +511,16 @@ def to_dict(self) -> dict: def load_from_file(cls, input_path: ReadablePathLike) -> "Endpoint": """Load an endpoint configuration from a JSON file. + This class method reads a JSON file containing an endpoint configuration, + determines the appropriate endpoint class, and instantiates it with the + loaded configuration. + Args: input_path (str | UPath): The path to the JSON configuration file. Returns: - Endpoint: An instance of the appropriate endpoint class. + Endpoint: An instance of the appropriate endpoint class, initialized + with the configuration from the file. """ input_path = ensure_path(input_path) with input_path.open("r") as f: @@ -526,17 +534,24 @@ def load_from_file(cls, input_path: ReadablePathLike) -> "Endpoint": @classmethod def load(cls, endpoint_config: dict) -> "Endpoint": # type: ignore - """Load an endpoint from a legacy ``{"endpoint_type": ...}`` dictionary. + """Load an endpoint configuration from a dictionary. + + This class method reads a dictionary containing an endpoint configuration, + determines the appropriate endpoint class, and instantiates it with the + loaded configuration. .. deprecated:: - New code should use :func:`~llmeter.serialization.load_object` with - dicts produced by :func:`~llmeter.serialization.dump_object`. + This supports the legacy ``{"endpoint_type": ...}`` format. New code should + use :func:`~llmeter.serialization.load_object` with dicts produced by + :func:`~llmeter.serialization.dump_object`. Args: - endpoint_config: Dictionary with at minimum an ``endpoint_type`` key. + endpoint_config (dict): A dictionary containing the endpoint configuration. + Must include at minimum an ``endpoint_type`` key. Returns: - Endpoint: An instance of the appropriate endpoint class. + Endpoint: An instance of the appropriate endpoint class, initialized + with the configuration from the dictionary. """ endpoint_type = endpoint_config.pop("endpoint_type") endpoint_module = importlib.import_module("llmeter.endpoints") diff --git a/llmeter/serialization.py b/llmeter/serialization.py index a3c4893..0e9f540 100644 --- a/llmeter/serialization.py +++ b/llmeter/serialization.py @@ -30,6 +30,7 @@ import base64 import importlib import inspect +import json import logging import os import re @@ -39,6 +40,9 @@ from typing import Any from upath import UPath as Path +from upath.types import ReadablePathLike, WritablePathLike + +from .utils import ensure_path logger = logging.getLogger(__name__) @@ -70,10 +74,25 @@ def str_to_datetime(s: str) -> datetime: def json_default(obj: Any) -> Any: - """Fallback serializer for :func:`json.dumps`. + """Serialize a single non-JSON-serializable object. + + Intended for use as the ``default`` argument to :func:`json.dumps` or + :func:`json.dump`. + + Type handling (checked in order): + + * ``bytes`` — wrapped in a ``{"__llmeter_bytes__": ""}`` marker so + that :func:`bytes_decoder` can restore them on the way back. + * ``datetime`` — converted to a UTC ISO-8601 string with a ``Z`` suffix. + * ``date`` / ``time`` — converted via ``.isoformat()``. + * ``os.PathLike`` — converted to a POSIX path string. + * Anything else — ``str()`` fallback (returns ``None`` if that also fails). - Handles ``bytes`` (→ base64 marker), ``datetime``, ``date``/``time``, - ``os.PathLike``, and falls back to ``str()``. + Args: + obj: The object that the default encoder could not handle. + + Returns: + A JSON-serializable representation of *obj*. """ if isinstance(obj, bytes): return {"__llmeter_bytes__": base64.b64encode(obj).decode("utf-8")} @@ -90,7 +109,18 @@ def json_default(obj: Any) -> Any: def bytes_decoder(dct: dict) -> dict | bytes: - """Object hook for :func:`json.loads` that restores ``__llmeter_bytes__`` markers.""" + """Decode ``__llmeter_bytes__`` marker objects back to ``bytes``. + + Intended for use as the ``object_hook`` argument to :func:`json.loads` or + :func:`json.load`. Marker objects produced by :func:`json_default` are detected + and converted back to ``bytes``; all other dicts pass through unchanged. + + Args: + dct: A dictionary produced by the JSON parser. + + Returns: + The original ``bytes`` if *dct* is a marker object, otherwise *dct* unchanged. + """ if "__llmeter_bytes__" in dct and len(dct) == 1: return base64.b64decode(dct["__llmeter_bytes__"]) return dct @@ -109,9 +139,18 @@ def _get_type_args(tp) -> tuple: def restore_dataclass_types(cls, data: dict) -> None: """Restore typed fields in a dict destined for a dataclass constructor. - Introspects ``cls`` (a dataclass) and converts JSON-native values back to - their annotated Python types. Only fields declared on ``cls`` are touched — - nested user payloads (e.g. ``input_payload``) are left unchanged. + Introspects ``cls`` (a dataclass) and converts JSON-native values back to their + annotated Python types. Currently handles: + + * ``datetime`` fields — parses ISO-8601 strings via :func:`str_to_datetime`. + * ``bytes`` fields — decodes ``__llmeter_bytes__`` markers via base64. + + Only fields declared on ``cls`` are touched — nested user payloads (e.g. + ``input_payload``) are left unchanged. Mutates *data* in place. + + Args: + cls: A dataclass type to introspect for field type annotations. + data: A dictionary of field values (e.g. from :func:`json.load`) to coerce. """ for f in fields(cls): val = data.get(f.name) @@ -161,6 +200,37 @@ def __setstate__(self, state: dict) -> None: deserialized = {k: _deserialize_value(v) for k, v in state.items()} self.__init__(**deserialized) + def save_to_file(self, path: WritablePathLike) -> None: + """Save this object to a JSON file. + + Uses the ``__getstate__`` protocol. Override ``__getstate__`` (not this method) + if custom serialization is needed. + + Args: + path: (Local or Cloud) path where the object will be saved. + """ + path = ensure_path(path) + path.parent.mkdir(parents=True, exist_ok=True) + data = dump_object(self) + with path.open("w") as f: + json.dump(data, f, indent=4, default=json_default) + + @classmethod + def load_from_file(cls, path: ReadablePathLike) -> "Serializable": + """Load an object from a JSON file. + + Detects the type from the ``__llmeter_class__`` field and reconstructs it. + + Args: + path: (Local or Cloud) path where the object was saved. + Returns: + The loaded instance. + """ + path = ensure_path(path) + with path.open("r") as f: + data = json.load(f) + return load_object(data) + # --------------------------------------------------------------------------- # Object serialization API @@ -170,7 +240,21 @@ def __setstate__(self, state: dict) -> None: def dump_object(obj: Any) -> dict: """Serialize an object to a type-tagged dict for round-trip persistence. - Returns ``{"__llmeter_class__": "module.Class", "__llmeter_state__": {...}}``. + The returned envelope has the form + ``{"__llmeter_class__": "module.Class", "__llmeter_state__": {...}}``. + + Serialization strategy (checked in order): + + 1. If the object has a custom ``__getstate__`` (not :func:`object.__getstate__`), + calls it to obtain the state dict. + 2. If the object is a dataclass, uses :func:`dataclasses.asdict`. + 3. Otherwise, takes all public (non-underscore-prefixed) entries from ``__dict__``. + + Args: + obj: The object to serialize. + + Returns: + A JSON-serializable dict that :func:`load_object` can reconstruct. """ class_path = f"{obj.__class__.__module__}.{obj.__class__.__qualname__}" if ( @@ -190,7 +274,20 @@ def dump_object(obj: Any) -> dict: def load_object(data: dict) -> Any: """Restore an object from a type-tagged dict produced by :func:`dump_object`. - .. warning:: Do not call on data from untrusted sources. + Imports the module identified by ``__llmeter_class__``, instantiates the class + (bypassing ``__init__`` via ``__new__``), and calls ``__setstate__`` with the + persisted state dict. + + Args: + data: A dict with ``__llmeter_class__`` and ``__llmeter_state__`` keys, as + produced by :func:`dump_object`. + + Returns: + The reconstructed object instance. + + .. warning:: + Do not call on data from untrusted sources — it imports and instantiates + arbitrary classes. """ class_path = data["__llmeter_class__"] module_path, class_name = class_path.rsplit(".", 1) @@ -216,7 +313,12 @@ def load_object(data: dict) -> Any: def _serialize_value(val: Any) -> Any: - """Recursively prepare a value for JSON persistence.""" + """Recursively prepare a value for JSON persistence. + + Handles primitives, known types (bytes, datetime, PathLike), nested + :class:`Serializable` objects (via :func:`dump_object`), dicts, and lists/tuples. + Raises :exc:`TypeError` for objects it cannot serialize. + """ if val is None or isinstance(val, (str, int, float, bool)): return val for types, fn in _SERIALIZERS: @@ -235,7 +337,12 @@ def _serialize_value(val: Any) -> Any: def _deserialize_value(val: Any) -> Any: - """Recursively restore a value from JSON persistence.""" + """Recursively restore a value from JSON persistence. + + Recognizes type-tagged dicts (``__llmeter_class__``), bytes markers + (``__llmeter_bytes__``), ISO-8601 datetime strings, and recursively processes + nested dicts and lists. + """ match val: case None | bool() | int() | float(): return val diff --git a/llmeter/tokenizers.py b/llmeter/tokenizers.py index 74ebf18..89e259d 100644 --- a/llmeter/tokenizers.py +++ b/llmeter/tokenizers.py @@ -4,7 +4,6 @@ from __future__ import annotations from abc import ABC, abstractmethod -from typing import Any from .serialization import Serializable @@ -43,25 +42,30 @@ def __subclasscheck__(cls, subclass): @staticmethod def load(tokenizer_info: dict) -> "Tokenizer": - """Load a tokenizer from a legacy ``{"tokenizer_module": ...}`` dictionary. + """Load a tokenizer from a dictionary. This supports configs saved before the unified ``dump_object``/``load_object`` serialization was introduced. New code should use :func:`~llmeter.serialization.load_object` instead. Args: - tokenizer_info: Dictionary with at minimum a ``tokenizer_module`` key. + tokenizer_info (dict): The tokenizer information to load. Must include at minimum + a ``tokenizer_module`` key. Returns: - A restored Tokenizer instance. + Tokenizer: The loaded tokenizer. """ return _load_tokenizer_from_info(tokenizer_info) def _load_tokenizer_from_info(tokenizer_info: dict) -> Tokenizer: - """Instantiate a tokenizer from a legacy info dict. + """Load a tokenizer from a legacy info dictionary. - Supports ``"transformers"``, ``"tiktoken"``, and ``"llmeter"`` modules. + Args: + tokenizer_info (dict): The tokenizer information to load. + + Returns: + Tokenizer: The loaded tokenizer. """ module = tokenizer_info["tokenizer_module"] From d395b936ac52d4162c427580816fe6e94c41d07b Mon Sep 17 00:00:00 2001 From: Alex Thewsey Date: Wed, 22 Jul 2026 17:19:30 +0800 Subject: [PATCH 12/20] test: Add test for CostModel save file round-trip Add tests to explicitly validate CostModel's save_to_file round trip (previously only the dump & load APIs were checked). --- tests/unit/callbacks/cost/test_model.py | 75 +++++++++++++++++++++++++ 1 file changed, 75 insertions(+) diff --git a/tests/unit/callbacks/cost/test_model.py b/tests/unit/callbacks/cost/test_model.py index 1a32f08..e05263b 100644 --- a/tests/unit/callbacks/cost/test_model.py +++ b/tests/unit/callbacks/cost/test_model.py @@ -1,5 +1,6 @@ # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # SPDX-License-Identifier: Apache-2.0 +import json from dataclasses import dataclass from unittest.mock import AsyncMock, Mock, NonCallableMock @@ -32,6 +33,80 @@ def test_cost_model_serialization(): assert "run_dims" in d +def test_cost_model_save_load_file_roundtrip(tmp_path): + """A CostModel can be saved to and loaded from a file, re-initialising its dimensions + + Exercises the inherited Serializable.save_to_file / load_from_file API (file I/O plus + json_default), which the in-memory dump_object/load_object test does not cover. + """ + from llmeter.callbacks.cost.dimensions import ( + EndpointTime, + InputTokens, + OutputTokens, + ) + + model = CostModel( + request_dims={ + "TokensIn": InputTokens(price_per_million=30, granularity=10), + "TokensOut": OutputTokens(price_per_million=60), + }, + run_dims={"ComputeSeconds": EndpointTime(price_per_hour=50)}, + ) + + path = tmp_path / "cost_model.json" + saved_path = model.save_to_file(path) + # save_to_file returns the (validated/normalized) path it wrote to + assert str(saved_path) == str(path) + assert path.is_file() + + # The file is valid JSON tagged with the CostModel class path + with open(path) as f: + raw = json.load(f) + assert raw["__llmeter_class__"] == "llmeter.callbacks.cost.model.CostModel" + + restored = CostModel.load_from_file(path) + + # The correct concrete type is reconstructed from the file + assert isinstance(restored, CostModel) + + # Each dimension is re-initialised as the correct type, under the same name, with values intact + assert set(restored.request_dims) == {"TokensIn", "TokensOut"} + assert set(restored.run_dims) == {"ComputeSeconds"} + + tokens_in = restored.request_dims["TokensIn"] + assert isinstance(tokens_in, InputTokens) + assert tokens_in.price_per_million == 30 + assert tokens_in.granularity == 10 + + tokens_out = restored.request_dims["TokensOut"] + assert isinstance(tokens_out, OutputTokens) + assert tokens_out.price_per_million == 60 + + compute = restored.run_dims["ComputeSeconds"] + assert isinstance(compute, EndpointTime) + assert compute.price_per_hour == 50 + + +def test_cost_model_load_from_file_dispatches_via_base_class(tmp_path): + """load_from_file resolves the concrete type from the file, even when called on Serializable + + The class is detected from the ``__llmeter_class__`` marker rather than the class the + classmethod is invoked on, so loading via the base mixin still yields a CostModel. + """ + from llmeter.callbacks.cost.dimensions import InputTokens + from llmeter.serialization import Serializable + + model = CostModel(request_dims={"TokensIn": InputTokens(price_per_million=15)}) + + path = tmp_path / "cost_model.json" + model.save_to_file(path) + + restored = Serializable.load_from_file(path) + assert isinstance(restored, CostModel) + assert restored.request_dims["TokensIn"].price_per_million == 15 + assert restored.run_dims == {} + + def test_cost_model_detects_duplicate_cost_dim_names(): """CostModel detects and raises errors when created with duplicate dimension names""" From 341e4567112488c9cdff3f3190316b12b48a9fb2 Mon Sep 17 00:00:00 2001 From: Alex Thewsey Date: Wed, 22 Jul 2026 17:22:50 +0800 Subject: [PATCH 13/20] refactor: save_to_file returns path, not None Incremental changes to resolve serialization API conflicts/confusion: - Endpoint.save is now a thin wrapper over save_to_file, and the latter is changed to return the saved path instead of None - Endpoint.load_from_file previously renamed the parent `path` param to `input_path`: Revert this. Both changes pose potential backward compatibility issues but assessed to be very minor in return for more standard API. --- llmeter/endpoints/base.py | 29 ++++++++++++++++------------- llmeter/serialization.py | 6 +++++- 2 files changed, 21 insertions(+), 14 deletions(-) diff --git a/llmeter/endpoints/base.py b/llmeter/endpoints/base.py index 0f9206f..47d7c78 100644 --- a/llmeter/endpoints/base.py +++ b/llmeter/endpoints/base.py @@ -11,6 +11,7 @@ import json import logging import time +import warnings from abc import ABC, abstractmethod from collections.abc import Callable from dataclasses import asdict, dataclass @@ -24,7 +25,6 @@ from ..serialization import ( Serializable, bytes_decoder, - dump_object, json_default, load_object, restore_dataclass_types, @@ -480,8 +480,10 @@ def __subclasshook__(cls, C: type) -> bool: def save(self, output_path: WritablePathLike) -> Path: """Save the endpoint configuration to a JSON file. - This method serializes the endpoint's configuration via the ``__getstate__`` protocol - to a JSON file at the specified path. + .. deprecated:: + Use :meth:`~llmeter.serialization.Serializable.save_to_file` instead, which + provides the same behavior with a consistent name across all serializable + LLMeter objects. This alias will be removed in a future major version. Args: output_path (str | UPath): The path where the configuration file will be saved. @@ -489,12 +491,13 @@ def save(self, output_path: WritablePathLike) -> Path: Returns: Path: The path the file was written to. """ - output_path = ensure_path(output_path) - output_path.parent.mkdir(parents=True, exist_ok=True) - data = dump_object(self) - with output_path.open("w") as f: - json.dump(data, f, indent=4, default=json_default) - return output_path + warnings.warn( + "Endpoint.save() is deprecated and will be removed in a future version; " + "use Endpoint.save_to_file() instead.", + DeprecationWarning, + stacklevel=2, + ) + return self.save_to_file(output_path) def to_dict(self) -> dict: """ @@ -508,7 +511,7 @@ def to_dict(self) -> dict: return endpoint_conf @classmethod - def load_from_file(cls, input_path: ReadablePathLike) -> "Endpoint": + def load_from_file(cls, path: ReadablePathLike) -> "Endpoint": """Load an endpoint configuration from a JSON file. This class method reads a JSON file containing an endpoint configuration, @@ -516,14 +519,14 @@ def load_from_file(cls, input_path: ReadablePathLike) -> "Endpoint": loaded configuration. Args: - input_path (str | UPath): The path to the JSON configuration file. + path (str | UPath): The path to the JSON configuration file. Returns: Endpoint: An instance of the appropriate endpoint class, initialized with the configuration from the file. """ - input_path = ensure_path(input_path) - with input_path.open("r") as f: + path = ensure_path(path) + with path.open("r") as f: data = json.load(f) if "__llmeter_class__" in data: return load_object(data) diff --git a/llmeter/serialization.py b/llmeter/serialization.py index 0e9f540..46f55e6 100644 --- a/llmeter/serialization.py +++ b/llmeter/serialization.py @@ -200,7 +200,7 @@ def __setstate__(self, state: dict) -> None: deserialized = {k: _deserialize_value(v) for k, v in state.items()} self.__init__(**deserialized) - def save_to_file(self, path: WritablePathLike) -> None: + def save_to_file(self, path: WritablePathLike) -> Path: """Save this object to a JSON file. Uses the ``__getstate__`` protocol. Override ``__getstate__`` (not this method) @@ -208,12 +208,16 @@ def save_to_file(self, path: WritablePathLike) -> None: Args: path: (Local or Cloud) path where the object will be saved. + + Returns: + The (validated/normalized) path the object was written to. """ path = ensure_path(path) path.parent.mkdir(parents=True, exist_ok=True) data = dump_object(self) with path.open("w") as f: json.dump(data, f, indent=4, default=json_default) + return path @classmethod def load_from_file(cls, path: ReadablePathLike) -> "Serializable": From 5a32dd44d9b58f065a011b72ea44a6535612a5af Mon Sep 17 00:00:00 2001 From: Alex Thewsey Date: Thu, 23 Jul 2026 00:22:22 +0800 Subject: [PATCH 14/20] fix: pickle conflict and legacy data loading Rename __{get|set}state__ serialization API to avoid conflicts with pickle & deepcopy. Fix parsing of older-format run config JSONs so users can still load Results generated by older versions that don't use __llmeter_class__ markers. Fix issue where CostModel used to push arbitrary extra fields directly onto InvocationResponse, but these no longer survive the serialization round trip - by moving them under an explicit 'annotations' map and providing a legacy loading method specific to InvocationResponse to pull unrecognised fields into `annotations`. Update user guide on serialization to be more user focussed and docstrings to include some extra design discussion. --- docs/serialization.md | 189 -------------- docs/user_guide/callbacks/index.md | 13 + docs/user_guide/serialization.md | 199 ++++++++++++++ llmeter/callbacks/base.py | 13 +- llmeter/callbacks/cost/dimensions.py | 22 +- llmeter/callbacks/cost/model.py | 46 ++-- llmeter/endpoints/base.py | 65 ++++- llmeter/runner.py | 23 +- llmeter/serialization.py | 101 ++++++-- mkdocs.yml | 2 +- .../cost/providers/test_sagemaker.py | 4 +- tests/unit/callbacks/cost/test_dimensions.py | 18 +- tests/unit/callbacks/cost/test_model.py | 66 ++++- tests/unit/callbacks/test_base.py | 6 +- tests/unit/endpoints/test_endpoints.py | 53 ++++ tests/unit/test_legacy_data_load.py | 245 ++++++++++++++++++ tests/unit/test_property_save_load.py | 34 +++ tests/unit/test_runner_backlog_ui.py | 5 +- tests/unit/test_serialize_value_roundtrip.py | 14 +- 19 files changed, 827 insertions(+), 291 deletions(-) delete mode 100644 docs/serialization.md create mode 100644 docs/user_guide/serialization.md create mode 100644 tests/unit/test_legacy_data_load.py diff --git a/docs/serialization.md b/docs/serialization.md deleted file mode 100644 index 7977915..0000000 --- a/docs/serialization.md +++ /dev/null @@ -1,189 +0,0 @@ -# Serialization - -LLMeter provides a unified serialization layer in a single module — -`llmeter.serialization` — that handles JSON encoding, datetime conversion, -binary content round-tripping, and full object persistence. - -## Module overview - -| Symbol | Purpose | -|--------|---------| -| `Serializable` | Mixin giving any class automatic `__getstate__`/`__setstate__` | -| `dump_object` / `load_object` | Full round-trip persistence via a type-tagged envelope | -| `json_default` | `json.dumps` fallback for bytes, datetime, PathLike | -| `bytes_decoder` | `json.loads` object hook to restore `__llmeter_bytes__` markers | -| `datetime_to_str` / `str_to_datetime` | UTC ISO-8601 with `Z` suffix, both directions | - -```python -from llmeter.serialization import ( - dump_object, load_object, json_default, bytes_decoder, - datetime_to_str, str_to_datetime, -) -``` - -## JSON encoding - -Use `json_default` as the `default` argument whenever you call `json.dumps` or -`json.dump` with LLMeter objects: - -```python -import json -from llmeter.serialization import json_default - -json.dump(my_data, f, default=json_default, indent=4) -``` - -It handles: - -- **bytes** — wrapped in `{"__llmeter_bytes__": ""}` markers -- **datetime** — UTC ISO-8601 string via `datetime_to_str` -- **date / time** — `.isoformat()` -- **os.PathLike** — POSIX path string -- **anything else** — `str()` fallback - -To decode bytes markers back on load: - -```python -data = json.load(f, object_hook=bytes_decoder) -``` - -## Datetime handling - -All datetime serialization goes through two functions: - -```python -from llmeter.serialization import datetime_to_str, str_to_datetime - -datetime_to_str(datetime(2024, 1, 1, tzinfo=timezone.utc)) -# → "2024-01-01T00:00:00Z" - -str_to_datetime("2024-01-01T00:00:00Z") -# → datetime(2024, 1, 1, 0, 0, tzinfo=timezone.utc) -``` - -Timezone-aware datetimes are normalized to UTC. Naive datetimes are serialized -as-is. The `Z` suffix is canonical. - -## Persisting runtime objects - -Objects with runtime state (boto3 clients, SDK connections, compiled -tokenizers) can't be naively serialized. LLMeter handles this via the -`__getstate__`/`__setstate__` protocol: - -```python -from llmeter.serialization import dump_object, load_object - -# Save an endpoint (config only — not the boto3 client) -data = dump_object(endpoint) -# → {"__llmeter_class__": "llmeter.endpoints.bedrock.BedrockConverse", -# "__llmeter_state__": {"model_id": "claude-3", "region": "us-west-2"}} - -# Restore it (boto3 client recreated automatically via __init__) -restored = load_object(data) -``` - -### How it works - -1. `dump_object(obj)` calls `obj.__getstate__()` to get a JSON-safe state - dict, then wraps it with the fully-qualified class path. -2. `load_object(data)` imports the class, creates an empty instance via - `cls.__new__`, and calls `obj.__setstate__(state)` to reconstruct it. - -### The `Serializable` mixin (zero boilerplate) - -Inherit from `Serializable` to get automatic `__getstate__`/`__setstate__`: - -```python -from llmeter.serialization import Serializable - -class MyEndpoint(Serializable, Endpoint): - def __init__(self, model_id: str, region: str = "us-east-1"): - super().__init__(endpoint_name=model_id, model_id=model_id, provider="custom") - self.region = region - self._client = create_client(region) # runtime — not serialized -``` - -The mixin introspects `__init__` parameters, matches them to instance -attributes (`self.name` or `self._name`), and serializes only what's needed to -recreate the object. Private attributes (prefixed with `_`) that don't -correspond to `__init__` params are skipped — they're assumed to be derived -runtime state. - -Nested `Serializable` objects are recursively handled: `__getstate__` wraps -them via `dump_object`, and `__setstate__` restores them via `load_object`. - -### When to override - -Override `__getstate__`/`__setstate__` only when: - -- An `__init__` parameter is consumed without being stored -- Reconstruction needs special logic beyond `__init__(**state)` -- You want to exclude large transient data from persistence - -```python -class SpecialEndpoint(Serializable, Endpoint): - def __getstate__(self) -> dict: - return {"model_id": self.model_id, "region": self.region} - - def __setstate__(self, state: dict): - self.__init__(**state) -``` - -## Callback persistence - -All callbacks support `save_to_file()` / `load_from_file()`: - -```python -from llmeter.callbacks.base import Callback -from llmeter.callbacks.mlflow import MlflowCallback - -cb = MlflowCallback(step=5, nested=True) -cb.save_to_file("/tmp/callback.json") - -# Polymorphic load — detects the type from __llmeter_class__ -restored = Callback.load_from_file("/tmp/callback.json") -# → MlflowCallback(step=5, nested=True) -``` - -## Runner config persistence - -`_RunConfig.save()` and `_RunConfig.load()` use `dump_object`/`load_object` -for all callable fields (endpoint, tokenizer, callbacks): - -```python -runner = Runner(endpoint=BedrockConverse(...), callbacks=[MlflowCallback(step=1)]) -runner.save(output_path="/tmp/run") - -# Full reconstruction: -restored = _RunConfig.load("/tmp/run") -``` - -When loading, the runner detects both formats: - -- **New format** (`__llmeter_class__` key) — uses `load_object` -- **Legacy format** (`endpoint_type` / `tokenizer_module` keys) — uses - `Endpoint.load()` / `Tokenizer.load()` for backward compatibility - -## Dataclass compatibility - -`@dataclass` classes work seamlessly with `Serializable`. The mixin -introspects the `__init__` that `@dataclass` generates: - -```python -@dataclass -class InputTokens(Serializable): - price_per_million: float - granularity: int = 1 - -data = dump_object(InputTokens(price_per_million=3.0)) -restored = load_object(data) -``` - -## Security - -`load_object` imports and instantiates whatever class path is found in the -`__llmeter_class__` field. This is the same trust model as Python's `pickle` — -**never load configs from untrusted sources**. - -This is appropriate for LLMeter's use case: developer-generated configs stored -on local disk or controlled cloud storage (S3 with IAM). diff --git a/docs/user_guide/callbacks/index.md b/docs/user_guide/callbacks/index.md index c1af74c..43fd3de 100644 --- a/docs/user_guide/callbacks/index.md +++ b/docs/user_guide/callbacks/index.md @@ -36,3 +36,16 @@ results = await runner.run( ``` Each callback will be processed in the **same order** as you provide them to the runner. This is important to remember and configure properly, if you're stacking multiple callbacks that access the same data (for example - transforming an invocation response *then* logging/exporting it somewhere, both using `after_invoke`). + +### Storing extra data on responses and results + +If your callback needs to attach **extra fields** on your `InvocationResponse`s (for example, a computed cost or a custom label), store them in the response's `annotations` dictionary: + +```python +async def after_invoke(self, response): + response.annotations["my_metric"] = compute_something(response) +``` + +Any extra fields added directly to the `InvocationResponse` object itself will *not* be preserved when the responses are saved to file. For legacy compatability, any unrecognized fields found in older responses.jsonl files are *currently* callected back to `annotations`, but this behaviour may be dropped in future. + +Likewise for extra **run-level** data, `after_run` can add numeric statistics to the `Result` via [`_update_contributed_stats`](../../reference/results/#llmeter.results.Result), which are persisted alongside the run's other stats. (This is how the built-in [cost callback](cost.md) records run-level cost metrics.) diff --git a/docs/user_guide/serialization.md b/docs/user_guide/serialization.md new file mode 100644 index 0000000..12c0cad --- /dev/null +++ b/docs/user_guide/serialization.md @@ -0,0 +1,199 @@ +# Save (and Reload) Results + +!!! warning "Important security warning" + Do not use LLMeter file load functions on data from untrusted sources! For more details on why, see the following section. + +With LLMeter, you can save comprehensive test configuration and result information to files - whether locally or on the Cloud - and load past runs back into Python for analysis later. For example: + +```python +from llmeter.results import Result +from llmeter.runner import Runner + +# Provide s3:// URI or local path: +base_output_path = f"s3://doc-example-bucket/llmeter/outputs/{endpoint.model_id}" + +runner = Runner( + endpoint, + # Configure the output path when creating a Runner... + output_path=base_output_path, +) +results = await runner.run( + payload=sample_payload, + n_requests=3, + clients=5, + # ...*Optionally* specify a specific path for an individual run: + # (Otherwise a subfolder will be created by run name, automatically) + output_path=f"{base_output_path}/my-cool-run" +) + +# At some later date, load your result back: +results = Result.load(f"{base_output_path}/my-cool-run") +``` + +!!! note "A note on performance" + While it's possible for LLMeter to write run outputs directly to Cloud object stores like + [Amazon S3](https://aws.amazon.com/s3/), remember it might reduce the maximum throughput you + can drive in *high-volume* tests, since it will consume network bandwidth. + +For the most part, generated files are JSON-based (or [JSON Lines](https://jsonlines.org/), for the +individual responses). However, LLMeter also handles some more complex data types including: + +- Binary data in request/response payloads (such as images) +- [Callbacks](./callbacks) configured on the test Run, including LLMeter built-ins as well as *your custom* callback classes + +Implementation of de/serialization functionality is centralized in the [`llmeter.serialization`](../reference/serialization) module. + + +## How complex data types are represented and loaded + +Objects that are not natively JSON serializable (but support LLMeter's serialization protocol) are saved to JSON-based formats something like the below: + +```json +{ + "__llmeter_class__": "llmeter.endpoints.bedrock.BedrockConverse", + "__llmeter_state__": {"model_id": "claude-3", "region": "us-west-2"} +} +``` + +- `datetime` objects are serialized to ISO8601 format strings converted to UTC timezone, like `2024-01-01T00:00:00Z`. +- `bytes` objects are serialized to base64 strings in a special `{"__llmeter_bytes__": ""}` wrapper. +- Objects implementing LLMeter's [Serializable](../reference/serialization.md#llmeter.serialization.Serializable) interface (including for example Endpoints and Callbacks) are represented as dicts with the `__llmeter_{class/state}__` properties as shown above + +When **loading** these objects back from file, LLMeter will import and instantiate whatever class path is found in the `__llmeter_class__` field. + +This is the same trust model as Python's native [pickle](https://docs.python.org/3/library/pickle.html) library: It is possible to construct malicious data which will run arbitrary code during loading, so **only load data that you trust**. + +## Under the hood: Serialization API components + +In many cases you'll be working with high-level classes like `Runner` and `Result` that already provide convenience methods to save to and load from file. However, building custom LLMeter extensions may require understanding how our [`llmeter.serialization`](../reference/serialization) components fit together: + +| Symbol | Purpose | +|--------|---------| +| [`Serializable`](../reference/serialization.md#llmeter.serialization.Serializable) | Mixin enabling your class to be serialized and loaded by LLMeter | +| [`dump_object`](../reference/serialization.md#llmeter.serialization.dump_object) / [`load_object`](../reference/serialization.md#llmeter.serialization.load_object) | Full round-trip persistence via a type-tagged envelope | +| [`json_default`](../reference/serialization.md#llmeter.serialization.json_default) | `json.dumps` fallback for bytes, datetime, PathLike | +| [`bytes_decoder`](../reference/serialization.md#llmeter.serialization.bytes_decoder) | `json.loads` object hook to restore `__llmeter_bytes__` markers | +| [`datetime_to_str`](../reference/serialization.md#llmeter.serialization.datetime_to_str) / [`str_to_datetime`](../reference/serialization.md#llmeter.serialization.str_to_datetime) | UTC ISO-8601 with `Z` suffix, both directions | + +```python +from llmeter.serialization import ( + dump_object, load_object, json_default, bytes_decoder, + datetime_to_str, str_to_datetime, +) +``` + +### Make custom classes serializable by LLMeter with the `Serializable` mixin + +The `Serializable` mixin provides two default methods: + +- `_get_llmeter_state`: Construct a JSON-ready dictionary of the state your class needs to be re-initialized + - The default implementation inspects the arguments of your `__init__` constructor and attempts to fetch those from the current object's fields, or with an `_` underscore prefix if the raw parameter name isn't present. +- `_set_llmeter_state`: Initialise an instance of your class, based on a state dictionary + - The default implementation calls your `__init__` with the arguments stored in the dictionary. + +Nested `Serializable` objects are recursively handled by default: `_get_llmeter_state` wraps them via `dump_object`, and `_set_llmeter_state` restores them via `load_object`. + +You'd only need to **override** these implementations if, for example: + +- An `__init__` parameter is consumed without being stored, or +- Reconstruction needs special logic beyond `__init__(**state)`, or +- You want to exclude large transient data from persistence + + +### Save and load your own objects + +Once your class inherits `Serializable`, it gets `save_to_file()` and `load_from_file()` for free - no extra code required: + +```python +from llmeter.callbacks.base import Callback +from llmeter.callbacks.mlflow import MlflowCallback + +cb = MlflowCallback(step=5, nested=True) +cb.save_to_file("/tmp/callback.json") + +# load_from_file is *polymorphic*: it reads the __llmeter_class__ recorded in the +# file and rebuilds the correct subclass, so you can call it on the base class. +restored = Callback.load_from_file("/tmp/callback.json") # -> MlflowCallback(step=5, nested=True) +``` + +Under the hood, those methods use two functions you can also call directly if you're managing the JSON yourself: + +- `dump_object(obj)` builds the type-tagged envelope (`{"__llmeter_class__": ..., "__llmeter_state__": ...}`). It reads the state from `obj._get_llmeter_state()`, or - for a plain `@dataclass` that *doesn't* inherit `Serializable` - from its fields. +- `load_object(data)` imports the class named in `__llmeter_class__`, creates a bare instance (via `__new__`, bypassing `__init__`), then repopulates it through `_set_llmeter_state()` - which by default re-runs `__init__` with the saved state. + +```python +from llmeter.serialization import dump_object, load_object + +data = dump_object(my_object) # -> plain dict +my_object = load_object(data) # -> reconstructed instance +``` + +Note the envelope produced by `dump_object` may still contain non-JSON values (like `bytes` or `datetime`) nested inside its state, so pair it with `json_default` when you actually write it out - see below. + +### Reading and writing the JSON yourself + +When you call `json.dump`/`json.dumps` directly on LLMeter data, pass `json_default` so the extra types are handled: + +```python +import json +from llmeter.serialization import json_default, bytes_decoder + +with open("my-file.json", "w") as f: + json.dump(my_data, f, default=json_default, indent=4) +``` + +`json_default` converts, in order: + +- **bytes** — wrapped in a `{"__llmeter_bytes__": ""}` marker +- **datetime** — UTC ISO-8601 string with a `Z` suffix +- **date / time** — `.isoformat()` +- **os.PathLike** — POSIX path string +- **anything else** — `str()` fallback + +To turn the `bytes` markers back into `bytes` on the way in, pass `bytes_decoder` as the `object_hook`: + +```python +with open("my-file.json") as f: + data = json.load(f, object_hook=bytes_decoder) +``` + +## How LLMeter's built-ins use this + +For everyday use you rarely touch the functions above directly, because the high-level classes wrap them for you. + +### Endpoints: config, not connections + +Endpoints are saved as *configuration only* - runtime state like a boto3 client is never written to file. When you load an endpoint back, `_set_llmeter_state` re-runs the constructor, which recreates that client for you: + +```python +from llmeter.serialization import dump_object, load_object + +data = dump_object(endpoint) +# → {"__llmeter_class__": "llmeter.endpoints.bedrock.BedrockConverse", +# "__llmeter_state__": {"model_id": "claude-3", "region": "us-west-2"}} + +restored = load_object(data) # boto3 client rebuilt via __init__ +``` + +### Runners: the whole run configuration + +When you give a `Runner` an `output_path`, it saves its full configuration - endpoint, tokenizer and callbacks included - as `run_config.json` at the start of each run. You can also trigger this yourself: + +```python +runner = Runner(endpoint=BedrockConverse(...), callbacks=[MlflowCallback(step=1)]) +runner.save(output_path="/tmp/run") # writes /tmp/run/run_config.json +``` + +If a callback (or any other configured object) can't be serialized because its class doesn't inherit `Serializable`, saving raises a `TypeError`. This is why custom callbacks and cost dimensions should subclass the relevant LLMeter base class - see [`Callback`](../reference/callbacks/base.md) and the [cost dimension base classes](../reference/callbacks/cost/dimensions.md). + +### Dataclasses work automatically + +Because the `Serializable` mixin introspects the constructor, any `@dataclass` that inherits it is serializable with no extra code - the generated `__init__` supplies the parameter names the mixin looks for. LLMeter's built-in cost dimensions are a good example: + +```python +from llmeter.callbacks.cost.dimensions import InputTokens +from llmeter.serialization import dump_object, load_object + +data = dump_object(InputTokens(price_per_million=3.0)) +restored = load_object(data) # -> InputTokens(price_per_million=3.0, granularity=1) +``` diff --git a/llmeter/callbacks/base.py b/llmeter/callbacks/base.py index 94d6aee..fc36e83 100644 --- a/llmeter/callbacks/base.py +++ b/llmeter/callbacks/base.py @@ -20,8 +20,11 @@ class Callback(Serializable, ABC): associated with test runs or individual model invocations. A Callback object may implement multiple of the defined lifecycle hooks (such as - `before_invoke`, `after_run`, etc). Serialization to/from file is inherited from - :class:`~llmeter.serialization.Serializable`. + `before_invoke`, `after_run`, etc) - which have no-op implementations by default. Serialization + to/from file is inherited from `llmeter.serialization.Serializable`, and is necessary so your + callback(s) can be saved to file (and restored) as part of a Run configuration. Any custom + callback class that is *not* LLMeter-serializable will raise a `TypeError` when the run config + it belongs to is saved. """ async def before_invoke(self, payload: dict) -> None: @@ -42,7 +45,10 @@ async def after_invoke(self, response: InvocationResponse) -> None: timing and token counts) Returns: None: If you'd like to add information to the `response` logged in the Run, modify it - in-place. + in-place. To attach **extra custom fields**, store them in `response.annotations` + (a dict) rather than setting arbitrary attributes on the response: `annotations` + is a declared field, so it is preserved through `Result` save/load, whereas loose + attributes are dropped on serialization. """ pass @@ -68,4 +74,3 @@ async def after_run(self, result: Result) -> None: None: If you'd like to modify the run `result`, edit the argument in-place. """ pass - diff --git a/llmeter/callbacks/cost/dimensions.py b/llmeter/callbacks/cost/dimensions.py index 64da593..7a12a26 100644 --- a/llmeter/callbacks/cost/dimensions.py +++ b/llmeter/callbacks/cost/dimensions.py @@ -31,6 +31,16 @@ class IRequestCostDimension(Protocol): can be used to model factors like cost-per-request, cost-per-input-tokens, cost-per-request-duration, etc. They're typically most relevant for serverless deployments like Amazon Bedrock, or estimating duration-based execution costs for AWS Lambda functions. + + !!! warning + This interface describes only the *calculation* behavior. It intentionally says nothing + about serialization: a dimension used purely in-memory need not be serializable. However, + to be usable inside a :class:`~llmeter.callbacks.cost.model.CostModel` that gets **saved** + (e.g. as a run callback, or via ``save_to_file``), a dimension must **additionally** be + LLMeter-serializable. The supported way to get that is to subclass + :class:`RequestCostDimensionBase` (which mixes in + :class:`~llmeter.serialization.Serializable`); implementing this Protocol alone is not + sufficient for persistence. """ async def calculate(self, response: InvocationResponse) -> float | None: @@ -45,6 +55,12 @@ class IRunCostDimension(Protocol): and then requested to `calculate()` at the end of the run. They're most relevant for provisioned-infrastructure based deployments like Amazon SageMaker, where factors like a (request-independent) cost-per-hour are important. + + !!! warning + As with :class:`IRequestCostDimension`, this interface covers only calculation behavior and + not serialization. To be usable inside a :class:`~llmeter.callbacks.cost.model.CostModel` + that gets saved, subclass :class:`RunCostDimensionBase` (which mixes in + :class:`~llmeter.serialization.Serializable`) rather than implementing this Protocol alone. """ async def before_run_start(self, run_config: _RunConfig) -> None: @@ -73,8 +89,7 @@ class RequestCostDimensionBase(Serializable, ABC): This class provides a default implementation of serialization (via :class:`~llmeter.serialization.Serializable`) and sets up an abstract method for - `calculate()`. It's fine if you don't want to derive from it directly - just be sure to - fully implement `IRequestCostDimension`! + `calculate()`. """ @abstractmethod @@ -89,8 +104,7 @@ class RunCostDimensionBase(Serializable, ABC): This class provides a default implementation of serialization (via :class:`~llmeter.serialization.Serializable`), a default empty `before_run_start` implementation, and abstract methods for the other requirements of the `IRunCostDimension` - protocol. It's fine if you don't want to derive from it directly - just make sure you fully - implement `IRunCostDimension`! + protocol. """ async def before_run_start(self, run_config: _RunConfig) -> None: diff --git a/llmeter/callbacks/cost/model.py b/llmeter/callbacks/cost/model.py index 2e23439..1255673 100644 --- a/llmeter/callbacks/cost/model.py +++ b/llmeter/callbacks/cost/model.py @@ -27,14 +27,10 @@ class CostModel(Callback): def __init__( self, request_dims: ( - dict[str, IRequestCostDimension] - | list[IRequestCostDimension] - | None + dict[str, IRequestCostDimension] | list[IRequestCostDimension] | None ) = None, run_dims: ( - dict[str, IRunCostDimension] - | list[IRunCostDimension] - | None + dict[str, IRunCostDimension] | list[IRunCostDimension] | None ) = None, ): """Create a CostModel @@ -94,8 +90,9 @@ async def calculate_request_cost( Args: response: The InvocationResponse to estimate costs for - save: Set `True` to also store the result in `response.cost`, in addition to returning - it. Defaults to `False` + save: Set `True` to also store the result in `response.annotations` (under `cost_` + prefixed keys), in addition to returning it. Because `annotations` is a persisted + field, saved costs survive `Result` save/load. Defaults to `False` """ dim_costs = CalculatedCostWithDimensions( **{ @@ -104,7 +101,10 @@ async def calculate_request_cost( } ) if save: - dim_costs.save_on_namespace(response, key_prefix="cost_") + # Store per-response costs in `response.annotations` (a declared field) so they + # persist through to_json/from_json and disk save/load, rather than as loose + # attributes on the response (which asdict-based serialization would drop). + dim_costs.save_on_namespace(response.annotations, key_prefix="cost_") return dim_costs async def calculate_run_cost( @@ -127,10 +127,7 @@ async def calculate_run_cost( it. Defaults to `False`. """ run_cost = CalculatedCostWithDimensions( - **{ - name: await dim.calculate(result) - for name, dim in self.run_dims.items() - } + **{name: await dim.calculate(result) for name, dim in self.run_dims.items()} ) if recalculate_request_costs: resp_costs = [ @@ -138,15 +135,17 @@ async def calculate_run_cost( for r in result.responses ] else: - resp_costs = list(filter( - lambda c: c, - ( - CalculatedCostWithDimensions.load_from_namespace( - r, key_prefix="cost_" - ) - for r in result.responses - ), - )) + resp_costs = list( + filter( + lambda c: c, + ( + CalculatedCostWithDimensions.load_from_namespace( + r.annotations, key_prefix="cost_" + ) + for r in result.responses + ), + ) + ) if len(resp_costs): run_cost.merge(sum(resp_costs)) # type: ignore if save: @@ -164,7 +163,8 @@ async def calculate_run_cost( async def after_invoke(self, response: InvocationResponse) -> None: """LLMeter Callback.after_invoke hook - Calls calculate_request_cost() with `save=True` to save the cost on the InvocationResponse. + Calls calculate_request_cost() with `save=True` to save the per-request cost into + `response.annotations` (under `cost_` prefixed keys), so it persists with the response. """ await self.calculate_request_cost(response, save=True) diff --git a/llmeter/endpoints/base.py b/llmeter/endpoints/base.py index 47d7c78..1914cc3 100644 --- a/llmeter/endpoints/base.py +++ b/llmeter/endpoints/base.py @@ -8,13 +8,14 @@ import copy import functools import importlib +import inspect import json import logging import time import warnings from abc import ABC, abstractmethod from collections.abc import Callable -from dataclasses import asdict, dataclass +from dataclasses import asdict, dataclass, field, fields from datetime import datetime, timezone from typing import Any, Generic, TypeVar from uuid import uuid4 @@ -57,6 +58,12 @@ class InvocationResponse: time_per_output_token (float): The average time taken to generate each token in the response. error (str): Any error that occurred during invocation. request_time: The wall-clock time when the request was sent. + annotations (dict): Free-form extra data attached to this response, for example by + callbacks. This is the **preferred** place for a `Callback` to store additional + per-response fields (rather than setting arbitrary attributes on the response), because + `annotations` is a declared field and therefore round-trips through + `to_json`/`from_json` and disk persistence. On load, any *unrecognized* top-level keys + (e.g. from older files or other LLMeter versions) are collected into `annotations` too. """ response_text: str | None @@ -73,6 +80,7 @@ class InvocationResponse: error: str | None = None retries: int | None = None request_time: datetime | None = None + annotations: dict[str, Any] = field(default_factory=dict) @classmethod def from_json(cls, json_str: str) -> "InvocationResponse": @@ -85,6 +93,12 @@ def from_json(cls, json_str: str) -> "InvocationResponse": * `datetime`-annotated fields are parsed from ISO-8601 strings back to Python `datetime` * `bytes`-typed fields and `__llmeter_bytes__` markers in nested payloads are restored + For legacy compatability, any top-level keys that are *not* recognized fields are currently + collected into [`annotations`][llmeter.endpoints.base.InvocationResponse] rather than being + dropped or raising an error. This supports older files where CostModel callbacks wrote + extra fields (such as `cost_*`) directly onto the response - but may be dropped in a future + version. + Args: json_str: A JSON string representation of an InvocationResponse (produced by `to_json` or similar). @@ -101,6 +115,18 @@ def from_json(cls, json_str: str) -> "InvocationResponse": """ data = json.loads(json_str, object_hook=bytes_decoder) restore_dataclass_types(cls, data) + # Route any unrecognized top-level keys into `annotations` rather than failing. This keeps + # forward/backward compatibility: e.g. older files where callbacks wrote extra fields (like + # `cost_*`) directly onto the response, or fields written by a different LLMeter version. + known = {f.name for f in fields(cls)} + extras = {k: data.pop(k) for k in list(data) if k not in known} + if extras: + data["annotations"] = {**extras, **(data.get("annotations") or {})} + logger.debug( + "Loaded %d unrecognized InvocationResponse field(s) into `annotations`: %s", + len(extras), + ", ".join(sorted(extras)), + ) return cls(**data) def to_json(self, default=json_default, **kwargs) -> str: @@ -533,7 +559,7 @@ def load_from_file(cls, path: ReadablePathLike) -> "Endpoint": endpoint_type = data.pop("endpoint_type") endpoint_module = importlib.import_module("llmeter.endpoints") endpoint_class = getattr(endpoint_module, endpoint_type) - return endpoint_class(**data) + return endpoint_class(**_filter_legacy_ctor_kwargs(endpoint_class, data)) @classmethod def load(cls, endpoint_config: dict) -> "Endpoint": # type: ignore @@ -559,4 +585,37 @@ def load(cls, endpoint_config: dict) -> "Endpoint": # type: ignore endpoint_type = endpoint_config.pop("endpoint_type") endpoint_module = importlib.import_module("llmeter.endpoints") endpoint_class = getattr(endpoint_module, endpoint_type) - return endpoint_class(**endpoint_config) + return endpoint_class( + **_filter_legacy_ctor_kwargs(endpoint_class, endpoint_config) + ) + + +def _filter_legacy_ctor_kwargs(endpoint_class: type, config: dict) -> dict: + """Drop legacy config keys that the target endpoint constructor won't accept. + + Older LLMeter configs persisted derived/read-only attributes (notably ``provider``, + which endpoints now set internally) alongside the real constructor arguments. Passing + those through to a modern ``__init__`` raises ``TypeError``, so we filter the dict down + to the parameters the constructor actually declares. + + If the constructor accepts ``**kwargs`` the dict is passed through unchanged. + + Args: + endpoint_class: The endpoint class about to be instantiated. + config: The legacy configuration dict (already stripped of ``endpoint_type``). + + Returns: + A copy of ``config`` containing only keys the constructor accepts. + """ + params = inspect.signature(endpoint_class.__init__).parameters + if any(p.kind == p.VAR_KEYWORD for p in params.values()): + return config + accepted = {name for name in params if name != "self"} + dropped = set(config) - accepted + if dropped: + logger.debug( + "Ignoring legacy config field(s) not accepted by %s.__init__: %s", + endpoint_class.__name__, + ", ".join(sorted(dropped)), + ) + return {k: v for k, v in config.items() if k in accepted} diff --git a/llmeter/runner.py b/llmeter/runner.py index e14b3d1..508579f 100644 --- a/llmeter/runner.py +++ b/llmeter/runner.py @@ -172,7 +172,8 @@ def save( ): """Save the configuration to a disk or cloud storage. - Uses the ``__getstate__``/``__setstate__`` protocol for Endpoint, Tokenizer, + Uses the LLMeter state-based serialization protocol + (``_get_llmeter_state``/``_set_llmeter_state``) for Endpoint, Tokenizer, and Callback serialization. Args: @@ -199,9 +200,7 @@ def save( config_copy.callbacks = [dump_object(cb) for cb in self.callbacks] with run_config_path.open("w") as f: - json.dump( - asdict(config_copy), f, default=json_default, indent=4 - ) + json.dump(asdict(config_copy), f, default=json_default, indent=4) @classmethod def load(cls, load_path: ReadablePathLike, file_name: str = "run_config.json"): @@ -218,21 +217,25 @@ def load(cls, load_path: ReadablePathLike, file_name: str = "run_config.json"): with (load_path / file_name).open() as f: config = json.load(f) - # Restore endpoint + # Restore endpoint and tokenizer. Only unwrap the new-format + # (``__llmeter_class__``) envelope here; legacy dicts (``endpoint_type`` / + # ``tokenizer_module``) are passed through untouched so ``__post_init__`` can + # route them to ``Endpoint.load`` / ``Tokenizer.load``. ep = config.get("endpoint") - if isinstance(ep, dict): + if isinstance(ep, dict) and "__llmeter_class__" in ep: config["endpoint"] = load_object(ep) - # Restore tokenizer tok = config.get("tokenizer") - if isinstance(tok, dict): + if isinstance(tok, dict) and "__llmeter_class__" in tok: config["tokenizer"] = load_object(tok) - # Restore callbacks + # Restore callbacks (new-format envelopes only; anything else is left as-is) cbs = config.get("callbacks") if isinstance(cbs, list): config["callbacks"] = [ - load_object(cb) if isinstance(cb, dict) else cb + load_object(cb) + if isinstance(cb, dict) and "__llmeter_class__" in cb + else cb for cb in cbs ] diff --git a/llmeter/serialization.py b/llmeter/serialization.py index 46f55e6..3b59713 100644 --- a/llmeter/serialization.py +++ b/llmeter/serialization.py @@ -6,8 +6,10 @@ Key components: -- :class:`Serializable` — Mixin that gives any class automatic - ``__getstate__``/``__setstate__`` by introspecting ``__init__``. +- :class:`Serializable` — Mixin that gives any class an automatic *state* protocol + (``_get_llmeter_state``/``_set_llmeter_state``) by introspecting ``__init__``. + Deliberately distinct from the pickle protocol so ``pickle`` / ``copy`` / + ``deepcopy`` keep their native behavior. - :func:`dump_object` / :func:`load_object` — Full round-trip persistence using a ``{"__llmeter_class__": ..., "__llmeter_state__": ...}`` envelope. @@ -21,10 +23,29 @@ - :func:`datetime_to_str` / :func:`str_to_datetime` — Standardized datetime ↔ string conversion (UTC ISO-8601 with ``Z`` suffix). -.. warning:: Security +!!! warning "A warning on security" :func:`load_object` imports and instantiates whatever class path is in the ``__llmeter_class__`` field. Do not load configs from untrusted sources. + +### Implementation details to be aware of + +**State vs. identity:** We split an object's serialized form into two layers. The +*state* (``_get_llmeter_state``) contains the object's own field values and does *not* +describe which class those values belong to. The **envelope** built by +:func:`dump_object` adds the identity (``__llmeter_class__``) around that state. +This is why the state methods are named ``_get_llmeter_state`` / +``_set_llmeter_state``: the dict they handle is not a self-describing representation +of the object, only of its state, so a for example a public ``to_dict`` name would +over-promise. Keeping identity out of the state dict also avoids polluting the user's +field namespace and lets nested polymorphic values each carry their own envelope. + +**Why a base-class mixin and not a Protocol:** The value :class:`Serializable` +provides is an *inherited implementation* (via ``__init__`` introspection), not +merely a method contract — a class becomes serializable by inheriting it, writing +zero serialization code. A structural ``Protocol`` could only describe the methods, +not hand over that implementation, so it would not actually help an author satisfy +the contract; they would inherit :class:`Serializable` regardless. """ import base64 @@ -173,7 +194,12 @@ def restore_dataclass_types(cls, data: dict) -> None: class Serializable: - """Mixin providing automatic ``__getstate__``/``__setstate__`` via __init__ introspection. + """Mixin providing a state extraction protocol compatible with LLMeter serialization. + + Serialization in LLMeter uses a state extraction protocol somewhat similar to, but deliberately + separate from, the (``__getstate__`` / ``__setstate__``) interface used by `pickle`, + `copy.copy`, and `copy.deepcopy`. Subclasses remain natively picklable/copyable, but can also + be saved to and loaded from LLMeter's JSON-based format. Works with plain classes, ``@dataclass``, and any class whose ``__init__`` parameters correspond to instance attributes (``self.x`` or ``self._x``). @@ -182,7 +208,17 @@ class Serializable: :func:`dump_object` / :func:`load_object`. """ - def __getstate__(self) -> dict: + def _get_llmeter_state(self) -> dict: + """Extract a JSON-serializable state dict by introspecting ``__init__`` parameters. + + Returns the object's constructor arguments (looked up as ``self.`` or + ``self._``), recursively serialized. Note that: + + 1. This is state only — it carries no class identity; :func:`dump_object` wraps it with + ``__llmeter_class__``. + 2. Any properties not exposed as `__init__` arguments will not be persisted. Override your + class' state get and set methods if you need different behaviour. + """ sig = inspect.signature(self.__init__) state = {} for name, param in sig.parameters.items(): @@ -196,15 +232,22 @@ def __getstate__(self) -> dict: state[name] = _serialize_value(getattr(self, f"_{name}")) return state - def __setstate__(self, state: dict) -> None: + def _set_llmeter_state(self, state: dict) -> None: + """Restore this instance from a state dict produced by :meth:`_get_llmeter_state`. + + Note this rebuilds the object by *calling the constructor* with the state as + keyword arguments (it is a ``from_dict``-style reconstruction, not pickle-style + state restoration onto an already-built instance). ``load_object`` first + creates a bare instance via ``__new__``, then calls this to populate it. + """ deserialized = {k: _deserialize_value(v) for k, v in state.items()} self.__init__(**deserialized) def save_to_file(self, path: WritablePathLike) -> Path: """Save this object to a JSON file. - Uses the ``__getstate__`` protocol. Override ``__getstate__`` (not this method) - if custom serialization is needed. + Uses the :meth:`_get_llmeter_state` protocol. Override + :meth:`_get_llmeter_state` (not this method) if custom serialization is needed. Args: path: (Local or Cloud) path where the object will be saved. @@ -249,8 +292,8 @@ def dump_object(obj: Any) -> dict: Serialization strategy (checked in order): - 1. If the object has a custom ``__getstate__`` (not :func:`object.__getstate__`), - calls it to obtain the state dict. + 1. If the object implements the LLMeter state protocol + (:meth:`Serializable._get_llmeter_state`), calls it to obtain the state dict. 2. If the object is a dataclass, uses :func:`dataclasses.asdict`. 3. Otherwise, takes all public (non-underscore-prefixed) entries from ``__dict__``. @@ -261,11 +304,8 @@ def dump_object(obj: Any) -> dict: A JSON-serializable dict that :func:`load_object` can reconstruct. """ class_path = f"{obj.__class__.__module__}.{obj.__class__.__qualname__}" - if ( - hasattr(obj, "__getstate__") - and type(obj).__getstate__ is not object.__getstate__ - ): - state = obj.__getstate__() + if hasattr(obj, "_get_llmeter_state"): + state = obj._get_llmeter_state() elif is_dataclass(obj) and not isinstance(obj, type): state = asdict(obj) elif hasattr(obj, "__dict__"): @@ -279,8 +319,9 @@ def load_object(data: dict) -> Any: """Restore an object from a type-tagged dict produced by :func:`dump_object`. Imports the module identified by ``__llmeter_class__``, instantiates the class - (bypassing ``__init__`` via ``__new__``), and calls ``__setstate__`` with the - persisted state dict. + (bypassing ``__init__`` via ``__new__``), and restores its state via + :meth:`Serializable._set_llmeter_state`. Objects that do not implement the + protocol are restored by assigning the (deserialized) state onto ``__dict__``. Args: data: A dict with ``__llmeter_class__`` and ``__llmeter_state__`` keys, as @@ -298,8 +339,19 @@ def load_object(data: dict) -> Any: module = importlib.import_module(module_path) cls = getattr(module, class_name) + state = data["__llmeter_state__"] obj = cls.__new__(cls) - obj.__setstate__(data["__llmeter_state__"]) + if hasattr(obj, "_set_llmeter_state"): + obj._set_llmeter_state(state) + else: + # Fallback for non-Serializable classes (state came from dump_object's dataclass / + # __dict__ branches). We assign attributes directly rather than calling __setstate__: + # `object` has no __setstate__ (so it would raise for a plain class), and if the class + # *does* define one it's the pickle hook, which may expect a different state shape than + # our JSON dict. Direct assignment (with per-value deserialization) is the safe generic + # restore, and mirrors pickle's own "restore __dict__ without calling __init__" default. + for key, value in state.items(): + setattr(obj, key, _deserialize_value(value)) return obj @@ -328,16 +380,19 @@ def _serialize_value(val: Any) -> Any: for types, fn in _SERIALIZERS: if isinstance(val, types): return fn(val) - if ( - hasattr(val, "__getstate__") - and type(val).__getstate__ is not object.__getstate__ - ): + if hasattr(val, "_get_llmeter_state"): return dump_object(val) if isinstance(val, dict): return {k: _serialize_value(v) for k, v in val.items()} if isinstance(val, (list, tuple)): return [_serialize_value(item) for item in val] - raise TypeError(f"Cannot serialize {type(val).__name__!r} object: {val!r}") + raise TypeError( + f"Cannot serialize {type(val).__name__!r} object: {val!r}. " + "Only JSON primitives, bytes, datetime, PathLike, dicts, lists/tuples, and " + "LLMeter-serializable objects are supported. To make a custom object " + "serializable (e.g. a custom cost dimension or callback), have its class " + "inherit from llmeter.serialization.Serializable." + ) def _deserialize_value(val: Any) -> Any: diff --git a/mkdocs.yml b/mkdocs.yml index 835ab3d..cb29416 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -42,7 +42,7 @@ nav: - user_guide/callbacks/cache_buster.md - user_guide/callbacks/cost.md - user_guide/callbacks/mlflow.md - - Serialization: serialization.md + - Save (and Reload) Results: user_guide/serialization.md - API Reference: - reference/index.md - callbacks: diff --git a/tests/unit/callbacks/cost/providers/test_sagemaker.py b/tests/unit/callbacks/cost/providers/test_sagemaker.py index 70d8cf5..3325eec 100644 --- a/tests/unit/callbacks/cost/providers/test_sagemaker.py +++ b/tests/unit/callbacks/cost/providers/test_sagemaker.py @@ -44,7 +44,7 @@ async def test_real_time_endpoint_compute_explicit_price(): result.total_test_time = None assert await dim.calculate(result) is None - dim_ser = dim.__getstate__() + dim_ser = dim._get_llmeter_state() assert dim_ser == { "instance_count": 1, "instance_type": "ml.doesnotexist", @@ -237,7 +237,7 @@ async def test_real_time_endpoint_storage_explicit_price(): result.total_test_time = None assert await dim.calculate(result) is None - dim_ser = dim.__getstate__() + dim_ser = dim._get_llmeter_state() assert dim_ser == { "gbs_provisioned": 5, "price_per_gb_hour": 9, diff --git a/tests/unit/callbacks/cost/test_dimensions.py b/tests/unit/callbacks/cost/test_dimensions.py index 6f16ee0..425de3f 100644 --- a/tests/unit/callbacks/cost/test_dimensions.py +++ b/tests/unit/callbacks/cost/test_dimensions.py @@ -33,7 +33,9 @@ async def test_cost_per_input_token(): "price_per_million": 30, "granularity": 10, } - dim_valid = InputTokens(price_per_million=spec["price_per_million"], granularity=spec["granularity"]) + dim_valid = InputTokens( + price_per_million=spec["price_per_million"], granularity=spec["granularity"] + ) success_response = InvocationResponse(response_text="hi", num_tokens_input=199999) assert await dim_valid.calculate(success_response) == 6 @@ -46,7 +48,7 @@ async def test_cost_per_input_token(): err_response = InvocationResponse.error_output() assert await dim_valid.calculate(err_response) is None - assert dim_valid.__getstate__() == { + assert dim_valid._get_llmeter_state() == { "price_per_million": 30, "granularity": 10, } @@ -58,7 +60,9 @@ async def test_cost_per_output_token(): "price_per_million": 40, "granularity": 10, } - dim_valid = OutputTokens(price_per_million=spec["price_per_million"], granularity=spec["granularity"]) + dim_valid = OutputTokens( + price_per_million=spec["price_per_million"], granularity=spec["granularity"] + ) success_response = InvocationResponse(response_text="hi", num_tokens_output=199999) assert await dim_valid.calculate(success_response) == 8 @@ -71,7 +75,7 @@ async def test_cost_per_output_token(): err_response = InvocationResponse.error_output() assert await dim_valid.calculate(err_response) is None - assert dim_valid.__getstate__() == { + assert dim_valid._get_llmeter_state() == { "price_per_million": 40, "granularity": 10, } @@ -83,7 +87,9 @@ async def test_cost_per_hour(): "price_per_hour": 30, "granularity_secs": 60, } - dim_valid = EndpointTime(price_per_hour=spec["price_per_hour"], granularity_secs=spec["granularity_secs"]) + dim_valid = EndpointTime( + price_per_hour=spec["price_per_hour"], granularity_secs=spec["granularity_secs"] + ) result = Result( [], @@ -100,7 +106,7 @@ async def test_cost_per_hour(): result.total_test_time = None assert await dim_valid.calculate(result) is None - assert dim_valid.__getstate__() == { + assert dim_valid._get_llmeter_state() == { "price_per_hour": 30, "granularity_secs": 60, } diff --git a/tests/unit/callbacks/cost/test_model.py b/tests/unit/callbacks/cost/test_model.py index e05263b..fa31532 100644 --- a/tests/unit/callbacks/cost/test_model.py +++ b/tests/unit/callbacks/cost/test_model.py @@ -9,6 +9,7 @@ from llmeter.serialization import dump_object, load_object from llmeter.callbacks.cost.model import CostModel from llmeter.callbacks.cost.results import CalculatedCostWithDimensions +from llmeter.endpoints.base import InvocationResponse def test_cost_model_serialization(): @@ -27,8 +28,8 @@ def test_cost_model_serialization(): assert restored.request_dims["TokensIn"].price_per_million == 30 assert restored.run_dims["ComputeSeconds"].price_per_hour == 50 - # __getstate__ produces a plain dict representation - d = model.__getstate__() + # get state produces a plain dict representation + d = model._get_llmeter_state() assert "request_dims" in d assert "run_dims" in d @@ -144,7 +145,7 @@ class DummyCostDimension: @pytest.mark.asyncio async def test_cost_model_callback_saves_request_costs(): - """By default, CostModel callbacks save request cost calculations to InvocationResponse""" + """By default, CostModel callbacks save request cost calculations to response.annotations""" dummy_req_dim = Mock() dummy_req_dim.calculate = AsyncMock(return_value=42) @@ -153,19 +154,56 @@ async def test_cost_model_callback_saves_request_costs(): run_dims=[], ) - response_mock = NonCallableMock() - assert await model.after_invoke(response_mock) is None - assert response_mock.cost_total == 42 - assert response_mock.cost_Mock == 42 # Class name is the default dimension name + response = InvocationResponse(response_text="hi") + assert await model.after_invoke(response) is None + # Costs are stored in the (persisted) annotations dict, not as loose attributes + assert response.annotations["cost_total"] == 42 + assert ( + response.annotations["cost_Mock"] == 42 + ) # Class name is the default dimension name # Check calculate_* fn produces same result as callback: assert await model.calculate_request_cost( - response_mock + response ) == CalculatedCostWithDimensions.load_from_namespace( - response_mock, key_prefix="cost_" + response.annotations, key_prefix="cost_" ) +@pytest.mark.asyncio +async def test_cost_model_request_costs_survive_response_roundtrip(): + """Per-response costs saved by CostModel persist through InvocationResponse to_json/from_json. + + This is the behavior that regressed when response serialization moved to ``asdict`` (which + drops loose attributes): costs stored in ``annotations`` must round-trip so a saved Result can + be reloaded with its per-request costs intact. + """ + from llmeter.callbacks.cost.dimensions import InputTokens, OutputTokens + + model = CostModel( + request_dims=[ + InputTokens(price_per_million=3.0), + OutputTokens(price_per_million=15.0), + ] + ) + response = InvocationResponse( + response_text="hello", num_tokens_input=1000, num_tokens_output=500 + ) + await model.after_invoke(response) + assert response.annotations["cost_total"] == pytest.approx(0.003 + 0.0075) + + # Round-trip through JSON (what gets written to responses.jsonl) + restored = InvocationResponse.from_json(response.to_json()) + assert restored.annotations == response.annotations + + # ...and the cost model can read the costs back off the restored response + reloaded = CalculatedCostWithDimensions.load_from_namespace( + restored.annotations, key_prefix="cost_" + ) + assert reloaded["InputTokens"] == pytest.approx(0.003) + assert reloaded["OutputTokens"] == pytest.approx(0.0075) + + @pytest.mark.asyncio async def test_cost_model_callback_saves_run_costs(): """By default, CostModel callbacks save run cost calculations to Result""" @@ -217,13 +255,17 @@ async def test_cost_model_combines_req_and_run_dims(): # Run the dummy test: run_mock = NonCallableMock() await model.before_run(run_mock) - response_mocks = [NonCallableMock(), NonCallableMock(), NonCallableMock()] - for r in response_mocks: + responses = [ + InvocationResponse(response_text="a"), + InvocationResponse(response_text="b"), + InvocationResponse(response_text="c"), + ] + for r in responses: await model.after_invoke(r) results_mock = NonCallableMock() update_contrib_stats_mock = Mock() results_mock._update_contributed_stats = update_contrib_stats_mock - results_mock.responses = response_mocks + results_mock.responses = responses results_mock.additional_metrics_for_aggregation = None await model.after_run(results_mock) diff --git a/tests/unit/callbacks/test_base.py b/tests/unit/callbacks/test_base.py index a648feb..2b450bf 100644 --- a/tests/unit/callbacks/test_base.py +++ b/tests/unit/callbacks/test_base.py @@ -35,12 +35,12 @@ def test_load_from_file_restores_callback(self, tmp_path): assert restored.nested is False def test_getstate_setstate_roundtrip(self): - """__getstate__ / __setstate__ round-trips correctly.""" + """get state / set state round-trips correctly.""" cb = MlflowCallback(step=7, nested=True) - state = cb.__getstate__() + state = cb._get_llmeter_state() restored = MlflowCallback.__new__(MlflowCallback) - restored.__setstate__(state) + restored._set_llmeter_state(state) assert restored.step == 7 assert restored.nested is True diff --git a/tests/unit/endpoints/test_endpoints.py b/tests/unit/endpoints/test_endpoints.py index bdaad99..2c54445 100644 --- a/tests/unit/endpoints/test_endpoints.py +++ b/tests/unit/endpoints/test_endpoints.py @@ -1,6 +1,8 @@ # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # SPDX-License-Identifier: Apache-2.0 +import json + import pytest import llmeter @@ -186,6 +188,57 @@ def test_endpoint_load_from_dict(): assert loaded_endpoint.provider == "test_provider" +def test_invocation_response_annotations_roundtrip(): + """The `annotations` dict round-trips through to_json/from_json.""" + resp = InvocationResponse( + response_text="hi", annotations={"cost_total": 0.01, "label": "x"} + ) + restored = InvocationResponse.from_json(resp.to_json()) + assert restored.annotations == {"cost_total": 0.01, "label": "x"} + + +def test_invocation_response_annotations_default_is_empty_dict(): + """annotations defaults to an independent empty dict per instance.""" + a = InvocationResponse(response_text="a") + b = InvocationResponse(response_text="b") + assert a.annotations == {} + a.annotations["x"] = 1 + assert b.annotations == {} # no shared mutable default + + +def test_from_json_collects_unknown_fields_into_annotations(): + """Unknown top-level keys (e.g. legacy callback-written fields) are captured into annotations.""" + raw = json.dumps( + { + "response_text": "hi", + "id": "r1", + "num_tokens_input": 10, + "cost_total": 0.02, # legacy callback-written field + "cost_InputTokens": 0.01, + "end_time": "2024-01-01T00:00:00Z", # since-removed field + } + ) + resp = InvocationResponse.from_json(raw) + + # Known fields are parsed normally + assert resp.response_text == "hi" + assert resp.id == "r1" + assert resp.num_tokens_input == 10 + # Unknown fields are preserved (not dropped, not raising) under annotations + assert resp.annotations == { + "cost_total": 0.02, + "cost_InputTokens": 0.01, + "end_time": "2024-01-01T00:00:00Z", + } + + +def test_from_json_merges_unknown_fields_with_existing_annotations(): + """Explicit annotations and stray unknown keys both survive load.""" + raw = json.dumps({"response_text": "hi", "annotations": {"kept": 1}, "stray": 2}) + resp = InvocationResponse.from_json(raw) + assert resp.annotations == {"kept": 1, "stray": 2} + + def test_invocation_response_to_dict(): response = InvocationResponse( id="test_id", diff --git a/tests/unit/test_legacy_data_load.py b/tests/unit/test_legacy_data_load.py new file mode 100644 index 0000000..7d80ce9 --- /dev/null +++ b/tests/unit/test_legacy_data_load.py @@ -0,0 +1,245 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Regression tests for loading *legacy-format* LLMeter data written by older releases. + +This covers artifacts produced before the current ``dump_object`` serialization, across the three +data types that can be reloaded: + +* **Run configs** (``run_config.json``) and **endpoint configs** - identified the endpoint class + with an ``endpoint_type`` key (rather than the current ``__llmeter_class__`` envelope) and the + tokenizer with a ``tokenizer_module`` key, and stored derived / read-only attributes such as + ``provider`` alongside real constructor arguments. Modern endpoint constructors compute + ``provider`` internally and reject it as a keyword argument. +* **Responses** (``responses.jsonl``, via ``Result.load``) - carried extra per-response fields that + callbacks wrote directly onto the response (e.g. ``cost_InputTokens`` / ``cost_total``), plus a + since-removed ``end_time`` field. Current ``InvocationResponse`` doesn't declare these, so a + strict load would raise; they are now collected into ``annotations`` instead. + +All of these previously broke loading of real, local data. These tests pin the backward-compatible +behavior so it can't silently regress on a minor-version release. If you intend to drop +legacy-format support, that is a breaking change and these tests should be updated deliberately. +""" + +import json + +import pytest + +from llmeter.endpoints.base import Endpoint, _filter_legacy_ctor_kwargs +from llmeter.endpoints.bedrock import BedrockConverse, BedrockConverseStream +from llmeter.results import Result +from llmeter.runner import _RunConfig +from llmeter.tokenizers import DummyTokenizer + + +def _legacy_endpoint_config() -> dict: + """A legacy endpoint config dict, matching the on-disk ``examples/outputs`` format. + + Note ``provider`` is present here but is *not* a parameter of + ``BedrockConverseStream.__init__`` (it is derived internally). + """ + return { + "endpoint_name": "amazon bedrock", + "model_id": "apac.amazon.nova-pro-v1:0", + "provider": "bedrock", + "region": "us-east-1", + "endpoint_type": "BedrockConverseStream", + } + + +def _legacy_run_config() -> dict: + """A full legacy ``run_config.json`` payload, mirroring real ``examples/outputs`` data.""" + return { + "endpoint": _legacy_endpoint_config(), + "output_path": "outputs/apac.amazon.nova-pro-v1:0/20251210-0115", + "tokenizer": {"tokenizer_module": "llmeter"}, + "clients": 20, + "n_requests": 10, + "payload": "outputs/apac.amazon.nova-pro-v1:0/20251210-0115/payload.jsonl", + "run_name": "20251210-0115", + "run_description": None, + "timeout": 60, + "callbacks": None, + } + + +class TestLegacyEndpointConfigLoad: + """Endpoint.load / load_from_file must accept legacy ``endpoint_type`` configs.""" + + def test_load_drops_derived_provider_field(self): + """A legacy dict carrying a derived ``provider`` must load, not raise TypeError.""" + endpoint = Endpoint.load(_legacy_endpoint_config()) + + assert isinstance(endpoint, BedrockConverseStream) + assert endpoint.model_id == "apac.amazon.nova-pro-v1:0" + assert endpoint.region == "us-east-1" + assert endpoint.endpoint_name == "amazon bedrock" + # provider is derived by the constructor, and should round-trip to the same value + assert endpoint.provider == "bedrock" + + def test_load_from_file_drops_derived_provider_field(self, tmp_path): + """The file-based legacy loader must behave the same as the dict loader.""" + path = tmp_path / "endpoint.json" + path.write_text(json.dumps(_legacy_endpoint_config())) + + endpoint = Endpoint.load_from_file(path) + + assert isinstance(endpoint, BedrockConverseStream) + assert endpoint.model_id == "apac.amazon.nova-pro-v1:0" + assert endpoint.region == "us-east-1" + assert endpoint.provider == "bedrock" + + def test_load_missing_endpoint_type_still_raises(self): + """Legacy loading still requires ``endpoint_type`` to identify the class.""" + config = _legacy_endpoint_config() + del config["endpoint_type"] + with pytest.raises(KeyError): + Endpoint.load(config) + + +class TestLegacyRunConfigLoad: + """_RunConfig.load must reconstruct endpoint + tokenizer from a legacy run_config.json.""" + + def test_load_reconstructs_legacy_run_config(self, tmp_path): + (tmp_path / "run_config.json").write_text(json.dumps(_legacy_run_config())) + + cfg = _RunConfig.load(tmp_path) + + # Endpoint: legacy endpoint_type -> concrete class, provider derived + assert isinstance(cfg._endpoint, BedrockConverseStream) + assert cfg._endpoint.model_id == "apac.amazon.nova-pro-v1:0" + assert cfg._endpoint.region == "us-east-1" + assert cfg._endpoint.provider == "bedrock" + + # Tokenizer: legacy tokenizer_module -> DummyTokenizer + assert isinstance(cfg._tokenizer, DummyTokenizer) + + # Scalar fields survive unchanged + assert cfg.clients == 20 + assert cfg.n_requests == 10 + assert cfg.timeout == 60 + assert cfg.run_name == "20251210-0115" + + def test_load_legacy_run_config_custom_filename(self, tmp_path): + """The legacy loader honors a custom config file name.""" + (tmp_path / "my_config.json").write_text(json.dumps(_legacy_run_config())) + + cfg = _RunConfig.load(tmp_path, file_name="my_config.json") + + assert isinstance(cfg._endpoint, BedrockConverseStream) + assert isinstance(cfg._tokenizer, DummyTokenizer) + + +class TestLegacyResponsesLoad: + """Result.load must tolerate older responses.jsonl carrying extra per-response keys. + + Older LLMeter versions let callbacks (e.g. the cost callback) write extra fields such as + ``cost_InputTokens`` / ``cost_total`` directly onto each response, and serialized them into + ``responses.jsonl``. A couple of runs also carry a since-removed ``end_time`` response field. + Current ``InvocationResponse`` doesn't declare these, so a strict ``cls(**data)`` load would + raise ``TypeError``. These are collected into ``annotations`` instead, so the data is preserved + and the file still loads. + """ + + def test_load_result_with_legacy_response_extra_keys(self, tmp_path): + summary = { + "total_requests": 2, + "clients": 1, + "n_requests": 2, + "model_id": "some-model", + } + (tmp_path / "summary.json").write_text(json.dumps(summary)) + + lines = [ + json.dumps( + { + "response_text": "a", + "id": "r1", + "num_tokens_input": 10, + "num_tokens_output": 5, + "cost_InputTokens": 0.01, + "cost_OutputTokens": 0.02, + "cost_total": 0.03, + } + ), + json.dumps( + { + "response_text": "b", + "id": "r2", + "num_tokens_input": 20, + "num_tokens_output": 8, + "cost_total": 0.05, + "end_time": "2024-01-01T00:00:00Z", # since-removed field + } + ), + ] + (tmp_path / "responses.jsonl").write_text("\n".join(lines) + "\n") + + # Previously raised TypeError: unexpected keyword argument 'cost_InputTokens' + result = Result.load(tmp_path) + + assert len(result.responses) == 2 + r1, r2 = result.responses + + # Core fields are still parsed normally + assert r1.num_tokens_input == 10 + assert r2.id == "r2" + + # Legacy extra keys are preserved under annotations rather than dropped + assert r1.annotations == { + "cost_InputTokens": 0.01, + "cost_OutputTokens": 0.02, + "cost_total": 0.03, + } + assert r2.annotations["cost_total"] == 0.05 + assert r2.annotations["end_time"] == "2024-01-01T00:00:00Z" + + def test_load_result_without_summary_recovers_legacy_responses(self, tmp_path): + """Metadata recovery (no summary.json) also streams responses and must tolerate extras.""" + lines = [ + json.dumps( + { + "response_text": "a", + "id": "r1", + "request_time": "2024-01-01T00:00:00Z", + "cost_total": 0.03, + } + ), + ] + (tmp_path / "responses.jsonl").write_text("\n".join(lines) + "\n") + + result = Result.load(tmp_path) # no summary.json -> _recover_metadata path + + assert len(result.responses) == 1 + assert result.responses[0].annotations["cost_total"] == 0.03 + + +class TestFilterLegacyCtorKwargs: + """Unit tests for the constructor-kwarg filtering helper.""" + + def test_drops_keys_not_accepted_by_constructor(self): + config = { + "model_id": "m", + "region": "us-east-1", + "provider": "bedrock", # not a BedrockConverseStream.__init__ param + } + filtered = _filter_legacy_ctor_kwargs(BedrockConverseStream, config) + + assert "provider" not in filtered + assert filtered == {"model_id": "m", "region": "us-east-1"} + # Original dict is not mutated + assert "provider" in config + + def test_keeps_keys_the_constructor_accepts(self): + config = {"model_id": "m", "endpoint_name": "n", "region": "us-east-1"} + filtered = _filter_legacy_ctor_kwargs(BedrockConverse, config) + assert filtered == config + + def test_passthrough_when_constructor_accepts_var_keyword(self): + class AcceptsKwargs: + def __init__(self, model_id, **kwargs): + self.model_id = model_id + + config = {"model_id": "m", "anything": 1, "provider": "x"} + # **kwargs means every key is acceptable, so nothing is dropped + assert _filter_legacy_ctor_kwargs(AcceptsKwargs, config) == config diff --git a/tests/unit/test_property_save_load.py b/tests/unit/test_property_save_load.py index b5198fd..804f334 100644 --- a/tests/unit/test_property_save_load.py +++ b/tests/unit/test_property_save_load.py @@ -273,6 +273,40 @@ def test_run_config_creates_parent_directories(self, config): assert config_file.exists() assert config_file.parent.exists() + def test_run_config_save_load_roundtrip(self): + """save/load should round-trip a real endpoint via the ``__llmeter_class__`` envelope. + + The property tests above only exercise the *save* side (with a mock endpoint), so this + is the coverage for ``_RunConfig.load`` reconstructing a real endpoint from the current + serialization format. It also pins that format (no legacy ``endpoint_type`` key), so a + future backward-compat change can't silently break current-format loading. Legacy-format + loading is covered separately in ``test_legacy_data_load.py``. + """ + from llmeter.endpoints.bedrock import BedrockConverseStream + + with tempfile.TemporaryDirectory() as tmpdir: + output_path = Path(tmpdir) + config = _RunConfig( + endpoint=BedrockConverseStream( + model_id="apac.amazon.nova-pro-v1:0", region="us-east-1" + ), + clients=3, + n_requests=7, + ) + config.save(output_path) + + # New saves use the type-tagged envelope, not the legacy endpoint_type key + saved = json.loads((output_path / "run_config.json").read_text()) + assert "__llmeter_class__" in saved["endpoint"] + assert "endpoint_type" not in saved["endpoint"] + + loaded = _RunConfig.load(output_path) + assert isinstance(loaded._endpoint, BedrockConverseStream) + assert loaded._endpoint.model_id == "apac.amazon.nova-pro-v1:0" + assert loaded._endpoint.region == "us-east-1" + assert loaded.clients == 3 + assert loaded.n_requests == 7 + # Result save/load property tests class TestResultSaveLoadProperties: diff --git a/tests/unit/test_runner_backlog_ui.py b/tests/unit/test_runner_backlog_ui.py index f47aac0..fd73c9a 100644 --- a/tests/unit/test_runner_backlog_ui.py +++ b/tests/unit/test_runner_backlog_ui.py @@ -37,7 +37,7 @@ def mock_endpoint(): @pytest.fixture def mock_tokenizer(): - # Tokenizer now serializes via the Serializable mixin (__getstate__), not a + # Tokenizer now serializes via the Serializable mixin (_get_llmeter_state), not a # to_dict() method. These tests run without an output_path, so the tokenizer # is never actually serialized — a bare spec'd mock is all that's needed. yield MagicMock(spec=Tokenizer) @@ -381,9 +381,6 @@ def slow_invoke(payload): mock_endpoint.invoke.side_effect = slow_invoke - backlog_bar_created = [False] - original_run = run._run - # Patch tqdm to track if a "Processing backlog" bar is created with patch("llmeter.runner.tqdm") as mock_tqdm: mock_tqdm.return_value = MagicMock() diff --git a/tests/unit/test_serialize_value_roundtrip.py b/tests/unit/test_serialize_value_roundtrip.py index f699d0a..68265cd 100644 --- a/tests/unit/test_serialize_value_roundtrip.py +++ b/tests/unit/test_serialize_value_roundtrip.py @@ -124,7 +124,7 @@ def test_path_serializes_to_string(self): class TestSerializableWithDatetimeAndBytes: """End-to-end: a Serializable class with datetime/bytes fields round-trips - via __getstate__/__setstate__.""" + via _get_llmeter_state/_set_llmeter_state.""" def test_serializable_with_datetime_field(self): class MyObj(Serializable): @@ -134,10 +134,10 @@ def __init__(self, created_at: datetime, name: str = "test"): dt = datetime(2024, 6, 15, 10, 30, 0, tzinfo=timezone.utc) obj = MyObj(created_at=dt, name="hello") - state = obj.__getstate__() + state = obj._get_llmeter_state() restored = MyObj.__new__(MyObj) - restored.__setstate__(state) + restored._set_llmeter_state(state) assert restored.created_at == dt assert isinstance(restored.created_at, datetime) @@ -150,10 +150,10 @@ def __init__(self, payload: bytes, label: str = "x"): self.label = label obj = MyObj(payload=b"\xde\xad\xbe\xef", label="binary") - state = obj.__getstate__() + state = obj._get_llmeter_state() restored = MyObj.__new__(MyObj) - restored.__setstate__(state) + restored._set_llmeter_state(state) assert restored.payload == b"\xde\xad\xbe\xef" assert isinstance(restored.payload, bytes) @@ -168,10 +168,10 @@ def __init__(self, ts: datetime, data: bytes, info: dict): dt = datetime(2024, 1, 1, tzinfo=timezone.utc) obj = MyObj(ts=dt, data=b"raw", info={"nested_bytes": b"\x01", "nested_ts": dt}) - state = obj.__getstate__() + state = obj._get_llmeter_state() restored = MyObj.__new__(MyObj) - restored.__setstate__(state) + restored._set_llmeter_state(state) assert restored.ts == dt assert restored.data == b"raw" From 6087a9318d637d3dc6df700c8de1073ec3c56833 Mon Sep 17 00:00:00 2001 From: Alex Thewsey Date: Thu, 23 Jul 2026 00:29:05 +0800 Subject: [PATCH 15/20] chore(lint): ruff format --- tests/unit/test_interrupted_run.py | 12 +++--------- tests/unit/test_serialization_properties.py | 14 +++++++------- 2 files changed, 10 insertions(+), 16 deletions(-) diff --git a/tests/unit/test_interrupted_run.py b/tests/unit/test_interrupted_run.py index 438878e..b813d38 100644 --- a/tests/unit/test_interrupted_run.py +++ b/tests/unit/test_interrupted_run.py @@ -29,9 +29,7 @@ def sample_responses(): time_to_last_token=0.3 * (i + 1), num_tokens_input=10 * (i + 1), num_tokens_output=20 * (i + 1), - request_time=datetime( - 2025, 6, 1, 10, 0, i * 2, tzinfo=timezone.utc - ), + request_time=datetime(2025, 6, 1, 10, 0, i * 2, tzinfo=timezone.utc), ) for i in range(5) ] @@ -123,9 +121,7 @@ def test_load_derives_timestamps_from_responses(self, interrupted_run_dir): assert result.last_request_time == datetime( 2025, 6, 1, 10, 0, 8, tzinfo=timezone.utc ) - assert result.end_time == datetime( - 2025, 6, 1, 10, 0, 8, tzinfo=timezone.utc - ) + assert result.end_time == datetime(2025, 6, 1, 10, 0, 8, tzinfo=timezone.utc) def test_load_sets_total_requests_from_responses(self, interrupted_run_dir): """total_requests should be set to the number of recovered responses.""" @@ -686,9 +682,7 @@ def slow_invoke(payload): import time time.sleep(0.05) - return InvocationResponse( - id="x", input_prompt="test", response_text="resp" - ) + return InvocationResponse(id="x", input_prompt="test", response_text="resp") mock_endpoint.invoke = slow_invoke diff --git a/tests/unit/test_serialization_properties.py b/tests/unit/test_serialization_properties.py index d38ded1..cd72a8e 100644 --- a/tests/unit/test_serialization_properties.py +++ b/tests/unit/test_serialization_properties.py @@ -444,9 +444,7 @@ def test_property_7_non_binary_payloads_are_backward_compatible(self, payload): # Feature: json-serialization-optimization, Property 7: Non-binary payloads are backward compatible """ # Serialize with the new encoder - serialized_with_encoder = json.dumps( - payload, default=json_default - ) + serialized_with_encoder = json.dumps(payload, default=json_default) # Serialize with standard json.dumps serialized_standard = json.dumps(payload) @@ -811,7 +809,9 @@ class CustomObject: def __str__(self): return "custom_string_representation" - decoded = json.loads(json.dumps({"custom": CustomObject()}, default=json_default)) + decoded = json.loads( + json.dumps({"custom": CustomObject()}, default=json_default) + ) assert decoded["custom"] == "custom_string_representation" def test_none_when_str_conversion_fails(self): @@ -821,7 +821,7 @@ class FailingObject: def __str__(self): raise RuntimeError("Cannot convert to string") - decoded = json.loads(json.dumps({"failing": FailingObject()}, default=json_default)) + decoded = json.loads( + json.dumps({"failing": FailingObject()}, default=json_default) + ) assert decoded["failing"] is None - - From b462fb55246f9a0a047145893956f094585e4e1b Mon Sep 17 00:00:00 2001 From: Alex Thewsey Date: Thu, 23 Jul 2026 02:52:12 +0800 Subject: [PATCH 16/20] doc: Small Zensical fixes docstrings, typings, uguide typos. --- docs/user_guide/callbacks/index.md | 4 ++-- docs/user_guide/serialization.md | 14 +++++++------- llmeter/callbacks/cost/dimensions.py | 26 ++++++++++++++------------ llmeter/endpoints/base.py | 28 +++++++++++++++------------- llmeter/results.py | 6 ++++-- llmeter/serialization.py | 2 +- 6 files changed, 43 insertions(+), 37 deletions(-) diff --git a/docs/user_guide/callbacks/index.md b/docs/user_guide/callbacks/index.md index 43fd3de..46e7e65 100644 --- a/docs/user_guide/callbacks/index.md +++ b/docs/user_guide/callbacks/index.md @@ -46,6 +46,6 @@ async def after_invoke(self, response): response.annotations["my_metric"] = compute_something(response) ``` -Any extra fields added directly to the `InvocationResponse` object itself will *not* be preserved when the responses are saved to file. For legacy compatability, any unrecognized fields found in older responses.jsonl files are *currently* callected back to `annotations`, but this behaviour may be dropped in future. +Any extra fields added directly to the `InvocationResponse` object itself will *not* be preserved when the responses are saved to file. For legacy compatability, any unrecognized fields found in older responses.jsonl files are *currently* collected back to `annotations`, but this behaviour may be dropped in future. -Likewise for extra **run-level** data, `after_run` can add numeric statistics to the `Result` via [`_update_contributed_stats`](../../reference/results/#llmeter.results.Result), which are persisted alongside the run's other stats. (This is how the built-in [cost callback](cost.md) records run-level cost metrics.) +Likewise for extra **run-level** data, `after_run` can add numeric statistics to the `Result` via [`_update_contributed_stats`](../../reference/results/#llmeter.results.Result), which are persisted alongside the run's other stats. (This is how the built-in [cost modelling callback](cost.md) records run-level cost metrics.) diff --git a/docs/user_guide/serialization.md b/docs/user_guide/serialization.md index 12c0cad..a29557b 100644 --- a/docs/user_guide/serialization.md +++ b/docs/user_guide/serialization.md @@ -3,7 +3,7 @@ !!! warning "Important security warning" Do not use LLMeter file load functions on data from untrusted sources! For more details on why, see the following section. -With LLMeter, you can save comprehensive test configuration and result information to files - whether locally or on the Cloud - and load past runs back into Python for analysis later. For example: +LLMeter can save your test configurations and results to files - whether locally or on the Cloud - and load past runs back into Python for analysis later. For example: ```python from llmeter.results import Result @@ -31,7 +31,7 @@ results = Result.load(f"{base_output_path}/my-cool-run") ``` !!! note "A note on performance" - While it's possible for LLMeter to write run outputs directly to Cloud object stores like + While it's *possible* for LLMeter to write run outputs directly to Cloud object stores like [Amazon S3](https://aws.amazon.com/s3/), remember it might reduce the maximum throughput you can drive in *high-volume* tests, since it will consume network bandwidth. @@ -41,7 +41,7 @@ individual responses). However, LLMeter also handles some more complex data type - Binary data in request/response payloads (such as images) - [Callbacks](./callbacks) configured on the test Run, including LLMeter built-ins as well as *your custom* callback classes -Implementation of de/serialization functionality is centralized in the [`llmeter.serialization`](../reference/serialization) module. +This core de/serialization functionality is implemented in the [`llmeter.serialization`](../reference/serialization) module. ## How complex data types are represented and loaded @@ -55,11 +55,11 @@ Objects that are not natively JSON serializable (but support LLMeter's serializa } ``` -- `datetime` objects are serialized to ISO8601 format strings converted to UTC timezone, like `2024-01-01T00:00:00Z`. +- `datetime` objects are stored as ISO-8601 format strings in UTC timezone, like `2024-01-01T00:00:00Z`. - `bytes` objects are serialized to base64 strings in a special `{"__llmeter_bytes__": ""}` wrapper. -- Objects implementing LLMeter's [Serializable](../reference/serialization.md#llmeter.serialization.Serializable) interface (including for example Endpoints and Callbacks) are represented as dicts with the `__llmeter_{class/state}__` properties as shown above +- Objects implementing LLMeter's [`Serializable`](../reference/serialization.md#llmeter.serialization.Serializable) interface (including for example Endpoints and Callbacks) are represented as dicts with the `__llmeter_{class/state}__` properties as shown above -When **loading** these objects back from file, LLMeter will import and instantiate whatever class path is found in the `__llmeter_class__` field. +When **loading** these objects back from file, LLMeter will try to import and instantiate **any** class path saved in the `__llmeter_class__` field. This is the same trust model as Python's native [pickle](https://docs.python.org/3/library/pickle.html) library: It is possible to construct malicious data which will run arbitrary code during loading, so **only load data that you trust**. @@ -184,7 +184,7 @@ runner = Runner(endpoint=BedrockConverse(...), callbacks=[MlflowCallback(step=1) runner.save(output_path="/tmp/run") # writes /tmp/run/run_config.json ``` -If a callback (or any other configured object) can't be serialized because its class doesn't inherit `Serializable`, saving raises a `TypeError`. This is why custom callbacks and cost dimensions should subclass the relevant LLMeter base class - see [`Callback`](../reference/callbacks/base.md) and the [cost dimension base classes](../reference/callbacks/cost/dimensions.md). +If a callback (or any other configured object) can't be serialized because its class doesn't inherit `Serializable`, saving raises a `TypeError`. This is why custom callbacks and cost dimensions should subclass the relevant LLMeter base class - see [`Callback`](../reference/callbacks/base.md#llmeter.callbacks.base.Callback) and the [cost dimension base classes](../reference/callbacks/cost/dimensions.md). ### Dataclasses work automatically diff --git a/llmeter/callbacks/cost/dimensions.py b/llmeter/callbacks/cost/dimensions.py index 7a12a26..466d7a3 100644 --- a/llmeter/callbacks/cost/dimensions.py +++ b/llmeter/callbacks/cost/dimensions.py @@ -35,12 +35,12 @@ class IRequestCostDimension(Protocol): !!! warning This interface describes only the *calculation* behavior. It intentionally says nothing about serialization: a dimension used purely in-memory need not be serializable. However, - to be usable inside a :class:`~llmeter.callbacks.cost.model.CostModel` that gets **saved** - (e.g. as a run callback, or via ``save_to_file``), a dimension must **additionally** be - LLMeter-serializable. The supported way to get that is to subclass - :class:`RequestCostDimensionBase` (which mixes in - :class:`~llmeter.serialization.Serializable`); implementing this Protocol alone is not - sufficient for persistence. + to be usable inside a [`CostModel`][llmeter.callbacks.cost.model.CostModel] that gets + **saved** (e.g. as a run callback, or via ``save_to_file``), a dimension must + **additionally** be LLMeter-serializable. The supported way to get that is to subclass + [`RequestCostDimensionBase`][llmeter.callbacks.cost.dimensions.RequestCostDimensionBase] + (which mixes in [`Serializable`][llmeter.serialization.Serializable]; implementing this + Protocol alone is not sufficient for persistence. """ async def calculate(self, response: InvocationResponse) -> float | None: @@ -57,10 +57,12 @@ class IRunCostDimension(Protocol): (request-independent) cost-per-hour are important. !!! warning - As with :class:`IRequestCostDimension`, this interface covers only calculation behavior and - not serialization. To be usable inside a :class:`~llmeter.callbacks.cost.model.CostModel` - that gets saved, subclass :class:`RunCostDimensionBase` (which mixes in - :class:`~llmeter.serialization.Serializable`) rather than implementing this Protocol alone. + As with `IRequestCostDimension`, this interface covers only calculation behavior and not + serialization. To be usable inside a + [`CostModel`][llmeter.callbacks.cost.model.CostModel] that gets saved, subclass + [`RunCostDimensionBase`][llmeter.callbacks.cost.dimensions.RunCostDimensionBase] (which + mixes in [`Serializable`][llmeter.serialization.Serializable]) rather than implementing + this Protocol alone. """ async def before_run_start(self, run_config: _RunConfig) -> None: @@ -88,7 +90,7 @@ class RequestCostDimensionBase(Serializable, ABC): """Base class for implementing per-request cost model dimensions This class provides a default implementation of serialization (via - :class:`~llmeter.serialization.Serializable`) and sets up an abstract method for + [`Serializable`][llmeter.serialization.Serializable]) and sets up an abstract method for `calculate()`. """ @@ -102,7 +104,7 @@ class RunCostDimensionBase(Serializable, ABC): """Base class for implementing per-run cost model dimensions This class provides a default implementation of serialization (via - :class:`~llmeter.serialization.Serializable`), a default empty `before_run_start` + [`Serializable`][llmeter.serialization.Serializable]), a default empty `before_run_start` implementation, and abstract methods for the other requirements of the `IRunCostDimension` protocol. """ diff --git a/llmeter/endpoints/base.py b/llmeter/endpoints/base.py index 1914cc3..1c4d80c 100644 --- a/llmeter/endpoints/base.py +++ b/llmeter/endpoints/base.py @@ -101,7 +101,7 @@ def from_json(cls, json_str: str) -> "InvocationResponse": Args: json_str: A JSON string representation of an InvocationResponse (produced by `to_json` - or similar). + or similar). Returns: InvocationResponse: The deserialized response. @@ -129,7 +129,9 @@ def from_json(cls, json_str: str) -> "InvocationResponse": ) return cls(**data) - def to_json(self, default=json_default, **kwargs) -> str: + def to_json( + self, default: Callable[[Any], Any] | None = json_default, **kwargs: Any + ) -> str: """Serialize this response to a JSON string. Uses [`json_default`][llmeter.serialization.json_default] by @@ -569,14 +571,14 @@ def load(cls, endpoint_config: dict) -> "Endpoint": # type: ignore determines the appropriate endpoint class, and instantiates it with the loaded configuration. - .. deprecated:: - This supports the legacy ``{"endpoint_type": ...}`` format. New code should - use :func:`~llmeter.serialization.load_object` with dicts produced by - :func:`~llmeter.serialization.dump_object`. + !!! warning "Deprecated" + This supports the legacy `{"endpoint_type": ...}` format. New code should + use [`load_object`][llmeter.serialization.load_object] with dicts produced by + [`dump_object`][llmeter.serialization.dump_object]. Args: endpoint_config (dict): A dictionary containing the endpoint configuration. - Must include at minimum an ``endpoint_type`` key. + Must include at minimum an `endpoint_type` key. Returns: Endpoint: An instance of the appropriate endpoint class, initialized @@ -593,19 +595,19 @@ def load(cls, endpoint_config: dict) -> "Endpoint": # type: ignore def _filter_legacy_ctor_kwargs(endpoint_class: type, config: dict) -> dict: """Drop legacy config keys that the target endpoint constructor won't accept. - Older LLMeter configs persisted derived/read-only attributes (notably ``provider``, - which endpoints now set internally) alongside the real constructor arguments. Passing - those through to a modern ``__init__`` raises ``TypeError``, so we filter the dict down - to the parameters the constructor actually declares. + Older LLMeter configs persisted derived/read-only attributes (notably `provider`, which + endpoints now set internally) alongside the real constructor arguments. Passing those through + to a modern `__init__` raises `TypeError`, so we filter the dict down to the parameters the + constructor actually declares. - If the constructor accepts ``**kwargs`` the dict is passed through unchanged. + If the constructor accepts `**kwargs` the dict is passed through unchanged. Args: endpoint_class: The endpoint class about to be instantiated. config: The legacy configuration dict (already stripped of ``endpoint_type``). Returns: - A copy of ``config`` containing only keys the constructor accepts. + A copy of `config` containing only keys the constructor accepts. """ params = inspect.signature(endpoint_class.__init__).parameters if any(p.kind == p.VAR_KEYWORD for p in params.values()): diff --git a/llmeter/results.py b/llmeter/results.py index af49e52..561ab67 100644 --- a/llmeter/results.py +++ b/llmeter/results.py @@ -3,7 +3,7 @@ import json import logging -from collections.abc import Sequence +from collections.abc import Callable, Sequence from dataclasses import asdict, dataclass from datetime import datetime from numbers import Number @@ -119,7 +119,9 @@ def save(self, output_path: WritablePathLike | None = None) -> None: for response in self.responses: f.write(json.dumps(asdict(response), default=json_default) + "\n") - def to_json(self, default=json_default, **kwargs) -> str: + def to_json( + self, default: Callable[[Any], Any] | None = json_default, **kwargs: Any + ) -> str: """Return the results as a JSON string. Args: diff --git a/llmeter/serialization.py b/llmeter/serialization.py index 3b59713..6e45469 100644 --- a/llmeter/serialization.py +++ b/llmeter/serialization.py @@ -157,7 +157,7 @@ def _get_type_args(tp) -> tuple: return (tp,) if isinstance(tp, type) else () -def restore_dataclass_types(cls, data: dict) -> None: +def restore_dataclass_types(cls: type, data: dict) -> None: """Restore typed fields in a dict destined for a dataclass constructor. Introspects ``cls`` (a dataclass) and converts JSON-native values back to their From 54b462b59217461d02ffceab32d8864c69d5b8d1 Mon Sep 17 00:00:00 2001 From: Alex Thewsey Date: Tue, 28 Jul 2026 15:02:59 +0800 Subject: [PATCH 17/20] fix(endpoint): Configs survive serialization Fix Anthropic, Bedrock, and OpenAI client configurations to survive Endpoint de/serialization. Previously, many config kwargs to these endpoint types were passed through to the underlying clients but not saved on the endpoint object itself which meant they were dropped during serialization. With this change, we modify the state behaviour for these classes to explicitly pull them through. --- llmeter/endpoints/anthropic_messages.py | 61 ++++++ llmeter/endpoints/bedrock.py | 2 + llmeter/endpoints/bedrock_invoke.py | 2 + llmeter/endpoints/openai.py | 83 ++++++- llmeter/endpoints/openai_response.py | 152 +++++++++++-- .../unit/endpoints/test_anthropic_messages.py | 178 ++++++++++++++- tests/unit/endpoints/test_bedrock.py | 103 +++++++++ tests/unit/endpoints/test_bedrock_invoke.py | 101 +++++++++ tests/unit/endpoints/test_openai.py | 203 ++++++++++++++++++ .../test_openai_response_serialization.py | 161 ++++++++++++++ 10 files changed, 1029 insertions(+), 17 deletions(-) diff --git a/llmeter/endpoints/anthropic_messages.py b/llmeter/endpoints/anthropic_messages.py index 4a4189d..c064536 100644 --- a/llmeter/endpoints/anthropic_messages.py +++ b/llmeter/endpoints/anthropic_messages.py @@ -66,10 +66,14 @@ after all thinking is complete. """ +# Python Built-Ins: +import inspect import logging import time from typing import Any, Generic, Iterable, TypeVar +# External Dependencies: +import httpx # (Indirect dependency of anthropic) import anthropic from anthropic.types import ( Message, @@ -77,6 +81,7 @@ RawMessageStreamEvent, ) +# Local Dependencies: from .base import Endpoint, InvocationResponse logger = logging.getLogger(__name__) @@ -142,8 +147,64 @@ def __init__( kwargs["api_key"] = api_key if aws_region is not None: kwargs["aws_region"] = aws_region + if isinstance(kwargs.get("timeout"), dict): + kwargs["timeout"] = httpx.Timeout(**kwargs["timeout"]) self._client = client_cls(**kwargs) + def _get_llmeter_state(self) -> dict: + """Extract serializable state by introspecting the underlying client. + + Since some constructor params are set on the underlying Anthropic client but *not* + persisted directly on the Endpoint object, we need to extend the default serialization + behaviour to preserve these extra config fields so they survive a serialization round trip. + + Rather than hardcoding per-provider fields (since there are different ANTHROPIC_CLIENTS + classes), we inspect the actual client class's __init__ signature and capture any matching + instance attributes — automatically adapting to whichever provider is in use. Secrets + (params containing 'key', 'token', or 'credential') and non-serializable objects are + excluded. + """ + skip_params = frozenset( + { + "http_client", + "_strict_response_validation", + # default_headers/default_query properties return merged SDK headers including + # secrets; we read _custom_headers/_custom_query directly below instead. + "default_headers", + "default_query", + } + ) + skip_substrings = ("key", "token", "credential") + + state = super()._get_llmeter_state() + sig = inspect.signature(type(self._client).__init__) + for name, param in sig.parameters.items(): + if name == "self" or param.kind in ( + param.VAR_POSITIONAL, + param.VAR_KEYWORD, + ): + continue + if name in skip_params or any(s in name for s in skip_substrings): + continue + if not hasattr(self._client, name): + continue + val = getattr(self._client, name) + if isinstance(val, httpx.URL): + val = str(val) + elif isinstance(val, httpx.Timeout): + val = { + "connect": val.connect, + "read": val.read, + "write": val.write, + "pool": val.pool, + } + state[name] = val + if self._client._custom_headers: + state["default_headers"] = dict(self._client._custom_headers) + if self._client._custom_query: + state["default_query"] = dict(self._client._custom_query) + return state + def _parse_payload(self, payload: MessageCreateParams | dict) -> str: """Extract user message text from an Anthropic Messages API payload. diff --git a/llmeter/endpoints/bedrock.py b/llmeter/endpoints/bedrock.py index 9253c6f..52cf0fc 100644 --- a/llmeter/endpoints/bedrock.py +++ b/llmeter/endpoints/bedrock.py @@ -161,6 +161,8 @@ def __init__( ) self.region = region or boto3.session.Session().region_name + # Persist extra config on the class just so de/serialization catches it: + self._max_attempts = max_attempts logger.info(f"Using AWS region: {self.region}") self._bedrock_client = bedrock_boto3_client diff --git a/llmeter/endpoints/bedrock_invoke.py b/llmeter/endpoints/bedrock_invoke.py index a98312f..04ec0e3 100644 --- a/llmeter/endpoints/bedrock_invoke.py +++ b/llmeter/endpoints/bedrock_invoke.py @@ -58,6 +58,8 @@ def __init__( self.generated_token_count_jmespath = generated_token_count_jmespath self.input_text_jmespath = input_text_jmespath self.input_token_count_jmespath = input_token_count_jmespath + # Persist extra config on the class just so de/serialization catches it: + self._max_attempts = max_attempts self.region = ( region diff --git a/llmeter/endpoints/openai.py b/llmeter/endpoints/openai.py index 79bcf72..13fd05d 100644 --- a/llmeter/endpoints/openai.py +++ b/llmeter/endpoints/openai.py @@ -2,13 +2,18 @@ # SPDX-License-Identifier: Apache-2.0 """LLMeter targets for testing OpenAI ChatCompletions-compatible endpoints (wherever they're hosted)""" +# Python Built-Ins: import base64 import logging import os import time +from collections.abc import Mapping from typing import Any, Literal, Generic, Iterable, TypeVar, cast +# External Dependencies: +import httpx # (Indirect dependency of OpenAI) from openai import OpenAI +from openai._constants import DEFAULT_MAX_RETRIES from openai.types.chat import ( ChatCompletion, ChatCompletionChunk, @@ -23,6 +28,7 @@ File as ChatCompletionFile, ) +# Local Dependencies: from ..prompt_utils import ( ContentItem, MediaContent, @@ -123,6 +129,14 @@ def __init__( endpoint_name: str = "openai", api_key: str | None = None, provider: str = "openai", + organization: str | None = None, + project: str | None = None, + base_url: str | httpx.URL | None = None, + websocket_base_url: str | httpx.URL | None = None, + timeout: float | httpx.Timeout | dict | None = None, + max_retries: int = DEFAULT_MAX_RETRIES, + default_headers: Mapping[str, str] | None = None, + default_query: Mapping[str, object] | None = None, **kwargs: Any, ): """Initialize OpenAI endpoint. @@ -130,12 +144,77 @@ def __init__( Args: model_id: ID of the OpenAI model to use endpoint_name: Name of the endpoint. Defaults to "openai". - api_key: OpenAI API key. Defaults to None. + api_key: OpenAI API key. Defaults to None (reads from OPENAI_API_KEY env var). provider: Provider name. Defaults to "openai". + organization: OpenAI organization ID. Defaults to None. + project: OpenAI project ID. Defaults to None. + base_url: Override the default base URL for the API. + websocket_base_url: Override the default base URL for websocket connections. + timeout: Request timeout in seconds, or an httpx.Timeout for granular control. If a + dict is passed (as is the case after serialization), it'll be interpreted as a set + of arguments to create an httpx.Timeout + max_retries: Maximum number of retries for failed requests. + default_headers: Additional headers to send with every request. + default_query: Additional query parameters to send with every request. **kwargs: Additional arguments passed to OpenAI client """ super().__init__(endpoint_name, model_id, provider=provider) - self._client = OpenAI(api_key=api_key, **kwargs) + client_kwargs: dict[str, Any] = { + "api_key": api_key, + "organization": organization, + "project": project, + "max_retries": max_retries, + **kwargs, + } + if base_url is not None: + client_kwargs["base_url"] = base_url + if websocket_base_url is not None: + client_kwargs["websocket_base_url"] = websocket_base_url + if timeout is not None: + if isinstance(timeout, dict): + timeout = httpx.Timeout(**timeout) + client_kwargs["timeout"] = timeout + if default_headers is not None: + client_kwargs["default_headers"] = default_headers + if default_query is not None: + client_kwargs["default_query"] = default_query + self._client = OpenAI(**client_kwargs) + + @property + def project(self) -> str | None: + """The OpenAI project ID configured on the client.""" + return self._client.project + + def _get_llmeter_state(self) -> dict: + """Extract serializable state from the endpoint and its underlying client. + + This override extends the default state (which pulls immediate properties of the Endpoint) + to include other params set on the OpenAI client but *not* persisted directly on the + Endpoint object. Any config we don't persist will get lost on a serialization round trip, + but anything we pull in should be restored because __init__ passes extra **kwargs to the + OpenAI client constructor. + """ + state = super()._get_llmeter_state() + state["organization"] = self._client.organization + state["base_url"] = str(self._client.base_url) + state["max_retries"] = self._client.max_retries + if self._client.websocket_base_url is not None: + state["websocket_base_url"] = str(self._client.websocket_base_url) + timeout = self._client.timeout + if isinstance(timeout, httpx.Timeout): + state["timeout"] = { + "connect": timeout.connect, + "read": timeout.read, + "write": timeout.write, + "pool": timeout.pool, + } + elif timeout is not None: + state["timeout"] = timeout + if self._client._custom_headers: + state["default_headers"] = dict(self._client._custom_headers) + if self._client._custom_query: + state["default_query"] = dict(self._client._custom_query) + return state def _parse_payload(self, payload: CompletionCreateParams | dict) -> str: """Extract the user message text from a ChatCompletions request payload diff --git a/llmeter/endpoints/openai_response.py b/llmeter/endpoints/openai_response.py index 5b78072..7b1c984 100644 --- a/llmeter/endpoints/openai_response.py +++ b/llmeter/endpoints/openai_response.py @@ -1,18 +1,23 @@ # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # SPDX-License-Identifier: Apache-2.0 +# Python Built-Ins: import logging import time -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from typing import Any, Generic, Iterable, TypeVar, cast +# External Dependencies: +import httpx # (Indirect dependency of OpenAI) from openai import OpenAI +from openai._constants import DEFAULT_MAX_RETRIES from openai.types.responses import Response, ResponseCreateParams, ResponseStreamEvent from openai.types.responses.response_create_params import ( ResponseCreateParamsNonStreaming, ResponseCreateParamsStreaming, ) +# Local Dependencies: from .base import Endpoint, InvocationResponse logger = logging.getLogger(__name__) @@ -31,6 +36,14 @@ def __init__( model_id: str, api_key: str | None = None, provider: str = "openai", + organization: str | None = None, + project: str | None = None, + base_url: str | httpx.URL | None = None, + websocket_base_url: str | httpx.URL | None = None, + timeout: float | httpx.Timeout | dict | None = None, + max_retries: int = DEFAULT_MAX_RETRIES, + default_headers: Mapping[str, str] | None = None, + default_query: Mapping[str, object] | None = None, **kwargs: Any, ): """Initialize Response API endpoint. @@ -40,10 +53,73 @@ def __init__( model_id: ID of the OpenAI model to use api_key: OpenAI API key (optional, uses OPENAI_API_KEY env var if not provided) provider: Provider name (default: "openai") + organization: OpenAI organization ID. Defaults to None. + project: OpenAI project ID. Defaults to None. + base_url: Override the default base URL for the API. + websocket_base_url: Override the default base URL for websocket connections. + timeout: Request timeout in seconds, or an httpx.Timeout for granular control. If a + dict is passed (as is the case after serialization), it'll be interpreted as a set + of arguments to create an httpx.Timeout. + max_retries: Maximum number of retries for failed requests. + default_headers: Additional headers to send with every request. + default_query: Additional query parameters to send with every request. **kwargs: Additional arguments passed to OpenAI client """ super().__init__(endpoint_name, model_id, provider=provider) - self._client = OpenAI(api_key=api_key, **kwargs) + client_kwargs: dict[str, Any] = { + "api_key": api_key, + "organization": organization, + "project": project, + "max_retries": max_retries, + **kwargs, + } + if base_url is not None: + client_kwargs["base_url"] = base_url + if websocket_base_url is not None: + client_kwargs["websocket_base_url"] = websocket_base_url + if timeout is not None: + if isinstance(timeout, dict): + timeout = httpx.Timeout(**timeout) + client_kwargs["timeout"] = timeout + if default_headers is not None: + client_kwargs["default_headers"] = default_headers + if default_query is not None: + client_kwargs["default_query"] = default_query + self._client = OpenAI(**client_kwargs) + + @property + def project(self) -> str | None: + """The OpenAI project ID configured on the client.""" + return self._client.project + + def _get_llmeter_state(self) -> dict: + """Extract serializable state from the endpoint and its underlying client. + + This override extends the default state (which pulls immediate properties of the Endpoint) + to include other params set on the OpenAI client but *not* persisted directly on the + Endpoint object. + """ + state = super()._get_llmeter_state() + state["organization"] = self._client.organization + state["base_url"] = str(self._client.base_url) + state["max_retries"] = self._client.max_retries + if self._client.websocket_base_url is not None: + state["websocket_base_url"] = str(self._client.websocket_base_url) + timeout = self._client.timeout + if isinstance(timeout, httpx.Timeout): + state["timeout"] = { + "connect": timeout.connect, + "read": timeout.read, + "write": timeout.write, + "pool": timeout.pool, + } + elif timeout is not None: + state["timeout"] = timeout + if self._client._custom_headers: + state["default_headers"] = dict(self._client._custom_headers) + if self._client._custom_query: + state["default_query"] = dict(self._client._custom_query) + return state @staticmethod def create_payload( @@ -129,6 +205,14 @@ def __init__( endpoint_name: str = "openai-response", api_key: str | None = None, provider: str = "openai", + organization: str | None = None, + project: str | None = None, + base_url: str | httpx.URL | None = None, + websocket_base_url: str | httpx.URL | None = None, + timeout: float | httpx.Timeout | dict | None = None, + max_retries: int = DEFAULT_MAX_RETRIES, + default_headers: Mapping[str, str] | None = None, + default_query: Mapping[str, object] | None = None, **kwargs: Any, ): """Initialize Response API endpoint. @@ -138,10 +222,30 @@ def __init__( endpoint_name: Name of the endpoint (default: "openai-response") api_key: OpenAI API key (optional, uses OPENAI_API_KEY env var if not provided) provider: Provider name (default: "openai") + organization: OpenAI organization ID. Defaults to None. + project: OpenAI project ID. Defaults to None. + base_url: Override the default base URL for the API. + websocket_base_url: Override the default base URL for websocket connections. + timeout: Request timeout in seconds, or an httpx.Timeout for granular control. + max_retries: Maximum number of retries for failed requests. + default_headers: Additional headers to send with every request. + default_query: Additional query parameters to send with every request. **kwargs: Additional arguments passed to OpenAI client """ super().__init__( - endpoint_name, model_id, api_key=api_key, provider=provider, **kwargs + endpoint_name, + model_id, + api_key=api_key, + provider=provider, + organization=organization, + project=project, + base_url=base_url, + websocket_base_url=websocket_base_url, + timeout=timeout, + max_retries=max_retries, + default_headers=default_headers, + default_query=default_query, + **kwargs, ) @OpenAIEndpointBase.llmeter_invoke @@ -216,6 +320,14 @@ def __init__( api_key: str | None = None, provider: str = "openai", ttft_visible_tokens_only: bool = True, + organization: str | None = None, + project: str | None = None, + base_url: str | httpx.URL | None = None, + websocket_base_url: str | httpx.URL | None = None, + timeout: float | httpx.Timeout | dict | None = None, + max_retries: int = DEFAULT_MAX_RETRIES, + default_headers: Mapping[str, str] | None = None, + default_query: Mapping[str, object] | None = None, **kwargs, ): """Initialize streaming Response API endpoint. @@ -227,13 +339,29 @@ def __init__( provider: Provider name (default: "openai") ttft_visible_tokens_only: When True (default), TTFT measures time to first visible text token. When False, TTFT includes reasoning token events. + organization: OpenAI organization ID. Defaults to None. + project: OpenAI project ID. Defaults to None. + base_url: Override the default base URL for the API. + websocket_base_url: Override the default base URL for websocket connections. + timeout: Request timeout in seconds, or an httpx.Timeout for granular control. + max_retries: Maximum number of retries for failed requests. + default_headers: Additional headers to send with every request. + default_query: Additional query parameters to send with every request. **kwargs: Additional arguments passed to OpenAI client """ super().__init__( - model_id=model_id, - endpoint_name=endpoint_name, + endpoint_name, + model_id, api_key=api_key, provider=provider, + organization=organization, + project=project, + base_url=base_url, + websocket_base_url=websocket_base_url, + timeout=timeout, + max_retries=max_retries, + default_headers=default_headers, + default_query=default_query, **kwargs, ) self.ttft_visible_tokens_only = ttft_visible_tokens_only @@ -270,10 +398,12 @@ def process_raw_response( `response.reasoning_text.delta`): when `ttft_visible_tokens_only` is ``False``, these set TTFT on the first reasoning token. """ - _REASONING_DELTA_TYPES = frozenset(( - "response.reasoning_summary_text.delta", - "response.reasoning_text.delta", - )) + _REASONING_DELTA_TYPES = frozenset( + ( + "response.reasoning_summary_text.delta", + "response.reasoning_text.delta", + ) + ) for event in raw_response: now = time.perf_counter() @@ -315,9 +445,7 @@ def process_raw_response( elif event.type == "response.failed": error_obj = getattr(event.response, "error", None) if error_obj is not None: - error_msg = ( - getattr(error_obj, "message", None) or str(error_obj) - ) + error_msg = getattr(error_obj, "message", None) or str(error_obj) error_code = getattr(error_obj, "code", None) if error_code: error_msg = f"{error_code}: {error_msg}" diff --git a/tests/unit/endpoints/test_anthropic_messages.py b/tests/unit/endpoints/test_anthropic_messages.py index 3273129..00b59a6 100644 --- a/tests/unit/endpoints/test_anthropic_messages.py +++ b/tests/unit/endpoints/test_anthropic_messages.py @@ -1,11 +1,16 @@ # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # SPDX-License-Identifier: Apache-2.0 +# Python Built-Ins: +from pathlib import Path +import tempfile import time from unittest.mock import Mock, patch +# External Dependencies: import pytest +# Local Dependencies: from llmeter.endpoints.anthropic_messages import ( _ANTHROPIC_CLIENTS, AnthropicMessages, @@ -457,9 +462,7 @@ def test_invoke_success(self, mock_client): assert result.error is None def test_invoke_api_error(self, mock_client): - mock_client.return_value.messages.create.side_effect = Exception( - "Stream error" - ) + mock_client.return_value.messages.create.side_effect = Exception("Stream error") endpoint = AnthropicMessagesStream(model_id="test-model") result = endpoint.invoke( @@ -705,3 +708,172 @@ def test_custom_endpoint_name(self, mock_client): model_id="test-model", endpoint_name="my-custom-endpoint" ) assert endpoint.endpoint_name == "my-custom-endpoint" + + +# --------------------------------------------------------------------------- +# Tests: Serialization +# --------------------------------------------------------------------------- + + +class TestAnthropicSerialization: + """Test that client configuration round-trips through serialization.""" + + def test_serializes_basic_fields(self): + """Test serialized state includes provider and model.""" + from llmeter.serialization import dump_object + + endpoint = AnthropicMessages(model_id="claude-opus-4-7", api_key="test") + state = dump_object(endpoint)["__llmeter_state__"] + assert state["model_id"] == "claude-opus-4-7" + assert state["provider"] == "anthropic" + assert state["endpoint_name"] == "anthropic-messages" + assert state["max_retries"] == 2 + assert "base_url" in state + + def test_secrets_excluded(self): + """Test that secrets (keys, tokens) are not persisted.""" + from llmeter.serialization import dump_object + + endpoint = AnthropicMessages(model_id="claude-opus-4-7", api_key="sk-secret") + state = dump_object(endpoint)["__llmeter_state__"] + for key in state: + assert "key" not in key, f"Secret-like key '{key}' found in state" + assert "token" not in key, f"Secret-like key '{key}' found in state" + assert "sk-secret" not in str(state.values()) + + def test_bedrock_mantle_serializes_region(self): + """Test Bedrock Mantle provider captures aws_region.""" + from llmeter.serialization import dump_object + + endpoint = AnthropicMessages( + model_id="anthropic.claude-opus-4-7", + provider="bedrock-mantle", + aws_region="us-west-2", + ) + state = dump_object(endpoint)["__llmeter_state__"] + assert state["aws_region"] == "us-west-2" + assert state["provider"] == "bedrock-mantle" + assert "aws_secret_key" not in state + assert "aws_access_key" not in state + assert "aws_session_token" not in state + + def test_custom_headers_serialized(self): + """Test that user-supplied headers round-trip but SDK headers don't.""" + from llmeter.serialization import dump_object + + endpoint = AnthropicMessages( + model_id="claude-opus-4-7", + api_key="test", + default_headers={"X-Custom": "val"}, + ) + state = dump_object(endpoint)["__llmeter_state__"] + assert state["default_headers"] == {"X-Custom": "val"} + + def test_no_headers_when_none_custom(self): + """Test that default_headers is absent when no custom headers set.""" + from llmeter.serialization import dump_object + + endpoint = AnthropicMessages(model_id="claude-opus-4-7", api_key="test") + state = dump_object(endpoint)["__llmeter_state__"] + assert "default_headers" not in state + + def test_timeout_serialized_as_dict(self): + """Test httpx.Timeout is serialized as a dict.""" + import httpx + + from llmeter.serialization import dump_object + + endpoint = AnthropicMessages( + model_id="claude-opus-4-7", + api_key="test", + timeout=httpx.Timeout(30.0, connect=5.0), + ) + state = dump_object(endpoint)["__llmeter_state__"] + assert state["timeout"] == { + "connect": 5.0, + "read": 30.0, + "write": 30.0, + "pool": 30.0, + } + + def test_round_trip_direct_provider(self): + """Test save/load round-trip for direct Anthropic provider.""" + from llmeter.endpoints.base import Endpoint + + with tempfile.TemporaryDirectory() as tmpdir: + original = AnthropicMessages( + model_id="claude-opus-4-7", + api_key="test", + max_retries=5, + base_url="https://custom.anthropic.com", + default_headers={"X-Foo": "bar"}, + ) + path = Path(tmpdir) / "endpoint.json" + original.save_to_file(path) + + loaded = Endpoint.load_from_file(path) + + assert isinstance(loaded, AnthropicMessages) + assert loaded.model_id == "claude-opus-4-7" + assert loaded.provider == "anthropic" + assert loaded._client.max_retries == 5 + assert "custom.anthropic.com" in str(loaded._client.base_url) + assert loaded._client._custom_headers == {"X-Foo": "bar"} + + def test_round_trip_bedrock_mantle(self): + """Test save/load round-trip for Bedrock Mantle provider.""" + from llmeter.endpoints.base import Endpoint + + with tempfile.TemporaryDirectory() as tmpdir: + original = AnthropicMessages( + model_id="anthropic.claude-opus-4-7", + provider="bedrock-mantle", + aws_region="eu-west-1", + ) + path = Path(tmpdir) / "endpoint.json" + original.save_to_file(path) + + loaded = Endpoint.load_from_file(path) + + assert isinstance(loaded, AnthropicMessages) + assert loaded.provider == "bedrock-mantle" + assert loaded._client.aws_region == "eu-west-1" + + def test_round_trip_stream_endpoint(self): + """Test save/load round-trip for streaming endpoint.""" + from llmeter.endpoints.base import Endpoint + + with tempfile.TemporaryDirectory() as tmpdir: + original = AnthropicMessagesStream( + model_id="claude-opus-4-7", + api_key="test", + ttft_visible_tokens_only=False, + max_retries=4, + ) + path = Path(tmpdir) / "endpoint.json" + original.save_to_file(path) + + loaded = Endpoint.load_from_file(path) + + assert isinstance(loaded, AnthropicMessagesStream) + assert loaded.ttft_visible_tokens_only is False + assert loaded._client.max_retries == 4 + + def test_round_trip_with_granular_timeout(self): + """Test httpx.Timeout round-trips correctly through dict serialization.""" + import httpx + + from llmeter.endpoints.base import Endpoint + + with tempfile.TemporaryDirectory() as tmpdir: + original = AnthropicMessages( + model_id="claude-opus-4-7", + api_key="test", + timeout=httpx.Timeout(10.0, connect=2.0), + ) + path = Path(tmpdir) / "endpoint.json" + original.save_to_file(path) + + loaded = Endpoint.load_from_file(path) + + assert loaded._client.timeout == httpx.Timeout(10.0, connect=2.0) diff --git a/tests/unit/endpoints/test_bedrock.py b/tests/unit/endpoints/test_bedrock.py index 1dd8066..b3949a0 100644 --- a/tests/unit/endpoints/test_bedrock.py +++ b/tests/unit/endpoints/test_bedrock.py @@ -833,3 +833,106 @@ def dropping_stream(): assert response.error is not None assert "connection" in response.error.lower() assert response.input_payload is not None + + +# --------------------------------------------------------------------------- +# Tests: Serialization +# --------------------------------------------------------------------------- + + +class TestBedrockSerialization: + """Test that Bedrock endpoint configuration round-trips through serialization.""" + + def test_serializes_basic_fields(self): + """Test serialized state when only required params are set.""" + from llmeter.serialization import dump_object + + endpoint = BedrockConverse(model_id="test-model") + state = dump_object(endpoint)["__llmeter_state__"] + assert state["model_id"] == "test-model" + assert state["max_attempts"] == 3 + assert state["inference_config"] is None + assert "region" in state + + def test_serializes_all_fields(self): + """Test serialized state includes all optional fields when set.""" + from llmeter.serialization import dump_object + + endpoint = BedrockConverse( + model_id="amazon.nova-pro-v1:0", + region="us-west-2", + max_attempts=7, + inference_config={"maxTokens": 512, "temperature": 0.7}, + ) + state = dump_object(endpoint)["__llmeter_state__"] + assert state["model_id"] == "amazon.nova-pro-v1:0" + assert state["region"] == "us-west-2" + assert state["max_attempts"] == 7 + assert state["inference_config"] == {"maxTokens": 512, "temperature": 0.7} + + def test_round_trip_converse_all_fields(self): + """Test save/load round-trip preserves all fields for BedrockConverse.""" + import tempfile + from pathlib import Path + + from llmeter.endpoints.base import Endpoint + + with tempfile.TemporaryDirectory() as tmpdir: + original = BedrockConverse( + model_id="amazon.nova-pro-v1:0", + region="us-west-2", + max_attempts=5, + inference_config={"maxTokens": 1024}, + ) + path = Path(tmpdir) / "endpoint.json" + original.save_to_file(path) + + loaded = Endpoint.load_from_file(path) + + assert isinstance(loaded, BedrockConverse) + assert loaded.model_id == "amazon.nova-pro-v1:0" + assert loaded.region == "us-west-2" + assert loaded._max_attempts == 5 + assert loaded._inference_config == {"maxTokens": 1024} + + def test_round_trip_converse_stream(self): + """Test save/load round-trip preserves all fields for BedrockConverseStream.""" + import tempfile + from pathlib import Path + + from llmeter.endpoints.base import Endpoint + + with tempfile.TemporaryDirectory() as tmpdir: + original = BedrockConverseStream( + model_id="amazon.nova-lite-v1:0", + region="eu-west-1", + max_attempts=10, + ) + path = Path(tmpdir) / "endpoint.json" + original.save_to_file(path) + + loaded = Endpoint.load_from_file(path) + + assert isinstance(loaded, BedrockConverseStream) + assert loaded.model_id == "amazon.nova-lite-v1:0" + assert loaded.region == "eu-west-1" + assert loaded._max_attempts == 10 + + def test_round_trip_without_optional_fields(self): + """Test save/load works with only required params.""" + import tempfile + from pathlib import Path + + from llmeter.endpoints.base import Endpoint + + with tempfile.TemporaryDirectory() as tmpdir: + original = BedrockConverse(model_id="test-model") + path = Path(tmpdir) / "endpoint.json" + original.save_to_file(path) + + loaded = Endpoint.load_from_file(path) + + assert isinstance(loaded, BedrockConverse) + assert loaded.model_id == "test-model" + assert loaded._max_attempts == 3 + assert loaded._inference_config is None diff --git a/tests/unit/endpoints/test_bedrock_invoke.py b/tests/unit/endpoints/test_bedrock_invoke.py index af101bc..2a76fbb 100644 --- a/tests/unit/endpoints/test_bedrock_invoke.py +++ b/tests/unit/endpoints/test_bedrock_invoke.py @@ -655,3 +655,104 @@ def dropping_body(): assert response.error is not None assert "connection" in response.error.lower() assert response.input_payload is not None + + +# --------------------------------------------------------------------------- +# Tests: Serialization +# --------------------------------------------------------------------------- + + +class TestBedrockInvokeSerialization: + """Test that BedrockInvoke endpoint configuration round-trips through serialization.""" + + def test_serializes_basic_fields(self): + """Test serialized state includes required fields with defaults.""" + from llmeter.serialization import dump_object + + endpoint = BedrockInvoke(model_id="test-model") + state = dump_object(endpoint)["__llmeter_state__"] + assert state["model_id"] == "test-model" + assert state["generated_text_jmespath"] == "choices[0].message.content" + assert state["input_text_jmespath"] == "messages[].content[].text" + assert state["max_attempts"] == 3 + assert state["generated_token_count_jmespath"] == "usage.completion_tokens" + assert state["input_token_count_jmespath"] == "usage.prompt_tokens" + assert "region" in state + + def test_serializes_all_fields(self): + """Test serialized state includes all optional fields when set.""" + from llmeter.serialization import dump_object + + endpoint = BedrockInvoke( + model_id="amazon.titan-text-v1", + generated_text_jmespath="results[0].outputText", + input_text_jmespath="inputText", + generated_token_count_jmespath="results[0].tokenCount", + input_token_count_jmespath="inputTokenCount", + region="us-west-2", + max_attempts=7, + ) + state = dump_object(endpoint)["__llmeter_state__"] + assert state["model_id"] == "amazon.titan-text-v1" + assert state["generated_text_jmespath"] == "results[0].outputText" + assert state["input_text_jmespath"] == "inputText" + assert state["generated_token_count_jmespath"] == "results[0].tokenCount" + assert state["input_token_count_jmespath"] == "inputTokenCount" + assert state["region"] == "us-west-2" + assert state["max_attempts"] == 7 + + def test_round_trip_invoke_all_fields(self): + """Test save/load round-trip preserves all fields for BedrockInvoke.""" + import tempfile + from pathlib import Path + + from llmeter.endpoints.base import Endpoint + + with tempfile.TemporaryDirectory() as tmpdir: + original = BedrockInvoke( + model_id="amazon.titan-text-v1", + generated_text_jmespath="results[0].outputText", + input_text_jmespath="inputText", + generated_token_count_jmespath="results[0].tokenCount", + region="us-west-2", + max_attempts=5, + ) + path = Path(tmpdir) / "endpoint.json" + original.save_to_file(path) + + loaded = Endpoint.load_from_file(path) + + assert isinstance(loaded, BedrockInvoke) + assert loaded.model_id == "amazon.titan-text-v1" + assert loaded.generated_text_jmespath == "results[0].outputText" + assert loaded.input_text_jmespath == "inputText" + assert loaded.generated_token_count_jmespath == "results[0].tokenCount" + assert loaded.region == "us-west-2" + assert loaded._max_attempts == 5 + + def test_round_trip_invoke_stream(self): + """Test save/load round-trip preserves all fields for BedrockInvokeStream.""" + import tempfile + from pathlib import Path + + from llmeter.endpoints.base import Endpoint + + with tempfile.TemporaryDirectory() as tmpdir: + original = BedrockInvokeStream( + model_id="amazon.titan-text-v1", + generated_text_jmespath="outputText", + input_text_jmespath="inputText", + region="eu-west-1", + max_attempts=8, + ) + path = Path(tmpdir) / "endpoint.json" + original.save_to_file(path) + + loaded = Endpoint.load_from_file(path) + + assert isinstance(loaded, BedrockInvokeStream) + assert loaded.model_id == "amazon.titan-text-v1" + assert loaded.generated_text_jmespath == "outputText" + assert loaded.input_text_jmespath == "inputText" + assert loaded.region == "eu-west-1" + assert loaded._max_attempts == 8 diff --git a/tests/unit/endpoints/test_openai.py b/tests/unit/endpoints/test_openai.py index 5934c13..a6fdc8e 100644 --- a/tests/unit/endpoints/test_openai.py +++ b/tests/unit/endpoints/test_openai.py @@ -1,6 +1,8 @@ # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # SPDX-License-Identifier: Apache-2.0 +from pathlib import Path +import tempfile import time from unittest.mock import MagicMock, Mock, patch @@ -35,6 +37,8 @@ def test_initialization(self): assert endpoint.endpoint_name == "test_openai" assert endpoint.provider == "openai" assert endpoint._client is not None + assert endpoint._client.project is None + assert endpoint.project is None def test_initialization_with_custom_provider(self): """Test OpenAI endpoint initialization with custom provider.""" @@ -47,6 +51,18 @@ def test_initialization_with_custom_provider(self): assert endpoint.provider == "custom_openai" assert endpoint.model_id == "gpt-4" + def test_initialization_with_project_id(self): + """Test OpenAI endpoint initialization with custom project ID.""" + endpoint = OpenAICompletionEndpoint( + model_id="openai.glm-5", + api_key="test_key", + project="proj_DUMMY", + ) + + assert endpoint.model_id == "openai.glm-5" + assert endpoint._client.project == "proj_DUMMY" + assert endpoint.project == "proj_DUMMY" + def test_initialization_without_api_key(self): """Test OpenAI endpoint initialization without API key.""" endpoint = OpenAICompletionEndpoint( @@ -758,6 +774,193 @@ def test_stream_response_usage_none_value(self): assert response.num_tokens_output is None +class TestOpenAIEndpointSerialization: + """Test serialization behaves as expected for both endpoint types.""" + + def test_completion_endpoint_serializes_basic_fields(self): + """Test sync endpoint serialized state when optional parameters not set.""" + from llmeter.serialization import dump_object + + endpoint = OpenAICompletionEndpoint(model_id="gpt-4") + state = dump_object(endpoint)["__llmeter_state__"] + assert "api_key" not in state + assert state["endpoint_name"] == "openai" + assert state["model_id"] == "gpt-4" + assert state.get("project") is None + assert state.get("organization") is None + assert state["provider"] == "openai" + assert "base_url" in state + assert state["max_retries"] == 2 + # timeout is always serialized (as dict when it's an httpx.Timeout) + assert "timeout" in state + + def test_completion_endpoint_serializes_all_client_fields(self): + """Test sync endpoint serializes all client configuration fields.""" + from llmeter.serialization import dump_object + + endpoint = OpenAICompletionEndpoint( + model_id="openai.gpt-oss-120b", + api_key="test_key", + endpoint_name="test-endpoint", + organization="org-abc", + project="proj_TEST", + base_url="https://custom.example.com/v1", + websocket_base_url="wss://custom.example.com/ws", + timeout=30.0, + max_retries=5, + default_headers={"X-Custom": "val"}, + default_query={"version": "2"}, + provider="test-provider", + ) + state = dump_object(endpoint)["__llmeter_state__"] + assert "api_key" not in state + assert state["endpoint_name"] == "test-endpoint" + assert state["model_id"] == "openai.gpt-oss-120b" + assert state["organization"] == "org-abc" + assert state["project"] == "proj_TEST" + assert "custom.example.com" in state["base_url"] + assert state["websocket_base_url"] == "wss://custom.example.com/ws" + assert state["timeout"] == 30.0 + assert state["max_retries"] == 5 + assert state["default_headers"] == {"X-Custom": "val"} + assert state["default_query"] == {"version": "2"} + assert state["provider"] == "test-provider" + + def test_timeout_serialized_as_dict_when_granular(self): + """Test httpx.Timeout is serialized as a dict with connect/read/write/pool.""" + import httpx + + from llmeter.serialization import dump_object + + endpoint = OpenAICompletionEndpoint( + model_id="gpt-4", + timeout=httpx.Timeout(10.0, connect=5.0), + ) + state = dump_object(endpoint)["__llmeter_state__"] + assert state["timeout"] == { + "connect": 5.0, + "read": 10.0, + "write": 10.0, + "pool": 10.0, + } + + def test_completion_endpoint_round_trip(self): + """Test save/load round-trip preserves all fields for completion endpoint.""" + from llmeter.endpoints.base import Endpoint + + with tempfile.TemporaryDirectory() as tmpdir: + original = OpenAICompletionEndpoint( + model_id="gpt-4", + api_key="test_key", + endpoint_name="proj-test", + organization="org-rt", + project="proj_roundtrip", + base_url="https://custom.test/v1", + websocket_base_url="wss://custom.test/ws", + timeout=45.0, + max_retries=3, + default_headers={"X-Foo": "bar"}, + default_query={"q": "1"}, + provider="test-provider", + ) + output_path = Path(tmpdir) / "endpoint.json" + original.save(output_path) + + loaded = Endpoint.load_from_file(output_path) + + assert isinstance(loaded, OpenAICompletionEndpoint) + assert loaded.model_id == "gpt-4" + assert loaded.endpoint_name == "proj-test" + assert loaded._client.organization == "org-rt" + assert loaded.project == "proj_roundtrip" + assert "custom.test" in str(loaded._client.base_url) + assert loaded._client.websocket_base_url == "wss://custom.test/ws" + assert loaded._client.timeout == 45.0 + assert loaded._client.max_retries == 3 + assert loaded._client._custom_headers == {"X-Foo": "bar"} + assert loaded._client._custom_query == {"q": "1"} + assert loaded.provider == "test-provider" + + def test_round_trip_with_granular_timeout(self): + """Test that httpx.Timeout round-trips correctly through dict serialization.""" + import httpx + + from llmeter.endpoints.base import Endpoint + + with tempfile.TemporaryDirectory() as tmpdir: + original = OpenAICompletionEndpoint( + model_id="gpt-4", + timeout=httpx.Timeout(10.0, connect=2.0), + ) + output_path = Path(tmpdir) / "endpoint.json" + original.save(output_path) + + loaded = Endpoint.load_from_file(output_path) + + assert isinstance(loaded, OpenAICompletionEndpoint) + assert loaded._client.timeout == httpx.Timeout(10.0, connect=2.0) + + def test_stream_endpoint_serializes_basic_fields(self): + """Test stream endpoint serialized state when optional parameters not set.""" + from llmeter.serialization import dump_object + + endpoint = OpenAICompletionStreamEndpoint(model_id="gpt-4") + state = dump_object(endpoint)["__llmeter_state__"] + assert "api_key" not in state + assert state["endpoint_name"] == "openai" + assert state["model_id"] == "gpt-4" + assert state.get("project") is None + assert state.get("organization") is None + assert state["provider"] == "openai" + + def test_stream_endpoint_round_trip(self): + """Test save/load round-trip preserves all fields for streaming endpoint.""" + from llmeter.endpoints.base import Endpoint + + with tempfile.TemporaryDirectory() as tmpdir: + original = OpenAICompletionStreamEndpoint( + model_id="zai.glm-5", + api_key="test_key", + endpoint_name="proj-test-stream", + organization="org-stream", + project="proj_roundtrip_stream", + base_url="https://stream.test/v1", + max_retries=4, + default_headers={"X-Stream": "yes"}, + provider="test-provider-stream", + ) + output_path = Path(tmpdir) / "endpoint.json" + original.save(output_path) + + loaded = Endpoint.load_from_file(output_path) + + assert isinstance(loaded, OpenAICompletionStreamEndpoint) + assert loaded.model_id == "zai.glm-5" + assert loaded.endpoint_name == "proj-test-stream" + assert loaded._client.organization == "org-stream" + assert loaded.project == "proj_roundtrip_stream" + assert "stream.test" in str(loaded._client.base_url) + assert loaded._client.max_retries == 4 + assert loaded._client._custom_headers == {"X-Stream": "yes"} + assert loaded.provider == "test-provider-stream" + + def test_round_trip_without_optional_fields(self): + """Test save/load round-trip works when only required params are set.""" + from llmeter.endpoints.base import Endpoint + + with tempfile.TemporaryDirectory() as tmpdir: + original = OpenAICompletionEndpoint(model_id="gpt-4", api_key="test_key") + output_path = Path(tmpdir) / "endpoint.json" + original.save(output_path) + + loaded = Endpoint.load_from_file(output_path) + + assert isinstance(loaded, OpenAICompletionEndpoint) + assert loaded.model_id == "gpt-4" + assert loaded.project is None + assert loaded._client.organization is None + + class TestStreamMidStreamErrors: """Verify that errors during stream consumption are caught by the invoke wrapper.""" diff --git a/tests/unit/endpoints/test_openai_response_serialization.py b/tests/unit/endpoints/test_openai_response_serialization.py index f98d82b..20e0f1a 100644 --- a/tests/unit/endpoints/test_openai_response_serialization.py +++ b/tests/unit/endpoints/test_openai_response_serialization.py @@ -224,6 +224,167 @@ def test_stream_endpoint_round_trip(self, mock_openai_class): assert loaded.provider == original.provider +class TestOpenAIResponseEndpointSerialization: + """Test that OpenAI client configuration round-trips through serialization.""" + + def test_response_endpoint_serializes_all_client_fields(self): + """Test all client config fields appear in serialized state.""" + from llmeter.serialization import dump_object + + endpoint = OpenAIResponseEndpoint( + model_id="gpt-4", + endpoint_name="test-response", + organization="org-abc", + project="proj_TEST", + base_url="https://custom.example.com/v1", + websocket_base_url="wss://custom.example.com/ws", + timeout=30.0, + max_retries=5, + default_headers={"X-Custom": "val"}, + default_query={"version": "2"}, + provider="test-provider", + ) + state = dump_object(endpoint)["__llmeter_state__"] + assert "api_key" not in state + assert state["model_id"] == "gpt-4" + assert state["organization"] == "org-abc" + assert state["project"] == "proj_TEST" + assert "custom.example.com" in state["base_url"] + assert state["websocket_base_url"] == "wss://custom.example.com/ws" + assert state["timeout"] == 30.0 + assert state["max_retries"] == 5 + assert state["default_headers"] == {"X-Custom": "val"} + assert state["default_query"] == {"version": "2"} + assert state["provider"] == "test-provider" + + def test_response_endpoint_serializes_defaults(self): + """Test serialized state is correct when only required params are set.""" + from llmeter.serialization import dump_object + + endpoint = OpenAIResponseEndpoint(model_id="gpt-4") + state = dump_object(endpoint)["__llmeter_state__"] + assert "api_key" not in state + assert state["model_id"] == "gpt-4" + assert state.get("project") is None + assert state.get("organization") is None + assert state["provider"] == "openai" + assert state["max_retries"] == 2 + assert "base_url" in state + + def test_timeout_serialized_as_dict_when_granular(self): + """Test httpx.Timeout is serialized as a dict with connect/read/write/pool.""" + import httpx + + from llmeter.serialization import dump_object + + endpoint = OpenAIResponseEndpoint( + model_id="gpt-4", + timeout=httpx.Timeout(10.0, connect=5.0), + ) + state = dump_object(endpoint)["__llmeter_state__"] + assert state["timeout"] == { + "connect": 5.0, + "read": 10.0, + "write": 10.0, + "pool": 10.0, + } + + def test_response_endpoint_round_trip_with_all_fields(self): + """Test save/load round-trip preserves all client config fields.""" + with tempfile.TemporaryDirectory() as tmpdir: + original = OpenAIResponseEndpoint( + model_id="gpt-4", + api_key="test_key", + endpoint_name="rt-test", + organization="org-rt", + project="proj_roundtrip", + base_url="https://custom.test/v1", + websocket_base_url="wss://custom.test/ws", + timeout=45.0, + max_retries=3, + default_headers={"X-Foo": "bar"}, + default_query={"q": "1"}, + provider="test-provider", + ) + output_path = Path(tmpdir) / "endpoint.json" + original.save(output_path) + + loaded = Endpoint.load_from_file(output_path) + + assert isinstance(loaded, OpenAIResponseEndpoint) + assert loaded.model_id == "gpt-4" + assert loaded._client.organization == "org-rt" + assert loaded.project == "proj_roundtrip" + assert "custom.test" in str(loaded._client.base_url) + assert loaded._client.websocket_base_url == "wss://custom.test/ws" + assert loaded._client.timeout == 45.0 + assert loaded._client.max_retries == 3 + assert loaded._client._custom_headers == {"X-Foo": "bar"} + assert loaded._client._custom_query == {"q": "1"} + assert loaded.provider == "test-provider" + + def test_stream_endpoint_round_trip_with_all_fields(self): + """Test save/load round-trip preserves all client config for streaming endpoint.""" + with tempfile.TemporaryDirectory() as tmpdir: + original = OpenAIResponseStreamEndpoint( + model_id="gpt-4", + api_key="test_key", + endpoint_name="stream-rt-test", + organization="org-stream", + project="proj_stream_rt", + base_url="https://stream.test/v1", + max_retries=4, + default_headers={"X-Stream": "yes"}, + provider="test-stream", + ttft_visible_tokens_only=False, + ) + output_path = Path(tmpdir) / "endpoint.json" + original.save(output_path) + + loaded = Endpoint.load_from_file(output_path) + + assert isinstance(loaded, OpenAIResponseStreamEndpoint) + assert loaded.model_id == "gpt-4" + assert loaded._client.organization == "org-stream" + assert loaded.project == "proj_stream_rt" + assert "stream.test" in str(loaded._client.base_url) + assert loaded._client.max_retries == 4 + assert loaded._client._custom_headers == {"X-Stream": "yes"} + assert loaded.provider == "test-stream" + assert loaded.ttft_visible_tokens_only is False + + def test_round_trip_with_granular_timeout(self): + """Test httpx.Timeout round-trips correctly through dict serialization.""" + import httpx + + with tempfile.TemporaryDirectory() as tmpdir: + original = OpenAIResponseEndpoint( + model_id="gpt-4", + timeout=httpx.Timeout(10.0, connect=2.0), + ) + output_path = Path(tmpdir) / "endpoint.json" + original.save(output_path) + + loaded = Endpoint.load_from_file(output_path) + + assert isinstance(loaded, OpenAIResponseEndpoint) + assert loaded._client.timeout == httpx.Timeout(10.0, connect=2.0) + + def test_round_trip_without_optional_fields(self): + """Test save/load works when only required params are set.""" + with tempfile.TemporaryDirectory() as tmpdir: + original = OpenAIResponseEndpoint(model_id="gpt-4", api_key="test_key") + output_path = Path(tmpdir) / "endpoint.json" + original.save(output_path) + + loaded = Endpoint.load_from_file(output_path) + + assert isinstance(loaded, OpenAIResponseEndpoint) + assert loaded.model_id == "gpt-4" + assert loaded.project is None + assert loaded._client.organization is None + + class TestSerializationPropertyTests: """Property-based tests for endpoint serialization. From c14ec83d2f4fc6571837d75100173dc618e92f60 Mon Sep 17 00:00:00 2001 From: Alex Thewsey Date: Thu, 30 Jul 2026 16:54:14 +0800 Subject: [PATCH 18/20] doc: Zensical & comment fixes Replace :func:, :class: references incompatible with our docs website generation. Add some missing types in files this PR has touched. Update commentary. Add warnings when Result loads Endpoint or Tokenizer from an old-style format since we'll likely drop support in future. --- docs/reference/serialization.md | 3 + llmeter/callbacks/base.py | 6 +- llmeter/callbacks/cost/model.py | 11 +- llmeter/endpoints/anthropic_messages.py | 2 +- llmeter/endpoints/bedrock_invoke.py | 2 +- llmeter/endpoints/openai_response.py | 4 +- llmeter/endpoints/sagemaker.py | 10 +- llmeter/results.py | 22 ++- llmeter/runner.py | 27 ++-- llmeter/serialization.py | 205 ++++++++++++------------ mkdocs.yml | 2 +- 11 files changed, 150 insertions(+), 144 deletions(-) diff --git a/docs/reference/serialization.md b/docs/reference/serialization.md index 9aa47cf..c14cb1e 100644 --- a/docs/reference/serialization.md +++ b/docs/reference/serialization.md @@ -1 +1,4 @@ ::: llmeter.serialization + options: + filters: + - ".*" # Allow private methods for `Serializable._{get|set}_llmeter_state` diff --git a/llmeter/callbacks/base.py b/llmeter/callbacks/base.py index fc36e83..6ab8463 100644 --- a/llmeter/callbacks/base.py +++ b/llmeter/callbacks/base.py @@ -45,10 +45,8 @@ async def after_invoke(self, response: InvocationResponse) -> None: timing and token counts) Returns: None: If you'd like to add information to the `response` logged in the Run, modify it - in-place. To attach **extra custom fields**, store them in `response.annotations` - (a dict) rather than setting arbitrary attributes on the response: `annotations` - is a declared field, so it is preserved through `Result` save/load, whereas loose - attributes are dropped on serialization. + in-place. To attach **extra custom fields** that you want preserved when responses + are saved to file, store them in the `response.annotations` map. """ pass diff --git a/llmeter/callbacks/cost/model.py b/llmeter/callbacks/cost/model.py index 1255673..a2a40f8 100644 --- a/llmeter/callbacks/cost/model.py +++ b/llmeter/callbacks/cost/model.py @@ -101,9 +101,6 @@ async def calculate_request_cost( } ) if save: - # Store per-response costs in `response.annotations` (a declared field) so they - # persist through to_json/from_json and disk save/load, rather than as loose - # attributes on the response (which asdict-based serialization would drop). dim_costs.save_on_namespace(response.annotations, key_prefix="cost_") return dim_costs @@ -137,7 +134,7 @@ async def calculate_run_cost( else: resp_costs = list( filter( - lambda c: c, + lambda c: c, # Skip responses where no cost data was found at all ( CalculatedCostWithDimensions.load_from_namespace( r.annotations, key_prefix="cost_" @@ -147,13 +144,19 @@ async def calculate_run_cost( ) ) if len(resp_costs): + # Merge the total request-level costs into the run-level costs: + # (Unless requests is empty, because sum([])=0 and not a CalculatedCostWithDimensions) run_cost.merge(sum(resp_costs)) # type: ignore if save: + # Save the overall run cost and breakdown on the main result object: run_cost.save_on_namespace(result, key_prefix="cost_") + # Contribute both 1/ the summary stats of request-level costs, and 2/ the overall run + # cost+breakdown, to result.stats: stats = CalculatedCostWithDimensions.summary_statistics( resp_costs, key_prefix="cost_", key_dim_name_suffix="_per_request", + # cost_total_per_request would be confusing, so skip 'total': key_total_name_and_suffix="per_request", ) run_cost.save_on_namespace(stats, key_prefix="cost_") diff --git a/llmeter/endpoints/anthropic_messages.py b/llmeter/endpoints/anthropic_messages.py index c064536..f314dfb 100644 --- a/llmeter/endpoints/anthropic_messages.py +++ b/llmeter/endpoints/anthropic_messages.py @@ -73,13 +73,13 @@ from typing import Any, Generic, Iterable, TypeVar # External Dependencies: -import httpx # (Indirect dependency of anthropic) import anthropic from anthropic.types import ( Message, MessageCreateParams, RawMessageStreamEvent, ) +import httpx # (Indirect dependency of anthropic) # Local Dependencies: from .base import Endpoint, InvocationResponse diff --git a/llmeter/endpoints/bedrock_invoke.py b/llmeter/endpoints/bedrock_invoke.py index 04ec0e3..af3ea86 100644 --- a/llmeter/endpoints/bedrock_invoke.py +++ b/llmeter/endpoints/bedrock_invoke.py @@ -350,7 +350,7 @@ def process_raw_response( """Parse the streaming response from Bedrock InvokeModelWithResponseStream API. Args: - client_response: The raw response from the Bedrock API. + raw_response: The raw response from the Bedrock API. start_t: The timestamp when the request was initiated. Returns: diff --git a/llmeter/endpoints/openai_response.py b/llmeter/endpoints/openai_response.py index 7b1c984..a9cb03f 100644 --- a/llmeter/endpoints/openai_response.py +++ b/llmeter/endpoints/openai_response.py @@ -126,7 +126,7 @@ def create_payload( user_message: str | Sequence[str], max_output_tokens: int = 256, instructions: str | None = None, - **kwargs, + **kwargs: Any, ) -> ResponseCreateParams: """Create a payload for the Responses API request. @@ -328,7 +328,7 @@ def __init__( max_retries: int = DEFAULT_MAX_RETRIES, default_headers: Mapping[str, str] | None = None, default_query: Mapping[str, object] | None = None, - **kwargs, + **kwargs: Any, ): """Initialize streaming Response API endpoint. diff --git a/llmeter/endpoints/sagemaker.py b/llmeter/endpoints/sagemaker.py index 5a82935..458c90a 100644 --- a/llmeter/endpoints/sagemaker.py +++ b/llmeter/endpoints/sagemaker.py @@ -5,7 +5,7 @@ import json import logging import time -from typing import Generic, TypeVar +from typing import Any, Generic, TypeVar import warnings import boto3 @@ -141,8 +141,8 @@ def create_payload( input_text: str | list[ContentItem], max_tokens: int = 256, inference_parameters: dict = {}, - **kwargs, - ): + **kwargs: Any, + ) -> dict: """Create a payload for the SageMaker API request. Args: @@ -291,8 +291,8 @@ def create_payload( input_text: str | list[ContentItem], max_tokens: int = 256, inference_parameters: dict = {}, - **kwargs, - ): + **kwargs: Any, + ) -> dict: """Create a payload for the SageMaker streaming API request. Args: diff --git a/llmeter/results.py b/llmeter/results.py index 561ab67..12fa3ab 100644 --- a/llmeter/results.py +++ b/llmeter/results.py @@ -126,8 +126,8 @@ def to_json( Args: default: Fallback serializer. Defaults to - :func:`~llmeter.serialization.json_default`. - **kwargs: Extra keyword arguments passed to :func:`json.dumps`. + [`json_default`][llmeter.serialization.json_default]. + **kwargs: Extra keyword arguments passed to `json.dumps`. """ summary = { k: o for k, o in asdict(self).items() if k not in ["responses", "stats"] @@ -137,20 +137,18 @@ def to_json( def to_dict(self, include_responses: bool = False) -> dict: """Return a dictionary representation of this result. - Returns a plain ``dict`` produced by :func:`dataclasses.asdict`, - preserving native Python types (``datetime``, ``UPath``, etc.). - This is suitable for programmatic access and internal data + Returns a plain `dict` produced by `dataclasses.asdict`, preserving native Python types + (`datetime`, `UPath`, etc.). This is suitable for programmatic access and internal data processing. - For JSON output, use :meth:`to_json` which delegates to - :func:`~llmeter.serialization.json_default` for - non-serializable types, or pass the dict through - ``json.dumps(result.to_dict(), default=json_default)``. + For JSON output, use `to_json` which delegates to + [`json_default`][llmeter.serialization.json_default] for non-serializable types, or pass + the dict through `json.dumps(result.to_dict(), default=json_default)`. Args: - include_responses: If ``True``, include the full list of - :class:`~llmeter.endpoints.base.InvocationResponse` dicts - and the ``stats`` key. Defaults to ``False``. + include_responses: Set `True` to include the full list of + [`InvocationResponse`][llmeter.endpoints.base.InvocationResponse] dicts and the + `stats` key. Returns: dict: A dictionary of result fields with native Python types. diff --git a/llmeter/runner.py b/llmeter/runner.py index 508579f..743d797 100644 --- a/llmeter/runner.py +++ b/llmeter/runner.py @@ -148,6 +148,11 @@ def __post_init__(self, disable_client_progress_bar, disable_clients_progress_ba if "__llmeter_class__" in self.endpoint: self._endpoint: Endpoint = load_object(self.endpoint) else: + # Legacy load path for <=v0.1.12 data files: + logger.warning( + "Loading endpoint config from a legacy data format. Support for this will be " + "removed in a future major version. Consider re-saving the file to update." + ) self._endpoint = Endpoint.load(self.endpoint) else: self._endpoint = self.endpoint @@ -161,6 +166,11 @@ def __post_init__(self, disable_client_progress_bar, disable_clients_progress_ba if "__llmeter_class__" in self.tokenizer: self._tokenizer: Tokenizer = load_object(self.tokenizer) else: + # Legacy load path for <=v0.1.12 data files: + logger.warning( + "Loading tokenizer from a legacy LLMeter data format. Support for this will " + "be removed in a future major version. Consider re-saving the file to update." + ) self._tokenizer = Tokenizer.load(self.tokenizer) else: self._tokenizer = self.tokenizer @@ -172,15 +182,14 @@ def save( ): """Save the configuration to a disk or cloud storage. - Uses the LLMeter state-based serialization protocol - (``_get_llmeter_state``/``_set_llmeter_state``) for Endpoint, Tokenizer, - and Callback serialization. - Args: output_path: Optional override for output folder. By default, self.output_path is used. file_name: File name to create under `output_path`. """ output_path = ensure_path(output_path or self.output_path) + if output_path is None: + logger.info("No output_path provided - skipping saving run config") + return output_path.mkdir(parents=True, exist_ok=True) run_config_path = output_path / file_name @@ -206,9 +215,6 @@ def save( def load(cls, load_path: ReadablePathLike, file_name: str = "run_config.json"): """Load a configuration from a (local or cloud-stored) JSON file. - Restores Endpoint, Tokenizer, and Callback objects via ``load_object``. - Also supports legacy format (``endpoint_type`` / ``tokenizer_module`` keys). - Args: load_path: Folder under which the configuration is stored file_name: File name within `load_path` for the run configuration JSON. @@ -217,10 +223,9 @@ def load(cls, load_path: ReadablePathLike, file_name: str = "run_config.json"): with (load_path / file_name).open() as f: config = json.load(f) - # Restore endpoint and tokenizer. Only unwrap the new-format - # (``__llmeter_class__``) envelope here; legacy dicts (``endpoint_type`` / - # ``tokenizer_module``) are passed through untouched so ``__post_init__`` can - # route them to ``Endpoint.load`` / ``Tokenizer.load``. + # Restore endpoint and tokenizer. Only unwrap the new-format (`__llmeter_class__`) envelope + # here; legacy dicts (`endpoint_type` / `tokenizer_module`) are passed through untouched so + # `__post_init__` can route them to `Endpoint.load` / `Tokenizer.load`. ep = config.get("endpoint") if isinstance(ep, dict) and "__llmeter_class__" in ep: config["endpoint"] = load_object(ep) diff --git a/llmeter/serialization.py b/llmeter/serialization.py index 6e45469..a6b9b85 100644 --- a/llmeter/serialization.py +++ b/llmeter/serialization.py @@ -1,51 +1,41 @@ # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # SPDX-License-Identifier: Apache-2.0 -"""Unified serialization for all LLMeter objects. +"""Unified JSON-based serialization for LLMeter objects. -This module is the single source of truth for serialization in LLMeter. +Most of the objects we'd like to save to and load from file in LLMeter are configuration-like and +*almost* JSON-compatible, but with some extensions (like datetimes, binary image/etc payloads, +and callback objects). Rather than falling back to Pickle (which is Python-specific, +non-human-readable, and Python version fragile) - we implement in this module a JSON-based scheme +for saving and loading compatible LLMeter objects. -Key components: - -- :class:`Serializable` — Mixin that gives any class an automatic *state* protocol - (``_get_llmeter_state``/``_set_llmeter_state``) by introspecting ``__init__``. - Deliberately distinct from the pickle protocol so ``pickle`` / ``copy`` / - ``deepcopy`` keep their native behavior. - -- :func:`dump_object` / :func:`load_object` — Full round-trip persistence - using a ``{"__llmeter_class__": ..., "__llmeter_state__": ...}`` envelope. - -- :func:`json_default` — A ``json.dumps``-compatible *default* handler for - ``bytes``, ``datetime``, ``PathLike``. +!!! warning "A warning on security" -- :func:`bytes_decoder` — A ``json.loads``-compatible *object_hook* that - restores binary content encoded by :func:`json_default`. + [`load_object`][llmeter.serialization.load_object] imports and instantiates whatever class path + is in the `__llmeter_class__` field. Like unpickle, this has the potential to run arbitrary + code. Do not load configs from untrusted sources! -- :func:`datetime_to_str` / :func:`str_to_datetime` — Standardized - datetime ↔ string conversion (UTC ISO-8601 with ``Z`` suffix). +### Key components -!!! warning "A warning on security" +- **[`Serializable`][llmeter.serialization.Serializable]**: Mixin to give any class an automatic + *state* protocol (`_get_llmeter_state`/`_set_llmeter_state`) by introspecting `__init__`. + Deliberately distinct from the pickle protocol so `pickle` / `copy` / `deepcopy` keep their + native behavior. +- **[`dump_object`][llmeter.serialization.dump_object]** and + **[`load_object`][llmeter.serialization.load_object]**: Full round-trip persistence for + `Serializable`-compatible objects, using a + `{"__llmeter_class__": ..., "__llmeter_state__": ...}` envelope. - :func:`load_object` imports and instantiates whatever class path is in the - ``__llmeter_class__`` field. Do not load configs from untrusted sources. ### Implementation details to be aware of -**State vs. identity:** We split an object's serialized form into two layers. The -*state* (``_get_llmeter_state``) contains the object's own field values and does *not* -describe which class those values belong to. The **envelope** built by -:func:`dump_object` adds the identity (``__llmeter_class__``) around that state. -This is why the state methods are named ``_get_llmeter_state`` / -``_set_llmeter_state``: the dict they handle is not a self-describing representation -of the object, only of its state, so a for example a public ``to_dict`` name would -over-promise. Keeping identity out of the state dict also avoids polluting the user's -field namespace and lets nested polymorphic values each carry their own envelope. - -**Why a base-class mixin and not a Protocol:** The value :class:`Serializable` -provides is an *inherited implementation* (via ``__init__`` introspection), not -merely a method contract — a class becomes serializable by inheriting it, writing -zero serialization code. A structural ``Protocol`` could only describe the methods, -not hand over that implementation, so it would not actually help an author satisfy -the contract; they would inherit :class:`Serializable` regardless. +**State vs. identity:** We split an object's serialized form into two layers. The *state* +(`_get_llmeter_state`) contains the object's own field values and does *not* describe which class +those values belong to. The **envelope** built by +[`dump_object`][llmeter.serialization.dump_object] adds the identity (`__llmeter_class__`) around +that state. The "state" dictionaries handled by `_get_llmeter_state` / `_set_llmeter_state` are +not a fully self-describing representation of the object - only of its state. Keeping identity out +of the state dict avoids polluting the user's field namespace and lets nested polymorphic values +each carry their own envelope. """ import base64 @@ -76,8 +66,8 @@ def datetime_to_str(dt: datetime) -> str: """Convert a datetime to a UTC ISO-8601 string with ``Z`` suffix. - Timezone-aware datetimes are converted to UTC first. Naive datetimes are - serialized as-is (assumed UTC). + Timezone-aware datetimes are converted to UTC first. Naive datetimes are serialized as-is + (assumed UTC). """ if dt.tzinfo is not None: dt = dt.astimezone(timezone.utc) @@ -95,22 +85,23 @@ def str_to_datetime(s: str) -> datetime: def json_default(obj: Any) -> Any: - """Serialize a single non-JSON-serializable object. + """Serialize a single non-natively-JSON-serializable object. - Intended for use as the ``default`` argument to :func:`json.dumps` or - :func:`json.dump`. + Intended for use as the `default` argument to `json.dump` or `json.dumps`. This **does not** + handle full recursive LLMeter serialization - see + [`dump_object`][llmeter.serialization.dump_object] instead. Type handling (checked in order): - * ``bytes`` — wrapped in a ``{"__llmeter_bytes__": ""}`` marker so - that :func:`bytes_decoder` can restore them on the way back. - * ``datetime`` — converted to a UTC ISO-8601 string with a ``Z`` suffix. - * ``date`` / ``time`` — converted via ``.isoformat()``. - * ``os.PathLike`` — converted to a POSIX path string. - * Anything else — ``str()`` fallback (returns ``None`` if that also fails). + * `bytes` — wrapped in a `{"__llmeter_bytes__": ""}` marker so that + [`bytes_decoder`][llmeter.serialization.bytes_decoder] can restore them on the way back. + * `datetime` — converted to a UTC ISO-8601 string with a `Z` suffix. + * `date` / `time` — converted via `.isoformat()`. + * `os.PathLike` — converted to a POSIX path string. + * Anything else — `str()` fallback (returns `None` if that also fails). Args: - obj: The object that the default encoder could not handle. + obj: The object that the default JSON encoder could not handle. Returns: A JSON-serializable representation of *obj*. @@ -130,17 +121,17 @@ def json_default(obj: Any) -> Any: def bytes_decoder(dct: dict) -> dict | bytes: - """Decode ``__llmeter_bytes__`` marker objects back to ``bytes``. + """Decode `__llmeter_bytes__` marker objects back to Python `bytes`. - Intended for use as the ``object_hook`` argument to :func:`json.loads` or - :func:`json.load`. Marker objects produced by :func:`json_default` are detected - and converted back to ``bytes``; all other dicts pass through unchanged. + Intended for use as the `object_hook` argument to `json.load` or `json.loads`. Marker objects + produced by [`json_default`][llmeter.serialization.json_default] are detected and converted + back to `bytes`; all other dicts pass through unchanged. Args: dct: A dictionary produced by the JSON parser. Returns: - The original ``bytes`` if *dct* is a marker object, otherwise *dct* unchanged. + The original `bytes` if *dct* is an LLMeter bytes marker object, otherwise *dct* unchanged. """ if "__llmeter_bytes__" in dct and len(dct) == 1: return base64.b64decode(dct["__llmeter_bytes__"]) @@ -148,7 +139,7 @@ def bytes_decoder(dct: dict) -> dict | bytes: def _get_type_args(tp) -> tuple: - """Return the members of a union type (e.g. ``datetime | None`` -> (datetime, NoneType)).""" + """Return the members of a union type (e.g. `datetime | None` -> (datetime, NoneType)).""" if isinstance(tp, _types.UnionType): return tp.__args__ origin = getattr(tp, "__origin__", None) @@ -160,18 +151,18 @@ def _get_type_args(tp) -> tuple: def restore_dataclass_types(cls: type, data: dict) -> None: """Restore typed fields in a dict destined for a dataclass constructor. - Introspects ``cls`` (a dataclass) and converts JSON-native values back to their - annotated Python types. Currently handles: + Introspects `cls` (a dataclass) and converts JSON-native values back to their annotated Python + types. Currently handles: - * ``datetime`` fields — parses ISO-8601 strings via :func:`str_to_datetime`. - * ``bytes`` fields — decodes ``__llmeter_bytes__`` markers via base64. + * `datetime` fields — parses ISO-8601 strings via `str_to_datetime`. + * `bytes` fields — decodes `__llmeter_bytes__` markers via base64. - Only fields declared on ``cls`` are touched — nested user payloads (e.g. - ``input_payload``) are left unchanged. Mutates *data* in place. + Only fields declared on `cls` are touched — nested user payloads (e.g. `input_payload`) are + left unchanged. Mutates *data* in place. Args: cls: A dataclass type to introspect for field type annotations. - data: A dictionary of field values (e.g. from :func:`json.load`) to coerce. + data: A dictionary of field values (e.g. from `json.load`) to coerce. """ for f in fields(cls): val = data.get(f.name) @@ -197,25 +188,28 @@ class Serializable: """Mixin providing a state extraction protocol compatible with LLMeter serialization. Serialization in LLMeter uses a state extraction protocol somewhat similar to, but deliberately - separate from, the (``__getstate__`` / ``__setstate__``) interface used by `pickle`, - `copy.copy`, and `copy.deepcopy`. Subclasses remain natively picklable/copyable, but can also + separate from, the (`__getstate__` / `__setstate__`) interface used by `pickle`, `copy.copy`, + and `copy.deepcopy`. Subclasses remain natively picklable/copyable, but can also be saved to and loaded from LLMeter's JSON-based format. - Works with plain classes, ``@dataclass``, and any class whose ``__init__`` - parameters correspond to instance attributes (``self.x`` or ``self._x``). + This default implementation works with plain classes, `@dataclass`, and any class whose + `__init__` parameters correspond to instance attributes (`self.x` or `self._x`). Nested + `Serializable`-like objects are recursively persisted and loaded via + [`dump_object`][llmeter.serialization.dump_object] and + [`load_object`][llmeter.serialization.load_object]. - Nested :class:`Serializable` objects are recursively persisted via - :func:`dump_object` / :func:`load_object`. + If your class needs more custom logic to represent and restore its state, customize the + provided methods by overriding or implementing your own from scratch. """ def _get_llmeter_state(self) -> dict: - """Extract a JSON-serializable state dict by introspecting ``__init__`` parameters. + """Extract a JSON-serializable state dict by introspecting `__init__` parameters. - Returns the object's constructor arguments (looked up as ``self.`` or - ``self._``), recursively serialized. Note that: + Returns the object's constructor arguments (looked up as `self.` or `self._`), + recursively serialized. Note that: - 1. This is state only — it carries no class identity; :func:`dump_object` wraps it with - ``__llmeter_class__``. + 1. This is state only — it carries no class identity; `dump_object` wraps it with + `__llmeter_class__`. 2. Any properties not exposed as `__init__` arguments will not be persisted. Override your class' state get and set methods if you need different behaviour. """ @@ -233,12 +227,11 @@ def _get_llmeter_state(self) -> dict: return state def _set_llmeter_state(self, state: dict) -> None: - """Restore this instance from a state dict produced by :meth:`_get_llmeter_state`. + """Restore this instance from a state dict produced by `_get_llmeter_state`. - Note this rebuilds the object by *calling the constructor* with the state as - keyword arguments (it is a ``from_dict``-style reconstruction, not pickle-style - state restoration onto an already-built instance). ``load_object`` first - creates a bare instance via ``__new__``, then calls this to populate it. + Note this rebuilds the object by *calling the constructor* with the state as keyword + arguments. [`load_object`][llmeter.serialization.load_object] first creates a bare instance + with `__new__`, then calls this method to populate it. """ deserialized = {k: _deserialize_value(v) for k, v in state.items()} self.__init__(**deserialized) @@ -246,8 +239,8 @@ def _set_llmeter_state(self, state: dict) -> None: def save_to_file(self, path: WritablePathLike) -> Path: """Save this object to a JSON file. - Uses the :meth:`_get_llmeter_state` protocol. Override - :meth:`_get_llmeter_state` (not this method) if custom serialization is needed. + Uses the `_get_llmeter_state` protocol. Override `_get_llmeter_state` (not this method) if + custom JSON-based serialization is needed. Args: path: (Local or Cloud) path where the object will be saved. @@ -266,7 +259,7 @@ def save_to_file(self, path: WritablePathLike) -> Path: def load_from_file(cls, path: ReadablePathLike) -> "Serializable": """Load an object from a JSON file. - Detects the type from the ``__llmeter_class__`` field and reconstructs it. + Detects the type from the `__llmeter_class__` field and reconstructs it. Args: path: (Local or Cloud) path where the object was saved. @@ -288,20 +281,22 @@ def dump_object(obj: Any) -> dict: """Serialize an object to a type-tagged dict for round-trip persistence. The returned envelope has the form - ``{"__llmeter_class__": "module.Class", "__llmeter_state__": {...}}``. + `{"__llmeter_class__": "module.Class", "__llmeter_state__": {...}}`. Serialization strategy (checked in order): - 1. If the object implements the LLMeter state protocol - (:meth:`Serializable._get_llmeter_state`), calls it to obtain the state dict. - 2. If the object is a dataclass, uses :func:`dataclasses.asdict`. - 3. Otherwise, takes all public (non-underscore-prefixed) entries from ``__dict__``. + 1. If the object implements the LLMeter state protocol (as in + [`Serializable._get_llmeter_state`][llmeter.serialization.Serializable._get_llmeter_state]) + this will call it to obtain the state dict. + 2. If the object is a dataclass, uses `dataclasses.asdict`. + 3. Otherwise, takes all public (non-underscore-prefixed) entries from `__dict__`. Args: obj: The object to serialize. Returns: - A JSON-serializable dict that :func:`load_object` can reconstruct. + dct: A JSON-serializable dict that [`load_object`][llmeter.serialization.load_object] can + reconstruct. """ class_path = f"{obj.__class__.__module__}.{obj.__class__.__qualname__}" if hasattr(obj, "_get_llmeter_state"): @@ -316,23 +311,25 @@ def dump_object(obj: Any) -> dict: def load_object(data: dict) -> Any: - """Restore an object from a type-tagged dict produced by :func:`dump_object`. + """Restore an object from a type-tagged dict produced by `dump_object`. + + !!! warning + This method imports and instantiates class paths specified by the input data, which (like + pickle) can enable arbitrary running arbitrary code. Do not run it on data from unstrusted + sources! - Imports the module identified by ``__llmeter_class__``, instantiates the class - (bypassing ``__init__`` via ``__new__``), and restores its state via - :meth:`Serializable._set_llmeter_state`. Objects that do not implement the - protocol are restored by assigning the (deserialized) state onto ``__dict__``. + Imports the module identified by `__llmeter_class__`, instantiates the class (bypassing + `__init__` via `__new__`), and restores its state via + [`Serializable._set_llmeter_state`][llmeter.serialization.Serializable._set_llmeter_state]. + Objects that do not implement the protocol are restored by assigning the (deserialized) state + onto their `__dict__`. Args: - data: A dict with ``__llmeter_class__`` and ``__llmeter_state__`` keys, as - produced by :func:`dump_object`. + data: A dict with `__llmeter_class__` and `__llmeter_state__` keys, as + produced by [`dump_object`][llmeter.serialization.dump_object]. Returns: The reconstructed object instance. - - .. warning:: - Do not call on data from untrusted sources — it imports and instantiates - arbitrary classes. """ class_path = data["__llmeter_class__"] module_path, class_name = class_path.rsplit(".", 1) @@ -372,8 +369,11 @@ def _serialize_value(val: Any) -> Any: """Recursively prepare a value for JSON persistence. Handles primitives, known types (bytes, datetime, PathLike), nested - :class:`Serializable` objects (via :func:`dump_object`), dicts, and lists/tuples. - Raises :exc:`TypeError` for objects it cannot serialize. + [`Serializable`][llmeter.serialization.Serializable] objects (via + [`dump_object`][llmeter.serialization.dump_object]), dicts, and lists/tuples. + + Raises: + TypeError: for objects it cannot serialize. """ if val is None or isinstance(val, (str, int, float, bool)): return val @@ -398,9 +398,8 @@ def _serialize_value(val: Any) -> Any: def _deserialize_value(val: Any) -> Any: """Recursively restore a value from JSON persistence. - Recognizes type-tagged dicts (``__llmeter_class__``), bytes markers - (``__llmeter_bytes__``), ISO-8601 datetime strings, and recursively processes - nested dicts and lists. + Recognizes type-tagged dicts (`__llmeter_class__`), bytes markers (`__llmeter_bytes__`), + ISO-8601 datetime strings, and recursively processes nested dicts and lists. """ match val: case None | bool() | int() | float(): diff --git a/mkdocs.yml b/mkdocs.yml index cb29416..a3028dc 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -68,12 +68,12 @@ nav: - openai_response: reference/endpoints/openai_response.md - sagemaker: reference/endpoints/sagemaker.md - experiments: reference/experiments.md - - serialization: reference/serialization.md - plotting: - reference/plotting/index.md - prompt_utils: reference/prompt_utils.md - results: reference/results.md - runner: reference/runner.md + - serialization: reference/serialization.md - tokenizers: reference/tokenizers.md - utils: reference/utils.md repo_url: https://github.com/awslabs/llmeter From 66bf76f1dab5be6e9446f05cc8ff4de5b13ef210 Mon Sep 17 00:00:00 2001 From: Alex Thewsey Date: Thu, 30 Jul 2026 20:34:51 +0800 Subject: [PATCH 19/20] fix(result): recover_metadata with new ep format Fix Result._recover_metadata (used when .load()-ing a result with no summary.json) to work with the new format of Endpoint serialization, as well as the old. This defect was found in the course of adding more extensive snapshot-based regression tests to catch any failure of result loading introduced by serialization changes. --- llmeter/results.py | 17 +- tests/unit/conftest.py | 16 + tests/unit/fixtures/_generate_fixtures.py | 901 ++++++++++++++++++ .../result_snapshots/base/payload.jsonl | 2 + .../result_snapshots/base/responses.jsonl | 5 + .../result_snapshots/base/run_config.json | 65 ++ .../fixtures/result_snapshots/base/stats.json | 45 + .../result_snapshots/base/summary.json | 16 + .../errors_and_annotations/responses.jsonl | 5 + .../errors_and_annotations/stats.json | 45 + .../errors_and_annotations/summary.json | 16 + .../interrupted_run/responses.jsonl | 5 + .../interrupted_run/run_config.json | 27 + .../legacy/v0_1_endpoint_type/responses.jsonl | 5 + .../legacy/v0_1_endpoint_type/run_config.json | 20 + .../legacy/v0_1_endpoint_type/stats.json | 41 + .../legacy/v0_1_endpoint_type/summary.json | 16 + .../legacy/v0_1_interrupted/responses.jsonl | 5 + .../legacy/v0_1_interrupted/run_config.json | 20 + .../legacy/v0_1_str_callbacks/responses.jsonl | 5 + .../legacy/v0_1_str_callbacks/run_config.json | 25 + .../legacy/v0_1_str_callbacks/stats.json | 41 + .../legacy/v0_1_str_callbacks/summary.json | 16 + .../load_test/00001-clients/responses.jsonl | 5 + .../load_test/00001-clients/stats.json | 41 + .../load_test/00001-clients/summary.json | 16 + .../load_test/00003-clients/responses.jsonl | 5 + .../load_test/00003-clients/stats.json | 41 + .../load_test/00003-clients/summary.json | 16 + tests/unit/test_snapshot_load.py | 354 +++++++ 30 files changed, 1834 insertions(+), 3 deletions(-) create mode 100644 tests/unit/conftest.py create mode 100644 tests/unit/fixtures/_generate_fixtures.py create mode 100644 tests/unit/fixtures/result_snapshots/base/payload.jsonl create mode 100644 tests/unit/fixtures/result_snapshots/base/responses.jsonl create mode 100644 tests/unit/fixtures/result_snapshots/base/run_config.json create mode 100644 tests/unit/fixtures/result_snapshots/base/stats.json create mode 100644 tests/unit/fixtures/result_snapshots/base/summary.json create mode 100644 tests/unit/fixtures/result_snapshots/errors_and_annotations/responses.jsonl create mode 100644 tests/unit/fixtures/result_snapshots/errors_and_annotations/stats.json create mode 100644 tests/unit/fixtures/result_snapshots/errors_and_annotations/summary.json create mode 100644 tests/unit/fixtures/result_snapshots/interrupted_run/responses.jsonl create mode 100644 tests/unit/fixtures/result_snapshots/interrupted_run/run_config.json create mode 100644 tests/unit/fixtures/result_snapshots/legacy/v0_1_endpoint_type/responses.jsonl create mode 100644 tests/unit/fixtures/result_snapshots/legacy/v0_1_endpoint_type/run_config.json create mode 100644 tests/unit/fixtures/result_snapshots/legacy/v0_1_endpoint_type/stats.json create mode 100644 tests/unit/fixtures/result_snapshots/legacy/v0_1_endpoint_type/summary.json create mode 100644 tests/unit/fixtures/result_snapshots/legacy/v0_1_interrupted/responses.jsonl create mode 100644 tests/unit/fixtures/result_snapshots/legacy/v0_1_interrupted/run_config.json create mode 100644 tests/unit/fixtures/result_snapshots/legacy/v0_1_str_callbacks/responses.jsonl create mode 100644 tests/unit/fixtures/result_snapshots/legacy/v0_1_str_callbacks/run_config.json create mode 100644 tests/unit/fixtures/result_snapshots/legacy/v0_1_str_callbacks/stats.json create mode 100644 tests/unit/fixtures/result_snapshots/legacy/v0_1_str_callbacks/summary.json create mode 100644 tests/unit/fixtures/result_snapshots/load_test/00001-clients/responses.jsonl create mode 100644 tests/unit/fixtures/result_snapshots/load_test/00001-clients/stats.json create mode 100644 tests/unit/fixtures/result_snapshots/load_test/00001-clients/summary.json create mode 100644 tests/unit/fixtures/result_snapshots/load_test/00003-clients/responses.jsonl create mode 100644 tests/unit/fixtures/result_snapshots/load_test/00003-clients/stats.json create mode 100644 tests/unit/fixtures/result_snapshots/load_test/00003-clients/summary.json create mode 100644 tests/unit/test_snapshot_load.py diff --git a/llmeter/results.py b/llmeter/results.py index 12fa3ab..0b01133 100644 --- a/llmeter/results.py +++ b/llmeter/results.py @@ -293,9 +293,20 @@ def _recover_metadata(cls, result_path) -> dict[str, Any]: metadata[passthru_field] = config[passthru_field] endpoint = config.get("endpoint", {}) if isinstance(endpoint, dict): - metadata["model_id"] = endpoint.get("model_id") - metadata["endpoint_name"] = endpoint.get("endpoint_name") - metadata["provider"] = endpoint.get("provider") + # Modern format nests endpoint fields under `__llmeter_state__`; legacy + # format keeps them at the top level. + if "__llmeter_class__" in endpoint: + endpoint_fields = endpoint.get("__llmeter_state__", {}) + else: + logger.warning( + "Recovering endpoint metadata from a legacy data format. Support for " + "this will be removed in a future major version. Consider re-saving " + "the file to update." + ) + endpoint_fields = endpoint + metadata["model_id"] = endpoint_fields.get("model_id") + metadata["endpoint_name"] = endpoint_fields.get("endpoint_name") + metadata["provider"] = endpoint_fields.get("provider") except (json.JSONDecodeError, KeyError) as e: logger.warning(f"Could not parse run_config.json: {e}") diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py new file mode 100644 index 0000000..871b09a --- /dev/null +++ b/tests/unit/conftest.py @@ -0,0 +1,16 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Shared fixtures for unit tests.""" + +from pathlib import Path + +import pytest + +RESULT_SNAPSHOTS_DIR = Path(__file__).parent / "fixtures" / "result_snapshots" + + +@pytest.fixture +def snapshots_dir() -> Path: + """Base directory holding the static synthetic result snapshots.""" + return RESULT_SNAPSHOTS_DIR diff --git a/tests/unit/fixtures/_generate_fixtures.py b/tests/unit/fixtures/_generate_fixtures.py new file mode 100644 index 0000000..30cee45 --- /dev/null +++ b/tests/unit/fixtures/_generate_fixtures.py @@ -0,0 +1,901 @@ +#!/usr/bin/env python3 +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 +"""One-shot script to generate fixture data with mathematically consistent stats. + +Run once from the repo root: + uv run python tests/unit/fixtures/_generate_fixtures.py + +This writes all fixture files under tests/unit/fixtures/result_snapshots/. +After running, verify the tests pass, then this script can be kept as documentation +of how the fixtures were generated (or deleted if preferred). +""" + +import json +import sys +from datetime import datetime, timezone +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[3])) + +from llmeter.endpoints.base import InvocationResponse +from llmeter.results import Result +from llmeter.serialization import json_default, restore_dataclass_types + +FIXTURES_DIR = Path(__file__).parent / "result_snapshots" + + +def _write_json(path: Path, data: dict): + path.write_text(json.dumps(data, indent=4, default=json_default) + "\n") + + +def _write_jsonl(path: Path, lines: list[str]): + path.write_text("\n".join(lines) + "\n") + + +def _compute_and_write_stats( + responses: list[InvocationResponse], summary: dict, path: Path +): + """Build a Result from the summary metadata, compute stats, write stats.json. + + Passing the full summary metadata (parsed to native types) makes the resulting + stats.json mirror what a real ``Result.save()`` produces. + """ + meta = {k: v for k, v in summary.items() if k not in ("total_requests",)} + restore_dataclass_types(Result, meta) + result = Result( + responses=responses, total_requests=summary["total_requests"], **meta + ) + stats = Result._compute_stats(result) + _write_json(path, stats) + return stats + + +# ============================================================================= +# Scenario: base (modern format, OpenAI endpoint, CostModel + Mlflow callbacks) +# ============================================================================= + + +def generate_base(): + out = FIXTURES_DIR / "base" + + responses = [ + InvocationResponse( + id="synth-base-001", + response_text="Lorem ipsum dolor sit amet, consectetur adipiscing elit.", + input_prompt="Summarize the key concepts of distributed systems.", + input_payload={ + "messages": [ + { + "role": "user", + "content": "Summarize the key concepts of distributed systems.", + } + ], + "max_tokens": 256, + "model": "synthetic.test-model-v1", + "stream": True, + "stream_options": {"include_usage": True}, + }, + time_to_first_token=0.452, + time_to_last_token=1.203, + num_tokens_input=42, + num_tokens_output=87, + num_tokens_input_cached=None, + num_tokens_output_reasoning=None, + time_per_output_token=0.00863, + error=None, + retries=0, + request_time=datetime(2025, 1, 15, 10, 0, 0, tzinfo=timezone.utc), + annotations={}, + ), + InvocationResponse( + id="synth-base-002", + response_text="Sed ut perspiciatis unde omnis iste natus error sit voluptatem.", + input_prompt="Explain the CAP theorem briefly.", + input_payload={ + "messages": [ + {"role": "user", "content": "Explain the CAP theorem briefly."} + ], + "max_tokens": 256, + "model": "synthetic.test-model-v1", + "stream": True, + "stream_options": {"include_usage": True}, + }, + time_to_first_token=0.389, + time_to_last_token=0.951, + num_tokens_input=38, + num_tokens_output=64, + num_tokens_input_cached=12, + num_tokens_output_reasoning=None, + time_per_output_token=0.00879, + error=None, + retries=0, + request_time=datetime(2025, 1, 15, 10, 0, 1, tzinfo=timezone.utc), + annotations={}, + ), + InvocationResponse( + id="synth-base-003", + response_text="Nemo enim ipsam voluptatem quia voluptas sit aspernatur.", + input_prompt="What is eventual consistency?", + input_payload={ + "messages": [ + {"role": "user", "content": "What is eventual consistency?"} + ], + "max_tokens": 256, + "model": "synthetic.test-model-v1", + "stream": True, + "stream_options": {"include_usage": True}, + }, + time_to_first_token=0.521, + time_to_last_token=1.445, + num_tokens_input=35, + num_tokens_output=102, + num_tokens_input_cached=None, + num_tokens_output_reasoning=8, + time_per_output_token=0.00906, + error=None, + retries=0, + request_time=datetime(2025, 1, 15, 10, 0, 2, tzinfo=timezone.utc), + annotations={}, + ), + InvocationResponse( + id="synth-base-004", + response_text="Ut enim ad minima veniam, quis nostrum exercitationem.", + input_prompt="Define partition tolerance.", + input_payload={ + "messages": [ + {"role": "user", "content": "Define partition tolerance."} + ], + "max_tokens": 256, + "model": "synthetic.test-model-v1", + "stream": True, + "stream_options": {"include_usage": True}, + }, + time_to_first_token=0.298, + time_to_last_token=0.812, + num_tokens_input=30, + num_tokens_output=55, + num_tokens_input_cached=None, + num_tokens_output_reasoning=None, + time_per_output_token=0.00934, + error=None, + retries=0, + request_time=datetime(2025, 1, 15, 10, 0, 3, tzinfo=timezone.utc), + annotations={}, + ), + InvocationResponse( + id="synth-base-005", + response_text="Quis autem vel eum iure reprehenderit qui in ea voluptate.", + input_prompt="Compare synchronous and asynchronous replication.", + input_payload={ + "messages": [ + { + "role": "user", + "content": "Compare synchronous and asynchronous replication.", + } + ], + "max_tokens": 256, + "model": "synthetic.test-model-v1", + "stream": True, + "stream_options": {"include_usage": True}, + }, + time_to_first_token=0.415, + time_to_last_token=1.102, + num_tokens_input=40, + num_tokens_output=78, + num_tokens_input_cached=None, + num_tokens_output_reasoning=None, + time_per_output_token=0.00881, + error=None, + retries=0, + request_time=datetime(2025, 1, 15, 10, 0, 4, tzinfo=timezone.utc), + annotations={}, + ), + ] + + # summary.json + summary = { + "total_requests": 5, + "clients": 2, + "n_requests": 3, + "total_test_time": 4.521, + "model_id": "synthetic.test-model-v1", + "output_path": None, + "endpoint_name": "synthetic-openai-stream", + "provider": "openai", + "run_name": "synthetic-base-test", + "run_description": "Synthetic fixture for deserialization regression tests", + "start_time": "2025-01-15T10:00:00Z", + "first_request_time": "2025-01-15T10:00:00Z", + "last_request_time": "2025-01-15T10:00:04Z", + "end_time": "2025-01-15T10:00:04Z", + } + _write_json(out / "summary.json", summary) + + # stats.json + _compute_and_write_stats(responses, summary, out / "stats.json") + + # responses.jsonl + _write_jsonl(out / "responses.jsonl", [r.to_json() for r in responses]) + + # run_config.json (modern __llmeter_class__ envelopes) + run_config = { + "endpoint": { + "__llmeter_class__": "llmeter.endpoints.openai.OpenAICompletionStreamEndpoint", + "__llmeter_state__": { + "model_id": "synthetic.test-model-v1", + "endpoint_name": "synthetic-openai-stream", + "provider": "openai", + "organization": None, + "base_url": "https://api.synthetic-provider.example.com/v1/", + "max_retries": 3, + "timeout": { + "connect": 5.0, + "read": 60.0, + "write": 5.0, + "pool": 5.0, + }, + "api_key": "sk-synthetic-test-key-not-real", + }, + }, + "output_path": "stale/path/that/should/be/overridden", + "tokenizer": { + "__llmeter_class__": "llmeter.tokenizers.DummyTokenizer", + "__llmeter_state__": {}, + }, + "clients": 2, + "n_requests": 3, + "run_duration": None, + "payload": "tests/unit/fixtures/result_snapshots/base/payload.jsonl", + "run_name": "synthetic-base-test", + "run_description": "Synthetic fixture for deserialization regression tests", + "timeout": 60, + "callbacks": [ + { + "__llmeter_class__": "llmeter.callbacks.cost.model.CostModel", + "__llmeter_state__": { + "request_dims": { + "InputTokens": { + "__llmeter_class__": "llmeter.callbacks.cost.dimensions.InputTokens", + "__llmeter_state__": { + "price_per_million": 3.0, + "granularity": 1, + }, + }, + "OutputTokens": { + "__llmeter_class__": "llmeter.callbacks.cost.dimensions.OutputTokens", + "__llmeter_state__": { + "price_per_million": 15.0, + "granularity": 1, + }, + }, + }, + "run_dims": {}, + }, + }, + { + "__llmeter_class__": "llmeter.callbacks.mlflow.MlflowCallback", + "__llmeter_state__": { + "step": None, + "nested": True, + }, + }, + ], + "low_memory": False, + "progress_bar_stats": None, + } + _write_json(out / "run_config.json", run_config) + + # payload.jsonl + payloads = [ + { + "messages": [ + { + "role": "user", + "content": "Summarize the key concepts of distributed systems.", + } + ], + "max_tokens": 256, + "model": "synthetic.test-model-v1", + "stream": True, + "stream_options": {"include_usage": True}, + }, + { + "messages": [ + {"role": "user", "content": "Explain the CAP theorem briefly."} + ], + "max_tokens": 256, + "model": "synthetic.test-model-v1", + "stream": True, + "stream_options": {"include_usage": True}, + }, + ] + _write_jsonl(out / "payload.jsonl", [json.dumps(p) for p in payloads]) + + print(f" base/ done ({len(responses)} responses)") + + +# ============================================================================= +# Scenario: legacy/v0_1_endpoint_type +# ============================================================================= + + +def generate_legacy_endpoint_type(): + out = FIXTURES_DIR / "legacy" / "v0_1_endpoint_type" + + responses = [ + InvocationResponse( + id="synth-legacy-001", + response_text="Consectetur adipiscing elit, sed do eiusmod tempor.", + input_prompt="What is a load balancer?", + time_to_first_token=0.612, + time_to_last_token=1.534, + num_tokens_input=28, + num_tokens_output=72, + time_per_output_token=0.01280, + error=None, + request_time=datetime(2025, 1, 10, 14, 0, 0, tzinfo=timezone.utc), + ), + InvocationResponse( + id="synth-legacy-002", + response_text="Duis aute irure dolor in reprehenderit in voluptate velit.", + input_prompt="Explain horizontal scaling.", + time_to_first_token=0.548, + time_to_last_token=1.289, + num_tokens_input=24, + num_tokens_output=65, + time_per_output_token=0.01140, + error=None, + request_time=datetime(2025, 1, 10, 14, 0, 1, tzinfo=timezone.utc), + ), + InvocationResponse( + id="synth-legacy-003", + response_text="Excepteur sint occaecat cupidatat non proident.", + input_prompt="What is auto-scaling?", + time_to_first_token=0.701, + time_to_last_token=1.678, + num_tokens_input=22, + num_tokens_output=58, + time_per_output_token=0.01684, + error=None, + request_time=datetime(2025, 1, 10, 14, 0, 2, tzinfo=timezone.utc), + ), + InvocationResponse( + id="synth-legacy-004", + response_text="Sunt in culpa qui officia deserunt mollit anim id est laborum.", + input_prompt="Define microservices architecture.", + time_to_first_token=0.489, + time_to_last_token=1.102, + num_tokens_input=26, + num_tokens_output=71, + time_per_output_token=0.00863, + error=None, + request_time=datetime(2025, 1, 10, 14, 0, 3, tzinfo=timezone.utc), + ), + InvocationResponse( + id="synth-legacy-005", + response_text="Neque porro quisquam est qui dolorem ipsum.", + input_prompt="What is service mesh?", + time_to_first_token=0.555, + time_to_last_token=1.345, + num_tokens_input=20, + num_tokens_output=48, + time_per_output_token=0.01646, + error=None, + request_time=datetime(2025, 1, 10, 14, 0, 4, tzinfo=timezone.utc), + ), + ] + + # summary.json + summary = { + "total_requests": 5, + "clients": 3, + "n_requests": 2, + "total_test_time": 5.102, + "model_id": "synthetic.legacy-model-v1", + "output_path": "stale/legacy/path", + "endpoint_name": "synthetic-legacy-endpoint", + "provider": "bedrock", + "run_name": "synthetic-legacy-test", + "run_description": None, + "start_time": "2025-01-10T14:00:00Z", + "first_request_time": "2025-01-10T14:00:00Z", + "last_request_time": "2025-01-10T14:00:04Z", + "end_time": "2025-01-10T14:00:05Z", + } + _write_json(out / "summary.json", summary) + + # stats.json + _compute_and_write_stats(responses, summary, out / "stats.json") + + # responses.jsonl — legacy format with flat cost annotations as top-level keys + legacy_responses = [] + costs = [ + { + "cost_InputTokens": 0.000084, + "cost_OutputTokens": 0.00108, + "cost_total": 0.001164, + }, + { + "cost_InputTokens": 0.000072, + "cost_OutputTokens": 0.000975, + "cost_total": 0.001047, + }, + { + "cost_InputTokens": 0.000066, + "cost_OutputTokens": 0.00087, + "cost_total": 0.000936, + }, + { + "cost_InputTokens": 0.000078, + "cost_OutputTokens": 0.001065, + "cost_total": 0.001143, + }, + { + "cost_InputTokens": 0.00006, + "cost_OutputTokens": 0.00072, + "cost_total": 0.00078, + }, + ] + for resp, cost in zip(responses, costs): + d = json.loads(resp.to_json()) + # Remove modern fields that didn't exist in v0.1 + d.pop("num_tokens_input_cached", None) + d.pop("num_tokens_output_reasoning", None) + d.pop("retries", None) + d.pop("annotations", None) + # Add legacy flat cost annotations + d.update(cost) + legacy_responses.append(json.dumps(d)) + _write_jsonl(out / "responses.jsonl", legacy_responses) + + # run_config.json — legacy format + run_config = { + "endpoint": { + "endpoint_name": "synthetic-legacy-endpoint", + "model_id": "synthetic.legacy-model-v1", + "provider": "bedrock", + "region": "us-west-2", + "endpoint_type": "BedrockConverse", + }, + "output_path": "stale/legacy/path", + "tokenizer": {"tokenizer_module": "llmeter"}, + "clients": 3, + "n_requests": 2, + "payload": "stale/legacy/payload.jsonl", + "run_name": "synthetic-legacy-test", + "run_description": None, + "timeout": 60, + "callbacks": None, + } + _write_json(out / "run_config.json", run_config) + + print(f" legacy/v0_1_endpoint_type/ done ({len(responses)} responses)") + + +# ============================================================================= +# Scenario: legacy/v0_1_str_callbacks +# ============================================================================= + + +def generate_legacy_str_callbacks(): + out = FIXTURES_DIR / "legacy" / "v0_1_str_callbacks" + + responses = [ + InvocationResponse( + id=f"synth-strcb-{i:03d}", + response_text=f"Synthetic response number {i} for string callback test.", + input_prompt=f"Synthetic prompt {i}.", + time_to_first_token=0.4 + 0.1 * i, + time_to_last_token=1.0 + 0.2 * i, + num_tokens_input=30 + 5 * i, + num_tokens_output=50 + 10 * i, + time_per_output_token=0.01, + error=None, + request_time=datetime(2025, 1, 12, 9, 0, i, tzinfo=timezone.utc), + ) + for i in range(1, 6) + ] + + # summary.json + summary = { + "total_requests": 5, + "clients": 1, + "n_requests": 5, + "total_test_time": 6.5, + "model_id": "synthetic.strcb-model-v1", + "output_path": "stale/strcb/path", + "endpoint_name": "synthetic-strcb-endpoint", + "provider": "bedrock", + "run_name": "synthetic-strcb-test", + "run_description": None, + "start_time": "2025-01-12T09:00:01Z", + "first_request_time": "2025-01-12T09:00:01Z", + "last_request_time": "2025-01-12T09:00:05Z", + "end_time": "2025-01-12T09:00:06Z", + } + _write_json(out / "summary.json", summary) + + # stats.json + _compute_and_write_stats(responses, summary, out / "stats.json") + + # responses.jsonl + _write_jsonl(out / "responses.jsonl", [r.to_json() for r in responses]) + + # run_config.json — legacy with string repr callbacks + run_config = { + "endpoint": { + "endpoint_name": "synthetic-strcb-endpoint", + "model_id": "synthetic.strcb-model-v1", + "provider": "bedrock", + "region": "us-east-1", + "endpoint_type": "BedrockConverseStream", + }, + "output_path": "stale/strcb/path", + "tokenizer": {"tokenizer_module": "llmeter"}, + "clients": 1, + "n_requests": 5, + "payload": "stale/strcb/payload.jsonl", + "run_name": "synthetic-strcb-test", + "run_description": None, + "timeout": 60, + "callbacks": [ + "<__main__.SyntheticCallback object at 0x1234567890>", + "", + ], + "low_memory": False, + "progress_bar_stats": None, + } + _write_json(out / "run_config.json", run_config) + + print(f" legacy/v0_1_str_callbacks/ done ({len(responses)} responses)") + + +# ============================================================================= +# Scenario: interrupted_run (no summary.json) +# ============================================================================= + + +def generate_interrupted_run(): + out = FIXTURES_DIR / "interrupted_run" + + responses = [ + InvocationResponse( + id=f"synth-intr-{i:03d}", + response_text=f"Interrupted run response {i}: tempor incididunt ut labore.", + input_prompt=f"Interrupted test prompt {i}.", + time_to_first_token=0.3 + 0.05 * i, + time_to_last_token=0.8 + 0.15 * i, + num_tokens_input=25 + 3 * i, + num_tokens_output=40 + 8 * i, + time_per_output_token=0.009, + error=None, + retries=0, + request_time=datetime(2025, 1, 20, 16, 30, i * 2, tzinfo=timezone.utc), + ) + for i in range(1, 6) + ] + + # NO summary.json — this is the point of this scenario + + # responses.jsonl + _write_jsonl(out / "responses.jsonl", [r.to_json() for r in responses]) + + # run_config.json — modern __llmeter_class__ format (what a *current* interrupted run leaves) + run_config = { + "endpoint": { + "__llmeter_class__": "llmeter.endpoints.bedrock.BedrockConverseStream", + "__llmeter_state__": { + # Mirrors real dump_object() output: fields the constructor accepts, with + # derived values like `provider` deliberately absent (recomputed on load). + "model_id": "synthetic.interrupted-model-v1", + "endpoint_name": "synthetic-interrupted-endpoint", + "region": "eu-west-1", + "inference_config": None, + "max_attempts": 3, + "ttft_visible_tokens_only": True, + }, + }, + "output_path": "stale/interrupted/path", + "tokenizer": { + "__llmeter_class__": "llmeter.tokenizers.DummyTokenizer", + "__llmeter_state__": {}, + }, + "clients": 4, + "n_requests": 20, + "payload": "stale/interrupted/payload.jsonl", + "run_name": "synthetic-interrupted-test", + "run_description": "This run was interrupted before completion", + "timeout": 120, + "callbacks": None, + "low_memory": False, + "progress_bar_stats": None, + } + _write_json(out / "run_config.json", run_config) + + print(f" interrupted_run/ done ({len(responses)} responses, no summary.json)") + + +# ============================================================================= +# Scenario: legacy/v0_1_interrupted (legacy-format config, no summary.json) +# ============================================================================= + + +def generate_legacy_interrupted_run(): + """A legacy-format interrupted run: no summary.json, top-level endpoint fields. + + Exercises the legacy branch of ``Result._recover_metadata`` - endpoint fields + (including the derived ``provider``) live at the top level of the endpoint dict, + and loading should recover them while warning that the format is deprecated. + """ + out = FIXTURES_DIR / "legacy" / "v0_1_interrupted" + + responses = [ + InvocationResponse( + id=f"synth-legacy-intr-{i:03d}", + response_text=f"Legacy interrupted response {i}.", + input_prompt=f"Legacy interrupted prompt {i}.", + time_to_first_token=0.35 + 0.05 * i, + time_to_last_token=0.9 + 0.1 * i, + num_tokens_input=20 + 4 * i, + num_tokens_output=45 + 6 * i, + time_per_output_token=0.01, + error=None, + request_time=datetime(2025, 1, 5, 12, 0, i, tzinfo=timezone.utc), + ) + for i in range(1, 6) + ] + + # NO summary.json — recovery path + + # responses.jsonl (legacy: no retries/annotations/cached/reasoning fields) + legacy_responses = [] + for resp in responses: + d = json.loads(resp.to_json()) + d.pop("num_tokens_input_cached", None) + d.pop("num_tokens_output_reasoning", None) + d.pop("retries", None) + d.pop("annotations", None) + legacy_responses.append(json.dumps(d)) + _write_jsonl(out / "responses.jsonl", legacy_responses) + + # run_config.json — legacy format with top-level endpoint fields (incl. derived provider) + run_config = { + "endpoint": { + "endpoint_name": "synthetic-legacy-intr-endpoint", + "model_id": "synthetic.legacy-interrupted-v1", + "provider": "bedrock", + "region": "us-east-1", + "endpoint_type": "BedrockConverse", + }, + "output_path": "stale/legacy/interrupted/path", + "tokenizer": {"tokenizer_module": "llmeter"}, + "clients": 2, + "n_requests": 4, + "payload": "stale/legacy/interrupted/payload.jsonl", + "run_name": "synthetic-legacy-interrupted-test", + "run_description": None, + "timeout": 60, + "callbacks": None, + } + _write_json(out / "run_config.json", run_config) + + print( + f" legacy/v0_1_interrupted/ done ({len(responses)} responses, no summary.json)" + ) + + +# ============================================================================= +# Scenario: errors_and_annotations +# ============================================================================= + + +def generate_errors_and_annotations(): + out = FIXTURES_DIR / "errors_and_annotations" + + responses = [ + InvocationResponse( + id="synth-err-001", + response_text="Successful response with custom annotations.", + input_prompt="Test prompt with annotations.", + time_to_first_token=0.350, + time_to_last_token=0.920, + num_tokens_input=32, + num_tokens_output=56, + time_per_output_token=0.01018, + error=None, + retries=0, + request_time=datetime(2025, 1, 18, 8, 0, 0, tzinfo=timezone.utc), + annotations={"custom_metric": 42.5, "experiment_tag": "baseline"}, + ), + InvocationResponse( + id="synth-err-002", + response_text="", + input_prompt="This request was throttled.", + time_to_first_token=None, + time_to_last_token=2.105, + num_tokens_input=25, + num_tokens_output=0, + time_per_output_token=None, + error="ThrottlingException: Rate exceeded", + retries=2, + request_time=datetime(2025, 1, 18, 8, 0, 1, tzinfo=timezone.utc), + annotations={}, + ), + InvocationResponse( + id="synth-err-003", + response_text="Another successful response after the error.", + input_prompt="Normal prompt after throttling.", + time_to_first_token=0.410, + time_to_last_token=1.050, + num_tokens_input=30, + num_tokens_output=62, + time_per_output_token=0.01032, + error=None, + retries=0, + request_time=datetime(2025, 1, 18, 8, 0, 3, tzinfo=timezone.utc), + annotations={}, + ), + InvocationResponse( + id="synth-err-004", + response_text="", + input_prompt="This request timed out.", + time_to_first_token=None, + time_to_last_token=None, + num_tokens_input=22, + num_tokens_output=0, + time_per_output_token=None, + error="ReadTimeoutError: timed out after 60s", + retries=0, + request_time=datetime(2025, 1, 18, 8, 0, 5, tzinfo=timezone.utc), + annotations={}, + ), + InvocationResponse( + id="synth-err-005", + response_text="Response with cached input tokens.", + input_prompt="Prompt that leverages prompt caching.", + time_to_first_token=0.180, + time_to_last_token=0.650, + num_tokens_input=45, + num_tokens_output=38, + num_tokens_input_cached=30, + num_tokens_output_reasoning=None, + time_per_output_token=0.01237, + error=None, + retries=1, + request_time=datetime(2025, 1, 18, 8, 0, 7, tzinfo=timezone.utc), + annotations={"cache_hit_ratio": 0.667}, + ), + ] + + # summary.json + summary = { + "total_requests": 5, + "clients": 1, + "n_requests": 5, + "total_test_time": 7.8, + "model_id": "synthetic.errors-model-v1", + "output_path": "stale/errors/path", + "endpoint_name": "synthetic-errors-endpoint", + "provider": "bedrock", + "run_name": "synthetic-errors-test", + "run_description": "Fixture with errors, retries, and annotations", + "start_time": "2025-01-18T08:00:00Z", + "first_request_time": "2025-01-18T08:00:00Z", + "last_request_time": "2025-01-18T08:00:07Z", + "end_time": "2025-01-18T08:00:07Z", + } + _write_json(out / "summary.json", summary) + + # responses.jsonl + _write_jsonl(out / "responses.jsonl", [r.to_json() for r in responses]) + + # stats.json — computed from the responses (including errors) + _compute_and_write_stats(responses, summary, out / "stats.json") + + print(f" errors_and_annotations/ done ({len(responses)} responses, 2 errors)") + + +# ============================================================================= +# Scenario: load_test (multi-concurrency) +# ============================================================================= + + +def generate_load_test(): + out = FIXTURES_DIR / "load_test" + + # 00001-clients + responses_1 = [ + InvocationResponse( + id=f"synth-lt1-{i:03d}", + response_text=f"Load test single-client response {i}.", + input_prompt=f"Load test prompt {i}.", + time_to_first_token=0.200 + 0.05 * i, + time_to_last_token=0.600 + 0.1 * i, + num_tokens_input=20 + 2 * i, + num_tokens_output=35 + 5 * i, + time_per_output_token=0.008, + error=None, + retries=0, + request_time=datetime(2025, 1, 22, 12, 0, i, tzinfo=timezone.utc), + ) + for i in range(1, 6) + ] + + out_1 = out / "00001-clients" + summary_1 = { + "total_requests": 5, + "clients": 1, + "n_requests": 5, + "total_test_time": 5.5, + "model_id": "synthetic.loadtest-model-v1", + "output_path": "stale/loadtest/00001-clients", + "endpoint_name": "synthetic-loadtest-endpoint", + "provider": "bedrock", + "run_name": "00001-clients", + "run_description": None, + "start_time": "2025-01-22T12:00:01Z", + "first_request_time": "2025-01-22T12:00:01Z", + "last_request_time": "2025-01-22T12:00:05Z", + "end_time": "2025-01-22T12:00:05Z", + } + _write_json(out_1 / "summary.json", summary_1) + _compute_and_write_stats(responses_1, summary_1, out_1 / "stats.json") + _write_jsonl(out_1 / "responses.jsonl", [r.to_json() for r in responses_1]) + + # 00003-clients + responses_3 = [ + InvocationResponse( + id=f"synth-lt3-{i:03d}", + response_text=f"Load test triple-client response {i}.", + input_prompt=f"Load test prompt {i}.", + time_to_first_token=0.350 + 0.08 * i, + time_to_last_token=0.900 + 0.15 * i, + num_tokens_input=20 + 2 * i, + num_tokens_output=35 + 5 * i, + time_per_output_token=0.012, + error=None, + retries=0, + request_time=datetime(2025, 1, 22, 12, 5, i, tzinfo=timezone.utc), + ) + for i in range(1, 6) + ] + + out_3 = out / "00003-clients" + summary_3 = { + "total_requests": 5, + "clients": 3, + "n_requests": 2, + "total_test_time": 3.2, + "model_id": "synthetic.loadtest-model-v1", + "output_path": "stale/loadtest/00003-clients", + "endpoint_name": "synthetic-loadtest-endpoint", + "provider": "bedrock", + "run_name": "00003-clients", + "run_description": None, + "start_time": "2025-01-22T12:05:01Z", + "first_request_time": "2025-01-22T12:05:01Z", + "last_request_time": "2025-01-22T12:05:05Z", + "end_time": "2025-01-22T12:05:05Z", + } + _write_json(out_3 / "summary.json", summary_3) + _compute_and_write_stats(responses_3, summary_3, out_3 / "stats.json") + _write_jsonl(out_3 / "responses.jsonl", [r.to_json() for r in responses_3]) + + print( + f" load_test/ done (2 subdirs, {len(responses_1)}+{len(responses_3)} responses)" + ) + + +# ============================================================================= + +if __name__ == "__main__": + print("Generating fixture data...") + generate_base() + generate_legacy_endpoint_type() + generate_legacy_str_callbacks() + generate_interrupted_run() + generate_legacy_interrupted_run() + generate_errors_and_annotations() + generate_load_test() + print("Done!") diff --git a/tests/unit/fixtures/result_snapshots/base/payload.jsonl b/tests/unit/fixtures/result_snapshots/base/payload.jsonl new file mode 100644 index 0000000..6168fbe --- /dev/null +++ b/tests/unit/fixtures/result_snapshots/base/payload.jsonl @@ -0,0 +1,2 @@ +{"messages": [{"role": "user", "content": "Summarize the key concepts of distributed systems."}], "max_tokens": 256, "model": "synthetic.test-model-v1", "stream": true, "stream_options": {"include_usage": true}} +{"messages": [{"role": "user", "content": "Explain the CAP theorem briefly."}], "max_tokens": 256, "model": "synthetic.test-model-v1", "stream": true, "stream_options": {"include_usage": true}} diff --git a/tests/unit/fixtures/result_snapshots/base/responses.jsonl b/tests/unit/fixtures/result_snapshots/base/responses.jsonl new file mode 100644 index 0000000..85d79db --- /dev/null +++ b/tests/unit/fixtures/result_snapshots/base/responses.jsonl @@ -0,0 +1,5 @@ +{"response_text": "Lorem ipsum dolor sit amet, consectetur adipiscing elit.", "input_payload": {"messages": [{"role": "user", "content": "Summarize the key concepts of distributed systems."}], "max_tokens": 256, "model": "synthetic.test-model-v1", "stream": true, "stream_options": {"include_usage": true}}, "id": "synth-base-001", "input_prompt": "Summarize the key concepts of distributed systems.", "time_to_first_token": 0.452, "time_to_last_token": 1.203, "num_tokens_input": 42, "num_tokens_output": 87, "num_tokens_input_cached": null, "num_tokens_output_reasoning": null, "time_per_output_token": 0.00863, "error": null, "retries": 0, "request_time": "2025-01-15T10:00:00Z", "annotations": {}} +{"response_text": "Sed ut perspiciatis unde omnis iste natus error sit voluptatem.", "input_payload": {"messages": [{"role": "user", "content": "Explain the CAP theorem briefly."}], "max_tokens": 256, "model": "synthetic.test-model-v1", "stream": true, "stream_options": {"include_usage": true}}, "id": "synth-base-002", "input_prompt": "Explain the CAP theorem briefly.", "time_to_first_token": 0.389, "time_to_last_token": 0.951, "num_tokens_input": 38, "num_tokens_output": 64, "num_tokens_input_cached": 12, "num_tokens_output_reasoning": null, "time_per_output_token": 0.00879, "error": null, "retries": 0, "request_time": "2025-01-15T10:00:01Z", "annotations": {}} +{"response_text": "Nemo enim ipsam voluptatem quia voluptas sit aspernatur.", "input_payload": {"messages": [{"role": "user", "content": "What is eventual consistency?"}], "max_tokens": 256, "model": "synthetic.test-model-v1", "stream": true, "stream_options": {"include_usage": true}}, "id": "synth-base-003", "input_prompt": "What is eventual consistency?", "time_to_first_token": 0.521, "time_to_last_token": 1.445, "num_tokens_input": 35, "num_tokens_output": 102, "num_tokens_input_cached": null, "num_tokens_output_reasoning": 8, "time_per_output_token": 0.00906, "error": null, "retries": 0, "request_time": "2025-01-15T10:00:02Z", "annotations": {}} +{"response_text": "Ut enim ad minima veniam, quis nostrum exercitationem.", "input_payload": {"messages": [{"role": "user", "content": "Define partition tolerance."}], "max_tokens": 256, "model": "synthetic.test-model-v1", "stream": true, "stream_options": {"include_usage": true}}, "id": "synth-base-004", "input_prompt": "Define partition tolerance.", "time_to_first_token": 0.298, "time_to_last_token": 0.812, "num_tokens_input": 30, "num_tokens_output": 55, "num_tokens_input_cached": null, "num_tokens_output_reasoning": null, "time_per_output_token": 0.00934, "error": null, "retries": 0, "request_time": "2025-01-15T10:00:03Z", "annotations": {}} +{"response_text": "Quis autem vel eum iure reprehenderit qui in ea voluptate.", "input_payload": {"messages": [{"role": "user", "content": "Compare synchronous and asynchronous replication."}], "max_tokens": 256, "model": "synthetic.test-model-v1", "stream": true, "stream_options": {"include_usage": true}}, "id": "synth-base-005", "input_prompt": "Compare synchronous and asynchronous replication.", "time_to_first_token": 0.415, "time_to_last_token": 1.102, "num_tokens_input": 40, "num_tokens_output": 78, "num_tokens_input_cached": null, "num_tokens_output_reasoning": null, "time_per_output_token": 0.00881, "error": null, "retries": 0, "request_time": "2025-01-15T10:00:04Z", "annotations": {}} diff --git a/tests/unit/fixtures/result_snapshots/base/run_config.json b/tests/unit/fixtures/result_snapshots/base/run_config.json new file mode 100644 index 0000000..4056dfd --- /dev/null +++ b/tests/unit/fixtures/result_snapshots/base/run_config.json @@ -0,0 +1,65 @@ +{ + "endpoint": { + "__llmeter_class__": "llmeter.endpoints.openai.OpenAICompletionStreamEndpoint", + "__llmeter_state__": { + "model_id": "synthetic.test-model-v1", + "endpoint_name": "synthetic-openai-stream", + "provider": "openai", + "organization": null, + "base_url": "https://api.synthetic-provider.example.com/v1/", + "max_retries": 3, + "timeout": { + "connect": 5.0, + "read": 60.0, + "write": 5.0, + "pool": 5.0 + }, + "api_key": "sk-synthetic-test-key-not-real" + } + }, + "output_path": "stale/path/that/should/be/overridden", + "tokenizer": { + "__llmeter_class__": "llmeter.tokenizers.DummyTokenizer", + "__llmeter_state__": {} + }, + "clients": 2, + "n_requests": 3, + "run_duration": null, + "payload": "tests/unit/fixtures/result_snapshots/base/payload.jsonl", + "run_name": "synthetic-base-test", + "run_description": "Synthetic fixture for deserialization regression tests", + "timeout": 60, + "callbacks": [ + { + "__llmeter_class__": "llmeter.callbacks.cost.model.CostModel", + "__llmeter_state__": { + "request_dims": { + "InputTokens": { + "__llmeter_class__": "llmeter.callbacks.cost.dimensions.InputTokens", + "__llmeter_state__": { + "price_per_million": 3.0, + "granularity": 1 + } + }, + "OutputTokens": { + "__llmeter_class__": "llmeter.callbacks.cost.dimensions.OutputTokens", + "__llmeter_state__": { + "price_per_million": 15.0, + "granularity": 1 + } + } + }, + "run_dims": {} + } + }, + { + "__llmeter_class__": "llmeter.callbacks.mlflow.MlflowCallback", + "__llmeter_state__": { + "step": null, + "nested": true + } + } + ], + "low_memory": false, + "progress_bar_stats": null +} diff --git a/tests/unit/fixtures/result_snapshots/base/stats.json b/tests/unit/fixtures/result_snapshots/base/stats.json new file mode 100644 index 0000000..01744b4 --- /dev/null +++ b/tests/unit/fixtures/result_snapshots/base/stats.json @@ -0,0 +1,45 @@ +{ + "total_requests": 5, + "clients": 2, + "n_requests": 3, + "total_test_time": 4.521, + "model_id": "synthetic.test-model-v1", + "output_path": null, + "endpoint_name": "synthetic-openai-stream", + "provider": "openai", + "run_name": "synthetic-base-test", + "run_description": "Synthetic fixture for deserialization regression tests", + "start_time": "2025-01-15T10:00:00Z", + "first_request_time": "2025-01-15T10:00:00Z", + "last_request_time": "2025-01-15T10:00:04Z", + "end_time": "2025-01-15T10:00:04Z", + "failed_requests": 0, + "failed_requests_rate": 0.0, + "requests_per_minute": 66.35700066357, + "total_input_tokens": 185, + "total_output_tokens": 386, + "total_cached_input_tokens": 12, + "total_reasoning_output_tokens": 8, + "average_input_tokens_per_minute": 2455.2090245520903, + "average_output_tokens_per_minute": 5122.760451227605, + "time_to_last_token-average": 1.1026, + "time_to_last_token-p50": 1.102, + "time_to_last_token-p90": 1.5417999999999998, + "time_to_last_token-p99": 1.67248, + "time_to_first_token-average": 0.415, + "time_to_first_token-p50": 0.415, + "time_to_first_token-p90": 0.5486000000000001, + "time_to_first_token-p99": 0.5858599999999999, + "num_tokens_output-average": 77.2, + "num_tokens_output-p50": 78, + "num_tokens_output-p90": 108.0, + "num_tokens_output-p99": 116.1, + "num_tokens_input-average": 37, + "num_tokens_input-p50": 38, + "num_tokens_input-p90": 42.8, + "num_tokens_input-p99": 43.88, + "num_tokens_input_cached-average": 12, + "num_tokens_input_cached-p50": 12, + "num_tokens_input_cached-p90": 12, + "num_tokens_input_cached-p99": 12 +} diff --git a/tests/unit/fixtures/result_snapshots/base/summary.json b/tests/unit/fixtures/result_snapshots/base/summary.json new file mode 100644 index 0000000..8ecc92e --- /dev/null +++ b/tests/unit/fixtures/result_snapshots/base/summary.json @@ -0,0 +1,16 @@ +{ + "total_requests": 5, + "clients": 2, + "n_requests": 3, + "total_test_time": 4.521, + "model_id": "synthetic.test-model-v1", + "output_path": null, + "endpoint_name": "synthetic-openai-stream", + "provider": "openai", + "run_name": "synthetic-base-test", + "run_description": "Synthetic fixture for deserialization regression tests", + "start_time": "2025-01-15T10:00:00Z", + "first_request_time": "2025-01-15T10:00:00Z", + "last_request_time": "2025-01-15T10:00:04Z", + "end_time": "2025-01-15T10:00:04Z" +} diff --git a/tests/unit/fixtures/result_snapshots/errors_and_annotations/responses.jsonl b/tests/unit/fixtures/result_snapshots/errors_and_annotations/responses.jsonl new file mode 100644 index 0000000..0b31f9c --- /dev/null +++ b/tests/unit/fixtures/result_snapshots/errors_and_annotations/responses.jsonl @@ -0,0 +1,5 @@ +{"response_text": "Successful response with custom annotations.", "input_payload": null, "id": "synth-err-001", "input_prompt": "Test prompt with annotations.", "time_to_first_token": 0.35, "time_to_last_token": 0.92, "num_tokens_input": 32, "num_tokens_output": 56, "num_tokens_input_cached": null, "num_tokens_output_reasoning": null, "time_per_output_token": 0.01018, "error": null, "retries": 0, "request_time": "2025-01-18T08:00:00Z", "annotations": {"custom_metric": 42.5, "experiment_tag": "baseline"}} +{"response_text": "", "input_payload": null, "id": "synth-err-002", "input_prompt": "This request was throttled.", "time_to_first_token": null, "time_to_last_token": 2.105, "num_tokens_input": 25, "num_tokens_output": 0, "num_tokens_input_cached": null, "num_tokens_output_reasoning": null, "time_per_output_token": null, "error": "ThrottlingException: Rate exceeded", "retries": 2, "request_time": "2025-01-18T08:00:01Z", "annotations": {}} +{"response_text": "Another successful response after the error.", "input_payload": null, "id": "synth-err-003", "input_prompt": "Normal prompt after throttling.", "time_to_first_token": 0.41, "time_to_last_token": 1.05, "num_tokens_input": 30, "num_tokens_output": 62, "num_tokens_input_cached": null, "num_tokens_output_reasoning": null, "time_per_output_token": 0.01032, "error": null, "retries": 0, "request_time": "2025-01-18T08:00:03Z", "annotations": {}} +{"response_text": "", "input_payload": null, "id": "synth-err-004", "input_prompt": "This request timed out.", "time_to_first_token": null, "time_to_last_token": null, "num_tokens_input": 22, "num_tokens_output": 0, "num_tokens_input_cached": null, "num_tokens_output_reasoning": null, "time_per_output_token": null, "error": "ReadTimeoutError: timed out after 60s", "retries": 0, "request_time": "2025-01-18T08:00:05Z", "annotations": {}} +{"response_text": "Response with cached input tokens.", "input_payload": null, "id": "synth-err-005", "input_prompt": "Prompt that leverages prompt caching.", "time_to_first_token": 0.18, "time_to_last_token": 0.65, "num_tokens_input": 45, "num_tokens_output": 38, "num_tokens_input_cached": 30, "num_tokens_output_reasoning": null, "time_per_output_token": 0.01237, "error": null, "retries": 1, "request_time": "2025-01-18T08:00:07Z", "annotations": {"cache_hit_ratio": 0.667}} diff --git a/tests/unit/fixtures/result_snapshots/errors_and_annotations/stats.json b/tests/unit/fixtures/result_snapshots/errors_and_annotations/stats.json new file mode 100644 index 0000000..8786a50 --- /dev/null +++ b/tests/unit/fixtures/result_snapshots/errors_and_annotations/stats.json @@ -0,0 +1,45 @@ +{ + "total_requests": 5, + "clients": 1, + "n_requests": 5, + "total_test_time": 7.8, + "model_id": "synthetic.errors-model-v1", + "output_path": "stale/errors/path", + "endpoint_name": "synthetic-errors-endpoint", + "provider": "bedrock", + "run_name": "synthetic-errors-test", + "run_description": "Fixture with errors, retries, and annotations", + "start_time": "2025-01-18T08:00:00Z", + "first_request_time": "2025-01-18T08:00:00Z", + "last_request_time": "2025-01-18T08:00:07Z", + "end_time": "2025-01-18T08:00:07Z", + "failed_requests": 2, + "failed_requests_rate": 0.4, + "requests_per_minute": 38.46153846153847, + "total_input_tokens": 154, + "total_output_tokens": 156, + "total_cached_input_tokens": 30, + "total_reasoning_output_tokens": 0, + "average_input_tokens_per_minute": 1184.6153846153848, + "average_output_tokens_per_minute": 1200.0, + "time_to_last_token-average": 1.18125, + "time_to_last_token-p50": 0.9850000000000001, + "time_to_last_token-p90": 2.6325, + "time_to_last_token-p99": 3.10725, + "time_to_first_token-average": 0.3133333333333333, + "time_to_first_token-p50": 0.35, + "time_to_first_token-p90": 0.446, + "time_to_first_token-p99": 0.46760000000000007, + "num_tokens_output-average": 31.2, + "num_tokens_output-p50": 38, + "num_tokens_output-p90": 64.4, + "num_tokens_output-p99": 67.64, + "num_tokens_input-average": 30.8, + "num_tokens_input-p50": 30, + "num_tokens_input-p90": 50.2, + "num_tokens_input-p99": 57.22, + "num_tokens_input_cached-average": 30, + "num_tokens_input_cached-p50": 30, + "num_tokens_input_cached-p90": 30, + "num_tokens_input_cached-p99": 30 +} diff --git a/tests/unit/fixtures/result_snapshots/errors_and_annotations/summary.json b/tests/unit/fixtures/result_snapshots/errors_and_annotations/summary.json new file mode 100644 index 0000000..3836355 --- /dev/null +++ b/tests/unit/fixtures/result_snapshots/errors_and_annotations/summary.json @@ -0,0 +1,16 @@ +{ + "total_requests": 5, + "clients": 1, + "n_requests": 5, + "total_test_time": 7.8, + "model_id": "synthetic.errors-model-v1", + "output_path": "stale/errors/path", + "endpoint_name": "synthetic-errors-endpoint", + "provider": "bedrock", + "run_name": "synthetic-errors-test", + "run_description": "Fixture with errors, retries, and annotations", + "start_time": "2025-01-18T08:00:00Z", + "first_request_time": "2025-01-18T08:00:00Z", + "last_request_time": "2025-01-18T08:00:07Z", + "end_time": "2025-01-18T08:00:07Z" +} diff --git a/tests/unit/fixtures/result_snapshots/interrupted_run/responses.jsonl b/tests/unit/fixtures/result_snapshots/interrupted_run/responses.jsonl new file mode 100644 index 0000000..bd2988a --- /dev/null +++ b/tests/unit/fixtures/result_snapshots/interrupted_run/responses.jsonl @@ -0,0 +1,5 @@ +{"response_text": "Interrupted run response 1: tempor incididunt ut labore.", "input_payload": null, "id": "synth-intr-001", "input_prompt": "Interrupted test prompt 1.", "time_to_first_token": 0.35, "time_to_last_token": 0.9500000000000001, "num_tokens_input": 28, "num_tokens_output": 48, "num_tokens_input_cached": null, "num_tokens_output_reasoning": null, "time_per_output_token": 0.009, "error": null, "retries": 0, "request_time": "2025-01-20T16:30:02Z", "annotations": {}} +{"response_text": "Interrupted run response 2: tempor incididunt ut labore.", "input_payload": null, "id": "synth-intr-002", "input_prompt": "Interrupted test prompt 2.", "time_to_first_token": 0.4, "time_to_last_token": 1.1, "num_tokens_input": 31, "num_tokens_output": 56, "num_tokens_input_cached": null, "num_tokens_output_reasoning": null, "time_per_output_token": 0.009, "error": null, "retries": 0, "request_time": "2025-01-20T16:30:04Z", "annotations": {}} +{"response_text": "Interrupted run response 3: tempor incididunt ut labore.", "input_payload": null, "id": "synth-intr-003", "input_prompt": "Interrupted test prompt 3.", "time_to_first_token": 0.45, "time_to_last_token": 1.25, "num_tokens_input": 34, "num_tokens_output": 64, "num_tokens_input_cached": null, "num_tokens_output_reasoning": null, "time_per_output_token": 0.009, "error": null, "retries": 0, "request_time": "2025-01-20T16:30:06Z", "annotations": {}} +{"response_text": "Interrupted run response 4: tempor incididunt ut labore.", "input_payload": null, "id": "synth-intr-004", "input_prompt": "Interrupted test prompt 4.", "time_to_first_token": 0.5, "time_to_last_token": 1.4, "num_tokens_input": 37, "num_tokens_output": 72, "num_tokens_input_cached": null, "num_tokens_output_reasoning": null, "time_per_output_token": 0.009, "error": null, "retries": 0, "request_time": "2025-01-20T16:30:08Z", "annotations": {}} +{"response_text": "Interrupted run response 5: tempor incididunt ut labore.", "input_payload": null, "id": "synth-intr-005", "input_prompt": "Interrupted test prompt 5.", "time_to_first_token": 0.55, "time_to_last_token": 1.55, "num_tokens_input": 40, "num_tokens_output": 80, "num_tokens_input_cached": null, "num_tokens_output_reasoning": null, "time_per_output_token": 0.009, "error": null, "retries": 0, "request_time": "2025-01-20T16:30:10Z", "annotations": {}} diff --git a/tests/unit/fixtures/result_snapshots/interrupted_run/run_config.json b/tests/unit/fixtures/result_snapshots/interrupted_run/run_config.json new file mode 100644 index 0000000..1d441a0 --- /dev/null +++ b/tests/unit/fixtures/result_snapshots/interrupted_run/run_config.json @@ -0,0 +1,27 @@ +{ + "endpoint": { + "__llmeter_class__": "llmeter.endpoints.bedrock.BedrockConverseStream", + "__llmeter_state__": { + "model_id": "synthetic.interrupted-model-v1", + "endpoint_name": "synthetic-interrupted-endpoint", + "region": "eu-west-1", + "inference_config": null, + "max_attempts": 3, + "ttft_visible_tokens_only": true + } + }, + "output_path": "stale/interrupted/path", + "tokenizer": { + "__llmeter_class__": "llmeter.tokenizers.DummyTokenizer", + "__llmeter_state__": {} + }, + "clients": 4, + "n_requests": 20, + "payload": "stale/interrupted/payload.jsonl", + "run_name": "synthetic-interrupted-test", + "run_description": "This run was interrupted before completion", + "timeout": 120, + "callbacks": null, + "low_memory": false, + "progress_bar_stats": null +} diff --git a/tests/unit/fixtures/result_snapshots/legacy/v0_1_endpoint_type/responses.jsonl b/tests/unit/fixtures/result_snapshots/legacy/v0_1_endpoint_type/responses.jsonl new file mode 100644 index 0000000..8d8dddf --- /dev/null +++ b/tests/unit/fixtures/result_snapshots/legacy/v0_1_endpoint_type/responses.jsonl @@ -0,0 +1,5 @@ +{"response_text": "Consectetur adipiscing elit, sed do eiusmod tempor.", "input_payload": null, "id": "synth-legacy-001", "input_prompt": "What is a load balancer?", "time_to_first_token": 0.612, "time_to_last_token": 1.534, "num_tokens_input": 28, "num_tokens_output": 72, "time_per_output_token": 0.0128, "error": null, "request_time": "2025-01-10T14:00:00Z", "cost_InputTokens": 8.4e-05, "cost_OutputTokens": 0.00108, "cost_total": 0.001164} +{"response_text": "Duis aute irure dolor in reprehenderit in voluptate velit.", "input_payload": null, "id": "synth-legacy-002", "input_prompt": "Explain horizontal scaling.", "time_to_first_token": 0.548, "time_to_last_token": 1.289, "num_tokens_input": 24, "num_tokens_output": 65, "time_per_output_token": 0.0114, "error": null, "request_time": "2025-01-10T14:00:01Z", "cost_InputTokens": 7.2e-05, "cost_OutputTokens": 0.000975, "cost_total": 0.001047} +{"response_text": "Excepteur sint occaecat cupidatat non proident.", "input_payload": null, "id": "synth-legacy-003", "input_prompt": "What is auto-scaling?", "time_to_first_token": 0.701, "time_to_last_token": 1.678, "num_tokens_input": 22, "num_tokens_output": 58, "time_per_output_token": 0.01684, "error": null, "request_time": "2025-01-10T14:00:02Z", "cost_InputTokens": 6.6e-05, "cost_OutputTokens": 0.00087, "cost_total": 0.000936} +{"response_text": "Sunt in culpa qui officia deserunt mollit anim id est laborum.", "input_payload": null, "id": "synth-legacy-004", "input_prompt": "Define microservices architecture.", "time_to_first_token": 0.489, "time_to_last_token": 1.102, "num_tokens_input": 26, "num_tokens_output": 71, "time_per_output_token": 0.00863, "error": null, "request_time": "2025-01-10T14:00:03Z", "cost_InputTokens": 7.8e-05, "cost_OutputTokens": 0.001065, "cost_total": 0.001143} +{"response_text": "Neque porro quisquam est qui dolorem ipsum.", "input_payload": null, "id": "synth-legacy-005", "input_prompt": "What is service mesh?", "time_to_first_token": 0.555, "time_to_last_token": 1.345, "num_tokens_input": 20, "num_tokens_output": 48, "time_per_output_token": 0.01646, "error": null, "request_time": "2025-01-10T14:00:04Z", "cost_InputTokens": 6e-05, "cost_OutputTokens": 0.00072, "cost_total": 0.00078} diff --git a/tests/unit/fixtures/result_snapshots/legacy/v0_1_endpoint_type/run_config.json b/tests/unit/fixtures/result_snapshots/legacy/v0_1_endpoint_type/run_config.json new file mode 100644 index 0000000..9d8ee73 --- /dev/null +++ b/tests/unit/fixtures/result_snapshots/legacy/v0_1_endpoint_type/run_config.json @@ -0,0 +1,20 @@ +{ + "endpoint": { + "endpoint_name": "synthetic-legacy-endpoint", + "model_id": "synthetic.legacy-model-v1", + "provider": "bedrock", + "region": "us-west-2", + "endpoint_type": "BedrockConverse" + }, + "output_path": "stale/legacy/path", + "tokenizer": { + "tokenizer_module": "llmeter" + }, + "clients": 3, + "n_requests": 2, + "payload": "stale/legacy/payload.jsonl", + "run_name": "synthetic-legacy-test", + "run_description": null, + "timeout": 60, + "callbacks": null +} diff --git a/tests/unit/fixtures/result_snapshots/legacy/v0_1_endpoint_type/stats.json b/tests/unit/fixtures/result_snapshots/legacy/v0_1_endpoint_type/stats.json new file mode 100644 index 0000000..50f3609 --- /dev/null +++ b/tests/unit/fixtures/result_snapshots/legacy/v0_1_endpoint_type/stats.json @@ -0,0 +1,41 @@ +{ + "total_requests": 5, + "clients": 3, + "n_requests": 2, + "total_test_time": 5.102, + "model_id": "synthetic.legacy-model-v1", + "output_path": "stale/legacy/path", + "endpoint_name": "synthetic-legacy-endpoint", + "provider": "bedrock", + "run_name": "synthetic-legacy-test", + "run_description": null, + "start_time": "2025-01-10T14:00:00Z", + "first_request_time": "2025-01-10T14:00:00Z", + "last_request_time": "2025-01-10T14:00:04Z", + "end_time": "2025-01-10T14:00:05Z", + "failed_requests": 0, + "failed_requests_rate": 0.0, + "requests_per_minute": 58.80047040376323, + "total_input_tokens": 120, + "total_output_tokens": 314, + "total_cached_input_tokens": 0, + "total_reasoning_output_tokens": 0, + "average_input_tokens_per_minute": 1411.2112896903175, + "average_output_tokens_per_minute": 3692.6695413563307, + "time_to_last_token-average": 1.3896, + "time_to_last_token-p50": 1.345, + "time_to_last_token-p90": 1.7355999999999998, + "time_to_last_token-p99": 1.8133599999999999, + "time_to_first_token-average": 0.581, + "time_to_first_token-p50": 0.555, + "time_to_first_token-p90": 0.7365999999999999, + "time_to_first_token-p99": 0.7846600000000001, + "num_tokens_output-average": 62.8, + "num_tokens_output-p50": 65, + "num_tokens_output-p90": 72.4, + "num_tokens_output-p99": 72.94, + "num_tokens_input-average": 24, + "num_tokens_input-p50": 24, + "num_tokens_input-p90": 28.8, + "num_tokens_input-p99": 29.88 +} diff --git a/tests/unit/fixtures/result_snapshots/legacy/v0_1_endpoint_type/summary.json b/tests/unit/fixtures/result_snapshots/legacy/v0_1_endpoint_type/summary.json new file mode 100644 index 0000000..d46b45d --- /dev/null +++ b/tests/unit/fixtures/result_snapshots/legacy/v0_1_endpoint_type/summary.json @@ -0,0 +1,16 @@ +{ + "total_requests": 5, + "clients": 3, + "n_requests": 2, + "total_test_time": 5.102, + "model_id": "synthetic.legacy-model-v1", + "output_path": "stale/legacy/path", + "endpoint_name": "synthetic-legacy-endpoint", + "provider": "bedrock", + "run_name": "synthetic-legacy-test", + "run_description": null, + "start_time": "2025-01-10T14:00:00Z", + "first_request_time": "2025-01-10T14:00:00Z", + "last_request_time": "2025-01-10T14:00:04Z", + "end_time": "2025-01-10T14:00:05Z" +} diff --git a/tests/unit/fixtures/result_snapshots/legacy/v0_1_interrupted/responses.jsonl b/tests/unit/fixtures/result_snapshots/legacy/v0_1_interrupted/responses.jsonl new file mode 100644 index 0000000..24246b0 --- /dev/null +++ b/tests/unit/fixtures/result_snapshots/legacy/v0_1_interrupted/responses.jsonl @@ -0,0 +1,5 @@ +{"response_text": "Legacy interrupted response 1.", "input_payload": null, "id": "synth-legacy-intr-001", "input_prompt": "Legacy interrupted prompt 1.", "time_to_first_token": 0.39999999999999997, "time_to_last_token": 1.0, "num_tokens_input": 24, "num_tokens_output": 51, "time_per_output_token": 0.01, "error": null, "request_time": "2025-01-05T12:00:01Z"} +{"response_text": "Legacy interrupted response 2.", "input_payload": null, "id": "synth-legacy-intr-002", "input_prompt": "Legacy interrupted prompt 2.", "time_to_first_token": 0.44999999999999996, "time_to_last_token": 1.1, "num_tokens_input": 28, "num_tokens_output": 57, "time_per_output_token": 0.01, "error": null, "request_time": "2025-01-05T12:00:02Z"} +{"response_text": "Legacy interrupted response 3.", "input_payload": null, "id": "synth-legacy-intr-003", "input_prompt": "Legacy interrupted prompt 3.", "time_to_first_token": 0.5, "time_to_last_token": 1.2000000000000002, "num_tokens_input": 32, "num_tokens_output": 63, "time_per_output_token": 0.01, "error": null, "request_time": "2025-01-05T12:00:03Z"} +{"response_text": "Legacy interrupted response 4.", "input_payload": null, "id": "synth-legacy-intr-004", "input_prompt": "Legacy interrupted prompt 4.", "time_to_first_token": 0.55, "time_to_last_token": 1.3, "num_tokens_input": 36, "num_tokens_output": 69, "time_per_output_token": 0.01, "error": null, "request_time": "2025-01-05T12:00:04Z"} +{"response_text": "Legacy interrupted response 5.", "input_payload": null, "id": "synth-legacy-intr-005", "input_prompt": "Legacy interrupted prompt 5.", "time_to_first_token": 0.6, "time_to_last_token": 1.4, "num_tokens_input": 40, "num_tokens_output": 75, "time_per_output_token": 0.01, "error": null, "request_time": "2025-01-05T12:00:05Z"} diff --git a/tests/unit/fixtures/result_snapshots/legacy/v0_1_interrupted/run_config.json b/tests/unit/fixtures/result_snapshots/legacy/v0_1_interrupted/run_config.json new file mode 100644 index 0000000..357eb32 --- /dev/null +++ b/tests/unit/fixtures/result_snapshots/legacy/v0_1_interrupted/run_config.json @@ -0,0 +1,20 @@ +{ + "endpoint": { + "endpoint_name": "synthetic-legacy-intr-endpoint", + "model_id": "synthetic.legacy-interrupted-v1", + "provider": "bedrock", + "region": "us-east-1", + "endpoint_type": "BedrockConverse" + }, + "output_path": "stale/legacy/interrupted/path", + "tokenizer": { + "tokenizer_module": "llmeter" + }, + "clients": 2, + "n_requests": 4, + "payload": "stale/legacy/interrupted/payload.jsonl", + "run_name": "synthetic-legacy-interrupted-test", + "run_description": null, + "timeout": 60, + "callbacks": null +} diff --git a/tests/unit/fixtures/result_snapshots/legacy/v0_1_str_callbacks/responses.jsonl b/tests/unit/fixtures/result_snapshots/legacy/v0_1_str_callbacks/responses.jsonl new file mode 100644 index 0000000..51fcf3f --- /dev/null +++ b/tests/unit/fixtures/result_snapshots/legacy/v0_1_str_callbacks/responses.jsonl @@ -0,0 +1,5 @@ +{"response_text": "Synthetic response number 1 for string callback test.", "input_payload": null, "id": "synth-strcb-001", "input_prompt": "Synthetic prompt 1.", "time_to_first_token": 0.5, "time_to_last_token": 1.2, "num_tokens_input": 35, "num_tokens_output": 60, "num_tokens_input_cached": null, "num_tokens_output_reasoning": null, "time_per_output_token": 0.01, "error": null, "retries": null, "request_time": "2025-01-12T09:00:01Z", "annotations": {}} +{"response_text": "Synthetic response number 2 for string callback test.", "input_payload": null, "id": "synth-strcb-002", "input_prompt": "Synthetic prompt 2.", "time_to_first_token": 0.6000000000000001, "time_to_last_token": 1.4, "num_tokens_input": 40, "num_tokens_output": 70, "num_tokens_input_cached": null, "num_tokens_output_reasoning": null, "time_per_output_token": 0.01, "error": null, "retries": null, "request_time": "2025-01-12T09:00:02Z", "annotations": {}} +{"response_text": "Synthetic response number 3 for string callback test.", "input_payload": null, "id": "synth-strcb-003", "input_prompt": "Synthetic prompt 3.", "time_to_first_token": 0.7000000000000001, "time_to_last_token": 1.6, "num_tokens_input": 45, "num_tokens_output": 80, "num_tokens_input_cached": null, "num_tokens_output_reasoning": null, "time_per_output_token": 0.01, "error": null, "retries": null, "request_time": "2025-01-12T09:00:03Z", "annotations": {}} +{"response_text": "Synthetic response number 4 for string callback test.", "input_payload": null, "id": "synth-strcb-004", "input_prompt": "Synthetic prompt 4.", "time_to_first_token": 0.8, "time_to_last_token": 1.8, "num_tokens_input": 50, "num_tokens_output": 90, "num_tokens_input_cached": null, "num_tokens_output_reasoning": null, "time_per_output_token": 0.01, "error": null, "retries": null, "request_time": "2025-01-12T09:00:04Z", "annotations": {}} +{"response_text": "Synthetic response number 5 for string callback test.", "input_payload": null, "id": "synth-strcb-005", "input_prompt": "Synthetic prompt 5.", "time_to_first_token": 0.9, "time_to_last_token": 2.0, "num_tokens_input": 55, "num_tokens_output": 100, "num_tokens_input_cached": null, "num_tokens_output_reasoning": null, "time_per_output_token": 0.01, "error": null, "retries": null, "request_time": "2025-01-12T09:00:05Z", "annotations": {}} diff --git a/tests/unit/fixtures/result_snapshots/legacy/v0_1_str_callbacks/run_config.json b/tests/unit/fixtures/result_snapshots/legacy/v0_1_str_callbacks/run_config.json new file mode 100644 index 0000000..15aca5e --- /dev/null +++ b/tests/unit/fixtures/result_snapshots/legacy/v0_1_str_callbacks/run_config.json @@ -0,0 +1,25 @@ +{ + "endpoint": { + "endpoint_name": "synthetic-strcb-endpoint", + "model_id": "synthetic.strcb-model-v1", + "provider": "bedrock", + "region": "us-east-1", + "endpoint_type": "BedrockConverseStream" + }, + "output_path": "stale/strcb/path", + "tokenizer": { + "tokenizer_module": "llmeter" + }, + "clients": 1, + "n_requests": 5, + "payload": "stale/strcb/payload.jsonl", + "run_name": "synthetic-strcb-test", + "run_description": null, + "timeout": 60, + "callbacks": [ + "<__main__.SyntheticCallback object at 0x1234567890>", + "" + ], + "low_memory": false, + "progress_bar_stats": null +} diff --git a/tests/unit/fixtures/result_snapshots/legacy/v0_1_str_callbacks/stats.json b/tests/unit/fixtures/result_snapshots/legacy/v0_1_str_callbacks/stats.json new file mode 100644 index 0000000..94eaf21 --- /dev/null +++ b/tests/unit/fixtures/result_snapshots/legacy/v0_1_str_callbacks/stats.json @@ -0,0 +1,41 @@ +{ + "total_requests": 5, + "clients": 1, + "n_requests": 5, + "total_test_time": 6.5, + "model_id": "synthetic.strcb-model-v1", + "output_path": "stale/strcb/path", + "endpoint_name": "synthetic-strcb-endpoint", + "provider": "bedrock", + "run_name": "synthetic-strcb-test", + "run_description": null, + "start_time": "2025-01-12T09:00:01Z", + "first_request_time": "2025-01-12T09:00:01Z", + "last_request_time": "2025-01-12T09:00:05Z", + "end_time": "2025-01-12T09:00:06Z", + "failed_requests": 0, + "failed_requests_rate": 0.0, + "requests_per_minute": 46.15384615384615, + "total_input_tokens": 225, + "total_output_tokens": 400, + "total_cached_input_tokens": 0, + "total_reasoning_output_tokens": 0, + "average_input_tokens_per_minute": 2076.9230769230767, + "average_output_tokens_per_minute": 3692.3076923076924, + "time_to_last_token-average": 1.6, + "time_to_last_token-p50": 1.6, + "time_to_last_token-p90": 2.08, + "time_to_last_token-p99": 2.1879999999999997, + "time_to_first_token-average": 0.7000000000000001, + "time_to_first_token-p50": 0.7000000000000001, + "time_to_first_token-p90": 0.9399999999999998, + "time_to_first_token-p99": 0.9939999999999999, + "num_tokens_output-average": 80, + "num_tokens_output-p50": 80, + "num_tokens_output-p90": 104.0, + "num_tokens_output-p99": 109.4, + "num_tokens_input-average": 45, + "num_tokens_input-p50": 45, + "num_tokens_input-p90": 57.0, + "num_tokens_input-p99": 59.7 +} diff --git a/tests/unit/fixtures/result_snapshots/legacy/v0_1_str_callbacks/summary.json b/tests/unit/fixtures/result_snapshots/legacy/v0_1_str_callbacks/summary.json new file mode 100644 index 0000000..7a0e968 --- /dev/null +++ b/tests/unit/fixtures/result_snapshots/legacy/v0_1_str_callbacks/summary.json @@ -0,0 +1,16 @@ +{ + "total_requests": 5, + "clients": 1, + "n_requests": 5, + "total_test_time": 6.5, + "model_id": "synthetic.strcb-model-v1", + "output_path": "stale/strcb/path", + "endpoint_name": "synthetic-strcb-endpoint", + "provider": "bedrock", + "run_name": "synthetic-strcb-test", + "run_description": null, + "start_time": "2025-01-12T09:00:01Z", + "first_request_time": "2025-01-12T09:00:01Z", + "last_request_time": "2025-01-12T09:00:05Z", + "end_time": "2025-01-12T09:00:06Z" +} diff --git a/tests/unit/fixtures/result_snapshots/load_test/00001-clients/responses.jsonl b/tests/unit/fixtures/result_snapshots/load_test/00001-clients/responses.jsonl new file mode 100644 index 0000000..174b71d --- /dev/null +++ b/tests/unit/fixtures/result_snapshots/load_test/00001-clients/responses.jsonl @@ -0,0 +1,5 @@ +{"response_text": "Load test single-client response 1.", "input_payload": null, "id": "synth-lt1-001", "input_prompt": "Load test prompt 1.", "time_to_first_token": 0.25, "time_to_last_token": 0.7, "num_tokens_input": 22, "num_tokens_output": 40, "num_tokens_input_cached": null, "num_tokens_output_reasoning": null, "time_per_output_token": 0.008, "error": null, "retries": 0, "request_time": "2025-01-22T12:00:01Z", "annotations": {}} +{"response_text": "Load test single-client response 2.", "input_payload": null, "id": "synth-lt1-002", "input_prompt": "Load test prompt 2.", "time_to_first_token": 0.30000000000000004, "time_to_last_token": 0.8, "num_tokens_input": 24, "num_tokens_output": 45, "num_tokens_input_cached": null, "num_tokens_output_reasoning": null, "time_per_output_token": 0.008, "error": null, "retries": 0, "request_time": "2025-01-22T12:00:02Z", "annotations": {}} +{"response_text": "Load test single-client response 3.", "input_payload": null, "id": "synth-lt1-003", "input_prompt": "Load test prompt 3.", "time_to_first_token": 0.35000000000000003, "time_to_last_token": 0.9, "num_tokens_input": 26, "num_tokens_output": 50, "num_tokens_input_cached": null, "num_tokens_output_reasoning": null, "time_per_output_token": 0.008, "error": null, "retries": 0, "request_time": "2025-01-22T12:00:03Z", "annotations": {}} +{"response_text": "Load test single-client response 4.", "input_payload": null, "id": "synth-lt1-004", "input_prompt": "Load test prompt 4.", "time_to_first_token": 0.4, "time_to_last_token": 1.0, "num_tokens_input": 28, "num_tokens_output": 55, "num_tokens_input_cached": null, "num_tokens_output_reasoning": null, "time_per_output_token": 0.008, "error": null, "retries": 0, "request_time": "2025-01-22T12:00:04Z", "annotations": {}} +{"response_text": "Load test single-client response 5.", "input_payload": null, "id": "synth-lt1-005", "input_prompt": "Load test prompt 5.", "time_to_first_token": 0.45, "time_to_last_token": 1.1, "num_tokens_input": 30, "num_tokens_output": 60, "num_tokens_input_cached": null, "num_tokens_output_reasoning": null, "time_per_output_token": 0.008, "error": null, "retries": 0, "request_time": "2025-01-22T12:00:05Z", "annotations": {}} diff --git a/tests/unit/fixtures/result_snapshots/load_test/00001-clients/stats.json b/tests/unit/fixtures/result_snapshots/load_test/00001-clients/stats.json new file mode 100644 index 0000000..77fe968 --- /dev/null +++ b/tests/unit/fixtures/result_snapshots/load_test/00001-clients/stats.json @@ -0,0 +1,41 @@ +{ + "total_requests": 5, + "clients": 1, + "n_requests": 5, + "total_test_time": 5.5, + "model_id": "synthetic.loadtest-model-v1", + "output_path": "stale/loadtest/00001-clients", + "endpoint_name": "synthetic-loadtest-endpoint", + "provider": "bedrock", + "run_name": "00001-clients", + "run_description": null, + "start_time": "2025-01-22T12:00:01Z", + "first_request_time": "2025-01-22T12:00:01Z", + "last_request_time": "2025-01-22T12:00:05Z", + "end_time": "2025-01-22T12:00:05Z", + "failed_requests": 0, + "failed_requests_rate": 0.0, + "requests_per_minute": 54.54545454545455, + "total_input_tokens": 130, + "total_output_tokens": 250, + "total_cached_input_tokens": 0, + "total_reasoning_output_tokens": 0, + "average_input_tokens_per_minute": 1418.1818181818182, + "average_output_tokens_per_minute": 2727.272727272727, + "time_to_last_token-average": 0.9, + "time_to_last_token-p50": 0.9, + "time_to_last_token-p90": 1.1400000000000001, + "time_to_last_token-p99": 1.194, + "time_to_first_token-average": 0.35000000000000003, + "time_to_first_token-p50": 0.35000000000000003, + "time_to_first_token-p90": 0.4699999999999999, + "time_to_first_token-p99": 0.49699999999999994, + "num_tokens_output-average": 50, + "num_tokens_output-p50": 50, + "num_tokens_output-p90": 62.0, + "num_tokens_output-p99": 64.7, + "num_tokens_input-average": 26, + "num_tokens_input-p50": 26, + "num_tokens_input-p90": 30.8, + "num_tokens_input-p99": 31.88 +} diff --git a/tests/unit/fixtures/result_snapshots/load_test/00001-clients/summary.json b/tests/unit/fixtures/result_snapshots/load_test/00001-clients/summary.json new file mode 100644 index 0000000..030ec9d --- /dev/null +++ b/tests/unit/fixtures/result_snapshots/load_test/00001-clients/summary.json @@ -0,0 +1,16 @@ +{ + "total_requests": 5, + "clients": 1, + "n_requests": 5, + "total_test_time": 5.5, + "model_id": "synthetic.loadtest-model-v1", + "output_path": "stale/loadtest/00001-clients", + "endpoint_name": "synthetic-loadtest-endpoint", + "provider": "bedrock", + "run_name": "00001-clients", + "run_description": null, + "start_time": "2025-01-22T12:00:01Z", + "first_request_time": "2025-01-22T12:00:01Z", + "last_request_time": "2025-01-22T12:00:05Z", + "end_time": "2025-01-22T12:00:05Z" +} diff --git a/tests/unit/fixtures/result_snapshots/load_test/00003-clients/responses.jsonl b/tests/unit/fixtures/result_snapshots/load_test/00003-clients/responses.jsonl new file mode 100644 index 0000000..2041ea1 --- /dev/null +++ b/tests/unit/fixtures/result_snapshots/load_test/00003-clients/responses.jsonl @@ -0,0 +1,5 @@ +{"response_text": "Load test triple-client response 1.", "input_payload": null, "id": "synth-lt3-001", "input_prompt": "Load test prompt 1.", "time_to_first_token": 0.43, "time_to_last_token": 1.05, "num_tokens_input": 22, "num_tokens_output": 40, "num_tokens_input_cached": null, "num_tokens_output_reasoning": null, "time_per_output_token": 0.012, "error": null, "retries": 0, "request_time": "2025-01-22T12:05:01Z", "annotations": {}} +{"response_text": "Load test triple-client response 2.", "input_payload": null, "id": "synth-lt3-002", "input_prompt": "Load test prompt 2.", "time_to_first_token": 0.51, "time_to_last_token": 1.2, "num_tokens_input": 24, "num_tokens_output": 45, "num_tokens_input_cached": null, "num_tokens_output_reasoning": null, "time_per_output_token": 0.012, "error": null, "retries": 0, "request_time": "2025-01-22T12:05:02Z", "annotations": {}} +{"response_text": "Load test triple-client response 3.", "input_payload": null, "id": "synth-lt3-003", "input_prompt": "Load test prompt 3.", "time_to_first_token": 0.59, "time_to_last_token": 1.35, "num_tokens_input": 26, "num_tokens_output": 50, "num_tokens_input_cached": null, "num_tokens_output_reasoning": null, "time_per_output_token": 0.012, "error": null, "retries": 0, "request_time": "2025-01-22T12:05:03Z", "annotations": {}} +{"response_text": "Load test triple-client response 4.", "input_payload": null, "id": "synth-lt3-004", "input_prompt": "Load test prompt 4.", "time_to_first_token": 0.6699999999999999, "time_to_last_token": 1.5, "num_tokens_input": 28, "num_tokens_output": 55, "num_tokens_input_cached": null, "num_tokens_output_reasoning": null, "time_per_output_token": 0.012, "error": null, "retries": 0, "request_time": "2025-01-22T12:05:04Z", "annotations": {}} +{"response_text": "Load test triple-client response 5.", "input_payload": null, "id": "synth-lt3-005", "input_prompt": "Load test prompt 5.", "time_to_first_token": 0.75, "time_to_last_token": 1.65, "num_tokens_input": 30, "num_tokens_output": 60, "num_tokens_input_cached": null, "num_tokens_output_reasoning": null, "time_per_output_token": 0.012, "error": null, "retries": 0, "request_time": "2025-01-22T12:05:05Z", "annotations": {}} diff --git a/tests/unit/fixtures/result_snapshots/load_test/00003-clients/stats.json b/tests/unit/fixtures/result_snapshots/load_test/00003-clients/stats.json new file mode 100644 index 0000000..ad2a2c5 --- /dev/null +++ b/tests/unit/fixtures/result_snapshots/load_test/00003-clients/stats.json @@ -0,0 +1,41 @@ +{ + "total_requests": 5, + "clients": 3, + "n_requests": 2, + "total_test_time": 3.2, + "model_id": "synthetic.loadtest-model-v1", + "output_path": "stale/loadtest/00003-clients", + "endpoint_name": "synthetic-loadtest-endpoint", + "provider": "bedrock", + "run_name": "00003-clients", + "run_description": null, + "start_time": "2025-01-22T12:05:01Z", + "first_request_time": "2025-01-22T12:05:01Z", + "last_request_time": "2025-01-22T12:05:05Z", + "end_time": "2025-01-22T12:05:05Z", + "failed_requests": 0, + "failed_requests_rate": 0.0, + "requests_per_minute": 93.75, + "total_input_tokens": 130, + "total_output_tokens": 250, + "total_cached_input_tokens": 0, + "total_reasoning_output_tokens": 0, + "average_input_tokens_per_minute": 2437.5, + "average_output_tokens_per_minute": 4687.5, + "time_to_last_token-average": 1.35, + "time_to_last_token-p50": 1.35, + "time_to_last_token-p90": 1.7099999999999997, + "time_to_last_token-p99": 1.7909999999999997, + "time_to_first_token-average": 0.59, + "time_to_first_token-p50": 0.59, + "time_to_first_token-p90": 0.782, + "time_to_first_token-p99": 0.8252000000000002, + "num_tokens_output-average": 50, + "num_tokens_output-p50": 50, + "num_tokens_output-p90": 62.0, + "num_tokens_output-p99": 64.7, + "num_tokens_input-average": 26, + "num_tokens_input-p50": 26, + "num_tokens_input-p90": 30.8, + "num_tokens_input-p99": 31.88 +} diff --git a/tests/unit/fixtures/result_snapshots/load_test/00003-clients/summary.json b/tests/unit/fixtures/result_snapshots/load_test/00003-clients/summary.json new file mode 100644 index 0000000..06f98af --- /dev/null +++ b/tests/unit/fixtures/result_snapshots/load_test/00003-clients/summary.json @@ -0,0 +1,16 @@ +{ + "total_requests": 5, + "clients": 3, + "n_requests": 2, + "total_test_time": 3.2, + "model_id": "synthetic.loadtest-model-v1", + "output_path": "stale/loadtest/00003-clients", + "endpoint_name": "synthetic-loadtest-endpoint", + "provider": "bedrock", + "run_name": "00003-clients", + "run_description": null, + "start_time": "2025-01-22T12:05:01Z", + "first_request_time": "2025-01-22T12:05:01Z", + "last_request_time": "2025-01-22T12:05:05Z", + "end_time": "2025-01-22T12:05:05Z" +} diff --git a/tests/unit/test_snapshot_load.py b/tests/unit/test_snapshot_load.py new file mode 100644 index 0000000..1cd55c8 --- /dev/null +++ b/tests/unit/test_snapshot_load.py @@ -0,0 +1,354 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Regression tests that load *static, synthetic* result snapshots from disk. + +Unlike the other load tests (which build data in-memory at test time), these load committed +fixture files under ``tests/unit/fixtures/result_snapshots/``. They pin the on-disk de/serialization +contract across the whole loading pipeline - ``Result.load``, ``_RunConfig.load``, +``LoadTestResult.load`` and the endpoint / tokenizer / callback restoration underneath - so format +drift can't silently break loading of previously-saved runs. + +The fixtures are deliberately synthetic (fake model IDs like ``synthetic.test-model-v1``, +lorem-ipsum responses, timestamps in the past) so nobody mistakes them for real benchmark data. +They cover a base "happy path" modern-format snapshot plus variations: legacy formats, an +interrupted run, error/annotation edge cases, and a multi-concurrency load test. + +To regenerate the fixtures (e.g. after an intentional format change), run: + uv run python tests/unit/fixtures/_generate_fixtures.py +""" + +from datetime import datetime, timezone + +import httpx +import pytest + +from llmeter.callbacks.cost.dimensions import InputTokens, OutputTokens +from llmeter.callbacks.cost.model import CostModel +from llmeter.callbacks.mlflow import MlflowCallback +from llmeter.endpoints.bedrock import BedrockConverse, BedrockConverseStream +from llmeter.endpoints.openai import OpenAICompletionStreamEndpoint +from llmeter.experiments import LoadTestResult +from llmeter.results import Result +from llmeter.runner import _RunConfig +from llmeter.tokenizers import DummyTokenizer + + +class TestBaseLoad: + """The modern-format baseline: OpenAI endpoint + CostModel + MlflowCallback.""" + + @pytest.fixture + def base_dir(self, snapshots_dir): + return snapshots_dir / "base" + + def test_result_loads(self, base_dir): + result = Result.load(base_dir) + assert len(result.responses) == 5 + assert result.total_requests == 5 + assert result.clients == 2 + assert result.n_requests == 3 + assert result.model_id == "synthetic.test-model-v1" + assert result.run_name == "synthetic-base-test" + + def test_null_output_path_filled_from_load_path(self, base_dir): + """A null output_path in summary.json must be filled with the real load path. + + This is what makes a relocated/copied result directory reloadable - the on-disk + path won't match wherever the data now lives. + """ + result = Result.load(base_dir) + assert result.output_path is not None + assert str(base_dir) in str(result.output_path) + + def test_timestamps_parsed_as_datetime(self, base_dir): + result = Result.load(base_dir) + assert result.start_time == datetime(2025, 1, 15, 10, 0, 0, tzinfo=timezone.utc) + assert result.end_time == datetime(2025, 1, 15, 10, 0, 4, tzinfo=timezone.utc) + assert isinstance(result.responses[0].request_time, datetime) + + def test_stats_match_frozen_values(self, base_dir): + """Stats loaded from stats.json (no responses) should match the frozen values.""" + result = Result.load(base_dir, load_responses=False) + stats = result.stats + assert stats["total_requests"] == 5 + assert stats["failed_requests"] == 0 + assert stats["total_input_tokens"] == 185 + assert stats["total_output_tokens"] == 386 + assert stats["time_to_first_token-p50"] == pytest.approx(0.415) + assert stats["requests_per_minute"] == pytest.approx(66.35700066357) + + def test_response_optional_token_fields(self, base_dir): + result = Result.load(base_dir) + by_id = {r.id: r for r in result.responses} + assert by_id["synth-base-002"].num_tokens_input_cached == 12 + assert by_id["synth-base-003"].num_tokens_output_reasoning == 8 + + # --- run_config.json / endpoint / tokenizer / callback restoration --- + + @pytest.fixture + def base_config(self, base_dir): + return _RunConfig.load(base_dir) + + def test_endpoint_restored(self, base_config): + ep = base_config._endpoint + assert isinstance(ep, OpenAICompletionStreamEndpoint) + assert ep.model_id == "synthetic.test-model-v1" + assert ep.endpoint_name == "synthetic-openai-stream" + assert ep.provider == "openai" + + def test_endpoint_client_config_restored(self, base_config): + """The OpenAI client's config surface (base_url, retries, timeout) must round-trip.""" + client = base_config._endpoint._client + assert str(client.base_url) == "https://api.synthetic-provider.example.com/v1/" + assert client.max_retries == 3 + assert isinstance(client.timeout, httpx.Timeout) + assert client.timeout.connect == 5.0 + assert client.timeout.read == 60.0 + assert client.timeout.write == 5.0 + assert client.timeout.pool == 5.0 + + def test_tokenizer_restored(self, base_config): + assert isinstance(base_config._tokenizer, DummyTokenizer) + + def test_cost_model_callback_restored(self, base_config): + cost_model = base_config.callbacks[0] + assert isinstance(cost_model, CostModel) + + input_dim = cost_model.request_dims["InputTokens"] + output_dim = cost_model.request_dims["OutputTokens"] + assert isinstance(input_dim, InputTokens) + assert isinstance(output_dim, OutputTokens) + assert input_dim.price_per_million == 3.0 + assert output_dim.price_per_million == 15.0 + assert cost_model.run_dims == {} + + def test_mlflow_callback_restored(self, base_config): + mlflow_cb = base_config.callbacks[1] + assert isinstance(mlflow_cb, MlflowCallback) + assert mlflow_cb.step is None + assert mlflow_cb.nested is True + + +class TestLegacyV01EndpointTypeLoad: + """Legacy v0.1 format: ``endpoint_type`` dispatch + flat cost annotations on responses.""" + + @pytest.fixture + def legacy_dir(self, snapshots_dir): + return snapshots_dir / "legacy" / "v0_1_endpoint_type" + + def test_result_loads(self, legacy_dir): + result = Result.load(legacy_dir) + assert len(result.responses) == 5 + assert result.model_id == "synthetic.legacy-model-v1" + assert result.run_name == "synthetic-legacy-test" + + def test_legacy_cost_keys_collected_into_annotations(self, legacy_dir): + """Flat ``cost_*`` keys that current InvocationResponse doesn't declare go to annotations.""" + result = Result.load(legacy_dir) + first = result.responses[0] + assert first.annotations["cost_InputTokens"] == pytest.approx(0.000084) + assert first.annotations["cost_OutputTokens"] == pytest.approx(0.00108) + assert first.annotations["cost_total"] == pytest.approx(0.001164) + + def test_endpoint_restored_via_endpoint_type(self, legacy_dir): + cfg = _RunConfig.load(legacy_dir) + ep = cfg._endpoint + assert isinstance(ep, BedrockConverse) + assert ep.model_id == "synthetic.legacy-model-v1" + assert ep.region == "us-west-2" + # provider is derived by the constructor, not read from the (legacy) config dict + assert ep.provider == "bedrock" + + def test_tokenizer_restored_via_tokenizer_module(self, legacy_dir): + cfg = _RunConfig.load(legacy_dir) + assert isinstance(cfg._tokenizer, DummyTokenizer) + + +class TestLegacyV01StrCallbacksLoad: + """Legacy v0.1 data where callbacks were serialized as Python repr strings.""" + + @pytest.fixture + def strcb_dir(self, snapshots_dir): + return snapshots_dir / "legacy" / "v0_1_str_callbacks" + + def test_result_loads(self, strcb_dir): + result = Result.load(strcb_dir) + assert len(result.responses) == 5 + assert result.model_id == "synthetic.strcb-model-v1" + + def test_config_loads_without_raising_on_str_callbacks(self, strcb_dir): + """Unrestorable string callbacks must not break run_config loading.""" + cfg = _RunConfig.load(strcb_dir) + assert isinstance(cfg._endpoint, BedrockConverseStream) + # String callbacks are left untouched (not restored to objects, but not fatal either) + assert cfg.callbacks is not None + assert len(cfg.callbacks) == 2 + assert all(isinstance(cb, str) for cb in cfg.callbacks) + + +class TestLegacyV01InterruptedRunLoad: + """A legacy-format interrupted run: no summary.json + top-level endpoint fields. + + Exercises the legacy branch of ``Result._recover_metadata``: older data kept endpoint + fields (including the derived ``provider``) at the top level rather than nested in a + ``__llmeter_state__`` envelope. Recovery must handle it - consistent with the legacy + handling in ``_RunConfig.load`` / ``Endpoint.load`` - and warn that the format is + deprecated. + """ + + @pytest.fixture + def legacy_interrupted_dir(self, snapshots_dir): + return snapshots_dir / "legacy" / "v0_1_interrupted" + + def test_no_summary_file_present(self, legacy_interrupted_dir): + assert not (legacy_interrupted_dir / "summary.json").exists() + + def test_recovers_top_level_endpoint_fields(self, legacy_interrupted_dir): + result = Result.load(legacy_interrupted_dir) + assert result.model_id == "synthetic.legacy-interrupted-v1" + assert result.endpoint_name == "synthetic-legacy-intr-endpoint" + # Legacy configs persist the derived provider at the top level, so it recovers + assert result.provider == "bedrock" + assert result.clients == 2 + assert result.n_requests == 4 + + def test_recovery_warns_legacy_format(self, legacy_interrupted_dir, caplog): + import logging + + with caplog.at_level(logging.WARNING, logger="llmeter.results"): + Result.load(legacy_interrupted_dir) + assert any("legacy data format" in r.message for r in caplog.records) + + +class TestInterruptedRunLoad: + """A snapshot with no summary.json - exercises metadata recovery.""" + + @pytest.fixture + def interrupted_dir(self, snapshots_dir): + return snapshots_dir / "interrupted_run" + + def test_no_summary_file_present(self, interrupted_dir): + assert not (interrupted_dir / "summary.json").exists() + + def test_recovers_responses(self, interrupted_dir): + result = Result.load(interrupted_dir) + assert len(result.responses) == 5 + assert result.total_requests == 5 + + def test_recovers_metadata_from_run_config(self, interrupted_dir): + result = Result.load(interrupted_dir) + assert result.clients == 4 + assert result.n_requests == 20 + assert result.run_name == "synthetic-interrupted-test" + + def test_recovers_endpoint_fields_from_modern_config(self, interrupted_dir): + """Endpoint fields must be read from the modern ``__llmeter_state__`` envelope. + + A *current* interrupted run writes run_config.json via ``dump_object`` (nested + state), so recovery must look inside ``__llmeter_state__`` and not just the top + level. ``provider`` is derived by the endpoint constructor and not persisted, so + it stays ``None`` on this lightweight recovery path (no endpoint instantiation). + """ + result = Result.load(interrupted_dir) + assert result.model_id == "synthetic.interrupted-model-v1" + assert result.endpoint_name == "synthetic-interrupted-endpoint" + assert result.provider is None + + def test_derives_timestamps_from_responses(self, interrupted_dir): + result = Result.load(interrupted_dir) + # request_time runs 16:30:02 .. 16:30:10 (i*2 for i in 1..5) + assert result.start_time == datetime( + 2025, 1, 20, 16, 30, 2, tzinfo=timezone.utc + ) + assert result.end_time == datetime(2025, 1, 20, 16, 30, 10, tzinfo=timezone.utc) + + def test_modern_config_recovery_does_not_warn_legacy(self, interrupted_dir, caplog): + """The modern-format recovery path must not emit the legacy-format warning.""" + import logging + + with caplog.at_level(logging.WARNING, logger="llmeter.results"): + Result.load(interrupted_dir) + assert not any("legacy data format" in r.message for r in caplog.records) + + def test_computes_stats_from_responses(self, interrupted_dir): + result = Result.load(interrupted_dir) + stats = result.stats + assert stats["failed_requests"] == 0 + assert "time_to_first_token-p50" in stats + + +class TestErrorsAndAnnotationsLoad: + """Responses carrying errors, retries, cached tokens, and custom annotations.""" + + @pytest.fixture + def errors_dir(self, snapshots_dir): + return snapshots_dir / "errors_and_annotations" + + def test_result_loads(self, errors_dir): + result = Result.load(errors_dir) + assert len(result.responses) == 5 + + def test_errors_preserved(self, errors_dir): + result = Result.load(errors_dir) + by_id = {r.id: r for r in result.responses} + assert by_id["synth-err-002"].error == "ThrottlingException: Rate exceeded" + assert by_id["synth-err-004"].error == "ReadTimeoutError: timed out after 60s" + assert by_id["synth-err-001"].error is None + + def test_failed_requests_counted(self, errors_dir): + result = Result.load(errors_dir) + assert result.stats["failed_requests"] == 2 + + def test_retries_preserved(self, errors_dir): + result = Result.load(errors_dir) + by_id = {r.id: r for r in result.responses} + assert by_id["synth-err-002"].retries == 2 + assert by_id["synth-err-005"].retries == 1 + assert by_id["synth-err-001"].retries == 0 + + def test_custom_annotations_preserved(self, errors_dir): + result = Result.load(errors_dir) + by_id = {r.id: r for r in result.responses} + assert by_id["synth-err-001"].annotations == { + "custom_metric": 42.5, + "experiment_tag": "baseline", + } + assert by_id["synth-err-005"].annotations["cache_hit_ratio"] == pytest.approx( + 0.667 + ) + + def test_cached_input_tokens_preserved(self, errors_dir): + result = Result.load(errors_dir) + by_id = {r.id: r for r in result.responses} + assert by_id["synth-err-005"].num_tokens_input_cached == 30 + + +class TestLoadTestResultLoad: + """A multi-concurrency snapshot loaded via LoadTestResult.load.""" + + @pytest.fixture + def load_test_dir(self, snapshots_dir): + return snapshots_dir / "load_test" + + def test_loads_all_concurrency_levels(self, load_test_dir): + lt = LoadTestResult.load(load_test_dir) + assert set(lt.results.keys()) == {1, 3} + + def test_sub_results_have_correct_data(self, load_test_dir): + lt = LoadTestResult.load(load_test_dir) + assert len(lt.results[1].responses) == 5 + assert len(lt.results[3].responses) == 5 + assert lt.results[1].clients == 1 + assert lt.results[3].clients == 3 + + def test_sub_result_stats_populated(self, load_test_dir): + lt = LoadTestResult.load(load_test_dir) + for clients in (1, 3): + stats = lt.results[clients].stats + assert stats["total_requests"] == 5 + assert "time_to_first_token-p50" in stats + + def test_load_without_responses(self, load_test_dir): + lt = LoadTestResult.load(load_test_dir, load_responses=False) + assert lt.results[1].responses == [] + assert lt.results[1].stats["total_requests"] == 5 From 14b3b12dc8bcd580039b4a044d220809d7dfcc94 Mon Sep 17 00:00:00 2001 From: Alex Thewsey Date: Thu, 30 Jul 2026 20:57:21 +0800 Subject: [PATCH 20/20] test(cost): Simplify cost dim test setups No need for us to manually key in these dictionary keys - just spread --- tests/unit/callbacks/cost/test_dimensions.py | 12 +++--------- 1 file changed, 3 insertions(+), 9 deletions(-) diff --git a/tests/unit/callbacks/cost/test_dimensions.py b/tests/unit/callbacks/cost/test_dimensions.py index 425de3f..e60ce3a 100644 --- a/tests/unit/callbacks/cost/test_dimensions.py +++ b/tests/unit/callbacks/cost/test_dimensions.py @@ -33,9 +33,7 @@ async def test_cost_per_input_token(): "price_per_million": 30, "granularity": 10, } - dim_valid = InputTokens( - price_per_million=spec["price_per_million"], granularity=spec["granularity"] - ) + dim_valid = InputTokens(**spec) success_response = InvocationResponse(response_text="hi", num_tokens_input=199999) assert await dim_valid.calculate(success_response) == 6 @@ -60,9 +58,7 @@ async def test_cost_per_output_token(): "price_per_million": 40, "granularity": 10, } - dim_valid = OutputTokens( - price_per_million=spec["price_per_million"], granularity=spec["granularity"] - ) + dim_valid = OutputTokens(**spec) success_response = InvocationResponse(response_text="hi", num_tokens_output=199999) assert await dim_valid.calculate(success_response) == 8 @@ -87,9 +83,7 @@ async def test_cost_per_hour(): "price_per_hour": 30, "granularity_secs": 60, } - dim_valid = EndpointTime( - price_per_hour=spec["price_per_hour"], granularity_secs=spec["granularity_secs"] - ) + dim_valid = EndpointTime(**spec) result = Result( [],