Fix cleanup of Environment class.

include/onnxruntime/core/session/environment.h

@@ -199,23 +199,6 @@ class Environment {
 
   using OrtAllocatorUniquePtr = std::unique_ptr<OrtAllocator, std::function<void(OrtAllocator*)>>;
 
-  // if the user calls CreateSharedAllocator and wraps the plugin EP's allocator with an arena we end up with
-  // OrtAllocator from EP -> wrapped in IAllocatorImplWrappingOrtAllocator -> inside a BFCArena IAllocator.
-  // we can put that in shared_allocators_ for sessions to use, but to have an OrtAllocator available in
-  // shared_ort_allocators_ that can be used outside of a session we need to additionally wrap that in an
-  // OrtAllocatorImplWrappingIAllocator. way too many levels of indirection but that is what it is currently.
-  // we need something to own that final OrtAllocator, so we add it to arena_ort_allocators_.
-  //
-  // TODO: we could split out the BFCArena implementation so it can be plugged into either an IAllocator
-  // or an OrtAllocator instance to reduce the indirection a little.
-  // with that we get an OrtAllocator from the EP, wrap it with an OrtAllocator based BFCArena, and wrap that with the
-  // IAllocatorImplWrappingOrtAllocator which takes ownership of the OrtAllocator and is in shared_allocators_.
-  //
-  // Alternatively we can disable wrapping an EP's allocator with a BFCArena and say the EP should provide the arena
-  // implementation directly. They're free to copy BFCArena as it came from TF originally. Or we could provide a
-  // cut-and-paste BFCArena implementation that works using the EP API that can be included in the EP source.
-  std::unordered_map<const OrtMemoryInfo*, std::unique_ptr<OrtAllocatorImplWrappingIAllocator>> arena_ort_allocators_;
-
 #if !defined(ORT_MINIMAL_BUILD)
   // register EPs that are built into the ORT binary so they can take part in AutoEP selection
   // added to ep_libraries

onnxruntime/core/session/environment.cc

@@ -182,11 +182,6 @@ Status Environment::UnregisterAllocatorImpl(const OrtMemoryInfo& mem_info, bool
     shared_ort_allocators_.erase(it2);
   }
 
-  // also remove an arena wrapped allocator from an EP if the user called CreateSharedAllocator to create one
-  if (auto it3 = arena_ort_allocators_.find(&mem_info); it3 != arena_ort_allocators_.end()) {
-    arena_ort_allocators_.erase(it3);
-  }
-
   if (found_shared_allocator) {
     shared_allocators_.erase(it);
   }
@@ -436,6 +431,10 @@ Environment::~Environment() {
   // instance and will call Release on it. If the plugin EP has been freed the Release will fail.
   shared_allocators_.clear();
 
+  // and as any OrtAllocator instances in shared_ort_allocators_ were owned by values in shared_allocators_ and have
+  // now been released we need to clear that too before calling UnregisterExecutionProviderLibrary().
+  shared_ort_allocators_.clear();
+
 #if !defined(ORT_MINIMAL_BUILD)
   // unregister any remaining EP libraries so they're cleaned up in a determistic way.
   while (!ep_libraries_.empty()) {
@@ -673,11 +672,6 @@ Status Environment::CreateSharedAllocatorImpl(const OrtEpDevice& ep_device,
     shared_ort_allocators_.erase(it);
   }
 
-  // if a previous call created an arena wrapped allocator for the EP's memory_info we also need to remove that
-  if (auto it = arena_ort_allocators_.find(&memory_info); it != arena_ort_allocators_.end()) {
-    arena_ort_allocators_.erase(it);
-  }
-
   // we only want one shared allocator for an OrtDevice in the shared_allocators_ so that it's deterministic which
   // one will be used for an inference session. ignore the name so that is the case.
   if (auto it = FindExistingAllocator(shared_allocators_, memory_info, /*match_name*/ false);