Cách theo dõi phân bổ bộ nhớ trong Python

Hướng dẫn này là phần giới thiệu về cách theo dõi phân bổ bộ nhớ trong Python bằng mô-đun tracemalloc tích hợp sẵn.

Khi viết mã bằng Python, bạn thường không cần phải bận tâm đến các chi tiết về phân bổ bộ nhớ. Tuy nhiên, việc theo dõi quá trình phân bổ bộ nhớ có thể hữu ích, đặc biệt nếu bạn đang làm việc với các hoạt động sử dụng nhiều bộ nhớ và tập dữ liệu lớn.

Mô-đun tracemalloc tích hợp của Python đi kèm với các chức năng giúp bạn hiểu cách sử dụng bộ nhớ và gỡ lỗi các ứng dụng. Với tracemalloc, bạn có thể biết vị trí và số lượng khối bộ nhớ đã được phân bổ, chụp ảnh nhanh, so sánh sự khác biệt giữa các ảnh chụp nhanh, v.v.

Chúng ta sẽ xem xét một số trong số này trong hướng dẫn này. Hãy bắt đầu.

Trước khi bạn bắt đầu

Chúng tôi sẽ sử dụng tập lệnh Python đơn giản để xử lý dữ liệu. Để làm điều này, chúng tôi sẽ tạo một tập dữ liệu mẫu và xử lý nó. Bên cạnh phiên bản Python gần đây, bạn cũng cần có gấu trúc và NumPy trong môi trường làm việc của mình.

Tạo một môi trường ảo và kích hoạt nó:

$ python3 -m venv v1
$ source v1/bin/activate

Và cài đặt các thư viện cần thiết:

$ pip3 install numpy pandas

Bạn có thể tìm thấy mã cho hướng dẫn này trên GitHub.

Tạo tập dữ liệu mẫu với chi tiết đơn hàng

Chúng tôi sẽ tạo tệp CSV mẫu có chi tiết đơn hàng. Bạn có thể chạy tập lệnh sau để tạo tệp CSV có bản ghi đơn hàng 100K:

# create_data.py
import pandas as pd
import numpy as np

# Create a sample dataset with order details
num_orders = 100000
data = {
	'OrderID': np.arange(1, num_orders + 1),
	'CustomerID': np.random.randint(1000, 5000, num_orders),
	'OrderAmount': np.random.uniform(10.0, 1000.0, num_orders).round(2),
	'OrderDate': pd.date_range(start='2023-01-01', periods=num_orders, freq='min')
}

df = pd.DataFrame(data)
df.to_csv('order_data.csv', index=False)

Tập lệnh này điền vào khung dữ liệu gấu trúc với các bản ghi 100K với bốn tính năng sau và xuất khung dữ liệu sang tệp CSV:

ID đơn hàng: Mã định danh duy nhất cho mỗi đơn hàng
CustomerID: ID của khách hàng
OrderAmount: Số tiền của mỗi đơn hàng
Ngày đặt hàng: Ngày và giờ đặt hàng

Phân bổ bộ nhớ theo dõi với tracemalloc

Bây giờ chúng ta sẽ tạo một tập lệnh Python để tải và xử lý tập dữ liệu. Chúng tôi cũng sẽ theo dõi việc phân bổ bộ nhớ.

Đầu tiên, chúng ta xác định các hàm load_data và process_data để tải và xử lý các bản ghi từ tệp CSV:

# main.py
import pandas as pd

def load_data(file_path):
    print("Loading data...")
    df = pd.read_csv(file_path)
    return df

def process_data(df):
    print("Processing data...")
    df['DiscountedAmount'] = df['OrderAmount'] * 0.9  # Apply a 10% discount
    df['OrderYear'] = pd.to_datetime(df['OrderDate']).dt.year  # Extract the order year
    return df

Sau đó, chúng ta có thể tiếp tục theo dõi quá trình phân bổ bộ nhớ bằng cách thực hiện như sau:

