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,7 +151,13 @@ 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)) {}

onnxruntime/python/onnxruntime_pybind_lora.cc

@@ -40,15 +40,15 @@ struct PyAdapterFormatReaderWriter {
       : format_version_(format_version),
         adapter_version_(adapter_version),
         model_version_(model_version),
-        loaded_adater_(std::move(loaded_adapter)),
+        loaded_adapter_(std::move(loaded_adapter)),
         parameters_(std::move(params)) {}
 
   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_;
+  std::optional<lora::LoraAdapter> loaded_adapter_;
   // This is a dictionary of string -> OrtValue
   // this is populated directly on write and
   // built on top of the loaded_adapter on read
@@ -65,11 +65,6 @@ void addAdapterFormatMethods(pybind11::module& m) {
           "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 +77,81 @@ void addAdapterFormatMethods(pybind11::module& m) {
           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;
-          },
-          R"pbdoc("Enables user to read/write adapter version stored in the file")pbdoc")
+          // The dict's values are pybind11 wrappers around raw OrtValue* pointers
+          // that alias storage owned by PyAdapterFormatReaderWriter::loaded_adapter_
+          // (see read_adapter below: entries are produced via py::cast(&ort_value),
+          // i.e. non-owning views into the LoraAdapter's params_values_ map).
+          //
+          // Without an explicit keep-alive, Python users following the natural
+          // pattern
+          //
+          //     params = ort.AdapterFormat.read_adapter(path).parameters
+          //
+          // would silently get a dict of dangling pointers: the temporary
+          // AdapterFormat (and with it the LoraAdapter that backs every OrtValue)
+          // is destroyed at the end of the expression, and any subsequent access
+          // such as params["x"].numpy() would dereference freed memory.
+          //
+          // Python callers reasonably expect refcounting to keep parents alive
+          // while derived objects are still referenced; this is the binding
+          // author's responsibility for non-owning views. py::keep_alive<0, 1>
+          // ties the returned dict (return value, index 0) to the owning
+          // PyAdapterFormatReaderWriter (self, index 1), so the parent — and
+          // therefore the underlying adapter buffer — survives at least as long
+          // as the dict the caller holds. def_property() does not accept
+          // keep_alive directly, so the policy is attached to the getter via
+          // py::cpp_function.
+          py::cpp_function(
+              [](const PyAdapterFormatReaderWriter* reader_writer) -> py::dict { return reader_writer->parameters_; },
+              py::keep_alive<0, 1>()),
+          py::cpp_function(
+              [](PyAdapterFormatReaderWriter* reader_writer, py::dict& parameters) -> void {
+                reader_writer->parameters_ = 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>();
+              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);
             }
 
+            // 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);
             }

onnxruntime/test/python/onnxruntime_test_python.py

@@ -2103,6 +2103,69 @@ def test_adater_export_read(self):
             self.assertEqual(expected_val.shape(), value.shape())
             np.testing.assert_allclose(value.numpy(), expected_val.numpy())
 
+    def test_adapter_parameters_keep_alive(self):
+        # Regression test: AdapterFormat.parameters returned a dict of OrtValue
+        # views into the parent adapter without a pybind11 keep_alive policy.
+        # The natural pattern below dropped the parent and left the dict with
+        # dangling pointers, causing a use-after-free on the next access.
+        adapter_version = 1
+        model_version = 1
+        file_path = pathlib.Path(os.path.realpath(__file__)).parent
+        file_path = str(file_path / "test_adapter_keep_alive.onnx_adapter")
+
+        float_data_type = 1
+        val = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
+        param_1 = np.array(val).astype(np.float32).reshape(5, 2)
+
+        ort_val_1 = onnxrt.OrtValue.ortvalue_from_numpy_with_onnx_type(param_1, float_data_type)
+
+        adapter_format = onnxrt.AdapterFormat()
+        adapter_format.set_adapter_version(adapter_version)
+        adapter_format.set_model_version(model_version)
+        adapter_format.set_parameters({"param_1": ort_val_1})
+        adapter_format.export_adapter(file_path)
+
+        try:
+            # Drop the AdapterFormat temporary; only `params` keeps a reference.
+            params = onnxrt.AdapterFormat.read_adapter(file_path).parameters
+            gc.collect()
+
+            self.assertIn("param_1", params)
+            value = params["param_1"]
+            self.assertTrue(value.is_tensor())
+            self.assertEqual(value.shape(), [5, 2])
+            np.testing.assert_allclose(value.numpy(), param_1)
+        finally:
+            os.remove(file_path)
+
+    def test_adapter_export_rejects_string_tensors(self):
+        # Regression test: export_adapter previously serialized Tensor::DataRaw()
+        # for SizeInBytes() bytes regardless of element type. For string tensors
+        # that copied the std::string object representation (heap pointers and
+        # uninitialized padding) into the adapter file, leaking runtime addresses
+        # (ASLR bypass) and producing an unloadable adapter. Export must reject
+        # string-typed parameters.
+        string_data_type = 8  # ONNX TensorProto.STRING
+        string_array = np.array(["hello", "world"], dtype=object)
+        ort_val_str = onnxrt.OrtValue.ortvalue_from_numpy_with_onnx_type(string_array, string_data_type)
+
+        adapter_format = onnxrt.AdapterFormat()
+        adapter_format.set_adapter_version(1)
+        adapter_format.set_model_version(1)
+        adapter_format.set_parameters({"str_param": ort_val_str})
+
+        file_path = pathlib.Path(os.path.realpath(__file__)).parent
+        file_path = str(file_path / "test_adapter_string_reject.onnx_adapter")
+
+        try:
+            with self.assertRaises(Exception) as ctx:
+                adapter_format.export_adapter(file_path)
+            self.assertIn("STRING", str(ctx.exception))
+            self.assertFalse(os.path.exists(file_path), "adapter file must not be created when export is rejected")
+        finally:
+            if os.path.exists(file_path):
+                os.remove(file_path)
+
     def test_run_with_adapter(self):
         model_path = get_name("lora/two_params_lora_model.onnx")
         file_path = os.getcwd() + "/" + get_name("lora/two_params_lora_model.onnx_adapter")