src/infer_request.h - main - rnekrasov/python_backend

src/infer_request.h
173 строки6 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.

#pragma once

#include <future>
#include <string>

#include "correlation_id.h"
#include "infer_response.h"
#include "infer_trace.h"
#include "pb_preferred_memory.h"
#include "pb_tensor.h"

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

namespace triton { namespace backend { namespace python {

class Stub;

//
// Inference Request
//
struct InferRequestShm {
  uint32_t input_count;
  uint32_t requested_output_count;
  int64_t model_version;
  uint32_t flags;
  intptr_t address;
  intptr_t response_factory_address;
  bool is_decoupled;
  uint64_t timeout;
  PreferredMemory preferred_memory;
  bi::managed_external_buffer::handle_t trace_shm_handle;
  uint32_t request_release_flags;
  bi::managed_external_buffer::handle_t correlation_id_shm_handle;
  bi::managed_external_buffer::handle_t model_name_shm_handle;
  bi::managed_external_buffer::handle_t request_id_shm_handle;
  bi::managed_external_buffer::handle_t parameters_shm_handle;
};

class InferRequest {
 public:
  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 = 0,
      const uint64_t timeout = 0, const intptr_t response_factory_address = 0,
      const intptr_t request_address = 0,
      const PreferredMemory& preferred_memory =
          PreferredMemory(PreferredMemory::kDefault, 0),
      const InferenceTrace& trace = InferenceTrace());

  const std::vector<std::shared_ptr<PbTensor>>& Inputs();
  const std::string& RequestId();
  const std::string& Parameters();
  CorrelationId& GetCorrelationId();
  const std::string& ModelName();
  int64_t ModelVersion();
  uint32_t Flags();
  void SetFlags(uint32_t flags);
  const std::set<std::string>& RequestedOutputNames();
  bi::managed_external_buffer::handle_t ShmHandle();
  uint64_t Timeout();
  bool IsDecoupled();
  void SetIsDecoupled(const bool is_decoupled);
  PreferredMemory& GetPreferredMemory();
  InferenceTrace& GetTrace();
  uint32_t ReleaseFlags();
  void SetReleaseFlags(const uint32_t& flags);

#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.
  /// \param shm_pool Shared memory pool to save the inference request.
  void SaveToSharedMemory(std::unique_ptr<SharedMemoryManager>& shm_pool);

  /// Create an Inference Request object from shared memory.
  /// \param shm_pool Shared memory pool
  /// \param request_handle Shared memory handle of the request.
  /// \param open_cuda_handle Determines if the tensor in the infer request
  /// object is a GPU tensor, to call the cudaIpcOpenMemHandle to obtain the
  /// tensor or not.
  /// \return Returns the infer request in the specified request_handle
  /// location.
  static std::unique_ptr<InferRequest> LoadFromSharedMemory(
      std::unique_ptr<SharedMemoryManager>& shm_pool,
      bi::managed_external_buffer::handle_t request_handle,
      bool open_cuda_handle);

  /// Disallow copying the inference request object.
  DISALLOW_COPY_AND_ASSIGN(InferRequest);

  intptr_t RequestAddress();
  ~InferRequest() {}

 private:
  InferRequest(
      AllocatedSharedMemory<char>& infer_request_shm,
      std::unique_ptr<PbString>& request_id_shm,
      std::unique_ptr<CorrelationId>& correlation_id,
      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);

  std::string request_id_;
  CorrelationId correlation_id_;
  std::vector<std::shared_ptr<PbTensor>> inputs_;
  std::set<std::string> requested_output_names_;
  std::string model_name_;
  int64_t model_version_;
  std::string parameters_;
  uint32_t flags_;
  uint64_t timeout_;
  intptr_t response_factory_address_;
  intptr_t request_address_;
  bool is_decoupled_;
  PreferredMemory preferred_memory_;
  InferenceTrace trace_;
  uint32_t request_release_flags_;

  // Shared Memory Data Structures
  AllocatedSharedMemory<char> infer_request_shm_;
  InferRequestShm* infer_request_shm_ptr_;

  std::unique_ptr<PbString> request_id_shm_;
  std::vector<std::unique_ptr<PbString>> requested_output_names_shm_;
  std::unique_ptr<PbString> model_name_shm_;
  bi::managed_external_buffer::handle_t* output_names_handle_shm_ptr_;
  bi::managed_external_buffer::handle_t* input_tensors_handle_ptr_;
  bi::managed_external_buffer::handle_t shm_handle_;
  std::unique_ptr<PbString> parameters_shm_;

#ifdef TRITON_PB_STUB
  std::shared_ptr<PbCancel> pb_cancel_;
  std::shared_ptr<ResponseSender> response_sender_;
#endif
};
}}};  // namespace triton::backend::python