Revert "feat: Ensemble asynchronous callback executions (#429)" (#436)

This reverts commit 56e97eb.

Author: dmitry-tokarev-nv

src/constants.h

@@ -62,7 +62,6 @@ constexpr char kPythonBackend[] = "python";
 
 #ifdef TRITON_ENABLE_ENSEMBLE
 constexpr char kEnsemblePlatform[] = "ensemble";
-constexpr uint64_t ENSEMBLE_CB_POOL_SIZE = 8u;
 #endif  // TRITON_ENABLE_ENSEMBLE
 
 constexpr char kTensorRTExecutionAccelerator[] = "tensorrt";

src/ensemble_scheduler/ensemble_scheduler.cc

@@ -1,4 +1,4 @@
-// Copyright 2019-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// Copyright 2019-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 //
 // Redistribution and use in source and binary forms, with or without
 // modification, are permitted provided that the following conditions
@@ -52,12 +52,10 @@ class RequestTracker {
   explicit RequestTracker(
       std::unique_ptr<InferenceRequest>&& request, uint64_t compute_start_ns,
       MetricModelReporter* metric_reporter,
-      InferenceStatsAggregator* stats_aggregator,
-      triton::common::ThreadPool* callback_pool)
+      InferenceStatsAggregator* stats_aggregator)
       : inflight_request_counter_(1), request_(std::move(request)),
         compute_start_ns_(compute_start_ns), metric_reporter_(metric_reporter),
-        stats_aggregator_(stats_aggregator), status_(Status::Success),
-        callback_pool_(callback_pool)
+        stats_aggregator_(stats_aggregator), status_(Status::Success)
   {
   }
 
@@ -72,8 +70,6 @@ class RequestTracker {
     return context_stats_aggregator_;
   }
 
-  triton::common::ThreadPool* CallbackPool() const { return callback_pool_; }
-
   void IncrementCounter()
   {
     std::lock_guard<std::mutex> lk(mtx_);
@@ -124,7 +120,6 @@ class RequestTracker {
   InferenceStatsAggregator* stats_aggregator_;
   InferenceStatsAggregator context_stats_aggregator_;
   Status status_;
-  triton::common::ThreadPool* const callback_pool_;
 };
 
 // Step is used as 'userp' and keeps ensemble context alive
@@ -242,7 +237,7 @@ class EnsembleContext {
       MetricModelReporter* metric_reporter,
       InferenceStatsAggregator* stats_aggregator, InferenceServer* is,
       EnsembleInfo* info, std::unique_ptr<InferenceRequest>& request,
-      cudaStream_t stream, triton::common::ThreadPool* callback_pool);
+      cudaStream_t stream);
 
   // Perform transition on 'context' state given the information of
   // 'completed_step'
@@ -331,8 +326,6 @@ class EnsembleContext {
   void CacheEnsembleTopLevelRequest(
       std::unique_ptr<InferenceResponse>& response);
 
-  triton::common::ThreadPool* CallbackPool() const { return callback_pool_; }
-
   InferenceServer* is_;
 
   EnsembleInfo* info_;
@@ -382,26 +375,20 @@ class EnsembleContext {
       TRITONSERVER_ResponseAllocator,
       decltype(&TRITONSERVER_ResponseAllocatorDelete)>
       allocator_;
-
-  // The thread pool used to execute ensemble callbacks and reduce e2e latency.
-  // The thread pool is managed by InferenceServer.
-  triton::common::ThreadPool* const callback_pool_;
 };
 
 EnsembleContext::EnsembleContext(
     MetricModelReporter* metric_reporter,
     InferenceStatsAggregator* stats_aggregator, InferenceServer* is,
     EnsembleInfo* info, std::unique_ptr<InferenceRequest>& request,
-    cudaStream_t stream, triton::common::ThreadPool* callback_pool)
+    cudaStream_t stream)
     : is_(is), info_(info), stream_(stream), inflight_step_counter_(0),
-      allocator_(nullptr, TRITONSERVER_ResponseAllocatorDelete),
-      callback_pool_(callback_pool)
+      allocator_(nullptr, TRITONSERVER_ResponseAllocatorDelete)
 {
   uint64_t compute_start_ns = 0;
   INFER_STATS_SET_TIMESTAMP(compute_start_ns);
   request_tracker_ = new RequestTracker(
-      std::move(request), compute_start_ns, metric_reporter, stats_aggregator,
-      callback_pool);
+      std::move(request), compute_start_ns, metric_reporter, stats_aggregator);
 
   auto& lrequest = request_tracker_->Request();
 
@@ -616,52 +603,29 @@ void
 EnsembleContext::RequestComplete(
     TRITONSERVER_InferenceRequest* request, const uint32_t flags, void* userp)
 {
-  auto request_tracker = reinterpret_cast<RequestTracker*>(userp);
-  auto pool = request_tracker->CallbackPool();
-  auto fn = [request, flags, request_tracker]() {
-    if ((flags & TRITONSERVER_REQUEST_RELEASE_ALL) != 0) {
-      LOG_TRITONSERVER_ERROR(
-          TRITONSERVER_InferenceRequestDelete(request),
-          "deleting ensemble inference request");
-      if (request_tracker->DecrementCounter()) {
-        delete request_tracker;
-      }
+  if ((flags & TRITONSERVER_REQUEST_RELEASE_ALL) != 0) {
+    LOG_TRITONSERVER_ERROR(
+        TRITONSERVER_InferenceRequestDelete(request),
+        "deleting ensemble inference request");
+    auto request_tracker = reinterpret_cast<RequestTracker*>(userp);
+    if (request_tracker->DecrementCounter()) {
+      delete request_tracker;
     }
-  };
-
-  // Attempt to enqueue the callback. If all workers are busy and queue is at
-  // capacity, execute the callback immediately.
-  if (pool->TaskQueueSize() < pool->Size()) {
-    pool->Enqueue(fn);
-  } else {
-    fn();
   }
 }
 
 void
 EnsembleContext::ResponseComplete(
     TRITONSERVER_InferenceResponse* response, const uint32_t flags, void* userp)
 {
-  auto step_raw_ptr = reinterpret_cast<Step*>(userp);
-  auto pool = step_raw_ptr->ctx_->CallbackPool();
-  auto fn = [response, flags, step_raw_ptr]() {
-    auto step_ptr = std::unique_ptr<Step>(step_raw_ptr);
-    step_ptr->response_flags_ = flags;
-    step_ptr->response_ = response;
-
-    EnsembleContext::Proceed(step_ptr->ctx_, step_ptr);
-    // Expecting more responses
-    if ((flags & TRITONSERVER_RESPONSE_COMPLETE_FINAL) == 0) {
-      step_ptr.release();
-    }
-  };
-
-  // Attempt to enqueue the callback. If all workers are busy and queue is at
-  // capacity, execute the callback immediately.
-  if (pool->TaskQueueSize() < pool->Size()) {
-    pool->Enqueue(fn);
-  } else {
-    fn();
+  auto step_ptr = std::unique_ptr<Step>(reinterpret_cast<Step*>(userp));
+  step_ptr->response_flags_ = flags;
+  step_ptr->response_ = response;
+
+  EnsembleContext::Proceed(step_ptr->ctx_, step_ptr);
+  // Expecting more responses
+  if ((flags & TRITONSERVER_RESPONSE_COMPLETE_FINAL) == 0) {
+    step_ptr.release();
   }
 }
 
@@ -1484,7 +1448,7 @@ EnsembleScheduler::Enqueue(std::unique_ptr<InferenceRequest>& request)
   RETURN_IF_ERROR(request->SetState(InferenceRequest::State::EXECUTING));
   std::shared_ptr<EnsembleContext> context(new EnsembleContext(
       metric_reporter_.get(), stats_aggregator_, is_, info_.get(), request,
-      stream_, callback_pool_));
+      stream_));
   EnsembleContext::Proceed(context);
   return Status::Success;
 }
@@ -1573,7 +1537,6 @@ EnsembleScheduler::EnsembleScheduler(
       info_->tensor_to_prev_step_.emplace(pair.second, step_idx);
     }
   }
