今天要介紹的是如何在 FastAPI 的框架啟動或結束時可以做一些前置作業的用法。
舉個例子來說, FastAPI 常常被拿來作為AI模型部署時的一個解決方案,那我們總不可能在每次呼叫推論端點時都需要去載入模型,而是需要在 APP 啟動時就先載入好了,那當我們今天應用程式關閉時,我們也需要將模型所占用的記憶體空間給釋放出來。
要做到上述的事情 FastAPI 的官方文件有詳細介紹我們如何在生命週期的事件中去插入我們需要執行的動作。
官方給出的方法是 lifespan function,我們可以透過定義這個函式與 context manager 做結合,來達到說啟動結束可以執行指定作業。
with open('test.txt', 'w', encoding='utf-8') as f:
print(f.read())
在 python 中我們可以使用 with 來打開一個資源,而在 with 當中的邏輯執行結束之後,打開的資源將會自動關閉,這種用法便是 context manager 的運用。
from contextlib import contextmanager
@contextmanager
def open_file(name, mode):
f = open(name, mode)
yield f
f.close()
而我們可以透過 自訂class 或函式庫的方式來自製一個 context manager,這裡我們使用函式庫的裝飾器來自訂一個函式,並在其中將要回傳的資源用 yield 回傳回去。
with open_file('test.txt', 'w') as f:
print(f.read())
而在使用時就如以上程式碼,使用 with 語法來使用這個 context manager,在執行完邏輯之後,便會去執行 yield 後面的程式,以達到釋放資源等後續處理。
from contextlib import asynccontextmanager
from fastapi import FastAPI
def fake_answer_to_everything_ml_model(x: float):
return x * 42
ml_models = {}
@asynccontextmanager
async def lifespan(app: FastAPI):
# Load the ML model
ml_models["answer_to_everything"] = fake_answer_to_everything_ml_model
yield
# Clean up the ML models and release the resources
ml_models.clear()
app = FastAPI(lifespan=lifespan)
@app.get("/predict")
async def predict(x: float):
result = ml_models["answer_to_everything"](x)
return {"result": result}
在官方的案例中使用 asynccontextmanager 來定義一個 lifespan 的函式,且參數需要帶入一個叫 app 的 FastAPI 實體。
裡面做的事就跟正常 context manager 相同,定義好開始邏輯, yield 返回,定義結束邏輯。
而在實體化 FastAPI 實體時,app = FastAPI(lifespan=lifespan)
,將該方法傳入 lifespan 參數中。
yield 前的程式碼將會在 APP 開始回應前執行,yield 後的程式碼將在 APP 關閉時執行。
如果想知道如何在其他生命週期中插入程式的話,可以參考:
我們在設計一個應用程式時,需要做一個全面的考量,在哪些生命週期該做哪些事情這些都需要去思考,FastAPI 對於這些事情基本上支援的很全面,如果官方文件沒有提到,那可以去看他的底層 Starlette 這個框架的官方文件,通常可以獲得你想要的答案。
Python 的 with 語法使用教學:Context Manager 資源管理器 - G. T. Wang (gtwang.org)
Lifespan Events - FastAPI (tiangolo.com)