AthInfer::TritonTool Class Reference

#include <TritonTool.h>

Inheritance diagram for AthInfer::TritonTool:

Collaboration diagram for AthInfer::TritonTool:

Public Member Functions
	TritonTool (const std::string &type, const std::string &name, const IInterface *parent)
StatusCode	initialize () override final
virtual StatusCode	inference (InputDataMap &inputData, OutputDataMap &outputData) const override final
void	print () const override final

Protected Member Functions
	TritonTool ()=delete
	TritonTool (const TritonTool &)=delete
TritonTool &	operator= (const TritonTool &)=delete

Protected Attributes
StringProperty	m_modelName {this, "ModelName", "", "Model name"}
IntegerProperty	m_port {this, "Port", 8001, "Port ID for Triton server"}
StringProperty	m_modelVersion {this, "ModelVersion", "", "Model version, empty for latest"}
FloatProperty	m_clientTimeout {this, "ClientTimeout", 0, "Client timeout in milliseconds, 0 for no timeout"}
StringProperty	m_url {this, "URL", "", "Triton URL"}
BooleanProperty	m_useSSL {this, "UseSSL", false, "Use SSL for Triton server connection"}

Private Member Functions
tc::InferenceServerGrpcClient *	getClient () const
template<typename T >
StatusCode	prepareInput (const std::string &name, const std::vector< int64_t > &shape, const std::vector< T > &data, std::vector< std::shared_ptr< tc::InferInput >> &inputs) const
template<typename T >
StatusCode	extractOutput (const std::string &name, const std::shared_ptr< tc::InferResult > &result, std::vector< T > &outputVec) const

Private Attributes
std::unique_ptr< tc::InferOptions >	m_options

Detailed Description

Definition at line 30 of file TritonTool.h.

Constructor & Destructor Documentation

TritonTool() [1/3]

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

Definition at line 7 of file TritonTool.cxx.

    : base_class(type, name, parent)
{
    declareInterface<AthInfer::IAthInferenceTool>(this);
}

TritonTool() [2/3]

AthInfer::TritonTool::TritonTool ( )

protecteddelete

TritonTool() [3/3]

AthInfer::TritonTool::TritonTool ( const TritonTool &  )

protecteddelete

Member Function Documentation

extractOutput()

template<typename T >

StatusCode AthInfer::TritonTool::extractOutput ( const std::string &  name, const std::shared_ptr< tc::InferResult > &  result, std::vector< T > &  outputVec  ) const

private

getClient()

tc::InferenceServerGrpcClient * AthInfer::TritonTool::getClient ( ) const

private

Definition at line 24 of file TritonTool.cxx.

                                                                 {
    thread_local std::unique_ptr<tc::InferenceServerGrpcClient> threadClient;
    if (!threadClient) {
        std::string url = m_url.value() + ":" + std::to_string(m_port); // always use the gRPC port
 
        bool verbose = false;
 
        tc::Error err = tc::InferenceServerGrpcClient::Create(&threadClient, url, verbose, m_useSSL);
        if (!err.IsOk()) {
            ATH_MSG_ERROR("Failed to create Triton gRPC client for model: " + m_modelName.value() + " at url: " + url);
            ATH_MSG_ERROR("Error message: " + err.Message());
            return nullptr;
        }
 
        ATH_MSG_INFO("Triton client created for model: "+ m_modelName.value() + " at url: "+ url);
 
    }
    return threadClient.get();
}

inference()

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

finaloverridevirtual

Definition at line 44 of file TritonTool.cxx.

                                                                                                 {
 
    // Create the tensor for the input data.
    // Use shared_ptr to manage the memory of the InferInput objects.
    std::vector<std::shared_ptr<tc::InferInput> > inputs_;
    inputs_.reserve(inputData.size());
 
    for (auto& [inputName, inputInfo]: inputData) {
        const std::vector<int64_t>& inputShape = inputInfo.first;
        const auto& variant = inputInfo.second;
 
        const auto status = std::visit([&](const auto& dataVec) {
            using T = std::decay_t<decltype(dataVec[0])>;
            return prepareInput<T>(inputName, inputShape, dataVec, inputs_);
        }, variant);
 
        if (status != StatusCode::SUCCESS) return status;
    }
 
   // construct raw points for inference
    std::vector<tc::InferInput*> rawInputs;
    for (auto& input: inputs_) {
        rawInputs.push_back(input.get());
    }
 
    // perform the inference.
    tc::InferResult* rawResultPtr = nullptr;
    tc::Headers http_headers;
    grpc_compression_algorithm compression_algorithm =
        grpc_compression_algorithm::GRPC_COMPRESS_NONE;
 
    FAIL_IF_ERR(
     getClient()->Infer(
        &rawResultPtr, *m_options, rawInputs, {}, http_headers, compression_algorithm),
        "unable to run model "+ m_modelName.value() + " error: " + err.Message()
    );
 
    std::shared_ptr<tc::InferResult> results(rawResultPtr);
 
    // Get the result of the inference.
    for (auto& [outputName, outputInfo]: outputData) {
        auto& variant = outputInfo.second;
 
        const auto status = std::visit([&](auto& dataVec) {
            using T = std::decay_t<decltype(dataVec[0])>;
            return extractOutput<T>(outputName, results, dataVec);
        }, variant);
 
        if (status != StatusCode::SUCCESS) return status;
    }
    return StatusCode::SUCCESS;
}

initialize()

StatusCode AthInfer::TritonTool::initialize ( )

finaloverride

Definition at line 15 of file TritonTool.cxx.

                                          {
 
    m_options = std::make_unique<tc::InferOptions>(m_modelName.value());
    m_options->model_version_ = m_modelVersion;
    m_options->client_timeout_ = m_clientTimeout;
 
    return getClient()? StatusCode::SUCCESS : StatusCode::FAILURE;
}

operator=()

TritonTool& AthInfer::TritonTool::operator= ( const TritonTool &  )

protecteddelete

prepareInput()

template<typename T >

StatusCode AthInfer::TritonTool::prepareInput ( const std::string &  name, const std::vector< int64_t > &  shape, const std::vector< T > &  data, std::vector< std::shared_ptr< tc::InferInput >> &  inputs  ) const

private

print()

void AthInfer::TritonTool::print ( ) const

inlinefinaloverride

Definition at line 40 of file TritonTool.h.

{} // nothing to print, but required by the interface.

Member Data Documentation

m_clientTimeout

FloatProperty AthInfer::TritonTool::m_clientTimeout {this, "ClientTimeout", 0, "Client timeout in milliseconds, 0 for no timeout"}

protected

Definition at line 50 of file TritonTool.h.

m_modelName

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

protected

Definition at line 47 of file TritonTool.h.

m_modelVersion

StringProperty AthInfer::TritonTool::m_modelVersion {this, "ModelVersion", "", "Model version, empty for latest"}

protected

Definition at line 49 of file TritonTool.h.

m_options

std::unique_ptr<tc::InferOptions> AthInfer::TritonTool::m_options

private

Definition at line 56 of file TritonTool.h.

m_port

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

protected

Definition at line 48 of file TritonTool.h.

m_url

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

protected

Definition at line 51 of file TritonTool.h.

m_useSSL

BooleanProperty AthInfer::TritonTool::m_useSSL {this, "UseSSL", false, "Use SSL for Triton server connection"}

protected

Definition at line 52 of file TritonTool.h.

---

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

• TritonTool.h
• TritonTool.cxx