// Copyright 2020-2021, 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 <pthread.h>

#include <sys/stat.h>

#include <sys/types.h>

#include <sys/vfs.h>

#include <sys/wait.h>

#include <unistd.h>

#include <array>

#include <atomic>

#include <boost/functional/hash.hpp>

#include <boost/interprocess/sync/interprocess_condition.hpp>

#include <boost/interprocess/sync/interprocess_mutex.hpp>

#include <boost/interprocess/sync/scoped_lock.hpp>

#include <boost/thread/thread_time.hpp>

#include <chrono>

#include <csignal>

#include <cstdio>

#include <cstring>

#include <ctime>

#include <functional>

#include <future>

#include <memory>

#include <numeric>

#include <sstream>

#include <string>

#include <thread>

#include <vector>

#include "infer_request.h"

#include "infer_response.h"

#include "ipc_message.h"

#include "message_queue.h"

#include "pb_env.h"

#include "pb_main_utils.h"

#include "pb_tensor.h"

#include "pb_utils.h"

#include "shm_manager.h"

#include "triton/backend/backend_common.h"

#include "triton/backend/backend_input_collector.h"

#include "triton/backend/backend_memory.h"

#include "triton/backend/backend_model.h"

#include "triton/backend/backend_model_instance.h"

#include "triton/common/triton_json.h"

#include "triton/core/tritonbackend.h"

#include "triton/core/tritonserver.h"

#define LOG_IF_EXCEPTION(X) \

do { \

try { \

(X); \

} \

catch (const PythonBackendException& pb_exception) { \

LOG_MESSAGE(TRITONSERVER_LOG_ERROR, pb_exception.what()); \

} \

} while (false)

#define RESPOND_ALL_AND_RETURN_IF_ERROR(RESPONSES, RESPONSES_COUNT, X) \

do { \

TRITONSERVER_Error* raarie_err__ = (X); \

if (raarie_err__ != nullptr) { \

SendErrorForResponses(RESPONSES, RESPONSES_COUNT, raarie_err__); \

return; \

} \

} while (false)

#define RESPOND_ALL_AND_RETURN_IF_EXCEPTION(RESPONSES, RESPONSES_COUNT, X) \

do { \

try { \

(X); \

} \

catch (const PythonBackendException& exception) { \

TRITONSERVER_Error* raarie_err__ = TRITONSERVER_ErrorNew( \

TRITONSERVER_ERROR_INTERNAL, exception.what()); \

SendErrorForResponses(RESPONSES, RESPONSES_COUNT, raarie_err__); \

return; \

} \

} while (false)

#define RESPOND_AND_RETURN_IF_ERROR(REQUEST, X) \

do { \

TRITONSERVER_Error* rarie_err__ = (X); \

if (rarie_err__ != nullptr) { \

TRITONBACKEND_Response* rarie_response__ = nullptr; \

LOG_IF_ERROR( \

TRITONBACKEND_ResponseNew(&rarie_response__, REQUEST), \

"failed to create response"); \

if (rarie_response__ != nullptr) { \

LOG_IF_ERROR( \

TRITONBACKEND_ResponseSend( \

rarie_response__, TRITONSERVER_RESPONSE_COMPLETE_FINAL, \

rarie_err__), \

"failed to send error response"); \

} \

return rarie_err__; \

} \

} while (false)

#define RESPOND_ALL_REQUESTS_AND_RETURN_IF_EXCEPTION( \

REQUESTS, REQUEST_COUNT, X) \

do { \

try { \

(X); \

} \

catch (const PythonBackendException& exception) { \

TRITONSERVER_Error* rarie_err__ = TRITONSERVER_ErrorNew( \

TRITONSERVER_ERROR_INTERNAL, exception.what()); \

RequestsRespondWithError(REQUESTS, REQUEST_COUNT, rarie_err__); \

return; \

} \

} while (false)

#define GUARDED_RESPOND_IF_ERROR(RESPONSES, IDX, X) \

do { \

if ((RESPONSES)[IDX] != nullptr) { \

TRITONSERVER_Error* err__ = (X); \

if (err__ != nullptr) { \

LOG_IF_ERROR( \

TRITONBACKEND_ResponseSend( \

(RESPONSES)[IDX], TRITONSERVER_RESPONSE_COMPLETE_FINAL, \

err__), \

"failed to send error response"); \

(RESPONSES)[IDX] = nullptr; \

TRITONSERVER_ErrorDelete(err__); \

} \

} \

} while (false)

