@@ -293,7 +293,7 @@ private void inferAgainstAllocatedModel(
293293
294294 assert nodes .stream ().mapToInt (Tuple ::v2 ).sum () == request .numberOfDocuments ()
295295 : "mismatch; sum of node requests does not match number of documents in request" ;
296- inferOnAssignmentNodes (nodes , request , responseBuilder , parentTaskId , listener );
296+ inferOnAssignmentNodes (assignment . getDeploymentId (), nodes , request , responseBuilder , parentTaskId , listener );
297297 }
298298
299299 private void inferOnBlockedRequest (InferenceWaitForAllocation .WaitingRequest request , TrainedModelAssignment assignment ) {
@@ -313,11 +313,19 @@ private void inferOnBlockedRequest(InferenceWaitForAllocation.WaitingRequest req
313313 );
314314 }
315315
316- inferOnAssignmentNodes (nodes , request .request (), request .responseBuilder (), request .parentTaskId (), request .listener ());
316+ inferOnAssignmentNodes (
317+ assignment .getDeploymentId (),
318+ nodes ,
319+ request .request (),
320+ request .responseBuilder (),
321+ request .parentTaskId (),
322+ request .listener ()
323+ );
317324 });
318325 }
319326
320327 private void inferOnAssignmentNodes (
328+ String deploymentId ,
321329 List <Tuple <String , Integer >> nodes ,
322330 Request request ,
323331 Response .Builder responseBuilder ,
@@ -334,14 +342,14 @@ private void inferOnAssignmentNodes(
334342 InferTrainedModelDeploymentAction .Request deploymentRequest ;
335343 if (request .getTextInput () == null ) {
336344 deploymentRequest = InferTrainedModelDeploymentAction .Request .forDocs (
337- request . getId () ,
345+ deploymentId ,
338346 request .getUpdate (),
339347 request .getObjectsToInfer ().subList (startPos , startPos + node .v2 ()),
340348 request .getInferenceTimeout ()
341349 );
342350 } else {
343351 deploymentRequest = InferTrainedModelDeploymentAction .Request .forTextInput (
344- request . getId () ,
352+ deploymentId ,
345353 request .getUpdate (),
346354 request .getTextInput ().subList (startPos , startPos + node .v2 ()),
347355 request .getInferenceTimeout ()
0 commit comments