-  callback_pool_ = is_->EnsembleCallbackPool();
 }
 
 EnsembleScheduler::~EnsembleScheduler()

src/ensemble_scheduler/ensemble_scheduler.h

@@ -1,4 +1,4 @@
-// Copyright 2019-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// Copyright 2019-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 //
 // Redistribution and use in source and binary forms, with or without
 // modification, are permitted provided that the following conditions
@@ -36,7 +36,6 @@
 #include "scheduler.h"
 #include "scheduler_utils.h"
 #include "status.h"
-#include "triton/common/thread_pool.h"
 
 #ifdef TRITON_ENABLE_GPU
 #include <cuda_runtime_api.h>
@@ -108,8 +107,6 @@ class EnsembleScheduler : public Scheduler {
   // \see Scheduler::Stop()
   void Stop() override {}
 
-  triton::common::ThreadPool* CallbackPool() const { return callback_pool_; }
-
  private:
   EnsembleScheduler(
       InferenceStatsAggregator* const stats_aggregator,
@@ -131,10 +128,6 @@ class EnsembleScheduler : public Scheduler {
   cudaStream_t stream_;
 
   std::atomic<size_t> inflight_count_;
-
-  // Fixed-size thread pool to run callbacks at end of each ensemble step.
-  // Managed by the server.
-  triton::common::ThreadPool* callback_pool_;
 };
 
 }}  // namespace triton::core

