Address Lora shortcomings

onnxruntime/core/session/lora_adapters.cc

@@ -37,17 +37,35 @@ void LoraAdapter::Load(const std::filesystem::path& file_path) {
 }
 
 void LoraAdapter::Load(std::vector<uint8_t> buffer) {
-  adapter_ = adapters::utils::ValidateAndGetAdapterFromBytes(buffer);
+  // Build everything that can throw against locals first, so a malformed adapter
+  // leaves *this unchanged (strong exception guarantee). Only after the new
+  // parameter map has been fully constructed do we release the previous backing
+  // storage and swap in the new state. ValidateAndGetAdapterFromBytes returns a
+  // pointer that aliases `buffer`'s data; std::vector's move preserves the data
+  // pointer, so the alias remains valid after we move `buffer` into buffer_.
+  const auto* new_adapter = adapters::utils::ValidateAndGetAdapterFromBytes(buffer);
+  auto new_params_values = BuildParamsValues(new_adapter);
+
+  // Commit. None of the operations below can throw.
   buffer_.emplace<BufferHolder>(std::move(buffer));
-  InitializeParamsValues();
+  adapter_ = new_adapter;
+  params_values_ = std::move(new_params_values);
 }
 
 void LoraAdapter::MemoryMap(const std::filesystem::path& file_path) {
   auto [mapped_memory, file_size] = adapters::utils::MemoryMapAdapterFile(file_path);
   auto u8_span = ReinterpretAsSpan<const uint8_t>(gsl::make_span(mapped_memory.get(), file_size));
-  adapter_ = adapters::utils::ValidateAndGetAdapterFromBytes(u8_span);
+
+  // Build everything that can throw against locals first; see Load() above.
+  // Moving Env::MappedMemoryPtr preserves the underlying pointer, so OrtValues
+  // built against `mapped_memory.get()` remain valid after the move into buffer_.
+  const auto* new_adapter = adapters::utils::ValidateAndGetAdapterFromBytes(u8_span);
+  auto new_params_values = BuildParamsValues(new_adapter);
+
+  // Commit. None of the operations below can throw.
   buffer_.emplace<MemMapHolder>(std::move(mapped_memory), file_size);
-  InitializeParamsValues();
+  adapter_ = new_adapter;
+  params_values_ = std::move(new_params_values);
 }
 
 static std::unique_ptr<IDataTransfer> GetDataTransfer(const OrtMemoryInfo& mem_info) {
@@ -82,10 +100,9 @@ static Status CreateOrtValueOnDevice(const OrtValue& ort_value_mapped,
   return Status::OK();
 }
 
-void LoraAdapter::InitializeParamsValues() {
-  if (adapter_ == nullptr) {
-    ORT_THROW("Adapter is not loaded yet.");
-  }
+std::unordered_map<std::string, LoraAdapter::Param>
+LoraAdapter::BuildParamsValues(const adapters::Adapter* adapter) const {
+  ORT_ENFORCE(adapter != nullptr, "Adapter is not loaded yet.");
 
   std::unique_ptr<IDataTransfer> data_transfer;
   if (device_allocator_) {
@@ -95,7 +112,7 @@ void LoraAdapter::InitializeParamsValues() {
     }
   }
 
-  const auto* params = adapter_->parameters();
+  const auto* params = adapter->parameters();
   ORT_ENFORCE(params != nullptr, "Params absent");
   std::unordered_map<std::string, Param> params_values;
   params_values.reserve(params->size());
@@ -117,7 +134,7 @@ void LoraAdapter::InitializeParamsValues() {
     }
   }
 
-  params_values_.swap(params_values);
+  return params_values;
 }
 
 }  // namespace lora

onnxruntime/core/session/lora_adapters.h

@@ -151,16 +151,26 @@ class LoraAdapter {
   }
 
  private:
-  void InitializeParamsValues();
+  // Build a fresh parameter map from the supplied adapter pointer (which must alias
+  // memory the caller is keeping alive). Throws on validation failure. Does not
+  // touch this->params_values_ / buffer_ / adapter_; callers are responsible for
+  // committing the result only after they have also committed the new backing
+  // storage. Keeping this function free of side effects on *this is what gives
+  // Load()/MemoryMap() the strong exception guarantee.
+  std::unordered_map<std::string, Param> BuildParamsValues(const adapters::Adapter* adapter) const;
 
   struct BufferHolder {
-    explicit BufferHolder(std::vector<uint8_t> buffer) : buffer_(std::move(buffer)) {}
+    explicit BufferHolder(std::vector<uint8_t> buffer) noexcept : buffer_(std::move(buffer)) {}
+    BufferHolder(BufferHolder&&) noexcept = default;
+    BufferHolder& operator=(BufferHolder&&) noexcept = default;
     std::vector<uint8_t> buffer_;
   };
 
   struct MemMapHolder {
-    MemMapHolder(Env::MappedMemoryPtr mapped_memory, size_t file_size)
+    MemMapHolder(Env::MappedMemoryPtr mapped_memory, size_t file_size) noexcept
         : mapped_memory_(std::move(mapped_memory)), file_size_(file_size) {}
+    MemMapHolder(MemMapHolder&&) noexcept = default;
+    MemMapHolder& operator=(MemMapHolder&&) noexcept = default;
     Env::MappedMemoryPtr mapped_memory_;
     size_t file_size_;
   };

onnxruntime/python/onnxruntime_inference_collection.py

@@ -104,9 +104,15 @@ def get_model_version(self) -> int:
         return self._adapter.model_version
 
     def set_parameters(self, params: dict[str, OrtValue]) -> None:
+        """Set adapter parameters for export."""
         self._adapter.parameters = {k: v._ortvalue for k, v in params.items()}
 
     def get_parameters(self) -> dict[str, OrtValue]:
+        """Get adapter parameters as a dict of name -> OrtValue.
+
+        On read instances, the returned OrtValues are zero-copy views; the
+        backing memory stays alive as long as any returned OrtValue is referenced.
+        """
         return {k: OrtValue(v) for k, v in self._adapter.parameters.items()}
 
 

onnxruntime/python/onnxruntime_pybind_lora.cc

@@ -29,29 +29,39 @@
 namespace {
 /// <summary>
 /// Class that supports writing and reading adapters
-/// in innxruntime format
+/// in onnxruntime format.
+///
+/// Design: Single-source architecture using py::capsule for memory ownership.
+///
+/// `parameters_` is always the single authoritative dict of name -> OrtValue.
+/// On the read path, OrtValues are zero-copy views into a LoraAdapter's buffer.
+/// The LoraAdapter is owned by a py::capsule that is pinned as a patient on each
+/// OrtValue via pybind11's add_patient. The object graph is:
+///
+///   PyAdapterFormatReaderWriter -> parameters_ dict -> OrtValues -> capsule -> LoraAdapter
+///
+/// No edge points back to PyAdapterFormatReaderWriter, so there is no reference
+/// cycle (pybind11 instances are not GC-traversable, so a cycle would leak).
+/// The capsule keeps the backing memory alive as long as any OrtValue referencing
+/// it is alive.
+///
+/// The setter replaces parameters_ freely. Old OrtValues (and their capsule
+/// patient) remain alive as long as Python references to them exist.
 /// </summary>
 struct PyAdapterFormatReaderWriter {
   PyAdapterFormatReaderWriter() = default;
   PyAdapterFormatReaderWriter(int format_version, int adapter_version,
                               int model_version,
-                              lora::LoraAdapter&& loaded_adapter,
-                              py::dict&& params)
+                              py::dict parameters)
       : format_version_(format_version),
         adapter_version_(adapter_version),
         model_version_(model_version),
-        loaded_adater_(std::move(loaded_adapter)),
-        parameters_(std::move(params)) {}
+        parameters_(std::move(parameters)) {}
 
   int format_version_{adapters::kAdapterFormatVersion};
   int adapter_version_{0};
   int model_version_{0};
-  // This container is used when reading the the file so
-  // OrtValue objects can be backed by it. Not exposed to Python
-  std::optional<lora::LoraAdapter> loaded_adater_;
-  // This is a dictionary of string -> OrtValue
-  // this is populated directly on write and
-  // built on top of the loaded_adapter on read
+  // Single source of truth: name -> OrtValue (views or owning).
   py::dict parameters_;
 };
 
