Add an option to override to the size of GPU system

PiperOrigin-RevId: 564804399
tensorflow · Sep 21, 2023 · 445a87b · 445a87b
1 parent 1f5d8d7
commit 445a87b
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
@@ -228,6 +228,8 @@ absl::Status TfrtSavedModelFactory::CreateTfrtSavedModelWithMetadata(
   options.graph_execution_options.run_placer_grappler_on_functions =
       config_.run_placer_grappler_on_functions();
   options.graph_execution_options.enable_tfrt_gpu = config_.enable_tfrt_gpu();
+  options.graph_execution_options.gpu_system_memory_size_in_mb =
+      config_.gpu_system_memory_size_in_mb();
   options.graph_execution_options.enable_grappler_function_optimizer =
       config_.enable_grappler_function_optimizer();
   options.graph_execution_options.enable_online_cost_analysis =

diff --git a/tensorflow_serving/servables/tensorflow/tfrt_saved_model_source_adapter.proto b/tensorflow_serving/servables/tensorflow/tfrt_saved_model_source_adapter.proto
@@ -186,6 +186,8 @@ message TfrtSavedModelConfig {
   }
   PredictResponseTensorSerializationOption
       predict_response_tensor_serialization_option = 2014;
+  // The size of reserved memory space for GPU system.
+  int32 gpu_system_memory_size_in_mb = 2015;
 }
 
 // Config proto for TfrtSavedModelSourceAdapter.