src/server.cc

@@ -1,4 +1,4 @@
-// Copyright 2018-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// Copyright 2018-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 //
 // Redistribution and use in source and binary forms, with or without
 // modification, are permitted provided that the following conditions
@@ -117,13 +117,6 @@ InferenceServer::InferenceServer()
 #endif  // TRITON_ENABLE_GPU
 
   inflight_request_counter_ = 0;
-
-#ifdef TRITON_ENABLE_ENSEMBLE
-  // TODO: Need to scale the thread pool size smarter, e.g. based on the
-  // instance_group count of composing models.
-  ensemble_cb_pool_.reset(
-      new triton::common::ThreadPool(ENSEMBLE_CB_POOL_SIZE));
-#endif
 }
 
 Status

src/server.h

@@ -1,4 +1,4 @@
-// Copyright 2018-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// Copyright 2018-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 //
 // Redistribution and use in source and binary forms, with or without
 // modification, are permitted provided that the following conditions
@@ -332,13 +332,6 @@ class InferenceServer {
     return cache_manager_;
   }
 
-#ifdef TRITON_ENABLE_ENSEMBLE
-  triton::common::ThreadPool* EnsembleCallbackPool() const
-  {
-    return ensemble_cb_pool_.get();
-  }
-#endif  // TRITON_ENABLE_ENSEMBLE
-
  private:
   const std::string version_;
   std::string id_;
@@ -382,12 +375,6 @@ class InferenceServer {
   std::unique_ptr<ModelRepositoryManager> model_repository_manager_;
   std::shared_ptr<TritonBackendManager> backend_manager_;
   std::shared_ptr<TritonCacheManager> cache_manager_;
-
-#ifdef TRITON_ENABLE_ENSEMBLE
-  // The thread pool for all ensemble models to execute callbacks
-  // asynchronously.
-  std::unique_ptr<triton::common::ThreadPool> ensemble_cb_pool_;
-#endif  // TRITON_ENABLE_ENSEMBLE
 };
 
 }}  // namespace triton::core