Serving ML Models in Production with FastAPI and Celery

Overview

Bạn đang không biết triển khai mô hình Machine Learning của bạn cho dự án thực tế bạn đang làm thì đọc bài này là đúng rồi đó (đùa thui). Hôm nay mình xin giới thiệu một ví dụ để deploy Machine Learning model bằng cách sử dụng Celery và FastAPI. Tất cả mã có thể được tìm thấy trong kho lưu trữ ở đây. Github

Trong bài này mình sử dụng một model object-detection được train bằng Tensorflow dựa trên bộ dữ liệu Coco Dataset. Nói chung là là nó có khoảng 80 class như chó, mèo, tờ rim, gà, vịt heo ... Mình không đề cập đến cách train model trong bài viết này nhé.

Asynchronous

Mục tiêu mình sử dụng Celery là vì Không đồng bộ: thay vì trực tiếp trả về một dự đoán, dịch vụ mô hình sẽ trả về một định danh duy nhất cho một nhiệm vụ. Trong khi nhiệm vụ dự đoán đang được hoàn thành bởi dịch vụ mô hình, khách hàng có thể tự do tiếp tục xử lý khác. Các bước dưới đây mô tả các hành động được thực hiện để xử lý một dự đoán:

1. Ứng dụng khách gửi yêu cầu POST tới FastAPI endpoint, với thông tin tính năng liên quan có trong nội dung yêu cầu trong trường hợp này là upload lên một hình ảnh cần dự đoán.
2. FastAPI xác thực thông tin và lưu hình vào kho lưu trử (vd: Nas, HDFS ...). Nếu xác thực thành công thì tác vụ dự đoán Celery sẽ được tạo và chuyển đến Broker - Message queue (ví dụ: RabbitMQ).
3. Đồng thời API service cũng tạo một id duy nhất và trả về client nếu tạo tác vụ Celery thành công.
4. Nhiệm vụ dự đoán được giao cho một Worker có sẵn bởi Broker. Sau khi được phân phối nhiệm vụ, worker sẽ tạo ra một dự đoán bằng cách sử dụng mô hình ML được đào tạo trước.
5. Khi một dự đoán đã được tạo, kết quả sẽ được lưu trữ bằng vào Celery backend (ví dụ: Redis).
6. Tại bất kỳ thời điểm nào sau bước 3, Client có thể bắt đầu thăm dò điểm cuối kết quả FastAPI bằng cách sử dụng id duy nhất nhận từ API service. Sau khi dự đoán sẵn sàng, nó sẽ được trả lại cho Client.

Bắt tay vào làm thoi!

Mình lại tiếp tực sử dụng Docker để phát triển ứng dụng này nhé.

Đầu tiên tiến hành sao chép repo về:

git clone https://github.com/apot-group/ml-models-in-production.git cd ml-models-in-production

Cấu trúc của dự án gồm:

├── ml-api -> Chứa code của FastAPI service
├── ml-celery -> Chứa code của Celery Worker service
├── ml-client -> Chứa code của Client
├── ml-storages -> Tạo một share storages chứa file upload và file object-detection result ├── upload ├── object-detection
├── docker-compose.dev.yaml -> run service bằng cái này!!!
...

Tiếp theo, chúng ta cần build toàn bộ service với docker-compose:

docker-compose -f docker-compose.dev.yaml build && docker-compose -f docker-compose.dev.yaml up

Vì xây dựng hình ảnh từ đầu, nên việc này có thể mất một lúc. Sau khi thực hiện lệnh docker-compose up, các dịch vụ sẽ khởi động. Có thể mất một khoảng thời gian trước khi mọi thứ bắt đầu và chạy.

Lưu ý là hình ảnh đầu tiên dự đoán sẻ mất một khoảng thời gian vì Celery cần thời gian để load ML model vào bộ nhớ. Từ hình ảnh thứ 2 trở đi dịch vụ chạy khá nhanh.