From 92cf6336a0939a49aa1b1627c8e1d5cf61e54cea Mon Sep 17 00:00:00 2001 From: Tao Peng Date: Thu, 13 Apr 2023 17:39:28 +0800 Subject: [PATCH] [Runtime] Fix memory leak when update full model in serving. Signed-off-by: Tao Peng --- tensorflow/core/common_runtime/direct_session_group.cc | 4 +++- tensorflow/core/common_runtime/local_device.cc | 2 +- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/tensorflow/core/common_runtime/direct_session_group.cc b/tensorflow/core/common_runtime/direct_session_group.cc index 40fa22b65a1..28659cfd48a 100644 --- a/tensorflow/core/common_runtime/direct_session_group.cc +++ b/tensorflow/core/common_runtime/direct_session_group.cc @@ -26,7 +26,9 @@ DirectSessionGroup::DirectSessionGroup() DirectSessionGroup::DirectSessionGroup( ResourceMgr* cpu_mgr, ResourceMgr* gpu_mgr) : cpu_shared_resource_mgr_(cpu_mgr) { - gpu_shared_resource_mgrs_.emplace_back(gpu_mgr); + if (gpu_mgr) { + gpu_shared_resource_mgrs_.emplace_back(gpu_mgr); + } leader_session_ids_.emplace_back(0); } diff --git a/tensorflow/core/common_runtime/local_device.cc b/tensorflow/core/common_runtime/local_device.cc index c783c633129..9eb8103a853 100644 --- a/tensorflow/core/common_runtime/local_device.cc +++ b/tensorflow/core/common_runtime/local_device.cc @@ -159,7 +159,7 @@ void LocalDevice::Init(const SessionOptions& options, tp_info = global_tp_info_[numa_node]; } else { if (opt.global_threadpool_num > 1) { - for (int i = 0; i < opt.global_threadpool_num; ++i) { + while (opt.global_threadpool_num > global_tp_info_.size()) { global_tp_info_.push_back(nullptr); } if (!global_tp_info_[opt.device_threadpool_index]) {