在本機上用十分鐘建立 Server 來跑模型

Abby Yeh
Abby Yeh
Nov 5 · 8 min read

Build an AI App — 1 : use Flask and Celery to build a simple model

Photo by Jordan Harrison on Unsplash

用 Flask 可以在短時間建立一個簡單的 server 來接 HTTP 的 request,但是Flask 是用來開發小型專案的輕量框架,並不支援異步的系統。而且,有些model 連 predict 都要跑一段時間,connection 一直連著直到結果出來實在不是一個好方法。這時就要加入 Celery 啦~ Celery 可以利用 in memory 的 DB 來把工作交給叫 worker 的小程式執行,並用 DB 和 server 溝通執行狀態和結果,來達到 Async 的 API。

— 環境架設 —

首先從官網下載並安裝 python3.6,用 pip 安裝 `virtualenv` 指定3.6的版本並啟動他 (用 3.7 也可以,但是會比較麻煩一點):

$ pip install virtualenv
$ virtualenv --python=python3.6 venv
$ . venv/bin/activate

安裝需要用的 packages:

(venv)$ pip install flask
(venv)$ pip install request
(venv)$ pip install celery
(venv)$ pip install numpy
(venv)$ pip install pillow # 這個 model 需要處理圖片

但是如果你用的是 python3.7 在裝 Celery 你必須用:

(venv)$ pip install -e git+https://github.com/celery/celery.git@master#egg=celery

最後還要記得安裝 Reddis 或 RabbitMQ 給 Celery 用。

— 建立 Flask Server—

建立 server 和基本的 API 請先看前一篇

— 新增 API —

建立 server 號,就可以加入 API 了。把前一篇 app/main.py 裡的 sync 的 API拆成三個,讓每一個API都可以馬上回傳。這三個 API 分別是叫 worker 執行、檢查 worker 的狀態、和取回 worker 的結果。先把所有在 celery_queue/worker.py 的 task import 進來 (之後會再完成 Celery worker 的部分):

from celery_queue import workers
  1. 第一個是叫 worker 執行的 api,前半部分和上一篇的 sync_transform 是一樣的,用 @app.route(URL, methods) 建立一個 api,然後得到 PIL Image 物件:
@app.route('/transform/<style>', methods=['POST'])
def transform(style):
if 'file' not in request.files:
return 'No file part', 500
file = request.files['file']
if file.filename == '':
return 'No selected file', 500
image = Image.open(file.stream).convert("RGB")

不一樣的的地方是不是所有物件都可以傳給 worker,我們要把圖片轉成字串或數列,在這裡是用 Numpy 把圖片轉成整數數列,並呼叫 task 的 apply_async 函式把工作和參數交給 worker 。

image = np.array(image).tolist()
task = workers.transform.apply_async(args=[image, style])

最後是回傳 task id ,讓使用者可以用這個 id 查詢進度。

return task.id

完整的 function 如下:

2. 接著是最簡單的檢查 worker status 的 api,一樣用角括號在網址中插入變數 task_id,但是這次要用 GET 因為不需要夾帶其他資訊。 用 celery 的 AsyncResult 函式並給他 task id 可以得到 task 的狀態物件,用 fail、ready 函式就可以知道 task 有沒有正確的執行完:

3. 最後用 check_status 確認 model 執行完了,還需要一個 GET 的 API 來下載結果。

@app.route('/transform/download/<task_id>', methods=['GET'])
def download_transformed_image(task_id):

還是先檢查一下 task 的狀況,來避免使用者沒有先確認 worker 的狀態直接下載:

res = workers.celery.AsyncResult(task_id)
if not res.ready():
return 'task hasn\'t been finished', 500

用 get 函式就可以得到 task 的執行結果 。如果結果是字串或數字,就可以直接回傳,但是這個 API 要回傳影像檔回去,所以先把結果存起來接:

result = res.get()

和傳資料給 worker 相同,一樣要用字串或數列,回傳的結果和傳給 worker 相同都是整數數列。所以要先用 numpy、pillow 這兩個 package 將結果轉成圖檔:

result = Image.fromarray(np.uint8(result))

接下來就和上一篇相同,轉成 file bytes 然後用 Flask 的 send_file 傳回 client端:

image_bytes = io.BytesIO()
result.save(image_bytes, format='png')
image_bytes.seek(0)
return send_file(image_bytes, attachment_filename='result.png', mimetype='image/png')

完整的 download API:

— 建立 Celery Worker —

接著要把 worker 要執行的 task 完成,在 app 資料夾裡建一個資料夾 celery_queue ,並建立一個檔案 app/celery_queue/worker.py:

from flask import Flask
from celery import Celery
# 建立一個 Celery 來分配 workers
# 並從 app/celery_queue/config.py 讀入設定
celery = Celery('tasks') celery.config_from_object('celery_queue.config')

再建立 Celery 的設定檔 app/celery_queue/config.py

BROKER_URL = 'redis://localhost:6379' # 分配 worker 所使用的 db
CELERY_RESULT_BACKEND = 'redis://localhost:6379' # 暫存結果所使用的 db

這樣就有在背景執行的 worker 了。

— 建立一個預測用的 task —

我們要呼叫的 transform 函數和上一篇相同:

def transform(image, style):

...

return output_image

前面在建 Celery物件的時候變數名是 celery ,所以用 @celery.task 這個 wrapper 就可以經鬆的建立一個 task 讓 worker 執行 transform 函數,就完成了:

@celery.task
def transform(image_list, style):
return transform(image_list, style)

— 執行 —

分別在不同分頁執行Celery worker 和 Flask server 就可以接收 requests 了。

執行 worker:

(venv)$ cd app
(venv)$ celery worker -A app.celery_queue.workers

執行 server:

(venv)$ export PYTHONPATH=app/
(venv)$ export FLASK_APP=app/main.py
(venv)$ export FLASK_DEBUG=1 # debug用,改code會重新啟動 server
(venv)$ python -m flask run

— Request 測試 —

最後用 request 套件來測試 Server 是否有成功。建立一個 test.py:

執行完後,result.png 可以正常顯示就代表成功了!接下來我們就可以用docker 包起來,deploy 成正式的 server。

— 參考資料—

[1] Flask Document Quickstart

[2] Install both python 3.6 and 3.7

[3] Celery Broken with Latest Python

[4] First Steps with Celery

Taiwan AI Academy

news, tech reviews and supplemental materials

Abby Yeh

Written by

Abby Yeh

Teaching assistants in Taiwan AI Academy, sharing what I learn when I assist students.

Taiwan AI Academy

news, tech reviews and supplemental materials

Welcome to a place where words matter. On Medium, smart voices and original ideas take center stage - with no ads in sight. Watch
Follow all the topics you care about, and we’ll deliver the best stories for you to your homepage and inbox. Explore
Get unlimited access to the best stories on Medium — and support writers while you’re at it. Just $5/month. Upgrade