@@ -65,11 +75,6 @@
           "format_version",
           [](const PyAdapterFormatReaderWriter* reader_writer) -> int { return reader_writer->format_version_; },
           R"pbdoc("Enables user to read format version stored in the file")pbdoc")
-      .def_property(
-          "adapter_version",
-          [](const PyAdapterFormatReaderWriter* reader_writer) -> int { return reader_writer->adapter_version_; },
-          [](PyAdapterFormatReaderWriter* reader_writer, int adapter_version) -> void { reader_writer->adapter_version_ = adapter_version; },
-          R"pbdoc("Enables user to read format version stored in the file")pbdoc")
       .def_property(
           "adapter_version",
           [](const PyAdapterFormatReaderWriter* reader_writer) -> int { return reader_writer->adapter_version_; },
@@ -82,35 +87,66 @@
           R"pbdoc("Enables user to read/write model version this adapter was created for")pbdoc")
       .def_property(
           "parameters",
-          [](const PyAdapterFormatReaderWriter* reader_writer) -> py::dict { return reader_writer->parameters_; },
-          [](PyAdapterFormatReaderWriter* reader_writer, py::dict& parameters) -> void {
-            reader_writer->parameters_ = parameters;
+          [](const PyAdapterFormatReaderWriter* reader_writer) -> py::dict {
+            return reader_writer->parameters_;
           },
-          R"pbdoc("Enables user to read/write adapter version stored in the file")pbdoc")
+          [](PyAdapterFormatReaderWriter* reader_writer, py::dict parameters) -> void {
+            reader_writer->parameters_ = std::move(parameters);
+          },
+          R"pbdoc("Enables user to read/write the dictionary of adapter parameters (name -> OrtValue)")pbdoc")
       .def(
           "export_adapter",
           [](const PyAdapterFormatReaderWriter* reader_writer, const std::wstring& path) {
             std::filesystem::path file_path(path);
-            std::ofstream file(file_path, std::ios::binary);
-            if (file.fail()) {
-              ORT_THROW("Failed to open file:", file_path, " for writing.");
-            }
 
             adapters::utils::AdapterFormatBuilder format_builder;
-            for (auto& [n, value] : reader_writer->parameters_) {
-              const std::string param_name = py::str(n);
-              const OrtValue* ort_value = value.cast<OrtValue*>();
-              const Tensor& tensor = ort_value->Get<Tensor>();
+
+            auto add_param = [&format_builder](const std::string& param_name, const OrtValue& ort_value) {
+              ORT_ENFORCE(ort_value.IsTensor(),
+                          "Lora adapter parameter '", param_name, "' must be a Tensor OrtValue.");
+              const Tensor& tensor = ort_value.Get<Tensor>();
+              ORT_ENFORCE(tensor.Location().device.Type() == OrtDevice::CPU,
+                          "Lora adapter parameter '", param_name, "' must reside on CPU to be exported.");
+              const auto element_type = tensor.GetElementType();
+              // Reject string tensors: Tensor::DataRaw() for a string tensor points to an
+              // array of std::string objects, and SizeInBytes() counts the sizeof(std::string)
+              // object representation (which contains heap pointers and uninitialized
+              // padding). Serializing those bytes would (a) leak runtime addresses
+              // (defeating ASLR) and uninitialized heap memory into the adapter file,
+              // and (b) produce an unloadable adapter, since reading the bytes back as
+              // std::string objects is undefined behavior. The adapter format has no
+              // representation for string tensors.
+              if (element_type == ONNX_NAMESPACE::TensorProto_DataType_STRING) {
+                ORT_THROW("Lora adapter parameter '", param_name,
+                          "' has element type STRING, which is not supported by the adapter format.");
+              }
               const auto data_span =
                   gsl::make_span<const uint8_t>(reinterpret_cast<const uint8_t*>(tensor.DataRaw()),
                                                 tensor.SizeInBytes());
               format_builder.AddParameter(
-                  param_name, static_cast<adapters::TensorDataType>(tensor.GetElementType()),
+                  param_name, static_cast<adapters::TensorDataType>(element_type),
                   tensor.Shape().GetDims(), data_span);
+            };
+
+            // Single source: always iterate parameters_ dict.
+            for (auto& [n, value] : reader_writer->parameters_) {
+              const std::string param_name = py::str(n);
+              const OrtValue* ort_value = value.cast<OrtValue*>();
+              add_param(param_name, *ort_value);
             }
 
+            // Build the entire adapter image in memory before touching the
+            // filesystem, so any failure (string-tensor rejection above, builder
+            // errors below, allocation failure inside FinishWithSpan, etc.)
+            // leaves no stray file behind.
             auto format_span = format_builder.FinishWithSpan(reader_writer->adapter_version_,
                                                              reader_writer->model_version_);
+
+            std::ofstream file(file_path, std::ios::binary);
+            if (file.fail()) {
+              ORT_THROW("Failed to open file:", file_path, " for writing.");
+            }
+
             if (file.write(reinterpret_cast<const char*>(format_span.data()), format_span.size()).fail()) {
               ORT_THROW("Failed to write :", std::to_string(format_span.size()), " bytes to ", file_path);
             }
@@ -123,22 +159,37 @@
 
       .def_static(
           "read_adapter", [](const std::wstring& file_path) -> std::unique_ptr<PyAdapterFormatReaderWriter> {
-            lora::LoraAdapter lora_adapter;
-            lora_adapter.Load(file_path);
-
-            auto [begin, end] = lora_adapter.GetParamIterators();
+            // Load into a heap-allocated LoraAdapter, then wrap it in a capsule.
+            // The capsule is pinned as a patient on each OrtValue view, so the
+            // backing memory stays alive as long as any Python reference to an
+            // OrtValue (or the dict containing them) exists.
+            auto adapter_ptr = std::make_unique<lora::LoraAdapter>();
+            adapter_ptr->Load(file_path);
+            const int format_version = adapter_ptr->FormatVersion();
+            const int adapter_version = adapter_ptr->AdapterVersion();
+            const int model_version = adapter_ptr->ModelVersion();
+
+            // Transfer ownership of the LoraAdapter to a capsule.
+            lora::LoraAdapter* raw_adapter = adapter_ptr.release();
+            py::capsule adapter_capsule(raw_adapter, [](void* p) {
+              delete static_cast<lora::LoraAdapter*>(p);
+            });
+
+            // Build the parameters dict with OrtValue views pinned to the capsule.
             py::dict params;
+            auto [begin, end] = raw_adapter->GetParamIterators();
             for (; begin != end; ++begin) {
               auto& [name, param] = *begin;
               OrtValue& ort_value = param.GetMapped();
-              params[py::str(name)] = py::cast(&ort_value);
+              // Cast with the capsule as parent: pybind11 attaches the capsule
+              // as a patient on the OrtValue handle (strong refcount pin).
+              py::object ort_value_obj = py::cast(
+                  &ort_value, py::return_value_policy::reference_internal, adapter_capsule);
+              params[py::str(name)] = std::move(ort_value_obj);
             }
 
-            auto py_adapter = std::make_unique<PyAdapterFormatReaderWriter>(
-                lora_adapter.FormatVersion(), lora_adapter.AdapterVersion(),
-                lora_adapter.ModelVersion(), std::move(lora_adapter), std::move(params));
-
-            return py_adapter;
+            return std::make_unique<PyAdapterFormatReaderWriter>(
+                format_version, adapter_version, model_version, std::move(params));
           },
           R"pbdoc(The function returns an instance of the class that contains a dictionary of name -> numpy arrays)pbdoc");
 
@@ -148,4 +199,4 @@
 }
 
 }  // namespace python
-}  // namespace onnxruntime
+}  // namespace onnxruntime