Update

Author: pskiran1

src/ensemble_scheduler/ensemble_scheduler.cc

@@ -375,9 +375,9 @@ class EnsembleContext {
   size_t inflight_step_counter_;
 
   // Backpressure support: Limits memory growth from decoupled models.
-  // Tracks inflight responses per step; blocks producers when downstream
-  // consumers are overloaded. Only active if max_inflight_responses_ > 0.
-  std::vector<size_t> step_inflight_response_counts_;
+  // Tracks inflight requests per step; blocks producers when downstream
+  // consumers are overloaded. Only active if max_inflight_requests_ > 0.
+  std::vector<size_t> step_inflight_request_counts_;
   std::vector<std::unique_ptr<std::mutex>> step_mutexes_;
   std::vector<std::unique_ptr<std::condition_variable>> step_cv_vec_;
 
@@ -518,9 +518,9 @@ EnsembleContext::EnsembleContext(
 
   // Initialize backpressure tracking if enabled.
   size_t num_steps = info_->steps_.size();
-  step_inflight_response_counts_.resize(num_steps, 0);
+  step_inflight_request_counts_.resize(num_steps, 0);
 
-  if (info_->max_inflight_responses_ > 0) {
+  if (info_->max_inflight_requests_ > 0) {
     step_mutexes_.resize(num_steps);
     step_cv_vec_.resize(num_steps);
 
@@ -935,11 +935,11 @@ EnsembleContext::UpdateEnsembleState(
 
       size_t completed_step_idx = completed_step->step_idx_;
 
-      // Decrement step_inflight_response_counts_, then notify any producer
+      // Decrement step_inflight_request_counts_, then notify any producer
       // threads blocked waiting for this step's capacity
-      if (info_->max_inflight_responses_ > 0 && !step_cv_vec_.empty()) {
+      if (info_->max_inflight_requests_ > 0 && !step_cv_vec_.empty()) {
         std::lock_guard<std::mutex> lk(*step_mutexes_[completed_step_idx]);
-        step_inflight_response_counts_[completed_step_idx]--;
+        step_inflight_request_counts_[completed_step_idx]--;
         step_cv_vec_[completed_step_idx]->notify_one();
       }
     }
@@ -985,13 +985,6 @@ EnsembleContext::GetNextSteps(
   for (const auto& idx : next_step_idx) {
     steps->emplace_back();
     RETURN_IF_ERROR(InitStep(idx.first, idx.second, &(steps->back())));
-
-    // Track as inflight. Checked by producers for backpressure; decremented on
-    // completion.
-    if (info_->max_inflight_responses_ > 0 && !step_mutexes_.empty()) {
-      std::lock_guard<std::mutex> lk(*step_mutexes_[idx.first]);
-      step_inflight_response_counts_[idx.first]++;
-    }
   }
   inflight_step_counter_ += steps->size();
 
@@ -1436,8 +1429,8 @@ EnsembleContext::ScheduleSteps(
     step->ctx_ = context;
     size_t this_step_idx = step->step_idx_;
 
-    // Block if this step is overloaded.
-    if (context->info_->max_inflight_responses_ > 0 &&
+    // Apply backpressure to downstream steps only, not the entry step
+    if ((this_step_idx != 0) && context->info_->max_inflight_requests_ > 0 &&
         !context->step_cv_vec_.empty()) {
       std::unique_lock<std::mutex> lk(*context->step_mutexes_[this_step_idx]);
 
@@ -1450,8 +1443,8 @@ EnsembleContext::ScheduleSteps(
       bool capacity_available =
           context->step_cv_vec_[this_step_idx]->wait_for(lk, timeout, [&] {
             return cancelled() ||
-                   (context->step_inflight_response_counts_[this_step_idx] <
-                    context->info_->max_inflight_responses_);
+                   (context->step_inflight_request_counts_[this_step_idx] <
+                    context->info_->max_inflight_requests_);
           });
 
       // Log error only if timeout occurred (not cancellation).
@@ -1460,8 +1453,8 @@ EnsembleContext::ScheduleSteps(
                   << context->info_->ensemble_name_
                   << "' unable to schedule step " << this_step_idx
                   << " (inflight: "
-                  << context->step_inflight_response_counts_[this_step_idx]
-                  << " >= limit: " << context->info_->max_inflight_responses_
+                  << context->step_inflight_request_counts_[this_step_idx]
+                  << " >= limit: " << context->info_->max_inflight_requests_
                   << ") for " << kMutexTimeoutSeconds
                   << " seconds. Proceeding to avoid deadlock.";
       }
@@ -1496,6 +1489,15 @@ EnsembleContext::ScheduleSteps(
       std::unique_ptr<InferenceRequest> request = std::move(step->request_);
       auto step_status = context->is_->InferAsync(request);
       if (step_status.IsOk()) {
+        // Increment inflight counter AFTER successful scheduling. Always
+        // increment for ALL steps (including step 0) to ensure symmetry with
+        // decrement and prevent underflow when steps complete.
+        if (context->info_->max_inflight_requests_ > 0 &&
+            !context->step_mutexes_.empty()) {
+          std::lock_guard<std::mutex> lk(
+              *context->step_mutexes_[this_step_idx]);
+          context->step_inflight_request_counts_[this_step_idx]++;
+        }
         step.release();
         continue;
       } else {
@@ -1678,15 +1680,15 @@ EnsembleScheduler::EnsembleScheduler(
   }
   callback_pool_ = is_->EnsembleCallbackPool();
 
-  // Backpressure configuration from protobuf field. Limits concurrent responses
+  // Backpressure configuration from protobuf field. Limits concurrent requests
   // from decoupled steps to prevent memory growth. Value of 0 means unlimited.
   if (config.has_ensemble_scheduling()) {
-    info_->max_inflight_responses_ =
-        config.ensemble_scheduling().max_inflight_responses();
-    if (info_->max_inflight_responses_ > 0) {
+    info_->max_inflight_requests_ =
+        config.ensemble_scheduling().max_inflight_requests();
+    if (info_->max_inflight_requests_ > 0) {
       LOG_INFO << "Ensemble model '" << config.name()
-               << "' configured with max_inflight_responses: "
-               << info_->max_inflight_responses_;
+               << "' configured with max_inflight_requests: "
+               << info_->max_inflight_requests_;
     }
   }
 }

src/ensemble_scheduler/ensemble_scheduler.h

@@ -84,11 +84,11 @@ struct EnsembleInfo {
   // backward path, ensemble tensor to the step that provides its data
   std::unordered_map<std::string, size_t> tensor_to_prev_step_;
 
-  // Maximum concurrent inflight responses from steps to downstream consumers.
+  // Maximum concurrent inflight requests from steps to downstream consumers.
   // Prevents memory growth by blocking producers when limit is reached.
   // Default value is 0, which indicates unlimited (no backpressure applied).
-  // Configured via 'max_inflight_responses' parameter in config.pbtxt.
-  size_t max_inflight_responses_ = 0;
+  // Configured via 'max_inflight_requests' field in ensemble_scheduling.
+  size_t max_inflight_requests_ = 0;
 };
 
 // Scheduler that implements ensemble scheduling.