Cache free var detection result based on function qualname and its module name

PiperOrigin-RevId: 480932683

tensorflow/core/function/capture/free_vars_detect.py

@@ -299,3 +299,30 @@ def generate_free_var_logging(fn, fn_threshold=5, var_threshold=10):
     logging_txt.append(ellipsis_line)
 
   return "\n".join(logging_txt)
+
+
+class FreevarDetector():
+  """Generate logging string for free vars detection and cache results."""
+
+  def __init__(self):
+    self.cache = dict()
+
+  def logging_free_vars(self, fn):
+    """Return logging string for free vars detection."""
+    if not (hasattr(fn, "__module__") and hasattr(fn, "__qualname__")):
+      return None
+
+    fn_key = (fn.__module__, fn.__qualname__)
+
+    # To prevent log spam, only generate logging once for each function
+    if fn_key in self.cache:
+      return None
+
+    try:
+      logging_txt = generate_free_var_logging(fn)
+    except Exception:  # pylint: disable=broad-except
+      # Only for logging purpose, do not raise errors to users
+      logging_txt = None
+    self.cache[fn_key] = logging_txt
+
+    return self.cache[fn_key]

tensorflow/core/function/capture/free_vars_detect_test.py

@@ -692,5 +692,20 @@ class GenerateLoggingTest(parameterized.TestCase):
     self.assertEqual(lines[2], "...")
 
 
+class FreevarDetectorTest(parameterized.TestCase):
+
+  def test_func_second_call_return_none(self):
+    x = 1
+
+    def f():
+      return x
+
+    detector = free_vars_detect.FreevarDetector()
+    logging_txt = detector.logging_free_vars(f)
+    self.assertIsNotNone(logging_txt)
+    logging_txt = detector.logging_free_vars(f)
+    self.assertIsNone(logging_txt)
+
+
 if __name__ == "__main__":
   unittest.main()