feat: Dynamic callback serialization via _callback_type marker#54
Open
acere wants to merge 9 commits intoawslabs:mainfrom
Open
feat: Dynamic callback serialization via _callback_type marker#54acere wants to merge 9 commits intoawslabs:mainfrom
acere wants to merge 9 commits intoawslabs:mainfrom
Conversation
Addresses awslabs#20 by implementing binary-safe serialization for payloads and results containing images. This prevents double base64 encoding and significantly reduces memory usage and serialization time. Changes: - Add binary serialization support to prompt_utils and results - Update endpoints to use binary-safe serialization - Add property-based tests for serialization correctness - Update integration tests to verify image handling
…s endpoints - Add multi-modal content handling (images, videos, audio, documents) to BedrockConverse, OpenAI, and SageMaker endpoints - Implement automatic format detection using puremagic with fallback to file extensions - Add multimodal utility functions for format conversion and content serialization - Support both file paths and raw bytes for multi-modal content - Add endpoint-specific format string handling (Bedrock short format, OpenAI MIME types) - Implement comprehensive unit tests for multi-modal serialization and properties across all endpoints - Add detailed README documentation with usage examples and security warnings for format detection - Fix Path serialization in JSONableBase to handle os.PathLike objects - Update pyproject.toml with optional multimodal extra for puremagic dependency - Improve integration tests with multi-modal payload examples - Enhance prompt utilities with multi-modal content handling
… compatibility - Replace built-in open() calls with UPath.open() across all file operations - Add UPath import to cost model for consistent path handling - Update cost model save_to_file() to create parent directories automatically - Standardize file reading in prompt_utils, results, runner, and tokenizers modules - Improves cross-platform file system support and enables cloud storage integration
Result.stats included raw datetime objects (start_time, end_time) from to_dict(), causing TypeError when users called json.dumps() without a custom serializer. Now to_dict() converts datetime fields via utc_datetime_serializer so stats is always directly serializable.
…path helper Replace scattered `os.PathLike | str` type annotations with the proper UPath type aliases from `upath.types`: - `ReadablePathLike` for parameters used to load/read data - `WritablePathLike` for parameters used to save/write data Add `ensure_path()` utility to `llmeter.utils` to centralize the `Path(x)` normalization boilerplate that was duplicated at the top of nearly every function accepting a path argument. The helper handles None passthrough and uses a lazy UPath import. Runtime `isinstance(obj, os.PathLike)` checks in serialization helpers are left unchanged since TypeAliases cannot be used for runtime checks.
…nstances Cloud-backed UPath instances (e.g. S3Path) do not implement os.PathLike, so isinstance checks against os.PathLike alone would miss them. This caused: - Serialization: cloud UPaths skipped the path branch in JSON serializers - runner.py: cloud UPath payloads not recognized as path references, leading to unnecessary re-saving or failure to load from path Fix by checking isinstance(obj, (os.PathLike, Path)) where Path is UPath, which catches both plain pathlib.Path (via os.PathLike) and cloud UPaths (via UPath). Serialization keeps .as_posix() for cross-platform safety.
Rationalize scattered serialization utilities into a single unified module: - Create llmeter/json_utils.py with LLMeterEncoder (handles bytes, datetime, date, time, PathLike, to_dict() objects, str() fallback) and llmeter_bytes_decoder (restores __llmeter_bytes__ markers to bytes). - Remove redundant encoders: LLMeterBytesEncoder (prompt_utils), InvocationResponseEncoder (results), utc_datetime_serializer (results), and inline _default_serializer lambdas (runner, endpoints/base). - Slim down callbacks/cost/serde.py to cost-specific helpers (JSONableBase, ISerializable, from_dict_with_class, from_dict_with_class_map, to_dict_recursive_generic). Update to Python 3.10 typing (dict/type builtins instead of Dict/Type). - Standardize all to_json() methods to default to cls=LLMeterEncoder via kwargs.setdefault(), ensuring consistent encoding across InvocationResponse, Result, and JSONableBase. - Remove serializer/deserializer/cls customization params from save_payloads, load_payloads, _load_data_file — hardcode LLMeterEncoder and llmeter_bytes_decoder since custom encoders produce files that can't be loaded back without metadata. - LLMeterEncoder.default() delegates to to_dict() for objects that implement it, enabling json.dump(self, f, cls=LLMeterEncoder) without manual to_dict() calls (used in Endpoint.save). - Convert all changed files to relative imports, run ruff check + format + import sorting. - Clean up llmeter/utils.py: move upath imports to top level (it's a hard dependency), remove unnecessary from __future__ import annotations. - Add docs/reference/json_utils.md and mkdocs.yml nav entry. - Add property-based tests for datetime, date, time, PathLike, and to_dict() encoding (TestDatetimeSerializationProperties, TestPathSerializationProperties, TestToDictSerializationProperties). - Update existing tests to use LLMeterEncoder/llmeter_bytes_decoder from llmeter.json_utils instead of old aliases. All 581 unit tests pass.
Move path and datetime string conversion out of to_dict() and to_dict_recursive_generic() — these are Python dict builders, not serializers. Type coercion to strings is now exclusively handled by LLMeterEncoder at JSON serialization time. - Endpoint.to_dict(): remove PathLike → as_posix() coercion, simplify to a dict comprehension. Remove unused os import. - to_dict_recursive_generic(): remove PathLike → as_posix() and datetime/date/time → isoformat() coercions. Keep structural recursion (nested dicts, lists, to_dict() delegation). Remove unused os, datetime imports.
Implement to_dict/from_dict on Callback base class using a "module:ClassName" type marker (_callback_type) for dynamic dispatch. Any Callback subclass — built-in or third-party — can now round-trip through JSON without a hardcoded registry. - Callback.to_dict() emits _callback_type with fully-qualified class path - Callback.from_dict() dynamically imports and instantiates the correct subclass, delegating to overridden from_dict when present - CostModel.to_dict() injects _callback_type alongside existing _type - CostModel.from_dict() strips both markers before construction - MlflowCallback gains to_dict/from_dict (replaces NotImplementedError stubs) - _RunConfig.save()/load() now serialize and restore callbacks - Replace old Callback base tests with serialization round-trip tests
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Implements dynamic serialization/deserialization for Callback objects using a
_callback_typemarker (module:ClassNameformat), enabling any Callback subclass — built-in or third-party — to round-trip through JSON without a hardcoded registry.Closes #53
Changes
Callbackbase class (callbacks/base.py): Addedto_dict(),from_dict(),to_json(),from_json()methods.to_dict()emits a_callback_typekey with the fully-qualified class path.from_dict()dynamically imports and instantiates the correct subclass, delegating to overriddenfrom_dict()when present (e.g.CostModel). Replaced thesave_to_file()/load_from_file()NotImplementedErrorstubs with working implementations.CostModel(callbacks/cost/model.py):to_dict()now injects_callback_typealongside the existing_typefromJSONableBase.from_dict()strips both markers before delegating to the constructor.MlflowCallback(callbacks/mlflow.py): Addedto_dict()/from_dict()overrides, replacing the oldNotImplementedErrorstubs._RunConfig(runner.py):save()now serializes callbacks viacb.to_dict().load()reconstructs them viaCallback.from_dict().Tests: Rewrote
tests/unit/callbacks/test_base.pywith round-trip serialization tests. Updatedtest_cost_model_serializationto expect the new_callback_typekey.Design
The approach mirrors the existing
Endpointserialization pattern (endpoint_typefield +importlib), but uses fully-qualifiedmodule:ClassNamepaths instead of bare class names. This means third-party callbacks from external packages work automatically — no central registry needed.