Find Python Memory Growth with tracemalloc Batch Snapshots

A Python batch job uses a little more memory after every round, but a rising system-monitor graph does not show which line created the data that remains alive. Use standard-library tracemalloc snapshots to find where retained data was allocated, then compare against a run that discards the payload.

This tutorial is for beginner-to-intermediate readers. It needs no third-party packages and does not exhaust server memory. For a job that stops progressing rather than growing in memory, start with faulthandler thread-stack troubleshooting.

1. Method: compare adjacent batches

A single total mixes initialization with ongoing accumulation. Capture a baseline, finish a batch, and compare snapshots. If the same allocation line keeps growing across rounds, inspect what retains those objects.

The official Python tracemalloc documentation describes snapshot comparisons, line statistics, and current/peak counters. Our report separates net growth attributed to this source file from current and peak values across all traced allocations.

Each round creates four 64 KiB bytearray objects. retain appends them to a list that survives across batches; discard only totals their lengths and retains no payload after the function returns. Both process 262144 bytes per round. The maximum is 20 rounds, bounding retained payload at 5 MiB, plus interpreter and diagnostic overhead.

2. Set up the environment

1. Use CPython 3.10 or newer and create a practice directory:

python3 --version
mkdir tracemalloc-demo
cd tracemalloc-demo

2. Save the complete code below as memory_demo.py in that directory. Run subsequent commands there; rerunning overwrites the two example JSONL reports.

The example was tested on Linux with CPython 3.11.4. This is a standalone diagnostic exercise that stops tracing on exit; do not transplant its entire entry point into a service with existing diagnostic configuration.

3. Complete implementation

"""Compare bounded retain/discard batches with CPython tracemalloc."""
import argparse
import gc
import json
from pathlib import Path
import tracemalloc

CHUNK_BYTES = 64 * 1024
BLOCKS_PER_ROUND = 4


def run_batch(cache, retain):
    processed = 0
    for _ in range(BLOCKS_PER_ROUND):
        chunk = bytearray(CHUNK_BYTES)
        processed += len(chunk)
        if retain:
            cache.append(chunk)
    return processed


def main():
    parser = argparse.ArgumentParser(description=__doc__, allow_abbrev=False)
    parser.add_argument("--mode", choices=("retain", "discard"), default="retain")
    parser.add_argument("--rounds", type=int, default=3)
    args = parser.parse_args()
    if not 1 <= args.rounds <= 20:
        parser.error("--rounds must be between 1 and 20")

    if not tracemalloc.is_tracing():
        tracemalloc.start(1)
    source = str(Path(__file__).resolve())
    filters = (tracemalloc.Filter(True, source),)
    cache = []
    gc.collect()
    before = tracemalloc.take_snapshot().filter_traces(filters)
    try:
        for number in range(1, args.rounds + 1):
            tracemalloc.reset_peak()
            processed = run_batch(cache, args.mode == "retain")
            gc.collect()
            current, peak = tracemalloc.get_traced_memory()
            after = tracemalloc.take_snapshot().filter_traces(filters)
            changes = after.compare_to(before, "lineno")
            positive = [item for item in changes if item.size_diff > 0]
            positive.sort(key=lambda item: item.size_diff, reverse=True)
            report = {
                "round": number,
                "mode": args.mode,
                "processed_bytes": processed,
                "retained_blocks": len(cache),
                "source_growth_kib": round(sum(s.size_diff for s in changes) / 1024, 1),
                "traced_current_kib": round(current / 1024, 1),
                "traced_peak_kib": round(peak / 1024, 1),
                "top_growth": [
                    {
                        "file": s.traceback[0].filename,
                        "line": s.traceback[0].lineno,
                        "size_diff_kib": round(s.size_diff / 1024, 1),
                        "count_diff": s.count_diff,
                    }
                    for s in positive[:3]
                ],
            }
            print(json.dumps(report, ensure_ascii=True), flush=True)
            before = after
    finally:
        cache.clear()
        tracemalloc.stop()


