chore: import upstream snapshot with attribution

tensorflow/python/eager/memory_tests/BUILD

@@ -0,0 +1,67 @@
+load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
+load("//tensorflow:tensorflow.default.bzl", "cuda_py_strict_test")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    licenses = ["notice"],
+)
+
+# NOTE: Do not add sharding to these tests. If tests run concurrently, they
+# seem to confuse the memory_profiler, and the tests begin to flake. Add new
+# test files as needed.
+
+py_library(
+    name = "memory_test_util",
+    srcs = ["memory_test_util.py"],
+    strict_deps = True,
+    visibility = ["//tensorflow:internal"],
+    deps = ["//tensorflow/python/eager:context"],
+)
+
+cuda_py_strict_test(
+    name = "memory_test",
+    size = "medium",
+    srcs = ["memory_test.py"],
+    tags = [
+        "manual",
+        "no_oss",
+        "notap",  #TODO(b/140640597): this test is flaky at the moment
+        "optonly",  # The test is too slow in non-opt mode
+    ],
+    # TODO(b/140065350): Re-enable
+    xla_enable_strict_auto_jit = False,
+    deps = [
+        ":memory_test_util",
+        "//tensorflow/python/eager:backprop",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/eager:test",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:gradients",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:variables",
+    ],
+)
+
+cuda_py_strict_test(
+    name = "remote_memory_test",
+    size = "medium",
+    srcs = ["remote_memory_test.py"],
+    tags = [
+        "no_gpu",  # TODO(b/168058741): Enable the test for GPU
+        "optonly",  # The test is too slow in non-opt mode
+    ],
+    xla_enable_strict_auto_jit = False,  # b/140261762
+    deps = [
+        ":memory_test_util",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/eager:remote",
+        "//tensorflow/python/eager:test",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/training:server_lib",
+    ],
+)

tensorflow/python/eager/memory_tests/memory_test.py

