Disable leaks warning if tracing Python allocators

news/492.feature.rst

@@ -0,0 +1 @@
+When generating a ``--leaks`` flamegraph, don't show a warning that the ``pymalloc`` allocator is in use if ``--trace-python-allocators`` was used when generating the capture file.

src/memray/_memray.pyx

@@ -673,6 +673,7 @@ cdef class Tracker:
                 command_line,
                 native_traces,
                 file_format,
+                trace_python_allocators,
             )
         )
 
@@ -753,6 +754,7 @@ cdef _create_metadata(header, peak_memory):
         pid=header["pid"],
         python_allocator=allocator_id_to_name[header["python_allocator"]],
         has_native_traces=header["native_traces"],
+        trace_python_allocators=header["trace_python_allocators"],
     )
 
 

src/memray/_memray/record_reader.cpp

@@ -87,7 +87,10 @@ RecordReader::readHeader(HeaderRecord& header)
                 sizeof(header.skipped_frames_on_main_tid))
         || !d_input->read(
                 reinterpret_cast<char*>(&header.python_allocator),
-                sizeof(header.python_allocator)))
+                sizeof(header.python_allocator))
+        || !d_input->read(
+                reinterpret_cast<char*>(&header.trace_python_allocators),
+                sizeof(header.trace_python_allocators)))
     {
         throw std::ios_base::failure("Failed to read input file header.");
     }
@@ -936,7 +939,7 @@ RecordReader::dumpAllRecords()
     printf("HEADER magic=%.*s version=%d native_traces=%s file_format=%s"
            " n_allocations=%zd n_frames=%zd start_time=%lld end_time=%lld"
            " pid=%d main_tid=%lu skipped_frames_on_main_tid=%zd"
-           " command_line=%s python_allocator=%s\n",
+           " command_line=%s python_allocator=%s trace_python_allocators=%s\n",
            (int)sizeof(d_header.magic),
            d_header.magic,
            d_header.version,
@@ -950,7 +953,8 @@ RecordReader::dumpAllRecords()
            d_header.main_tid,
            d_header.skipped_frames_on_main_tid,
            d_header.command_line.c_str(),
-           python_allocator.c_str());
+           python_allocator.c_str(),
+           d_header.trace_python_allocators ? "true" : "false");
 
     switch (d_header.file_format) {
         case FileFormat::ALL_ALLOCATIONS:

src/memray/_memray/record_writer.cpp

@@ -49,7 +49,8 @@ class StreamingRecordWriter : public RecordWriter
     explicit StreamingRecordWriter(
             std::unique_ptr<memray::io::Sink> sink,
             const std::string& command_line,
-            bool native_traces);
+            bool native_traces,
+            bool trace_python_allocators);
 
     StreamingRecordWriter(StreamingRecordWriter& other) = delete;
     StreamingRecordWriter(StreamingRecordWriter&& other) = delete;
@@ -90,7 +91,8 @@ class AggregatingRecordWriter : public RecordWriter
     explicit AggregatingRecordWriter(
             std::unique_ptr<memray::io::Sink> sink,
             const std::string& command_line,
-            bool native_traces);
+            bool native_traces,
+            bool trace_python_allocators);
 
     AggregatingRecordWriter(StreamingRecordWriter& other) = delete;
     AggregatingRecordWriter(StreamingRecordWriter&& other) = delete;
@@ -138,16 +140,22 @@ createRecordWriter(
         std::unique_ptr<memray::io::Sink> sink,
         const std::string& command_line,
         bool native_traces,
-        FileFormat file_format)
+        FileFormat file_format,
+        bool trace_python_allocators)
 {
     switch (file_format) {
         case FileFormat::ALL_ALLOCATIONS:
-            return std::make_unique<StreamingRecordWriter>(std::move(sink), command_line, native_traces);
+            return std::make_unique<StreamingRecordWriter>(
+                    std::move(sink),
+                    command_line,
+                    native_traces,
+                    trace_python_allocators);
         case FileFormat::AGGREGATED_ALLOCATIONS:
             return std::make_unique<AggregatingRecordWriter>(
                     std::move(sink),
                     command_line,
-                    native_traces);
+                    native_traces,
+                    trace_python_allocators);
         default:
             throw std::runtime_error("Invalid file format enumerator");
     }
