bpo-40630: Add tracemalloc.reset_peak (GH-20102)

The reset_peak function sets the peak memory size to the current size,
representing a resetting of that metric. This allows for recording the
peak of specific sections of code, ignoring other code that may have
had a higher peak (since the most recent `tracemalloc.start()` or
tracemalloc.clear_traces()` call).

Doc/library/tracemalloc.rst

@@ -249,6 +249,47 @@ Example of output of the Python test suite::
 
 See :meth:`Snapshot.statistics` for more options.
 
+Record the current and peak size of all traced memory blocks
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+The following code computes two sums like ``0 + 1 + 2 + ...`` inefficiently, by
+creating a list of those numbers. This list consumes a lot of memory
+temporarily. We can use :func:`get_traced_memory` and :func:`reset_peak` to
+observe the small memory usage after the sum is computed as well as the peak
+memory usage during the computations::
+
+  import tracemalloc
+
+  tracemalloc.start()
+
+  # Example code: compute a sum with a large temporary list
+  large_sum = sum(list(range(100000)))
+
+  first_size, first_peak = tracemalloc.get_traced_memory()
+
+  tracemalloc.reset_peak()
+
+  # Example code: compute a sum with a small temporary list
+  small_sum = sum(list(range(1000)))
+
+  second_size, second_peak = tracemalloc.get_traced_memory()
+
+  print(f"{first_size=}, {first_peak=}")
+  print(f"{second_size=}, {second_peak=}")
+
+Output::
+
+  first_size=664, first_peak=3592984
+  second_size=804, second_peak=29704
+
+Using :func:`reset_peak` ensured we could accurately record the peak during the
+computation of ``small_sum``, even though it is much smaller than the overall
+peak size of memory blocks since the :func:`start` call. Without the call to
+:func:`reset_peak`, ``second_peak`` would still be the peak from the
+computation ``large_sum`` (that is, equal to ``first_peak``). In this case,
+both peaks are much higher than the final memory usage, and which suggests we
+could optimise (by removing the unnecessary call to :class:`list`, and writing
+``sum(range(...))``).
 
 API
 ---
@@ -289,6 +330,24 @@ Functions
    :mod:`tracemalloc` module as a tuple: ``(current: int, peak: int)``.
 
 
+.. function:: reset_peak()
+
+   Set the peak size of memory blocks traced by the :mod:`tracemalloc` module
+   to the current size.
+
+   Do nothing if the :mod:`tracemalloc` module is not tracing memory
+   allocations.
+
+   This function only modifies the recorded peak size, and does not modify or
+   clear any traces, unlike :func:`clear_traces`. Snapshots taken with
+   :func:`take_snapshot` before a call to :func:`reset_peak` can be
+   meaningfully compared to snapshots taken after the call.
+
+   See also :func:`get_traced_memory`.
+
+   .. versionadded:: 3.10
+
+
 .. function:: get_tracemalloc_memory()
 
    Get the memory usage in bytes of the :mod:`tracemalloc` module used to store

Doc/whatsnew/3.10.rst

@@ -86,6 +86,12 @@ New Modules
 Improved Modules
 ================
 
+tracemalloc
+-----------
+
+Added :func:`tracemalloc.reset_peak` to set the peak size of traced memory
+blocks to the current size, to measure the peak of specific pieces of code.
+(Contributed by Huon Wilson in :issue:`40630`.)
 
 Optimizations
 =============

Lib/test/test_tracemalloc.py

@@ -246,6 +246,30 @@ class TestTracemallocEnabled(unittest.TestCase):
         traceback2 = tracemalloc.get_object_traceback(obj)
         self.assertIsNone(traceback2)
 
+    def test_reset_peak(self):
+        # Python allocates some internals objects, so the test must tolerate
+        # a small difference between the expected size and the real usage
+        tracemalloc.clear_traces()
+
+        # Example: allocate a large piece of memory, temporarily
+        large_sum = sum(list(range(100000)))
+        size1, peak1 = tracemalloc.get_traced_memory()
+
+        # reset_peak() resets peak to traced memory: peak2 < peak1
+        tracemalloc.reset_peak()
+        size2, peak2 = tracemalloc.get_traced_memory()
+        self.assertGreaterEqual(peak2, size2)
+        self.assertLess(peak2, peak1)
+
+        # check that peak continue to be updated if new memory is allocated:
+        # peak3 > peak2
+        obj_size = 1024 * 1024
+        obj, obj_traceback = allocate_bytes(obj_size)
+        size3, peak3 = tracemalloc.get_traced_memory()
+        self.assertGreaterEqual(peak3, size3)
+        self.assertGreater(peak3, peak2)
+        self.assertGreaterEqual(peak3 - peak2, obj_size)
+
     def test_is_tracing(self):
         tracemalloc.stop()
         self.assertFalse(tracemalloc.is_tracing())

Misc/ACKS

@@ -1863,6 +1863,7 @@ Alex Willmer
 David Wilson
 Geoff Wilson
 Greg V. Wilson
+Huon Wilson
 J Derek Wilson
 Paul Winkler
 Jody Winston

Misc/NEWS.d/next/Library/2020-05-15-13-40-15.bpo-40630.YXEX_M.rst

@@ -0,0 +1,2 @@
+Added :func:`tracemalloc.reset_peak` to set the peak size of traced memory
+blocks to the current size, to measure the peak of specific pieces of code.

Modules/_tracemalloc.c

@@ -1643,6 +1643,30 @@ _tracemalloc_get_traced_memory_impl(PyObject *module)
     return Py_BuildValue("nn", size, peak_size);
 }
 
+/*[clinic input]
+_tracemalloc.reset_peak
+
+Set the peak size of memory blocks traced by tracemalloc to the current size.
+
+Do nothing if the tracemalloc module is not tracing memory allocations.
+
+[clinic start generated code]*/
+
+static PyObject *
+_tracemalloc_reset_peak_impl(PyObject *module)
+/*[clinic end generated code: output=140c2870f691dbb2 input=18afd0635066e9ce]*/
+{
+    if (!_Py_tracemalloc_config.tracing) {
+        Py_RETURN_NONE;
+    }
+
+    TABLES_LOCK();
+    tracemalloc_peak_traced_memory = tracemalloc_traced_memory;
+    TABLES_UNLOCK();
+
+    Py_RETURN_NONE;
+}
+
 
 static PyMethodDef module_methods[] = {
     _TRACEMALLOC_IS_TRACING_METHODDEF
@@ -1654,6 +1678,7 @@ static PyMethodDef module_methods[] = {
     _TRACEMALLOC_GET_TRACEBACK_LIMIT_METHODDEF
     _TRACEMALLOC_GET_TRACEMALLOC_MEMORY_METHODDEF
     _TRACEMALLOC_GET_TRACED_MEMORY_METHODDEF
+    _TRACEMALLOC_RESET_PEAK_METHODDEF
     /* sentinel */
     {NULL, NULL}
 };

Modules/clinic/_tracemalloc.c.h

@@ -197,4 +197,24 @@ _tracemalloc_get_traced_memory(PyObject *module, PyObject *Py_UNUSED(ignored))
 {
     return _tracemalloc_get_traced_memory_impl(module);
 }
-/*[clinic end generated code: output=1bc96dc569706afa input=a9049054013a1b77]*/
+
+PyDoc_STRVAR(_tracemalloc_reset_peak__doc__,
+"reset_peak($module, /)\n"
+"--\n"
+"\n"
+"Set the peak size of memory blocks traced by tracemalloc to the current size.\n"
+"\n"
+"Do nothing if the tracemalloc module is not tracing memory allocations.");
+
+#define _TRACEMALLOC_RESET_PEAK_METHODDEF    \
+    {"reset_peak", (PyCFunction)_tracemalloc_reset_peak, METH_NOARGS, _tracemalloc_reset_peak__doc__},
+
+static PyObject *
+_tracemalloc_reset_peak_impl(PyObject *module);
+
+static PyObject *
+_tracemalloc_reset_peak(PyObject *module, PyObject *Py_UNUSED(ignored))
+{
+    return _tracemalloc_reset_peak_impl(module);
+}
+/*[clinic end generated code: output=a130117b1af821da input=a9049054013a1b77]*/