memray — tìm memory leak Python như dân anh Bloomberg

2h sáng, Sentry nổ, một worker Python OOM restart liên tục. Mình SSH vào container, gõ top, thấy RSS nhảy từ 200MB lên 3.8GB trong 15 phút rồi bị kill. Không có crash log, không có stack trace — chỉ có dòng Killed khô khan từ kernel.

Lúc đó mình mới nhận ra: print(gc.get_objects()) và tracemalloc built-in không đủ để debug một app đang chạy với hàng trăm nghìn object. Thứ mình cần là một profiler thấy được cả C extension, chạy được trên production và cho ra flamegraph đọc được trong 5 phút.

Và rồi mình tìm được memray — memory profiler open-source của Bloomberg. Bài này mình chia sẻ cách dùng nó qua những case thực tế mình từng gặp.

Vì sao memray, mà không phải tracemalloc hay pympler

Python có sẵn tracemalloc từ 3.4. Nó ổn cho việc check nhanh, nhưng có 2 hạn chế lớn:

Không thấy C extension. Nếu leak nằm trong numpy, pandas, psycopg2 hay thư viện Rust/C++ nào đó — tracemalloc đứng nhìn.
Overhead cao nếu bật full trace. Không chạy được trên workload production.

memray thì:

Trace được cả Python và native allocations (C, C++, Rust thông qua malloc, mmap, calloc…)
Overhead ~20-30% khi record, có thể chấp nhận cho staging
Output ra nhiều report: flamegraph, tree, table, stats, summary — mỗi loại trả lời một câu hỏi khác nhau
Có live mode (xem real-time) và pytest plugin (chặn leak ngay trong CI)

Cài đặt và chạy lần đầu

pip install memray

memray chỉ chạy trên Linux và macOS — Windows anh em phải qua WSL. Bloomberg dùng nó nội bộ nên tập trung vào Unix.

Chạy app với memray attach vào:

# Record toàn bộ allocations ra file .bin
memray run -o out.bin my_app.py

# Hoặc run module
memray run -o out.bin -m gunicorn app:create_app

File .bin là binary, khá nhẹ — mình từng record một process chạy 10 phút, file ra ~150MB. Dùng --compress-on-exit sẽ giảm còn ~40MB.

Flamegraph — “công cụ số một mình dùng”

Sau khi có file .bin, sinh flamegraph:

memray flamegraph out.bin
# → tạo file memray-flamegraph-out.html

Mở file HTML trong browser. Bạn sẽ thấy một cây allocation theo call stack — chiều rộng của mỗi box là lượng memory đã cấp ở frame đó.

Cái hay của flamegraph memray:

Hover để xem chính xác bao nhiêu byte, bao nhiêu lần gọi
Click vào một frame để zoom vào nhánh đó
Toggle giữa resident set size (đang giữ) và allocated size (tổng đã cấp)

Case thực tế: mình debug một FastAPI service, flamegraph chỉ ngay một frame rộng chiếm 62% — hóa ra là pd.read_sql() load nguyên bảng 2 triệu row vào RAM rồi filter bằng Python. Fix bằng cách đẩy filter xuống SQL — RSS giảm từ 3.8GB xuống 280MB.

Các report khác — mỗi cái giải một câu hỏi

`summary` — nhìn tổng quan nhanh

memray summary out.bin

In ra terminal top các function allocate nhiều nhất. Dùng khi chưa muốn mở browser.

`tree` — call tree theo allocation

memray tree out.bin

Tốt khi bạn biết mơ hồ leak nằm đâu và muốn đi từ root xuống.

`stats` — phân bố size

memray stats out.bin

Cho biết có bao nhiêu allocation <100 byte, 100B-1KB, 1KB-1MB, >1MB. Rất hữu dụng khi nghi ngờ fragmentation — nhiều allocation nhỏ li ti cũng là leak.

`flamegraph --leaks`

memray flamegraph --leaks out.bin

Chỉ hiện những allocation chưa được free tại thời điểm stop. Đây mới thật sự là leak, không phải peak memory.

Live mode — nhìn leak xảy ra real-time

Nếu bạn muốn nhìn memory tăng theo thời gian (rất hữu ích khi debug một endpoint cụ thể):

# Terminal 1
memray run --live my_app.py

# Terminal 2 (hoặc bạn gõ tay): bật app, gọi endpoint nghi ngờ
curl localhost:8000/api/suspicious-endpoint

Terminal 1 sẽ hiện một TUI kiểu htop, top các function theo memory đang giữ, refresh mỗi giây. Bạn gọi endpoint → thấy function nào sưng lên ngay lập tức.

Mình dùng cách này để debug một bug: sau mỗi request, dict cache nội bộ cứ to thêm. Hóa ra là @functools.lru_cache(maxsize=None) trên một method lấy self làm key — mỗi request tạo instance mới, cache không bao giờ evict.

# ❌ Leak: self vào key → mỗi instance một cache entry
class UserService:
    @functools.lru_cache(maxsize=None)
    def get_permissions(self, user_id):
        return db.query(...)

# ✅ Đẩy cache ra module-level, hoặc dùng maxsize cụ thể
@functools.lru_cache(maxsize=1024)
def _get_permissions(user_id):
    return db.query(...)

class UserService:
    def get_permissions(self, user_id):
        return _get_permissions(user_id)

pytest plugin — chặn leak ngay trong CI

Đây là feature mình thích nhất. Cài thêm:

pip install pytest-memray

Rồi trong pytest.ini hoặc pyproject.toml:

[tool.pytest.ini_options]
addopts = "--memray"

Mỗi test sẽ có bảng memory profile. Nhưng tính năng mạnh hơn là assert:

import pytest

@pytest.mark.limit_memory("100 MB")
def test_process_large_file():
    result = process_file("sample.csv")
    assert result.ok

@pytest.mark.limit_leaks("1 MB")
def test_no_leak_in_hot_loop():
    for _ in range(1000):
        handle_request()

Test fail nếu allocate quá ngưỡng hoặc leak vượt giới hạn. Mình đặt limit_leaks trên các integration test — nếu ai đó merge code mới gây leak, CI đỏ ngay.

Attach vào process đang chạy

Đây là “superpower”: bạn có thể attach memray vào một Python process đang chạy mà không cần restart:

memray attach <PID> -o out.bin --duration 60

Cực kỳ hữu ích khi app production đang leak mà bạn không muốn/không thể restart. Yêu cầu: process phải chạy cùng user, và kernel cho phép ptrace (có thể cần sudo sysctl kernel.yama.ptrace_scope=0).

Bài học rút ra

Đừng đoán, đo. tracemalloc ổn cho dev nhưng không thấy C extension — production bug 90% thời gian nằm ở đó.
flamegraph --leaks trước, rồi mới flamegraph tổng. Leak và peak là hai vấn đề khác nhau.
lru_cache trên method là cái bẫy nếu instance ngắn hạn. Luôn check maxsize.
Tích hợp pytest-memray từ sớm — chặn leak ở CI rẻ hơn debug lúc 2h sáng rất nhiều.
Live mode rất mạnh để reproduce bug. Gọi endpoint, nhìn memory nhảy — bug hiện ra trước mắt.
Attach vào PID khi không thể restart production. Ghi 60s đủ để thấy pattern.

memray không phải viên đạn bạc, nhưng nó lấp vào đúng khoảng trống mà tracemalloc để lại: thấy native memory, chạy được trên workload thực, và output đẹp đủ để sếp nhìn cũng hiểu.

Bạn đã từng gặp memory leak nào khó debug chưa? Dùng tool gì để tìm ra? Kể mình nghe trong phần comment nhé.