weave
weave copied to clipboard
feat: Pluggable object serialization
This enables users to register Serializers with Weave, like so:
import faiss
from weave.trace import serializer
def save_instance(obj: faiss.Index, artifact, name: str) -> None:
with artifact.writeable_file_path(f"{name}.faissindex") as write_path:
faiss.write_index(obj, write_path)
def load_instance(
artifact,
name: str,
) -> faiss.Index:
return faiss.read_index(artifact.path(f"{name}.faissindex"))
serializer.register_serializer(faiss.Index, save_instance, load_instance)
After the register_serializer
call, if the user tries to save an object that has a faiss.Index
attribute, Weave will store the full FAISS index instead of just a repr string.
This replaces the prior mechanism that relied on the older weave.Type. We were only relying on that for Op saving. The new implementation is fully compatible and generates objects with the same digests.
We will also save the load_instance method as an op, and add a reference to the load op from the saved object, so that the object can be correctly deserialized in a Python runtime that does not have the serializer registered.
So the effect is that objects with custom typed attributes are fully relocatable.
This PR also improves op serialization, notably it makes it so that ops that classes refered to by ops are captured correctly.
TODO:
- [ ] Fix anonymous op issues, they are now annotated with @weave.op so I need to fix digests
- [ ] Consider if we want
<module>.<name>
in the encoded weave type string, or just<name>
- [ ] Register Serializers by ID string instead of class, so we don't need to import libraries to register serializers.