// Copyright 2022-2024, 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 <sys/stat.h>

#include <sys/types.h>

#include <array>

#include <atomic>

#include <boost/asio.hpp>

#include <boost/asio/post.hpp>

#include <boost/asio/thread_pool.hpp>

#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 <regex>

#include <sstream>

#include <string>

#include <thread>

#include <unordered_map>

#include <vector>

#include "infer_request.h"

#include "infer_response.h"

#include "ipc_message.h"

#include "memory_manager.h"

#include "message_queue.h"

#include "metric.h"

#include "metric_family.h"

#include "pb_env.h"

#include "pb_map.h"

#include "pb_metric_reporter.h"

#include "pb_utils.h"

#include "request_executor.h"

#include "scoped_defer.h"

#include "shm_manager.h"

#include "stub_launcher.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/nvtx.h"

#include "triton/common/triton_json.h"

#include "triton/core/tritonbackend.h"

#include "triton/core/tritonserver.h"

#ifdef _WIN32

#define NOMINMAX

#include <windows.h>

#else

#include <sys/wait.h>

#include <unistd.h>

#endif

#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* raasnie_err__ = (X); \

if (raasnie_err__ != nullptr) { \

for (size_t ridx = 0; ridx < RESPONSES_COUNT; ++ridx) { \

if ((*RESPONSES)[ridx] != nullptr) { \

LOG_IF_ERROR( \

TRITONBACKEND_ResponseSend( \

(*RESPONSES)[ridx], TRITONSERVER_RESPONSE_COMPLETE_FINAL, \

raasnie_err__), \

"failed to send error response"); \

(*RESPONSES)[ridx] = nullptr; \

} \

} \

TRITONSERVER_ErrorDelete(raasnie_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()); \

for (size_t ridx = 0; ridx < RESPONSES_COUNT; ++ridx) { \

if ((*RESPONSES)[ridx] != nullptr) { \

LOG_IF_ERROR( \

TRITONBACKEND_ResponseSend( \

(*RESPONSES)[ridx], TRITONSERVER_RESPONSE_COMPLETE_FINAL, \

raarie_err__), \

"failed to send error response"); \

(*RESPONSES)[ridx] = nullptr; \

} \

} \

TRITONSERVER_ErrorDelete(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 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;

int64_t thread_pool_size;

// FIXME [DLIS-5969]: Enable for Windows when custom execution environments

// are supported.

#ifndef _WIN32

std::unique_ptr<EnvironmentManager> env_manager;

#endif

std::string runtime_modeldir;

};

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_; }

// Is decoupled API being used.

bool IsDecoupled() { return decoupled_; }

// Set decoupled mode

void SetDecoupled(bool decoupled) { decoupled_ = decoupled; }

// Returns the value in the `runtime_modeldir_` field

std::string RuntimeModelDir() { return runtime_modeldir_; }

// Launch auto-complete stub process.

TRITONSERVER_Error* LaunchAutoCompleteStubProcess();

// Validate Model Configuration

TRITONSERVER_Error* ValidateModelConfig();

// Overrides `BackendModel::SetModelConfig` to also

// set `ModelState::decoupled_`

TRITONSERVER_Error* SetModelConfig();

// Auto-complete stub

std::unique_ptr<StubLauncher>& Stub() { return auto_complete_stub_; }

private:

ModelState(TRITONBACKEND_Model* triton_model);

BackendState* backend_state_;

std::string python_execution_env_;

bool force_cpu_only_input_tensors_;

bool decoupled_;

std::string runtime_modeldir_;

std::unique_ptr<StubLauncher> auto_complete_stub_;

};

