Repository: TileDB-Inc/TileDB-Py

Issue:
Reading dense array doesn't free memory
Hi,

I'm wondering if this is expected behavior or if you have any tips to fix. On Ubuntu 16, Python 3.7, and  _tiledb_ 0.4.1:

Create toy array:
```
x = np.ones(10000000)
ctx = tiledb.Ctx()
path = 'test_tile_db'
d1 = tiledb.Dim(
    'test_domain', domain=(0, x.shape[0] - 1), tile=10000, dtype="uint32"
)
domain = tiledb.Domain(d1)
v = tiledb.Attr(
    'test_value',
    dtype="float32",
)
schema = tiledb.ArraySchema(
    domain=domain, attrs=(v,), cell_order="row-major", tile_order="row-major"
)
A = tiledb.DenseArray.create(path, schema)
values = x.astype(np.float32)
with tiledb.DenseArray(path, mode="w", ctx=ctx) as A:
    A[:] = {'test_value': values}
```

Read from array:
```
for i in range(10):
    with tiledb.DenseArray(path, mode='r') as data:
        data[:]
    print('Gigs:', round(psutil.virtual_memory().used / (10**9), 2))
Gigs: 0.84
Gigs: 0.89
Gigs: 0.93
Gigs: 0.97
Gigs: 1.01
Gigs: 1.05
Gigs: 1.1
Gigs: 1.14
Gigs: 1.18
Gigs: 1.22
```

Basically memory never seems to get released even when I don't assign the `data[:]` to any variable. I've tried playing around with garbage collection (`import gc; gc.collect()`) but it seems Python is not aware. Have also tried doing some explicit closing of the DenseArray. Eventually have to reset Jupyter notebook to get memory to free. 

In my real use case I am iterating over several tileDBs and pulling full array data out from each, doing some transforms, and writing new tileDBs with transformed data. Works okay except every read call adds around 2GBs to the used memory and never releases it, causing the machine to eventually run out of memory. Current work around is to spin up new process for every iteration.

Thanks!


Discussion:

Resolving patch:
diff --git a/setup.py b/setup.py
--- a/setup.py
+++ b/setup.py
@@ -36,7 +36,7 @@ def get_zipfile(url):
 from sys import version_info as ver
 
 # Target branch
-TILEDB_VERSION = "dev"
+TILEDB_VERSION = "1.5.1"
 
 # Use `setup.py [] --debug` for a debug build of libtiledb
 TILEDB_DEBUG_BUILD = False


Test patch:
diff --git a/tiledb/tests/common.py b/tiledb/tests/common.py
--- a/tiledb/tests/common.py
+++ b/tiledb/tests/common.py
@@ -4,9 +4,11 @@
 import os
 import shutil
 import tempfile
+import traceback
 from unittest import TestCase
 
 class DiskTestCase(TestCase):
+    pathmap = dict()
 
     def setUp(self):
         prefix = 'tiledb-' + self.__class__.__name__
@@ -20,8 +22,14 @@ def tearDown(self):
             except OSError as exc:
                 print("test '{}' error deleting '{}'".format(self.__class__.__name__,
                                                              dirpath))
-                raise
+                print("registered paths and originating functions:")
+                for path,frame in self.pathmap.items():
+                    print("  '{}' <- '{}'".format(path,frame))
+                raise exc
 
     def path(self, path):
-        return os.path.abspath(os.path.join(self.rootdir, path))
+        out = os.path.abspath(os.path.join(self.rootdir, path))
+        frame = traceback.extract_stack(limit=2)[-2][2]
+        self.pathmap[out] = frame
+        return out
 
diff --git a/tiledb/tests/test_libtiledb.py b/tiledb/tests/test_libtiledb.py
--- a/tiledb/tests/test_libtiledb.py
+++ b/tiledb/tests/test_libtiledb.py
@@ -1046,14 +1046,13 @@ def test_varlen_write_floats(self):
         att = tiledb.Attr(dtype=np.float64, var=True, ctx=ctx)
 
         schema = tiledb.ArraySchema(dom, (att,), ctx=ctx)
