Skip to content

Commit

Permalink
Add Python backend request cancellation (#304)
Browse files Browse the repository at this point in the history
* Add cancelled response status

* Add request cancellation

* Check cancellation on response factory if available

* Remove unnecessary wrapping

* Throw error instead of log error

* Add is cancelled check at response sender

* Enable more reuse on request cancellation and improve model interface

* Documentation wording updates

* Copyright year update

* Rollback response sender auto close on cancel

* Rollback non-decoupled any response on cancel

* Decoupled final flag docs update
  • Loading branch information
kthui authored Oct 6, 2023
1 parent b136bf3 commit 67ca860
Show file tree
Hide file tree
Showing 14 changed files with 329 additions and 10 deletions.
2 changes: 2 additions & 0 deletions CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -208,6 +208,8 @@ set(
src/pb_stub.cc
src/pb_response_iterator.h
src/pb_response_iterator.cc
src/pb_cancel.cc
src/pb_cancel.h
)

list(APPEND
Expand Down
36 changes: 36 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,7 @@ any C++ code.
- [`execute`](#execute)
- [Default Mode](#default-mode)
- [Error Handling](#error-handling)
- [Request Cancellation Handling](#request-cancellation-handling)
- [Decoupled mode](#decoupled-mode)
- [Use Cases](#use-cases)
- [Known Issues](#known-issues)
Expand Down Expand Up @@ -502,6 +503,36 @@ Supported error codes:
* `pb_utils.TritonError.UNAVAILABLE`
* `pb_utils.TritonError.UNSUPPORTED`
* `pb_utils.TritonError.ALREADY_EXISTS`
* `pb_utils.TritonError.CANCELLED` (since 23.10)

#### Request Cancellation Handling

One or more requests may be cancelled by the client during execution. Starting
from 23.10, `request.is_cancelled()` returns whether the request is cancelled or
not. For example:

```python
import triton_python_backend_utils as pb_utils

class TritonPythonModel:
...

def execute(self, requests):
responses = []

for request in requests:
if request.is_cancelled():
responses.append(pb_utils.InferenceResponse(
error=pb_utils.TritonError("Message", pb_utils.TritonError.CANCELLED)))
else:
...

return responses
```

Although checking for request cancellation is optional, it is recommended to
check for cancellation at strategic request execution stages that can early
terminate the execution in the event of its response is no longer needed.

#### Decoupled mode

Expand Down Expand Up @@ -543,6 +574,11 @@ request. After setting errors for an pb_utils.InferenceResponse
object, use InferenceResponseSender.send() to send response with the
error back to the user.

Starting from 23.10, request cancellation can be checked directly on the
`InferenceResponseSender` object using `response_sender.is_cancelled()`. Sending
the TRITONSERVER_RESPONSE_COMPLETE_FINAL flag at the end of response is still
needed even the request is cancelled.

##### Use Cases

The decoupled mode is powerful and supports various other use cases:
Expand Down
14 changes: 12 additions & 2 deletions src/infer_request.cc
Original file line number Diff line number Diff line change
Expand Up @@ -71,9 +71,11 @@ InferRequest::InferRequest(
inputs_ = inputs;
requested_output_names_ = requested_output_names;
#ifdef TRITON_PB_STUB
pb_cancel_ =
std::make_shared<PbCancel>(response_factory_address_, request_address_);
response_sender_ = std::make_shared<ResponseSender>(
request_address_, response_factory_address_,
Stub::GetOrCreateInstance()->SharedMemory());
Stub::GetOrCreateInstance()->SharedMemory(), pb_cancel_);
#endif
}

Expand Down Expand Up @@ -379,9 +381,11 @@ InferRequest::InferRequest(
trace_ = infer_request_shm_ptr_->trace;

#ifdef TRITON_PB_STUB
pb_cancel_ =
std::make_shared<PbCancel>(response_factory_address_, request_address_);
response_sender_ = std::make_shared<ResponseSender>(
request_address_, response_factory_address_,
Stub::GetOrCreateInstance()->SharedMemory());
Stub::GetOrCreateInstance()->SharedMemory(), pb_cancel_);
#endif
}

Expand All @@ -400,6 +404,12 @@ InferRequest::DeleteResponseFactory()
#endif

#ifdef TRITON_PB_STUB
bool
InferRequest::IsCancelled()
{
return pb_cancel_->IsCancelled();
}

std::shared_ptr<ResponseSender>
InferRequest::GetResponseSender()
{
Expand Down
3 changes: 3 additions & 0 deletions src/infer_request.h
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,7 @@
#include "pb_tensor.h"

#ifdef TRITON_PB_STUB
#include "pb_cancel.h"
#include "response_sender.h"
#endif

Expand Down Expand Up @@ -107,6 +108,7 @@ class InferRequest {
#ifdef TRITON_PB_STUB
std::shared_ptr<InferResponse> Exec(const bool is_decoupled);
std::shared_ptr<ResponseSender> GetResponseSender();
bool IsCancelled();
#endif

/// Save an Inference Request to shared memory.
Expand Down Expand Up @@ -173,6 +175,7 @@ class InferRequest {
std::unique_ptr<PbString> parameters_shm_;

#ifdef TRITON_PB_STUB
std::shared_ptr<PbCancel> pb_cancel_;
std::shared_ptr<ResponseSender> response_sender_;
#endif
};
Expand Down
3 changes: 2 additions & 1 deletion src/ipc_message.h
Original file line number Diff line number Diff line change
Expand Up @@ -62,7 +62,8 @@ typedef enum PYTHONSTUB_commandtype_enum {
PYTHONSTUB_MetricRequestSet,
PYTHONSTUB_LoadModelRequest,
PYTHONSTUB_UnloadModelRequest,
PYTHONSTUB_ModelReadinessRequest
PYTHONSTUB_ModelReadinessRequest,
PYTHONSTUB_IsRequestCancelled
} PYTHONSTUB_CommandType;

///
Expand Down
90 changes: 90 additions & 0 deletions src/pb_cancel.cc
Original file line number Diff line number Diff line change
@@ -0,0 +1,90 @@
// Copyright 2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions
// are met:
// * Redistributions of source code must retain the above copyright
// notice, this list of conditions and the following disclaimer.
// * Redistributions in binary form must reproduce the above copyright
// notice, this list of conditions and the following disclaimer in the
// documentation and/or other materials provided with the distribution.
// * Neither the name of NVIDIA CORPORATION nor the names of its
// contributors may be used to endorse or promote products derived
// from this software without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
// EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR
// CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
// EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
// PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
// PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
// OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

#include "pb_cancel.h"

#include "pb_stub.h"

namespace triton { namespace backend { namespace python {

void
PbCancel::SaveToSharedMemory(std::unique_ptr<SharedMemoryManager>& shm_pool)
{
cancel_shm_ = shm_pool->Construct<IsCancelledMessage>();
new (&(cancel_shm_.data_->mu)) bi::interprocess_mutex;
new (&(cancel_shm_.data_->cv)) bi::interprocess_condition;
cancel_shm_.data_->waiting_on_stub = false;
cancel_shm_.data_->response_factory_address = response_factory_address_;
cancel_shm_.data_->request_address = request_address_;
cancel_shm_.data_->is_cancelled = is_cancelled_;
}

bi::managed_external_buffer::handle_t
PbCancel::ShmHandle()
{
return cancel_shm_.handle_;
}

IsCancelledMessage*
PbCancel::ShmPayload()
{
return cancel_shm_.data_.get();
}

bool
PbCancel::IsCancelled()
{
std::unique_lock<std::mutex> lk(mu_);
// The cancelled flag can only move from false to true, not the other way, so
// it is checked on each query until cancelled and then implicitly cached.
if (is_cancelled_) {
return is_cancelled_;
}
if (!updating_) {
std::unique_ptr<Stub>& stub = Stub::GetOrCreateInstance();
if (!stub->StubToParentServiceActive()) {
LOG_ERROR << "Cannot communicate with parent service";
return false;
}
stub->EnqueueIsCancelled(this);
updating_ = true;
}
cv_.wait(lk, [this] { return !updating_; });
return is_cancelled_;
}

void
PbCancel::ReportIsCancelled(bool is_cancelled)
{
{
std::lock_guard<std::mutex> lk(mu_);
is_cancelled_ = is_cancelled;
updating_ = false;
}
cv_.notify_all();
}

}}} // namespace triton::backend::python
64 changes: 64 additions & 0 deletions src/pb_cancel.h
Original file line number Diff line number Diff line change
@@ -0,0 +1,64 @@
// Copyright 2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions
// are met:
// * Redistributions of source code must retain the above copyright
// notice, this list of conditions and the following disclaimer.
// * Redistributions in binary form must reproduce the above copyright
// notice, this list of conditions and the following disclaimer in the
// documentation and/or other materials provided with the distribution.
// * Neither the name of NVIDIA CORPORATION nor the names of its
// contributors may be used to endorse or promote products derived
// from this software without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
// EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR
// CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
// EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
// PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
// PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
// OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

#pragma once

#include <condition_variable>
#include <mutex>

#include "pb_utils.h"

namespace triton { namespace backend { namespace python {

class PbCancel {
public:
PbCancel(intptr_t response_factory_address, intptr_t request_address)
: updating_(false), response_factory_address_(response_factory_address),
request_address_(request_address), is_cancelled_(false)
{
}
DISALLOW_COPY_AND_ASSIGN(PbCancel);

void SaveToSharedMemory(std::unique_ptr<SharedMemoryManager>& shm_pool);
bi::managed_external_buffer::handle_t ShmHandle();
IsCancelledMessage* ShmPayload();

bool IsCancelled();
void ReportIsCancelled(bool is_cancelled);

private:
AllocatedSharedMemory<IsCancelledMessage> cancel_shm_;

std::mutex mu_;
std::condition_variable cv_;
bool updating_;

intptr_t response_factory_address_;
intptr_t request_address_;
bool is_cancelled_;
};

}}}; // namespace triton::backend::python
51 changes: 49 additions & 2 deletions src/pb_stub.cc
Original file line number Diff line number Diff line change
Expand Up @@ -945,6 +945,9 @@ Stub::ServiceStubToParentRequests()
SendLogMessage(utils_msg_payload);
} else if (utils_msg_payload->command_type == PYTHONSTUB_CleanupRequest) {
SendCleanupId(utils_msg_payload);
} else if (
utils_msg_payload->command_type == PYTHONSTUB_IsRequestCancelled) {
SendIsCancelled(utils_msg_payload);
} else {
std::cerr << "Error when sending message via stub_to_parent message "
"buffer - unknown command\n";
Expand Down Expand Up @@ -1028,6 +1031,44 @@ Stub::EnqueueCleanupId(void* id)
}
}

void
Stub::EnqueueIsCancelled(PbCancel* pb_cancel)
{
std::unique_ptr<UtilsMessagePayload> utils_msg_payload =
std::make_unique<UtilsMessagePayload>(
PYTHONSTUB_IsRequestCancelled, reinterpret_cast<void*>(pb_cancel));
EnqueueUtilsMessage(std::move(utils_msg_payload));
}

void
Stub::SendIsCancelled(std::unique_ptr<UtilsMessagePayload>& utils_msg_payload)
{
PbCancel* pb_cancel =
reinterpret_cast<PbCancel*>(utils_msg_payload->utils_message_ptr);
pb_cancel->SaveToSharedMemory(shm_pool_);

IsCancelledMessage* message_payload = pb_cancel->ShmPayload();
std::unique_ptr<IPCMessage> ipc_message =
IPCMessage::Create(shm_pool_, false /* inline_response */);
ipc_message->Command() = utils_msg_payload->command_type;
ipc_message->Args() = pb_cancel->ShmHandle();

bool is_cancelled = false;
{
bi::scoped_lock<bi::interprocess_mutex> lk(message_payload->mu);

SendIPCUtilsMessage(ipc_message);
while (!message_payload->waiting_on_stub) {
message_payload->cv.wait(lk);
}

is_cancelled = message_payload->is_cancelled;
message_payload->waiting_on_stub = false;
message_payload->cv.notify_all();
}
pb_cancel->ReportIsCancelled(is_cancelled);
}

bool
Stub::StubToParentServiceActive()
{
Expand Down Expand Up @@ -1364,6 +1405,7 @@ PYBIND11_EMBEDDED_MODULE(c_python_backend_utils, module)
.value(
"ALREADY_EXISTS",
TRITONSERVER_Error_Code::TRITONSERVER_ERROR_ALREADY_EXISTS)
.value("CANCELLED", TRITONSERVER_Error_Code::TRITONSERVER_ERROR_CANCELLED)
.export_values();
triton_error.def_property_readonly_static(
"UNKNOWN",
Expand All @@ -1386,6 +1428,9 @@ PYBIND11_EMBEDDED_MODULE(c_python_backend_utils, module)
triton_error.def_property_readonly_static(
"ALREADY_EXISTS",
[](py::object /* self */) { return TRITONSERVER_ERROR_ALREADY_EXISTS; });
triton_error.def_property_readonly_static(
"CANCELLED",
[](py::object /* self */) { return TRITONSERVER_ERROR_CANCELLED; });
triton_error.def(
py::init<const std::string&, TRITONSERVER_Error_Code>(),
py::arg("message").none(false),
Expand Down Expand Up @@ -1501,7 +1546,8 @@ PYBIND11_EMBEDDED_MODULE(c_python_backend_utils, module)
.def(
"requested_output_names", &InferRequest::RequestedOutputNames,
py::return_value_policy::reference_internal)
.def("get_response_sender", &InferRequest::GetResponseSender);
.def("get_response_sender", &InferRequest::GetResponseSender)
.def("is_cancelled", &InferRequest::IsCancelled);

py::class_<PbTensor, std::shared_ptr<PbTensor>>(module, "Tensor")
.def(py::init(&PbTensor::FromNumpy))
Expand Down Expand Up @@ -1539,7 +1585,8 @@ PYBIND11_EMBEDDED_MODULE(c_python_backend_utils, module)
module, "InferenceResponseSender")
.def(
"send", &ResponseSender::Send, py::arg("response") = nullptr,
py::arg("flags") = 0);
py::arg("flags") = 0)
.def("is_cancelled", &ResponseSender::IsCancelled);

py::class_<ResponseIterator, std::shared_ptr<ResponseIterator>>(
module, "ResponseIterator")
Expand Down
Loading

0 comments on commit 67ca860

Please sign in to comment.