diff --git a/onnxruntime/core/session/lora_adapters.cc b/onnxruntime/core/session/lora_adapters.cc index de3acb22e12f8..0f25dcd85355d 100644 --- a/onnxruntime/core/session/lora_adapters.cc +++ b/onnxruntime/core/session/lora_adapters.cc @@ -37,17 +37,35 @@ void LoraAdapter::Load(const std::filesystem::path& file_path) { } void LoraAdapter::Load(std::vector 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(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(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(std::move(mapped_memory), file_size); - InitializeParamsValues(); + adapter_ = new_adapter; + params_values_ = std::move(new_params_values); } static std::unique_ptr 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 +LoraAdapter::BuildParamsValues(const adapters::Adapter* adapter) const { + ORT_ENFORCE(adapter != nullptr, "Adapter is not loaded yet."); std::unique_ptr 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 params_values; params_values.reserve(params->size()); @@ -117,7 +134,7 @@ void LoraAdapter::InitializeParamsValues() { } } - params_values_.swap(params_values); + return params_values; } } // namespace lora diff --git a/onnxruntime/core/session/lora_adapters.h b/onnxruntime/core/session/lora_adapters.h index 77534b2bb7d15..02c37a55726eb 100644 --- a/onnxruntime/core/session/lora_adapters.h +++ b/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 BuildParamsValues(const adapters::Adapter* adapter) const; struct BufferHolder { - explicit BufferHolder(std::vector buffer) : buffer_(std::move(buffer)) {} + explicit BufferHolder(std::vector buffer) noexcept : buffer_(std::move(buffer)) {} + BufferHolder(BufferHolder&&) noexcept = default; + BufferHolder& operator=(BufferHolder&&) noexcept = default; std::vector 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_; }; diff --git a/onnxruntime/python/onnxruntime_inference_collection.py b/onnxruntime/python/onnxruntime_inference_collection.py index 59ddfc49d71a6..70f699768e92c 100644 --- a/onnxruntime/python/onnxruntime_inference_collection.py +++ b/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()} diff --git a/onnxruntime/python/onnxruntime_pybind_lora.cc b/onnxruntime/python/onnxruntime_pybind_lora.cc index af8365418e5ea..7c6608a4f9052 100644 --- a/onnxruntime/python/onnxruntime_pybind_lora.cc +++ b/onnxruntime/python/onnxruntime_pybind_lora.cc @@ -29,29 +29,39 @@ namespace py = pybind11; namespace { /// /// 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. /// 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 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 @@ 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 +87,66 @@ 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; + [](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(); - const Tensor& tensor = ort_value->Get(); + + 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(); + 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(reinterpret_cast(tensor.DataRaw()), tensor.SizeInBytes()); format_builder.AddParameter( - param_name, static_cast(tensor.GetElementType()), + param_name, static_cast(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(); + 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(format_span.data()), format_span.size()).fail()) { ORT_THROW("Failed to write :", std::to_string(format_span.size()), " bytes to ", file_path); } @@ -123,24 +159,41 @@ void addAdapterFormatMethods(pybind11::module& m) { .def_static( "read_adapter", [](const std::wstring& file_path) -> std::unique_ptr { - 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(); + 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. + // Construct capsule while unique_ptr still owns the pointer, so + // if capsule allocation throws, unique_ptr cleans up. + py::capsule adapter_capsule(adapter_ptr.get(), [](void* p) { + delete static_cast(p); + }); + lora::LoraAdapter* raw_adapter = adapter_ptr.release(); + + // 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( - lora_adapter.FormatVersion(), lora_adapter.AdapterVersion(), - lora_adapter.ModelVersion(), std::move(lora_adapter), std::move(params)); - - return py_adapter; + return std::make_unique( + 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"); + R"pbdoc(The function returns an instance of the class that contains a dictionary of name -> OrtValue)pbdoc"); py::class_ lora_adapter_binding(m, "LoraAdapter"); lora_adapter_binding.def(py::init()) @@ -148,4 +201,4 @@ void addAdapterFormatMethods(pybind11::module& m) { } } // namespace python -} // namespace onnxruntime \ No newline at end of file +} // namespace onnxruntime diff --git a/onnxruntime/test/python/onnxruntime_test_python.py b/onnxruntime/test/python/onnxruntime_test_python.py index 1368f6d53deb7..bebd290bbe810 100644 --- a/onnxruntime/test/python/onnxruntime_test_python.py +++ b/onnxruntime/test/python/onnxruntime_test_python.py @@ -2103,6 +2103,143 @@ 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_read_modify_export(self): + # Verify that an instance created by read_adapter can have its + # parameters replaced and re-exported (single-source architecture). + adapter_version = 1 + model_version = 1 + file_path = pathlib.Path(os.path.realpath(__file__)).parent + original_path = str(file_path / "test_adapter_read_modify_original.onnx_adapter") + modified_path = str(file_path / "test_adapter_read_modify_new.onnx_adapter") + + float_data_type = 1 + original_data = np.array([1, 2, 3, 4]).astype(np.float32).reshape(2, 2) + modified_data = np.array([10, 20, 30, 40, 50, 60]).astype(np.float32).reshape(3, 2) + + ort_original = onnxrt.OrtValue.ortvalue_from_numpy_with_onnx_type(original_data, float_data_type) + ort_modified = onnxrt.OrtValue.ortvalue_from_numpy_with_onnx_type(modified_data, float_data_type) + + # Write original adapter + adapter_format = onnxrt.AdapterFormat() + adapter_format.set_adapter_version(adapter_version) + adapter_format.set_model_version(model_version) + adapter_format.set_parameters({"param": ort_original}) + adapter_format.export_adapter(original_path) + + try: + # Read, replace parameters, and re-export. + adapter_read = onnxrt.AdapterFormat.read_adapter(original_path) + adapter_read.set_parameters({"new_param": ort_modified}) + adapter_read.export_adapter(modified_path) + + # Verify the re-exported file has the NEW parameters. + adapter_verify = onnxrt.AdapterFormat.read_adapter(modified_path) + params = adapter_verify.get_parameters() + self.assertIn("new_param", params) + self.assertNotIn("param", params) + self.assertEqual(params["new_param"].shape(), [3, 2]) + np.testing.assert_allclose(params["new_param"].numpy(), modified_data) + finally: + if os.path.exists(original_path): + os.remove(original_path) + if os.path.exists(modified_path): + os.remove(modified_path) + + def test_adapter_parameters_keep_alive(self): + # Regression test: AdapterFormat.read_adapter returned OrtValue views + # over storage owned by the C AdapterFormat with nothing keeping the + # parent alive. The natural pattern below dropped the parent and left + # the dict with dangling pointers, causing a use-after-free on the + # next access. read_adapter now pins the owning C AdapterFormat on + # every OrtValue it produces (pybind11 add_patient via + # keep_alive_impl), so the dict and any individual value keep the + # backing adapter alive on their own. Mirrors the strong-ref pattern + # used by the SparseTensor view bindings. + 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).get_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) + + # Also drop the dict; an individual OrtValue must keep the adapter + # alive on its own (we pin it on every value in get_parameters()). + single_value = params["param_1"] + del params + gc.collect() + np.testing.assert_allclose(single_value.numpy(), param_1) + finally: + if os.path.exists(file_path): + 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. + # + # There is no public Python API to construct a string-typed OrtValue + # directly (ortvalue_from_numpy rejects non-numeric arrays), so we + # obtain one by running a tiny Constant model whose only output is a + # string tensor. + from onnx import TensorProto, helper # noqa: PLC0415 + + const_node = helper.make_node( + "Constant", + inputs=[], + outputs=["str_out"], + value=helper.make_tensor("v", TensorProto.STRING, dims=[2], vals=[b"hello", b"world"]), + ) + graph = helper.make_graph( + [const_node], + "string_const", + inputs=[], + outputs=[helper.make_tensor_value_info("str_out", TensorProto.STRING, [2])], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) + sess = onnxrt.InferenceSession(model.SerializeToString(), providers=onnxrt.get_available_providers()) + ort_val_str = sess.run_with_ort_values(["str_out"], {})[0] + + 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")