Add pybind support for all memory config options in OrtArenaCfg (#12658)

* Add support for initial_growth_chunk_size_bytes setting in OrtArenaCfg pybind

* Add overloaded constructor for KVP, UT still in progress

* Fix class member access in pybind, fix unit test

* Resolve linter warnings

* Improve formatting

* Simplify UT

* Fix linter formatting

Co-authored-by: Peter Mcaughan <petermca@microsoft.com>

onnxruntime/python/onnxruntime_pybind_state.cc

@@ -1083,6 +1083,8 @@ void addObjectMethods(py::module& m, Environment& env, ExecutionProviderRegistra
       .def_static("default_memory", []() { return OrtDevice::MemType::DEFAULT; });
 
   py::class_<OrtArenaCfg> ort_arena_cfg_binding(m, "OrtArenaCfg");
+  // Note: Doesn't expose initial_growth_chunk_sizes_bytes option. This constructor kept for
+  // backwards compatibility, key-value pair constructor overload exposes all options
   // There is a global var: arena_extend_strategy, which means we can't use that var name here
   // See docs/C_API.md for details on what the following parameters mean and how to choose these values
   ort_arena_cfg_binding.def(py::init([](size_t max_mem, int arena_extend_strategy_local,
@@ -1093,7 +1095,32 @@ void addObjectMethods(py::module& m, Environment& env, ExecutionProviderRegistra
     ort_arena_cfg->initial_chunk_size_bytes = initial_chunk_size_bytes;
     ort_arena_cfg->max_dead_bytes_per_chunk = max_dead_bytes_per_chunk;
     return ort_arena_cfg;
-  }));
+  }))
+      .def(py::init([](const py::dict& feeds) {
+    auto ort_arena_cfg = std::make_unique<OrtArenaCfg>();
+    for (const auto kvp : feeds) {
+      std::string key = kvp.first.cast<std::string>();
+      if (key == "max_mem") {
+        ort_arena_cfg->max_mem = kvp.second.cast<size_t>();
+      } else if (key == "arena_extend_strategy") {
+        ort_arena_cfg->arena_extend_strategy = kvp.second.cast<int>();
+      } else if (key == "initial_chunk_size_bytes") {
+        ort_arena_cfg->initial_chunk_size_bytes = kvp.second.cast<int>();
+      } else if (key == "max_dead_bytes_per_chunk") {
+        ort_arena_cfg->max_dead_bytes_per_chunk = kvp.second.cast<int>();
+      } else if (key == "initial_growth_chunk_size_bytes") {
+        ort_arena_cfg->initial_growth_chunk_size_bytes = kvp.second.cast<int>();
+        } else {
+        ORT_THROW("Invalid OrtArenaCfg option: ", key);
+      }
+    }
+    return ort_arena_cfg;
+  }))
+      .def_readwrite("max_mem", &OrtArenaCfg::max_mem)
+      .def_readwrite("arena_extend_strategy", &OrtArenaCfg::arena_extend_strategy)
+      .def_readwrite("initial_chunk_size_bytes", &OrtArenaCfg::initial_chunk_size_bytes)
+      .def_readwrite("max_dead_bytes_per_chunk", &OrtArenaCfg::max_dead_bytes_per_chunk)
+      .def_readwrite("initial_growth_chunk_size_bytes", &OrtArenaCfg::initial_growth_chunk_size_bytes);
 
   py::class_<OrtMemoryInfo> ort_memory_info_binding(m, "OrtMemoryInfo");
   ort_memory_info_binding.def(py::init([](const char* name, OrtAllocatorType type, int id, OrtMemType mem_type) {

onnxruntime/test/python/onnxruntime_test_python.py

@@ -1154,7 +1154,9 @@ class TestInferenceSession(unittest.TestCase):
     def testSharedAllocatorUsingCreateAndRegisterAllocator(self):
         # Create and register an arena based allocator
 
-        # ort_arena_cfg = onnxrt.OrtArenaCfg(0, -1, -1, -1) (create an OrtArenaCfg like this template if you want to use non-default parameters)
+        # To create an OrtArenaCfg using non-default parameters, use one of below templates:
+        # ort_arena_cfg = onnxrt.OrtArenaCfg(0, -1, -1, -1) - Note: doesn't expose initial_growth_chunk_size_bytes option
+        # ort_arena_cfg = onnxrt.OrtArenaCfg({"max_mem": -1, ""arena_extend_strategy": 1, etc..})
         ort_memory_info = onnxrt.OrtMemoryInfo(
             "Cpu",
             onnxrt.OrtAllocatorType.ORT_ARENA_ALLOCATOR,
@@ -1308,6 +1310,43 @@ class TestInferenceSession(unittest.TestCase):
         )
         print("Create session with customize execution provider successfully!")
 
+    def testCreateAllocator(self):
+        def verify_allocator(allocator, expected_config):
+            for key, val in expected_config.items():
+                if key == "max_mem":
+                    self.assertEqual(allocator.max_mem, val)
+                elif key == "arena_extend_strategy":
+                    self.assertEqual(allocator.arena_extend_strategy, val)
+                elif key == "initial_chunk_size_bytes":
+                    self.assertEqual(allocator.initial_chunk_size_bytes, val)
+                elif key == "max_dead_bytes_per_chunk":
+                    self.assertEqual(allocator.max_dead_bytes_per_chunk, val)
+                elif key == "initial_growth_chunk_size_bytes":
+                    self.assertEqual(allocator.initial_growth_chunk_size_bytes, val)
+                else:
+                    raise ValueError("Invalid OrtArenaCfg option: " + key)
+
+        # Verify ordered parameter initialization
+        ort_arena_cfg = onnxrt.OrtArenaCfg(8, 0, 4, 2)
+        expected_allocator = {
+            "max_mem": 8,
+            "arena_extend_strategy": 0,
+            "initial_chunk_size_bytes": 4,
+            "max_dead_bytes_per_chunk": 2,
+        }
+        verify_allocator(ort_arena_cfg, expected_allocator)
+
+        # Verify key-value pair initialization
+        expected_kvp_allocator = {
+            "max_mem": 16,
+            "arena_extend_strategy": 1,
+            "initial_chunk_size_bytes": 8,
+            "max_dead_bytes_per_chunk": 4,
+            "initial_growth_chunk_size_bytes": 2,
+        }
+        ort_arena_cfg_kvp = onnxrt.OrtArenaCfg(expected_kvp_allocator)
+        verify_allocator(ort_arena_cfg_kvp, expected_kvp_allocator)
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=1)