Khởi tạo việc theo dõi bộ nhớ bằng tracemalloc.start().
Hàm load_data() đọc tệp CSV vào khung dữ liệu. Chúng tôi chụp ảnh nhanh mức sử dụng bộ nhớ sau bước này.
Hàm process_data() thêm hai cột mới vào khung dữ liệu: 'DiscountedAmount' và 'OrderYear'. Chúng tôi chụp một ảnh chụp nhanh khác sau khi xử lý.
Chúng tôi so sánh hai ảnh chụp nhanh để tìm ra sự khác biệt về mức sử dụng bộ nhớ và in ra các dòng tiêu tốn nhiều bộ nhớ nhất.
Sau đó in mức sử dụng bộ nhớ hiện tại và cao nhất để hiểu tác động tổng thể.

Đây là mã tương ứng:

import tracemalloc

def main():
    # Start tracing memory allocations
    tracemalloc.start()

    # Load data
    df = load_data('order_data.csv')

    # Take a snapshot
    snapshot1 = tracemalloc.take_snapshot()

    # Process data
    df = process_data(df)

    # Take another snapshot
    snapshot2 = tracemalloc.take_snapshot()

    # Compare snapshots
    top_stats = snapshot2.compare_to(snapshot1, 'lineno')

    print("[ Top memory-consuming lines ]")
    for stat in top_stats[:10]:
        print(stat)

    # Current and peak memory usage
    current, peak = tracemalloc.get_traced_memory()
    print(f"Current memory usage: {current / 1024 / 1024:.1f} MB")
    print(f"Peak usage: {peak / 1024 / 1024:.1f} MB")

    tracemalloc.stop()

if __name__ == "__main__":
    main()

Bây giờ hãy chạy tập lệnh Python:

$ python3 main.py

Điều này đưa ra các dòng tiêu thụ bộ nhớ nhiều nhất cũng như mức sử dụng bộ nhớ hiện tại và cao nhất:

Loading data...
Processing data...
[ Top 3 memory-consuming lines ]
/home/balapriya/trace_malloc/v1/lib/python3.11/site-packages/pandas/core/frame.py:12683: size=1172 KiB (+1172 KiB), count=4 (+4), average=293 KiB
/home/balapriya/trace_malloc/v1/lib/python3.11/site-packages/pandas/core/arrays/datetimelike.py:2354: size=781 KiB (+781 KiB), count=3 (+3), average=260 KiB
:123: size=34.6 KiB (+15.3 KiB), count=399 (+180), average=89 B
Current memory usage: 10.8 MB
Peak usage: 13.6 MB

Kết thúc

Việc sử dụng tracemalloc để theo dõi quá trình phân bổ bộ nhớ giúp xác định các hoạt động sử dụng nhiều bộ nhớ và có khả năng tối ưu hóa hiệu suất bằng cách sử dụng dấu vết bộ nhớ và số liệu thống kê được trả về.

Bạn sẽ có thể xem liệu bạn có thể sử dụng các cấu trúc dữ liệu và phương pháp xử lý hiệu quả hơn để giảm thiểu việc sử dụng bộ nhớ hay không. Đối với các ứng dụng chạy lâu, bạn có thể sử dụng tracemalloc định kỳ để theo dõi việc sử dụng bộ nhớ. Điều đó có nghĩa là bạn luôn có thể sử dụng tracemalloc kết hợp với các công cụ lập hồ sơ khác để có cái nhìn toàn diện về việc sử dụng bộ nhớ.

Nếu bạn quan tâm đến việc tìm hiểu cách lập hồ sơ bộ nhớ bằng trình phân tích bộ nhớ, hãy đọc Giới thiệu về Lập hồ sơ bộ nhớ trong Python.

Bala Priya C là nhà phát triển và nhà văn kỹ thuật đến từ Ấn Độ. Cô ấy thích làm việc ở lĩnh vực giao thoa giữa toán học, lập trình, khoa học dữ liệu và sáng tạo nội dung. Các lĩnh vực chuyên môn và quan tâm của cô bao gồm DevOps, khoa học dữ liệu và xử lý ngôn ngữ tự nhiên. Cô ấy thích đọc, viết, viết mã và cà phê! Hiện tại, cô ấy đang nỗ lực học hỏi và chia sẻ kiến thức của mình với cộng đồng nhà phát triển bằng cách viết các hướng dẫn, hướng dẫn cách thực hiện, các ý kiến, v.v. Bala cũng tạo ra các hướng dẫn viết mã và tổng quan tài nguyên hấp dẫn.