src/infer_request.cc - main - rnekrasov/python_backend

src/infer_request.cc
615 строк21 KB
Olga Andreeva
Exposing trace context to python backend (#346)
16 мар 2024, 00:51
16 мар 2024, 00:51474fb30
// Copyright 2021-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.

#include "infer_request.h"

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

#include "gpu_buffers.h"
#include "pb_utils.h"
#include "scoped_defer.h"
#ifdef TRITON_PB_STUB
#include "pb_stub.h"
#endif

namespace triton { namespace backend { namespace python {

InferRequest::InferRequest(
    const std::string& request_id, const CorrelationId& correlation_id,
    const std::vector<std::shared_ptr<PbTensor>>& inputs,
    const std::set<std::string>& requested_output_names,
    const std::string& model_name, const int64_t model_version,
    const std::string& parameters, const uint32_t flags, const uint64_t timeout,
    const intptr_t response_factory_address, const intptr_t request_address,
    const PreferredMemory& preferred_memory, const InferenceTrace& trace)
    : request_id_(request_id), correlation_id_(correlation_id), inputs_(inputs),
      requested_output_names_(requested_output_names), model_name_(model_name),
      model_version_(model_version), parameters_(parameters), flags_(flags),
      timeout_(timeout), response_factory_address_(response_factory_address),
      request_address_(request_address), preferred_memory_(preferred_memory),
      trace_(trace), request_release_flags_(TRITONSERVER_REQUEST_RELEASE_ALL)
{
  for (auto& input : inputs) {
    if (!input) {
      throw PythonBackendException(
          "Input tensor for request with id '" + request_id +
          "' and model name '" + model_name + "' should not be empty.");
    }
  }

  for (auto& requested_output_name : requested_output_names) {
    if (requested_output_name == "") {
      throw PythonBackendException(
          "Requested output name for request with id '" + request_id +
          "' and model name '" + model_name + "' should not be empty.");
    }
  }

  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(), pb_cancel_);
#endif
}

const std::vector<std::shared_ptr<PbTensor>>&
InferRequest::Inputs()
{
  return inputs_;
}

const std::string&
InferRequest::Parameters()
{
  return parameters_;
}

const std::string&
InferRequest::RequestId()
{
  return request_id_;
}

CorrelationId&
InferRequest::GetCorrelationId()
{
  return correlation_id_;
}

const std::set<std::string>&
InferRequest::RequestedOutputNames()
{
  return requested_output_names_;
}

const std::string&
InferRequest::ModelName()
{
  return model_name_;
}

int64_t
InferRequest::ModelVersion()
{
  return model_version_;
}

uint32_t
InferRequest::Flags()
{
  return flags_;
}

intptr_t
InferRequest::RequestAddress()
{
  return request_address_;
}

void
InferRequest::SetFlags(uint32_t flags)
{
  flags_ = flags;
}

bi::managed_external_buffer::handle_t
InferRequest::ShmHandle()
{
  return shm_handle_;
}

uint64_t
InferRequest::Timeout()
{
  return timeout_;
}

void
InferRequest::SetIsDecoupled(const bool is_decoupled)
{
  is_decoupled_ = is_decoupled;
}

bool
InferRequest::IsDecoupled()
{
  return is_decoupled_;
}

PreferredMemory&
InferRequest::GetPreferredMemory()
{
  return preferred_memory_;
}

InferenceTrace&
InferRequest::GetTrace()
{
  return trace_;
}

uint32_t
InferRequest::ReleaseFlags()
{
  request_release_flags_ = infer_request_shm_ptr_->request_release_flags;
  return request_release_flags_;
}

void
InferRequest::SetReleaseFlags(const uint32_t& flags)
{
  request_release_flags_ = flags;
  infer_request_shm_ptr_->request_release_flags = request_release_flags_;
}

void
InferRequest::SaveToSharedMemory(std::unique_ptr<SharedMemoryManager>& shm_pool)
{
  AllocatedSharedMemory<char> infer_request_shm = shm_pool->Construct<char>(
      sizeof(InferRequestShm) +
      (RequestedOutputNames().size() *
       sizeof(bi::managed_external_buffer::handle_t)) +
      (Inputs().size() * sizeof(bi::managed_external_buffer::handle_t)));

  infer_request_shm_ptr_ =
      reinterpret_cast<InferRequestShm*>(infer_request_shm.data_.get());
  infer_request_shm_ptr_->input_count = Inputs().size();
  infer_request_shm_ptr_->model_version = model_version_;
  infer_request_shm_ptr_->requested_output_count =
      RequestedOutputNames().size();
  infer_request_shm_ptr_->flags = Flags();
  infer_request_shm_ptr_->address = request_address_;
  infer_request_shm_ptr_->response_factory_address = response_factory_address_;
  infer_request_shm_ptr_->is_decoupled = is_decoupled_;
  infer_request_shm_ptr_->timeout = timeout_;
  infer_request_shm_ptr_->preferred_memory = preferred_memory_;
  infer_request_shm_ptr_->request_release_flags = request_release_flags_;

  output_names_handle_shm_ptr_ =
      reinterpret_cast<bi::managed_external_buffer::handle_t*>(
          reinterpret_cast<char*>(infer_request_shm_ptr_) +
          sizeof(InferRequestShm));

  // [FIXME] This could also be a part of the single allocated memory for this
  // object.
  size_t i = 0;
  std::vector<std::unique_ptr<PbString>> requested_output_names_shm;
  for (auto& requested_output_name : requested_output_names_) {
    std::unique_ptr<PbString> requested_output_name_shm =
        PbString::Create(shm_pool, requested_output_name);
    output_names_handle_shm_ptr_[i] = requested_output_name_shm->ShmHandle();
    requested_output_names_shm.emplace_back(
        std::move(requested_output_name_shm));
    i++;
  }

  input_tensors_handle_ptr_ =
      reinterpret_cast<bi::managed_external_buffer::handle_t*>(
          reinterpret_cast<char*>(output_names_handle_shm_ptr_) +
          sizeof(bi::managed_external_buffer::handle_t) *
              RequestedOutputNames().size());
  i = 0;
  for (auto& input : Inputs()) {
    input_tensors_handle_ptr_[i] = input->ShmHandle();
    i++;
  }

  correlation_id_.SaveToSharedMemory(shm_pool);
  infer_request_shm_ptr_->correlation_id_shm_handle =
      correlation_id_.ShmHandle();

  std::unique_ptr<PbString> model_name_shm =
      PbString::Create(shm_pool, ModelName());
  infer_request_shm_ptr_->model_name_shm_handle = model_name_shm->ShmHandle();

  std::unique_ptr<PbString> request_id_shm =
      PbString::Create(shm_pool, RequestId());
  infer_request_shm_ptr_->request_id_shm_handle = request_id_shm->ShmHandle();

  std::unique_ptr<PbString> parameters_shm =
      PbString::Create(shm_pool, Parameters());
  infer_request_shm_ptr_->parameters_shm_handle = parameters_shm->ShmHandle();

  trace_.SaveToSharedMemory(shm_pool);
  infer_request_shm_ptr_->trace_shm_handle = trace_.ShmHandle();

  // Save the references to shared memory.
  infer_request_shm_ = std::move(infer_request_shm);
  request_id_shm_ = std::move(request_id_shm);
  model_name_shm_ = std::move(model_name_shm);
  parameters_shm_ = std::move(parameters_shm);
  shm_handle_ = infer_request_shm_.handle_;
  requested_output_names_shm_ = std::move(requested_output_names_shm);
}

std::unique_ptr<InferRequest>
InferRequest::LoadFromSharedMemory(
    std::unique_ptr<SharedMemoryManager>& shm_pool,
    bi::managed_external_buffer::handle_t request_handle, bool open_cuda_handle)
{
  AllocatedSharedMemory<char> infer_request_shm =
      shm_pool->Load<char>(request_handle);
  InferRequestShm* infer_request_shm_ptr =
      reinterpret_cast<InferRequestShm*>(infer_request_shm.data_.get());

  std::vector<std::unique_ptr<PbString>> requested_output_names_shm;
  uint32_t requested_output_count =
      infer_request_shm_ptr->requested_output_count;

  bi::managed_external_buffer::handle_t* output_names_handle_shm_ptr =
      reinterpret_cast<bi::managed_external_buffer::handle_t*>(
          (reinterpret_cast<char*>(infer_request_shm_ptr) +
           sizeof(InferRequestShm)));

  for (size_t output_idx = 0; output_idx < requested_output_count;
       ++output_idx) {
    std::unique_ptr<PbString> pb_string = PbString::LoadFromSharedMemory(
        shm_pool, output_names_handle_shm_ptr[output_idx]);
    requested_output_names_shm.emplace_back(std::move(pb_string));
  }

  bi::managed_external_buffer::handle_t* input_names_handle_shm_ptr =
      reinterpret_cast<bi::managed_external_buffer::handle_t*>(
          (reinterpret_cast<char*>(infer_request_shm_ptr) +
           sizeof(InferRequestShm) +
           (infer_request_shm_ptr->requested_output_count *
            sizeof(bi::managed_external_buffer::handle_t))));

  std::vector<std::shared_ptr<PbTensor>> input_tensors;
  for (size_t input_idx = 0; input_idx < infer_request_shm_ptr->input_count;
       ++input_idx) {
    std::shared_ptr<PbTensor> input_tensor = PbTensor::LoadFromSharedMemory(
        shm_pool, input_names_handle_shm_ptr[input_idx], open_cuda_handle);
    input_tensors.emplace_back(std::move(input_tensor));
  }

  std::unique_ptr<CorrelationId> correlation_id_shm =
      CorrelationId::LoadFromSharedMemory(
          shm_pool, infer_request_shm_ptr->correlation_id_shm_handle);

  std::unique_ptr<InferenceTrace> infer_trace_shm =
      InferenceTrace::LoadFromSharedMemory(
          shm_pool, infer_request_shm_ptr->trace_shm_handle);

  std::unique_ptr<PbString> model_name_shm = PbString::LoadFromSharedMemory(
      shm_pool, infer_request_shm_ptr->model_name_shm_handle);
  std::unique_ptr<PbString> request_id_shm = PbString::LoadFromSharedMemory(
      shm_pool, infer_request_shm_ptr->request_id_shm_handle);
  std::unique_ptr<PbString> parameters_shm = PbString::LoadFromSharedMemory(
      shm_pool, infer_request_shm_ptr->parameters_shm_handle);

  return std::unique_ptr<InferRequest>(new InferRequest(
      infer_request_shm, request_id_shm, correlation_id_shm,
      requested_output_names_shm, model_name_shm, input_tensors, parameters_shm,
      infer_trace_shm));
}

InferRequest::InferRequest(
    AllocatedSharedMemory<char>& infer_request_shm,
    std::unique_ptr<PbString>& request_id_shm,
    std::unique_ptr<CorrelationId>& correlation_id_shm,
    std::vector<std::unique_ptr<PbString>>& requested_output_names_shm,
    std::unique_ptr<PbString>& model_name_shm,
    std::vector<std::shared_ptr<PbTensor>>& input_tensors,
    std::unique_ptr<PbString>& parameters_shm,
    std::unique_ptr<InferenceTrace>& infer_trace_shm)
    : infer_request_shm_(std::move(infer_request_shm)),
      request_id_shm_(std::move(request_id_shm)),
      requested_output_names_shm_(std::move(requested_output_names_shm)),
      model_name_shm_(std::move(model_name_shm)),
      parameters_shm_(std::move(parameters_shm))
{
  infer_request_shm_ptr_ =
      reinterpret_cast<InferRequestShm*>(infer_request_shm_.data_.get());
  output_names_handle_shm_ptr_ =
      reinterpret_cast<bi::managed_external_buffer::handle_t*>(
          reinterpret_cast<char*>(infer_request_shm_ptr_) +
          sizeof(InferRequestShm));
  input_tensors_handle_ptr_ =
      reinterpret_cast<bi::managed_external_buffer::handle_t*>(
          reinterpret_cast<char*>(infer_request_shm_ptr_) +
          sizeof(InferRequestShm) +
          sizeof(bi::managed_external_buffer::handle_t) *
              infer_request_shm_ptr_->requested_output_count);
  inputs_ = std::move(input_tensors);

  std::set<std::string> requested_output_names;
  for (size_t output_idx = 0;
       output_idx < infer_request_shm_ptr_->requested_output_count;
       ++output_idx) {
    auto& pb_string = requested_output_names_shm_[output_idx];
    requested_output_names.emplace(pb_string->String());
  }

  correlation_id_ = CorrelationId(correlation_id_shm);
  request_id_ = request_id_shm_->String();
  parameters_ = parameters_shm_->String();
  requested_output_names_ = std::move(requested_output_names);
  model_name_ = model_name_shm_->String();
  flags_ = infer_request_shm_ptr_->flags;
  model_version_ = infer_request_shm_ptr_->model_version;
  request_address_ = infer_request_shm_ptr_->address;
  response_factory_address_ = infer_request_shm_ptr_->response_factory_address;
  is_decoupled_ = infer_request_shm_ptr_->is_decoupled;
  timeout_ = infer_request_shm_ptr_->timeout;
  preferred_memory_ = infer_request_shm_ptr_->preferred_memory;
  trace_ = InferenceTrace(infer_trace_shm);
  request_release_flags_ = infer_request_shm_ptr_->request_release_flags;

#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(), pb_cancel_);
#endif
}

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

std::shared_ptr<ResponseSender>
InferRequest::GetResponseSender()
{
  std::unique_ptr<Stub>& stub = Stub::GetOrCreateInstance();
  if (!stub->IsDecoupled()) {
    throw PythonBackendException(
        "'get_response_sender' function must be called only when the model is "
        "using the decoupled transaction policy.");
  }

  return response_sender_;
}

std::shared_ptr<InferResponse>
InferRequest::Exec(const bool is_decoupled)
{
  // Release the GIL. This avoids a potential deadlock situation in the parent
  // process, where every thread in the thread pool is indirectly waiting for a
  // function in the stub process that acquires the GIL. Meanwhile, the current
  // thread, which holds the GIL, is also waiting for the parent side to have
  // the next available thread to pick up the job during resource contention.
  py::gil_scoped_release release;

  // BLS should not be used in "initialize" or "finalize" function.
  std::unique_ptr<Stub>& stub = Stub::GetOrCreateInstance();
  if (!stub->IsInitialized() || stub->IsFinalizing()) {
    throw PythonBackendException(
        "BLS is only supported during the 'execute' function.");
  }

  ResponseBatch* response_batch = nullptr;
  bool responses_is_set = false;
  std::unique_ptr<SharedMemoryManager>& shm_pool = stub->SharedMemory();
  bi::managed_external_buffer::handle_t* response_handle = nullptr;

  PythonBackendException pb_exception(std::string{});
  std::unique_ptr<IPCMessage> ipc_message;

  AllocatedSharedMemory<char> request_batch;
  ScopedDefer data_load_complete([&ipc_message] {
    bi::scoped_lock<bi::interprocess_mutex> lock{
        *(ipc_message->ResponseMutex())};
    ipc_message->ResponseCondition()->notify_all();
  });

  try {
    ipc_message = IPCMessage::Create(shm_pool, true /* inline_response */);
    bool has_exception = false;
    PythonBackendException pb_exception(std::string{});

    if (is_decoupled) {
      ipc_message->Command() =
          PYTHONSTUB_CommandType::PYTHONSTUB_InferStreamExecRequest;
    } else {
      ipc_message->Command() =
          PYTHONSTUB_CommandType::PYTHONSTUB_InferExecRequest;
    }

    request_batch = shm_pool->Construct<char>(
        sizeof(RequestBatch) + sizeof(bi::managed_external_buffer::handle_t));

    RequestBatch* request_batch_shm_ptr =
        reinterpret_cast<RequestBatch*>(request_batch.data_.get());
    request_batch_shm_ptr->batch_size = 1;
    ipc_message->Args() = request_batch.handle_;

    bi::managed_external_buffer::handle_t* requests_shm =
        reinterpret_cast<bi::managed_external_buffer::handle_t*>(
            request_batch.data_.get() + sizeof(RequestBatch));
    request_batch_shm_ptr->batch_size = 1;

    bool has_gpu_tensor = false;
    size_t i = 0;
    for (auto& input_tensor : inputs_) {
      input_tensor->SaveToSharedMemory(shm_pool, false /* copy_gpu */);
      if (!input_tensor->IsCPU()) {
        has_gpu_tensor = true;
      }
      ++i;
    }

    SaveToSharedMemory(shm_pool);

    // Save the shared memory offset of the request.
    *requests_shm = ShmHandle();

    // Send the BLS request to the parent process and wait for the response.
    {
      bi::scoped_lock<bi::interprocess_mutex> lock{
          *(ipc_message->ResponseMutex())};
      stub->SendIPCMessage(ipc_message);
      ipc_message->ResponseCondition()->wait(lock);
    }

    // Additional round trip required for asking the stub process
    // to fill in the GPU tensor buffers
    if (has_gpu_tensor) {
      AllocatedSharedMemory<GPUBuffersShm> gpu_buffers_shm =
          shm_pool->Load<GPUBuffersShm>(
              request_batch_shm_ptr->gpu_buffers_handle);
      AllocatedSharedMemory<bi::managed_external_buffer::handle_t>
          gpu_buffers_handle =
              shm_pool->Load<bi::managed_external_buffer::handle_t>(
                  gpu_buffers_shm.data_->buffers);
      try {
        if (!gpu_buffers_shm.data_->success) {
          std::unique_ptr<PbString> error = PbString::LoadFromSharedMemory(
              shm_pool, gpu_buffers_shm.data_->error);
          throw PythonBackendException(error->String());
        }
#ifdef TRITON_ENABLE_GPU
        size_t i = 0;
        for (auto& input_tensor : this->Inputs()) {
          if (!input_tensor->IsCPU()) {
            std::unique_ptr<PbMemory> dst_buffer =
                PbMemory::LoadFromSharedMemory(
                    shm_pool, (gpu_buffers_handle.data_.get())[i],
                    true /* open cuda handle */);
            PbMemory::CopyBuffer(dst_buffer, input_tensor->Memory());
            ++i;
          }
        }
#endif  // TRITON_ENABLE_GPU
      }
      catch (const PythonBackendException& exception) {
        // We need to catch the exception here. Otherwise, we will not notify
        // the main process and it will wait for the response forever.
        pb_exception = exception;
        has_exception = true;
      }

      {
        bi::scoped_lock<bi::interprocess_mutex> lock{
            *(ipc_message->ResponseMutex())};
        ipc_message->ResponseCondition()->notify_all();
        ipc_message->ResponseCondition()->wait(lock);
      }
    }

    // The exception will be thrown after the message was sent to the main
    // process.
    if (has_exception) {
      throw pb_exception;
    }

    // Get the response for the current message.
    std::unique_ptr<IPCMessage> bls_response = IPCMessage::LoadFromSharedMemory(
        shm_pool, ipc_message->ResponseHandle());

    AllocatedSharedMemory<char> response_batch_shm =
        shm_pool->Load<char>(bls_response->Args());
    response_batch =
        reinterpret_cast<ResponseBatch*>(response_batch_shm.data_.get());
    response_handle = reinterpret_cast<bi::managed_external_buffer::handle_t*>(
        response_batch_shm.data_.get() + sizeof(ResponseBatch));

    responses_is_set = true;
    if (response_batch->has_error) {
      if (response_batch->is_error_set) {
        std::unique_ptr<PbString> pb_string =
            PbString::LoadFromSharedMemory(shm_pool, response_batch->error);
        auto error_response = std::make_unique<InferResponse>(
            std::vector<std::shared_ptr<PbTensor>>{},
            std::make_shared<PbError>(pb_string->String()));

        return error_response;
      } else {
        auto error_response = std::make_unique<InferResponse>(
            std::vector<std::shared_ptr<PbTensor>>{},
            std::make_shared<PbError>(
                "An error occurred while performing BLS request."));

        return error_response;
      }
    }
  }
  catch (const PythonBackendException& pb_exception) {
    auto error_response = std::make_unique<InferResponse>(
        std::vector<std::shared_ptr<PbTensor>>{},
        std::make_shared<PbError>(pb_exception.what()));

    return error_response;
  }

  if (responses_is_set) {
    auto& memory_manager_message_queue = stub->MemoryManagerQueue();
    std::unique_ptr<InferResponse> return_response =
        InferResponse::LoadFromSharedMemory(
            shm_pool, *response_handle, true /* open cuda handle */);

    for (auto& output_tensor : return_response->OutputTensors()) {
      if (!output_tensor->IsCPU()) {
        uint64_t memory_release_id = output_tensor->Memory()->MemoryReleaseId();
        output_tensor->Memory()->SetMemoryReleaseCallback(
            [&memory_manager_message_queue, memory_release_id, &shm_pool]() {
              memory_manager_message_queue->Push(memory_release_id);
            });
      }
    }

    return return_response;
  } else {
    auto error_response = std::make_unique<InferResponse>(
        std::vector<std::shared_ptr<PbTensor>>{},
        std::make_shared<PbError>(
            "An error occurred while performing BLS request."));

    return error_response;
  }
}

#endif

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