Merge pull request #26618 from trevor-m:tmorris_tftrt_dont_rebuild_failed_engines

PiperOrigin-RevId: 238143344
diff --git a/tensorflow/compiler/tf2tensorrt/kernels/trt_engine_op.cc b/tensorflow/compiler/tf2tensorrt/kernels/trt_engine_op.cc
index 0f800d7..30f2990 100644
--- a/tensorflow/compiler/tf2tensorrt/kernels/trt_engine_op.cc
+++ b/tensorflow/compiler/tf2tensorrt/kernels/trt_engine_op.cc
@@ -376,9 +376,9 @@
   }
   EngineContext* engine_context = GetEngine(input_shapes, ctx);
   if (!engine_context->cuda_engine) {
-    LOG(WARNING) << "Engine retrieval for input shapes: "
-                 << TensorShapeUtils::ShapeListString(input_shapes)
-                 << " failed. Running native segment for " << name();
+    VLOG(1) << "Engine retrieval for input shapes: "
+            << TensorShapeUtils::ShapeListString(input_shapes)
+            << " failed. Running native segment for " << name();
     ExecuteNativeSegment(ctx, helper);
     return;
   }
@@ -625,14 +625,12 @@
         partial_shapes, &logger, allocator, calibrator_.get(), &engine,
         use_calibration_, &convert_successfully);
     if (!status.ok()) {
-      if (convert_successfully) {
-        // This means it fail to build the engine even when the network is built
-        // successfully, probably due to internal issues. In this case we don't
-        // retry in the future.
-        cache.emplace(engine_input_shapes, absl::make_unique<EngineContext>());
-      }
-      LOG(WARNING) << "Engine creation for batch size " << batch_size
-                   << " failed " << status;
+      LOG(WARNING) << "Engine creation for " << name() << " failed. "
+                   << "The native segment will be used instead. "
+                   << "Reason: " << status;
+      // Store an empty engine in the cache for these input shapes so we don't
+      // try to build the same failing engine again.
+      cache.emplace(engine_input_shapes, absl::make_unique<EngineContext>());
       return &empty_context;
     }
     VLOG(1) << "Conversion is done";