AthInfer::TritonTool Class Reference

#include <TritonTool.h>

Inheritance diagram for AthInfer::TritonTool:

Collaboration diagram for AthInfer::TritonTool:

Classes
struct	Impl

Public Member Functions
	TritonTool (const std::string &type, const std::string &name, const IInterface *parent)
	Constructor.
virtual	~TritonTool ()
	Destructor.
Function(s) inherited from @c AthAlgTool
virtual StatusCode	initialize () override
	Initialize the tool.
Function(s) inherited from @c IAthInferenceTool
virtual StatusCode	inference (InputDataMap &inputData, OutputDataMap &outputData) const override final
	Run inference with multiple inputs and multiple outputs.
virtual void	print () const override
	Print the tool's properties and configuration.

Private Attributes
std::unique_ptr< Impl >	m_impl
	Pointer to the implementation details.
Tool properties
StringProperty	m_modelName {this, "ModelName", "", "Model name"}
IntegerProperty	m_port {this, "Port", 8001, "Port ID for Triton server"}
StringProperty	m_modelVersion
FloatProperty	m_clientTimeout
StringProperty	m_url {this, "URL", "", "Triton URL"}
BooleanProperty	m_useSSL

Detailed Description

Definition at line 14 of file TritonTool.h.

Constructor & Destructor Documentation

TritonTool()

AthInfer::TritonTool::TritonTool	(	const std::string &	type,
		const std::string &	name,
		const IInterface *	parent )

Constructor.

Definition at line 123 of file TritonTool.cxx.

    : base_class(type, name, parent) {}

~TritonTool()

AthInfer::TritonTool::~TritonTool ( )

virtualdefault

Destructor.

Member Function Documentation

inference()

StatusCode AthInfer::TritonTool::inference ( InputDataMap & inputData, OutputDataMap & outputData ) const

finaloverridevirtual

Run inference with multiple inputs and multiple outputs.

Definition at line 168 of file TritonTool.cxx.

                                                                            {
 
  assert(m_impl);
 
  // Create the tensor for the input data.
  // Use shared_ptr to manage the memory of the InferInput objects.
  std::vector<std::unique_ptr<tc::InferInput>> inputs;
  inputs.reserve(inputData.size());
 
  for (auto& [inputName, inputInfo] : inputData) {
 
    const std::vector<int64_t>& inputShape = inputInfo.first;
    const DataVariant& variant = inputInfo.second;
 
    ATH_CHECK(std::visit(
        [&](const auto& dataVec) {
          using T = std::decay_t<decltype(dataVec[0])>;
          return m_impl->prepareInput<T>(inputName, inputShape, dataVec,
                                         inputs);
        },
        variant));
  }
 
  // construct raw points for inference
  std::vector<tc::InferInput*> rawInputs;
  for (auto& input : inputs) {
    rawInputs.push_back(input.get());
  }
 
  // Get the triton client object.
  tc::InferenceServerGrpcClient* client = nullptr;
  ATH_CHECK(m_impl->getClient(client, m_url, m_port, m_useSSL));
  assert(client != nullptr);
 
  // perform the inference.
  std::shared_ptr<tc::InferResult> results;
  tc::Headers http_headers;
  grpc_compression_algorithm compression_algorithm =
      grpc_compression_algorithm::GRPC_COMPRESS_NONE;
 
  if (m_impl->m_parentAsyncAlg == nullptr) {
    tc::InferResult* rawResultPtr = nullptr;
    TRITON_CHECK(client->Infer(&rawResultPtr, *(m_impl->m_options), rawInputs,
                               {}, http_headers, compression_algorithm));
    assert(rawResultPtr != nullptr);
    results.reset(rawResultPtr);
  } else {
    // If m_impl->m_parentAsyncAlg is set, use asynchronous inference
    using Promise_t = boost::fibers::promise<tc::InferResult*>;
    using Future_t = boost::fibers::future<tc::InferResult*>;
    Promise_t promise{};
    Future_t future = promise.get_future();
    auto callback = [&promise](tc::InferResult* resultPtr) {
      assert(resultPtr != nullptr);
      promise.set_value(resultPtr);
    };
    TRITON_CHECK(client->AsyncInfer(callback, *(m_impl->m_options), rawInputs,
                                    {}, http_headers, compression_algorithm));
    results.reset(future.get());
    ATH_CHECK(m_impl->m_parentAsyncAlg->restoreAfterSuspend());
  }
 
  // Get the result of the inference.
  for (auto& [outputName, outputInfo] : outputData) {
 
    DataVariant& variant = outputInfo.second;
 
    ATH_CHECK(std::visit(
        [&](auto& dataVec) {
          using T = std::decay_t<decltype(dataVec[0])>;
          return m_impl->extractOutput<T>(outputName, *results, dataVec);
        },
        variant));
  }
 
  // Return gracefully.
  return StatusCode::SUCCESS;
}