-
         tiledb.DenseArray.create(self.path("foo"), schema)
         with tiledb.DenseArray(self.path("foo"), mode='w', ctx=ctx) as T:
             T[:] = A
 
         with tiledb.DenseArray(self.path("foo"), mode='r', ctx=ctx) as T:
             T_ = T[:]
-            self.assertEqual(len(A), len(T))
+            self.assertEqual(len(A), len(T_))
             # can't use assert_array_equal w/ np.object array
             self.assertTrue(all(np.array_equal(x,A[i]) for i,x in enumerate(T_)))
 
@@ -1560,9 +1559,9 @@ def test_pickle_roundtrip(self):
             with io.BytesIO() as buf, tiledb.DenseArray(uri) as V:
                 pickle.dump(V, buf)
                 buf.seek(0)
-                V2 = pickle.load(buf)
-                # make sure anonymous view pickles and round-trips
-                assert_array_equal(V, V2)
+                with pickle.load(buf) as V2:
+                    # make sure anonymous view pickles and round-trips
+                    assert_array_equal(V, V2)
 
     def test_pickle_with_config(self):
         import io, pickle
@@ -1606,13 +1605,13 @@ def test_view_multiattr(self):
         anon_ar = np.random.rand(3, 3)
         named_ar = np.random.rand(3, 3)
 
-        with tiledb.DenseArray(uri, 'w') as T:
+        with tiledb.DenseArray(uri, 'w', ctx=ctx) as T:
             T[:] = {'': anon_ar, 'named': named_ar}
 
         with self.assertRaises(KeyError):
-            T = tiledb.DenseArray(uri, 'r', attr="foo111")
+            T = tiledb.DenseArray(uri, 'r', attr="foo111", ctx=ctx)
 
-        with tiledb.DenseArray(uri, 'r', attr="named") as T:
+        with tiledb.DenseArray(uri, 'r', attr="named", ctx=ctx) as T:
             assert_array_equal(T, named_ar)
             # make sure each attr view can pickle and round-trip
             with io.BytesIO() as buf:
@@ -1621,7 +1620,7 @@ def test_view_multiattr(self):
                 with pickle.load(buf) as T_rt:
                     assert_array_equal(T, T_rt)
 
-        with tiledb.DenseArray(uri, 'r', attr="") as T:
+        with tiledb.DenseArray(uri, 'r', attr="", ctx=ctx) as T:
             assert_array_equal(T, anon_ar)
 
             with io.BytesIO() as buf:
@@ -1632,10 +1631,10 @@ def test_view_multiattr(self):
 
         # set subarray on multi-attribute
         range_ar = np.arange(0,9).reshape(3,3)
-        with tiledb.DenseArray(uri, 'w', attr='named') as V_named:
+        with tiledb.DenseArray(uri, 'w', attr='named', ctx=ctx) as V_named:
             V_named[1:3,1:3] = range_ar[1:3,1:3]
 
-        with tiledb.DenseArray(uri, 'r', attr='named') as V_named:
+        with tiledb.DenseArray(uri, 'r', attr='named', ctx=ctx) as V_named:
             assert_array_equal(V_named[1:3,1:3], range_ar[1:3,1:3])
 
 
@@ -1749,8 +1748,7 @@ def test_kv_write_schema_load(self):
         a1 = tiledb.Attr("value", dtype=bytes, ctx=ctx)
         schema = tiledb.KVSchema(ctx, attrs=(a1,))
         # persist kv schema
-        kv = tiledb.KV.create(self.path("foo"), schema, ctx=ctx)
-        self.assertNotEqual(kv, None)
+        tiledb.KV.create(self.path("foo"), schema, ctx=ctx)
         self.assertEqual(tiledb.KVSchema.load(self.path("foo"), ctx=ctx), schema)
 
     def test_kv_contains(self):
@@ -1798,8 +1796,7 @@ def test_kv_write_consolidate(self):
         schema = tiledb.KVSchema(attrs=(a1,), ctx=ctx)
 
         # persist kv schema
-        kv = tiledb.KV.create(self.path("foo1"), schema, ctx=ctx)
-        kv.close()
+        tiledb.KV.create(self.path("foo1"), schema, ctx=ctx)
 
         def append_kv(path, k, v):
             kv = tiledb.KV(path, mode='w', ctx=ctx)
