model_client.cpp

#include <iostream>
#include <mutex>
#include <numeric>
#include <opencv2/opencv.hpp>
#include <string>
#include <thread>
#include <vector>

#include "model_client.hpp"
#include <cuda_runtime_api.h>

namespace tc = triton::client;

auto fail_on_error = [](const tc::Error &err, const std::string &msg) {
    if (!err.IsOk()) {
        throw std::runtime_error("model client: error: " + msg + ": " + err.Message());
    }
};

auto fail_on_cuda_error = [](cudaError_t result) {
    if (result != cudaSuccess) {
        throw std::runtime_error("model client: CUDA exception (line " + std::to_string(__LINE__) +
                                 "): " + cudaGetErrorName(result) + " (" + cudaGetErrorString(result) + ")");
    }
};

template <typename S, typename T>
ModelClient<S, T>::ModelClient(const ModelClientSettings &settings) : settings_(settings)
{
    bool verbose = false;
    fail_on_error(tc::InferenceServerGrpcClient::Create(&client_, settings_.model_url, verbose),
                  "unable to create grpc client");
}

template <typename S, typename T>
ModelClient<S, T>::~ModelClient()
{
    if (settings_.cuda_shared_mem) {
        fail_on_error(client_->UnregisterSystemSharedMemory(), "unable to unregister all system shared memory regions");
        fail_on_error(client_->UnregisterCudaSharedMemory(), "unable to unregister all cuda shared memory regions");
    }
}

template <typename S, typename T>
std::any ModelClient<S, T>::RunInference(std::vector<uint8_t> input_data)
{
    std::shared_ptr<tc::InferInput> input = CreateInput();
    std::shared_ptr<tc::InferRequestedOutput> output = CreateRequestedOutput();

    const size_t input_size = CalculateInputMemorySize();
    const size_t output_size = CalculateOutputMemorySize();

    if (input_size != input_data.size() * sizeof(uint8_t)) {
        throw std::runtime_error(
            "model client: input data is not the expected size, actual = " + std::to_string(input_data.size()) +
            ", expected = " + std::to_string(input_size / sizeof(uint8_t)));
    }

    GPUMemoryHandles handles;
    SharedMemoryRegistration registration = {"input_data" + GetThisThreadID(), "output_data" + GetThisThreadID(),
                                             input_size, output_size};
    if (settings_.cuda_shared_mem) {
        AllocateGPUMemory(handles, input_size, output_size);
        CreateIPCHandles(handles);

        RegisterSharedMemoryOnServer(registration, handles);
        RegisterSharedMemoryOnClient(input, output, registration);

        WriteInputToGPUMemory(handles.input_d_ptr, input_data, input_size);
    }
    else {
        WriteInputToCPUMemory(input, input_data, input_size);
    }

    std::shared_ptr<tc::InferResult> result = Infer(input, output);
    ValidateShapeAndDatatype(result);

    std::vector<T> output_data;
    if (settings_.cuda_shared_mem) {
        output_data = ReadResultFromGPUMemory(handles.output_d_ptr, output_size);
        UnRegisterSharedMemoryOnServer(registration);
        ReleaseGPUMemory(handles);
    }
    else {
        output_data = ReadResultFromCPUMemory(result, output_size);
    }

    return output_data;
}

template <typename S, typename T>
std::shared_ptr<tc::InferInput> ModelClient<S, T>::CreateInput()
{
    tc::InferInput *input = nullptr;
    fail_on_error(
        tc::InferInput::Create(&input, settings_.input_name, settings_.input_shape, TensorRTType<S>::GetType()),
        "unable to get " + settings_.input_name);
    std::shared_ptr<tc::InferInput> input_ptr;
    input_ptr.reset(input);
    return input_ptr;
}

template <typename S, typename T>
std::shared_ptr<tc::InferRequestedOutput> ModelClient<S, T>::CreateRequestedOutput()
{
    tc::InferRequestedOutput *output;
    fail_on_error(tc::InferRequestedOutput::Create(&output, settings_.output_name),
                  "unable to get " + settings_.output_name);
    std::shared_ptr<tc::InferRequestedOutput> output_ptr;
    output_ptr.reset(output);
    return output_ptr;
}

