Repository: google/jax

Issue:
Cyclic references could lead to memory leaks
Cyclic references (for example, in graph data structures) don't get garbage collected by python immediately. This can cause tensors to not be deallocated when all references are removed, causing memory leaks on device and possible OOM crashes.


Discussion:
Can you give a MWE of what you mean by in-graph data structures? You might be able to fix this with Python weakrefs if the cycles are created in data structures your code manages.
Yeah that's what we ended up doing. But weakref is a fairly low level API to be using.
Can you clarify what you mean? Is your concern that (a) references don't get collected by Python *immediately* or (b) that references don't get collected at all?

I expect you mean the former, in which case there isn't a whole lot we can do about it at our level — it's something you'll have to fix in your code. The only thing I can possibly suggest is that we could add a Delete() method to our DeviceArrays to give you more explicit control over destruction.

I'd be surprised if the latter is happening, but if you have a repro I'd be happy to look into it.
Put another way: is this an issue with JAX specifically, or would e.g. regular NumPy have this problem?
Yes adding a `Delete` method would be great. 

The issue is amplified with JAX when compared to numpy since python calls garbage collection when it starts to run out of runtime memory, but not when JAX runs out of device memory.

Note this cyclic issue is only true outside of a `jit`.

Resolving patch:
diff --git a/jax/interpreters/xla.py b/jax/interpreters/xla.py
--- a/jax/interpreters/xla.py
+++ b/jax/interpreters/xla.py
@@ -469,6 +469,8 @@ def __init__(self, result_shape, device_buffer):
   # TODO make device_buffer a property, make the _npy_value writeable, invalidate
   @property
   def _value(self):
+    if self.device_buffer is None:
+      raise ValueError("Cannot fetch the value of a deleted DeviceArray.")
     if self._npy_value is None:
       self._npy_value = self.device_buffer.to_py()
       self._npy_value.flags.writeable = False
@@ -478,6 +480,21 @@ def copy(self):
     """Returns an ndarray (backed by host memory, not device memory)."""
     return onp.asarray(self)
 
+  def delete(self):
+    """Deletes the device array and any cached copy on the host.
+
+    It is an error to access the contents of a `DeviceArray` after it has
+    been deleted.
+
+    Use of this method is optional; device buffers will be reclaimed
+    automatically by Python when a DeviceArray object is garbage collected.
+    However, it is sometimes useful to have more explicit control over the
+    time of deletion.
+    """
+    self.device_buffer.delete()
+    self.device_buffer = None
+    self._npy_value = None
+
   def __repr__(self):
     return onp.array_repr(self)
 


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -571,6 +571,12 @@ def test_devicearray_repr(self):
     self.assertIsInstance(x, DeviceArray)
     repr(x)  # doesn't crash
 
+  def test_devicearray_delete(self):
+    x = device_put(1.)
+    x.delete()
+    jtu.check_raises_regexp(lambda: repr(x), ValueError,
+                            "Cannot fetch the value of a deleted DeviceArray.")
+
   def test_namedtuple_transparency(self):
     # See https://github.com/google/jax/issues/446
     Point = collections.namedtuple("Point", ["x", "y"])
@@ -587,6 +593,5 @@ def f(pt):
     f_jit = api.jit(f)
     self.assertAllClose(f(pt), f_jit(pt), check_dtypes=False)
 
-
 if __name__ == '__main__':
   absltest.main()