[Serving] Fix memory leak when update full model in serving. (#809)

shanshanpt · web-flow · commit 99e28041decf · 2023-04-14T08:06:28.000+08:00
Signed-off-by: Tao Peng &lt;jiankeng.pt@alibaba-inc.com&gt;
diff --git 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
@@ -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]) {

Original file line number	Diff line number	Diff line change
`@@ -159,7 +159,7 @@ void LocalDevice::Init(const SessionOptions& options,`
`159`	`159`	`tp_info = global_tp_info_[numa_node];`
`160`	`160`	`} else {`
`161`	`161`	`if (opt.global_threadpool_num > 1) {`
`162`		`- for (int i = 0; i < opt.global_threadpool_num; ++i) {`
	`162`	`+ while (opt.global_threadpool_num > global_tp_info_.size()) {`
`163`	`163`	`global_tp_info_.push_back(nullptr);`
`164`	`164`	`}`
`165`	`165`	`if (!global_tp_info_[opt.device_threadpool_index]) {`