Address Lora shortcomings
Merged
YU
5mo ago+277+277+012345678901234567890123456789-54-54-01234567890123456789
This pull request significantly improves the safety, correctness, and memory management of LoRA adapter handling in ONNX Runtime, especially around Python bindings and adapter file export/import. The main focus is on ensuring strong exception safety, preventing use-after-free bugs by improving object lifetimes, and rejecting unsupported tensor types during export. Additionally, comprehensive regression tests are added to guard against these issues.
Key changes include:
Exception Safety & Parameter Handling
- Refactored
LoraAdapter::LoadandMemoryMapto provide a strong exception guarantee: all potentially-throwing operations are performed using local variables before committing to the object's state, ensuring no partial updates occur on failure. The newBuildParamsValuesmethod builds the parameter map without side effects, replacing the oldInitializeParamsValues. [1] [2] [3] [4] [5]
Python Bindings & Memory Management
- Improved the Python adapter format bindings so that every
OrtValuereturned from adapter parameter getters is pinned to its owning C++ adapter object via pybind11'skeep_alivemechanism. This prevents use-after-free errors if the parentAdapterFormatobject is dropped while references to its parameters remain. The getter now builds the parameter dictionary on demand and avoids reference cycles that would leak memory. [1] [2] [3]
Adapter Export Robustness
- Enhanced the adapter export logic to reject string tensors, preventing the leaking of memory addresses and creation of unloadable adapter files. The export path now builds the adapter image entirely in memory before writing to disk, ensuring no partial files are left behind on error.
Clean-up & Consistency
- Simplified the construction and usage of the
PyAdapterFormatReaderWriterclass, ensuring that its internal state is only populated as appropriate for read or write operations, and removed unnecessary parameter passing. [1] [2] - Minor cleanup in property definitions and comments for clarity and maintainability.
Regression Tests
- Added thorough regression tests to verify that adapter parameter lifetimes are managed correctly and that exporting string tensors is properly rejected, with checks to ensure no files are created on failure.
These changes collectively make adapter handling safer and more robust, especially when interacting with Python, and add critical safeguards against subtle memory and serialization bugs.