@@ -156,7 +164,8 @@ createRecordWriter(
 StreamingRecordWriter::StreamingRecordWriter(
         std::unique_ptr<memray::io::Sink> sink,
         const std::string& command_line,
-        bool native_traces)
+        bool native_traces,
+        bool trace_python_allocators)
 : RecordWriter(std::move(sink))
 , d_stats({0, 0, duration_cast<milliseconds>(system_clock::now().time_since_epoch()).count()})
 {
@@ -170,7 +179,8 @@ StreamingRecordWriter::StreamingRecordWriter(
             ::getpid(),
             0,
             0,
-            getPythonAllocator()};
+            getPythonAllocator(),
+            trace_python_allocators};
     strncpy(d_header.magic, MAGIC, sizeof(d_header.magic));
 }
 
@@ -357,7 +367,7 @@ RecordWriter::writeHeaderCommon(const HeaderRecord& header)
         or !writeSimpleType(header.stats) or !writeString(header.command_line.c_str())
         or !writeSimpleType(header.pid) or !writeSimpleType(header.main_tid)
         or !writeSimpleType(header.skipped_frames_on_main_tid)
-        or !writeSimpleType(header.python_allocator))
+        or !writeSimpleType(header.python_allocator) or !writeSimpleType(header.trace_python_allocators))
     {
         return false;
     }
@@ -383,13 +393,15 @@ StreamingRecordWriter::cloneInChildProcess()
     return std::make_unique<StreamingRecordWriter>(
             std::move(new_sink),
             d_header.command_line,
-            d_header.native_traces);
+            d_header.native_traces,
+            d_header.trace_python_allocators);
 }
 
 AggregatingRecordWriter::AggregatingRecordWriter(
         std::unique_ptr<memray::io::Sink> sink,
         const std::string& command_line,
-        bool native_traces)
+        bool native_traces,
+        bool trace_python_allocators)
 : RecordWriter(std::move(sink))
 {
     memcpy(d_header.magic, MAGIC, sizeof(d_header.magic));
@@ -399,6 +411,7 @@ AggregatingRecordWriter::AggregatingRecordWriter(
     d_header.command_line = command_line;
     d_header.pid = ::getpid();
     d_header.python_allocator = getPythonAllocator();
+    d_header.trace_python_allocators = trace_python_allocators;
 
     d_stats.start_time = duration_cast<milliseconds>(system_clock::now().time_since_epoch()).count();
 }
@@ -512,7 +525,8 @@ AggregatingRecordWriter::cloneInChildProcess()
     return std::make_unique<AggregatingRecordWriter>(
             std::move(new_sink),
             d_header.command_line,
-            d_header.native_traces);
+            d_header.native_traces,
+            d_header.trace_python_allocators);
 }
 
 bool

src/memray/_memray/record_writer.h

@@ -62,7 +62,8 @@ createRecordWriter(
         std::unique_ptr<memray::io::Sink> sink,
         const std::string& command_line,
         bool native_traces,
-        FileFormat file_format);
+        FileFormat file_format,
+        bool trace_python_allocators);
 
 template<typename T>
 bool inline RecordWriter::writeSimpleType(const T& item)

src/memray/_memray/record_writer.pxd

@@ -13,4 +13,5 @@ cdef extern from "record_writer.h" namespace "memray::api":
         string command_line,
         bool native_trace,
         FileFormat file_format,
+        bool trace_python_allocators,
     ) except+

src/memray/_memray/records.h

@@ -18,7 +18,7 @@
 namespace memray::tracking_api {
 
 extern const char MAGIC[7];  // Value assigned in records.cpp
-const int CURRENT_HEADER_VERSION = 10;
+const int CURRENT_HEADER_VERSION = 11;
 
 using frame_id_t = size_t;
 using thread_id_t = unsigned long;
@@ -118,6 +118,7 @@ struct HeaderRecord
     thread_id_t main_tid{};
     size_t skipped_frames_on_main_tid{};
     PythonAllocatorType python_allocator{};
+    bool trace_python_allocators{};
 };
 
 struct MemoryRecord

src/memray/_memray/records.pxd

@@ -34,6 +34,7 @@ cdef extern from "records.h" namespace "memray::tracking_api":
        size_t main_tid
        size_t skipped_frames_on_main_tid
        int python_allocator
+       bool trace_python_allocators
 
    cdef cppclass Allocation:
        thread_id_t tid

src/memray/_metadata.py

@@ -13,3 +13,4 @@ class Metadata:
     pid: int
     python_allocator: str
     has_native_traces: bool
+    trace_python_allocators: bool

src/memray/reporters/templates/base.html

@@ -55,16 +55,21 @@
   <main class="container-fluid">
     <div class="row">
       <div class="col bg-light py-3">
