[2023] Dùng PyO3 để làm Python nhanh hơn 100 lần

Gần đây khi tìm hiểu về Python free-threading, tôi bắt đầu quan tâm đến PyO3 nên dù đây là bài viết từ 2 năm trước vẫn muốn chia sẻ lại.

Making Python 100× Faster with <100 Lines of Rust – tóm tắt

Bối cảnh

• Thư viện Python cốt lõi trong pipeline xử lý 3-D nội bộ trở thành nút thắt cổ chai khi số người dùng đồng thời tăng lên.
• Việc viết lại toàn bộ bằng Rust quá rủi ro và tốn thời gian, vì vậy đã chọn tối ưu hóa từng phần.

Cách tiếp cận

1. Đo đạc trước: dùng profiler lấy mẫu py-spy để xác định nút thắt cổ chai.
2. Đưa Rust vào dần dần
   • Kết nối Python ↔ Rust bằng PyO3 + maturin.
   • Trước tiên chỉ chuyển hàm find_close_polygons sang Rust.
   • Sau đó chuyển cả cấu trúc dữ liệu Polygon sang Rust và cho phép subclass từ Python.
3. Lặp lại chu trình profiling-cải tiến
   • Giảm thiểu các chuyển đổi NumPy → Rust không cần thiết.
   • Giảm cấp phát và sao chép, đồng thời vi tối ưu bằng cách tính khoảng cách trực tiếp.

Thay đổi hiệu năng

Giai đoạn	Thời gian chạy trung bình (ms)	Mức cải thiện
Ban đầu, Python thuần	293.41	1×
v1 – chỉ hàm sang Rust (--release)	23.44	12.5×
v2 – cả Polygon cũng sang Rust	6.29	46.5×
v3 – loại bỏ cấp phát, tính trực tiếp	2.90	101×

Công nghệ cốt lõi

• PyO3 : FFI Python ↔ Rust an toàn.
• maturin : tự động hóa build và phân phối.
• ndarray / numpy crate : mảng và đại số tuyến tính phía Rust.
• py-spy : profiler có thể nhìn tới cả native stack.

Bài học rút ra

• Nếu profiling trước, chỉ với thay đổi mã nhỏ cũng có thể thu được lợi ích lớn.
• Ngay cả khi giữ nguyên Python API, chỉ thay module Rust cũng có thể áp dụng ngay vào dịch vụ thực tế.
• Rust vẫn rất hiệu quả ngay cả khi chỉ được đưa vào mỏng ở “vùng hiệu năng”.

[allinux]

Việc tạo phần mở rộng Python bằng C/C++ làm năng suất giảm quá nhiều, còn PyO3 thì trước hết đã có maturin, cargo nên rất tiện.
Ngoài ra, vì module Python cũng bắt buộc phải hỗ trợ cross-compile, mà Rust thì việc cross-compile cũng rất thuận tiện.

[lamanus]

maturin... đau khổ...

[aer0700]

Cố gắng cầm cự tối đa bằng vector hóa NumPy, nếu không được thì gắn GPU rồi chuyển sang CuPy hoặc torch, mà vẫn không ổn thì viết native bằng Cython các kiểu... nhưng có vẻ tốt nhất là nếu không thật sự cần thì đừng làm native. Cực lắm.