Tối ưu hóa ast.walk: Cách tăng tốc độ xử lý lên 220 lần

Trong công cụ xây dựng ứng dụng AI reflex-app của chúng tôi, hệ thống tạo ra lượng lớn mã Python. Đôi khi, quá trình này gặp phải những lỗi khá sơ đẳng như tham số vị trí đặt sau tham số từ khóa, sử dụng cú pháp cũ của các phiên bản framework trước, v.v.

Việc chạy lệnh reflex compile cuối cùng sẽ tìm ra tất cả các lỗi này, nhưng nó chỉ phát hiện một vấn đề mỗi lần. Điều này có nghĩa là nếu AI mắc nhiều lỗi sai, độ trễ sẽ tăng lên rất lớn đối với những bản sửa lỗi tương đối đơn giản. Do đó, chúng tôi quyết định sử dụng một trình linter (linter) là cách tiếp cận tốt nhất để khắc phục điều này. Vì cần thêm các quy tắc riêng cho Reflex, chúng tôi không thể sử dụng công cụ có sẵn mà phải tự xây dựng một cái.

Tối ưu hóa hiệu năng Python

Trình linter ban đầu của chúng tôi trông có vẻ đơn giản, nhưng nhanh chóng gặp phải vấn đề về hiệu năng vì chúng tôi đang xử lý lượng mã được tạo ra rất lớn. Đáng chú ý, phần chậm nhất trong đoạn mã trên không phải là các kiểm tra isinstance, mà là ast.walk.

Tất nhiên, có nhiều cách để tối ưu hóa mà không cần làm cho ast.walk nhanh hơn, và chúng tôi đã triển khai những thay đổi "dễ thực hiện" đó trước tiên. Tuy nhiên, chúng tôi nhanh chóng nhận ra rằng rất khó để làm cho trình linter nhanh hơn đáng kể nếu không giải quyết thách thức làm cho chính hàm walk nhanh hơn.

Tại sao ast.walk lại chậm?

Việc duyệt qua một cây trừu tượng (AST) không nhất thiết phải chậm. Việc duyệt qua module difflib mất khoảng 2ms trên thiết bị của tôi, với khoảng 7.000 nút. Bản thân nó không quá tệ, nhưng nó cộng dồn rất nhanh. Một tính toán sơ bộ cho chúng ta khoảng 285 nanosecond cho mỗi nút, tương đương hàng nghìn chu kỳ CPU - nhiều hơn nhiều so với những gì một thao tác duyệt đơn giản cần.

Điểm đầu tiên cần lưu ý ở đây là việc sử dụng yield. Generators và cú pháp yield là một tính năng mạnh mẽ của Python, nhưng chúng đi kèm với cái giá phải trả: việc treo thực thi của vòng lặp và khôi phục nó liên tục trong một đường dẫn nóng (hot path) nơi chúng ta đang tiêu thụ toàn bộ danh sách anyway. Chắc chắn, nó tiết kiệm bộ nhớ, nhưng đó không phải là vấn đề của chúng tôi lúc này.

Bằng cách lưu trữ một danh sách và liên tục thêm vào nó, chúng tôi có thể giảm thiểu điều này. Tuy nhiên, sau khi chạy, tôi nhận ra nó chỉ mang lại cải thiện 5%. Ít hơn nhiều so với tôi mong đợi.

Tối ưu hóa sâu hơn với Python

Vấn đề tiếp theo nằm ở iter_child_nodes. À, một generator khác. Nếu chúng ta nội tuyến hóa (inline) nó, chúng ta sẽ thấy một số cải thiện hiệu suất khá tốt. Chúng ta đạt được mức cải thiện khoảng 25%.

Điều đặt ra câu hỏi là: 75% còn lại nằm ở đâu? Và trong khi đang bàn về chủ đề này, iter_fields là gì? Đây lại là một generator khác. Chúng ta cũng đang trả về một tuple mà chúng ta không sử dụng, vì chúng ta chỉ quan tâm đến giá trị, không phải tên.

Hàm getattr(node, field, None) sẽ nhanh hơn, vì việc xử lý ngoại lệ có thể nhanh hơn khi được ủy quyền cho CPython. Kết hợp cả hai thay đổi này mang lại cho chúng ta khoảng 50% cải thiện tích lũy.

Gấp đôi tốc độ cũng không tệ, nhưng việc đi bộ với tốc độ gấp 2 có được gọi là chạy nước rút (sprinting) không? Tôi không nghĩ vậy. Hãy đẩy điều này xa hơn nữa.

Chuyển sang Rust và tối ưu hóa bộ nhớ

Tại thời điểm này, tôi đã chạm đến giới hạn của những gì Python có thể làm. Python có một thủ thuật cuối cùng trong tay áo: bindings. Nó cho phép chúng ta viết logic này trong mã máy gốc. Tôi chọn sử dụng Rust thay vì C vì quen thuộc hơn.

Việc chuyển logic sang Rust mang lại khoảng 78% cải thiện tích lũy. Điều này cũng cho phép chúng ta làm một điều thú vị hơn. Vì chúng ta đang thực hiện rất nhiều getattr, biên dịch xuống việc đọc từ điển, chúng ta có thể chỉ cần lặp qua chính từ điển đó (__dict__).

Chúng ta cũng có thể cải thiện việc kiểm tra lớp con. Chỉ có 132 lớp con của ast.AST, vì vậy thay vì gọi isinstance thực sự, chúng ta có thể lưu địa chỉ bộ nhớ của tất cả các lớp đó trong một tập hợp và chỉ cần kiểm tra tư cách thành viên.

Kết hợp cả hai điều này mang lại mức cải thiện ~93%. Tổng cộng, nhanh hơn khoảng 14 lần.

Đẩy đến giới hạn tuyệt đối

Cuối cùng, chúng tôi viết lại hàm PyDict_Next của CPython để loại bỏ các kiểm tra và đếm tham chiếu không cần thiết. Một quan sát khác là tất cả các lớp con ast.AST đều nằm trong một chuỗi tương đối ngắn.

Chúng tôi đạt được khoảng 99% cải thiện. Đó là hai cấp độ độ lớn! Chúng tôi đang đến gần giới hạn rồi.

Quan sát cuối cùng là __dict__ của các lớp con AST rất dễ dự đoán; các giá trị đầu tiên là _fields rồi đến _attributes. Vì vậy, nếu chúng ta chỉ cần quét các mục đầu tiên len(_fields) (đã được tính toán trước), chúng ta tiết kiệm thời gian kiểm tra lineno/col_offset và có thể bỏ qua các tham chiếu ngược .parent do người dùng đính kèm.

Tất cả những điều kết hợp lại, chúng ta đạt được ~99,5%, tương đương với mức cải thiện khoảng 220 lần.

Bạn có thể tìm hiểu thêm về mã nguồn tại https://github.com/reflex-dev/fast-walk/.