feat: Dynamic callback serialization via __llmeter_class__ marker#54
feat: Dynamic callback serialization via __llmeter_class__ marker#54acere wants to merge 16 commits into
Conversation
1a41bfe to
8c07b89
Compare
athewsey
left a comment
athewsey
left a comment
There was a problem hiding this comment.
First a minor/easy point: Since we standardized on __llmeter_bytes__ bytes serialization in payloads, should we follow that convention and always use special fields like __llmeter_{...}__, rather than e.g. _callback_type as used here?
But second and more holistically: I'm not a fan of implementing this without reconciling & aligning it with how cost model serialization should work... Can we really not share a common mechanism here, instead of rebuilding it at multiple levels of the hierarchy?
The cost module JSONableBase.from_dict implementation didn't promote unsupervised recursion, but did provide enough flexibility for classes like CostModel to choose to also instantiate Python sub-objects while loading from config JSON.
I see valid arguments between usability & security for 1/ allowing open-ended class import (like this PR) versus 2/ defining a default allowlist of target classes and requiring users to explicitly specify extras (like JSONableBase tried)... But if the outermost class anyway has control over init proceeds and whether sub-modules of config should be allowed to create more python objects or not - wouldn't it be much simpler to just e.g. enable an __llmeter_class__ hook in a renamed version of llmeter_bytes_decoder (llmeter_default_deserializer?)
8c07b89 to
1ede455
Compare
acere
changed the title
9dee344 to
de9f263
Compare
|
@athewsey this commit should address most of the concerns on serialization, and for all callables, not only callbacks. |
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)
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.
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.
- 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.
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.
- 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.
ae42e23 to
358f35d
Compare
Some unit tests were still referencing old json_utils module
athewsey
left a comment
athewsey
left a comment
There was a problem hiding this comment.
It looks like Serializable itself doesn't actually handle datetime or bytes properties nicely, which seems like a surprising gap?
_serialize_valuedoesn't specifically detect datetimes, so they just fall back tostr(...)which produces non-ISO strings that stay as strings when the object is re-loaded- Likewise there's no explicit handling for bytes, so those become the repr string e.g.
"b''\\xff\\xd8\\xff\\xe0'"and stay that way after loading.
…alue 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.
- Optional[X] → X | None - Sequence/Callable/Iterator from collections.abc instead of typing
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.
|
That should be fixed, and the serialization should now be consistent across modules. |
| class IRequestCostDimension(ISerializable): | ||
| """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. | ||
| """ | ||
| class IRequestCostDimension(Protocol): | ||
| """Interface for one dimension of a per-request cost model.""" |
There was a problem hiding this comment.
I've noticed that Kiro has a real propensity to strip out / cut down comments when doing significant file re-writes... Which makes it a bit difficult to guess how many of the docstring updates in these files you're actually invested in vs was just automatic and haven't checked?
Not specifically this one but I think several of the docstring comments affected by these changes might've been a bit more meaningful/useful before than after?
There was a problem hiding this comment.
Kiro did, Claude undid.
| 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) | ||
| case {"__llmeter_bytes__": b64} if len(val) == 1: | ||
| return base64.b64decode(b64) |
There was a problem hiding this comment.
By choosing to de/serialize datetimes directly as ISO strings rather than a guarded object e.g. {"__llmeter_datetime__": "..."}, we prevent the deserializer from working nicely with any datetime-like strings that are supposed to stay as strings on load.
I'm not aware of any real cases of this so far, and don't expect we'd introduce any in LLMeter itself, but could be a potential risk if we're using the same de/ser for objects not fully under our control e.g. payloads?
So maybe this is fine to leave as-is as produces a nicer (and more backward-compatible) format, but just wanted to make sure we'd thought through the trade-off here.
There was a problem hiding this comment.
yeah, i've been going back and forth on how to deal with that. From the original sin of JSON not supporting datetime as native type, i think the class argument type introspection to be the escape hatch for this issue. As long as the callback property types the datetimes fields, the serializer should be able to handle the round trip without issues.
| :func:`dump_object` / :func:`load_object`. | ||
| """ | ||
|
|
||
| def __getstate__(self) -> dict: |
There was a problem hiding this comment.
By re-using the same protocols as Pickle for this, we affect Serializable subclasses' pickle behavior in ways that seem mostly neutral or negative:
- We narrow what pickle natively supports (e.g.
datefields or any non-allow-listed types cause failure) - We replace pickle's native state handling by funneling our JSON through it
- We make
__init__run on unpickle, which is non-standard behaviour - We lose shared-reference identity: Standard pickle is capable of tracking when two fields point to the same object, whereas our system is not
Although I don't think we're likely to promote using pickle with LLMeter directly, I do wonder if there could be potential impacts because multiprocessing uses pickle and deepcopy uses the same semantics.
Might it be better to use a different protocol and avoid interfering with pickle behaviour here?
There was a problem hiding this comment.
ok, i don't see any harm in using a different protocol. suggestion for the names?
There was a problem hiding this comment.
I've just pushed an update proposing _get_llmeter_state and _set_llmeter_state? but 🤷♂️ not a super strong opinion.
…ocstrings 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.
Add tests to explicitly validate CostModel's save_to_file round trip (previously only the dump & load APIs were checked).
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.
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.
|
Per 5a32dd4, I noticed a couple of issues loading older data with our new code:
|
docstrings, typings, uguide typos.
2 participants
Summary
Implements dynamic serialization/deserialization for Callback objects using an
__llmeter_class__marker (module:ClassNameformat), enabling any Callback subclass — built-in or third-party — to round-trip through JSON without a hardcoded registry.Closes #53
Changes from previous revision
Addresses all review feedback from the first submission:
Naming convention: Replaced
_callback_typewith__llmeter_class__, following the established__llmeter_*__double-underscore convention (consistent with__llmeter_bytes__).Unified serialization: Instead of building a parallel mechanism in
Callback, the dynamic class dispatch is integrated intojson_utils.pyvia a newllmeter_default_deserializer()that handles both__llmeter_bytes__and__llmeter_class__markers. This is a single mechanism shared across the codebase.CostModel compatibility:
CostModelparticipates through the same__llmeter_class__mechanism while preserving its existing_type+JSONableBasepattern for dimension class-map dispatch.Detailed changes
json_utils.py: Added_resolve_class()helper andllmeter_default_deserializer()— a superset ofllmeter_bytes_decoderthat also handles__llmeter_class__markers.Callbackbase class (callbacks/base.py): ReplacedNotImplementedErrorstubs with workingto_dict(),from_dict(),save_to_file(),load_from_file().to_dict()emits__llmeter_class__with the fully-qualified class path.from_dict()dynamically imports and dispatches to the correct subclass.CostModel(callbacks/cost/model.py):to_dict()injects__llmeter_class__alongside the existing_type.from_dict()strips the marker before delegating toJSONableBase.MlflowCallback(callbacks/mlflow.py): ReplacedNotImplementedErrorstubs withto_dict()/from_dict()overrides. Removed unused imports._RunConfig(runner.py):save()serializes callbacks viacb.to_dict().load()reconstructs them viaCallback.from_dict().Tests: Rewrote
tests/unit/callbacks/test_base.pywith round-trip serialization tests. Added MlflowCallback serialization unit tests. Added integration tests using real mlflow library. Updatedtest_cost_model_serializationto expect__llmeter_class__.