template <typename S, typename T>
size_t ModelClient<S, T>::CalculateInputMemorySize()
{
    size_t input_nof_elements = std::accumulate(settings_.input_shape.begin(), settings_.input_shape.end(),
                                                static_cast<size_t>(1), std::multiplies<size_t>());
    const size_t input_size = input_nof_elements * sizeof(S);
    return input_size;
}

template <typename S, typename T>
size_t ModelClient<S, T>::CalculateOutputMemorySize()
{
    size_t output_nof_elements = std::accumulate(settings_.output_shape.begin(), settings_.output_shape.end(),
                                                 static_cast<size_t>(1), std::multiplies<size_t>());
    size_t output_size = output_nof_elements * sizeof(T);
    return output_size;
}

template <typename S, typename T>
void ModelClient<S, T>::WriteInputToGPUMemory(int *input_d_ptr, const std::vector<uint8_t> &input_data,
                                              const size_t input_size)
{
    fail_on_cuda_error(
        cudaMemcpy((void *) input_d_ptr, (void *) input_data.data(), input_size, cudaMemcpyHostToDevice));
}

template <typename S, typename T>
void ModelClient<S, T>::WriteInputToCPUMemory(std::shared_ptr<tc::InferInput> &input,
                                              const std::vector<uint8_t> &input_data, const size_t input_size)
{
    fail_on_error(input->AppendRaw(input_data.data(), input_size), "unable to set data for " + settings_.input_name);
}

template <typename S, typename T>
std::shared_ptr<tc::InferResult> ModelClient<S, T>::Infer(std::shared_ptr<tc::InferInput> &input,
                                                          std::shared_ptr<tc::InferRequestedOutput> &output)
{
    std::vector<tc::InferInput *> inputs = {input.get()};
    std::vector<const tc::InferRequestedOutput *> outputs = {output.get()};

    tc::InferOptions options(settings_.model_name);
    options.model_version_ = settings_.model_version;
    options.client_timeout_ = client_timeout_;
    tc::Headers http_headers;
    grpc_compression_algorithm compression_algorithm = grpc_compression_algorithm::GRPC_COMPRESS_NONE;

    tc::InferResult *results;
    std::shared_ptr<tc::InferResult> result;
    {
        const std::lock_guard<std::mutex> lock(mutex_);
        fail_on_error(client_->Infer(&results, options, inputs, outputs, http_headers, compression_algorithm),
                      "unable to run model");
    }
    result.reset(results);
    return result;
}

template <typename S, typename T>
std::vector<T> ModelClient<S, T>::ReadResultFromGPUMemory(int *output_d_ptr, const size_t output_size)
{
    uint8_t *output_data = new uint8_t[output_size];  // allocated on heap to avoid stack overflow on large data amounts
    fail_on_cuda_error(cudaMemcpy(output_data, output_d_ptr, output_size, cudaMemcpyDeviceToHost));

    size_t num_elems = output_size / sizeof(T);
    std::vector<T> output_vec(num_elems);
    std::memcpy(output_vec.data(), output_data, output_size);

    delete[] output_data;

    return output_vec;
}

template <typename S, typename T>
std::vector<T> ModelClient<S, T>::ReadResultFromCPUMemory(std::shared_ptr<tc::InferResult> &result,
                                                          const size_t output_size)
{
    uint8_t *output_data;
    size_t received_size;
    fail_on_error(result->RawData(settings_.output_name, (const uint8_t **) &output_data, &received_size),
                  "unable to get result data for " + settings_.output_name);
    if (received_size != output_size) {
        throw std::runtime_error("error: received incorrect byte size for " + settings_.output_name + " : " +
                                 std::to_string(received_size));
    }

    size_t num_elems = output_size / sizeof(T);
    std::vector<T> output_vec(num_elems);
    std::memcpy(output_vec.data(), output_data, received_size);

    return output_vec;
}

template <typename S, typename T>
void ModelClient<S, T>::ValidateShapeAndDatatype(std::shared_ptr<tc::InferResult> &result)
{
    const std::string name = settings_.output_name;
    std::vector<int64_t> shape;
    fail_on_error(result->Shape(name, &shape), "unable to get shape for '" + name + "'");
    if (shape.size() != settings_.output_shape.size() ||
        !std::equal(shape.begin(), shape.end(), settings_.output_shape.begin())) {
        throw std::runtime_error("model-client: received incorrect shapes for " + name);
    }

    std::string datatype;
    fail_on_error(result->Datatype(name, &datatype), "unable to get datatype for '" + name + "'");
    if (datatype.compare(TensorRTType<T>::GetType()) != 0) {
        throw std::runtime_error("model-client: received incorrect datatype for " + name + " : " + datatype);
    }
}

