Improved/fixed memory-leak tests #398

python_bindings/tests/bindings_test.py

@@ -304,6 +304,155 @@ class SparseTestCase(TestCaseBase):
         self.assertEqual(len(index), 4)
         self.assertEqual(index[3], [(3, 1.0)])
 
+class MemoryLeak1TestCase(TestCaseBase):
+    def testMemoryLeak1(self):
+        process = psutil.Process(os.getpid())
+
+        np.random.seed(23)
+        data = np.random.randn(MEM_TEST_DATA_QTY, MEM_TEST_DATA_DIM).astype(np.float32)
+        query = np.random.randn(MEM_TEST_QUERY_QTY, MEM_TEST_DATA_DIM).astype(np.float32)
+        space_name = 'l2'
+
+        num_threads=4
+
+        index_time_params = {'M': 20, 
+                             'efConstruction': 100, 
+                             'indexThreadQty': num_threads,
+                             'post' : 0,
+                             'skip_optimized_index' : 1} # using non-optimized index!
+
+        query_time_params = {'efSearch': 100}
+
+        fail_qty = 0
+        test_qty = 0
+        delta_first = None
+
+        gc.collect()
+        time.sleep(0.25)
+
+        init_mem = process.memory_info().rss
+
+        for tid in range(MEM_TEST_REPEAT_QTY1):
+
+            with tempfile.NamedTemporaryFile() as tmp:
+
+                index = nmslib.init(method='hnsw', space=space_name, data_type=nmslib.DataType.DENSE_VECTOR)
+                index.addDataPointBatch(data)
+
+                index.createIndex(index_time_params) 
+                index.saveIndex(tmp.name, save_data=True)
+
+                index = None
+                gc.collect()
+
+                for iter_id in range(MEM_TEST_ITER1):
+
+                    index = nmslib.init(method='hnsw', space=space_name, data_type=nmslib.DataType.DENSE_VECTOR)
+                    index.loadIndex(tmp.name, load_data=True)
+                    index.setQueryTimeParams(query_time_params)
+
+                    if iter_id == 0 and tid == 0:
+                      delta_first = process.memory_info().rss - init_mem
+
+                    delta_curr = process.memory_info().rss - init_mem
+
+                    #print('Step %d mem deltas current: %d first: %d ratio %f' % (iter_id, delta_curr, delta_first, float(delta_curr)/max(delta_first, 1)))
+
+                    nbrs = index.knnQueryBatch(query, k = 10, num_threads = num_threads)
+
+                    nbrs = None 
+                    index = None
+
+                    gc.collect()
+
+                gc.collect()
+                time.sleep(0.25)
+                delta_last = process.memory_info().rss - init_mem
+                #print('Delta last %d' % delta_last)
+
+                test_qty += 1
+                if delta_last >= delta_first * MEM_GROW_COEFF:
+                  fail_qty += 1
+
+
+        print('')
+        print('Fail qty %d out of %d' % (fail_qty, test_qty))
+        self.assertTrue(fail_qty < MEM_TEST_ITER1 * MEM_TEST_CRIT_FAIL_RATE)
+
+class MemoryLeak2TestCase(TestCaseBase):
+    def testMemoryLeak2(self):
+        process = psutil.Process(os.getpid())
+
+        np.random.seed(23)
+        data = np.random.randn(MEM_TEST_DATA_QTY, 10).astype(np.float32)
+        query = np.random.randn(MEM_TEST_QUERY_QTY, 10).astype(np.float32)
+        space_name = 'l2'
+
+        num_threads=4
+
+        index_time_params = {'M': 20, 
+                             'efConstruction': 100, 
+                             'indexThreadQty': num_threads,
+                             'post' : 0,
+                             'skip_optimized_index' : 1} # using non-optimized index!
+
+        query_time_params = {'efSearch': 100}
+
+        gc.collect()
+        time.sleep(0.25)
+
+        init_mem = process.memory_info().rss
+
+
+        fail_qty = 0
+        test_qty = 0
+        delta_first = None
+
+        for tid in range(MEM_TEST_REPEAT_QTY2):
+            with tempfile.NamedTemporaryFile() as tmp:
+
+                gc.collect()
+                init_mem = process.memory_info().rss
+
+                delta1 = None
+
+                for iter_id in range(MEM_TEST_ITER2):
+
+                    index = nmslib.init(method='hnsw', space=space_name, data_type=nmslib.DataType.DENSE_VECTOR)
+                    index.addDataPointBatch(data)
+                    index.createIndex(index_time_params) 
+
+                    if iter_id == 0 and tid == 0:
+                      delta_first = process.memory_info().rss - init_mem
+
+                    delta_curr = process.memory_info().rss - init_mem
+
+                    #print('Step %d mem deltas current: %d first: %d ratio %f' % (iter_id, delta_curr, delta_first, float(delta_curr)/max(delta_first, 1)))
+
+                    index.setQueryTimeParams(query_time_params)
+                    nbrs = index.knnQueryBatch(query, k = 10, num_threads = num_threads)
+
+                    nbrs = None 
+                    index = None
+
+                    gc.collect()
+
+
+                gc.collect()
+                time.sleep(0.25)
+                delta_last = process.memory_info().rss - init_mem
+                #print('Delta last %d' % delta_last)
+
+                test_qty += 1
+                if delta_last >= delta_first * MEM_GROW_COEFF:
+                  fail_qty += 1
+
+
+        print('')
+        print('Fail qty %d out of %d' % (fail_qty, test_qty))
+        self.assertTrue(fail_qty < MEM_TEST_ITER2 * MEM_TEST_CRIT_FAIL_RATE)
+
+
 
 class GlobalTestCase(TestCaseBase):
     def testGlobal(self):