infer_grpc_client.h

Go to the documentation of this file.

/*
 * SPDX-FileCopyrightText: Copyright (c) 2021-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 * SPDX-License-Identifier: LicenseRef-NvidiaProprietary
 *
 * NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
 * property and proprietary rights in and to this material, related
 * documentation and any modifications thereto. Any use, reproduction,
 * disclosure or distribution of this material and related documentation
 * without an express license agreement from NVIDIA CORPORATION or
 * its affiliates is strictly prohibited.
 */
 
#ifndef __INFER_GRPC_CLIENT_H__
#define __INFER_GRPC_CLIENT_H__
 
#include <stdarg.h>
#include <condition_variable>
#include <functional>
#include <list>
#include <memory>
#include <mutex>
#include <queue>
 
#include "infer_icontext.h"
#include "infer_post_datatypes.h"
#include "infer_common.h"
#include "grpc_client.h"
 
 
namespace tc = triton::client;
 
namespace nvdsinferserver {
 
typedef std::map<std::string, std::string> Headers;
 
class TritonGrpcRequest;
 
using SharedGrpcRequest = std::shared_ptr<TritonGrpcRequest>;
 
using TritonGrpcAsyncDone = std::function<void(NvDsInferStatus, SharedBatchArray)>;
 
 
class TritonGrpcRequest {
public:
    ~TritonGrpcRequest();
    NvDsInferStatus appendInput(const std::shared_ptr<tc::InferInput> &input);
    NvDsInferStatus setOutput(const std::vector<std::shared_ptr<tc::InferRequestedOutput>> &output);
    NvDsInferStatus setOption(std::shared_ptr<tc::InferOptions> &option);
    std::vector<std::shared_ptr<tc::InferInput>> inputs() { return m_InferInputs; }
    std::vector<std::shared_ptr<tc::InferRequestedOutput>> outputs() { return m_RequestOutputs; }
    std::shared_ptr<tc::InferOptions> getOption() {return m_InferOptions;}
    SharedIBatchArray inputBatchArray() { return m_InputBatchArray;}
    void setInputBatchArray(SharedIBatchArray inputBatch) { m_InputBatchArray = inputBatch;}
    std::vector<std::string> getOutNames() { return m_OutputNames;}
    std::vector<std::string> getInputCudaBufNames() { return m_InputCudaBufNames;}
    void setOutNames(std::vector<std::string> outnames) {m_OutputNames = outnames;}
    void attachData(void *data) {
        m_CpuData.push_back(data);
    }
 
 
private:
    std::vector<std::shared_ptr<tc::InferInput>> m_InferInputs;
    std::vector<std::shared_ptr<tc::InferRequestedOutput>> m_RequestOutputs;
    std::shared_ptr<tc::InferOptions> m_InferOptions;
    std::vector<std::string> m_OutputNames;
    SharedIBatchArray m_InputBatchArray;
    std::vector<void*> m_CpuData;
    std::vector<std::string> m_InputCudaBufNames;
};
 
class InferGrpcClient {
public:
    InferGrpcClient (std::string url, bool enableCudaBufferSharing);
    ~InferGrpcClient();
    NvDsInferStatus Initialize();
    NvDsInferStatus getModelMetadata(inference::ModelMetadataResponse *model_metadata,
        std::string &model_name, std::string &model_version);
    NvDsInferStatus getModelConfig(inference::ModelConfigResponse *config, const std::string &name,
        const std::string &version = "", const Headers &headers = Headers());
    bool isServerLive();
    bool isServerReady();
    bool isModelReady(const std::string &model, const std::string version = "");
    NvDsInferStatus LoadModel(const std::string& model_name, const Headers& headers = Headers());
    NvDsInferStatus UnloadModel(const std::string& model_name, const Headers& headers = Headers());
    SharedGrpcRequest createRequest(const std::string& model, const std::string &version, SharedIBatchArray input,
        const std::vector<std::string> &outputs,
        const std::vector<TritonClassParams>& classList = std::vector<TritonClassParams>());
 
    NvDsInferStatus inferAsync(SharedGrpcRequest request, TritonGrpcAsyncDone done);
 
private:
    void InferComplete (tc::InferResult *result, SharedGrpcRequest request,
        TritonGrpcAsyncDone done);
    NvDsInferStatus parseOptions(tc::InferOptions *outOpt, const IOptions *inOpt);
    tc::Error SetInputCudaSharedMemory(tc::InferInput *inferInput,
        const SharedBatchBuf& inbuf, SharedGrpcRequest request, uint64_t bufId);
 
private:
    std::string m_Url;
    bool m_EnableCudaBufferSharing;
    std::unique_ptr<tc::InferenceServerGrpcClient> m_GrpcClient;
    std::atomic<uint64_t> m_LastRequestId{UINT64_C(0)};
    std::unordered_map<std::string, std::string> m_RegisteredBuffersHash;
};
 
}  // namespace nvdsinferserver
 
#endif