class ModelInstanceState : public BackendModelInstance {

ModelInstanceState(

ModelState* model_state, TRITONBACKEND_ModelInstance* model_instance);

TRITONBACKEND_Model* triton_model_;

std::unique_ptr<StubLauncher> model_instance_stub_;

std::vector<intptr_t> closed_requests_;

std::mutex closed_requests_mutex_;

std::thread stub_to_parent_queue_monitor_;

bool stub_to_parent_thread_;

std::mutex mu_;

std::condition_variable cv_;

std::unique_ptr<IPCMessage> received_message_;

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

std::unique_ptr<boost::asio::thread_pool> thread_pool_;

std::unordered_map<intptr_t, std::shared_ptr<InferPayload>> infer_payload_;

std::mutex infer_payload_mu_;

std::unique_ptr<RequestExecutor> request_executor_;

public:

static TRITONSERVER_Error* Create(

ModelState* model_state, TRITONBACKEND_ModelInstance* model_instance,

ModelInstanceState** model_instance_state);

~ModelInstanceState();

// Launch stub process.

TRITONSERVER_Error* LaunchStubProcess();

void ResponseSendDecoupled(std::shared_ptr<IPCMessage> response_send_message);

// The parent message queue is monitored only by this function during the

// execute phase. No other thread should pop any message from the message

// queue.

void MessageQueueMonitor();

// This function is executed on a separate thread and monitors the queue for

// message sent from stub to parent process.

void StubToParentMQMonitor();

// Process the log request.

void ProcessLogRequest(const std::unique_ptr<IPCMessage>& message);

// Convert TRITONBACKEND_Input to Python backend tensors.

TRITONSERVER_Error* GetInputTensor(

const uint32_t input_idx, std::shared_ptr<PbTensor>& input_tensor,

TRITONBACKEND_Request* request,

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

// Process all the requests in the decoupled mode.

TRITONSERVER_Error* ProcessRequests(

TRITONBACKEND_Request** requests, const uint32_t request_count,

std::vector<std::unique_ptr<InferRequest>>& pb_infer_requests,

PbMetricReporter& pb_metric_reporter);

bool ExistsInClosedRequests(intptr_t closed_request);

// Execute a BLS Request

void ExecuteBLSRequest(

std::shared_ptr<IPCMessage> ipc_message, const bool is_stream);

// Cleanup BLS responses

void CleanupBLSResponses();

// Check the incoming requests for errors

TRITONSERVER_Error* CheckIncomingRequests(

TRITONBACKEND_Request** requests, const uint32_t request_count,

size_t& total_batch_size);

// Set error for response send message

void SetErrorForResponseSendMessage(

ResponseSendMessage* response_send_message,

std::shared_ptr<TRITONSERVER_Error*> error,

std::unique_ptr<PbString>& error_message);

TRITONSERVER_Error* SaveRequestsToSharedMemory(

TRITONBACKEND_Request** requests, const uint32_t request_count,

std::vector<std::unique_ptr<InferRequest>>& pb_inference_requests,

AllocatedSharedMemory<char>& request_batch,

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

void SendMessageAndReceiveResponse(

bi::managed_external_buffer::handle_t message,

bi::managed_external_buffer::handle_t& response,

std::shared_ptr<std::vector<TRITONBACKEND_Response*>>& responses,

TRITONBACKEND_Request** requests, const uint32_t request_count);

void RespondErrorToAllRequests(

const char* message,

std::shared_ptr<std::vector<TRITONBACKEND_Response*>>& responses,

TRITONBACKEND_Request** requests, const uint32_t request_count);

// void SendMessageToStub(bi::managed_external_buffer::handle_t message);

TRITONSERVER_Error* SendMessageToStub(

bi::managed_external_buffer::handle_t message);

// Checks whether the stub process is live

bool IsStubProcessAlive();

// Model instance stub

std::unique_ptr<StubLauncher>& Stub() { return model_instance_stub_; }

// Stop the stub_to_parent_queue_monitor thread

void TerminateMonitor();

// Start the stub_to_parent_queue_monitor thread

void StartMonitor();

// Send bls decoupled response to the stub process

void SendBLSDecoupledResponse(std::unique_ptr<InferResponse> infer_response);

// Prepare the response batch object

void PrepareResponseBatch(

ResponseBatch** response_batch,

AllocatedSharedMemory<char>& response_batch_shm,

std::unique_ptr<IPCMessage>* ipc_message,

bi::managed_external_buffer::handle_t** response_handle);

// Prepare the response handle

void PrepareResponseHandle(

std::unique_ptr<InferResponse>* infer_response,

bi::managed_external_buffer::handle_t* response_handle);

// Process the decoupled cleanup request for InferPayload and ResponseFactory

void ProcessCleanupRequest(const std::unique_ptr<IPCMessage>& message);

// Process request cancellation query

void ProcessIsRequestCancelled(const std::unique_ptr<IPCMessage>& message);

// Process a message. The function 'request_handler' is invoked

// to handle the request. T should be either 'MetricFamily', 'Metric' or

// 'ModelLoader', and MessageType should be either 'MetricFamilyMessage',

// 'MetricMessage' or 'ModelLoaderMessage'.

template <typename T, typename MessageType>

void ProcessMessage(

const std::unique_ptr<IPCMessage>& message,

std::function<void(std::unique_ptr<T>&, MessageType*)> request_handler);

// Process a metric family request

void ProcessMetricFamilyRequest(const std::unique_ptr<IPCMessage>& message);

// Process a metric request

void ProcessMetricRequest(const std::unique_ptr<IPCMessage>& message);

// Process a model control request

void ProcessModelControlRequest(const std::unique_ptr<IPCMessage>& message);

// Attempt to share CUDA memory pool with the stub process

void ShareCUDAMemoryPool(const int32_t device_id);

};

}}} // namespace triton::backend::python