Porting the grappler batch prioritization rewriter to TFRT.

PiperOrigin-RevId: 776869098

Author: tensorflower-gardener

tensorflow_serving/servables/tensorflow/tfrt_saved_model_factory.cc

@@ -222,6 +222,8 @@ absl::StatusOr<tfrt::SavedModel::Options> CreateCommonSavedModelOptions(
   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_queue_global_prioritization_num_threads =
+      config.tfrt_batch_queue_global_prioritization_num_threads();
   compile_options.batch_padding_policy = config.batch_padding_policy();
   compile_options.batch_options = config.in_graph_batching_parameters();
 

tensorflow_serving/servables/tensorflow/tfrt_saved_model_source_adapter.proto

@@ -234,6 +234,10 @@ message TfrtSavedModelConfig {
 
   // Batching parameters to be rewritten in the existing BatchFunction ops.
   BatchingOptions in_graph_batching_parameters = 2028;
+
+  // If non-zero, all models on this server are switched to use a prioritized
+  // batching function using this number of global threads.
+  int64 tfrt_batch_queue_global_prioritization_num_threads = 2029;
 }
 
 // Config proto for TfrtSavedModelSourceAdapter.