template <typename S, typename T>
void ModelClient<S, T>::AllocateGPUMemory(GPUMemoryHandles &handles, size_t input_size, size_t output_size)
{
    fail_on_cuda_error(cudaMalloc((void **) &handles.input_d_ptr, input_size));
    fail_on_cuda_error(cudaMalloc((void **) &handles.output_d_ptr, output_size));
}

template <typename S, typename T>
void ModelClient<S, T>::CreateIPCHandles(GPUMemoryHandles &handles)
{
    CreateCUDAIPCHandle(&handles.input_cuda_handle, (void *) handles.input_d_ptr);
    CreateCUDAIPCHandle(&handles.output_cuda_handle, (void *) handles.output_d_ptr);
}

template <typename S, typename T>
void ModelClient<S, T>::CreateCUDAIPCHandle(cudaIpcMemHandle_t *cuda_handle, void *device_ptr, int device_id)
{
    fail_on_cuda_error(cudaSetDevice(device_id));
    fail_on_cuda_error(cudaIpcGetMemHandle(cuda_handle, device_ptr));
}

template <typename S, typename T>
void ModelClient<S, T>::RegisterSharedMemoryOnServer(const SharedMemoryRegistration &registration,
                                                     const GPUMemoryHandles &handles)
{
    fail_on_error(client_->RegisterCudaSharedMemory(registration.input_register_name, handles.input_cuda_handle,
                                                    0 /* device_id */, registration.input_size),
                  "failed to register input shared memory region");
    fail_on_error(client_->RegisterCudaSharedMemory(registration.output_register_name, handles.output_cuda_handle,
                                                    0 /* device_id */, registration.output_size),
                  "failed to register output shared memory region");
}

template <typename S, typename T>
void ModelClient<S, T>::UnRegisterSharedMemoryOnServer(const SharedMemoryRegistration &registration)
{
    fail_on_error(client_->UnregisterCudaSharedMemory(registration.input_register_name),
                  "unable to unregister shared memory input region");
    fail_on_error(client_->UnregisterCudaSharedMemory(registration.output_register_name),
                  "unable to unregister shared memory output region");
}

template <typename S, typename T>
void ModelClient<S, T>::RegisterSharedMemoryOnClient(std::shared_ptr<tc::InferInput> &input,
                                                     std::shared_ptr<tc::InferRequestedOutput> &output,
                                                     const SharedMemoryRegistration &registration)
{
    fail_on_error(input->SetSharedMemory(registration.input_register_name, registration.input_size, 0 /* offset */),
                  "unable to set shared memory for " + settings_.input_name);
    fail_on_error(output->SetSharedMemory(registration.output_register_name, registration.output_size, 0 /* offset */),
                  "unable to set shared memory for " + settings_.output_name);
}

template <typename S, typename T>
void ModelClient<S, T>::ReleaseGPUMemory(const GPUMemoryHandles &handles)
{
    fail_on_cuda_error(cudaFree(handles.input_d_ptr));
    fail_on_cuda_error(cudaFree(handles.output_d_ptr));
}

template <typename S, typename T>
std::string ModelClient<S, T>::GetThisThreadID()
{
    std::thread::id thread_id = std::this_thread::get_id();
    std::stringstream ss;
    ss << thread_id;
    std::string thread_id_str = ss.str();

    return thread_id_str;
}

template <>
struct TensorRTType<float> {
    static std::string GetType()
    {
        return "FP32";
    }
};

template <>
struct TensorRTType<uint8_t> {
    static std::string GetType()
    {
        return "UINT8";
    }
};

template <>
struct TensorRTType<int32_t> {
    static std::string GetType()
    {
        return "INT32";
    }
};

// Supported types are float and uint8_t on both input and output
template class ModelClient<float, uint8_t>;
template class ModelClient<float, float>;
template class ModelClient<uint8_t, uint8_t>;
template class ModelClient<uint8_t, float>;
template class ModelClient<uint8_t, int32_t>;