Add RAII wrappers for ORT Model Editor API types #1953
Conversation
There was a problem hiding this comment.
Pull request overview
This PR adds RAII (Resource Acquisition Is Initialization) wrappers for ONNX Runtime Model Editor API types and refactors the graph builder code to use these wrappers for improved resource management. Additionally, it removes mutex protection from the graph session cache based on the assumption that ort-genai execution is single-threaded.
Changes:
- Added RAII wrapper structs (OrtGraph, OrtModel, OrtValueInfo, OrtNode) in onnxruntime_api.h following the existing OrtOpAttr pattern
- Removed SessionCache struct and mutex from graph_session_cache_ in generators.h, replacing it with a simple unordered_map
- Refactored graph_builder.cpp Build() function to use std::unique_ptr for automatic resource cleanup instead of manual try/catch blocks
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 5 comments.
| File | Description |
|---|---|
| src/models/onnxruntime_api.h | Adds RAII wrapper structs for OrtGraph, OrtModel, OrtValueInfo, and OrtNode with custom delete operators for automatic cleanup |
| src/generators.h | Removes SessionCache struct with mutex, replaces with simple unordered_map based on single-threaded execution assumption |
| src/models/graph_executor.cpp | Removes mutex include and simplifies cache access after SessionCache removal |
| src/models/graph_builder.cpp | Refactors Build() to use RAII patterns with std::unique_ptr instead of manual try/catch exception handling |
src/generators.h
Outdated
| // Note: No mutex needed since ort-genai execution is single-threaded | ||
| std::unordered_map<uint64_t, std::unique_ptr<OrtSession>> graph_session_cache_; |
There was a problem hiding this comment.
The comment stating "No mutex needed since ort-genai execution is single-threaded" may not be accurate. The codebase contains evidence of multi-threaded execution:
- The Engine class documentation (src/engine/engine.h:24) states it's "designed to handle multiple requests concurrently"
- WorkerThread class exists for asynchronous task execution (src/worker_thread.h)
- Other global caches use mutexes for thread safety (e.g., TopkBenchmarkCacheManager in src/cuda/cuda_topk_benchmark_cache.h:70)
- The tracing sink uses a mutex for thread-safe logging (src/tracing.cpp:85)
If graph_session_cache_ can be accessed from multiple threads (e.g., through ExecuteCastOp called from different Engine requests), removing the mutex could lead to race conditions during cache lookups and insertions. Consider verifying whether GetOrCreateSession can be called concurrently before removing the mutex protection.
| // Create node attributes | ||
| for (const auto& attr : config.attributes) { | ||
| node_attributes.push_back(CreateOpAttr(attr)); | ||
| } |
There was a problem hiding this comment.
The node_attributes vector contains raw pointers that are manually released at line 218-220. However, if any Ort::ThrowOnError call throws an exception between lines 174 (where attributes are created) and 220 (where they're released), these attributes will leak.
Consider using std::unique_ptr for node_attributes as well, similar to how graph_inputs_owned and graph_outputs_owned are managed. This would ensure automatic cleanup on exception. The raw pointers for passing to CreateNode can be extracted when needed.
There was a problem hiding this comment.
CreateNode does NOT copy the attributes - it stores references/pointers to them. Therefore:
The OrtOpAttr* pointers must remain valid until AFTER the model is fully built
We cannot use RAII (std::unique_ptr) which would destroy them immediately when the vector goes out of scope
Manual cleanup is required, but only AFTER AddGraphToModel completes
|
@baijumeswani This is the follow-up PR to resolve your remaining comments in #1895. The pipeline failures may not related with my changes. My another PR #1952 meets the similar issue after I rebase the code to latest without any code changes. However, it passed before. |
The pipeline failures are because Hugging Face released |
No description provided.