@@ -1850,18 +1847,19 @@ def test_kv_write_load_read_encrypted(self):
 
     def test_kv_update_reload(self):
         # create a kv array
-        ctx = tiledb.Ctx()
-        a1 = tiledb.Attr("val", ctx=ctx, dtype=bytes)
+        ctx1 = tiledb.Ctx()
+        ctx2 = tiledb.Ctx()
+        a1 = tiledb.Attr("val", ctx=ctx1, dtype=bytes)
         # persist kv schema
-        schema = tiledb.KVSchema(attrs=(a1,), ctx=ctx)
-        tiledb.KV.create(self.path("foo"), schema, ctx=ctx)
+        schema = tiledb.KVSchema(attrs=(a1,), ctx=ctx1)
+        tiledb.KV.create(self.path("foo"), schema, ctx=ctx1)
 
         # load kv array
-        with tiledb.KV(self.path("foo"), mode='w', ctx=ctx) as kv1:
+        with tiledb.KV(self.path("foo"), mode='w', ctx=ctx1) as kv1:
             kv1['foo'] = 'bar'
             kv1.flush()
 
-            with tiledb.KV(self.path("foo"), mode='r', ctx=ctx) as kv2:
+            with tiledb.KV(self.path("foo"), mode='r', ctx=ctx2) as kv2:
                 self.assertTrue('foo' in kv2)
                 kv1['bar'] = 'baz'
                 kv1.flush()
@@ -2113,6 +2111,75 @@ def test_io(self):
         self.assertEqual(io.readall(), b"")
 
 
+class MemoryTest(DiskTestCase):
+    # sanity check that memory usage doesn't increase more than 10% reading 40MB 100x
+    # https://github.com/TileDB-Inc/TileDB-Py/issues/150
+
+    def setUp(self):
+        super(MemoryTest, self).setUp()
+        import sys
+        if not sys.platform.startswith("linux"):
+            self.skipTest("Only run MemoryTest on linux")
+
+    @staticmethod
+    def use_many_buffers(path):
+        import psutil, os
+        # https://stackoverflow.com/questions/938733/total-memory-used-by-python-process
+        process = psutil.Process(os.getpid())
+
+        x = np.ones(10000000, dtype=np.float32)
+        ctx = tiledb.Ctx()
+        d1 = tiledb.Dim(
+            'test_domain', domain=(0, x.shape[0] - 1), tile=10000, dtype="uint32")
+        domain = tiledb.Domain(d1)
+        v = tiledb.Attr(
+            'test_value',
+            dtype="float32")
+
+        schema = tiledb.ArraySchema(
+            domain=domain, attrs=(v,), cell_order="row-major", tile_order="row-major")
+
+        A = tiledb.DenseArray.create(path, schema)
+
+        with tiledb.DenseArray(path, mode="w", ctx=ctx) as A:
+            A[:] = {'test_value': x}
+
+        with tiledb.DenseArray(path, mode='r') as data:
+            data[:]
+            initial = process.memory_info().rss
+            print("  initial RSS: {}".format(round(initial / (10 ** 6)), 2))
+            for i in range(100):
+                # read but don't store: this memory should be freed
+                data[:]
+
+                if i % 10 == 0:
+                    print('    read iter {}, RSS (MB): {}'.format(
+                        i, round(process.memory_info().rss / (10 ** 6), 2)))
+
+        return initial
+
+    def test_memory_cleanup(self):
+        import tiledb, numpy as np
+        import psutil, os
+
+        # run function which reads 100x from a 40MB test array
+        # TODO: RSS is too loose to do this end-to-end, so should use instrumentation.
+        print("Starting TileDB-Py memory test:")
+        initial = self.use_many_buffers(self.path('test_memory_cleanup'))
+
+        process = psutil.Process(os.getpid())
+        final = process.memory_info().rss
+        print("  final RSS: {}".format(round(final / (10 ** 6)), 2))
+
+        import gc
+        gc.collect()
+
+        final_gc = process.memory_info().rss
+        print("  final RSS after forced GC: {}".format(round(final_gc / (10 ** 6)), 2))
+
+        self.assertTrue((final - initial) < (.1 * initial))
+
+
 #if __name__ == '__main__':
 #    # run a single example for in-process debugging
 #    # better to use `pytest --gdb` if available