-        {% if show_memory_leaks and metadata.python_allocator == "pymalloc" %}
+        {% if show_memory_leaks and metadata.python_allocator == "pymalloc" and not metadata.trace_python_allocators %}
         <div class="alert alert-warning alert-dismissible fade show" role="alert">
-          <p><strong>Report generated using "--leaks" using pymalloc
-              allocator</strong></p>
-          <p>This report for memory leaks was generated with the
-            pymalloc allocator active. This can show confusing results because
-            the Python allocator doesn't necessarily release memory to
-            the system when Python objects are deallocated and these can still
-            appear as "leaks". If you want to exclude these, you can run your
-            application with the `PYTHONMALLOC=malloc` environment variable set.
+          <p><strong>Report generated using "--leaks" using pymalloc allocator</strong></p>
+          <p>
+            This report for memory leaks was generated with the pymalloc
+            allocator active, but without tracking enabled for calls to the
+            pymalloc allocator. This will show confusing results because the
+            pymalloc allocator will retain memory in memory pools even after
+            the objects that requested that memory are deallocated, and Memray
+            won't be able to distinguish memory set aside for reuse from leaked
+            memory. You should rerun your application with the
+            `PYTHONMALLOC=malloc` environment variable set or pass the
+            `--trace-python-allocators` flag when profiling your application.
+            <a href="https://bloomberg.github.io/memray/python_allocators.html">
+            Click here</a> for more information.
           </p>
           <button type="button" class="close" data-dismiss="alert" aria-label="Close">
             <span aria-hidden="true">&times;</span>

tests/integration/test_main.py

@@ -112,15 +112,25 @@ def _wait_until_process_blocks(pid: int) -> None:
         time.sleep(0.1)
 
 
-def generate_sample_results(tmp_path, code, *, native=False):
+def generate_sample_results(
+    tmp_path,
+    code,
+    *,
+    native=False,
+    trace_python_allocators=False,
+    disable_pymalloc=False,
+):
     results_file = tmp_path / "result.bin"
+    env = os.environ.copy()
+    env["PYTHONMALLOC"] = "malloc" if disable_pymalloc else "pymalloc"
     subprocess.run(
         [
             sys.executable,
             "-m",
             "memray",
             "run",
             *(["--native"] if native else []),
+            *(["--trace-python-allocators"] if trace_python_allocators else []),
             "--output",
             str(results_file),
             str(code),
@@ -129,6 +139,7 @@ def generate_sample_results(tmp_path, code, *, native=False):
         check=True,
         capture_output=True,
         text=True,
+        env=env,
     )
     return results_file, code
 
@@ -750,6 +761,48 @@ def test_split_threads_subcommand(self, tmp_path, simple_test_file):
         assert output_file.exists()
         assert str(source_file) in output_file.read_text()
 
+    @pytest.mark.parametrize("trace_python_allocators", [True, False])
+    @pytest.mark.parametrize("disable_pymalloc", [True, False])
+    def test_leaks_with_pymalloc_warning(
+        self,
+        tmp_path,
+        simple_test_file,
+        trace_python_allocators,
+        disable_pymalloc,
+    ):
+        results_file, _ = generate_sample_results(
+            tmp_path,
+            simple_test_file,
+            native=True,
+            trace_python_allocators=trace_python_allocators,
+            disable_pymalloc=disable_pymalloc,
+        )
+        output_file = tmp_path / "output.html"
+        warning_expected = not trace_python_allocators and not disable_pymalloc
+
+        # WHEN
+        subprocess.run(
+            [
+                sys.executable,
+                "-m",
+                "memray",
+                "flamegraph",
+                "--leaks",
+                str(results_file),
+            ],
+            check=True,
+            capture_output=True,
+            text=True,
+        )
+
+        # THEN
+        output_file = tmp_path / "memray-flamegraph-result.html"
+        assert output_file.exists()
+        assert warning_expected == (
+            'Report generated using "--leaks" using pymalloc allocator'
+            in output_file.read_text()
+        )
+
 
 class TestSummarySubCommand:
     def test_summary_generated(self, tmp_path, simple_test_file):

tests/unit/test_stats_reporter.py

@@ -99,6 +99,7 @@ def fake_stats():
             pid=123456,
             python_allocator="pymalloc",
             has_native_traces=False,
+            trace_python_allocators=True,
         ),
         total_num_allocations=20,
         total_memory_allocated=sum(mem_allocation_list),
@@ -436,6 +437,7 @@ def test_stats_output_json(fake_stats, tmp_path):
             "pid": 123456,
             "python_allocator": "pymalloc",
             "has_native_traces": False,
+            "trace_python_allocators": True,
         },
     }
     actual = json.loads(output_file.read_text())