Limit number of entries in the cache.

At times the memory usage is high due to the usage of creating a new Namedtuple
type within some loop.

PiperOrigin-RevId: 199503489

1 files changed, 8 insertions, 3 deletions

diff --git a/tensorflow/python/util/util.cc b/tensorflow/python/util/util.cc
index 0dd406aa4e..c79d8a8445 100644
--- a/tensorflow/python/util/util.cc
+++ b/tensorflow/python/util/util.cc
@@ -33,6 +33,8 @@ namespace {
 PyObject* CollectionsSequenceType = nullptr;
 PyTypeObject* SparseTensorValueType = nullptr;
 
+const int kMaxItemsInCache = 1024;
+
 bool WarnedThatSetIsNotSequence = false;
 
 bool IsString(PyObject* o) {
@@ -196,11 +198,14 @@ int IsSequenceHelper(PyObject* o) {
   // NOTE: This is never decref'd, but we don't want the type to get deleted
   // as long as it is in the map. This should not be too much of a
   // leak, as there should only be a relatively small number of types in the
-  // map, and an even smaller number that are eligible for decref.
-  Py_INCREF(type);
+  // map, and an even smaller number that are eligible for decref. As a
+  // precaution, we limit the size of the map to 1024.
   {
     mutex_lock l(g_type_to_sequence_map);
-    type_to_sequence_map->insert({type, is_sequence});
+    if (type_to_sequence_map->size() < kMaxItemsInCache) {
+      Py_INCREF(type);
+      type_to_sequence_map->insert({type, is_sequence});
+    }
   }
 
   return is_sequence;