Passing parameters to batched requests from data-flow engine (#374)

SeldonIO · Jul 29, 2022 · 5550c43 · 5550c43
1 parent 4bdcc6d
commit 5550c43
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/scheduler/data-flow/src/main/kotlin/io/seldon/dataflow/kafka/BatchProcessor.kt b/scheduler/data-flow/src/main/kotlin/io/seldon/dataflow/kafka/BatchProcessor.kt
@@ -1,6 +1,7 @@
 package io.seldon.dataflow.kafka
 
 import io.klogging.noCoLogger
+import io.seldon.mlops.inference.v2.V2Dataplane.InferParameter
 import io.seldon.mlops.inference.v2.V2Dataplane.InferTensorContents
 import io.seldon.mlops.inference.v2.V2Dataplane.ModelInferRequest
 import io.seldon.mlops.inference.v2.V2Dataplane.ModelInferRequest.InferInputTensor
@@ -57,11 +58,13 @@ class BatchProcessor(private val threshold: Int) : Transformer<String, ModelInfe
                 val batchSize = v.sumOf { it.getShape(0) }
                 val dataShape = v.first().shapeList.drop(1)
                 val shape = mutableListOf(batchSize) + dataShape
-
+                val parameters = mutableMapOf<String, InferParameter>()
+                v.map { it.parametersMap }.forEach {parameters.putAll(it)}
                 InferInputTensor
                     .newBuilder()
                     .setName(k)
                     .addAllShape(shape)
+                    .putAllParameters(parameters)
                     .setContents(
                         InferTensorContents
                             .newBuilder()
@@ -84,6 +87,7 @@ class BatchProcessor(private val threshold: Int) : Transformer<String, ModelInfe
             .setId(batchReferenceRequest.id)
             .setModelName(batchReferenceRequest.modelName)
             .setModelVersion(batchReferenceRequest.modelVersion)
+            .putAllParameters(batchReferenceRequest.parametersMap)
             .addAllInputs(tensorNames.values)
             .build()
     }