Repository: numba/numba

Issue:
Use CPython allocator in NRT
NRT should optionally use the CPython memory allocation functions (when imported from CPython). This would allow Numba-allocated memory to be seen by other utilities such as `sys.getallocatedblocks()`, `sys.debugmallocstats()`, and `tracemalloc`.



Discussion:
Actually, since the NRT must be able to run GIL-less, only `tracemalloc` will see NRT-allocated data.


Resolving patch:
diff --git a/numba/runtime/nrt.py b/numba/runtime/nrt.py
--- a/numba/runtime/nrt.py
+++ b/numba/runtime/nrt.py
@@ -98,7 +98,8 @@ def get_allocation_stats(self):
 # Alias to _nrt_python._MemInfo
 MemInfo = _nrt._MemInfo
 
-# Create uninitialized runtime
+# Create runtime
+_nrt.memsys_use_cpython_allocator()
 rtsys = _Runtime()
 
 # Install finalizer


Test patch:
diff --git a/numba/tests/test_nrt.py b/numba/tests/test_nrt.py
--- a/numba/tests/test_nrt.py
+++ b/numba/tests/test_nrt.py
@@ -1,7 +1,13 @@
 from __future__ import absolute_import, division, print_function
 
+import math
+import os
+import sys
+
 import numpy as np
+
 from numba import unittest_support as unittest
+from numba import njit
 from numba.runtime import rtsys
 from numba.config import PYVERSION
 from .support import MemoryLeakMixin
@@ -164,15 +170,71 @@ def test_buffer(self):
         # consumed by another thread.
 
 
+@unittest.skipUnless(sys.version_info >= (3, 4),
+                     "need Python 3.4+ for the tracemalloc module")
+class TestTracemalloc(unittest.TestCase):
+    """
+    Test NRT-allocated memory can be tracked by tracemalloc.
+    """
+
+    def measure_memory_diff(self, func):
+        import tracemalloc
+        tracemalloc.start()
+        try:
+            before = tracemalloc.take_snapshot()
+            # Keep the result and only delete it after taking a snapshot
+            res = func()
+            after = tracemalloc.take_snapshot()
+            del res
+            return after.compare_to(before, 'lineno')
+        finally:
+            tracemalloc.stop()
+
+    def test_snapshot(self):
+        N = 1000000
+        dtype = np.int8
+
+        @njit
+        def alloc_nrt_memory():
+            """
+            Allocate and return a large array.
+            """
+            return np.empty(N, dtype)
+
+        def keep_memory():
+            return alloc_nrt_memory()
+
+        def release_memory():
+            alloc_nrt_memory()
+
+        alloc_lineno = keep_memory.__code__.co_firstlineno + 1
+
+        # Warmup JIT
+        alloc_nrt_memory()
+
+        # The large NRT-allocated array should appear topmost in the diff
+        diff = self.measure_memory_diff(keep_memory)
+        stat = diff[0]
+        # There is a slight overhead, so the allocated size won't exactly be N
+        self.assertGreaterEqual(stat.size, N)
+        self.assertLess(stat.size, N * 1.01)
+        frame = stat.traceback[0]
+        self.assertEqual(os.path.basename(frame.filename), "test_nrt.py")
+        self.assertEqual(frame.lineno, alloc_lineno)
+
+        # If NRT memory is released before taking a snapshot, it shouldn't
+        # appear.
+        diff = self.measure_memory_diff(release_memory)
+        stat = diff[0]
+        # Something else appears, but nothing the magnitude of N
+        self.assertLess(stat.size, N * 0.01)
+
+
 class TestNRTIssue(MemoryLeakMixin, unittest.TestCase):
     def test_issue_with_refct_op_pruning(self):
         """
         GitHub Issue #1244 https://github.com/numba/numba/issues/1244
         """
-        from numba import njit
-        import numpy as np
-        import math
-
         @njit
         def calculate_2D_vector_mag(vector):
             x, y = vector