Statistical Profiling in Python, Part 2: Memory
In Part 1 I wrote about sampling where a Python program spends its time. Memory usage is the other half of the question. Like time profiling, you can sample exhaustively (which is what the standard libarry tracemalloc does), but this is far too expensive to leave enabled in a production server. Unfortunately this means you often tend to enable it after you have a problem, on a process that no longer contains the interesting state.
Go’s heap profiler instead uses sampling, which reduces the overhead enough to have it on all the time. I wanted that for Python, so I wrote mprofile. It presents an API modeled after tracemalloc, but uses statistical sampling to estimate the live heap.
Sampling a stream of bytes
There is no timer interrupt for memory. Time arrives at a steady rate and the operating system will tap you on the shoulder as it passes; allocations arrive in bursts of wildly varying size, and the only place to observe them is inside the allocator, on the hot path of every running program.
The key idea comes from tcmalloc’s sampler, which is also the basis of the Go heap profiler, and which mprofile uses more or less directly. Imagine all allocated memory as one long stream of bytes. Independently mark bytes in that stream according to a Poisson process, with an average distance of \(R\) bytes between marks. An allocation is sampled if its range contains a mark.
For an allocation of \(x\) bytes, the probability of seeing at least one mark is \[P(\text{sample}) = 1 - e^{-x/R}.\]
This is useful because it samples large allocations more often than small ones while remaining unbiased with respect to the total number of bytes. A one-megabyte allocation is very likely to be selected with a 128-kilobyte sampling period; a four-kilobyte allocation is selected only occasionally. When a sample is selected, mprofile records its address, size, and Python traceback. Later, we can scale the observed samples by the inverse of the selection probability.
This approach is substantially more efficent than tracemalloc. The bookkeeping is not proportional to the allocation rate; it is proportional to the number of samples, which you choose. A program allocating furiously in a tight loop and one allocating slowly pay roughly the same profiling cost per megabyte, and you set that cost with a single parameter. In the pyperformance benchmark, using a 128-kilobyte sample period (Go’s default) incurred overhead in the tornado_http of roughly five percent.
Adapting heapprof to CPython
MProfile is a C++ extension. When profiling starts, it wraps CPython’s raw, memory, and object allocators using PyMem_SetAllocator. The wrappers call the underlying allocator first, then pass the resulting pointer and size to the profiler. free removes a sampled pointer from the live set; realloc is treated as a free followed by a new allocation. A small thread-local reentrancy guard prevents allocator calls made by the profiler itself from recursively appearing as application allocations — without it, recording an allocation allocates, which records an allocation.
There are a few details that make this less mechanical than it sounds. Some raw allocations happen without the GIL, so mprofile reacquires it before retaining references to Python code objects in a traceback.
The live pointer table is protected by a small spinlock rather than a mutex, because the critical section is a handful of instructions and putting a thread to sleep would cost more than the work it is waiting for. Each live sampled pointer costs one entry holding a trace handle and a size, stored in tcmalloc’s AddressMap — a structure built for exactly this job, keyed by pointer and allocating its own storage outside the heap it is measuring.
Stack traces are interned as a tree. Each interned frame records its location and a pointer to its parent, so a trace handle is just a pointer to a leaf frame, and two tracebacks sharing a prefix share those frames physically; walking the parent chain recovers the full trace on demand. The set holding them has to be a node-based hash set rather than a flat one, because those handles are interior pointers that must stay valid across a rehash — a small correctness constraint that quietly dictates the data structure. Filenames and function names are interned separately into a string table, so the thousandth sample from the same line stores one pointer and no strings.
The public Python objects deliberately resemble tracemalloc: snapshots can be grouped by traceback, filename, or line number; differences can be computed between snapshots; and an object’s allocation traceback can be queried while it is alive. The result is not an exact census. It is an estimate with a known sampling rate, which is usually a better trade for a long-running process.
Sampling profilers are for production
Most memory problems in a long-lived service are not mysteries about one object; they are a slow upward drift that somebody notices on a dashboard a week later. Sampling profilers make profiling efficient enough to leave them on so that you have the data you need.
Next: wiring both profilers up to HTTP, so a running Python service can be profiled the way a Go service can.
