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

@@ -107,6 +107,10 @@ def set_parameters(self, params: dict[str, OrtValue]) -> None:
         self._adapter.parameters = {k: v._ortvalue for k, v in params.items()}
 
     def get_parameters(self) -> dict[str, OrtValue]:
+        # Each call to the C getter rebuilds the dict and pins the owning
+        # C AdapterFormat onto every returned OrtValue (pybind11 add_patient
+        # via keep_alive_impl), so callers can drop this AdapterFormat and
+        # the dict (or any individual value) keeps the backing adapter alive.
         return {k: OrtValue(v) for k, v in self._adapter.parameters.items()}
 
 

onnxruntime/python/onnxruntime_pybind_lora.cc

@@ -35,23 +35,25 @@
   PyAdapterFormatReaderWriter() = default;
   PyAdapterFormatReaderWriter(int format_version, int adapter_version,
                               int model_version,
-                              lora::LoraAdapter&& loaded_adapter,
-                              py::dict&& params)
+                              lora::LoraAdapter&& loaded_adapter)
       : format_version_(format_version),
         adapter_version_(adapter_version),
         model_version_(model_version),
-        loaded_adater_(std::move(loaded_adapter)),
-        parameters_(std::move(params)) {}
+        loaded_adapter_(std::move(loaded_adapter)) {}
 
   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
+  // Populated only on the read path (read_adapter). When present, the
+  // `parameters` getter builds OrtValue views over its params_values_ map
+  // and pins this object as the keep_alive patient on each value.
+  std::optional<lora::LoraAdapter> loaded_adapter_;
+  // Populated only on the write path (user-supplied via the `parameters`
+  // setter) and consumed by export_adapter. On the read path this stays
+  // empty; the getter builds a fresh dict from loaded_adapter_ each call,
+  // which avoids an AdapterFormat -> dict -> OrtValue -> AdapterFormat
+  // reference cycle that pybind11's patient list would not let the GC
+  // break.
   py::dict parameters_;
 };
 
@@ -65,11 +67,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 +79,124 @@
           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_; },
+          // Getter: build the dict on demand. On the read path each
+          // returned OrtValue is a non-owning view over storage owned by
+          // loaded_adapter_, so we pin `self` (the AdapterFormat Python
+          // object) onto every OrtValue. The patient list bumps self's
+          // Python refcount, so callers that keep any individual OrtValue
+          // keep the backing adapter alive.
+          //
+          // Why per-element pinning, not a property-level
+          // return_value_policy: pybind11 policies attach to the function's
+          // single return value. Setting reference_internal on this getter
+          // would try to pin `self` onto the returned `py::dict`, but
+          // (a) dict is not a pybind-registered type so keep_alive falls
+          // back to weakref and dicts are not weak-referenceable -> runtime
+          // TypeError, and (b) even if that worked, it would not propagate
+          // to the dict's elements, so a caller doing
+          //     value = params["x"]; del params
+          // would still dangle. We therefore apply the policy per element,
+          // using the per-call form of py::cast: it requests
+          // reference_internal with `self` as parent, which routes through
+          // pybind11's add_patient on the OrtValue (a registered type) --
+          // a strong refcount-based pin, not a weakref.
+          //
+          // The dict is built in a local first and only returned at the
+          // end, so a throw mid-build leaves no half-populated state.
+          //
+          // We deliberately do NOT cache the dict on `self`: caching would
+          // create a self -> dict -> OrtValue -> self reference cycle
+          // (via the patient list), and pybind11 instances are not
+          // traversable by the cyclic GC, so the cycle would leak.
+          [](py::object self) -> py::dict {
+            auto* rw = self.cast<PyAdapterFormatReaderWriter*>();
+            if (!rw->loaded_adapter_.has_value()) {
+              // Write path: return whatever the user previously set.
+              return rw->parameters_;
+            }
+            py::dict params;
+            auto [begin, end] = rw->loaded_adapter_->GetParamIterators();
+            for (; begin != end; ++begin) {
+              auto& [name, param] = *begin;
+              OrtValue& ort_value = param.GetMapped();
+              // reference_internal == "borrowed reference owned by `self`";
+              // pybind11 attaches `self` as a keep_alive patient on the
+              // returned OrtValue handle.
+              py::object ort_value_obj = py::cast(
+                  &ort_value, py::return_value_policy::reference_internal, self);
+              params[py::str(name)] = std::move(ort_value_obj);
+            }
+            return params;
+          },
           [](PyAdapterFormatReaderWriter* reader_writer, py::dict& parameters) -> void {
             reader_writer->parameters_ = parameters;
           },
-          R"pbdoc("Enables user to read/write adapter version stored in the file")pbdoc")
+          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>();
+
+            // Visit every (name, OrtValue) pair from the appropriate source.
+            // We avoid building any intermediate Python dict here so the
+            // export path is exception-safe by construction: the builder
+            // accumulates in memory, and we only touch the filesystem
+            // after FinishWithSpan succeeds.
+            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);
+            };
+
+            if (reader_writer->loaded_adapter_.has_value()) {
+              auto [begin, end] = reader_writer->loaded_adapter_->GetParamIterators();
+              for (; begin != end; ++begin) {
+                auto& [name, param] = *begin;
+                add_param(name, param.GetMapped());
+              }
+            } else {
+              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);
             }
@@ -125,20 +211,11 @@
           "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();
-            py::dict params;
-            for (; begin != end; ++begin) {
-              auto& [name, param] = *begin;
-              OrtValue& ort_value = param.GetMapped();
-              params[py::str(name)] = py::cast(&ort_value);
-            }
-
-            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;
+            const int format_version = lora_adapter.FormatVersion();
+            const int adapter_version = lora_adapter.AdapterVersion();
+            const int model_version = lora_adapter.ModelVersion();
+            return std::make_unique<PyAdapterFormatReaderWriter>(
+                format_version, adapter_version, model_version, std::move(lora_adapter));
           },
           R"pbdoc(The function returns an instance of the class that contains a dictionary of name -> numpy arrays)pbdoc");