Provide a runtime option to lower bound the number of batch threads.

PiperOrigin-RevId: 626118642
tensorflow · Apr 18, 2024 · 50b07e4 · 50b07e4
1 parent 2235e43
commit 50b07e4
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 0 deletions.
diff --git a/tensorflow_serving/servables/tensorflow/tfrt_saved_model_factory.cc b/tensorflow_serving/servables/tensorflow/tfrt_saved_model_factory.cc
@@ -230,6 +230,7 @@ absl::Status TfrtSavedModelFactory::CreateTfrtSavedModelWithMetadata(
       ToTpuAllowUnpaddedBatch(config_.tpu_unpadded_batch_mode());
   compile_options.use_gpu_compile_and_execute_op =
       config_.tfrt_use_fused_gpu_op();
+  compile_options.min_num_batch_threads = config_.tfrt_min_num_batch_threads();
 
   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
@@ -194,6 +194,9 @@ message TfrtSavedModelConfig {
 
   // Whether to use fused op for GPU compile, execute and data transfer.
   bool tfrt_use_fused_gpu_op = 2018;
+
+  // The minimum number of batch threads.
+  int64 tfrt_min_num_batch_threads = 2019;
 }
 
 // Config proto for TfrtSavedModelSourceAdapter.