@@ -0,0 +1,117 @@
+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for memory leaks in eager execution.
+
+It is possible that this test suite will eventually become flaky due to taking
+too long to run (since the tests iterate many times), but for now they are
+helpful for finding memory leaks since not all PyObject leaks are found by
+introspection (test_util decorators). Please be careful adding new tests here.
+"""
+
+from tensorflow.python.eager import backprop
+from tensorflow.python.eager import def_function
+from tensorflow.python.eager import test
+from tensorflow.python.eager.memory_tests import memory_test_util
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import gradients as gradient_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops.variables import Variable
+
+
+class MemoryTest(test.TestCase):
+
+  def testMemoryLeakAnonymousVariable(self):
+    if not memory_test_util.memory_profiler_is_available():
+      self.skipTest("memory_profiler required to run this test")
+
+    def f():
+      inputs = Variable(array_ops.zeros([32, 100], dtypes.float32))
+      del inputs
+
+    memory_test_util.assert_no_leak(
+        f, num_iters=10000, increase_threshold_absolute_mb=10)
+
+  def testMemoryLeakInFunction(self):
+    if not memory_test_util.memory_profiler_is_available():
+      self.skipTest("memory_profiler required to run this test")
+
+    def f():
+
+      @def_function.function
+      def graph(x):
+        return x * x + x
+
+      graph(constant_op.constant(42))
+
+    memory_test_util.assert_no_leak(
+        f, num_iters=1000, increase_threshold_absolute_mb=30)
+
+  @test_util.assert_no_new_pyobjects_executing_eagerly()
+  def testNestedFunctionsDeleted(self):
+
+    @def_function.function
+    def f(x):
+
+      @def_function.function
+      def my_sin(x):
+        return math_ops.sin(x)
+
+      return my_sin(x)
+
+    x = constant_op.constant(1.)
+    with backprop.GradientTape() as t1:
+      t1.watch(x)
+      with backprop.GradientTape() as t2:
+        t2.watch(x)
+        y = f(x)
+      dy_dx = t2.gradient(y, x)
+    dy2_dx2 = t1.gradient(dy_dx, x)
+
+    self.assertAllClose(0.84147096, y.numpy())  # sin(1.)
+    self.assertAllClose(0.54030230, dy_dx.numpy())  # cos(1.)
+    self.assertAllClose(-0.84147096, dy2_dx2.numpy())  # -sin(1.)
+
+  def testMemoryLeakInGlobalGradientRegistry(self):
+    # Past leak: b/139819011
+
+    if not memory_test_util.memory_profiler_is_available():
+      self.skipTest("memory_profiler required to run this test")
+
+    def f():
+
+      @def_function.function(autograph=False)
+      def graph(x):
+
+        @def_function.function(autograph=False)
+        def cubed(a):
+          return a * a * a
+
+        y = cubed(x)
+        # To ensure deleting the function does not affect the gradient
+        # computation.
+        del cubed
+        return gradient_ops.gradients(gradient_ops.gradients(y, x), x)
+
+      return graph(constant_op.constant(1.5))[0].numpy()
+
+    memory_test_util.assert_no_leak(
+        f, num_iters=300, increase_threshold_absolute_mb=50)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/eager/memory_tests/memory_test_util.py

@@ -0,0 +1,73 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Utils for memory tests."""
+
+import collections
+import gc
+import time
+
+from tensorflow.python.eager import context
+
+# memory_profiler might not be available in the OSS version of TensorFlow.
+try:
+  import memory_profiler  # pylint:disable=g-import-not-at-top
+except ImportError:
+  memory_profiler = None
+
+
+def _instance_count_by_class():
+  counter = collections.Counter()
+
+  for obj in gc.get_objects():
+    try:
+      counter[obj.__class__.__name__] += 1
+    except Exception:  # pylint:disable=broad-except
+      pass
+
+  return counter
+
+
+def assert_no_leak(f, num_iters=100000, increase_threshold_absolute_mb=25):
+  """Assert memory usage doesn't increase beyond given threshold for f."""
+
+  with context.eager_mode():
+    # Warm up.
+    f()
+
+    # Wait for background threads to start up and take over memory.
+    # FIXME: The nature of this test leaves few other options. Maybe there
+    # is a better way to do this.
+    time.sleep(4)
+
+    gc.collect()
+    initial = memory_profiler.memory_usage(-1)[0]
+    instance_count_by_class_before = _instance_count_by_class()
+
+    for _ in range(num_iters):
+      f()
+
+    gc.collect()
+    increase = memory_profiler.memory_usage(-1)[0] - initial
+
+    assert increase < increase_threshold_absolute_mb, (
+        "Increase is too high. Initial memory usage: %f MB. Increase: %f MB. "
+        "Maximum allowed increase: %f MB. "
+        "Instance count diff before/after: %s") % (
+            initial, increase, increase_threshold_absolute_mb,
+            _instance_count_by_class() - instance_count_by_class_before)
+
+
+def memory_profiler_is_available():
+  return memory_profiler is not None

tensorflow/python/eager/memory_tests/remote_memory_test.py

@@ -0,0 +1,60 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for memory leaks in remote eager execution."""
+
+from tensorflow.python.eager import def_function
+from tensorflow.python.eager import remote
+from tensorflow.python.eager import test
+from tensorflow.python.eager.memory_tests import memory_test_util
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.training import server_lib
+
+
+class RemoteWorkerMemoryTest(test.TestCase):
+
+  def __init__(self, method):
+    super(RemoteWorkerMemoryTest, self).__init__(method)
+
+    # used for remote worker tests
+    self._cached_server = server_lib.Server.create_local_server()
+    self._cached_server_target = self._cached_server.target[len("grpc://"):]
+
+  def testMemoryLeakInLocalCopy(self):
+    if not memory_test_util.memory_profiler_is_available():
+      self.skipTest("memory_profiler required to run this test")
+
+    remote.connect_to_remote_host(self._cached_server_target)
+
+    # Run a function locally with the input on a remote worker and ensure we
+    # do not leak a reference to the remote tensor.
+
+    @def_function.function
+    def local_func(i):
+      return i
+
+    def func():
+      with ops.device("job:worker/replica:0/task:0/device:CPU:0"):
+        x = array_ops.zeros([1000, 1000], dtypes.int32)
+
+      local_func(x)
+
+    memory_test_util.assert_no_leak(
+        func, num_iters=100, increase_threshold_absolute_mb=50)
+
+
+if __name__ == "__main__":
+  test.main()