Support top level response caching for ensemble models (#338)

Author: lkomali

src/dynamic_batch_scheduler.cc

@@ -39,14 +39,6 @@
 
 namespace triton { namespace core {
 
-uint64_t
-CaptureTimeNs()
-{
-  return std::chrono::duration_cast<std::chrono::nanoseconds>(
-             std::chrono::steady_clock::now().time_since_epoch())
-      .count();
-}
-
 bool
 IsStaleState(Payload::State payload_state)
 {
@@ -753,32 +745,9 @@ DynamicBatchScheduler::CacheLookUp(
     std::unique_ptr<InferenceRequest>& request,
     std::unique_ptr<InferenceResponse>& cached_response)
 {
-  Status status;
   auto cache = model_->Server()->CacheManager()->Cache();
-  std::unique_ptr<InferenceResponse> local_response;
-  request->ResponseFactory()->CreateResponse(&local_response);
-  // Hash request into cache key
-  std::string key = "";
-  if (!request->CacheKeyIsSet()) {
-    status = cache->Hash(*request, &key);
-    if (!status.IsOk()) {
-      LOG_ERROR << "Failed to hash request: " << status.Message();
-      return;
-    }
-    request->SetCacheKey(key);
-  } else {
-    key = request->CacheKey();
-  }
-
-  // Lookup and capture timestamps
-  {
-    request->CaptureCacheLookupStartNs();
-    status = cache->Lookup(local_response.get(), key);
-    request->CaptureCacheLookupEndNs();
-  }
-
-  if (status.IsOk() && (local_response != nullptr)) {
-    cached_response = std::move(local_response);
+  bool is_lookup_success = CacheLookUpUtil(request, cached_response, cache);
+  if (is_lookup_success) {
 #ifdef TRITON_ENABLE_STATS
     // Update model metrics/stats on cache hits
     // Backends will update metrics as normal on cache misses

src/ensemble_scheduler/ensemble_scheduler.cc

@@ -61,6 +61,10 @@ class RequestTracker {
 
   std::unique_ptr<InferenceRequest>& Request() { return request_; }
 
+  InferenceStatsAggregator* StatsAggregator() { return stats_aggregator_; }
+
+  MetricModelReporter* MetricReporter() { return metric_reporter_; }
+
   InferenceStatsAggregator& ContextStatsAggregator()
   {
     return context_stats_aggregator_;
@@ -316,6 +320,9 @@ class EnsembleContext {
       const std::set<std::pair<std::string, IterationCount>>& updated_tensors,
       std::unique_ptr<InferenceResponse>* response);
 
+  void CacheEnsembleTopLevelRequest(
+      std::unique_ptr<InferenceResponse>& response);
+
   InferenceServer* is_;
 
   EnsembleInfo* info_;
@@ -1033,6 +1040,50 @@ EnsembleContext::ReshapeTensorDims(
   return res;
 }
 
+// Caching function
+void
+EnsembleContext::CacheEnsembleTopLevelRequest(
+    std::unique_ptr<InferenceResponse>& response)
+{
+  const std::string key = request_tracker_->Request()->CacheKey();
+  const bool is_key_set = request_tracker_->Request()->CacheKeyIsSet();
+
+#ifdef TRITON_ENABLE_STATS
+  const uint64_t lookup_end_ns =
+      request_tracker_->Request()->CacheLookupEndNs();
+  const uint64_t lookup_start_ns =
+      request_tracker_->Request()->CacheLookupStartNs();
+#endif
+
+  if (!is_key_set) {
+    LOG_ERROR << "Request cache key was not set correctly.";
+  }
+
+  auto cache = is_->CacheManager()->Cache();
+#ifdef TRITON_ENABLE_STATS
+  const uint64_t insert_start_ns = CaptureTimeNs();
+#endif
+  auto status = cache->Insert(response.get(), key);
+  if (!status.IsOk()) {
+    LOG_ERROR << "Failed to insert key [" << key
+              << "] into response cache: " << status.Message();
+  }
+
+#ifdef TRITON_ENABLE_STATS
+  const uint64_t insert_end_ns = CaptureTimeNs();
+  uint64_t lookup_ns = lookup_end_ns - lookup_start_ns;
+  if (lookup_start_ns > lookup_end_ns) {
+    lookup_ns = 0;
+    LOG_ERROR << "Request lookup duration was not set correctly.";
+  }
+  uint64_t insert_ns = insert_end_ns - insert_start_ns;
+  uint64_t cache_miss_ns = lookup_ns + insert_ns;
+  request_tracker_->StatsAggregator()->UpdateSuccessCacheMiss(
+      request_tracker_->MetricReporter(), cache_miss_ns);
+#endif
+}
+
+
 Status
 EnsembleContext::FinishEnsemble(std::unique_ptr<InferenceResponse>&& response)
 {
@@ -1053,6 +1104,10 @@ EnsembleContext::FinishEnsemble(std::unique_ptr<InferenceResponse>&& response)
                      ? TRITONSERVER_RESPONSE_COMPLETE_FINAL
                      : 0;
     if (response != nullptr) {
+      // Cache the request if caching is enabled.
+      if (info_->is_cache_enabled_) {
+        CacheEnsembleTopLevelRequest(response);
+      }
       InferenceResponse::Send(std::move(response), flags);
       response_sent_ = true;
     } else if (flags != 0) {
@@ -1319,6 +1374,21 @@ EnsembleScheduler::Create(
   return Status::Success;
 }
 
+
+void
+EnsembleScheduler::CacheLookUp(
+    std::unique_ptr<InferenceRequest>& request,
+    std::unique_ptr<InferenceResponse>& cached_response)
+{
+  auto cache = is_->CacheManager()->Cache();
+  bool is_lookup_success = CacheLookUpUtil(request, cached_response, cache);
+  if (is_lookup_success) {
+#ifdef TRITON_ENABLE_STATS
+    request->ReportStatisticsCacheHit(metric_reporter_.get());
+#endif
+  }
+}
+
 Status
 EnsembleScheduler::Enqueue(std::unique_ptr<InferenceRequest>& request)
 {
@@ -1333,6 +1403,19 @@ EnsembleScheduler::Enqueue(std::unique_ptr<InferenceRequest>& request)
       TRITONSERVER_TRACE_TENSOR_QUEUE_INPUT, "EnsembleScheduler Enqueue");
 #endif  // TRITON_ENABLE_TRACING
 
+  std::unique_ptr<InferenceResponse> cached_response;
+  if (info_->is_cache_enabled_) {
+    CacheLookUp(request, cached_response);
+  }
+
+  if (cached_response != nullptr) {
+    InferenceResponse::Send(
+        std::move(cached_response), TRITONSERVER_RESPONSE_COMPLETE_FINAL);
+    InferenceRequest::Release(
+        std::move(request), TRITONSERVER_REQUEST_RELEASE_ALL);
+    return Status::Success;
+  }
+
   // Add additional callback to keep track of in-flight count
   ++inflight_count_;
   request->AddInternalReleaseCallback(
@@ -1387,6 +1470,10 @@ EnsembleScheduler::EnsembleScheduler(
   // This config field is filled internally for ensemble models
   info_->is_decoupled_ = config.model_transaction_policy().decoupled();
 
+  // field to check if response cache enabled in the ensemble model config.
+  info_->is_cache_enabled_ =
+      config.response_cache().enable() && is_->ResponseCacheEnabled();
+
   for (const auto& input : config.input()) {
     info_->tensor_to_step_.emplace(input.name(), std::set<size_t>());
     if (input.optional()) {

src/ensemble_scheduler/ensemble_scheduler.h

@@ -34,6 +34,7 @@
 #include "model_config.pb.h"
 #include "model_config_utils.h"
 #include "scheduler.h"
+#include "scheduler_utils.h"
 #include "status.h"
 
 #ifdef TRITON_ENABLE_GPU
@@ -65,6 +66,8 @@ struct EnsembleInfo {
 
   bool is_decoupled_;
 
+  bool is_cache_enabled_;
+
   // the ensemble output (re)shape expected by the ensemble
   std::unordered_map<std::string, triton::common::DimsList>
       ensemble_output_shape_;
@@ -97,6 +100,7 @@ class EnsembleScheduler : public Scheduler {
   // \see Scheduler::Enqueue()
   Status Enqueue(std::unique_ptr<InferenceRequest>& request) override;
 
+
   // \see Scheduler::InflightInferenceCount()
   size_t InflightInferenceCount() override { return inflight_count_; }
 
@@ -109,6 +113,10 @@ class EnsembleScheduler : public Scheduler {
       InferenceServer* const server, const ModelIdentifier& model_id,
       const inference::ModelConfig& config);
 
+  void CacheLookUp(
+      std::unique_ptr<InferenceRequest>& request,
+      std::unique_ptr<InferenceResponse>& cached_response);
+
   std::shared_ptr<MetricModelReporter> metric_reporter_;
   InferenceStatsAggregator* const stats_aggregator_;
   InferenceServer* const is_;

src/scheduler_utils.cc

@@ -33,6 +33,44 @@
 
 namespace triton { namespace core {
 
+uint64_t
+CaptureTimeNs()
+{
+  return std::chrono::duration_cast<std::chrono::nanoseconds>(
+             std::chrono::steady_clock::now().time_since_epoch())
+      .count();
+}
+
+bool
+CacheLookUpUtil(
+    std::unique_ptr<InferenceRequest>& request,
+    std::unique_ptr<InferenceResponse>& cached_response,
+    std::shared_ptr<TritonCache> cache)
+{
+  Status status;
+  std::unique_ptr<InferenceResponse> local_response;
+  request->ResponseFactory()->CreateResponse(&local_response);
+  std::string key = "";
+  if (!request->CacheKeyIsSet()) {
+    status = cache->Hash(*request, &key);
+    if (!status.IsOk()) {
+      LOG_ERROR << "Failed to hash request: " << status.Message();
+      return false;
+    }
+    request->SetCacheKey(key);
+  } else {
+    key = request->CacheKey();
+  }
+  request->CaptureCacheLookupStartNs();
+  status = cache->Lookup(local_response.get(), key);
+  request->CaptureCacheLookupEndNs();
+  if (status.IsOk() && (local_response != nullptr)) {
+    cached_response = std::move(local_response);
+    return true;
+  }
+  return false;
+}
+
 Status
 RequiredEqualInputs::Initialize(
     const std::unique_ptr<InferenceRequest>& request,

src/scheduler_utils.h

@@ -28,10 +28,19 @@
 #include <deque>
 #include <unordered_map>
 
+#include "cache_manager.h"
 #include "scheduler.h"
 
 namespace triton { namespace core {
 
+uint64_t CaptureTimeNs();
+// Utility function called by the scheduler to lookup if the request is in the
+// cache and get the response.
+bool CacheLookUpUtil(
+    std::unique_ptr<InferenceRequest>& request,
+    std::unique_ptr<InferenceResponse>& cached_response,
+    std::shared_ptr<TritonCache> cache);
+
 struct RequiredEqualInputs {
  public:
   RequiredEqualInputs() : init_(false), has_optional_input_(false) {}