资源管理新视角：利用 FastAPI Lifespan 事件优化你的应用II

news2024/9/21 14:24:47

本文说明在 FastAPI 应用程序中使用 lifespan 事件来管理资源的加载和卸载。lifespan 事件允许你在应用启动时执行一些初始化代码，并在应用关闭时执行一些清理代码。这是通过使用异步上下文管理器实现的，具体来说，是通过 asynccontextmanager 装饰器创建一个异步生成器函数。
在这里插入图片描述

# 导入asynccontextmanager用于创建异步上下文管理器
from contextlib import asynccontextmanager
# 导入FastAPI，用于创建和管理Web应用
from fastapi import FastAPI

# 定义一个模拟的机器学习模型函数，它接受一个浮点数x并返回x乘以42的结果
def fake_answer_to_everything_ml_model(x: float):
    return x * 42

# 创建一个空字典，用于存储机器学习模型
ml_models = {}

# 使用asynccontextmanager装饰器定义一个异步上下文管理器函数lifespan
@asynccontextmanager
async def lifespan(app: FastAPI):
    # 在异步上下文管理器中，"进入上下文"时加载机器学习模型
    ml_models["answer_to_everything"] = fake_answer_to_everything_ml_model
    # yield
    yield
    # 在异步上下文管理器中，"退出上下文"时清理机器学习模型，释放资源
    ml_models.clear()

# 创建FastAPI应用实例，并将lifespan函数作为其生命周期管理器
app = FastAPI(lifespan=lifespan)

# 定义一个GET请求路由，用于处理预测请求
@app.get("/predict")
async def predict(x: float):
    # 从ml_models字典中获取模型，并使用它对输入值x进行预测
    result = ml_models["answer_to_everything"](x)
    # 返回预测结果
    return {"result": result}

在这段代码中，lifespan 函数作为异步上下文管理器，负责在FastAPI应用启动时加载机器学习模型，并在应用关闭时清理模型。这样的设计模式确保了应用在运行时能够有效地管理资源。predict 路由函数提供了一个接口，允许用户通过HTTP请求发送数据，并获取模型的预测结果。

代码解释：

导入必要的库：asynccontextmanager 用于创建异步上下文管理器，FastAPI 是 FastAPI 框架的核心类。
定义一个模拟的机器学习模型函数 fake_answer_to_everything_ml_model，它接受一个浮点数 x 并返回 x * 42。
创建一个全局字典 ml_models，用于存储机器学习模型。
使用 @asynccontextmanager 装饰器定义一个异步上下文管理器 lifespan，它接受一个 FastAPI 实例作为参数。
在 lifespan 函数中，使用 yield 语句之前的代码块来加载模型。在这个例子中，模型加载操作被模拟为将模型函数添加到 ml_models 字典中。
yield 语句之后的代码块用于在应用关闭时执行清理操作，这里它清除了 ml_models 字典。
创建一个 FastAPI 实例，并将其 lifespan 参数设置为 lifespan 函数。
定义一个路由处理函数 predict，它使用加载的模型来处理预测请求，并返回结果。

实际作用：

在应用启动时，lifespan 函数中的代码在 yield 之前被执行，模拟加载模型的操作。
应用开始接收请求后，predict 路由可以被调用，使用已加载的模型来处理请求。
当应用关闭时，lifespan 函数中的代码在 yield 之后被执行，进行资源的清理。

在 Python 的异步上下文管理器中，yield 关键字扮演着非常重要的角色。它用于定义上下文管理器的执行流程，特别是在异步上下文管理器（使用 asynccontextmanager 装饰的生成器函数）中。

理解 `yield` 在异步上下文管理器中的作用：

定义执行点：在异步上下文管理器中，yield 语句定义了异步生成器在执行时的暂停和恢复点。当进入上下文时，代码执行到 yield 语句，生成器会暂停，等待外部代码执行。一旦退出上下文，生成器会继续执行 yield 之后的代码。
上下文的分割：yield 将上下文管理器的代码分为两部分：
- 进入上下文：yield 之前的代码，用于设置必要的资源或状态。在异步上下文管理器中，这通常包括资源的初始化和分配。
- 退出上下文：yield 之后的代码，用于清理资源或恢复状态。在异步上下文管理器中，这通常包括资源的释放和清理。

示例解释：

在你提供的代码中：

@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()

进入上下文：在 yield 之前的部分，ml_models 字典被用来存储一个模拟的机器学习模型。这模拟了在应用启动时加载模型的过程。
退出上下文：yield 之后的代码在上下文管理器退出时执行，这里它清除 ml_models 字典，模拟了在应用关闭时卸载模型和释放资源的过程。

异步上下文管理器的使用：

当使用 async with 语句块时，进入点是 async with 后面的代码块开始执行的地方，而退出点是该代码块执行完毕的地方。例如：

async with lifespan(app) as manager:
    # 这里可以访问和使用上下文中的资源
    pass

在这个 async with 块中，进入上下文的代码在 lifespan 函数被调用时执行，直到遇到 yield。yield 之后的代码在 async with 块的代码执行完毕后执行。yield 在异步上下文管理器中是一个关键的执行控制点，它定义了资源的生命周期，确保资源在使用前后都能得到妥善的处理。

简单的两个 Demo：

当然，以下是上述两个示例代码：

Demo 1：数据库连接

# 导入必要的库
from contextlib import asynccontextmanager
from fastapi import FastAPI
import asyncpg

# 定义数据库连接的URL
DATABASE_URL = "postgres://user:password@localhost/dbname"

# 使用asynccontextmanager装饰器定义一个异步上下文管理器
@asynccontextmanager
async def lifespan(app: FastAPI):
    # 异步连接数据库
    connection = await asyncpg.connect(DATABASE_URL)
    # 将数据库连接存储在app.state中，以便在应用中其他地方使用
    app.state.db = connection
    yield
    # 在上下文结束时关闭数据库连接
    await app.state.db.close()

# 创建FastAPI应用实例，并传入lifespan函数
app = FastAPI(lifespan=lifespan)

# 定义一个路由，用于获取数据库中的数据
@app.get("/get_data")
async def get_data():
    # 从app.state中获取数据库连接
    db = app.state.db
    # 执行SQL查询并返回结果
    return {"data": await db.fetch("SELECT * FROM my_table")}

Demo 2：启动时加载配置文件

# 导入必要的库
from contextlib import asynccontextmanager
from fastapi import FastAPI

# 创建一个空字典，用于存储配置信息
config = {}

# 使用asynccontextmanager装饰器定义一个异步上下文管理器
@asynccontextmanager
async def lifespan(app: FastAPI):
    # 模拟加载配置文件，这里只是简单地将一个配置项设置为一个值
    config['setting'] = 'some_value'
    # 将配置信息存储在app.state中，以便在应用中其他地方使用
    app.state.config = config
    yield
    # 在上下文结束时清理配置信息（如果有必要的话）
    config.clear()

# 创建FastAPI应用实例，并传入lifespan函数
app = FastAPI(lifespan=lifespan)

# 定义一个路由，用于获取当前的配置信息
@app.get("/get_config")
async def get_config():
    # 从app.state中返回配置信息
    return app.state.config