Program Listing for File grpc_health.h

↰ Return to documentation for file (riva/services/grpc_health.h)

/*
 * Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
 *
 * NVIDIA CORPORATION and its licensors retain all intellectual property
 * and proprietary rights in and to this software, related documentation
 * and any modifications thereto.  Any use, reproduction, disclosure or
 * distribution of this software and related documentation without an express
 * license agreement from NVIDIA CORPORATION is strictly prohibited.
 */

#ifndef RIVA_SRC_SERVICES_GRPC_HEALTH_H_
#define RIVA_SRC_SERVICES_GRPC_HEALTH_H_

#include <string>

#include "riva/proto/health.grpc.pb.h"
#include "riva/trtis/grpc_client.h"

using grpc::ServerContext;
using grpc::ServerWriter;
using grpc::Status;

namespace nrh = grpc::health::v1;
namespace ni = nvidia::inferenceserver;
namespace nic = nvidia::inferenceserver::client;

class HealthImpl final : public nrh::Health::Service {
 public:
  explicit HealthImpl(std::string uri);

  Status Check(
      ServerContext* context, const nrh::HealthCheckRequest* request,
      nrh::HealthCheckResponse* response) override;

  Status Watch(
      ServerContext* context, const nrh::HealthCheckRequest* request,
      ServerWriter<nrh::HealthCheckResponse>* stream) override;

 private:
  std::string uri_;
  std::map<const std::string, nrh::HealthCheckResponse::ServingStatus> status_map_;
  // using client object to query for Triton's liveness
  std::unique_ptr<nic::InferenceServerGrpcClient> client_;
  bool init_;
};


#endif