From 8303d635954cdda840afc7151598402f7600b679 Mon Sep 17 00:00:00 2001
From: Shanqing Cai <cais@google.com>
Date: Thu, 10 Oct 2019 09:00:23 -0700
Subject: [PATCH] [tfdbg] Blanket replacement of shutil.rmtree() w/
 file_io.delete_recursively()

- file_io.delete_recursively() works better across platforms (e.g., on Windows)

This will fix the ongoing test breakage of debug_v2_ops_test on Windows.

PiperOrigin-RevId: 273968373
---
 tensorflow/python/debug/BUILD                        |  1 -
 tensorflow/python/debug/cli/analyzer_cli_test.py     | 12 ++++++------
 tensorflow/python/debug/cli/cli_config_test.py       |  4 ++--
 tensorflow/python/debug/cli/readline_ui_test.py      |  4 ++--
 tensorflow/python/debug/lib/debug_data_test.py       |  4 ++--
 .../python/debug/lib/debug_events_writer_test.py     |  4 ++--
 tensorflow/python/debug/lib/debug_gradients_test.py  |  4 ++--
 .../debug/lib/debug_graph_reconstruction_test.py     |  4 ++--
 tensorflow/python/debug/lib/debug_grappler_test.py   |  4 ++--
 tensorflow/python/debug/lib/debug_v2_ops_test.py     |  4 ++--
 .../python/debug/lib/grpc_debug_test_server.py       |  4 ++--
 .../python/debug/lib/session_debug_file_test.py      |  4 ++--
 .../python/debug/lib/session_debug_grpc_test.py      |  6 +++---
 .../python/debug/lib/session_debug_multi_gpu_test.py |  4 ++--
 tensorflow/python/debug/lib/session_debug_testlib.py | 10 +++++-----
 tensorflow/python/debug/lib/source_utils_test.py     |  6 +++---
 .../python/debug/wrappers/dumping_wrapper_test.py    |  4 ++--
 tensorflow/python/debug/wrappers/framework_test.py   |  4 ++--
 .../python/debug/wrappers/local_cli_wrapper.py       |  4 ++--
 .../python/debug/wrappers/local_cli_wrapper_test.py  |  4 ++--
 20 files changed, 47 insertions(+), 48 deletions(-)

