Skip to content

Commit b5a7410

Browse files
committed
fix crash during torch model reload with inter op thread parameter set
torch allows setting the interop thread count only once and before any inter-op parallel work is started because it creates threadpool statically. To follow this constraint, check the interop thread count and set only if it's not yet set.
1 parent db70751 commit b5a7410

File tree

1 file changed

+9
-7
lines changed

1 file changed

+9
-7
lines changed

src/libtorch.cc

Lines changed: 9 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -532,15 +532,17 @@ ModelState::ParseParameters()
532532
TRITONSERVER_ErrorDelete(err);
533533
}
534534
} else {
535-
if (inter_op_thread_count > 0) {
535+
if ((inter_op_thread_count > 0) &&
536+
(inter_op_thread_count != at::get_num_interop_threads()) &&
537+
(at::get_num_interop_threads() == std::thread::hardware_concurrency())) {
536538
at::set_num_interop_threads(inter_op_thread_count);
537-
LOG_MESSAGE(
538-
TRITONSERVER_LOG_INFO,
539-
(std::string("Inter op thread count is set to ") +
540-
std::to_string(inter_op_thread_count) + " for model instance '" +
541-
Name() + "'")
542-
.c_str());
543539
}
540+
LOG_MESSAGE(
541+
TRITONSERVER_LOG_INFO,
542+
(std::string("Inter op thread count is set to ") +
543+
std::to_string(at::get_num_interop_threads()) + " for model instance '" +
544+
Name() + "'")
545+
.c_str());
544546
}
545547
}
546548

0 commit comments

Comments
 (0)