Khắc phục lỗi Python Multiprocessing với Class tùy chỉnh: Tại sao Pickle lại gây crash?

Mới đây, tôi có một kịch bản Python cần xử lý một lượng lớn hình ảnh. Xử lý ảnh đúng là một trường hợp sử dụng điển hình cho đa tiến trình (multiprocessing). Tôi dự định chia nhỏ khối lượng công việc này cho 8 nhân của CPU để tăng tốc độ xử lý.

Lần thử đầu tiên trông có vẻ rất chuẩn:

from multiprocessing import Pool

def process_image(image_path):
    # thực hiện xử lý
    return processed

with Pool(8) as pool:
    results = pool.map(process_image, image_paths)

Mọi thứ có vẻ ổn, nhưng khi chạy thì chương trình lập tức bị sập. Một lỗi Pickle error hiện ra:

Cannot pickle <class 'MyImageProcessor'>

Lại cái gì nữa thế này?

Vấn đề nằm ở chỗ tôi đã quên một kiến thức cơ bản. Multiprocessing trong Python sử dụng pickle để gửi dữ liệu giữa các tiến trình. Các hàm được pickle, các đối số được pickle, và cả kết quả trả về cũng được pickle. Tất cả mọi thứ đều phải tuần tự hóa được.

Và lớp xử lý ảnh của tôi thì... không thể pickle được.

Tôi đã thử vài cách:

• Đơn giản hóa lớp đó đi.
• Chỉ sử dụng hàm đơn thuần.
• Chuyển sang dùng threading (luồng).

Threading thì chạy được nhưng tốc độ chậm khủng khiếp do GIL (Global Interpreter Lock). Xử lý ảnh là tác vụ tốn nhiều tài nguyên CPU, nên threading chẳng giúp ích gì cả.

Cuối cùng, tôi cũng tìm ra giải pháp. Lớp của tôi giữ lại một số trạng thái không thể pickle. Giải pháp là chuyển những thứ không pickle được đó vào các phương thức __getstate__ và __setstate__, hoặc đơn giản hơn là khởi tạo chúng bên trong hàm worker thay vì truyền cả đối tượng vào.

Phiên bản hoạt động tốt:

from multiprocessing import Pool

def init_worker():
    # Khởi tạo các đối tượng không thể pickle tại đây
    global processor
    processor = MyImageProcessor()

def process_image(image_path):
    global processor
    return processor.process(image_path)

if __name__ == "__main__":
    with Pool(8, initializer=init_worker) as pool:
        results = pool.map(process_image, image_paths)

Việc sử dụng mẫu khởi tạo (initializer pattern) có nghĩa là mỗi tiến trình worker sẽ thiết lập bộ xử lý (processor) của riêng nó. Không cần quá trình pickle nào cả, và không còn lỗi crash.

Ngoài ra, nếu đối tượng của bạn đơn giản hơn, bạn có thể sử dụng functools.partial:

from functools import partial
from multiprocessing import Pool

def process_image(processor, image_path):
    return processor.process(image_path)

processor = MyImageProcessor()
bounded_func = partial(process_image, processor)

with Pool(8) as pool:
    results = pool.map(bounded_func, image_paths)

Tuy nhiên, cách này chỉ hoạt động nếu chính bộ xử lý (processor) đó có thể pickle được.

Vẫn chưa hiểu tại sao Multiprocessing trong Python lại mặc định dùng pickle thay vì một thứ something dễ chịu hơn. Thông báo lỗi cũng thật tệ, nó chỉ chỉ vào lớp mà không giải thích cái gì đang chặn tuần tự hóa thực sự.

Dù sao thì kịch bản cũng đã chạy được rồi. Chỉ hơi bực mình vì mất vài giờ đồng hồ để sửa cái lỗi này thôi.