From 23d51b67d092412a02cd0908ff2f819b2a988bc0 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 13 Apr 2020 20:13:39 -0700
Subject: [PATCH] For long (multi-day) training runs, temporary files that got
 created at the start of the run might have gotten cleaned up by the OS before
 the process finishes. So don't assume they're still around, in the atexit
 handlers.

PiperOrigin-RevId: 306364079
Change-Id: Ic456f05027d45d6b7699de2e1028da1e31dc0f8a
---
 tensorflow/python/autograph/pyct/loader.py      | 16 +++++++++++++++-
 tensorflow/python/autograph/pyct/loader_test.py |  7 +++++++
 2 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/tensorflow/python/autograph/pyct/loader.py b/tensorflow/python/autograph/pyct/loader.py
index 8dff536dbb8..6eb925bca45 100644
--- a/tensorflow/python/autograph/pyct/loader.py
+++ b/tensorflow/python/autograph/pyct/loader.py
@@ -23,6 +23,7 @@ from __future__ import division
 from __future__ import print_function
 
 import atexit
+import errno
 import importlib
 import os
 import sys
@@ -33,6 +34,19 @@ from tensorflow.python.autograph.pyct import parser
 from tensorflow.python.autograph.utils import compat_util
 
 
+def _remove_file(file_name):
+  """Remove a file, if it exists."""
+  try:
+    os.remove(file_name)
+  except OSError as e:
+    if e.errno == errno.ENOENT:
+      # The file disappeared. Ignore this. Temporary files might get
+      # cleaned up, especially if they reside in /tmp.
+      pass
+    else:
+      raise
+
+
 def load_source(source, delete_on_exit):
   """Loads the given source code as a Python module."""
   # TODO(mdan): Drop the linter verride once the CI stops running Py2.
@@ -43,7 +57,7 @@ def load_source(source, delete_on_exit):
     f.write(source)
 
   if delete_on_exit:
-    atexit.register(lambda: os.remove(file_name))
+    atexit.register(lambda: _remove_file(file_name))
 
   spec = importlib.util.spec_from_file_location(module_name, file_name)
   module = importlib.util.module_from_spec(spec)
diff --git a/tensorflow/python/autograph/pyct/loader_test.py b/tensorflow/python/autograph/pyct/loader_test.py
index dba974354b0..dcc32233bfd 100644
--- a/tensorflow/python/autograph/pyct/loader_test.py
+++ b/tensorflow/python/autograph/pyct/loader_test.py
@@ -19,6 +19,7 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
+import os
 import textwrap
 
 import gast
@@ -105,6 +106,12 @@ class LoaderTest(test.TestCase):
     self.assertEqual(
         module.f.__doc__, '日本語 Δθₜ ← Δθₜ₋₁ + ∇Q(sₜ, aₜ)(rₜ + γₜ₊₁ max Q(⋅))')
 
+  def test_cleanup(self):
+    test_source = textwrap.dedent('')
+    _, filename = loader.load_source(test_source, delete_on_exit=True)
+    # Clean up the file before loader.py tries to remove it, to check that the
+    # latter can deal with that situation.
+    os.unlink(filename)
 
 if __name__ == '__main__':
   test.main()