Skip to content

Commit 1b6fa77

Browse files
krishung5krishung5
authored
Add healthiness check to avoid hanging during model initialization (triton-inference-server#221)
* Add healthiness check to avoid hanging during model initialization * Address comment * Fix the type of the argument
1 parent 4f23086 commit 1b6fa77

4 files changed

Lines changed: 65 additions & 46 deletions

File tree

‎src/python_be.cc‎

Lines changed: 2 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -165,7 +165,7 @@ ModelInstanceState::SendMessageAndReceiveResponse(
165165
}
166166

167167
bi::managed_external_buffer::handle_t response_message;
168-
error = ReceiveMessageFromStub(response_message);
168+
error = Stub()->ReceiveMessageFromStub(response_message);
169169
if (error != nullptr) {
170170
restart = true;
171171
RespondErrorToAllRequests(
@@ -215,44 +215,6 @@ ModelInstanceState::SendMessageToStub(
215215
return nullptr; // success
216216
}
217217

218-
TRITONSERVER_Error*
219-
ModelInstanceState::ReceiveMessageFromStub(
220-
bi::managed_external_buffer::handle_t& message)
221-
{
222-
bool success = false;
223-
while (!success) {
224-
uint64_t timeout_miliseconds = 1000;
225-
{
226-
boost::posix_time::ptime timeout =
227-
boost::get_system_time() +
228-
boost::posix_time::milliseconds(timeout_miliseconds);
229-
230-
bi::scoped_lock<bi::interprocess_mutex> lock(
231-
*Stub()->HealthMutex(), timeout);
232-
233-
// Check if lock has been acquired.
234-
if (lock) {
235-
Stub()->IpcControl()->stub_health = false;
236-
} else {
237-
// If it failed to obtain the lock, it means that the stub has been
238-
// stuck or exited while holding the health mutex lock.
239-
return TRITONSERVER_ErrorNew(
240-
TRITONSERVER_ERROR_INTERNAL, "Failed to obtain the health mutex.");
241-
}
242-
}
243-
244-
message = Stub()->ParentMessageQueue()->Pop(
245-
timeout_miliseconds /* duration ms */, success);
246-
247-
if (!success && !IsStubProcessAlive()) {
248-
return TRITONSERVER_ErrorNew(
249-
TRITONSERVER_ERROR_INTERNAL, "Stub process is not healthy.");
250-
}
251-
}
252-
253-
return nullptr; // success
254-
}
255-
256218
void
257219
ModelInstanceState::RespondErrorToAllRequests(
258220
const char* message,
@@ -1251,7 +1213,7 @@ ModelInstanceState::ProcessRequests(
12511213
boost::asio::post(*thread_pool_, std::move(task));
12521214
futures_.emplace_back(std::move(future));
12531215

1254-
auto error = ReceiveMessageFromStub(response_message);
1216+
auto error = Stub()->ReceiveMessageFromStub(response_message);
12551217
if (error != nullptr) {
12561218
restart = true;
12571219
RespondErrorToAllRequests(

‎src/python_be.h‎

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -289,9 +289,6 @@ class ModelInstanceState : public BackendModelInstance {
289289
// Checks whether the stub process is live
290290
bool IsStubProcessAlive();
291291

292-
// Get a message from the stub process
293-
TRITONSERVER_Error* ReceiveMessageFromStub(off_t& message);
294-
295292
// Get a message from the stub process
296293
void SendMessageAndReceiveResponse(
297294
off_t message, off_t& response, bool& restart,
@@ -365,10 +362,10 @@ class ModelInstanceState : public BackendModelInstance {
365362
// Model instance stub
366363
std::unique_ptr<StubLauncher>& Stub() { return model_instance_stub_; }
367364

368-
// Stop the log monitor threads
365+
// Stop the stub_to_parent_queue_monitor thread
369366
void TerminateMonitor();
370367

371-
// Start the log monitor threads
368+
// Start the stub_to_parent_queue_monitor thread
372369
void StartMonitor();
373370

374371
// Send bls decoupled response to the stub process

‎src/stub_launcher.cc‎

Lines changed: 57 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -428,8 +428,11 @@ StubLauncher::ModelInstanceStubProcess()
428428
initialize_message->Args() = initialize_map_handle;
429429
stub_message_queue_->Push(initialize_message->ShmHandle());
430430

431+
bi::managed_external_buffer::handle_t message;
432+
RETURN_IF_ERROR(ReceiveMessageFromStub(message));
433+
431434
std::unique_ptr<IPCMessage> initialize_response_message =
432-
IPCMessage::LoadFromSharedMemory(shm_pool_, parent_message_queue_->Pop());
435+
IPCMessage::LoadFromSharedMemory(shm_pool_, message);
433436

434437
if (initialize_response_message->Command() != PYTHONSTUB_InitializeResponse) {
435438
return TRITONSERVER_ErrorNew(
@@ -535,4 +538,57 @@ StubLauncher::KillStubProcess()
535538
stub_pid_ = 0;
536539
}
537540

541+
TRITONSERVER_Error*
542+
StubLauncher::ReceiveMessageFromStub(
543+
bi::managed_external_buffer::handle_t& message)
544+
{
545+
bool success = false;
546+
while (!success) {
547+
uint64_t timeout_miliseconds = 1000;
548+
{
549+
boost::posix_time::ptime timeout =
550+
boost::get_system_time() +
551+
boost::posix_time::milliseconds(timeout_miliseconds);
552+
553+
bi::scoped_lock<bi::interprocess_mutex> lock(*health_mutex_, timeout);
554+
555+
// Check if lock has been acquired.
556+
if (lock) {
557+
ipc_control_->stub_health = false;
558+
} else {
559+
// If it failed to obtain the lock, it means that the stub has been
560+
// stuck or exited while holding the health mutex lock.
561+
return TRITONSERVER_ErrorNew(
562+
TRITONSERVER_ERROR_INTERNAL, "Failed to obtain the health mutex.");
563+
}
564+
}
565+
566+
message = parent_message_queue_->Pop(
567+
timeout_miliseconds /* duration ms */, success);
568+
569+
bool is_stub_alive = false;
570+
{
571+
boost::posix_time::ptime timeout =
572+
boost::get_system_time() + boost::posix_time::seconds(1);
573+
bi::scoped_lock<bi::interprocess_mutex> lock(*health_mutex_, timeout);
574+
if (lock) {
575+
is_stub_alive = ipc_control_->stub_health;
576+
} else {
577+
// If It failed to obtain the lock, it means that the stub has been
578+
// stuck or exited while holding the health mutex lock.
579+
is_stub_alive = false;
580+
}
581+
}
582+
583+
if (!success && !is_stub_alive) {
584+
return TRITONSERVER_ErrorNew(
585+
TRITONSERVER_ERROR_INTERNAL,
586+
(std::string("Stub process '") + model_instance_name_ +
587+
"' is not healthy.")
588+
.c_str());
589+
}
590+
}
591+
592+
return nullptr; // success
593+
}
538594
}}}; // namespace triton::backend::python

‎src/stub_launcher.h‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -145,6 +145,10 @@ class StubLauncher {
145145
// Kill stub process
146146
void KillStubProcess();
147147

148+
// Get a message from the stub process
149+
TRITONSERVER_Error* ReceiveMessageFromStub(
150+
bi::managed_external_buffer::handle_t& message);
151+
148152
private:
149153
pid_t parent_pid_;
150154
pid_t stub_pid_;

0 commit comments

Comments
 (0)