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

deqiangc · tensorflow-copybara · commit 50b07e4ceac4 · 2024-04-18T12:39:54.000-07:00
PiperOrigin-RevId: 626118642
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.

Original file line number	Diff line number	Diff line change
`@@ -194,6 +194,9 @@ message TfrtSavedModelConfig {`
`194`	`194`
`195`	`195`	`// Whether to use fused op for GPU compile, execute and data transfer.`
`196`	`196`	`bool tfrt_use_fused_gpu_op = 2018;`
	`197`	`+`
	`198`	`+ // The minimum number of batch threads.`
	`199`	`+ int64 tfrt_min_num_batch_threads = 2019;`
`197`	`200`	`}`
`198`	`201`
`199`	`202`	`// Config proto for TfrtSavedModelSourceAdapter.`