Allow missing ensemble output(s) (#399)

triton-inference-server · Oct 16, 2024 · 571c5dc · 571c5dc
1 parent 84868f1
commit 571c5dc
Showing 1 changed file with 6 additions and 5 deletions.
diff --git a/src/ensemble_scheduler/ensemble_scheduler.cc b/src/ensemble_scheduler/ensemble_scheduler.cc
@@ -1200,11 +1200,6 @@ EnsembleContext::CheckAndSetEnsembleOutput(
         ready = false;
         break;
       }
-      // check if the output is provided
-      else if (tensor[iteration_count].data_ == nullptr) {
-        ready = false;
-        break;
-      }
     }
   }
   if (!ready) {
@@ -1224,6 +1219,12 @@ EnsembleContext::CheckAndSetEnsembleOutput(
     auto& tensor_data = tensor_data_[output_pair.first];
     auto& tensor = tensor_data.tensor_[iteration_count];
 
+    if (tensor.data_ == nullptr) {
+      LOG_VERBOSE(1) << "Composing models did not output tensor "
+                     << output_pair.first;
+      continue;
+    }
+
     auto shape = ReshapeTensorDims(
         output_pair.second, (lrequest->BatchSize() != 0),
         tensor_data.batch_size_, tensor.data_->OriginalShape());