initialize()

StatusCode AthInfer::TritonTool::initialize ( )

overridevirtual

Initialize the tool.

Definition at line 129 of file TritonTool.cxx.

                                  {
 
  // Set up the implementation object.
  m_impl = std::make_unique<Impl>(name() + "::Impl");
  m_impl->m_options = std::make_unique<tc::InferOptions>(m_modelName.value());
  m_impl->m_options->model_version_ = m_modelVersion;
  m_impl->m_options->client_timeout_ = m_clientTimeout;
 
  // Figure out if parent is an AthAsynchronousAlgorithm, and set pointer if it
  // is
  const IAlgTool* p = dynamic_cast<const IAlgTool*>(this);
  // Follow chain of parents up until we hit one that can't be converted to an
  // IAlgTool
  const IInterface* myParent = nullptr;
  while (p != nullptr) {
    myParent = p->parent();
    p = dynamic_cast<const IAlgTool*>(myParent);
  }
  // If this ultimate ancestor can be converted to an AthAsynchronousAlgorithm,
  // set the member variable
  m_impl->m_parentAsyncAlg =
      dynamic_cast<const AthAsynchronousAlgorithm*>(myParent);
  if (m_impl->m_parentAsyncAlg != nullptr) {
    ATH_MSG_INFO(
        "Owned by an AthAsynchronousAlgorithm, using asynchronous inference");
  } else {
    ATH_MSG_INFO(
        "Not owned by an AthAsynchronousAlgorithm, not using asynchronous "
        "inference");
  }
 
  // Make sure already during initialization that a client can be created.
  tc::InferenceServerGrpcClient* dummyClient = nullptr;
  ATH_CHECK(m_impl->getClient(dummyClient, m_url, m_port, m_useSSL));
 
  // Return gracefully.
  return StatusCode::SUCCESS;
}

print()

void AthInfer::TritonTool::print ( ) const

overridevirtual

Print the tool's properties and configuration.

Definition at line 248 of file TritonTool.cxx.

{}

Member Data Documentation

m_clientTimeout

FloatProperty AthInfer::TritonTool::m_clientTimeout

private

Initial value:

{
      this, "ClientTimeout", 0,
      "Client timeout in milliseconds, 0 for no timeout"}

Definition at line 51 of file TritonTool.h.

                               {
      this, "ClientTimeout", 0,
      "Client timeout in milliseconds, 0 for no timeout"};

m_impl

std::unique_ptr<Impl> AthInfer::TritonTool::m_impl

private

Pointer to the implementation details.

Definition at line 63 of file TritonTool.h.

m_modelName

StringProperty AthInfer::TritonTool::m_modelName {this, "ModelName", "", "Model name"}

private

Definition at line 47 of file TritonTool.h.

{this, "ModelName", "", "Model name"};

m_modelVersion

StringProperty AthInfer::TritonTool::m_modelVersion

private

Initial value:

{this, "ModelVersion", "",
                                "Model version, empty for latest"}

Definition at line 49 of file TritonTool.h.

                               {this, "ModelVersion", "",
                                "Model version, empty for latest"};

m_port

IntegerProperty AthInfer::TritonTool::m_port {this, "Port", 8001, "Port ID for Triton server"}

private

Definition at line 48 of file TritonTool.h.

{this, "Port", 8001, "Port ID for Triton server"};

m_url

StringProperty AthInfer::TritonTool::m_url {this, "URL", "", "Triton URL"}

private

Definition at line 54 of file TritonTool.h.

{this, "URL", "", "Triton URL"};

m_useSSL

BooleanProperty AthInfer::TritonTool::m_useSSL

private

Initial value:

{this, "UseSSL", false,
                           "Use SSL for Triton server connection"}

Definition at line 55 of file TritonTool.h.

                          {this, "UseSSL", false,
                           "Use SSL for Triton server connection"};

---

The documentation for this class was generated from the following files:

• TritonTool.h
• TritonTool.cxx