fix: RequestTracker counter mismatch (triton-inference-server#483)

* Fix RequestTracker counter mismatch in ScheduleSteps with parallel failures

Author: mattwittwer

src/ensemble_scheduler/ensemble_scheduler.cc

@@ -87,8 +87,24 @@ class RequestTracker {
   {
   }
 
+  // Accessed without additional synchronization while protected by
+  // EnsembleContext::mutex_.
   std::unique_ptr<InferenceRequest>& Request() { return request_; }
 
+  // Used from paths where request_ may be released concurrently.
+  bool IsCancelled()
+  {
+    std::lock_guard<std::mutex> lk(mtx_);
+    return (request_ == nullptr) || request_->IsCancelled();
+  }
+
+  std::string LogRequest()
+  {
+    std::lock_guard<std::mutex> lk(mtx_);
+    return (request_ != nullptr) ? request_->LogRequest()
+                                 : std::string("[request released] ");
+  }
+
   InferenceStatsAggregator* StatsAggregator() { return stats_aggregator_; }
 
   MetricModelReporter* MetricReporter() { return metric_reporter_; }
@@ -141,6 +157,23 @@ class RequestTracker {
     status_ = status;
   }
 
+  void RespondIfError(const Status& status, FailureReason reason)
+  {
+    std::lock_guard<std::mutex> lk(mtx_);
+    if (request_ != nullptr) {
+      InferenceRequest::RespondIfError(
+          request_, status, false /* release_request */, reason);
+    }
+  }
+
+  void SendFlags(const uint32_t flags)
+  {
+    std::lock_guard<std::mutex> lk(mtx_);
+    if (request_ != nullptr) {
+      request_->ResponseFactory()->SendFlags(flags);
+    }
+  }
+
  private:
   std::mutex mtx_;
   uint32_t inflight_request_counter_;
@@ -153,6 +186,7 @@ class RequestTracker {
   triton::common::ThreadPool* const callback_pool_;
 };
 
+using RequestTrackerReference = std::shared_ptr<RequestTracker>;
 // Step is used as 'userp' and keeps ensemble context alive
 // until no more internal requests are inflight.
 // Step contains metadata, and status for the
@@ -188,6 +222,11 @@ struct Step {
   const bool preserve_responses_order_;
 
   size_t step_idx_;
+
+  // Heap allocation passed as the release-callback userp. The allocation
+  // stores a shared_ptr<RequestTracker> so the tracker stays alive until the
+  // release callback or local failure cleanup drops this reference.
+  RequestTrackerReference* callback_tracker_ref_{nullptr};
 };
 
 struct TensorData {
@@ -396,7 +435,7 @@ class EnsembleContext {
 
   // Objects related to the ensemble infer request
   Status ensemble_status_;
-  RequestTracker* request_tracker_;
+  std::shared_ptr<RequestTracker> request_tracker_;
   // Use in conjunction with 'is_decoupled_' in EnsembleInfo to
   // better distinguish ensemble ending behavior (see annotation in
   // FinishEnsemble for details).
@@ -429,7 +468,7 @@ EnsembleContext::EnsembleContext(
 {
   uint64_t compute_start_ns = 0;
   INFER_STATS_SET_TIMESTAMP(compute_start_ns);
-  request_tracker_ = new RequestTracker(
+  request_tracker_ = std::make_shared<RequestTracker>(
       std::move(request), compute_start_ns, metric_reporter, stats_aggregator,
       callback_pool);
 
@@ -646,16 +685,17 @@ void
 EnsembleContext::RequestComplete(
     TRITONSERVER_InferenceRequest* request, const uint32_t flags, void* userp)
 {
-  auto request_tracker = reinterpret_cast<RequestTracker*>(userp);
+  auto callback_tracker_ref = reinterpret_cast<RequestTrackerReference*>(userp);
+  auto request_tracker = *callback_tracker_ref;
   auto pool = request_tracker->CallbackPool();
-  auto fn = [request, flags, request_tracker]() {
+  auto fn = [request, flags, request_tracker, callback_tracker_ref]() {
     if ((flags & TRITONSERVER_REQUEST_RELEASE_ALL) != 0) {
+      std::unique_ptr<RequestTrackerReference> managed_callback_tracker_ref(
+          callback_tracker_ref);
       LOG_TRITONSERVER_ERROR(
           TRITONSERVER_InferenceRequestDelete(request),
           "deleting ensemble inference request");
-      if (request_tracker->DecrementCounter()) {
-        delete request_tracker;
-      }
+      request_tracker->DecrementCounter();
     }
   };
 
@@ -1070,12 +1110,19 @@ EnsembleContext::InitStep(
   irequest->SetSecondaryStatsAggregator(
       &request_tracker_->ContextStatsAggregator());
 #endif
+  // Heap-allocate the release-callback userp because the C callback API only
+  // stores a raw void*. The heap object itself is single-owner here, while the
+  // object stored inside it is a shared_ptr<RequestTracker> that keeps the
+  // tracker alive until the callback or local cleanup drops this reference.
+  auto callback_tracker_ref =
+      std::make_unique<RequestTrackerReference>(request_tracker_);
   irequest->SetResponseCallback(
       reinterpret_cast<ResponseAllocator*>(allocator_.get()), step->get(),
       ResponseComplete, step->get());
-  irequest->SetReleaseCallback(RequestComplete, request_tracker_);
+  irequest->SetReleaseCallback(RequestComplete, callback_tracker_ref.get());
 
   RETURN_IF_ERROR(irequest->PrepareForInference());
+  (*step)->callback_tracker_ref_ = callback_tracker_ref.release();
 
 #ifdef TRITON_ENABLE_TRACING
   auto& parent_trace = request_tracker_->Request()->TraceProxy();
@@ -1220,16 +1267,14 @@ EnsembleContext::FinishEnsemble(std::unique_ptr<InferenceResponse>&& response)
         ensemble_status_ = Status(
             Status::Code::INVALID_ARG,
             "in ensemble '" + info_->ensemble_name_ + "', " +
-                request_tracker_->Request()->LogRequest() +
+                request_tracker_->LogRequest() +
                 "unexpected deadlock, at least one output is not set while no "
                 "more "
                 "ensemble steps can be made");
-        InferenceRequest::RespondIfError(
-            request_tracker_->Request(), ensemble_status_,
-            false /* release_requests */, FailureReason::OTHER);
+        request_tracker_->RespondIfError(
+            ensemble_status_, FailureReason::OTHER);
       } else {
-        request_tracker_->Request()->ResponseFactory()->SendFlags(
-            TRITONSERVER_RESPONSE_COMPLETE_FINAL);
+        request_tracker_->SendFlags(TRITONSERVER_RESPONSE_COMPLETE_FINAL);
       }
     }
   } else {
@@ -1239,9 +1284,8 @@ EnsembleContext::FinishEnsemble(std::unique_ptr<InferenceResponse>&& response)
             std::move(response), TRITONSERVER_RESPONSE_COMPLETE_FINAL,
             ensemble_status_);
       } else {
-        InferenceRequest::RespondIfError(
-            request_tracker_->Request(), ensemble_status_,
-            false /* release_requests */, FailureReason::OTHER);
+        request_tracker_->RespondIfError(
+            ensemble_status_, FailureReason::OTHER);
       }
       error_response_sent_ = true;
     }
@@ -1251,10 +1295,8 @@ EnsembleContext::FinishEnsemble(std::unique_ptr<InferenceResponse>&& response)
     // Reach here when the ensemble execution comes to the end,
     // 'ensemble_status_' at this point is representative.
     request_tracker_->SetStatus(ensemble_status_);
-    if (request_tracker_->DecrementCounter()) {
-      delete request_tracker_;
-    }
-    request_tracker_ = nullptr;
+    request_tracker_->DecrementCounter();
+    request_tracker_.reset();
   }
   return ensemble_status_;
 }
@@ -1436,7 +1478,7 @@ EnsembleContext::ScheduleSteps(
     if (should_schedule) {
       // If the ensemble request is cancelled, propagate the cancellation to the
       // next request step.
-      if (context->request_tracker_->Request()->IsCancelled()) {
+      if (context->request_tracker_->IsCancelled()) {
         step->request_->Cancel();
       }
       // Acquire a slot from the per-step shared limiter only for steps that
@@ -1469,16 +1511,20 @@ EnsembleContext::ScheduleSteps(
     // Reaching here means the step is not being scheduled, update corresponding
     // counters and attempt to finish ensemble if it is the last step.
 
-
     // Release the limiter slot if one was acquired, and update counters.
     if (should_schedule &&
         !context->info_->step_inflight_request_limiters_.empty()) {
       context->info_->step_inflight_request_limiters_[this_step_idx]->Release();
     }
 
     std::lock_guard<std::mutex> lock(context->mutex_);
-    // Decrement only when IncrementCounter was called. An unconditional
-    // decrement would underflow the counter and cause a use-after-free.
+    // The request never reaches the callback-owned release path, so drop the
+    // heap-allocated callback userp here.
+    delete step->callback_tracker_ref_;
+    step->callback_tracker_ref_ = nullptr;
+    // Only undo IncrementCounter() for steps that actually reached the
+    // scheduling path. Otherwise the counter can underflow and release the
+    // top-level request while FinishEnsemble is still using it.
     if (should_schedule) {
       context->request_tracker_->DecrementCounter();
     }

src/ensemble_scheduler/ensemble_scheduler.h

@@ -27,7 +27,9 @@
 
 #ifdef TRITON_ENABLE_ENSEMBLE
 
+#include <condition_variable>
 #include <memory>
+#include <mutex>
 
 #include "metric_model_reporter.h"
 #include "model.h"