diff --git a/tensorflow/python/debug/BUILD b/tensorflow/python/debug/BUILD
index 13958b2d498..9caeec40226 100644
--- a/tensorflow/python/debug/BUILD
+++ b/tensorflow/python/debug/BUILD
@@ -751,7 +751,6 @@ py_test(
     python_version = "PY2",
     srcs_version = "PY2AND3",
     tags = [
-        "no_windows",
         "oss_serial",
     ],
     deps = [
diff --git a/tensorflow/python/debug/cli/analyzer_cli_test.py b/tensorflow/python/debug/cli/analyzer_cli_test.py
index 75bd47c5a80..58a0f3546f8 100644
--- a/tensorflow/python/debug/cli/analyzer_cli_test.py
+++ b/tensorflow/python/debug/cli/analyzer_cli_test.py
@@ -18,7 +18,6 @@ from __future__ import division
 from __future__ import print_function
 
 import os
-import shutil
 import tempfile
 
 import numpy as np
@@ -38,6 +37,7 @@ from tensorflow.python.debug.lib import debug_utils
 from tensorflow.python.debug.lib import source_utils
 from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import test_util
+from tensorflow.python.lib.io import file_io
 from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import control_flow_ops
 from tensorflow.python.ops import math_ops
@@ -636,8 +636,8 @@ class AnalyzerCLISimpleMulAddTest(test_util.TensorFlowTestCase):
   @classmethod
   def tearDownClass(cls):
     # Tear down temporary dump directory.
-    shutil.rmtree(cls._dump_root)
-    shutil.rmtree(cls._dump_root_for_unique)
+    file_io.delete_recursively(cls._dump_root)
+    file_io.delete_recursively(cls._dump_root_for_unique)
 
   def testMeasureTensorListColumnWidthsGivesRightAnswerForEmptyData(self):
     timestamp_col_width, dump_size_col_width, op_type_col_width = (
@@ -1665,7 +1665,7 @@ class AnalyzerCLIPrintLargeTensorTest(test_util.TensorFlowTestCase):
   @classmethod
   def tearDownClass(cls):
     # Tear down temporary dump directory.
-    shutil.rmtree(cls._dump_root)
+    file_io.delete_recursively(cls._dump_root)
 
   def testPrintLargeTensorWithoutAllOption(self):
     out = self._registry.dispatch_command(
@@ -1745,7 +1745,7 @@ class AnalyzerCLIControlDepTest(test_util.TensorFlowTestCase):
   @classmethod
   def tearDownClass(cls):
     # Tear down temporary dump directory.
-    shutil.rmtree(cls._dump_root)
+    file_io.delete_recursively(cls._dump_root)
 
   def testNodeInfoWithControlDependencies(self):
     # Call node_info on a node with control inputs.
@@ -2061,7 +2061,7 @@ class AnalyzerCLIWhileLoopTest(test_util.TensorFlowTestCase):
   @classmethod
   def tearDownClass(cls):
     # Tear down temporary dump directory.
-    shutil.rmtree(cls._dump_root)
+    file_io.delete_recursively(cls._dump_root)
 
   def testMultipleDumpsPrintTensorNoNumber(self):
     output = self._registry.dispatch_command("pt", ["while/Identity:0"])
diff --git a/tensorflow/python/debug/cli/cli_config_test.py b/tensorflow/python/debug/cli/cli_config_test.py
index f2b44aa6370..e8785b4b65c 100644
--- a/tensorflow/python/debug/cli/cli_config_test.py
+++ b/tensorflow/python/debug/cli/cli_config_test.py
@@ -19,11 +19,11 @@ from __future__ import print_function
 
 import json
 import os
-import shutil
 import tempfile
 
 from tensorflow.python.debug.cli import cli_config
 from tensorflow.python.framework import test_util
+from tensorflow.python.lib.io import file_io
 from tensorflow.python.platform import gfile
 from tensorflow.python.platform import googletest
 
@@ -37,7 +37,7 @@ class CLIConfigTest(test_util.TensorFlowTestCase):
     super(CLIConfigTest, self).setUp()
 
   def tearDown(self):
-    shutil.rmtree(self._tmp_dir)
+    file_io.delete_recursively(self._tmp_dir)
     super(CLIConfigTest, self).tearDown()
 
   def testConstructCLIConfigWithoutFile(self):
diff --git a/tensorflow/python/debug/cli/readline_ui_test.py b/tensorflow/python/debug/cli/readline_ui_test.py
index a3cfc7b35f2..267a158edf4 100644
--- a/tensorflow/python/debug/cli/readline_ui_test.py
+++ b/tensorflow/python/debug/cli/readline_ui_test.py
@@ -18,7 +18,6 @@ from __future__ import division
 from __future__ import print_function
 
 import os
-import shutil
 import argparse
 import tempfile
 
@@ -27,6 +26,7 @@ from tensorflow.python.debug.cli import debugger_cli_common
 from tensorflow.python.debug.cli import readline_ui
 from tensorflow.python.debug.cli import ui_factory
 from tensorflow.python.framework import test_util
+from tensorflow.python.lib.io import file_io
 from tensorflow.python.platform import gfile
 from tensorflow.python.platform import googletest
 
@@ -62,7 +62,7 @@ class CursesTest(test_util.TensorFlowTestCase):
     super(CursesTest, self).setUp()
 
   def tearDown(self):
-    shutil.rmtree(self._tmp_dir)
+    file_io.delete_recursively(self._tmp_dir)
     super(CursesTest, self).tearDown()
 
   def _babble(self, args, screen_info=None):
diff --git a/tensorflow/python/debug/lib/debug_data_test.py b/tensorflow/python/debug/lib/debug_data_test.py
index 7ce7ef6a979..6796187ac66 100644
--- a/tensorflow/python/debug/lib/debug_data_test.py
+++ b/tensorflow/python/debug/lib/debug_data_test.py
@@ -19,7 +19,6 @@ from __future__ import print_function
 
 import os
 import platform
-import shutil
 import tempfile
 
 import numpy as np
@@ -28,6 +27,7 @@ from tensorflow.core.framework import graph_pb2
 from tensorflow.core.framework import tensor_pb2
 from tensorflow.python.debug.lib import debug_data
 from tensorflow.python.framework import test_util
+from tensorflow.python.lib.io import file_io
 from tensorflow.python.platform import gfile
 from tensorflow.python.platform import googletest
 from tensorflow.python.platform import test
@@ -156,7 +156,7 @@ class DebugDumpDirTest(test_util.TensorFlowTestCase):
 
   def tearDown(self):
     # Tear down temporary dump directory.
-    shutil.rmtree(self._dump_root)
+    file_io.delete_recursively(self._dump_root)
 
   def _makeDataDirWithMultipleDevicesAndDuplicateNodeNames(self):
     cpu_0_dir = os.path.join(
diff --git a/tensorflow/python/debug/lib/debug_events_writer_test.py b/tensorflow/python/debug/lib/debug_events_writer_test.py
index 4a9fa10e2ca..97aadc38a29 100644
--- a/tensorflow/python/debug/lib/debug_events_writer_test.py
+++ b/tensorflow/python/debug/lib/debug_events_writer_test.py
@@ -20,7 +20,6 @@ from __future__ import print_function
 
 import glob
 import os
-import shutil
 import tempfile
 import threading
 
@@ -28,6 +27,7 @@ from tensorflow.core.protobuf import debug_event_pb2
 from tensorflow.python.debug.lib import debug_events_writer
 from tensorflow.python.framework import ops
 from tensorflow.python.framework import test_util
+from tensorflow.python.lib.io import file_io
 from tensorflow.python.lib.io import tf_record
 from tensorflow.python.platform import googletest
 
@@ -53,7 +53,7 @@ class PywrapeventsWriterTest(test_util.TensorFlowTestCase):
 
   def tearDown(self):
     if os.path.isdir(self.dump_root):
-      shutil.rmtree(self.dump_root, ignore_errors=True)
+      file_io.delete_recursively(self.dump_root)
     super(PywrapeventsWriterTest, self).tearDown()
 
   def testMultiThreadedConstructorCallWorks(self):
diff --git a/tensorflow/python/debug/lib/debug_gradients_test.py b/tensorflow/python/debug/lib/debug_gradients_test.py
index e592e46095c..92d31171133 100644
--- a/tensorflow/python/debug/lib/debug_gradients_test.py
+++ b/tensorflow/python/debug/lib/debug_gradients_test.py
@@ -18,7 +18,6 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
-import shutil
 import tempfile
 
 from tensorflow.core.protobuf import config_pb2
@@ -29,6 +28,7 @@ from tensorflow.python.debug.lib import debug_gradients
 from tensorflow.python.debug.lib import debug_utils
 from tensorflow.python.framework import ops
 from tensorflow.python.framework import test_util
+from tensorflow.python.lib.io import file_io
 from tensorflow.python.ops import gradients_impl
 from tensorflow.python.ops import math_ops
 from tensorflow.python.ops import variables
@@ -377,7 +377,7 @@ class IdentifyGradientTest(test_util.TensorFlowTestCase):
       debug_gradients.gradient_values_from_dump(grad_debugger, self.v, dump)
 
     # Cleanup.
-    shutil.rmtree(dump_dir)
+    file_io.delete_recursively(dump_dir)
 
 
 if __name__ == "__main__":
diff --git a/tensorflow/python/debug/lib/debug_graph_reconstruction_test.py b/tensorflow/python/debug/lib/debug_graph_reconstruction_test.py
index 34030c0adca..85d52503b7f 100644
--- a/tensorflow/python/debug/lib/debug_graph_reconstruction_test.py
+++ b/tensorflow/python/debug/lib/debug_graph_reconstruction_test.py
@@ -17,7 +17,6 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
-import shutil
 import tempfile
 
 from tensorflow.core.framework import graph_pb2
@@ -30,6 +29,7 @@ from tensorflow.python.debug.lib import debug_utils
 from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import ops
 from tensorflow.python.framework import test_util
+from tensorflow.python.lib.io import file_io
 from tensorflow.python.ops import control_flow_ops
 from tensorflow.python.ops import math_ops
 from tensorflow.python.ops import variables
@@ -57,7 +57,7 @@ class ReconstructNonDebugGraphTest(test_util.TensorFlowTestCase):
     ops.reset_default_graph()
 
   def tearDown(self):
-    shutil.rmtree(self._dump_dir)
+    file_io.delete_recursively(self._dump_dir)
     super(ReconstructNonDebugGraphTest, self).tearDown()
 
   def _graphDefWithoutBlacklistedNodes(self, graph_def):
diff --git a/tensorflow/python/debug/lib/debug_grappler_test.py b/tensorflow/python/debug/lib/debug_grappler_test.py
index 7a3bf90a616..b96e090f378 100644
--- a/tensorflow/python/debug/lib/debug_grappler_test.py
+++ b/tensorflow/python/debug/lib/debug_grappler_test.py
@@ -18,7 +18,6 @@ from __future__ import division
 from __future__ import print_function
 
 import os
-import shutil
 import tempfile
 
 from tensorflow.core.protobuf import config_pb2
@@ -29,6 +28,7 @@ from tensorflow.python.debug.lib import debug_utils
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import ops
 from tensorflow.python.framework import test_util
+from tensorflow.python.lib.io import file_io
 from tensorflow.python.ops import math_ops
 from tensorflow.python.ops import variables
 from tensorflow.python.platform import googletest
@@ -58,7 +58,7 @@ class SessionDebugGrapplerInteractionTest(test_util.TensorFlowTestCase):
   def tearDown(self):
     ops.reset_default_graph()
     if os.path.isdir(self._dump_root):
-      shutil.rmtree(self._dump_root)
+      file_io.delete_recursively(self._dump_root)
     super(SessionDebugGrapplerInteractionTest, self).tearDown()
 
   def testArithmeticOptimizationActive(self):
diff --git a/tensorflow/python/debug/lib/debug_v2_ops_test.py b/tensorflow/python/debug/lib/debug_v2_ops_test.py
index 69c68da65e0..a3229670569 100644
--- a/tensorflow/python/debug/lib/debug_v2_ops_test.py
+++ b/tensorflow/python/debug/lib/debug_v2_ops_test.py
@@ -20,7 +20,6 @@ from __future__ import print_function
 
 import glob
 import os
-import shutil
 import tempfile
 
 import numpy as np
@@ -33,6 +32,7 @@ from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import ops
 from tensorflow.python.framework import tensor_util
 from tensorflow.python.framework import test_util
+from tensorflow.python.lib.io import file_io
 from tensorflow.python.lib.io import tf_record
 from tensorflow.python.ops import gen_debug_ops
 from tensorflow.python.ops import math_ops
@@ -89,7 +89,7 @@ class DebugIdentityV2OpTest(test_util.TensorFlowTestCase):
   def tearDown(self):
     self.writer.Close()
     if os.path.isdir(self.dump_root):
-      shutil.rmtree(self.dump_root)
+      file_io.delete_recursively(self.dump_root)
     super(DebugIdentityV2OpTest, self).tearDown()
 
   @test_util.run_in_graph_and_eager_modes
diff --git a/tensorflow/python/debug/lib/grpc_debug_test_server.py b/tensorflow/python/debug/lib/grpc_debug_test_server.py
index 9c0e5390bd8..d8c3b9cde06 100644
--- a/tensorflow/python/debug/lib/grpc_debug_test_server.py
+++ b/tensorflow/python/debug/lib/grpc_debug_test_server.py
@@ -24,7 +24,6 @@ import hashlib
 import json
 import os
 import re
-import shutil
 import tempfile
 import threading
 import time
@@ -40,6 +39,7 @@ from tensorflow.python.debug.lib import debug_utils
 from tensorflow.python.debug.lib import grpc_debug_server
 from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import errors
+from tensorflow.python.lib.io import file_io
 from tensorflow.python.ops import variables
 from tensorflow.python.util import compat
 
@@ -471,7 +471,7 @@ def _poll_server_till_success(max_attempts,
       if dump_dir:
         if os.path.isdir(
             dump_dir) and debug_data.DebugDumpDir(dump_dir).size > 0:
-          shutil.rmtree(dump_dir)
+          file_io.delete_recursively(dump_dir)
           print("Poll succeeded.")
           return True
         else:
diff --git a/tensorflow/python/debug/lib/session_debug_file_test.py b/tensorflow/python/debug/lib/session_debug_file_test.py
index 16ab815d92d..45c9fe308b3 100644
--- a/tensorflow/python/debug/lib/session_debug_file_test.py
+++ b/tensorflow/python/debug/lib/session_debug_file_test.py
@@ -18,7 +18,6 @@ from __future__ import division
 from __future__ import print_function
 
 import os
-import shutil
 import tempfile
 
 from tensorflow.core.protobuf import config_pb2
@@ -29,6 +28,7 @@ from tensorflow.python.debug.lib import session_debug_testlib
 from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import ops
 from tensorflow.python.framework import test_util
+from tensorflow.python.lib.io import file_io
 from tensorflow.python.ops import math_ops
 from tensorflow.python.ops import variables
 from tensorflow.python.platform import googletest
@@ -128,7 +128,7 @@ class SessionDebugConcurrentTest(
     ops.reset_default_graph()
     for dump_root in self._dump_roots:
       if os.path.isdir(dump_root):
-        shutil.rmtree(dump_root)
+        file_io.delete_recursively(dump_root)
 
   def _get_concurrent_debug_urls(self):
     return [("file://%s" % dump_root) for dump_root in self._dump_roots]
diff --git a/tensorflow/python/debug/lib/session_debug_grpc_test.py b/tensorflow/python/debug/lib/session_debug_grpc_test.py
index f3b187cb3e9..a0bfa543345 100644
--- a/tensorflow/python/debug/lib/session_debug_grpc_test.py
+++ b/tensorflow/python/debug/lib/session_debug_grpc_test.py
@@ -22,7 +22,6 @@ from __future__ import division
 from __future__ import print_function
 
 import os
-import shutil
 
 from six.moves import xrange  # pylint: disable=redefined-builtin
 
@@ -38,6 +37,7 @@ from tensorflow.python.debug.wrappers import hooks
 from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import ops
 from tensorflow.python.framework import test_util
+from tensorflow.python.lib.io import file_io
 from tensorflow.python.ops import math_ops
 from tensorflow.python.ops import state_ops
 from tensorflow.python.ops import variables
@@ -115,7 +115,7 @@ class SessionDebugGrpcTest(session_debug_testlib.SessionDebugTestBase):
 
   def tearDown(self):
     if os.path.isdir(self._server_dump_dir):
-      shutil.rmtree(self._server_dump_dir)
+      file_io.delete_recursively(self._server_dump_dir)
     session_debug_testlib.SessionDebugTestBase.tearDown(self)
 
   def _debug_urls(self, run_number=None):
@@ -345,7 +345,7 @@ class SessionDebugConcurrentTest(
   def tearDown(self):
     ops.reset_default_graph()
     if os.path.isdir(self._server_dump_dir):
-      shutil.rmtree(self._server_dump_dir)
+      file_io.delete_recursively(self._server_dump_dir)
 
   def _get_concurrent_debug_urls(self):
     urls = []
diff --git a/tensorflow/python/debug/lib/session_debug_multi_gpu_test.py b/tensorflow/python/debug/lib/session_debug_multi_gpu_test.py
index 8eef45392f2..e41d76fb76d 100644
--- a/tensorflow/python/debug/lib/session_debug_multi_gpu_test.py
+++ b/tensorflow/python/debug/lib/session_debug_multi_gpu_test.py
@@ -18,7 +18,6 @@ from __future__ import division
 from __future__ import print_function
 
 import os
-import shutil
 import tempfile
 
 from tensorflow.core.protobuf import config_pb2
@@ -29,6 +28,7 @@ from tensorflow.python.debug.lib import debug_utils
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import ops
 from tensorflow.python.framework import test_util
+from tensorflow.python.lib.io import file_io
 from tensorflow.python.ops import math_ops
 from tensorflow.python.ops import variables
 from tensorflow.python.platform import googletest
@@ -44,7 +44,7 @@ class SessionDebugMultiGPUTest(test_util.TensorFlowTestCase):
 
     # Tear down temporary dump directory.
     if os.path.isdir(self._dump_root):
-      shutil.rmtree(self._dump_root)
+      file_io.delete_recursively(self._dump_root)
 
   def testMultiGPUSessionRun(self):
     local_devices = device_lib.list_local_devices()
diff --git a/tensorflow/python/debug/lib/session_debug_testlib.py b/tensorflow/python/debug/lib/session_debug_testlib.py
index e2740d80705..16f92085baa 100644
--- a/tensorflow/python/debug/lib/session_debug_testlib.py
+++ b/tensorflow/python/debug/lib/session_debug_testlib.py
@@ -21,7 +21,6 @@ import collections
 import functools
 import glob
 import os
-import shutil
 import tempfile
 import threading
 
@@ -40,6 +39,7 @@ from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import errors
 from tensorflow.python.framework import ops
 from tensorflow.python.framework import test_util
+from tensorflow.python.lib.io import file_io
 from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import control_flow_ops
 from tensorflow.python.ops import data_flow_ops
@@ -112,7 +112,7 @@ class SessionDebugTestBase(test_util.TensorFlowTestCase):
 
     # Tear down temporary dump directory.
     if os.path.isdir(self._dump_root):
-      shutil.rmtree(self._dump_root)
+      file_io.delete_recursively(self._dump_root)
 
   def _debug_urls(self, run_number=None):
     raise NotImplementedError(
@@ -1145,7 +1145,7 @@ class SessionDebugTestBase(test_util.TensorFlowTestCase):
       self.assertEqual([ph.name], dump1.core_metadata.input_names)
       self.assertEqual([x.name], dump1.core_metadata.output_names)
       self.assertEqual([], dump1.core_metadata.target_nodes)
-      shutil.rmtree(self._dump_root)
+      file_io.delete_recursively(self._dump_root)
 
       # Calling run() with the same feed, same output and same debug watch
       # options should increment both session_run_index and
@@ -1160,7 +1160,7 @@ class SessionDebugTestBase(test_util.TensorFlowTestCase):
       self.assertEqual([ph.name], dump2.core_metadata.input_names)
       self.assertEqual([x.name], dump2.core_metadata.output_names)
       self.assertEqual([], dump2.core_metadata.target_nodes)
-      shutil.rmtree(self._dump_root)
+      file_io.delete_recursively(self._dump_root)
 
       run_options = config_pb2.RunOptions(output_partition_graphs=True)
       debug_utils.watch_graph(
@@ -1388,7 +1388,7 @@ class SessionDebugTestBase(test_util.TensorFlowTestCase):
 
       # Another run with the default mute_if_healthy (false) value should
       # dump all the tensors.
-      shutil.rmtree(self._dump_root)
+      file_io.delete_recursively(self._dump_root)
       _, dump = self._debug_run_and_get_dump(
           sess, y, debug_ops=["DebugNumericSummary()"])
       self.assertLessEqual(8, dump.size)
diff --git a/tensorflow/python/debug/lib/source_utils_test.py b/tensorflow/python/debug/lib/source_utils_test.py
index 07a5fc744ba..b18c9b9781e 100644
--- a/tensorflow/python/debug/lib/source_utils_test.py
+++ b/tensorflow/python/debug/lib/source_utils_test.py
@@ -19,7 +19,6 @@ from __future__ import division
 from __future__ import print_function
 
 import os
-import shutil
 import tempfile
 
 import numpy as np
@@ -32,6 +31,7 @@ from tensorflow.python.debug.lib import source_utils
 from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import ops
 from tensorflow.python.framework import test_util
+from tensorflow.python.lib.io import file_io
 from tensorflow.python.ops import control_flow_ops
 from tensorflow.python.ops import math_ops
 # Import resource_variable_ops for the variables-to-tensor implicit conversion.
@@ -147,7 +147,7 @@ class SourceHelperTest(test_util.TensorFlowTestCase):
 
   def tearDown(self):
     if os.path.isdir(self.dump_root):
-      shutil.rmtree(self.dump_root)
+      file_io.delete_recursively(self.dump_root)
     ops.reset_default_graph()
 
   def testAnnotateWholeValidSourceFileGivesCorrectResult(self):
@@ -270,7 +270,7 @@ class ListSourceAgainstDumpTest(test_util.TensorFlowTestCase):
 
   def tearDown(self):
     if os.path.isdir(self.dump_root):
-      shutil.rmtree(self.dump_root)
+      file_io.delete_recursively(self.dump_root)
     ops.reset_default_graph()
 
   def testGenerateSourceList(self):
diff --git a/tensorflow/python/debug/wrappers/dumping_wrapper_test.py b/tensorflow/python/debug/wrappers/dumping_wrapper_test.py
index ea8f31a1ae8..4a14a15562d 100644
--- a/tensorflow/python/debug/wrappers/dumping_wrapper_test.py
+++ b/tensorflow/python/debug/wrappers/dumping_wrapper_test.py
@@ -19,7 +19,6 @@ from __future__ import print_function
 
 import glob
 import os
-import shutil
 import tempfile
 import threading
 
@@ -32,6 +31,7 @@ from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import ops
 from tensorflow.python.framework import test_util
+from tensorflow.python.lib.io import file_io
 from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import state_ops
 from tensorflow.python.ops import variables
@@ -61,7 +61,7 @@ class DumpingDebugWrapperSessionTest(test_util.TensorFlowTestCase):
   def tearDown(self):
     ops.reset_default_graph()
     if os.path.isdir(self.session_root):
-      shutil.rmtree(self.session_root)
+      file_io.delete_recursively(self.session_root)
 
   def _assert_correct_run_subdir_naming(self, run_subdir):
     self.assertStartsWith(run_subdir, "run_")
diff --git a/tensorflow/python/debug/wrappers/framework_test.py b/tensorflow/python/debug/wrappers/framework_test.py
index aa070d44284..0265501e625 100644
--- a/tensorflow/python/debug/wrappers/framework_test.py
+++ b/tensorflow/python/debug/wrappers/framework_test.py
@@ -18,7 +18,6 @@ from __future__ import division
 from __future__ import print_function
 
 import os
-import shutil
 import tempfile
 import threading
 
@@ -34,6 +33,7 @@ from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import errors
 from tensorflow.python.framework import ops
 from tensorflow.python.framework import test_util
+from tensorflow.python.lib.io import file_io
 from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import math_ops
 # Import resource_variable_ops for the variables-to-tensor implicit conversion.
@@ -197,7 +197,7 @@ class DebugWrapperSessionTest(test_util.TensorFlowTestCase):
   def tearDown(self):
     # Tear down temporary dump directory.
     if os.path.isdir(self._dump_root):
-      shutil.rmtree(self._dump_root)
+      file_io.delete_recursively(self._dump_root)
 
     ops.reset_default_graph()
 
diff --git a/tensorflow/python/debug/wrappers/local_cli_wrapper.py b/tensorflow/python/debug/wrappers/local_cli_wrapper.py
index 5f7fec5bfab..d8f4b3b8e8a 100644
--- a/tensorflow/python/debug/wrappers/local_cli_wrapper.py
+++ b/tensorflow/python/debug/wrappers/local_cli_wrapper.py
@@ -19,7 +19,6 @@ from __future__ import print_function
 
 import argparse
 import os
-import shutil
 import sys
 import tempfile
 
@@ -34,6 +33,7 @@ from tensorflow.python.debug.cli import ui_factory
 from tensorflow.python.debug.lib import common
 from tensorflow.python.debug.lib import debug_data
 from tensorflow.python.debug.wrappers import framework
+from tensorflow.python.lib.io import file_io
 
 
 _DUMP_ROOT_PREFIX = "tfdbg_"
@@ -375,7 +375,7 @@ class LocalCLIDebugWrapperSession(framework.BaseDebugWrapperSession):
 
   def _remove_dump_root(self):
     if os.path.isdir(self._dump_root):
-      shutil.rmtree(self._dump_root)
+      file_io.delete_recursively(self._dump_root)
 
   def _prep_debug_cli_for_run_end(self,
                                   debug_dump,
diff --git a/tensorflow/python/debug/wrappers/local_cli_wrapper_test.py b/tensorflow/python/debug/wrappers/local_cli_wrapper_test.py
index 83222f2cbe1..f252ba9bf9e 100644
--- a/tensorflow/python/debug/wrappers/local_cli_wrapper_test.py
+++ b/tensorflow/python/debug/wrappers/local_cli_wrapper_test.py
@@ -18,7 +18,6 @@ from __future__ import division
 from __future__ import print_function
 
 import os
-import shutil
 import tempfile
 
 import numpy as np
@@ -36,6 +35,7 @@ from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import errors
 from tensorflow.python.framework import ops
 from tensorflow.python.framework import test_util
+from tensorflow.python.lib.io import file_io
 from tensorflow.python.keras import backend
 from tensorflow.python.keras.engine import sequential
 from tensorflow.python.keras.layers import core
@@ -175,7 +175,7 @@ class LocalCLIDebugWrapperSessionTest(test_util.TensorFlowTestCase):
   def tearDown(self):
     ops.reset_default_graph()
     if os.path.isdir(self._tmp_dir):
-      shutil.rmtree(self._tmp_dir)
+      file_io.delete_recursively(self._tmp_dir)
 
   def testConstructWrapper(self):
     local_cli_wrapper.LocalCLIDebugWrapperSession(