Add the batch_padding_policy attribute the tensorflow serving api.

tf-marissaw · tensorflow-copybara · commit ea02141a0008 · 2024-08-01T14:58:29.000-07:00
Original author is piatov@

PiperOrigin-RevId: 658548835
diff --git a/tensorflow_serving/servables/tensorflow/tfrt_saved_model_factory.cc b/tensorflow_serving/servables/tensorflow/tfrt_saved_model_factory.cc
@@ -238,6 +238,7 @@ absl::Status TfrtSavedModelFactory::CreateTfrtSavedModelWithMetadata(
   compile_options.min_num_batch_threads = config_.tfrt_min_num_batch_threads();
   compile_options.min_max_enqueued_batches =
       config_.tfrt_min_max_enqueued_batches();
+  compile_options.batch_padding_policy = config_.batch_padding_policy();
 
   options.graph_execution_options.run_placer_grappler_on_functions =
       config_.run_placer_grappler_on_functions();
diff --git a/tensorflow_serving/servables/tensorflow/tfrt_saved_model_source_adapter.proto b/tensorflow_serving/servables/tensorflow/tfrt_saved_model_source_adapter.proto
@@ -224,6 +224,11 @@ message TfrtSavedModelConfig {
 
   // The minimum of the maximum number of outstanding enqueue batches
   int64 tfrt_min_max_enqueued_batches = 2025;
+
+  // The policy used by a BatchScheduler to pad (or split) batches.
+  // See the description of the batch_padding_policy attribute of
+  // BatchFunction in core/ops/batch_ops.cc.
+  string batch_padding_policy = 2026;
 }
 
 // Config proto for TfrtSavedModelSourceAdapter.