if __name__ == "__main__":
    main()

The cache deliberately survives across rounds. Calling gc.collect() before each measurement reduces collection-timing differences in this controlled comparison. Data still referenced by the list remains alive; collection is not the fix for the production job.

The filter includes allocations made directly in this file, making the bytearray line easy to find. For a real project, adjust the filter to cover relevant application files. Allocations inside third-party libraries may be excluded by this filter.

4. Run both experiments and read the reports

1. Inspect the arguments and run two separate processes:

python3 memory_demo.py --help
python3 memory_demo.py --mode retain --rounds 3 > retain.jsonl
python3 memory_demo.py --mode discard --rounds 3 > discard.jsonl
cat retain.jsonl
cat discard.jsonl

2. Each report contains three JSON lines. Focus on these fields:

Line numbers, object overhead, and small fluctuations depend on the interpreter. Do not turn one exact byte count into a business threshold. source_growth_kib is a signed net change, while top_growth includes only positive changes, so their sums can differ. The peak counter is read before taking the second snapshot; it does not cover the entire subsequent reporting process.

3. Verify the experiment automatically, using broad ranges rather than exact object-header sizes:

python3 - <<'PY_CHECK'
import json
from pathlib import Path

def read(name):
    return [json.loads(line) for line in Path(name).read_text().splitlines()]

kept = read("retain.jsonl")
discarded = read("discard.jsonl")
assert len(kept) == len(discarded) == 3
assert [r["retained_blocks"] for r in kept] == [4, 8, 12]
assert all(r["processed_bytes"] == 262144 for r in kept + discarded)
assert all(r["retained_blocks"] == 0 for r in discarded)
assert all(r["source_growth_kib"] > 240 for r in kept)
assert all(abs(r["source_growth_kib"]) < 32 for r in discarded)
assert all(r["top_growth"][0]["size_diff_kib"] > 240 for r in kept)
print("PASS: retained growth found; discarded payload does not accumulate")
PY_CHECK

Expect PASS: retained growth found; discarded payload does not accumulate. These assertions validate this small controlled experiment, not an arbitrary memory-leak detector. For real jobs, compare equivalent input sizes and work stages across several rounds after module loading and necessary warm-up.

5. Apply the measurement and fix retention

Replace run_batch() with one batch of real work, keeping the measurements around it. Once an allocation line is identified, inspect cross-batch lists, unbounded dictionary caches, and unconsumed task results. Release references only when the application no longer needs the data.

If the cause is reading an entire CSV at once, use streaming CSV processing to aggregate as you read. If too many Futures are submitted at once, use bounded parallel batch processing to limit submissions and collect results. Rerun the same input after changing the code and verify that business results match.

To observe allocations made during startup, enable tracing earlier:

python3 -X tracemalloc=5 memory_demo.py --mode retain --rounds 3

The CPython command-line documentation documents -X tracemalloc=NFRAME. Our example does not restart tracing when it is already active. More saved frames support further analysis, although this report still groups by allocation line. The startup option does not print reports by itself.

6. Troubleshoot common misreadings

Traced size is not whole-process RSS. Python’s allocator can keep freed memory for reuse, and native extensions can allocate through untraced allocators; see the PEP 454 design rationale. An RSS value that does not immediately drop does not by itself prove that references remain. This experiment also cannot rule out every native-memory issue.

Invalid arguments fail explicitly, for example:

python3 memory_demo.py --rounds 0
printf 'exit=%s\n' "$?"

The final line should be exit=2. The example does not adjust system limits or terminate other processes.

7. Summary

Compare adjacent batches with equivalent work, then inspect object retention around allocation lines that keep growing. Reduce unnecessary references with streaming, bounded caches, or prompt result consumption, and repeat measurements while checking business output. Keep snapshot growth, batch peaks, and process RSS separate when deciding what to optimize next.