diff --git a/tensorflow/python/framework/memory_checker.py b/tensorflow/python/framework/memory_checker.py index 5def0a9cd86..ff99182966a 100644 --- a/tensorflow/python/framework/memory_checker.py +++ b/tensorflow/python/framework/memory_checker.py @@ -24,9 +24,9 @@ from tensorflow.python.profiler.traceme import traceme_wrapper from tensorflow.python.util import tf_inspect try: - from tensorflow.python.platform.cpp_memory_checker import _CppMemoryChecker as CppMemoryChecker # pylint:disable=g-import-not-at-top + from tensorflow.python.platform.cpp_memory_checker import _CppMemoryChecker # pylint:disable=g-import-not-at-top except ImportError: - CppMemoryChecker = None + _CppMemoryChecker = None def _get_test_name_best_effort(): @@ -77,13 +77,13 @@ class MemoryChecker(object): self._trace_me = TraceMe('with MemoryChecker():') self._trace_me.__enter__() self._python_memory_checker = _PythonMemoryChecker() - if CppMemoryChecker: - self._cpp_memory_checker = CppMemoryChecker(_get_test_name_best_effort()) + if _CppMemoryChecker: + self._cpp_memory_checker = _CppMemoryChecker(_get_test_name_best_effort()) return self @traceme_wrapper def __exit__(self, exc_type, exc_value, traceback): - if CppMemoryChecker: + if _CppMemoryChecker: self._cpp_memory_checker.stop() self._trace_me.__exit__(exc_type, exc_value, traceback) @@ -99,7 +99,7 @@ class MemoryChecker(object): code complexity and the allcoation pattern. """ self._python_memory_checker.record_snapshot() - if CppMemoryChecker: + if _CppMemoryChecker: self._cpp_memory_checker.record_snapshot() @traceme_wrapper @@ -111,7 +111,7 @@ class MemoryChecker(object): directory provided the infra instead. """ self._python_memory_checker.report() - if CppMemoryChecker: + if _CppMemoryChecker: self._cpp_memory_checker.report() @traceme_wrapper @@ -124,7 +124,7 @@ class MemoryChecker(object): """ self._python_memory_checker.assert_no_leak_if_all_possibly_except_one() - if CppMemoryChecker: + if _CppMemoryChecker: self._cpp_memory_checker.assert_no_leak_if_all_possibly_except_one() @traceme_wrapper diff --git a/tensorflow/python/kernel_tests/BUILD b/tensorflow/python/kernel_tests/BUILD index f731b743f4b..a1b6417ec98 100644 --- a/tensorflow/python/kernel_tests/BUILD +++ b/tensorflow/python/kernel_tests/BUILD @@ -926,8 +926,6 @@ cuda_py_test( name = "resource_variable_ops_test", size = "medium", srcs = ["resource_variable_ops_test.py"], - # TODO(kkb): CppMemoryChecker is flaky without these two flags, investigate. - # # TODO(b/128347673): Re-enable. tags = ["no_windows"], tfrt_enabled = True, @@ -938,7 +936,6 @@ cuda_py_test( "//tensorflow/python:errors", "//tensorflow/python:framework_for_generated_wrappers", "//tensorflow/python:framework_test_lib", - "//tensorflow/python:memory_checker", "//tensorflow/python:resource_variable_ops", "//tensorflow/python:variables", "@absl_py//absl/testing:parameterized", diff --git a/tensorflow/python/kernel_tests/resource_variable_ops_test.py b/tensorflow/python/kernel_tests/resource_variable_ops_test.py index 1411dcf70d1..83fdfc7a33b 100644 --- a/tensorflow/python/kernel_tests/resource_variable_ops_test.py +++ b/tensorflow/python/kernel_tests/resource_variable_ops_test.py @@ -34,7 +34,6 @@ from tensorflow.python.framework import constant_op from tensorflow.python.framework import cpp_shape_inference_pb2 from tensorflow.python.framework import dtypes from tensorflow.python.framework import errors -from tensorflow.python.framework import memory_checker from tensorflow.python.framework import ops from tensorflow.python.framework import tensor_shape from tensorflow.python.framework import tensor_util @@ -1569,26 +1568,6 @@ class ResourceVariableOpsTest(test_util.TensorFlowTestCase, var.handle, indices, dtype=dtype) self.assertAllEqual(expected, result) - @test_util.run_v2_only - def testUninitializedVariableMemoryUsage(self): - # TODO(kkb): Python memory checker complains continuous `weakref` - # allocations, investigate. - if memory_checker.CppMemoryChecker is None: - self.skipTest("Requires the C++ memory checker") - - def _create_and_delete_variable(): - resource_variable_ops.UninitializedVariable( - shape=[100, 100], - dtype=dtypes.float32) - - _create_and_delete_variable() - checker = memory_checker.CppMemoryChecker( - "ResourceVariableOps.testUninitializedVariableMemoryUsage") - for _ in range(2): - _create_and_delete_variable() - checker.record_snapshot() - checker.report() - checker.assert_no_leak_if_all_possibly_except_one() if __name__ == "__main__": test.main() diff --git a/tensorflow/python/ops/resource_variable_ops.py b/tensorflow/python/ops/resource_variable_ops.py index 4806617c1af..b12a2023d54 100644 --- a/tensorflow/python/ops/resource_variable_ops.py +++ b/tensorflow/python/ops/resource_variable_ops.py @@ -1954,7 +1954,7 @@ class UninitializedVariable(BaseResourceVariable): unique_id = shared_name else: unique_id = "%s_%d" % (handle_name, ops.uid()) - shared_name = context.shared_name() + shared_name = context.shared_name(unique_id) handle = _variable_handle_from_shape_and_dtype( shape=shape, dtype=dtype,