#define GUARDED_RESPOND_IF_EXCEPTION(RESPONSES, IDX, X) \

do { \

if ((RESPONSES)[IDX] != nullptr) { \

try { \

(X); \

} \

catch (const PythonBackendException& pb_exception) { \

TRITONSERVER_Error* err__ = TRITONSERVER_ErrorNew( \

TRITONSERVER_ERROR_INTERNAL, pb_exception.what()); \

LOG_IF_ERROR( \

TRITONBACKEND_ResponseSend( \

(RESPONSES)[IDX], TRITONSERVER_RESPONSE_COMPLETE_FINAL, \

err__), \

"failed to send error response"); \

(RESPONSES)[IDX] = nullptr; \

TRITONSERVER_ErrorDelete(err__); \

} \

} \

} while (false)

#define RETURN_IF_EXCEPTION(X) \

do { \

try { \

(X); \

} \

catch (const PythonBackendException& pb_exception) { \

TRITONSERVER_Error* rarie_err__ = TRITONSERVER_ErrorNew( \

TRITONSERVER_ERROR_INTERNAL, pb_exception.what()); \

return rarie_err__; \

} \

} while (false)

namespace triton { namespace backend { namespace python {

namespace bi = boost::interprocess;

struct BackendState {

std::string python_lib;

int64_t shm_default_byte_size;

int64_t shm_growth_byte_size;

int64_t stub_timeout_seconds;

int64_t shm_message_queue_size;

std::atomic<int> number_of_instance_inits;

std::string shared_memory_region_prefix;

std::unique_ptr<EnvironmentManager> env_manager;

};

class ModelState : public BackendModel {

public:

static TRITONSERVER_Error* Create(

TRITONBACKEND_Model* triton_model, ModelState** state);

// Get backend state

BackendState* StateForBackend() { return backend_state_; }

// Get the Python execution environment

std::string PythonExecutionEnv() { return python_execution_env_; }

// Force CPU only tensors

bool ForceCPUOnlyInputTensors() { return force_cpu_only_input_tensors_; }

private:

ModelState(TRITONBACKEND_Model* triton_model);

BackendState* backend_state_;

std::string python_execution_env_;

bool force_cpu_only_input_tensors_;

};

class ModelInstanceState : public BackendModelInstance {

ModelInstanceState(

ModelState* model_state, TRITONBACKEND_ModelInstance* model_instance);

TRITONBACKEND_Model* triton_model_;

bi::interprocess_mutex* health_mutex_;

std::unique_ptr<MessageQueue> stub_message_queue_;

std::unique_ptr<MessageQueue> parent_message_queue_;

off_t ipc_control_offset_;

std::string model_path_;

IPCControl* ipc_control_;

std::vector<std::future<void>> bls_futures_;

std::vector<TRITONSERVER_InferenceResponse*> bls_inference_responses_;

std::mutex bls_responses_mutex_;

std::unique_ptr<SharedMemory> shm_pool_;

std::string shm_region_name_;

off_t shm_reset_offset_;

std::vector<std::unique_ptr<InferResponse>> infer_responses_;

std::vector<std::unique_ptr<RequestExecutor>> request_executors_;

std::vector<std::future<void>> handles_;

// Stub process pid

pid_t stub_pid_;

// Parent process pid

pid_t parent_pid_;

bool initialized_;

// Path to python execution environment

std::string path_to_libpython_;

std::string path_to_activate_;

public:

static TRITONSERVER_Error* Create(

ModelState* model_state, TRITONBACKEND_ModelInstance* model_instance,

ModelInstanceState** model_instance_state);

~ModelInstanceState();

// Load Triton inputs to the appropriate Protobufs

TRITONSERVER_Error* GetInputTensor(

const uint32_t input_idx, Tensor* input_tensor_shm,

std::shared_ptr<PbTensor>& input_tensor, TRITONBACKEND_Request* request,

std::vector<TRITONBACKEND_Response*>& responses);

void ProcessRequests(

TRITONBACKEND_Request** requests, const uint32_t request_count,

bool& restart, std::atomic<bool>& cleanup);

// Create the stub process.

TRITONSERVER_Error* SetupStubProcess();

TRITONSERVER_Error* SendMessageToStub(off_t message);

void CleanupBLSResponses();

void WaitForBLSRequestsToFinish();

// Checks whether the stub process is live

bool IsStubProcessAlive();

// Get a message from the stub process

TRITONSERVER_Error* ReceiveMessageFromStub(off_t& message);

// Get a message from the stub process

void SendMessageAndReceiveResponse(

off_t message, off_t& response, bool& restart,

std::vector<TRITONBACKEND_Response*>& responses,

TRITONBACKEND_Request** requests, const uint32_t request_count);

// Responds to all the requests with an error message.

void RespondErrorToAllRequests(

const char* message, std::vector<TRITONBACKEND_Response*>& responses,

TRITONBACKEND_Request** requests, const uint32_t request_count);

// Kill stub process

void KillStubProcess();

// Start stub process

TRITONSERVER_Error* StartStubProcess();

// Reset the shared memory offset

void ResetSharedMemoryOffset();

void ExecuteBLSRequest(

std::unique_ptr<SharedMemory>& shm_pool, off_t message_offset,

std::atomic<bool>& cleanup);

};

ModelInstanceState::ModelInstanceState(

ModelState* model_state, TRITONBACKEND_ModelInstance* triton_model_instance)

: BackendModelInstance(model_state, triton_model_instance), stub_pid_(0),

initialized_(false)

{

}

TRITONSERVER_Error*

ModelInstanceState::Create(

ModelState* model_state, TRITONBACKEND_ModelInstance* triton_model_instance,

ModelInstanceState** state)

{

try {

*state = new ModelInstanceState(model_state, triton_model_instance);

}

catch (const BackendModelInstanceException& ex) {

RETURN_ERROR_IF_TRUE(

ex.err_ == nullptr, TRITONSERVER_ERROR_INTERNAL,

std::string("unexpected nullptr in BackendModelInstanceException"));

RETURN_IF_ERROR(ex.err_);

}

return nullptr;

}

void

ModelInstanceState::KillStubProcess()

{

kill(stub_pid_, SIGKILL);

int status;

waitpid(stub_pid_, &status, 0);

stub_pid_ = 0;

}

void

ModelInstanceState::WaitForBLSRequestsToFinish()

{

bls_futures_.clear();

}

void

ModelInstanceState::SendMessageAndReceiveResponse(

off_t message, off_t& response, bool& restart,

std::vector<TRITONBACKEND_Response*>& responses,

TRITONBACKEND_Request** requests, const uint32_t request_count)

{

auto error = SendMessageToStub(message);

if (error != nullptr) {

restart = true;

RespondErrorToAllRequests(

TRITONSERVER_ErrorMessage(error), responses, requests, request_count);

return;

}

off_t response_message;

error = ReceiveMessageFromStub(response_message);

if (error != nullptr) {

restart = true;

RespondErrorToAllRequests(

TRITONSERVER_ErrorMessage(error), responses, requests, request_count);

return;

}

response = response_message;

}

TRITONSERVER_Error*

ModelInstanceState::SendMessageToStub(off_t message)

{

bool success = false;

while (!success) {

uint64_t timeout_miliseconds = 1000;

{

boost::posix_time::ptime timeout =

boost::get_system_time() +

boost::posix_time::milliseconds(timeout_miliseconds);

bi::scoped_lock<bi::interprocess_mutex> lock(*health_mutex_, timeout);

// Check if lock has been acquired.

if (lock) {

ipc_control_->stub_health = false;

} else {

// If it failed to obtain the lock, it means that the stub has been

// stuck or exited while holding the health mutex lock.

return TRITONSERVER_ErrorNew(

TRITONSERVER_ERROR_INTERNAL, "Failed to obtain the health mutex.");

}

}

stub_message_queue_->Push(

message, timeout_miliseconds /* duration ms */, success);

if (!success && !IsStubProcessAlive()) {

return TRITONSERVER_ErrorNew(

TRITONSERVER_ERROR_INTERNAL, "Stub process is not healthy.");

}

}

return nullptr; // success

}

TRITONSERVER_Error*

ModelInstanceState::ReceiveMessageFromStub(off_t& message)

{

bool success = false;

while (!success) {

uint64_t timeout_miliseconds = 1000;

{

boost::posix_time::ptime timeout =

boost::get_system_time() +

boost::posix_time::milliseconds(timeout_miliseconds);

bi::scoped_lock<bi::interprocess_mutex> lock(*health_mutex_, timeout);

// Check if lock has been acquired.

if (lock) {

ipc_control_->stub_health = false;

} else {

// If it failed to obtain the lock, it means that the stub has been

// stuck or exited while holding the health mutex lock.

return TRITONSERVER_ErrorNew(

TRITONSERVER_ERROR_INTERNAL, "Failed to obtain the health mutex.");

}

}

message = parent_message_queue_->Pop(

timeout_miliseconds /* duration ms */, success);

if (!success && !IsStubProcessAlive()) {

return TRITONSERVER_ErrorNew(

TRITONSERVER_ERROR_INTERNAL, "Stub process is not healthy.");

}

}

return nullptr; // success

}

void

ModelInstanceState::RespondErrorToAllRequests(

const char* message, std::vector<TRITONBACKEND_Response*>& responses,

TRITONBACKEND_Request** requests, const uint32_t request_count)

{

for (uint32_t r = 0; r < request_count; ++r) {

if (responses[r] == nullptr)

continue;

std::string err_message =

std::string(

"Failed to process the request(s) for model instance '" + Name() +

"', message: ") +

message;

TRITONSERVER_Error* err =

TRITONSERVER_ErrorNew(TRITONSERVER_ERROR_INTERNAL, err_message.c_str());

LOG_IF_ERROR(

TRITONBACKEND_ResponseSend(

responses[r], TRITONSERVER_RESPONSE_COMPLETE_FINAL, err),

"failed sending response");

responses[r] = nullptr;

TRITONSERVER_ErrorDelete(err);

}

}

void

ModelInstanceState::ResetSharedMemoryOffset()

{

shm_pool_->SetOffset(shm_reset_offset_);

}

void

ModelInstanceState::ExecuteBLSRequest(

std::unique_ptr<SharedMemory>& shm_pool, off_t message_offset,

std::atomic<bool>& cleanup)

{

ModelState* model_state = reinterpret_cast<ModelState*>(Model());

auto request_executor =

std::make_unique<RequestExecutor>(model_state->TritonServer());

std::unique_ptr<IPCMessage> ipc_message =

IPCMessage::LoadFromSharedMemory(shm_pool, message_offset);

bool is_response_batch_set = false;

ResponseBatch* response_batch;

TRITONSERVER_InferenceResponse* inference_response = nullptr;

try {

std::unique_ptr<IPCMessage> bls_response =

std::make_unique<IPCMessage>(shm_pool_, false /* inline_response */);

RequestBatch* request_batch;

shm_pool_->MapOffset((char**)&request_batch, ipc_message->Args());

bls_response->Command() = PYTHONSTUB_InferExecResponse;

ipc_message->RequestOffset() = bls_response->SharedMemoryOffset();

shm_pool_->Map(

(char**)&response_batch, sizeof(ResponseBatch), bls_response->Args());

response_batch->batch_size = 1;

response_batch->has_error = false;

response_batch->is_error_set = false;

response_batch->cleanup = false;

is_response_batch_set = true;

bool has_gpu_tensor = false;

PythonBackendException pb_exception(std::string{});

if (request_batch->batch_size == 1) {

std::unique_ptr<InferRequest> infer_request;

std::shared_ptr<std::mutex> cuda_ipc_mutex;

infer_request = InferRequest::LoadFromSharedMemory(

shm_pool_, request_batch->requests, cuda_ipc_mutex, cuda_ipc_mutex);

std::unique_ptr<InferResponse> infer_response;

// If the BLS inputs are in GPU an additional round trip between the

// stub process and the main process is required. The reason is that we

// need to first allocate the GPU memory from the memory pool and then

// ask the stub process to fill in those allocated buffers.

for (auto& input_tensor : infer_request->Inputs()) {

if (!input_tensor->IsCPU()) {

#ifdef TRITON_ENABLE_GPU

BackendMemory* backend_memory;

std::unique_ptr<BackendMemory> lbackend_memory;

has_gpu_tensor = true;

TRITONSERVER_Error* error = BackendMemory::Create(

Model()->TritonMemoryManager(),

{BackendMemory::AllocationType::GPU_POOL,

BackendMemory::AllocationType::GPU},

input_tensor->MemoryTypeId(), input_tensor->ByteSize(),

&backend_memory);

if (error != nullptr) {

LOG_MESSAGE(

TRITONSERVER_LOG_ERROR, TRITONSERVER_ErrorMessage(error));

break;

}

lbackend_memory.reset(backend_memory);

input_tensor->SetBackendMemory(std::move(lbackend_memory), shm_pool_);

#endif // TRITON_ENABLE_GPU

}

}

// Wait for the extra round trip to complete. The stub process will fill

// in the data for the GPU tensors. If there is an error, the extra round

// trip must be still completed, otherwise the stub process will always be

// waiting for a message from the parent process.

if (has_gpu_tensor) {

bi::scoped_lock<bi::interprocess_mutex> lock{

*(ipc_message->ResponseMutex())};

ipc_message->ResponseCondition()->notify_all();

ipc_message->ResponseCondition()->wait(lock);

}

if (pb_exception.what() != nullptr) {

infer_response = request_executor->Infer(

infer_request, shm_pool_, &inference_response);

if (inference_response != nullptr) {

std::lock_guard<std::mutex> lock{bls_responses_mutex_};

bls_inference_responses_.push_back(inference_response);

}

Response* response;

shm_pool_->Map(

(char**)&response, sizeof(Response), response_batch->responses);

infer_response->SaveToSharedMemory(

shm_pool_, response, false /* copy_cpu */, true /* copy_gpu */);

} else {

throw pb_exception;

}

}

}

catch (const PythonBackendException& pb_exception) {

if (is_response_batch_set) {

response_batch->has_error = true;

off_t string_offset = 0;

LOG_IF_EXCEPTION(SaveStringToSharedMemory(

shm_pool_, string_offset, pb_exception.what()));

if (string_offset != 0) {

response_batch->is_error_set = true;

response_batch->error = string_offset;

}

} else {

LOG_MESSAGE(TRITONSERVER_LOG_ERROR, pb_exception.what());

}

}

{

bi::scoped_lock<bi::interprocess_mutex> lock{

*(ipc_message->ResponseMutex())};

ipc_message->ResponseCondition()->notify_all();

}

request_executors_.emplace_back(std::move(request_executor));

}

void

ModelInstanceState::ProcessRequests(

TRITONBACKEND_Request** requests, const uint32_t request_count,

bool& restart, std::atomic<bool>& cleanup)

{

ModelState* model_state = reinterpret_cast<ModelState*>(Model());

int max_batch_size = model_state->MaxBatchSize();

std::string name = model_state->Name();

// For each request collect the total batch size for this inference

// execution. The batch-size, number of inputs, and size of each

// input has already been checked so don't need to do that here.

size_t total_batch_size = 0;

for (size_t i = 0; i < request_count; i++) {

// If we get a nullptr request then something is badly wrong. Fail

// and release all requests.

if (requests[i] == nullptr) {

RequestsRespondWithError(

requests, request_count,

TRITONSERVER_ErrorNew(

TRITONSERVER_ERROR_INTERNAL,

std::string(

"null request given to Python backend for '" + name + "'")

.c_str()));

return;

}

if (max_batch_size > 0) {

// Retrieve the batch size from one of the inputs, if the model

// supports batching, the first dimension size is batch size

TRITONBACKEND_Input* input;

TRITONSERVER_Error* err =

TRITONBACKEND_RequestInputByIndex(requests[i], 0 /* index */, &input);

if (err == nullptr) {

const int64_t* shape;

err = TRITONBACKEND_InputProperties(

input, nullptr, nullptr, &shape, nullptr, nullptr, nullptr);

total_batch_size += shape[0];

}

if (err != nullptr) {

RequestsRespondWithError(requests, request_count, err);

return;

}

} else {

++total_batch_size;

}

}

// If there are no valid payloads then no need to run the inference.

if (total_batch_size == 0) {

return;

}

// Make sure the maximum batch size is not exceeded. The

// total_batch_size must be 1 for models that don't support batching

// (i.e. max_batch_size == 0). If max_batch_size is exceeded then

// scheduler has done something badly wrong so fail and release all

// requests.

if ((total_batch_size != 1) && (total_batch_size > (size_t)max_batch_size)) {

RequestsRespondWithError(

requests, request_count,

TRITONSERVER_ErrorNew(

TRITONSERVER_ERROR_INTERNAL,

std::string(

"batch size " + std::to_string(total_batch_size) + " for '" +

name + "', max allowed is " + std::to_string(max_batch_size))

.c_str()));

return;

}

LOG_MESSAGE(

TRITONSERVER_LOG_VERBOSE,

(std::string("model ") + model_state->Name() + ", instance " + Name() +

", executing " + std::to_string(request_count) + " requests")

.c_str());

uint64_t exec_start_ns = 0;

SET_TIMESTAMP(exec_start_ns);

std::unique_ptr<IPCMessage> ipc_message =

std::make_unique<IPCMessage>(shm_pool_, false /*inline_resposne*/);

ipc_message->Command() = PYTHONSTUB_CommandType::PYTHONSTUB_ExecuteRequest;

RequestBatch* request_batch;

off_t request_batch_offset;

RESPOND_ALL_REQUESTS_AND_RETURN_IF_EXCEPTION(

requests, request_count,

shm_pool_->Map(

(char**)&request_batch, sizeof(RequestBatch), request_batch_offset));

ipc_message->Args() = request_batch_offset;

request_batch->batch_size = request_count;

Request* requests_shm;

off_t requests_shm_offset;

RESPOND_ALL_REQUESTS_AND_RETURN_IF_EXCEPTION(

requests, request_count,

shm_pool_->Map(

(char**)&requests_shm, sizeof(Request) * request_count,

requests_shm_offset));

request_batch->requests = requests_shm_offset;

// We take the responsibilty of the responses.

std::vector<TRITONBACKEND_Response*> responses;

responses.reserve(request_count);

for (size_t i = 0; i < request_count; i++) {

TRITONBACKEND_Response* response;

auto err = TRITONBACKEND_ResponseNew(&response, requests[i]);

if (err == nullptr) {

responses.emplace_back(response);

} else {

responses.emplace_back(nullptr);

LOG_MESSAGE(TRITONSERVER_LOG_ERROR, "Fail to create response.");

TRITONSERVER_ErrorDelete(err);

}

}

for (uint32_t r = 0; r < request_count; ++r) {

TRITONBACKEND_Request* request = requests[r];

Request* python_infer_request = &requests_shm[r];

uint32_t requested_input_count = 0;

RESPOND_ALL_AND_RETURN_IF_ERROR(

&responses, request_count,

TRITONBACKEND_RequestInputCount(request, &requested_input_count));

uint32_t requested_output_count = 0;

RESPOND_ALL_AND_RETURN_IF_ERROR(

&responses, request_count,

TRITONBACKEND_RequestOutputCount(request, &requested_output_count));

python_infer_request->requested_output_count = requested_output_count;

Tensor* input_tensors;

off_t input_tensors_offset;

RESPOND_ALL_AND_RETURN_IF_EXCEPTION(

&responses, request_count,

shm_pool_->Map(

(char**)&input_tensors, sizeof(Tensor) * requested_input_count,

input_tensors_offset));

python_infer_request->inputs = input_tensors_offset;

std::vector<std::shared_ptr<PbTensor>> pb_input_tensors;

for (size_t iidx = 0; iidx < requested_input_count; ++iidx) {

Tensor* input_tensor = &input_tensors[iidx];

std::shared_ptr<PbTensor> pb_input_tensor;

RESPOND_ALL_AND_RETURN_IF_ERROR(

&responses, request_count,

GetInputTensor(

iidx, input_tensor, pb_input_tensor, request, responses));

pb_input_tensors.emplace_back(std::move(pb_input_tensor));

}

std::vector<std::string> requested_output_names;

// Append the list of requested outputs to the inference_request

for (size_t iidx = 0; iidx < requested_output_count; ++iidx) {

const char* requested_output_name;

RESPOND_ALL_AND_RETURN_IF_ERROR(

&responses, request_count,

TRITONBACKEND_RequestOutputName(

request, iidx, &requested_output_name));

requested_output_names.emplace_back(requested_output_name);

}

// request id

const char* id;

RESPOND_ALL_AND_RETURN_IF_ERROR(

&responses, request_count, TRITONBACKEND_RequestId(request, &id));

uint64_t correlation_id;

RESPOND_ALL_AND_RETURN_IF_ERROR(

&responses, request_count,

TRITONBACKEND_RequestCorrelationId(request, &correlation_id));

InferRequest infer_request(

id, correlation_id, pb_input_tensors, requested_output_names,

model_state->Name(), model_state->Version());

RESPOND_ALL_AND_RETURN_IF_EXCEPTION(

&responses, request_count,

infer_request.SaveToSharedMemory(shm_pool_, python_infer_request));

}

uint64_t compute_start_ns = 0;

SET_TIMESTAMP(compute_start_ns);

// This means that the stub process has exited and Python

// backend failed to restart the stub process.

if (stub_pid_ == 0) {

const char* error_message = "The stub process has exited unexpectedly.";

RespondErrorToAllRequests(

error_message, responses, requests, request_count);

return;

}

off_t response_message;

SendMessageAndReceiveResponse(

ipc_message->SharedMemoryOffset(), response_message, restart, responses,

requests, request_count);

if (restart) {

return;

}

RESPOND_ALL_AND_RETURN_IF_EXCEPTION(

&responses, request_count,

ipc_message =

IPCMessage::LoadFromSharedMemory(shm_pool_, response_message));

// If the stub command is no longer PYTHONSTUB_InferExecRequest, it

// indicates that inference request exeuction has finished.

while (ipc_message->Command() ==

PYTHONSTUB_CommandType::PYTHONSTUB_InferExecRequest) {

off_t current_message = response_message;

// Launch the BLS request in a future.

bls_futures_.emplace_back(

std::async(std::launch::async, [this, current_message, &cleanup]() {

this->ExecuteBLSRequest(this->shm_pool_, current_message, cleanup);

}));

auto error = ReceiveMessageFromStub(response_message);

if (error != nullptr) {

restart = true;

RespondErrorToAllRequests(

TRITONSERVER_ErrorMessage(error), responses, requests, request_count);

return;

}

RESPOND_ALL_AND_RETURN_IF_EXCEPTION(

&responses, request_count,

ipc_message =

IPCMessage::LoadFromSharedMemory(shm_pool_, response_message));

}

uint64_t compute_end_ns = 0;

SET_TIMESTAMP(compute_end_ns);

// Parsing the request response

ResponseBatch* response_batch;

RESPOND_ALL_AND_RETURN_IF_EXCEPTION(

&responses, request_count,

shm_pool_->MapOffset((char**)&response_batch, ipc_message->Args()));

// If inference fails, release all the requests and send an error response.

// If inference fails at this stage, it usually indicates a bug in the model

// code

if (response_batch->has_error) {

if (response_batch->is_error_set) {

char* error_message;

RESPOND_ALL_AND_RETURN_IF_EXCEPTION(

&responses, request_count,

LoadStringFromSharedMemory(

shm_pool_, response_batch->error, error_message));

RespondErrorToAllRequests(

error_message, responses, requests, request_count);

} else {

const char* error_message =

"Failed to fetch the error in response batch.";

RespondErrorToAllRequests(

error_message, responses, requests, request_count);

}

return;

}

Response* responses_shm;

RESPOND_ALL_AND_RETURN_IF_EXCEPTION(

&responses, request_count,

shm_pool_->MapOffset((char**)&responses_shm, response_batch->responses));

bool has_gpu_output = false;

// The vector that stores the tensor pairs for the tensors that the stub

// provides in GPU but the output buffer provided by Triton is in CPU.

std::vector<std::pair<std::shared_ptr<PbTensor>, std::pair<void*, uint32_t>>>

tensor_buffer_pairs;

for (uint32_t r = 0; r < request_count; ++r) {

TRITONBACKEND_Response* response = responses[r];

TRITONBACKEND_Request* request = requests[r];

uint32_t requested_output_count = 0;

std::unique_ptr<InferResponse> infer_response;

try {

std::shared_ptr<std::mutex> cuda_ipc_mutex;

infer_response = InferResponse::LoadFromSharedMemory(

shm_pool_, response_batch->responses + sizeof(Response) * r,

cuda_ipc_mutex /* cuda_ipc_open_mutex */,

cuda_ipc_mutex /* cuda_ipc_close_mutex */);

if (infer_response->HasError()) {

if (infer_response->IsErrorMessageSet()) {

TRITONSERVER_Error* err = TRITONSERVER_ErrorNew(

TRITONSERVER_ERROR_INTERNAL,

infer_response->Error()->Message().c_str());

LOG_IF_ERROR(

TRITONBACKEND_ResponseSend(

responses[r], TRITONSERVER_RESPONSE_COMPLETE_FINAL, err),

"failed sending response");

TRITONSERVER_ErrorDelete(err);

} else {

const char* err_string = "Failed to process response.";

TRITONSERVER_Error* err =

TRITONSERVER_ErrorNew(TRITONSERVER_ERROR_INTERNAL, err_string);

LOG_IF_ERROR(

TRITONBACKEND_ResponseSend(

responses[r], TRITONSERVER_RESPONSE_COMPLETE_FINAL, err),

"failed sending response");

TRITONSERVER_ErrorDelete(err);

}

responses[r] = nullptr;

// If has_error is true, we do not look at the response even if the

// response is set.

continue;

}

}

catch (const PythonBackendException& pb_exception) {

TRITONSERVER_Error* err = TRITONSERVER_ErrorNew(

TRITONSERVER_ERROR_INTERNAL, pb_exception.what());

LOG_IF_ERROR(

TRITONBACKEND_ResponseSend(

responses[r], TRITONSERVER_RESPONSE_COMPLETE_FINAL, err),

"failed sending response");

TRITONSERVER_ErrorDelete(err);

responses[r] = nullptr;

continue;

}

GUARDED_RESPOND_IF_ERROR(

responses, r,

TRITONBACKEND_RequestOutputCount(request, &requested_output_count));

bool cuda_copy = false;

std::set<std::string> requested_output_names;

for (size_t j = 0; j < requested_output_count; ++j) {

const char* output_name;

GUARDED_RESPOND_IF_ERROR(

responses, r,

TRITONBACKEND_RequestOutputName(request, j, &output_name));

requested_output_names.insert(output_name);

}

for (auto& output_tensor : infer_response->OutputTensors()) {

if (requested_output_names.find(output_tensor->Name()) ==

requested_output_names.end()) {

continue;

}

TRITONSERVER_MemoryType src_memory_type = output_tensor->MemoryType();

int64_t src_memory_type_id = output_tensor->MemoryTypeId();

TRITONSERVER_MemoryType actual_memory_type = src_memory_type;

int64_t actual_memory_type_id = src_memory_type_id;

if (actual_memory_type == TRITONSERVER_MEMORY_GPU)

has_gpu_output = true;

TRITONBACKEND_Output* response_output;

GUARDED_RESPOND_IF_ERROR(

responses, r,

TRITONBACKEND_ResponseOutput(

response, &response_output, output_tensor->Name().c_str(),

static_cast<TRITONSERVER_DataType>(output_tensor->TritonDtype()),

output_tensor->Dims().data(), output_tensor->Dims().size()));

void* buffer;

bool cuda_used = false;

GUARDED_RESPOND_IF_ERROR(

responses, r,

TRITONBACKEND_OutputBuffer(

response_output, &buffer, output_tensor->ByteSize(),

&actual_memory_type, &actual_memory_type_id));

if (src_memory_type == TRITONSERVER_MEMORY_GPU &&

actual_memory_type == TRITONSERVER_MEMORY_GPU) {

#ifdef TRITON_ENABLE_GPU

cudaSetDevice(output_tensor->MemoryTypeId());

cudaIpcMemHandle_t cuda_ipc_mem_handle;

cudaError_t err = cudaIpcGetMemHandle(&cuda_ipc_mem_handle, buffer);

output_tensor->SetCudaIpcMemHandle(&cuda_ipc_mem_handle);

if (err != cudaSuccess) {

GUARDED_RESPOND_IF_ERROR(

responses, r,

TRITONSERVER_ErrorNew(

TRITONSERVER_ERROR_INTERNAL,

std::string(

"failed to get cuda ipc handle: " +

std::string(cudaGetErrorString(err)))

.c_str()));

} else {

output_tensor->SetDataPtr(buffer);

output_tensor->RawDataShm()->offset =

output_tensor->GetGPUPointerOffset();

}

#endif

}

if (src_memory_type == TRITONSERVER_MEMORY_GPU &&

(actual_memory_type == TRITONSERVER_MEMORY_CPU ||