Running multiple TensorRT optimized models in Tensorflow

My project uses multiple Keras models. Those models can have an input with different batch size, that varies from 1 to 24. I decided to optimize those models using TF-TRT.

I tried 2 conversion approaches:

from tensorflow.python.compiler.tensorrt import trt_convert as trt

First approach converts the model but does not create a TensorRT engines for the model:

conversion_params = trt.DEFAULT_TRT_CONVERSION_PARAMS._replace(
                precision_mode=trt.TrtPrecisionMode.FP32)
converter = trt.TrtGraphConverterV2(
                input_saved_model_dir=saved_model_path,
                conversion_params=conversion_params)
converter.convert()
converter.save(output_saved_model_dir=trt_fp32_model_path)

Second approach converts the model and and builds TensorRT engine for all possible input shapes:

def input_function():
    def input_function():
    input_shapes = [(x, MODEL_INPUT_H, MODEL_INPUT_W, 3) for x in range(1, 25)]
    for shape in input_shapes:
         yield [np.random.normal(size=shape).astype(np.float32)]

conversion_params = trt.DEFAULT_TRT_CONVERSION_PARAMS._replace(
    precision_mode=trt.TrtPrecisionMode.FP32,
    maximum_cached_engines=100
)

converter = trt.TrtGraphConverterV2(
    input_saved_model_dir=saved_model_path,
    conversion_params=conversion_params)

converter.convert()
converter.build(input_fn=input_function)
converter.save(output_saved_model_dir=trt_fp32_model_path)

In script that uses my models, I use those models consecutively:

some loop:
    model1.predict(model1_input)
    model2.predict(model2_input)
    model3.predict(model3_input)

When the first conversion approach is used to optimize the models, I am able to load all models, but at runtime Tensorflow rebuilds TensorRT engines every time a model execution context changes. This causes a large performance overhead, which I was trying to overcome by caching TensorRT engines for those models (second conversion approach).

The problem is that when I am trying to load more than one TensorRT optimized model with pre-built engines, Tensorflow throws the following error:

2020-04-01 09:11:44.820866: W tensorflow/core/common_runtime/base_collective_executor.cc:216] BaseCollectiveExecutor::StartAbort Internal: Expect engine cache to be empty, but got 24 entries.
     [[{{node StatefulPartitionedCall/InitializeTRTResource}}]]
Error -  Expect engine cache to be empty, but got 24 entries.
     [[{{node StatefulPartitionedCall/InitializeTRTResource}}]] [Op:__inference_restored_function_body_64832]

Function call stack:
restored_function_body

The same error occurs when only one engine is saved for each model.

I use the following code to load TensorRT optimized SavedModel:

saved_model_loaded = tf.saved_model.load(
            trt_fp32_model_path,
            tags=[tag_constants.SERVING]
        )
graph_func = saved_model_loaded.signatures['serving_default']

I also tried to convert graph_func to frozen_func, but this didn't make any difference:

graph_func = saved_model_loaded.signatures[signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY]

frozen_func = convert_to_constants.convert_variables_to_constants_v2(
     graph_func)

I am using nvcr.io/nvidia/tensorflow:19.12-tf2-py3 docker container to optimize/run the models.

Is it possible at all to run simultaneously multiple TensorRT-optimized models with pre-built engines using Tensorflow? Or this can only be done using TensorRT inference server?

If case it is a valid usage scenario, what am I missing in my workflow?