FastAPI-数据科学应用构建指南-四-

57 阅读1小时+

FastAPI 数据科学应用构建指南(四)

原文:annas-archive.org/md5/a1f4ad3f5a4649378151351d58ad6e73

译者:飞龙

协议:CC BY-NC-SA 4.0

第十二章:使用 FastAPI 创建高效的预测 API 端点

在上一章中,我们介绍了 Python 社区中广泛使用的最常见的数据科学技术和库。得益于这些工具,我们现在可以构建能够高效预测和分类数据的机器学习模型。当然,我们现在需要考虑一个便捷的接口,以便能够充分利用它们的智能。这样,微服务或前端应用程序就可以请求我们的模型进行预测,从而改善用户体验或业务运营。在本章中,我们将学习如何使用 FastAPI 实现这一点。

正如我们在本书中看到的,FastAPI 允许我们使用清晰而轻量的语法实现非常高效的 REST API。在本章中,你将学习如何以最有效的方式使用它们,以便处理成千上万的预测请求。为了帮助我们完成这项任务,我们将引入另一个库——Joblib,它提供了帮助我们序列化已训练模型和缓存预测结果的工具。

本章我们将涵盖以下主要内容:

  • 使用 Joblib 持久化已训练的模型

  • 实现高效的预测端点

  • 使用 Joblib 缓存结果

技术要求

本章需要你设置一个 Python 虚拟环境,就像我们在第一章中设置的那样,Python 开发环境 设置

你可以在专门的 GitHub 仓库中找到本章的所有代码示例,地址为github.com/PacktPublishing/Building-Data-Science-Applications-with-FastAPI-Second-Edition/tree/main/chapter12

使用 Joblib 持久化已训练的模型

在上一章中,你学习了如何使用 scikit-learn 训练一个估计器。当构建这样的模型时,你可能会获得一个相当复杂的 Python 脚本来加载训练数据、进行预处理,并用最佳的参数集来训练模型。然而,在将模型部署到像 FastAPI 这样的 Web 应用程序时,你不希望在服务器启动时重复执行这些脚本和所有操作。相反,你需要一个现成的已训练模型表示,只需要加载并使用它即可。

这就是 Joblib 的作用。这个库旨在提供高效保存 Python 对象到磁盘的工具,例如大型数据数组或函数结果:这个操作通常被称为持久化。Joblib 已经是 scikit-learn 的一个依赖,因此我们甚至不需要安装它。实际上,scikit-learn 本身也在内部使用它来加载打包的示例数据集。

如我们所见,使用 Joblib 持久化已训练的模型只需要一行代码。

持久化已训练的模型

在这个示例中,我们使用了我们在第十一章中看到的 newsgroups 示例,Python 中的数据科学入门一节的链式预处理器和估算器部分。为了提醒一下,我们加载了newsgroups数据集中 20 个类别中的 4 个,并构建了一个模型,将新闻文章自动分类到这些类别中。一旦完成,我们将模型导出到一个名为newsgroups_model.joblib的文件中:

chapter12_dump_joblib.py


# Make the pipelinemodel = make_pipeline(
           TfidfVectorizer(),
           MultinomialNB(),
)
# Train the model
model.fit(newsgroups_training.data, newsgroups_training.target)
# Serialize the model and the target names
model_file = "newsgroups_model.joblib"
model_targets_tuple = (model, newsgroups_training.target_names)
joblib.dump(model_targets_tuple, model_file)

github.com/PacktPublishing/Building-Data-Science-Applications-with-FastAPI-Second-Edition/tree/main/chapter12/chapter12_dump_joblib.py

如你所见,Joblib 提供了一个名为dump的函数,它仅需要两个参数:要保存的 Python 对象和文件路径。

请注意,我们并没有单独导出model变量:相反,我们将它和类别名称target_names一起封装在一个元组中。这使得我们能够在预测完成后检索类别的实际名称,而不必重新加载训练数据集。

如果你运行这个脚本,你会看到newsgroups_model.joblib文件已被创建:


(venv) $ python chapter12/chapter12_dump_joblib.py$ ls -lh *.joblib
-rw-r--r--    1 fvoron    staff       3,0M 10 jan 08:27 newsgroups_model.joblib

注意,这个文件相当大:它超过了 3 MB!它存储了每个词在每个类别中的概率,这些概率是通过多项式朴素贝叶斯模型计算得到的。

这就是我们需要做的。这个文件现在包含了我们 Python 模型的静态表示,它将易于存储、共享和加载。现在,让我们学习如何加载它并检查我们是否可以对其进行预测。

加载导出的模型

现在我们已经有了导出的模型文件,让我们学习如何使用 Joblib 再次加载它,并检查一切是否正常工作。在下面的示例中,我们将加载位于chapter12目录下的 Joblib 导出文件,并进行预测:

chapter12_load_joblib.py


import osimport joblib
from sklearn.pipeline import Pipeline
# Load the model
model_file = os.path.join(os.path.dirname(__file__), "newsgroups_model.joblib")
loaded_model: tuple[Pipeline, list[str]] = joblib.load(model_file)
model, targets = loaded_model
# Run a prediction
p = model.predict(["computer cpu memory ram"])
print(targets[p[0]])

github.com/PacktPublishing/Building-Data-Science-Applications-with-FastAPI-Second-Edition/tree/main/chapter12/chapter12_load_joblib.py

在这里,我们只需要调用 Joblib 的load函数,并将导出文件的有效路径传递给它。这个函数的结果是我们导出的相同 Python 对象。在这里,它是一个包含 scikit-learn 估算器和类别列表的元组。

注意,我们添加了一些类型提示:虽然这不是必须的,但它帮助 mypy 或你使用的任何 IDE 识别加载对象的类型,并受益于类型检查和自动补全功能。

最后,我们对模型进行了预测:它是一个真正的 scikit-learn 估算器,包含所有必要的训练参数。

就这样!如你所见,Joblib 的使用非常直接。尽管如此,它是一个重要工具,用于导出你的 scikit-learn 模型,并能够在外部服务中使用这些模型,而无需重复训练阶段。现在,我们可以在 FastAPI 项目中使用这些已保存的文件。

实现一个高效的预测端点

现在我们已经有了保存和加载机器学习模型的方法,是时候在 FastAPI 项目中使用它们了。正如你所看到的,如果你跟随本书的内容进行操作,实施过程应该不会太令你惊讶。实现的主要部分是类依赖,它将处理加载模型并进行预测。如果你需要复习类依赖的内容,可以查看第五章FastAPI 中的依赖注入

开始吧!我们的示例将基于上一节中提到的newgroups模型。我们将从展示如何实现类依赖开始,这将处理加载模型并进行预测:

chapter12_prediction_endpoint.py


class PredictionInput(BaseModel):           text: str
class PredictionOutput(BaseModel):
           category: str
class NewsgroupsModel:
           model: Pipeline | None = None
           targets: list[str] | None = None
           def load_model(self) -> None:
                         """Loads the model"""
                         model_file = os.path.join(os.path.dirname(__file__), "newsgroups_model.joblib")
                         loaded_model: tuple[Pipeline, list[str]] = joblib.load(model_file)
                         model, targets = loaded_model
                         self.model = model
                         self.targets = targets
           async def predict(self, input: PredictionInput) -> PredictionOutput:
                         """Runs a prediction"""
                         if not self.model or not self.targets:
                                       raise RuntimeError("Model is not loaded")
                         prediction = self.model.predict([input.text])
                         category = self.targets[prediction[0]]
                         return PredictionOutput(category=category)

github.com/PacktPublishing/Building-Data-Science-Applications-with-FastAPI-Second-Edition/tree/main/chapter12/chapter12_prediction_endpoint.py

首先,我们定义了两个 Pydantic 模型:PredictionInputPredictionOutput。按照纯 FastAPI 的理念,它们将帮助我们验证请求负载并返回结构化的 JSON 响应。在这里,作为输入,我们仅期望一个包含我们想要分类的文本的text属性。作为输出,我们期望一个包含预测类别的category属性。

这个代码片段中最有趣的部分是NewsgroupsModel类。它实现了两个方法:load_modelpredict

load_model方法使用 Joblib 加载模型,如我们在上一节中所见,并将模型和目标存储在类的属性中。因此,它们将可以在predict方法中使用。

另一方面,predict方法将被注入到路径操作函数中。如你所见,它直接接受PredictionInput,这个输入将由 FastAPI 注入。在这个方法中,我们进行预测,就像我们通常在 scikit-learn 中做的那样。我们返回一个PredictionOutput对象,包含我们预测的类别。

你可能已经注意到,首先我们在进行预测之前,会检查模型及其目标是否在类属性中被分配。当然,我们需要确保在进行预测之前,load_model已经被调用。你可能会想,为什么我们不把这个逻辑放在初始化函数__init__中,这样我们可以确保模型在类实例化时就加载。这种做法完全可行,但也会带来一些问题。正如我们所看到的,我们在 FastAPI 之后立即实例化了一个NewsgroupsModel实例,以便可以在路由中使用它。如果加载逻辑放在__init__中,那么每次我们从这个文件中导入一些变量(比如app实例)时,模型都会被加载,例如在单元测试中。在大多数情况下,这样会导致不必要的 I/O 操作和内存消耗。正如我们所看到的,最好使用 FastAPI 的生命周期处理器,在应用运行时加载模型。

以下摘录展示了其余的实现,并包含处理预测的 FastAPI 路由:

chapter12_prediction_endpoint.py


newgroups_model = NewsgroupsModel()@contextlib.asynccontextmanager
async def lifespan(app: FastAPI):
           newgroups_model.load_model()
           yield
app = FastAPI(lifespan=lifespan)
@app.post("/prediction")
async def prediction(
           output: PredictionOutput = Depends(newgroups_model.predict),
) -> PredictionOutput:
           return output

github.com/PacktPublishing/Building-Data-Science-Applications-with-FastAPI-Second-Edition/tree/main/chapter12/chapter12_prediction_endpoint.py

正如我们之前提到的,我们创建了一个NewsgroupsModel实例,以便将其注入到路径操作函数中。此外,我们正在实现一个生命周期处理器来调用load_model。通过这种方式,我们确保在应用程序启动时加载模型,并使其随时可用。

预测端点非常简单:正如你所看到的,我们直接依赖于predict方法,它会处理注入和验证负载。我们只需要返回输出即可。

就这样!再次感谢 FastAPI,让我们的生活变得更加轻松,它让我们能够编写简单且可读的代码,即使是面对复杂的任务。我们可以像往常一样使用 Uvicorn 来运行这个应用:


(venv) $ uvicorn chapter12.chapter12_prediction_endpoint:app

现在,我们可以尝试使用 HTTPie 进行一些预测:


$ http POST http://localhost:8000/prediction text="computer cpu memory ram"HTTP/1.1 200 OK
content-length: 36
content-type: application/json
date: Tue, 10 Jan 2023 07:37:22 GMT
server: uvicorn
{
           "category": "comp.sys.mac.hardware"
}

我们的机器学习分类器已经启动!为了进一步推动这一点,让我们看看如何使用 Joblib 实现一个简单的缓存机制。

使用 Joblib 缓存结果

如果你的模型需要一定时间才能进行预测,缓存结果可能会非常有用:如果某个特定输入的预测已经完成,那么返回我们保存在磁盘上的相同结果比再次运行计算更有意义。在本节中,我们将学习如何借助 Joblib 来实现这一点。

Joblib 为我们提供了一个非常方便且易于使用的工具,因此实现起来非常简单。主要的关注点是我们应该选择标准函数还是异步函数来实现端点和依赖关系。这样,我们可以更详细地解释 FastAPI 的一些技术细节。

我们将在前一节中提供的示例基础上进行构建。我们必须做的第一件事是初始化一个 Joblib 的Memory类,它是缓存函数结果的辅助工具。然后,我们可以为想要缓存的函数添加一个装饰器。你可以在以下示例中看到这一点:

chapter12_caching.py


memory = joblib.Memory(location="cache.joblib")@memory.cache(ignore=["model"])
def predict(model: Pipeline, text: str) -> int:
           prediction = model.predict([text])
           return prediction[0]

github.com/PacktPublishing/Building-Data-Science-Applications-with-FastAPI-Second-Edition/tree/main/chapter12/chapter12_caching.py

在初始化memory时,主要参数是location,它是 Joblib 存储结果的目录路径。Joblib 会自动将缓存结果保存在硬盘上。

然后,你可以看到我们实现了一个predict函数,它接受我们的 scikit-learn 模型和一些文本输入,并返回预测的类别索引。这与我们之前看到的预测操作相同。在这里,我们将其从NewsgroupsModel依赖类中提取出来,因为 Joblib 缓存主要是为了与常规函数一起使用的。缓存类方法并不推荐。正如你所看到的,我们只需在这个函数上方添加一个@memory.cache装饰器来启用 Joblib 缓存。

每当这个函数被调用时,Joblib 会检查它是否已经有相同参数的结果保存在磁盘上。如果有,它会直接返回该结果。否则,它会继续进行常规的函数调用。

正如你所看到的,我们为装饰器添加了一个ignore参数,这允许我们告诉 Joblib 在缓存机制中忽略某些参数。这里,我们排除了model参数。Joblib 无法存储复杂对象,比如 scikit-learn 估算器。但这不是问题:因为模型在多个预测之间是不会变化的,所以我们不关心是否缓存它。如果我们对模型进行了改进并部署了一个新的模型,我们只需要清除整个缓存,这样较早的预测就会使用新的模型重新计算。

现在,我们可以调整NewsgroupsModel依赖类,使其与这个新的predict函数兼容。你可以在以下示例中看到这一点:

chapter12_caching.py


class NewsgroupsModel:           model: Pipeline | None = None
           targets: list[str] | None = None
           def load_model(self) -> None:
                         """Loads the model"""
                         model_file = os.path.join(os.path.dirname(__file__), "newsgroups_model.joblib")
                         loaded_model: tuple[Pipeline, list[str]] = joblib.load(model_file)
                         model, targets = loaded_model
                         self.model = model
                         self.targets = targets
           def predict(self, input: PredictionInput) -> PredictionOutput:
                         """Runs a prediction"""
                         if not self.model or not self.targets:
                                       raise RuntimeError("Model is not loaded")
                         prediction = predict(self.model, input.text)
                         category = self.targets[prediction]
                         return PredictionOutput(category=category)

github.com/PacktPublishing/Building-Data-Science-Applications-with-FastAPI-Second-Edition/tree/main/chapter12/chapter12_caching.py

predict方法中,我们调用了外部的predict函数,而不是直接在方法内部调用,并且注意将模型和输入文本作为参数传递。之后我们只需要做的就是获取对应的类别名称并构建一个PredictionOutput对象。

最后,我们有了 REST API 端点。在这里,我们添加了一个delete/cache路由,以便通过 HTTP 请求清除整个 Joblib 缓存。你可以在以下示例中看到这一点:

chapter12_caching.py


@app.post("/prediction")def prediction(
           output: PredictionOutput = Depends(newgroups_model.predict),
) -> PredictionOutput:
           return output
@app.delete("/cache", status_code=status.HTTP_204_NO_CONTENT)
def delete_cache():
           memory.clear()

github.com/PacktPublishing/Building-Data-Science-Applications-with-FastAPI-Second-Edition/tree/main/chapter12/chapter12_caching.py

memory对象上的clear方法会删除硬盘上所有的 Joblib 缓存文件。

我们的 FastAPI 应用程序现在正在缓存预测结果。如果你使用相同的输入发出两次请求,第二次的响应将显示缓存结果。在这个例子中,我们的模型非常快,所以你不会注意到执行时间上的差异;然而,对于更复杂的模型,这可能会变得很有趣。

选择标准函数或异步函数

你可能注意到我们已经修改了predict方法以及predictiondelete_cache路径操作函数,使它们成为标准的非异步的函数。

自从本书开始以来,我们已经向你展示了 FastAPI 如何完全拥抱异步 I/O,以及这对应用程序性能的好处。我们还推荐了能够异步工作的库,例如数据库驱动程序,以便利用这一优势。

然而,在某些情况下,这是不可能的。在这种情况下,Joblib 被实现为同步工作。然而,它执行的是长时间的 I/O 操作:它在硬盘上读取和写入缓存文件。因此,它会阻塞进程,在此期间无法响应其他请求,正如我们在第二章异步 I/O部分中所解释的那样,Python 编程特性

为了解决这个问题,FastAPI 实现了一个巧妙的机制:如果你将路径操作函数或依赖项定义为标准的、非异步的函数,它将在一个单独的线程中运行。这意味着阻塞操作,例如同步文件读取,不会阻塞主进程。从某种意义上说,我们可以说它模仿了一个异步操作。

为了理解这一点,我们将进行一个简单的实验。在以下示例中,我们构建了一个包含三个端点的虚拟 FastAPI 应用程序:

  • /fast,它直接返回响应

  • /slow-async,一个定义为async的路径操作,创建一个同步阻塞操作,需要 10 秒钟才能完成

  • /slow-sync,一个作为标准方法定义的路径操作,创建一个同步阻塞操作,需要 10 秒钟才能完成

你可以在这里查看相应的代码:

chapter12_async_not_async.py


import timefrom fastapi import FastAPI
app = FastAPI()
@app.get("/fast")
async def fast():
           return {"endpoint": "fast"}
@app.get("/slow-async")
async def slow_async():
           """Runs in the main process"""
           time.sleep(10)    # Blocking sync operation
           return {"endpoint": "slow-async"}
@app.get("/slow-sync")
def slow_sync():
           """Runs in a thread"""
           time.sleep(10)    # Blocking sync operation
           return {"endpoint": "slow-sync"}

github.com/PacktPublishing/Building-Data-Science-Applications-with-FastAPI-Second-Edition/tree/main/chapter12/chapter12_async_not_async.py

通过这个简单的应用程序,我们的目标是观察那些阻塞操作如何阻塞主进程。让我们使用 Uvicorn 运行这个应用程序:


(venv) $ uvicorn chapter12.chapter12_async_not_async:app

接下来,打开两个新的终端。在第一个终端中,向/``slow-async端点发出请求:


$ http GET http://localhost:8000/slow-async

在没有等待响应的情况下,在第二个终端向/``fast端点发出请求:


$ http GET http://localhost:8000/fast

您会看到,您必须等待 10 秒钟才能收到/fast端点的响应。这意味着/slow-async阻塞了进程,导致服务器无法在此期间响应其他请求。

现在,让我们在/``slow-sync端点进行相同的实验:


$ http GET http://localhost:8000/slow-sync

再次运行以下命令:


$ http GET http://localhost:8000/fast

您将立即收到/fast的响应,而无需等待/slow-sync完成。由于它被定义为标准的非异步函数,FastAPI 会在一个线程中运行它,以避免阻塞。但是请记住,将任务发送到单独的线程会带来一些开销,因此在解决当前问题时,需要考虑最佳的处理方式。

那么,在使用 FastAPI 进行开发时,如何在路径操作和依赖之间选择标准函数和异步函数呢?以下是一些经验法则:

  • 如果函数不涉及长时间的 I/O 操作(例如文件读取、网络请求等),请将其定义为async

  • 如果涉及 I/O 操作,请参考以下内容:

    • 尝试选择与异步 I/O 兼容的库,正如我们在数据库或 HTTP 客户端中看到的那样。在这种情况下,您的函数将是async

    • 如果不可能,如 Joblib 缓存的情况,请将它们定义为标准函数。FastAPI 将会在单独的线程中运行它们。

由于 Joblib 在进行 I/O 操作时是完全同步的,我们将路径操作和依赖方法切换为同步的标准方法。

在这个示例中,差异不太明显,因为 I/O 操作较小且快速。然而,如果您需要实现更慢的操作,例如将文件上传到云存储时,记得考虑这个问题。

总结

恭喜!您现在可以构建一个快速高效的 REST API 来服务您的机器学习模型。感谢 Joblib,您学会了如何将训练好的 scikit-learn 估计器保存到一个文件中,以便轻松加载并在您的应用程序中使用。我们还展示了使用 Joblib 缓存预测结果的方法。最后,我们讨论了 FastAPI 如何通过将同步操作发送到单独的线程来处理,以避免阻塞。虽然这有点技术性,但在处理阻塞 I/O 操作时,牢记这一点是很重要的。

我们的 FastAPI 之旅接近尾声。在让您独立构建令人惊叹的数据科学应用之前,我们将提供三个章节,进一步推动这一进程并研究更复杂的使用案例。我们将从一个可以执行实时物体检测的应用开始,得益于 WebSocket 和计算机视觉模型。

第十三章:使用 WebSockets 和 FastAPI 实现实时物体检测系统

在上一章中,你学习了如何创建高效的 REST API 端点,用于通过训练好的机器学习模型进行预测。这种方法涵盖了很多使用场景,假设我们有一个单一的观测值需要处理。然而,在某些情况下,我们可能需要持续对一系列输入进行预测——例如,一个实时处理视频输入的物体检测系统。这正是我们将在本章中构建的内容。怎么做?如果你还记得,除了 HTTP 端点,FastAPI 还具备处理 WebSockets 端点的能力,这让我们能够发送和接收数据流。在这种情况下,浏览器会通过 WebSocket 发送来自摄像头的图像流,我们的应用程序会运行物体检测算法,并返回图像中每个检测到的物体的坐标和标签。为此任务,我们将依赖于Hugging Face,它既是一组工具,也是一个预训练 AI 模型的库。

本章我们将涵盖以下主要内容:

  • 使用 Hugging Face 库的计算机视觉模型

  • 实现一个 HTTP 端点,执行单张图像的物体检测

  • 在 WebSocket 中从浏览器发送图像流

  • 在浏览器中显示物体检测结果

技术要求

本章你将需要一个 Python 虚拟环境,就像我们在第一章中所设置的那样,Python 开发 环境设置

你可以在这个专门的 GitHub 仓库中找到本章的所有代码示例:github.com/PacktPublishing/Building-Data-Science-Applications-with-FastAPI-Second-Edition/tree/main/chapter13

使用 Hugging Face 库的计算机视觉模型

计算机视觉是一个研究和技术领域,致力于使计算机能够从数字图像或视频中提取有意义的信息,从而模拟人类视觉能力。它涉及基于统计方法或机器学习开发算法,使机器能够理解、分析和解释视觉数据。计算机视觉的一个典型应用是物体检测:一个能够在图像中检测和识别物体的系统。这正是我们将在本章中构建的系统。

为了帮助我们完成这个任务,我们将使用 Hugging Face 提供的一组工具。Hugging Face 是一家旨在让开发者能够快速、轻松地使用最新、最强大的 AI 模型的公司。为此,它构建了两个工具:

  • 一套基于机器学习库(如 PyTorch 和 TensorFlow)构建的开源 Python 工具集。我们将在本章中使用其中的一些工具。

  • 一个在线库,用于分享和下载各种机器学习任务的预训练模型,例如计算机视觉或图像生成。

你可以在其官方网站上了解更多它的功能:huggingface.co/

你会看到,这将极大地帮助我们在短时间内构建一个强大且精确的目标检测系统!首先,我们将安装项目所需的所有库:


(venv) $ pip install "transformers[torch]" Pillow

Hugging Face 的 transformers 库将允许我们下载并运行预训练的机器学习模型。请注意,我们通过可选的 torch 依赖项来安装它。Hugging Face 工具可以与 PyTorch 或 TensorFlow 一起使用,这两者都是非常强大的机器学习框架。在这里,我们选择使用 PyTorch。Pillow 是一个广泛使用的 Python 库,用于处理图像。稍后我们会看到为什么需要它。

在开始使用 FastAPI 之前,让我们先实现一个简单的脚本来运行一个目标检测算法。它包括四个主要步骤:

  1. 使用 Pillow 从磁盘加载图像。

  2. 加载一个预训练的目标检测模型。

  3. 在我们的图像上运行模型。

  4. 通过在检测到的物体周围绘制矩形来显示结果。

我们将一步一步地进行实现:

chapter13_object_detection.py


from pathlib import Pathimport torch
from PIL import Image, ImageDraw, ImageFont
from transformers import YolosForObjectDetection, YolosImageProcessor
root_directory = Path(__file__).parent.parent
picture_path = root_directory / "assets" / "coffee-shop.jpg"
image = Image.open(picture_path)

github.com/PacktPublishing/Building-Data-Science-Applications-with-FastAPI-Second-Edition/tree/main/chapter13/chapter13_object_detection.py

如你所见,第一步是从磁盘加载我们的图像。在这个示例中,我们使用名为 coffee-shop.jpg 的图像,该图像可以在我们的示例仓库中找到,地址是 github.com/PacktPublishing/Building-Data-Science-Applications-with-FastAPI-Second-Edition/blob/main/assets/coffee-shop.jpg

chapter13_object_detection.py


image_processor = YolosImageProcessor.from_pretrained("hustvl/yolos-tiny")model = YolosForObjectDetection.from_pretrained("hustvl/yolos-tiny")

github.com/PacktPublishing/Building-Data-Science-Applications-with-FastAPI-Second-Edition/tree/main/chapter13/chapter13_object_detection.py

接下来,我们从 Hugging Face 加载一个模型。在这个示例中,我们选择了 YOLOS 模型。它是一种先进的目标检测方法,已在 118K 个带注释的图像上进行训练。你可以在以下 Hugging Face 文章中了解更多关于该技术的方法:huggingface.co/docs/transformers/model_doc/yolos。为了限制下载大小并节省计算机磁盘空间,我们选择使用精简版,这是原始模型的一个更轻量化版本,可以在普通机器上运行,同时保持良好的精度。这个版本在 Hugging Face 上有详细描述:huggingface.co/hustvl/yolos-tiny

请注意,我们实例化了两个东西:图像处理器模型。如果你还记得我们在第十一章《Python 中的数据科学入门》中提到的内容,你会知道我们需要一组特征来供我们的机器学习算法使用。因此,图像处理器的作用是将原始图像转换为对模型有意义的一组特征。

这正是我们在接下来的代码行中所做的:我们通过调用image_processor处理图像,创建一个inputs变量。请注意,return_tensors参数被设置为pt,因为我们选择了 PyTorch 作为我们的底层机器学习框架。然后,我们可以将这个inputs变量输入到模型中以获得outputs

chapter13_object_detection.py


inputs = image_processor(images=image, return_tensors="pt")outputs = model(**inputs)

github.com/PacktPublishing/Building-Data-Science-Applications-with-FastAPI-Second-Edition/tree/main/chapter13/chapter13_object_detection.py

你可能会认为这就是预测阶段的全部内容,我们现在可以展示结果了。然而,事实并非如此。此类算法的结果是一组多维矩阵,著名的post_process_object_detection操作由image_processor提供:

chapter13_object_detection.py


target_sizes = torch.tensor([image.size[::-1]])results = image_processor.post_process_object_detection(
    outputs, target_sizes=target_sizes
)[0]

github.com/PacktPublishing/Building-Data-Science-Applications-with-FastAPI-Second-Edition/tree/main/chapter13/chapter13_object_detection.py

这个操作的结果是一个字典,包含以下内容:

  • labels:每个检测到的物体的标签列表

  • boxes:每个检测到的物体的边界框坐标

  • scores:算法对于每个检测到的物体的置信度分数

我们需要做的就是遍历这些对象,以便利用 Pillow 绘制矩形和相应的标签。最后,我们只展示处理后的图像。请注意,我们只考虑那些得分大于0.7的物体,以减少假阳性的数量:

chapter13_object_detection.py


draw = ImageDraw.Draw(image)font_path = root_directory / "assets" / "OpenSans-ExtraBold.ttf"
font = ImageFont.truetype(str(font_path), 24)
for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
    if score > 0.7:
        box_values = box.tolist()
        label = model.config.id2label[label.item()]
        draw.rectangle(box_values, outline="red", width=5)
        draw.text(box_values[0:2], label, fill="red", font=font)
image.show()

github.com/PacktPublishing/Building-Data-Science-Applications-with-FastAPI-Second-Edition/tree/main/chapter13/chapter13_object_detection.py

多亏了 Pillow,我们能够绘制矩形并在检测到的物体上方添加标签。请注意,我们加载了一个自定义字体——Open Sans,它是一个可以在网上获得的开源字体:fonts.google.com/specimen/Open+Sans。让我们尝试运行这个脚本,看看结果:


(venv) $ python chapter13/chapter13_object_detection.py

首次运行时,你会看到模型被下载。根据你的计算机性能,预测过程可能需要几秒钟。完成后,生成的图像应该会自动打开,如图 13.1所示。

图 13.1 – 在示例图像上的目标检测结果

图 13.1 – 在示例图像上的目标检测结果

你可以看到,模型检测到图像中的几个人物,以及各种物体,如沙发和椅子。就这样!不到 30 行代码就能实现一个可运行的目标检测脚本!Hugging Face 使我们能够高效地利用最新 AI 技术的强大功能。

当然,我们在这一章的目标是将所有这些智能功能部署到远程服务器上,以便能够为成千上万的用户提供这一体验。再次强调,FastAPI 将是我们在这里的得力助手。

实现一个 REST 端点,用于在单张图像上执行目标检测

在使用 WebSockets 之前,我们先从简单开始,利用 FastAPI 实现一个经典的 HTTP 端点,接受图像上传并对其进行目标检测。正如你所看到的,与之前的示例的主要区别在于我们如何获取图像:不是从磁盘读取,而是通过文件上传获取,之后需要将其转换为 Pillow 图像对象。

此外,我们还将使用我们在第十二章中看到的完全相同的模式,使用 FastAPI 创建高效的预测 API 端点——也就是为我们的预测模型创建一个专门的类,该类将在生命周期处理程序中加载。

在此实现中,我们做的第一件事是定义 Pydantic 模型,以便正确地结构化我们预测模型的输出。你可以看到如下所示:

chapter13_api.py


class Object(BaseModel):    box: tuple[float, float, float, float]
    label: str
class Objects(BaseModel):
    objects: list[Object]

github.com/PacktPublishing/Building-Data-Science-Applications-with-FastAPI-Second-Edition/tree/main/chapter13/chapter13_api.py

我们有一个表示单个检测到的目标的模型,它包括box,一个包含四个数字的元组,描述边界框的坐标,以及label,表示检测到的物体类型。Objects模型是一个简单的结构,包含物体列表。

我们不会详细介绍模型预测类,因为它与我们在上一章和上一节中看到的非常相似。相反,我们直接关注 FastAPI 端点的实现:

chapter13_api.py


object_detection = ObjectDetection()@contextlib.asynccontextmanager
async def lifespan(app: FastAPI):
    object_detection.load_model()
    yield
app = FastAPI(lifespan=lifespan)
@app.post("/object-detection", response_model=Objects)
async def post_object_detection(image: UploadFile = File(...)) -> Objects:
    image_object = Image.open(image.file)
    return object_detection.predict(image_object)

github.com/PacktPublishing/Building-Data-Science-Applications-with-FastAPI-Second-Edition/tree/main/chapter13/chapter13_api.py

这里没有什么特别令人惊讶的!需要关注的重点是正确使用UploadFileFile依赖项,以便获取上传的文件。如果你需要复习这一部分内容,可以查看第三章中关于表单数据和文件上传的章节,快速开发 RESTful API 使用 FastAPI。然后,我们只需将其实例化为合适的 Pillow 图像对象,并调用我们的预测模型。

如我们所说,别忘了在生命周期处理程序中加载模型。

你可以使用常规的 Uvicorn 命令来运行这个示例:


(venv) $ uvicorn chapter13.chapter13_api:app

我们将使用上一节中看到的相同的咖啡店图片。让我们用 HTTPie 上传它到我们的端点:


$ http --form POST http://localhost:8000/object-detection image@./assets/coffee-shop.jpg{
    "objects": [
        {
            "box": [659.8709716796875, 592.8882446289062, 792.0460815429688, 840.2132568359375],
            "label": "person"
        },
        {
            "box": [873.5499267578125, 875.7918090820312, 1649.1378173828125, 1296.362548828125],
            "label": "couch"
        }
    ]
}

我们正确地得到了检测到的对象列表,每个对象都有它的边界框和标签。太棒了!我们的目标检测系统现在已经作为一个 Web 服务器可用。然而,我们的目标仍然是创建一个实时系统:借助 WebSockets,我们将能够处理图像流。

实现 WebSocket 以对图像流进行目标检测

WebSockets 的主要优势之一,正如我们在第八章中看到的,在 FastAPI 中定义用于双向交互通信的 WebSockets,是它在客户端和服务器之间打开了一个全双工通信通道。一旦连接建立,消息可以快速传递,而不需要经过 HTTP 协议的所有步骤。因此,它更适合实时传输大量数据。

这里的关键是实现一个 WebSocket 端点,能够接收图像数据并进行目标检测。这里的主要挑战是处理一个被称为背压的现象。简单来说,我们将从浏览器接收到的图像比服务器能够处理的要多,因为运行检测算法需要一定的时间。因此,我们必须使用一个有限大小的队列(或缓冲区),并在处理流时丢弃一些图像,以便接近实时地处理。

我们将逐步讲解实现过程:

app.py


async def receive(websocket: WebSocket, queue: asyncio.Queue):    while True:
        bytes = await websocket.receive_bytes()
        try:
            queue.put_nowait(bytes)
        except asyncio.QueueFull:
            pass
async def detect(websocket: WebSocket, queue: asyncio.Queue):
    while True:
        bytes = await queue.get()
        image = Image.open(io.BytesIO(bytes))
        objects = object_detection.predict(image)
        await websocket.send_json(objects.dict())

github.com/PacktPublishing/Building-Data-Science-Applications-with-FastAPI-Second-Edition/tree/main/chapter13/websocket_object_detection/app.py

我们定义了两个任务:receivedetect。第一个任务是等待从 WebSocket 接收原始字节,而第二个任务则执行检测并发送结果,正如我们在上一节中看到的那样。

这里的关键是使用asyncio.Queue对象。这是一个便捷的结构,允许我们在内存中排队一些数据,并以先进先出FIFO)策略来检索它。我们可以设置存储在队列中的元素数量限制:这就是我们限制处理图像数量的方式。

receive 函数接收数据并将其放入队列末尾。在使用 asyncio.Queue 时,我们有两个方法可以将新元素放入队列:putput_nowait。如果队列已满,第一个方法会等待直到队列有空间。这不是我们在这里想要的:我们希望丢弃那些无法及时处理的图像。使用 put_nowait 时,如果队列已满,会抛出 QueueFull 异常。在这种情况下,我们只需跳过并丢弃数据。

另一方面,detect 函数从队列中提取第一个消息,并在发送结果之前运行检测。请注意,由于我们直接获取的是原始图像字节,我们需要用 io.BytesIO 将它们包装起来,才能让 Pillow 处理。

WebSocket 的实现本身类似于我们在第八章中看到的内容,在 FastAPI 中定义 WebSocket 进行双向交互通信。我们正在调度这两个任务并等待其中一个任务停止。由于它们都运行一个无限循环,因此当 WebSocket 断开连接时,这个情况会发生:

app.py


@app.websocket("/object-detection")async def ws_object_detection(websocket: WebSocket):
    await websocket.accept()
    queue: asyncio.Queue = asyncio.Queue(maxsize=1)
    receive_task = asyncio.create_task(receive(websocket, queue))
    detect_task = asyncio.create_task(detect(websocket, queue))
    try:
        done, pending = await asyncio.wait(
            {receive_task, detect_task},
            return_when=asyncio.FIRST_COMPLETED,
        )
        for task in pending:
            task.cancel()
        for task in done:
            task.result()
    except WebSocketDisconnect:
        pass

github.com/PacktPublishing/Building-Data-Science-Applications-with-FastAPI-Second-Edition/tree/main/chapter13/websocket_object_detection/app.py

提供静态文件

如果你查看前面示例的完整实现,你会注意到我们在服务器中定义了另外两个东西:一个 index 端点,它仅返回 index.html 文件,以及一个 StaticFiles 应用,它被挂载在 /assets 路径下。这两个功能的存在是为了让我们的 FastAPI 应用直接提供 HTML 和 JavaScript 代码。这样,浏览器就能够在同一个服务器上查询这些文件。

这部分的关键点是,尽管 FastAPI 是为构建 REST API 设计的,但它同样可以完美地提供 HTML 和静态文件。

我们的后端现在已经准备好!接下来,让我们看看如何在浏览器中使用它的功能。

通过 WebSocket 从浏览器发送图像流

在本节中,我们将展示如何在浏览器中捕捉来自摄像头的图像并通过 WebSocket 发送。由于这主要涉及 JavaScript 代码,坦率来说,它有点超出了本书的范围,但它对于让应用程序正常工作是必要的。

第一步是在浏览器中启用摄像头输入,打开 WebSocket 连接,捕捉摄像头图像并通过 WebSocket 发送。基本上,它会像这样工作:通过 MediaDevices 浏览器 API,我们将能够列出设备上所有可用的摄像头输入。借此,我们将构建一个选择表单,供用户选择他们想要使用的摄像头。你可以在以下代码中看到具体的 JavaScript 实现:

script.js


window.addEventListener('DOMContentLoaded', (event) => {  const video = document.getElementById('video');
  const canvas = document.getElementById('canvas');
  const cameraSelect = document.getElementById('camera-select');
  let socket;
  // List available cameras and fill select
  navigator.mediaDevices.getUserMedia({ audio: true, video: true }).then(() => {
    navigator.mediaDevices.enumerateDevices().then((devices) => {
      for (const device of devices) {
        if (device.kind === 'videoinput' && device.deviceId) {
          const deviceOption = document.createElement('option');
          deviceOption.value = device.deviceId;
          deviceOption.innerText = device.label;
          cameraSelect.appendChild(deviceOption);
        }
      }
    });
  });

github.com/PacktPublishing/Building-Data-Science-Applications-with-FastAPI-Second-Edition/tree/main/chapter13/websocket_object_detection/assets/script.js

一旦用户提交表单,我们会调用一个startObjectDetection函数,并传入选定的摄像头。大部分实际的检测逻辑是在这个函数中实现的:

script.js


  // Start object detection on the selected camera on submit  document.getElementById('form-connect').addEventListener('submit', (event) => {
    event.preventDefault();
    // Close previous socket is there is one
    if (socket) {
      socket.close();
    }
    const deviceId = cameraSelect.selectedOptions[0].value;
    socket = startObjectDetection(video, canvas, deviceId);
  });
});

github.com/PacktPublishing/Building-Data-Science-Applications-with-FastAPI-Second-Edition/tree/main/chapter13/websocket_object_detection/assets/script.js

让我们看一下下面代码块中的startObjectDetection函数。首先,我们与 WebSocket 建立连接。连接打开后,我们可以开始从选定的摄像头获取图像流。为此,我们使用MediaDevices API 来启动视频捕获,并将输出显示在一个 HTML 的<video>元素中。你可以在 MDN 文档中阅读有关MediaDevices API 的所有细节:developer.mozilla.org/en-US/docs/Web/API/MediaDevices

script.js


const startObjectDetection = (video, canvas, deviceId) => {  const socket = new WebSocket(`ws://${location.host}/object-detection`);
  let intervalId;
  // Connection opened
  socket.addEventListener('open', function () {
    // Start reading video from device
    navigator.mediaDevices.getUserMedia({
      audio: false,
      video: {
        deviceId,
        width: { max: 640 },
        height: { max: 480 },
      },
    }).then(function (stream) {
      video.srcObject = stream;
      video.play().then(() => {
        // Adapt overlay canvas size to the video size
        canvas.width = video.videoWidth;
        canvas.height = video.videoHeight;

github.com/PacktPublishing/Building-Data-Science-Applications-with-FastAPI-Second-Edition/tree/main/chapter13/websocket_object_detection/assets/script.js

然后,正如下一个代码块所示,我们启动一个重复的任务,捕获来自视频输入的图像并将其发送到服务器。为了实现这一点,我们必须使用一个<canvas>元素,这是一个专门用于图形绘制的 HTML 标签。它提供了完整的 JavaScript API,允许我们以编程方式在其中绘制图像。在这里,我们可以绘制当前的视频图像,并将其转换为有效的 JPEG 字节。如果你想了解更多关于这个内容,MDN 提供了一个非常详细的<canvas>教程:developer.mozilla.org/en-US/docs/Web/API/Canvas_API/Tutorial

script.js


        // Send an image in the WebSocket every 42 ms        intervalId = setInterval(() => {
          // Create a virtual canvas to draw current video image
          const canvas = document.createElement('canvas');
          const ctx = canvas.getContext('2d');
          canvas.width = video.videoWidth;
          canvas.height = video.videoHeight;
          ctx.drawImage(video, 0, 0);
          // Convert it to JPEG and send it to the WebSocket
          canvas.toBlob((blob) => socket.send(blob), 'image/jpeg');
        }, IMAGE_INTERVAL_MS);
      });
    });
  });

github.com/PacktPublishing/Building-Data-Science-Applications-with-FastAPI-Second-Edition/tree/main/chapter13/websocket_object_detection/assets/script.js

请注意,我们将视频输入的大小限制为 640x480 像素,以防止上传过大的图像使服务器崩溃。此外,我们将间隔设置为每 42 毫秒执行一次(该值在IMAGE_INTERVAL_MS常量中设置),大约相当于每秒 24 帧图像。

最后,我们将事件监听器连接起来,以处理从 WebSocket 接收到的消息。它调用了drawObjects函数,我们将在下一节中详细介绍:

script.js


  // Listen for messages  socket.addEventListener('message', function (event) {
    drawObjects(video, canvas, JSON.parse(event.data));
  });
  // Stop the interval and video reading on close
  socket.addEventListener('close', function () {
    window.clearInterval(intervalId);
    video.pause();
  });
  return socket;
};

github.com/PacktPublishing/Building-Data-Science-Applications-with-FastAPI-Second-Edition/tree/main/chapter13/websocket_object_detection/assets/script.js

在浏览器中展示物体检测结果

现在我们能够将输入图像发送到服务器,我们需要在浏览器中展示检测结果。与我们在使用 Hugging Face 计算机视觉模型部分中展示的类似,我们将围绕检测到的对象绘制一个绿色矩形,并标注它们的标签。因此,我们需要找到一种方式,将服务器发送的矩形坐标在浏览器中绘制出来。

为了实现这一点,我们将再次使用<canvas>元素。这次,它将对用户可见,我们将使用它来绘制矩形。关键是使用 CSS 使这个元素覆盖视频:这样,矩形就会直接显示在视频和对应对象的上方。你可以在这里看到 HTML 代码:

index.html


<body>  <div class="container">
    <h1 class="my-3">Chapter 13 - Real time object detection</h1>
    <form id="form-connect">
      <div class="input-group mb-3">
        <select id="camera-select"></select>
        <button class="btn btn-success" type="submit" id="button-start">Start</button>
      </div>
    </form>
    <div class="position-relative" style="width: 640px; height: 480px;">
      <video id="video"></video>
      <canvas id="canvas" class="position-absolute top-0 start-0"></canvas>
    </div>
  </div>
  <script src="img/script.js"></script>
</body>

github.com/PacktPublishing/Building-Data-Science-Applications-with-FastAPI-Second-Edition/tree/main/chapter13/websocket_object_detection/index.html

我们使用了来自 Bootstrap 的 CSS 类,Bootstrap 是一个非常常见的 CSS 库,提供了很多类似这样的辅助工具。基本上,我们通过绝对定位设置了 canvas,并将其放置在左上角,这样它就能覆盖视频元素。

关键在于使用 Canvas API 根据接收到的坐标绘制矩形。这正是drawObjects函数的目的,下面的示例代码块展示了这一点:

script.js


const drawObjects = (video, canvas, objects) => {  const ctx = canvas.getContext('2d');
  ctx.width = video.videoWidth;
  ctx.height = video.videoHeight;
  ctx.beginPath();
  ctx.clearRect(0, 0, ctx.width, ctx.height);
  for (const object of objects.objects) {
    const [x1, y1, x2, y2] = object.box;
    const label = object.label;
    ctx.strokeStyle = '#49fb35';
    ctx.beginPath();
    ctx.rect(x1, y1, x2 - x1, y2 - y1);
    ctx.stroke();
    ctx.font = 'bold 16px sans-serif';
    ctx.fillStyle = '#ff0000';
    ctx.fillText(label, x1 - 5 , y1 - 5);
  }
};

github.com/PacktPublishing/Building-Data-Science-Applications-with-FastAPI-Second-Edition/tree/main/chapter13/websocket_object_detection/assets/script.js

使用<canvas>元素,我们可以使用 2D 上下文在对象中绘制内容。请注意,我们首先清除所有内容,以移除上次检测的矩形。然后,我们遍历所有检测到的对象,并使用给定的坐标x1y1x2y2绘制一个矩形。最后,我们会在矩形上方稍微绘制标签。

我们的系统现在完成了!图 13.2 给出了我们实现的文件结构概览。

图 13.2 – 对象检测应用结构

图 13.2 – 对象检测应用结构

现在是时候尝试一下了!我们可以使用常见的 Uvicorn 命令启动它:


(venv) $ uvicorn chapter13.websocket_object_detection.app:app

您可以通过地址http://localhost:8000在浏览器中访问应用程序。正如我们在前一部分所说,index端点将被调用并返回我们的index.html文件。

您将看到一个界面,邀请您选择要使用的摄像头,如图 13.3所示:

图 13.3 – 用于对象检测网页应用的摄像头选择

图 13.3 – 用于对象检测网页应用的摄像头选择

选择您想要使用的摄像头并点击开始。视频输出将显示出来,通过 WebSocket 开始对象检测,并且绿色矩形将会围绕检测到的对象绘制。我们在图 13.4中展示了这一过程:

图 13.4 – 运行对象检测网页应用

图 13.4 – 运行对象检测网页应用

它成功了!我们将我们 Python 系统的智能带到了用户的网页浏览器中。这只是使用 WebSockets 和机器学习算法可以实现的一种示例,但它绝对可以让您为用户创建接近实时的体验。

总结

在本章中,我们展示了 WebSockets 如何帮助我们为用户带来更具互动性的体验。得益于 Hugging Face 社区提供的预训练模型,我们能够迅速实现一个对象检测系统。接着,在 FastAPI 的帮助下,我们将其集成到一个 WebSocket 端点。最后,通过使用现代 JavaScript API,我们直接在浏览器中发送视频输入并显示算法结果。总的来说,像这样的项目乍一看可能显得很复杂,但我们看到强大的工具,如 FastAPI,能够让我们在非常短的时间内并且通过易于理解的源代码实现结果。

到目前为止,在我们的不同示例和项目中,我们假设我们使用的机器学习模型足够快,可以直接在 API 端点或 WebSocket 任务中运行。然而,情况并非总是如此。在某些情况下,算法如此复杂,以至于运行需要几分钟。如果我们直接在 API 端点内部运行这种算法,用户将不得不等待很长时间才能得到响应。这不仅会让用户感到困惑,还会迅速堵塞整个服务器,阻止其他用户使用 API。为了解决这个问题,我们需要为 API 服务器配备一个助手:一个工作者。

在下一章中,我们将研究这个挑战的一个具体例子:我们将构建我们自己的 AI 系统,从文本提示生成图像!

第十四章:使用 Stable Diffusion 模型创建分布式文本到图像 AI 系统

到目前为止,在这本书中,我们构建的 API 中所有操作都是在请求处理内部计算的。换句话说,用户必须等待服务器完成我们定义的所有操作(如请求验证、数据库查询、ML 预测等),才能收到他们的响应。然而,并非总是希望或可能要求这种行为。

典型例子是电子邮件通知。在 Web 应用程序中,我们经常需要向用户发送电子邮件,因为他们刚刚注册或执行了特定操作。为了做到这一点,服务器需要向电子邮件服务器发送请求,以便发送电子邮件。此操作可能需要几毫秒时间。如果我们在请求处理中执行此操作,响应将延迟直到我们发送电子邮件。这不是一个很好的体验,因为用户并不真正关心电子邮件是如何何时发送的。这个例子是我们通常所说的后台操作的典型例子:需要在我们的应用程序中完成的事情,但不需要直接用户交互。

另一种情况是当用户请求一个耗时的操作,在合理的时间内无法完成。这通常是复杂数据导出或重型 AI 模型的情况。在这种情况下,用户希望直接获取结果,但如果在请求处理程序中执行此操作,将会阻塞服务器进程,直到完成。如果大量用户请求这种操作,会迅速使我们的服务器无响应。此外,某些网络基础设施,如代理或 Web 客户端(如浏览器),具有非常严格的超时设置,这意味着如果响应时间过长,它们通常会取消操作。

为了解决这个问题,我们将引入一个典型的 Web 应用程序架构:web-queue-worker。正如我们将在本章中看到的,我们将把最昂贵、耗时最长的操作推迟到后台进程,即worker。为了展示这种架构的运行方式,我们将建立我们自己的 AI 系统,使用Stable Diffusion模型根据文本提示生成图像。

本章中,我们将涵盖以下主要话题:

  • 使用 Stable Diffusion 模型与 Hugging Face Diffusers 生成图像的文本提示

  • 使用 Dramatiq 实现工作进程和图像生成任务

  • 存储和服务于对象存储中的文件

技术要求

对于本章,您将需要一个 Python 虚拟环境,就像我们在第一章中设置的那样,Python 开发环境设置

为了正确运行 Stable Diffusion 模型,我们建议你使用配备至少 16 GB RAM 的最新计算机,理想情况下还应配备 8 GB VRAM 的专用 GPU。对于 Mac 用户,配备 M1 Pro 或 M2 Pro 芯片的最新型号也非常适合。如果你没有这种机器,也不用担心:我们会告诉你如何以其他方式运行系统——唯一的缺点是图像生成会变慢并且效果较差。

要运行工作程序,你需要在本地计算机上运行Redis 服务器。最简单的方法是将其作为 Docker 容器运行。如果你以前从未使用过 Docker,我们建议你阅读官方文档中的入门教程,网址为docs.docker.com/get-started/。完成后,你将能够通过以下简单命令运行 Redis 服务器:


$ docker run -d --name worker-redis -p 6379:6379 redis

你可以在专用的 GitHub 仓库中找到本章的所有代码示例,地址为github.com/PacktPublishing/Building-Data-Science-Applications-with-FastAPI-Second-Edition/tree/main/chapter14

使用 Stable Diffusion 从文本提示生成图像

最近,一种新一代的 AI 工具引起了全世界的关注:图像生成模型,例如 DALL-E 或 Midjourney。这些模型是在大量图像数据上进行训练的,能够从简单的文本提示中生成全新的图像。这些 AI 模型非常适合作为后台工作程序:它们的处理时间为几秒钟甚至几分钟,并且需要大量的 CPU、RAM 甚至 GPU 资源。

为了构建我们的系统,我们将依赖于 Stable Diffusion,这是一种非常流行的图像生成模型,发布于 2022 年。该模型是公开的,可以在现代游戏计算机上运行。正如我们在上一章中所做的,我们将依赖 Hugging Face 工具来下载和运行该模型。

首先,让我们安装所需的工具:


(venv) $ pip install accelerate diffusers

现在,我们已经准备好通过 Hugging Face 使用扩散模型了。

在 Python 脚本中实现模型

在下面的示例中,我们将展示一个能够实例化模型并运行图像生成的类的实现。再次提醒,我们将应用懒加载模式,使用单独的 load_modelgenerate 方法。首先,让我们专注于 load_model

text_to_image.py


class TextToImage:    pipe: StableDiffusionPipeline | None = None
    def load_model(self) -> None:
        # Enable CUDA GPU
        if torch.cuda.is_available():
            device = "cuda"
        # Enable Apple Silicon (M1) GPU
        elif torch.backends.mps.is_available():
            device = "mps"
        # Fallback to CPU
        else:
            device = "cpu"
        pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5")
        pipe.to(device)
        self.pipe = pipe

github.com/PacktPublishing/Building-Data-Science-Applications-with-FastAPI-Second-Edition/tree/main/chapter14/basic/text_to_image.py

该方法的第一部分旨在根据你的计算机找到最有效的运行模型的方式。当在 GPU 上运行时,这些扩散模型的速度更快——这就是为什么我们首先检查是否有 CUDA(NVIDIA GPU)或 MPS(Apple Silicon)设备可用。如果没有,我们将退回到 CPU。

然后,我们只需创建一个由 Hugging Face 提供的StableDiffusionPipeline管道。我们只需要设置我们想要从 Hub 下载的模型。对于这个例子,我们选择了runwayml/stable-diffusion-v1-5。你可以在 Hugging Face 上找到它的详细信息:huggingface.co/runwayml/stable-diffusion-v1-5

我们现在可以专注于generate方法:

text_to_image.py


    def generate(        self,
        prompt: str,
        *,
        negative_prompt: str | None = None,
        num_steps: int = 50,
        callback: Callable[[int, int, torch.FloatTensor], None] | None = None,
    )    Image.Image:
        if not self.pipe:
            raise RuntimeError("Pipeline is not loaded")
        return self.pipe(
            prompt,
            negative_prompt=negative_prompt,
            num_inference_steps=num_steps,
            guidance_scale=9.0,
            callback=callback,
        ).images[0]

github.com/PacktPublishing/Building-Data-Science-Applications-with-FastAPI-Second-Edition/tree/main/chapter14/basic/text_to_image.py

你可以看到它接受四个参数:

  • prompt,当然,这是描述我们想要生成的图像的文本提示。

  • negative_prompt,这是一个可选的提示,用于告诉模型我们绝对不希望出现的内容。

  • num_steps,即模型应执行的推理步骤数。更多步骤会导致更好的图像,但每次迭代都会延迟推理。默认值50应该在速度和质量之间提供良好的平衡。

  • callback,这是一个可选的函数,它将在每次迭代步骤中被调用。这对于了解生成进度并可能执行更多逻辑(如将进度保存到数据库中)非常有用。

方法签名中的星号(*)是什么意思?

你可能已经注意到方法签名中的星号(*)。它告诉 Python,星号后面的参数应该仅作为关键字参数处理。换句话说,你只能像这样调用它们:.generate("PROMPT", negative_prompt="NEGATIVE", num_steps=10)

尽管不是必须的,但这是一种保持函数清晰且自解释的方式。如果你开发的是供其他开发者使用的类或函数,这尤其重要。

还有一种语法可以强制参数仅作为位置参数传递,方法是使用斜杠(/)符号。你可以在这里阅读更多相关内容:docs.python.org/3/whatsnew/3.8.html#positional-only-parameters

然后,我们只需要将这些参数传递给pipe。如果需要的话,还有更多的参数可以调节,但默认的参数应该会给你不错的结果。你可以在 Hugging Face 文档中找到完整的参数列表:huggingface.co/docs/diffusers/api/pipelines/stable_diffusion/text2img#diffusers.StableDiffusionPipeline.__call__。这个pipe对象能够为每个提示生成多张图像,因此该操作的结果是一个 Pillow 图像列表。这里的默认行为是生成一张图像,所以我们直接返回第一张。

就这些!再次感谢 Hugging Face,通过允许我们在几十行代码内运行最前沿的模型,真的是让我们的生活变得更轻松!

执行 Python 脚本

我们敢打赌你急于自己试一试——所以我们在示例的底部添加了一个小的main脚本:

text_to_image.py


if __name__ == "__main__":    text_to_image = TextToImage()
    text_to_image.load_model()
    def callback(step: int, _timestep, _tensor):
        print(f"🚀 Step {step}")
    image = text_to_image.generate(
        "A Renaissance castle in the Loire Valley",
        negative_prompt="low quality, ugly",
        callback=callback,
    )
    image.save("output.png")

github.com/PacktPublishing/Building-Data-Science-Applications-with-FastAPI-Second-Edition/tree/main/chapter14/basic/text_to_image.py

这个小脚本实例化了我们的TextToImage类,加载了模型,并在保存到磁盘之前生成了图像。我们还定义了一个虚拟回调函数,让你能看到它是如何工作的。

当你第一次运行这个脚本时,你会注意到 Hugging Face 会将几个 GB 的文件下载到你的计算机上:那就是稳定扩散模型,确实相当庞大!

然后,推理开始了。你会看到一个进度条,显示剩余的推理步骤数,并显示我们回调函数中的print语句,如图 14.1所示。

图 14.1 – 稳定扩散生成图像

图 14.1 – 稳定扩散生成图像

生成一张图像需要多长时间?

我们在不同类型的计算机上进行了多次测试。在配备 8 GB RAM 的现代 NVIDIA GPU 或 M1 Pro 芯片的 Mac 上,模型能够在大约一分钟内生成一张图像,并且内存使用合理。而在 CPU 上运行时,大约需要5 到 10 分钟,并且会占用多达 16 GB 的内存。

如果你的计算机上推理速度确实太慢,你可以尝试减少num_steps参数。

当推理完成后,你会在磁盘上找到生成的图像和你的脚本。图 14.2展示了这种结果的一个例子。不错吧?

图 14.2 – 稳定扩散图像生成结果

图 14.2 – 稳定扩散图像生成结果

现在,我们已经拥有了我们 AI 系统的基础构件。接下来,我们需要构建一个 API,供用户生成自己的图像。正如我们刚刚看到的,生成一张图像需要一些时间。正如我们在介绍中所说的,我们需要引入一个 Web 队列工作进程架构,使得这个系统既可靠又具有可扩展性。

创建 Dramatiq 工作进程并定义图像生成任务

正如我们在本章的介绍中提到的,直接在我们的 REST API 服务器上运行图像生成模型是不可行的。正如我们在上一节所见,这一操作可能需要几分钟,并消耗大量内存。为了解决这个问题,我们将定义一个独立于服务器进程的其他进程来处理图像生成任务:工作进程。本质上,工作进程可以是任何一个在后台执行任务的程序。

在 Web 开发中,这个概念通常意味着比这更多的内容。工作进程是一个持续运行在后台的进程,等待接收任务。这些任务通常由 Web 服务器发送,服务器会根据用户的操作请求执行特定的操作。

因此,我们可以看到,我们需要一个通信通道来连接 Web 服务器和工作进程。这就是队列的作用。队列会接收并堆积来自 Web 服务器的消息,然后将这些消息提供给工作进程读取。这就是 Web 队列工作进程架构。为了更好地理解这一点,图 14.4 展示了这种架构的示意图。

图 14.3 – Web 队列工作进程架构示意图

图 14.3 – Web 队列工作进程架构示意图

这是不是让你想起了什么?是的,这与我们在第八章中看到的非常相似,在处理多个 WebSocket 连接并广播消息这一节。实际上,这是同一个原理:我们通过一个中央数据源来解决有多个进程的问题。

这种架构的一个伟大特性是它非常容易扩展。试想你的应用程序取得了巨大成功,成千上万的用户想要生成图像:单个工作进程根本无法满足这种需求。事实上,我们所需要做的就是启动更多的工作进程。由于架构中有一个单独的消息代理,每个工作进程会在收到消息时进行拉取,从而实现任务的并行处理。它们甚至不需要位于同一台物理机器上。图 14.4 展示了这一点。

图 14.4 – 带有多个工作进程的 Web 队列工作进程架构

图 14.4 – 带有多个工作进程的 Web 队列工作进程架构

在 Python 中,有多种库可以帮助实现工作进程。它们提供了定义任务、将任务调度到队列中并运行进程、拉取并执行任务所需的工具。在本书中,我们将使用 Dramatiq,一个轻量级但强大且现代的后台任务处理库。正如我们在 第八章 中所做的,我们将使用 Redis 作为消息代理。

实现一个工作进程

和往常一样,我们首先安装所需的依赖项。运行以下命令:


(venv) $ pip install "dramatiq[redis]"

这将安装 Dramatiq,并安装与 Redis 代理通信所需的依赖项。

在一个最小的示例中,设置 Dramatiq 工作进程涉及两件事:

  1. 设置代理类型和 URL。

  2. 通过使用 @``dramatiq.actor 装饰器来定义任务。

它非常适合绝大多数任务,比如发送电子邮件或生成导出文件。

然而,在我们的案例中,我们需要加载庞大的 Stable Diffusion 模型。正如我们通常在 FastAPI 服务器中通过 startup 事件做的那样,我们希望只有在进程实际启动时才执行这一操作。为了使用 Dramatiq 实现这一点,我们需要实现一个中间件。它们允许我们在工作进程生命周期中的几个关键事件插入自定义逻辑,包括当工作进程启动时。

你可以在以下示例中看到我们自定义中间件的实现:

worker.py


class TextToImageMiddleware(Middleware):    def __init__(self) -> None:
        super().__init__()
        self.text_to_image = TextToImage()
    def after_process_boot(self, broker):
        self.text_to_image.load_model()
        return super().after_process_boot(broker)
text_to_image_middleware = TextToImageMiddleware()
redis_broker = RedisBroker(host="localhost")
redis_broker.add_middleware(text_to_image_middleware)
dramatiq.set_broker(redis_broker)

github.com/PacktPublishing/Building-Data-Science-Applications-with-FastAPI-Second-Edition/tree/main/chapter14/basic/worker.py

我们定义了一个 TextToImageMiddleware 类,它的作用是承载 TextToImage 的实例,这是我们在上一节中定义的图像生成服务。它继承自 Dramatiq 的 Middleware 类。这里的关键是 after_process_boot 方法。它是 Dramatiq 提供的事件钩子之一,允许我们插入自定义逻辑。在这里,我们告诉它在工作进程启动后加载 Stable Diffusion 模型。你可以在官方文档中查看支持的钩子列表:dramatiq.io/reference.html#middleware

接下来的几行代码让我们可以配置我们的工作进程。我们首先实例化我们自定义中间件的一个实例。然后,我们创建一个与我们选择的技术相对应的代理类;在我们的案例中是 Redis。在告诉 Dramatiq 使用它之前,我们需要将中间件添加到这个代理中。我们的工作进程现在已经完全配置好,可以连接到 Redis 代理,并在启动时加载我们的模型。

现在,让我们来看一下如何定义一个任务来生成图像:

worker.py


@dramatiq.actor()def text_to_image_task(
    prompt: str, *, negative_prompt: str | None = None, num_steps: int = 50
):
    image = text_to_image_middleware.text_to_image.generate(
        prompt, negative_prompt=negative_prompt, num_steps=num_steps
    )
    image.save(f"{uuid.uuid4()}.png")

github.com/PacktPublishing/Building-Data-Science-Applications-with-FastAPI-Second-Edition/tree/main/chapter14/basic/worker.py

实现是直接的:Dramatiq 任务实际上是我们用 @dramatiq.actor 装饰的普通函数。我们可以像定义其他函数一样定义参数。然而,这里有一个重要的陷阱需要避免:当我们从服务器调度任务时,参数将必须存储在队列存储中。因此,Dramatiq 会将参数内部序列化为 JSON。这意味着你的任务参数必须是可序列化的数据——你不能有任意的 Python 对象,比如类实例或函数。

函数体在将图像保存到磁盘之前,会调用我们在 text_to_image_middleware 中加载的 TextToImage 实例。为了避免文件覆盖,我们选择在这里生成一个UUID,即通用唯一标识符。它是一个大的随机字符串,保证每次生成时都是唯一的。凭借这个,我们可以安全地将其作为文件名,并确保它不会在磁盘上已存在。

这就是 worker 实现的内容。

启动 worker

我们还没有代码来调用它,但我们可以手动尝试。首先,确保你已经启动了一个 Redis 服务器,正如在技术要求部分中所解释的那样。然后,我们可以使用以下命令启动 Dramatiq worker:


(venv) $ dramatiq -p 1 -t 1 chapter14.basic.worker

Dramatiq 提供了命令行工具来启动 worker 进程。主要的位置参数是 worker 模块的点路径。这类似于我们在使用 Uvicorn 时的操作。我们还设置了两个可选参数,-p-t。它们控制 Dramatiq 启动的进程和线程的数量。默认情况下,它启动 10 个进程,每个进程有 8 个线程。这意味着将会有 80 个 worker 来拉取并执行任务。虽然这个默认配置适合常见需求,但由于两个原因,它不适用于我们的 Stable Diffusion 模型:

  • 进程中的每个线程共享相同的内存空间。这意味着,如果两个(或更多)线程尝试生成图像,它们将对内存中的同一对象进行读写操作。对于我们的模型来说,这会导致并发问题。我们说它是非线程安全的。因此,每个进程应该仅启动一个线程:这就是-t 1选项的意义所在。

  • 每个进程都应该将模型加载到内存中。这意味着,如果我们启动 8 个进程,我们将加载 8 次模型。正如我们之前所看到的,它需要相当大的内存,所以这样做可能会使你的计算机内存爆炸。为了安全起见,我们仅启动一个进程,使用-p 1选项。如果你想尝试并行化并查看我们的 worker 能否并行生成两张图像,你可以尝试-p 2来启动两个进程。但要确保你的计算机能够处理!

如果你运行前面的命令,你应该会看到类似这样的输出:


[2023-02-02 08:52:11,479] [PID 44348] [MainThread] [dramatiq.MainProcess] [INFO] Dramatiq '1.13.0' is booting up.Fetching 19 files:   0%|          | 0/19 [00:00<?, ?it/s]
Fetching 19 files: 100%|██████████| 19/19 [00:00<00:00, 13990.83it/s]
[2023-02-02 08:52:11,477] [PID 44350] [MainThread] [dramatiq.WorkerProcess(0)] [INFO] Worker process is ready for action.
[2023-02-02 08:52:11,578] [PID 44355] [MainThread] [dramatiq.ForkProcess(0)] [INFO] Fork process 'dramatiq.middleware.prometheus:_run_exposition_server' is ready for action.

你可以通过查看 Stable Diffusion 流水线的输出,检查模型文件是否已经下载,直到 worker 完全启动。这意味着它已经正确加载。

在 worker 中调度任务

现在我们可以尝试在工作线程中调度任务了。为此,我们可以启动一个 Python 交互式 Shell 并导入task函数。打开一个新的命令行并运行以下命令(确保你已启用 Python 虚拟环境):


(venv) $ python>>> from chapter14.basic.worker import text_to_image_task
>>> text_to_image_task.send("A Renaissance castle in the Loire Valley")
Message(queue_name='default', actor_name='text_to_image_task', args=('A Renaissance castle in the Loire Valley',), kwargs={}, options={'redis_message_id': '663df44a-cfc1-4f13-8457-05d8181290c1'}, message_id='bf57d112-6c20-49bc-a926-682ca43ea7ea', message_timestamp=1675324585644)

就是这样——我们在工作线程中安排了一个任务!注意我们在task函数上使用了send方法,而不是直接调用它:这是告诉 Dramatiq 将其发送到队列中的方式。

如果你回到工作线程终端,你会看到 Stable Diffusion 正在生成图像。过一会儿,你的图像将保存在磁盘上。你还可以尝试在短时间内连续发送两个任务。你会发现 Dramatiq 会一个接一个地处理它们。

干得好!我们的后台进程已经准备好,甚至能够在其中调度任务。下一步就是实现 REST API,以便用户可以自己请求图像生成。

实现 REST API

要在工作线程中调度任务,我们需要一个用户可以交互的安全接口。REST API 是一个不错的选择,因为它可以轻松集成到任何软件中,如网站或移动应用。在这一节中,我们将快速回顾一下我们实现的简单 API 端点,用于将图像生成任务发送到队列中。以下是实现代码:

api.py


class ImageGenerationInput(BaseModel):    prompt: str
    negative_prompt: str | None
    num_steps: int = Field(50, gt=0, le=50)
class ImageGenerationOutput(BaseModel):
    task_id: UUID4
app = FastAPI()
@app.post(
    "/image-generation",
    response_model=ImageGenerationOutput,
    status_code=status.HTTP_202_ACCEPTED,
)
async def post_image_generation(input: ImageGenerationInput) -> ImageGenerationOutput:
    task: Message = text_to_image_task.send(
        input.prompt, negative_prompt=input.negative_prompt, num_steps=input.num_steps
    )
    return ImageGenerationOutput(task_id=task.message_id)

github.com/PacktPublishing/Building-Data-Science-Applications-with-FastAPI-Second-Edition/tree/main/chapter14/basic/api.py

如果你从这本书的开头一直跟到现在,这不应该让你感到惊讶。我们已经妥善地定义了合适的 Pydantic 模型来构建和验证端点负载。然后,这些数据会直接用于发送任务到 Dramatiq,正如我们在前一节看到的那样。

在这个简单的实现中,输出仅包含消息 ID,Dramatiq 会自动为每个任务分配这个 ID。注意我们将 HTTP 状态码设置为202,表示已接受。从语义上讲,这意味着服务器已理解并接受了请求,但处理尚未完成,甚至可能还没有开始。它专门用于处理在后台进行的情况,这正是我们在这里的情况。

如果你同时启动工作线程和这个 API,你将能够通过 HTTP 调用触发图像生成。

你可能在想:这不错……但是用户怎么才能获取结果呢?他们怎么知道任务是否完成? 你说得对——我们完全没有讨论这个问题!实际上,这里有两个方面需要解决:我们如何跟踪待处理任务及其执行情况?我们如何存储并提供生成的图像?这就是下一节的内容。

将结果存储在数据库和对象存储中

在上一节中,我们展示了如何实现一个后台工作程序来执行繁重的计算,以及一个 API 来调度任务给这个工作程序。然而,我们仍然缺少两个重要方面:用户没有任何方式了解任务的进度,也无法获取最终结果。让我们来解决这个问题!

在工作程序和 API 之间共享数据

正如我们所见,工作程序是一个在后台运行的程序,执行 API 请求它做的计算。然而,工作程序并没有与 API 服务器通信的任何方式。这是预期中的:因为可能有任意数量的服务器进程,且它们甚至可能运行在不同的物理服务器上,因此进程之间不能直接通信。始终是同样的问题:需要有一个中央数据源,供进程写入和读取数据。

事实上,解决 API 和工作程序之间缺乏通信的第一种方法是使用我们用来调度任务的相同代理:工作程序可以将结果写入代理,API 可以从中读取。这在大多数后台任务库中都是可能的,包括 Dramatiq。然而,这个解决方案有一些局限性,其中最主要的是我们能保留数据的时间有限。像 Redis 这样的代理并不适合长时间可靠地存储数据。在某些时候,我们需要删除最古老的数据以限制内存使用。

然而,我们已经知道有一些东西能够高效地存储结构化数据:当然是数据库!这就是我们在这里展示的方法。通过拥有一个中央数据库,我们可以在其中存储图像生成请求和结果,这样就能在工作程序和 API 之间共享信息。为此,我们将重用我们在《第六章》的使用 SQLAlchemy ORM 与 SQL 数据库通信部分中展示的很多技巧。我们开始吧!

定义一个 SQLAlchemy 模型

第一步是定义一个 SQLAlchemy 模型来存储单个图像生成任务。你可以如下所示查看它:

models.py


class GeneratedImage(Base):    __tablename__ = "generated_images"
    id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
    created_at: Mapped[datetime] = mapped_column(
        DateTime, nullable=False, default=datetime.now
    )
    progress: Mapped[int] = mapped_column(Integer, nullable=False, default=0)
    prompt: Mapped[str] = mapped_column(Text, nullable=False)
    negative_prompt: Mapped[str | None] = mapped_column(Text, nullable=True)
    num_steps: Mapped[int] = mapped_column(Integer, nullable=False)
    file_name: Mapped[str | None] = mapped_column(String(255), nullable=True)

github.com/PacktPublishing/Building-Data-Science-Applications-with-FastAPI-Second-Edition/tree/main/chapter14/complete/models.py

像往常一样,我们定义一个自增的 ID 作为主键。我们还添加了 promptnegative_promptnum_steps 列,这些列对应我们传递给工作程序任务的参数。这样,我们就可以直接将 ID 传递给工作程序,它会直接从对象中获取参数。此外,这还允许我们存储并记住用于特定生成的参数。

progress 列是一个整数,用来存储当前生成任务的进度。

最后,file_name 将存储我们在系统中保存的实际文件名。我们将在下一节中关于对象存储的部分看到如何使用它。

将 API 调整为在数据库中保存图像生成任务

有了这个模型后,我们对 API 中图像生成任务的调度方式稍微做了些调整。我们不再直接将任务发送给工作进程,而是首先在数据库中创建一行数据,并将该对象的 ID 作为输入传递给工作进程任务。端点的实现如下所示:

api.py


@app.post(    "/generated-images",
    response_model=schemas.GeneratedImageRead,
    status_code=status.HTTP_201_CREATED,
)
async def create_generated_image(
    generated_image_create: schemas.GeneratedImageCreate,
    session: AsyncSession = Depends(get_async_session),
)    GeneratedImage:
    image = GeneratedImage(**generated_image_create.dict())
    session.add(image)
    await session.commit()
    text_to_image_task.send(image.id)
    return image

github.com/PacktPublishing/Building-Data-Science-Applications-with-FastAPI-Second-Edition/tree/main/chapter14/complete/api.py

我们不会深入讨论如何使用 SQLAlchemy ORM 在数据库中创建对象。如果你需要复习,可以参考《第六章》中的使用 SQLAlchemy ORM 与 SQL 数据库通信部分。

在这个代码片段中,主要需要注意的是我们将新创建对象的 ID 作为text_to_image_task的参数传递。正如我们稍后看到的,工作进程会从数据库中重新读取这个 ID,以检索生成参数。

该端点的响应仅仅是我们GeneratedImage模型的表示,使用了 Pydantic 架构GeneratedImageRead。因此,用户将会收到类似这样的响应:


{    "created_at": "2023-02-07T10:17:50.992822",
    "file_name": null,
    "id": 6,
    "negative_prompt": null,
    "num_steps": 50,
    "progress": 0,
    "prompt": "a sunset over a beach"
}

它展示了我们在请求中提供的提示,最重要的是,它给了一个 ID。这意味着用户将能够再次查询此特定请求以检索数据,并查看是否完成。这就是下面定义的get_generated_image端点的目的。我们不会在这里展示它,但你可以在示例仓库中阅读到它。

将工作进程调整为从数据库中读取和更新图像生成任务

你可能已经猜到,我们需要改变任务的实现,以便它能从数据库中检索对象,而不是直接读取参数。让我们一步步来进行调整。

我们做的第一件事是使用在任务参数中获得的 ID 从数据库中检索一个GeneratedImage

worker.py


@dramatiq.actor()def text_to_image_task(image_id: int):
    image = get_image(image_id)

github.com/PacktPublishing/Building-Data-Science-Applications-with-FastAPI-Second-Edition/tree/main/chapter14/complete/worker.py

为了实现这一点,你会看到我们使用了一个名为get_image的辅助函数。它定义在任务的上方。让我们来看一下:

worker.py


def get_image(id: int) -> GeneratedImage:    async def _get_image(id: int) -> GeneratedImage:
        async with async_session_maker() as session:
            select_query = select(GeneratedImage).where(GeneratedImage.id == id)
            result = await session.execute(select_query)
            image = result.scalar_one_or_none()
            if image is None:
                raise Exception("Image does not exist")
            return image
    return asyncio.run(_get_image(id))

github.com/PacktPublishing/Building-Data-Science-Applications-with-FastAPI-Second-Edition/tree/main/chapter14/complete/worker.py

它看起来可能有些奇怪,但实际上,你已经对其大部分逻辑非常熟悉了。如果你仔细观察,你会发现它定义了一个嵌套的私有函数,在其中我们定义了实际的逻辑来使用 SQLAlchemy ORM 获取和保存对象。请注意,它是异步的,并且我们在其中大量使用了异步 I/O 模式,正如本书中所展示的那样。

这正是我们需要像这样的辅助函数的原因。事实上,Dramatiq 并未原生设计为运行异步函数,因此我们需要手动使用asyncio.run来调度其执行。我们已经在第二章中看到过这个函数,那里介绍了异步 I/O。它的作用是运行异步函数并返回其结果。这就是我们如何在任务中同步调用包装函数而不出现任何问题。

其他方法也可以解决异步 I/O 问题。

我们在这里展示的方法是解决异步工作者问题最直接且稳健的方法。

另一种方法可能是为 Dramatiq 设置装饰器或中间件,使其能够原生支持运行异步函数,但这种方法复杂且容易出现 BUG。

我们也可以考虑拥有另一个同步工作的 SQLAlchemy 引擎和会话生成器。然而,这会导致代码中出现大量重复的内容。而且,如果我们有除了 SQLAlchemy 之外的其他异步函数,这也无法提供帮助。

现在,让我们回到text_to_image_task的实现:

worker.py


@dramatiq.actor()def text_to_image_task(image_id: int):
    image = get_image(image_id)
    def callback(step: int, _timestep, _tensor):
        update_progress(image, step)

github.com/PacktPublishing/Building-Data-Science-Applications-with-FastAPI-Second-Edition/tree/main/chapter14/complete/worker.py

我们为 Stable Diffusion 管道定义了一个callback函数。它的作用是将当前的进度保存到数据库中,针对当前的GeneratedImage。为此,我们再次使用了一个辅助函数update_progress

worker.py


def update_progress(image: GeneratedImage, step: int):    async def _update_progress(image: GeneratedImage, step: int):
        async with async_session_maker() as session:
            image.progress = int((step / image.num_steps) * 100)
            session.add(image)
            await session.commit()
    asyncio.run(_update_progress(image, step))

github.com/PacktPublishing/Building-Data-Science-Applications-with-FastAPI-Second-Edition/tree/main/chapter14/complete/worker.py

我们使用与get_image相同的方法来包装异步函数。

回到text_to_image_task,我们现在可以调用我们的TextToImage模型来生成图像。这与前一节中展示的调用完全相同。唯一的区别是,我们从image对象中获取参数。我们还使用 UUID 生成一个随机的文件名:

worker.py


    image_output = text_to_image_middleware.text_to_image.generate(        image.prompt,
        negative_prompt=image.negative_prompt,
        num_steps=image.num_steps,
        callback=callback,
    )
    file_name = f"{uuid.uuid4()}.png"

github.com/PacktPublishing/Building-Data-Science-Applications-with-FastAPI-Second-Edition/tree/main/chapter14/complete/worker.py

以下部分用于将图像上传到对象存储。我们将在下一部分中更详细地解释这一点:

worker.py


    storage = Storage()    storage.upload_image(image_output, file_name, settings.storage_bucket)

github.com/PacktPublishing/Building-Data-Science-Applications-with-FastAPI-Second-Edition/tree/main/chapter14/complete/worker.py

最后,我们调用另一个辅助函数update_file_name,将随机文件名保存到数据库中。它将允许我们为用户检索该文件:

worker.py


    update_file_name(image, file_name)

github.com/PacktPublishing/Building-Data-Science-Applications-with-FastAPI-Second-Edition/tree/main/chapter14/complete/worker.py

如你所见,这个实现的重点是我们从数据库中读取和写入GeneratedImage的信息。这就是我们如何在 API 服务器和工作进程之间进行同步。工作进程的部分就到此为止!有了这个逻辑,我们就可以从 API 调度一个图像生成任务,而工作进程则能够在设置最终文件名之前定期更新任务进度。因此,通过 API,一个简单的GET请求就能让我们看到任务的状态。

在对象存储中存储和服务文件

我们必须解决的最后一个挑战是关于存储生成的图像。我们需要一种可靠的方式来存储它们,同时让用户能够轻松地从互联网中检索它们。

传统上,Web 应用程序的处理方式非常简单。它们将文件直接存储在服务器的硬盘中,在指定的目录下,并配置其 Web 服务器,当访问某个 URL 时提供这些文件。这实际上是我们在第十三章中的 WebSocket 示例中做的:我们使用了 StaticFiles 中间件来静态地提供我们磁盘上的 JavaScript 脚本。

虽然这种方式适用于静态文件,比如每个服务器都有自己副本的 JavaScript 或 CSS 文件,但对于用户上传或后台生成的动态文件来说并不合适,尤其是在多个进程运行在不同物理机器上的复杂架构中。问题再次出现,即不同进程读取的中央数据源问题。在前面的部分,我们看到消息代理和数据库可以在多个场景中解决这个问题。而对于任意的二进制文件,无论是图像、视频还是简单的文本文件,我们需要其他解决方案。让我们来介绍对象存储

对象存储与我们日常在计算机中使用的标准文件存储有所不同,后者中的磁盘是以目录和文件的层次结构组织的。而对象存储将每个文件作为一个对象进行存储,其中包含实际数据及其所有元数据,如文件名、大小、类型和唯一 ID。这种概念化的主要好处是,它更容易将这些文件分布到多个物理机器上:我们可以将数十亿个文件存储在同一个对象存储中。从用户的角度来看,我们只需请求一个特定的文件,存储系统会负责从实际的物理磁盘加载该文件。

在云时代,这种方法显然获得了广泛的关注。2006 年,亚马逊网络服务AWS)推出了其自有实现的对象存储——Amazon S3。它为开发人员提供了几乎无限的磁盘空间,允许通过一个简单的 API 存储文件,并且价格非常低廉。Amazon S3 因其广泛的流行,其 API 成为行业事实上的标准。如今,大多数云对象存储,包括微软 Azure 或 Google Cloud 等竞争对手的存储,都与 S3 API 兼容。开源实现也应运而生,如 MinIO。这个通用的 S3 API 的主要好处是,您可以在项目中使用相同的代码和库与任何对象存储提供商进行交互,并在需要时轻松切换。

总结一下,对象存储是一种非常方便的方式,用于大规模存储和提供文件,无论有多少个进程需要访问这些数据。在本节结束时,我们项目的全球架构将像图 14.5中所示。

图 14.5 – Web-队列-工作者架构和对象存储

图 14.5 – Web-队列-工作者架构和对象存储

值得注意的是,对象存储会直接将文件提供给用户。不会有一个端点,服务器在从对象存储下载文件后再将其发送给用户。以这种方式操作并没有太大好处,即使在认证方面也是如此。我们将看到,兼容 S3 的存储具有内建的机制来保护文件不被未授权访问。

实现一个对象存储助手

那么我们开始写代码吧!我们将使用 MinIO 的 Python 客户端库,这是一个与任何兼容 S3 的存储进行交互的库。让我们先安装它:


(venv) $ pip install minio

我们现在可以实现一个类,以便手头有我们需要的所有操作。我们先从初始化器开始:

storage.py


class Storage:    def __init__(self) -> None:
        self.client = Minio(
            settings.storage_endpoint,
            access_key=settings.storage_access_key,
            secret_key=settings.storage_secret_key,
        )

github.com/PacktPublishing/Building-Data-Science-Applications-with-FastAPI-Second-Edition/tree/main/chapter14/complete/storage.py

在该类的初始化函数中,我们创建了一个 Minio 客户端实例。你会看到我们使用一个 settings 对象来提取存储 URL 和凭证。因此,使用环境变量就能非常轻松地切换它们。

然后我们将实现一些方法,帮助我们处理对象存储。第一个方法是 ensure_bucket

storage.py


    def ensure_bucket(self, bucket_name: str):        bucket_exists = self.client.bucket_exists(bucket_name)
        if not bucket_exists:
            self.client.make_bucket(bucket_name)

github.com/PacktPublishing/Building-Data-Science-Applications-with-FastAPI-Second-Edition/tree/main/chapter14/complete/storage.py

该方法的作用是确保在我们的对象存储中创建了正确的存储桶。在 S3 实现中,存储桶就像是你拥有的文件夹,你可以将文件存储在其中。你上传的每个文件都必须放入一个现有的存储桶中。

然后,我们定义了 upload_image

storage.py


    def upload_image(self, image: Image, object_name: str, bucket_name: str):        self.ensure_bucket(bucket_name)
        image_data = io.BytesIO()
        image.save(image_data, format="PNG")
        image_data.seek(0)
        image_data_length = len(image_data.getvalue())
        self.client.put_object(
            bucket_name,
            object_name,
            image_data,
            length=image_data_length,
            content_type="image/png",
        )

github.com/PacktPublishing/Building-Data-Science-Applications-with-FastAPI-Second-Edition/tree/main/chapter14/complete/storage.py

这是用于将图像上传到存储的。为了简化操作,该方法接受一个 Pillow Image 对象,因为这是我们在 Stable Diffusion 流水线的最后得到的结果。我们实现了一些逻辑,将这个 Image 对象转换为适合 S3 上传的原始字节流。该方法还期望接收 object_name,即存储中实际的文件名,以及 bucket_name。请注意,我们首先确保存储桶已经正确创建,然后再尝试上传文件。

最后,我们添加了 get_presigned_url 方法:

storage.py


    def get_presigned_url(        self,
        object_name: str,
        bucket_name: str,
        *,
        expires: timedelta = timedelta(days=7)
    )    str:
        return self.client.presigned_get_object(
            bucket_name, object_name, expires=expires
        )

github.com/PacktPublishing/Building-Data-Science-Applications-with-FastAPI-Second-Edition/tree/main/chapter14/complete/storage.py

这种方法将帮助我们安全地将文件提供给用户。出于安全原因,S3 存储中的文件默认对互联网用户不可访问。为了给予文件访问权限,我们可以执行以下任一操作:

  • 将文件设置为公开状态,这样任何拥有 URL 的人都能访问它。这个适合公开文件,但对于私密的用户文件则不适用。

  • 生成一个带有临时访问密钥的 URL。这样,我们就可以将文件访问权限提供给用户,即使 URL 被窃取,访问也会在一段时间后被撤销。这带来的巨大好处是,URL 生成发生在我们的 API 服务器上,使用 S3 客户端。因此,在生成文件 URL 之前,我们可以根据自己的逻辑检查用户是否通过身份验证,并且是否有权访问特定的文件。这就是我们在这里采用的方法,并且此方法会在特定存储桶中的特定文件上生成预签名 URL,且有效期为一定时间。

如你所见,我们的类只是 MinIO 客户端的一个薄包装。现在我们要做的就是用它来上传图像并从 API 获取预签名 URL。

在工作者中使用对象存储助手

在上一节中,我们展示了任务实现中的以下几行代码:

worker.py


    storage = Storage()    storage.upload_image(image_output, file_name, settings.storage_bucket)

github.com/PacktPublishing/Building-Data-Science-Applications-with-FastAPI-Second-Edition/tree/main/chapter14/complete/worker.py

现在我们已经谈到了 Storage 类,你应该能猜到我们在这里做的事情:我们获取生成的图像及其随机名称,并将其上传到 settings 中定义的存储桶。就这样!

在服务器上生成预签名 URL

在 API 端,我们实现了一个新端点,角色是返回给定 GeneratedImage 的预签名 URL:

server.py


@app.get("/generated-images/{id}/url")async def get_generated_image_url(
    image: GeneratedImage = Depends(get_generated_image_or_404),
    storage: Storage = Depends(get_storage),
)    schemas.GeneratedImageURL:
    if image.file_name is None:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail="Image is not available yet. Please try again later.",
        )
    url = storage.get_presigned_url(image.file_name, settings.storage_bucket)
    return schemas.GeneratedImageURL(url=url)

github.com/PacktPublishing/Building-Data-Science-Applications-with-FastAPI-Second-Edition/tree/main/chapter14/complete/server.py

在生成 URL 之前,我们首先检查 GeneratedImage 对象上是否设置了 file_name 属性。如果没有,意味着工作者任务尚未完成。如果有,我们就可以继续调用 Storage 类的 get_presigned_url 方法。

请注意,我们已经定义了依赖注入来获取 Storage 实例。正如本书中所展示的那样,在处理外部服务时,FastAPI 中使用依赖是一个非常好的实践。

好的,看来我们一切准备就绪!让我们看看它如何运行。

运行图像生成系统

首先,我们需要为项目填充环境变量,特别是数据库 URL 和 S3 凭据。为了简化,我们将使用一个简单的 SQLite 数据库和 MinIO 的示例平台作为 S3 存储。MinIO 是一个免费的开源对象存储平台,非常适合示例和玩具项目。当进入生产环境时,你可以轻松切换到任何兼容 S3 的提供商。让我们在项目根目录下创建一个 .env 文件:


DATABASE_URL=sqlite+aiosqlite:///chapter14.dbSTORAGE_ENDPOINT=play.min.io
STORAGE_ACCESS_KEY=Q3AM3UQ867SPQQA43P2F
STORAGE_SECRET_KEY=zuf+tfteSlswRu7BJ86wekitnifILbZam1KYY3TG
STORAGE_BUCKET=fastapi-book-text-to-image

存储端点、访问密钥和秘钥是 MinIO 演示环境的参数。确保查看它们的官方文档,以了解自我们编写本书以来是否有所更改:min.io/docs/minio/linux/developers/python/minio-py.html#id5

我们的Settings类将自动加载此文件,以填充我们在代码中使用的设置。如果你需要复习这一概念,确保查看第十章中的设置和使用环境变量部分。

现在我们可以运行系统了。确保你的 Redis 服务器仍在运行,正如在技术要求部分所解释的那样。首先,让我们启动 FastAPI 服务器:


(venv) $ uvicorn chapter14.complete.api:app

然后,启动工作进程:


(venv) $ dramatiq -p 1 -t 1 chapter14.complete.worker

堆栈现在已准备好生成图像。让我们使用 HTTPie 发起请求,开始一个新的任务:


$ http POST http://localhost:8000/generated-images prompt="a sunset over a beach"HTTP/1.1 201 Created
content-length: 151
content-type: application/json
date: Mon, 13 Feb 2023 07:24:44 GMT
server: uvicorn
{
    "created_at": "2023-02-13T08:24:45.954240",
    "file_name": null,
    "id": 1,
    "negative_prompt": null,
    "num_steps": 50,
    "progress": 0,
    "prompt": "a sunset over a beach"
}

一个新的GeneratedImage已在数据库中创建,分配的 ID 为1。进度为0%;处理尚未开始。让我们尝试通过 API 查询它:


http GET http://localhost:8000/generated-images/1HTTP/1.1 200 OK
content-length: 152
content-type: application/json
date: Mon, 13 Feb 2023 07:25:04 GMT
server: uvicorn
{
    "created_at": "2023-02-13T08:24:45.954240",
    "file_name": null,
    "id": 1,
    "negative_prompt": null,
    "num_steps": 50,
    "progress": 36,
    "prompt": "a sunset over a beach"
}

API 返回相同的对象及其所有属性。注意,进度已更新,现在为36%。过一会儿,我们可以再次尝试相同的请求:


$ http GET http://localhost:8000/generated-images/1HTTP/1.1 200 OK
content-length: 191
content-type: application/json
date: Mon, 13 Feb 2023 07:25:34 GMT
server: uvicorn
{
    "created_at": "2023-02-13T08:24:45.954240",
    "file_name": "affeec65-5d9b-480e-ac08-000c74e22dc9.png",
    "id": 1,
    "negative_prompt": null,
    "num_steps": 50,
    "progress": 100,
    "prompt": "a sunset over a beach"
}

这次,进度为100%,文件名已经填写。图像准备好了!现在我们可以请求 API 为该图像生成一个预签名 URL:


$ http GET http://localhost:8000/generated-images/1/urlHTTP/1.1 200 OK
content-length: 366
content-type: application/json
date: Mon, 13 Feb 2023 07:29:53 GMT
server: uvicorn
{
    "url": "https://play.min.io/fastapi-book-text-to-image/affeec65-5d9b-480e-ac08-000c74e22dc9.png?X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Credential=Q3AM3UQ867SPQQA43P2F%2F20230213%2Fus-east-1%2Fs3%2Faws4_request&X-Amz-Date=20230213T072954Z&X-Amz-Expires=604800&X-Amz-SignedHeaders=host&X-Amz-Signature=6ffddb81702bed6aac50786578eb75af3c1f6a3db28e4990467c973cb3b457a9"
}

我们在 MinIO 服务器上得到了一个非常长的 URL。如果你在浏览器中打开它,你会看到刚刚由我们的系统生成的图像,如图 14.6所示。

图 14.6 – 生成的图像托管在对象存储中

图 14.6 – 生成的图像托管在对象存储中

很不错,不是吗?我们现在拥有一个功能齐全的系统,用户能够执行以下操作:

  • 请求根据他们自己的提示和参数生成图像

  • 获取请求进度的信息

  • 从可靠存储中获取生成的图像

我们在这里看到的架构已经可以在具有多台机器的云环境中部署。通常,我们可能会有一台标准的便宜服务器来提供 API 服务,而另一台则是更昂贵的服务器,配有专用 GPU 和充足的 RAM 来运行工作进程。代码无需更改就可以处理这种部署,因为进程间的通信是由中央元素——消息代理、数据库和对象存储来处理的。

总结

太棒了!你可能还没有意识到,但在这一章中,你已经学习了如何架构和实现一个非常复杂的机器学习系统,它能与你在外面看到的现有图像生成服务相媲美。我们在这里展示的概念是至关重要的,且是所有你能想象的分布式系统的核心,无论它们是设计用来运行机器学习模型、提取管道,还是数学计算。通过使用像 FastAPI 和 Dramatiq 这样的现代工具,你将能够在短时间内用最少的代码实现这种架构,最终得到一个非常快速且稳健的结果。

我们的旅程即将结束。在让你用 FastAPI 开始自己的冒险之前,我们将研究构建数据科学应用程序时的最后一个重要方面:日志记录和监控。

第十五章:监控数据科学系统的健康状况和性能

在本章中,我们将深入探讨,以便你能够构建稳健、适用于生产的系统。实现这一目标的最重要方面之一是拥有所有必要的数据,以确保系统正确运行,并尽早检测到问题,以便采取纠正措施。在本章中,我们将展示如何设置适当的日志设施,以及如何实时监控我们软件的性能和健康状况。

我们即将结束 FastAPI 数据科学之旅。到目前为止,我们主要关注的是我们实现的程序的功能。然而,还有一个方面常常被开发者忽视,但实际上非常重要:评估系统是否在生产环境中正确且可靠地运行,并在系统出现问题时尽早收到警告。

为此,存在许多工具和技术,我们可以收集尽可能多的数据,了解我们的程序如何运行。这就是我们在本章中要回顾的内容。

我们将涵盖以下主要主题:

  • 配置并使用 Loguru 日志设施

  • 配置 Prometheus 指标并在 Grafana 中监控它们

  • 配置 Sentry 用于报告错误

技术要求

对于本章,你将需要一个 Python 虚拟环境,就像我们在第一章中设置的那样,Python 开发环境设置

要运行 Dramatiq 工作程序,你需要在本地计算机上运行 Redis 服务器。最简单的方式是将其作为 Docker 容器运行。如果你以前从未使用过 Docker,我们建议你阅读官方文档中的入门教程docs.docker.com/get-started/。完成后,你可以通过以下简单命令运行 Redis 服务器:


$ docker run -d --name worker-redis -p 6379:6379 redis

你可以在专门的 GitHub 仓库中找到本章的所有代码示例:github.com/PacktPublishing/Building-Data-Science-Applications-with-FastAPI-Second-Edition/tree/main/chapter15

关于截图的说明

在本章中,我们将展示一些截图,特别是 Grafana 界面的截图。它们的目的是帮助你了解界面的整体布局,帮助你识别不同的部分。如果你在阅读实际内容时遇到困难,不用担心:周围的解释将帮助你找到需要关注的地方并了解该与哪些部分交互。

配置并使用 Loguru 日志设施

在软件开发中,日志可能是控制系统行为最简单但最强大的方式。它们通常由程序中特定位置打印的纯文本行组成。通过按时间顺序阅读这些日志,我们可以追踪程序的行为,确保一切顺利进行。实际上,在本书中我们已经看到过日志行。当你使用 Uvicorn 运行 FastAPI 应用并发出一些请求时,你会在控制台输出中看到这些日志行:


INFO:     Started server process [94918]INFO:     Waiting for application startup.
INFO:     Application startup complete.
INFO:     Uvicorn running on http://127.0.0.1:8000 (Press CTRL+C to quit)
INFO:     127.0.0.1:60736 - "POST /generated-images HTTP/1.1" 201 Created

这些是 Uvicorn 生成的日志,告诉我们它何时启动,以及何时处理了一个请求。正如你所见,日志可以帮助我们了解程序发生了什么以及执行了哪些操作。它们还可以告诉我们何时出现问题,这可能是一个需要解决的 bug。

理解日志级别

请注意,在每个日志行之前,我们都有INFO关键字。这就是我们所说的日志级别。它是分类日志重要性的方式。一般来说,定义了以下几种级别:

  • DEBUG

  • INFO

  • WARNING

  • ERROR

你可以将其视为重要性等级DEBUG是关于程序执行的非常具体的信息,这有助于调试代码,而ERROR意味着程序中发生了问题,可能需要你采取行动。关于这些级别的好处是,我们可以配置日志记录器应输出的最小级别。即使日志函数调用仍然存在于代码中,如果它不符合最小级别,日志记录器也会忽略它。

通常,我们可以在本地开发中设置DEBUG级别,这样可以获取所有信息以帮助我们开发和修复程序。另一方面,我们可以在生产环境中将级别设置为INFOWARNING,以便只获取最重要的消息。

使用 Loguru 添加日志

使用标准库中提供的logging模块,向 Python 程序添加日志非常容易。你可以像这样做:


>>> import logging>>> logging.warning("This is my log")
WARNING:root:This is my log

如你所见,这只是一个带有字符串参数的函数调用。通常,日志模块将不同的级别作为方法暴露,就像这里的warning一样。

标准的logging模块非常强大,允许你精细定制日志的处理、打印和格式。如果你浏览官方文档中的日志教程,docs.python.org/3/howto/logging.html,你会发现它很快会变得非常复杂,甚至对于简单的情况也是如此。

这就是为什么 Python 开发者通常使用封装了logging模块并提供更友好函数和接口的库。在本章中,我们将回顾如何使用和配置Loguru,一种现代而简单的日志处理方法。

和往常一样,首先需要在我们的 Python 环境中安装它:


(venv) $ pip install loguru

我们可以立即在 Python shell 中尝试:


>>> from loguru import logger>>> logger.debug("This is my log!")
2023-02-21 08:44:00.168 | DEBUG    | __main__:<module>:1 - This is my log!

你可能会认为这与我们使用标准的 logging 模块没什么不同。然而,注意到生成的日志已经包含了时间戳、级别以及函数调用的位置。这就是 Loguru 的主要优势之一:它自带合理的默认设置,开箱即用。

让我们在一个更完整的脚本中看看它的实际效果。我们将定义一个简单的函数,检查一个整数 n 是否为奇数。我们将添加一行调试日志,让我们知道函数开始执行逻辑。然后,在计算结果之前,我们将首先检查 n 是否确实是一个整数,如果不是,就记录一个错误。这个函数的实现如下:

chapter15_logs_01.py


from loguru import loggerdef is_even(n) -> bool:
    logger.debug("Check if {n} is even", n=n)
    if not isinstance(n, int):
        logger.error("{n} is not an integer", n=n)
        raise TypeError()
    return n % 2 == 0
if __name__ == "__main__":
    is_even(2)
    is_even("hello")

github.com/PacktPublishing/Building-Data-Science-Applications-with-FastAPI-Second-Edition/tree/main/chapter15/chapter15_logs_01.py

如你所见,它的使用非常简单:我们只需要导入 logger 并在需要记录日志的地方调用它。还注意到我们如何可以添加变量来格式化字符串:只需要在字符串中添加大括号内的占位符,然后通过关键字参数将每个占位符映射到其值。这个语法实际上类似于标准的 str.format 方法。你可以在官方的 Python 文档中了解更多内容:docs.python.org/fr/3/library/stdtypes.html#str.format

如果我们运行这个简单的脚本,我们将在控制台输出中看到我们的日志行:


(venv) $ python chapter15/chapter15_logs_01.py2023-03-03 08:16:40.145 | DEBUG    | __main__:is_even:5 - Check if 2 is even
2023-03-03 08:16:40.145 | DEBUG    | __main__:is_even:5 - Check if hello is even
2023-03-03 08:16:40.145 | ERROR    | __main__:is_even:7 - hello is not an integer
Traceback (most recent call last):
  File "/Users/fvoron/Development/Building-Data-Science-Applications-with-FastAPI-Second-Edition/chapter15/chapter15_logs_01.py", line 14, in <module>
    is_even("hello")
  File "/Users/fvoron/Development/Building-Data-Science-Applications-with-FastAPI-Second-Edition/chapter15/chapter15_logs_01.py", line 8, in is_even
    raise TypeError()
TypeError

我们的日志行在实际抛出异常之前已经正确添加到输出中。注意,Loguru 能够准确告诉我们日志调用来自代码的哪个位置:我们有函数名和行号。

理解和配置 sinks

我们已经看到,默认情况下,日志会添加到控制台输出。默认情况下,Loguru 定义了一个指向标准错误的sink。Sink 是 Loguru 引入的一个概念,用于定义日志行应该如何由日志记录器处理。我们不限于控制台输出:我们还可以将它们保存到文件、数据库,甚至发送到 Web 服务!

好的一点是,你并不只限于使用一个 sink;你可以根据需要使用多个!然后,每个日志调用都会通过每个 sink 进行处理。你可以在图 15.1中看到这种方法的示意图。

图 15.1 – Loguru sinks 架构

图 15.1 – Loguru sinks 架构

每个sink 都与 一个日志级别 相关联。这意味着我们可以根据 sink 使用不同的日志级别。例如,我们可以选择将所有日志输出到文件中,并且只在控制台保留最重要的警告和错误日志。我们再次以之前的示例为例,使用这种方式配置 Loguru:

chapter15_logs_02.py


logger.remove()logger.add(sys.stdout, level="WARNING")
logger.add("file.log", level="DEBUG", rotation="1 day")

github.com/PacktPublishing/Building-Data-Science-Applications-with-FastAPI-Second-Edition/tree/main/chapter15/chapter15_logs_02.py

loggerremove方法有助于删除先前定义的接收器。当这样调用时,没有参数传递,所有定义的接收器都会被移除。通过这样做,我们可以从没有默认接收器的全新状态开始。

接着,我们调用add来定义新的接收器。第一个参数,像sys.stdout或这里的file.log,定义了日志调用应该如何处理。这个参数可以是很多东西,比如一个可调用的函数,但为了方便,Loguru 允许我们直接传递类似文件的对象,如sys.stdout,或被解释为文件名的字符串。接收器的所有方面都可以通过多个参数进行定制,尤其是日志级别。

正如我们所说,标准输出接收器只会记录至少为WARNING级别的消息,而文件接收器会记录所有消息。

请注意,我们为文件接收器添加了rotation参数。由于日志会不断附加到文件中,文件大小会在应用程序生命周期内迅速增长。因此,我们提供了一些选项供您选择:

  • “轮换”文件:这意味着当前文件将被重命名,并且新的日志会添加到一个新文件中。此操作可以配置为在一段时间后发生(例如每天,如我们的示例)或当文件达到一定大小时。

  • 删除旧文件:经过一段时间后,保留占用磁盘空间的旧日志可能就不太有用了。

您可以在 Loguru 的官方文档中阅读有关这些功能的所有详细信息:loguru.readthedocs.io/en/stable/api/logger.html#file

现在,如果我们运行这个示例,我们将在控制台输出中看到以下内容:


(venv) $ python chapter15/chapter15_logs_02.py2023-03-03 08:15:16.804 | ERROR    | __main__:is_even:12 - hello is not an integer
Traceback (most recent call last):
  File "/Users/fvoron/Development/Building-Data-Science-Applications-with-FastAPI-Second-Edition/chapter15/chapter15_logs_02.py", line 19, in <module>
    is_even("hello")
  File "/Users/fvoron/Development/Building-Data-Science-Applications-with-FastAPI-Second-Edition/chapter15/chapter15_logs_02.py", line 13, in is_even
    raise TypeError()
TypeError

DEBUG级别的日志不再出现了。然而,如果我们读取file.log文件,我们将看到两者:


$ cat file.log2023-03-03 08:15:16.803 | DEBUG    | __main__:is_even:10 - Check if 2 is even
2023-03-03 08:15:16.804 | DEBUG    | __main__:is_even:10 - Check if hello is even
2023-03-03 08:15:16.804 | ERROR    | __main__:is_even:12 - hello is not an integer

就这样!接收器非常有用,可以根据日志的性质或重要性,将日志路由到不同的位置。

日志结构化和添加上下文

在最简单的形式下,日志由自由格式的文本组成。虽然这样很方便,但我们已经看到,通常需要记录变量值,以便更好地理解发生了什么。仅用字符串时,通常会导致多个连接值拼接成的混乱字符串。

更好的处理方式是采用结构化日志记录。目标是为每个日志行提供清晰且适当的结构,这样我们就可以在不牺牲可读性的前提下嵌入所有需要的信息。Loguru 本身通过上下文支持这种方法。下一个示例展示了如何使用它:

chapter15_logs_03.py


def is_even(n) -> bool:    logger_context = logger.bind(n=n)
    logger_context.debug("Check if even")
    if not isinstance(n, int):
        logger_context.error("Not an integer")
        raise TypeError()
    return n % 2 == 0

github.com/PacktPublishing/Building-Data-Science-Applications-with-FastAPI-Second-Edition/tree/main/chapter15/chapter15_logs_03.py

我们再次使用之前的相同示例。如你所见,我们使用了 logger 的 bind 方法来保留额外信息。在这里,我们设置了 n 变量。这个方法返回一个新的 logger 实例,并附加了这些属性。然后,我们可以正常使用这个实例来记录日志。我们不需要在格式化字符串中再添加 n 了。

然而,如果你直接运行这个示例,你将不会在日志中看到 n 的值。这是正常的:默认情况下,Loguru 不会将上下文信息添加到格式化的日志行中。我们需要自定义它!让我们看看如何操作:

chapter15_logs_04.py


logger.add(    sys.stdout,
    level="DEBUG",
    format="<green>{time:YYYY-MM-DD HH:mm:ss.SSS}</green> | "
    "<level>{level: <8}</level> | "
    "<cyan>{name}</cyan>:<cyan>{function}</cyan>:<cyan>{line}</cyan> - <level>{message}</level>"
    " - {extra}",
)

github.com/PacktPublishing/Building-Data-Science-Applications-with-FastAPI-Second-Edition/tree/main/chapter15/chapter15_logs_04.py

要格式化日志输出,我们必须在配置 sink 时使用 format 参数。它期望一个模板字符串。在这里,我们复制并粘贴了默认的 Loguru 格式,并添加了一个包含 extra 变量的部分。extra 是一个字典,Loguru 在其中存储所有你在上下文中添加的值。在这里,我们只是直接输出它,这样我们就能看到所有变量。

格式语法和可用变量

你可以在 Loguru 文档中找到所有可用的变量,这些变量可以在格式字符串中输出,如 extralevel,网址为:loguru.readthedocs.io/en/stable/api/logger.html#record

格式字符串支持标准的格式化指令,这些指令对于提取值、格式化数字、填充字符串等非常有用。你可以在 Python 文档中阅读更多相关内容:docs.python.org/3/library/string.html#format-string-syntax

此外,Loguru 还添加了特殊的标记语法,你可以用它来为输出着色。你可以在这里了解更多内容:loguru.readthedocs.io/en/stable/api/logger.html#color

这次,如果你运行这个示例,你会看到额外的上下文信息已经被添加到日志行中:


(venv) $ python chapter15/chapter15_logs_04.py2023-03-03 08:30:10.905 | DEBUG    | __main__:is_even:18 - Check if even - {'n': 2}
2023-03-03 08:30:10.905 | DEBUG    | __main__:is_even:18 - Check if even - {'n': 'hello'}
2023-03-03 08:30:10.905 | ERROR    | __main__:is_even:20 - Not an integer - {'n': 'hello'}

这种方法非常方便且强大:如果你想在日志中追踪一个你关心的值,只需添加一次。

以 JSON 对象形式记录日志

另一种结构化日志的方法是将日志的所有数据序列化为一个 JSON 对象。通过在配置 sink 时设置 serialize=True,可以轻松启用此功能。如果你计划使用日志摄取服务,如 Logstash 或 Datadog,这种方法可能会很有用:它们能够解析 JSON 数据并使其可供查询。

现在你已经掌握了使用 Loguru 添加和配置日志的基本知识。接下来,让我们看看如何在 FastAPI 应用中利用它们。

配置 Loguru 作为中央日志记录器

向你的 FastAPI 应用添加日志非常有用,它能帮助你了解不同路由和依赖项中发生了什么。

让我们来看一个来自 第五章 的例子,我们在其中添加了一个全局依赖项,用于检查应该在头部设置的密钥值。在这个新版本中,我们将添加一个调试日志,以跟踪 secret_header 依赖项何时被调用,并添加一个警告日志,告知我们此密钥缺失或无效:

chapter15_logs_05.py


from loguru import loggerdef secret_header(secret_header: str | None = Header(None)) -> None:
    logger.debug("Check secret header")
    if not secret_header or secret_header != "SECRET_VALUE":
        logger.warning("Invalid or missing secret header")
        raise HTTPException(status.HTTP_403_FORBIDDEN)

github.com/PacktPublishing/Building-Data-Science-Applications-with-FastAPI-Second-Edition/tree/main/chapter15/chapter15_logs_05.py

如果你一直跟随我们的教程,到这里应该没有什么令人惊讶的!现在,让我们用 Uvicorn 运行这个应用,并发出一个带有无效头部的请求:


INFO:     Started server process [47073]INFO:     Waiting for application startup.
INFO:     Application startup complete.
INFO:     Uvicorn running on http://127.0.0.1:8000 (Press CTRL+C to quit)
2023-03-03 09:00:47.324 | DEBUG    | chapter15.chapter15_logs_05:secret_header:6 - Check secret header
2023-03-03 09:00:47.324 | WARNING  | chapter15.chapter15_logs_05:secret_header:8 - Invalid or missing secret header
INFO:     127.0.0.1:58190 - "GET /route1 HTTP/1.1" 403 Forbidden

我们自己的日志在这里,但有一个问题:Uvicorn 也添加了它自己的日志,但是它没有遵循我们的格式!实际上,这是可以预料的:其他库,比如 Uvicorn,可能有自己的日志和设置。因此,它们不会遵循我们用 Loguru 定义的格式。这有点让人烦恼,因为如果我们有一个复杂且经过深思熟虑的设置,我们希望每个日志都能遵循它。幸运的是,还是有一些方法可以配置它。

首先,我们将创建一个名为 logger.py 的模块,在其中放置所有的日志配置。在你的项目中创建这个模块是一个很好的做法,这样你的配置就能集中在一个地方。我们在这个文件中做的第一件事是配置 Loguru:

logger.py


LOG_LEVEL = "DEBUG"logger.remove()
logger.add(
    sys.stdout,
    level=LOG_LEVEL,
    format="<green>{time:YYYY-MM-DD HH:mm:ss.SSS}</green> | "
    "<level>{level: <8}</level> | "
    "<cyan>{name}</cyan>:<cyan>{function}</cyan>:<cyan>{line}</cyan> - <level>{message}</level>"
    " - {extra}",
)

github.com/PacktPublishing/Building-Data-Science-Applications-with-FastAPI-Second-Edition/tree/main/chapter15/logger.py

就像我们在上一节中所做的那样,我们移除了默认的处理器并定义了我们自己的。注意,我们通过一个名为 LOG_LEVEL 的常量来设置级别。我们在这里硬编码了它,但更好的做法是从 Settings 对象中获取这个值,就像我们在 第十章 中所示的那样。这样,我们可以直接从环境变量中设置级别!

之后,我们在名为 InterceptHandler 的类中有一段相当复杂的代码。它是一个自定义处理器,针对标准日志模块,会将每个标准日志调用转发到 Loguru。这段代码直接取自 Loguru 文档。我们不会深入讲解它的工作原理,但只需要知道它会获取日志级别并遍历调用栈来获取原始调用者,然后将这些信息转发给 Loguru。

然而,最重要的部分是我们如何使用这个类。让我们在这里看看:

logger.py


logging.basicConfig(handlers=[InterceptHandler()], level=0, force=True)for uvicorn_logger_name in ["uvicorn.error", "uvicorn.access"]:
    uvicorn_logger = logging.getLogger(uvicorn_logger_name)
    uvicorn_logger.propagate = False
    uvicorn_logger.handlers = [InterceptHandler()]

github.com/PacktPublishing/Building-Data-Science-Applications-with-FastAPI-Second-Edition/tree/main/chapter15/logger.py

这里的技巧是调用标准日志模块中的basicConfig方法来设置我们的自定义拦截处理程序。这样,通过根日志记录器发出的每个日志调用,即使是来自外部库的,也会通过它并由 Loguru 处理。

然而,在某些情况下,这种配置是不够的。一些库定义了自己的日志记录器和处理程序,因此它们不会使用根配置。这对于 Uvicorn 来说就是这样,它定义了两个主要的日志记录器:uvicorn.erroruvicorn.access。通过获取这些日志记录器并更改其处理程序,我们强制它们也通过 Loguru。

如果你使用其他像 Uvicorn 一样定义自己日志记录器的库,你可能需要应用相同的技巧。你需要做的就是确定它们日志记录器的名称,这应该很容易在库的源代码中找到。

它与 Dramatiq 开箱即用

如果你实现了一个 Dramatiq 工作程序,正如我们在第十四章中展示的那样,你会看到,如果你使用logger模块,Dramatiq 的默认日志将会被 Loguru 正确处理。

最后,我们在模块的末尾处理设置__all__变量:

logger.py


__all__ = ["logger"]

github.com/PacktPublishing/Building-Data-Science-Applications-with-FastAPI-Second-Edition/tree/main/chapter15/logger.py

__all__是一个特殊变量,告诉 Python 在导入此模块时应该公开哪些变量。在这里,我们将暴露 Loguru 中的logger,以便在项目中任何需要的地方都能轻松导入它。

请记住,使用__all__并不是严格必要的:我们完全可以在没有它的情况下导入logger,但它是一种干净的方式来隐藏我们希望保持私有的其他内容,例如InterceptHandler

最后,我们可以像之前在代码中看到的那样使用它:

logger.py


from chapter15.logger import loggerdef secret_header(secret_header: str | None = Header(None))    None:
    logger.debug("Check secret header")
    if not secret_header or secret_header != "SECRET_VALUE":
        logger.warning("Invalid or missing secret header")
        raise HTTPException(status.HTTP_403_FORBIDDEN)

github.com/PacktPublishing/Building-Data-Science-Applications-with-FastAPI-Second-Edition/tree/main/chapter15/logger.py

如果我们用 Uvicorn 运行它,你会发现我们所有的日志现在都以相同的格式显示:


2023-03-03 09:06:16.196 | INFO     | uvicorn.server:serve:75 - Started server process [47534] - {}2023-03-03 09:06:16.196 | INFO     | uvicorn.lifespan.on:startup:47 - Waiting for application startup. - {}
2023-03-03 09:06:16.196 | INFO     | uvicorn.lifespan.on:startup:61 - Application startup complete. - {}
2023-03-03 09:06:16.196 | INFO     | uvicorn.server:_log_started_message:209 - Uvicorn running on http://127.0.0.1:8000 (Press CTRL+C to quit) - {}
2023-03-03 09:06:18.500 | DEBUG    | chapter15.chapter15_logs_06:secret_header:7 - Check secret header - {}
2023-03-03 09:06:18.500 | WARNING  | chapter15.chapter15_logs_06:secret_header:9 - Invalid or missing secret header - {}
2023-03-03 09:06:18.500 | INFO     | uvicorn.protocols.http.httptools_impl:send:489 - 127.0.0.1:59542 - "GET /route1 HTTP/1.1" 403 - {}

太好了!现在,每当你需要在应用程序中添加日志时,所要做的就是从logger模块中导入logger

现在,你已经掌握了将日志添加到应用程序的基本知识,并有许多选项可以微调如何以及在哪里输出日志。日志对于监控你的应用程序在微观层面上的行为非常有用,逐操作地了解它在做什么。监控的另一个重要方面是获取更一般层面的信息,以便获取大的数据图并快速发现问题。这就是我们现在要通过指标来实现的目标。

添加 Prometheus 指标

在上一节中,我们看到日志如何帮助我们通过精细地追踪程序随时间执行的操作,来理解程序的行为。然而,大多数时候,你不能整天盯着日志看:它们对于理解和调试特定情况非常有用,但对于获取全球性洞察力、在出现问题时发出警报却要差得多。

为了解决这个问题,我们将在本节中学习如何将指标添加到我们的应用程序中。它们的作用是衡量程序执行中重要的事项:发出的请求数量、响应时间、工作队列中待处理任务的数量、机器学习预测的准确性……任何我们可以轻松地随时间监控的事情——通常通过图表和图形——这样我们就能轻松监控系统的健康状况。我们称之为为应用程序添加监控

为了完成这个任务,我们将使用两种在行业中广泛使用的技术:Prometheus 和 Grafana。

理解 Prometheus 和不同的指标

Prometheus 是一种帮助你为应用程序添加监控的技术。它由三部分组成:

  • 各种编程语言的库,包括 Python,用于向应用程序添加指标

  • 一个服务器,用于聚合并存储这些指标随时间变化的值

  • 一种查询语言 PromQL,用于将这些指标中的数据提取到可视化工具中

Prometheus 对如何定义指标有非常精确的指南和约定。实际上,它定义了四种不同类型的指标。

计数器指标

计数器指标是一种衡量随着时间推移上升的值的方法。例如,这可以是已答复的请求数量或完成的预测数量。它不能用于可以下降的值。对于这种情况,有仪表指标。

图 15.2 – 计数器的可能表示

图 15.2 – 计数器的可能表示

仪表指标

仪表指标是一种衡量随着时间的推移可以上升或下降的值的方法。例如,这可以是当前的内存使用量或工作队列中待处理任务的数量。

图 15.3 – 仪表的可能表示

图 15.3 – 仪表的可能表示

直方图指标

与计数器和仪表不同,直方图将测量值并将其计入桶中。通常,如果我们想测量 API 的响应时间,我们可以统计处理时间少于 10 毫秒、少于 100 毫秒和少于 1 秒的请求数量。例如,做这个比仅获取一个简单的平均值或中位数要有洞察力得多。

使用直方图时,我们有责任定义所需的桶以及它们的值阈值。

图 15.4 – 直方图的可能表示

图 15.4 – 直方图的可能表示

Prometheus 定义了第四种类型的指标——摘要。它与直方图指标非常相似,但它使用滑动分位数而不是定义的桶。由于在 Python 中支持有限,我们不会详细介绍。此外,在本章的 Grafana 部分,我们将看到能够使用直方图指标计算分位数。

您可以在官方 Prometheus 文档中阅读有关这些指标的更多详细信息:

prometheus.io/docs/concepts/metric_types/

测量和暴露指标

一旦定义了指标,我们就可以开始在程序生命周期中进行测量。与我们记录日志的方式类似,指标暴露了方法,使我们能够在应用程序执行期间存储值。然后,Prometheus 会将这些值保存在内存中,以便构建指标。

那么,我们如何访问这些指标以便实际分析和监控呢?很简单,使用 Prometheus 的应用程序通常会暴露一个名为 /metrics 的 HTTP 端点,返回所有指标的当前值,格式是特定的。您可以在图 15.5中查看它的样子。

图 15.5 – Prometheus 指标端点的输出

图 15.5 – Prometheus 指标端点的输出

该端点可以由 Prometheus 服务器定期轮询,Prometheus 会随着时间推移存储这些指标,并通过 PromQL 提供访问。

当您的应用程序重启时,指标会被重置。

值得注意的是,每次重启应用程序时(如 FastAPI 服务器),指标值都会丢失,并且从零开始。这可能有些令人惊讶,但理解指标值仅保存在应用程序的内存中是非常重要的。永久保存它们的责任属于 Prometheus 服务器。

现在我们已经对它们的工作原理有了一个大致了解,接下来让我们看看如何将指标添加到 FastAPI 和 Dramatiq 应用程序中。

将 Prometheus 指标添加到 FastAPI

正如我们所说,Prometheus 为各种语言(包括 Python)维护了官方库。

我们完全可以单独使用它,并手动定义各种指标来监控我们的 FastAPI 应用。我们还需要编写一些逻辑,将其挂钩到 FastAPI 请求处理程序中,以便我们可以衡量诸如请求计数、响应时间、负载大小等指标。

虽然完全可以实现,但我们将采取捷径,再次依赖开源社区,它提供了一个现成的库,用于将 Prometheus 集成到 FastAPI 项目中:/metrics 端点。

首先,当然需要通过 pip 安装它。运行以下命令:


(venv) $ pip install prometheus_fastapi_instrumentator

在下面的示例中,我们实现了一个非常简单的 FastAPI 应用,并启用了仪表监控器:

chapter15_metrics_01.py


from fastapi import FastAPIfrom prometheus_fastapi_instrumentator import Instrumentator, metrics
app = FastAPI()
@app.get("/")
async def hello():
    return {"hello": "world"}
instrumentator = Instrumentator()
instrumentator.add(metrics.default())
instrumentator.instrument(app).expose(app)

github.com/PacktPublishing/Building-Data-Science-Applications-with-FastAPI-Second-Edition/tree/main/chapter15/chapter15_metrics_01.py

启用仪表监控器只需要三行代码:

  1. 实例化 Instrumentator 类。

  2. 启用库中提供的默认指标。

  3. 将其与我们的 FastAPI app 连接并暴露 /metrics 端点。

就这样!FastAPI 已经集成了 Prometheus!

让我们用 Uvicorn 运行这个应用并访问 hello 端点。内部,Prometheus 将会测量有关这个请求的一些信息。现在让我们访问 /metrics 来查看结果。如果你滚动查看这个长长的指标列表,你应该能看到以下几行:


# HELP http_requests_total Total number of requests by method, status and handler.# TYPE http_requests_total counter
http_requests_total{handler="/",method="GET",status="2xx"} 1.0

这是计数请求数量的指标。我们看到总共有一个请求,这对应于我们对hello的调用。请注意,仪表监控工具足够智能,可以根据路径、方法,甚至状态码为指标打上标签。这非常方便,因为它使我们能够根据请求的特征提取有趣的数据。

添加自定义指标

内置的指标是一个不错的开始,但我们可能需要根据我们应用的特定需求来定义自己的指标。

假设我们想要实现一个掷骰子的函数,骰子有六个面,并通过 REST API 暴露它。我们希望定义一个指标,允许我们计算每个面出现的次数。对于这个任务,计数器是一个很好的选择。让我们看看如何在代码中声明它:

chapter15_metrics_02.py


DICE_COUNTER = Counter(    "app_dice_rolls_total",
    "Total number of dice rolls labelled per face",
    labelnames=["face"],
)

github.com/PacktPublishing/Building-Data-Science-Applications-with-FastAPI-Second-Edition/tree/main/chapter15/chapter15_metrics_02.py

我们必须实例化一个 Counter 对象。前两个参数分别是指标的名称和描述。名称将由 Prometheus 用来唯一标识这个指标。因为我们想要统计每个面出现的次数,所以我们还添加了一个名为 face 的标签。每次我们统计骰子的投掷次数时,都需要将此标签设置为相应的面值。

度量命名规范

Prometheus 为度量命名定义了非常精确的规范。特别是,它应该以度量所属的领域开始,例如 http_app_,并且如果仅是一个值计数,则应该以单位结尾,例如 _seconds_bytes_total。我们强烈建议您阅读 Prometheus 的命名规范:prometheus.io/docs/practices/naming/

现在我们可以在代码中使用这个度量了。在下面的代码片段中,您将看到 roll_dice 函数的实现:

chapter15_metrics_02.py


def roll_dice() -> int:    result = random.randint(1, 6)
    DICE_COUNTER.labels(result).inc()
    return result

github.com/PacktPublishing/Building-Data-Science-Applications-with-FastAPI-Second-Edition/tree/main/chapter15/chapter15_metrics_02.py

您可以看到,我们直接使用度量实例 DICE_COUNTER,首先调用 labels 方法来设置骰子面数,然后调用 inc 来实际增加计数器。

这就是我们需要做的:我们的度量已经自动注册到 Prometheus 客户端,并将通过 /metrics 端点开始暴露。在 图 15.6 中,您可以看到此度量在 Grafana 中的可能可视化。

图 15.6 – 在 Grafana 中表示骰子投掷度量

图 15.6 – 在 Grafana 中表示骰子投掷度量

如您所见,声明和使用新度量是非常简单的:我们只需在想要监控的代码中直接调用它。

处理多个进程

第十章中,我们在 为部署添加 Gunicorn 作为服务器进程 部分提到过,在生产部署中,FastAPI 应用程序通常会使用多个工作进程运行。基本上,它会启动多个相同应用程序的进程,并在它们之间平衡传入的请求。这使得我们可以并发处理更多请求,避免因某个操作阻塞进程而导致的阻塞。

不要混淆 Gunicorn 工作进程和 Dramatiq 工作进程

当我们谈论 Gunicorn 部署中的工作进程时,我们指的是通过启动多个进程来并发处理 API 请求的方式。我们不是指 Dramatiq 中的工作进程,这些进程是在后台处理任务。

对于同一应用程序,拥有多个进程在 Prometheus 度量方面是有点问题的。事实上,正如我们之前提到的,这些度量仅存储在内存中,并通过 /``metrics 端点暴露。

如果我们有多个进程来处理请求,每个进程都会有自己的一组度量值。然后,当 Prometheus 服务器请求 /metrics 时,我们将获得响应我们请求的进程的度量值,而不是其他进程的度量值。这些值在下次轮询时可能会发生变化!显然,这将完全破坏我们最初的目标。

为了绕过这个问题,Prometheus 客户端有一个特殊的多进程模式。基本上,它不会将值存储在内存中,而是将它们存储在专用文件夹中的文件里。当调用 /metrics 时,它会负责加载所有文件并将所有进程的值进行合并。

启用此模式需要我们设置一个名为 PROMETHEUS_MULTIPROC_DIR 的环境变量。它应该指向文件系统中一个有效的文件夹,存储指标文件。以下是如何设置这个变量并启动带有四个工作进程的 Gunicorn 的命令示例:


(venv) $ PROMETHEUS_MULTIPROC_DIR=./prometheus-tmp gunicorn -w 4 -k uvicorn.workers.UvicornWorker chapter15.chapter15_metrics_01:app

当然,在生产环境部署时,你应该在平台上全局设置环境变量,正如我们在第十章中所解释的那样。

如果你尝试这个命令,你会看到 Prometheus 会开始在文件夹内存储一些 .db 文件,每个文件对应一个指标和一个进程。副作用是,在重启进程时,指标不会被清除。如果你更改了指标定义,或者运行了完全不同的应用程序,可能会导致意外的行为。确保为每个应用选择一个专用的文件夹,并在运行新版本时清理该文件夹。

我们现在能够精确地对 FastAPI 应用进行监控。然而,正如我们在前一章中所看到的,数据科学应用可能包含一个独立的工作进程,其中运行着大量的逻辑和智能。因此,对应用的这一部分进行监控也至关重要。

向 Dramatiq 添加 Prometheus 指标

第十四章中,我们实现了一个复杂的应用,包含一个独立的工作进程,该进程负责加载并执行 Stable Diffusion 模型来生成图像。因此,架构中的这一部分非常关键,需要进行监控,以确保一切顺利。

在这一部分,我们将学习如何向 Dramatiq 工作进程添加 Prometheus 指标。好消息是,Dramatiq 已经内置了指标,并且默认暴露了 /metrics 端点。实际上,几乎不需要做什么!

让我们来看一个非常基础的 Dramatiq 工作进程的例子,里面包含一个虚拟任务:

chapter15_metrics_03.py


import timeimport dramatiq
from dramatiq.brokers.redis import RedisBroker
redis_broker = RedisBroker(host="localhost")
dramatiq.set_broker(redis_broker)
@dramatiq.actor()
def addition_task(a: int, b: int):
    time.sleep(2)
    print(a + b)

github.com/PacktPublishing/Building-Data-Science-Applications-with-FastAPI-Second-Edition/tree/main/chapter15/chapter15_metrics_03.py

正如你现在可能已经理解的,Dramatiq 本质上是一个多进程程序:它会启动多个工作进程来并发处理任务。因此,我们需要确保 Prometheus 处于多进程模式,正如我们在处理多个进程部分中提到的那样。因此,我们需要设置PROMETHEUS_MULTIPROC_DIR环境变量,正如我们之前解释的那样,还需要设置dramatiq_prom_db。事实上,Dramatiq 实现了自己的机制来启用 Prometheus 的多进程模式,这应该是开箱即用的,但根据我们的经验,明确指出这一点会更好。

以下命令展示了如何启动带有PROMETHEUS_MULTIPROC_DIRdramatiq_prom_db设置的工作进程:


(venv) $ PROMETHEUS_MULTIPROC_DIR=./prometheus-tmp-dramatiq dramatiq_prom_db=./prometheus-tmp-dramatiq dramatiq chapter15.chapter15_metrics_03

为了让你能轻松在这个工作进程中调度任务,我们添加了一个小的__name__ == "__main__"指令。在另一个终端中,运行以下命令:


(venv) $ python -m chapter15.chapter15_metrics_03

它将在工作进程中调度一个任务。你可能会在工作进程日志中看到它的执行情况。

现在,尝试在浏览器中打开以下 URL:http://localhost:9191/metrics。你将看到类似于我们在图 15.7中展示的结果。

图 15.7 – Dramatiq Prometheus 度量端点的输出

图 15.7 – Dramatiq Prometheus 度量端点的输出

我们已经看到几个度量指标,包括一个用于统计 Dramatiq 处理的消息总数的计数器,一个用于测量任务执行时间的直方图,以及一个用于衡量当前正在进行的任务数量的仪表。你可以在 Dramatiq 的官方文档中查看完整的度量指标列表:dramatiq.io/advanced.html#prometheus-metrics

添加自定义指标

当然,对于 FastAPI,我们可能也希望向 Dramatiq 工作进程添加我们自己的指标。事实上,这与我们在上一节中看到的非常相似。让我们再次以掷骰子为例:

chapter15_metrics_04.py


DICE_COUNTER = Counter(    "worker_dice_rolls_total",
    "Total number of dice rolls labelled per face",
    labelnames=["face"],
)
@dramatiq.actor()
def roll_dice_task():
    result = random.randint(1, 6)
    time.sleep(2)
    DICE_COUNTER.labels(result).inc()
    print(result)

github.com/PacktPublishing/Building-Data-Science-Applications-with-FastAPI-Second-Edition/tree/main/chapter15/chapter15_metrics_04.py

我们所需要做的只是创建我们的Counter对象,正如我们之前所做的那样,并在任务中使用它。如果你尝试运行工作进程并请求/metrics端点,你会看到这个新指标出现。

我们现在可以对我们的 FastAPI 和 Dramatiq 应用进行指标收集。正如我们之前多次提到的那样,我们现在需要将这些指标汇总到 Prometheus 服务器中,并在 Grafana 中进行可视化。这就是我们将在下一节中讨论的内容。

在 Grafana 中监控指标

拥有度量指标固然不错,但能够可视化它们更好!在本节中,我们将看到如何收集 Prometheus 度量指标,将它们发送到 Grafana,并创建仪表板来监控它们。

Grafana 是一个开源的 Web 应用程序,用于数据可视化和分析。它能够连接到各种数据源,比如时间序列数据库,当然也包括 Prometheus。其强大的查询和图形构建器使我们能够创建详细的仪表板,在其中实时监控我们的数据。

配置 Grafana 收集指标

由于它是开源的,你可以在自己的机器或服务器上运行它。详细的安装说明可以在官方文档中找到:grafana.com/docs/grafana/latest/setup-grafana/installation/。不过,为了加快进程并快速开始,我们这里依赖的是 Grafana Cloud,这是一个官方托管平台。它提供了一个免费的计划,足以让你开始使用。你可以在这里创建账户:grafana.com/auth/sign-up/create-user。完成后,你将被要求创建自己的实例,即一个“Grafana Stack”,通过选择子域名和数据中心区域,如图 15**.8 所示。请选择一个靠近你地理位置的区域。

图 15.8 – 在 Grafana Cloud 上创建实例

图 15.8 – 在 Grafana Cloud 上创建实例

然后,你将看到一组常见的操作,帮助你开始使用 Grafana。我们要做的第一件事是添加 Prometheus 指标。点击扩展和集中现有数据,然后选择托管 Prometheus 指标。你将进入一个配置 Prometheus 指标收集的页面。在顶部点击名为配置详情的选项卡。页面将呈现如图 15**.9所示。

图 15.9 – 在 Grafana 上配置托管 Prometheus 指标

图 15.9 – 在 Grafana 上配置托管 Prometheus 指标

你可以看到,我们有两种方式来转发指标:通过 Grafana Agent 或通过 Prometheus 服务器。

如前所述,Prometheus 服务器负责收集我们所有应用程序的指标,并将数据存储在数据库中。这是标准的做法。你可以在官方文档中找到如何安装它的说明:prometheus.io/docs/prometheus/latest/installation/。不过,请注意,它是一个专用的应用服务器,需要适当的备份,因为它会存储所有的指标数据。

最直接的方式是使用 Grafana Agent。它由一个小型命令行程序和一个配置文件组成。当它运行时,它会轮询每个应用程序的指标,并将数据发送到 Grafana Cloud。所有数据都会存储在 Grafana Cloud 上,因此即使停止或删除代理,数据也不会丢失。这就是我们在这里使用的方法。

Grafana 会在页面上显示下载、解压并执行 Grafana Agent 程序的命令。执行这些命令,以便将其放在项目的根目录中。

然后,在最后一步,你需要创建一个 API 令牌,以便 Grafana Agent 可以将数据发送到你的实例。给它起个名字,然后点击创建 API 令牌。一个新的文本区域将出现,显示一个新的命令,用于创建代理的配置文件,正如你在图 15.10中看到的那样。

图 15.10 – 创建 Grafana Agent 配置的命令

图 15.10 – 创建 Grafana Agent 配置的命令

执行 ./grafana-agent-linux-amd64 –config.file=agent-config.yaml 命令。一个名为 agent-config.yaml 的文件将被创建在你的项目中。我们现在需要编辑它,以便配置我们的实际 FastAPI 和 Dramatiq 应用程序。你可以在以下代码片段中看到结果:

agent-config.yaml


metrics:  global:
    scrape_interval: 60s
  configs:
  - name: hosted-prometheus
    scrape_configs:
      - job_name: app
        static_configs:
        - targets: ['localhost:8000']
      - job_name: worker
        static_configs:
        - targets: ['localhost:9191']
    remote_write:
      - url: https://prometheus-prod-01-eu-west-0.grafana.net/api/prom/push
        basic_auth:
          username: 811873
          password: __YOUR_API_TOKEN__

github.com/PacktPublishing/Building-Data-Science-Applications-with-FastAPI-Second-Edition/tree/main/chapter15/agent-config.yaml

这是一个 YAML 配置文件,我们可以在其中设置 Grafana Agent 的各种选项。最重要的部分是 scrape_configs 键。如你所见,我们可以定义所有要收集指标的应用程序列表,并指定它们的主机名,对于 FastAPI 应用程序是“目标”:localhost:8000,而 Dramatiq worker 是 localhost:9191。当然,这个配置适用于本地开发,但在生产环境中,你需要根据实际应用程序的主机名进行调整。

我们现在准备启动 Grafana Agent 并收集指标数据了!确保你的 FastAPI 和 Dramatiq 应用程序正在运行,然后启动 Grafana Agent。根据你的系统,执行文件的名称会有所不同,但大致如下所示:


$ ./grafana-agent-linux-amd64 --config.file=agent-config.yaml

Grafana Agent 将启动并定期收集指标数据,然后将其发送到 Grafana。我们现在可以开始绘制一些数据了!

在 Grafana 中可视化指标

我们的指标数据现在已发送到 Grafana。我们准备好查询它并构建一些图表了。第一步是创建一个新的仪表板,这是一个可以创建和组织多个图表的地方。点击右上角的加号按钮,然后选择新建仪表板

一个新的空白仪表板将出现,正如你在图 15.11中看到的那样。

图 15.11 – 在 Grafana 中创建新仪表板

图 15.11 – 在 Grafana 中创建新仪表板

点击添加新面板。将会出现一个用于构建新图表的界面。主要有三个部分:

  • 左上角的图表预览。在开始时,它是空的。

  • 左下角的查询构建器。这是我们查询指标数据的地方。

  • 右侧的图表设置。这是我们选择图表类型并精细配置其外观和感觉的地方,类似于电子表格软件中的操作。

让我们尝试为我们 FastAPI 应用中的 HTTP 请求时长创建一个图表。在名为指标的选择菜单中,你将能访问到我们应用所报告的所有 Prometheus 指标。选择http_request_duration_seconds_bucket。这是 Prometheus FastAPI Instrumentator 默认定义的直方图指标,用于衡量我们端点的响应时间。

然后,点击运行查询。在后台,Grafana 会构建并执行 PromQL 查询来检索数据。

在图表的右上角,我们选择一个较短的时间跨度,比如过去 15 分钟。由于我们还没有太多数据,如果只看几分钟的数据,而不是几小时的数据,图表会更加清晰。你应该会看到一个类似图 15.12的图表。

图 15.12 – Grafana 中直方图指标的基本图

图 15.12 – Grafana 中直方图指标的基本图

Grafana 已绘制出多个系列:对于每个handler(对应于端点模式),我们有多个桶,le。每条线大致代表了我们在少于“le”秒内处理handler请求的次数

这是指标的原始表示。然而,你可能会发现它不太方便阅读和分析。如果我们能以另一种方式查看这些数据,按分位数排列的响应时间,可能会更好。

幸运的是,PromQL 包含一些数学运算,这样我们就可以对原始数据进行处理。在指标菜单下方的部分允许我们添加这些运算。我们甚至可以看到 Grafana 建议我们使用添加 histogram_quantile。如果点击这个蓝色按钮,Grafana 会自动添加三种操作:速率按 le 求和,最后是直方图分位数,默认设置为0.95

通过这样做,我们现在可以看到响应时间的变化情况:95%的时间,我们的响应时间少于x秒。

默认的y轴单位不太方便。由于我们知道我们使用的是秒,接下来在图表选项中选择这个单位。在右侧,找到标准选项部分,然后在单位菜单中,在时间组下选择秒(s)。现在你的图表应该像图 15.13一样。

图 15.13 – Grafana 中直方图指标的分位数表示

图 15.13 – Grafana 中直方图指标的分位数表示

现在情况更具洞察力了:我们可以看到,我们几乎处理了所有的请求(95%)都在 100 毫秒以内。如果我们的服务器开始变慢,我们会立即在图表中看到上升,这能提醒我们系统出现了问题。

如果我们希望在同一个图表上显示其他分位数,可以通过点击复制按钮(位于运行查询上方)来复制这个查询。然后,我们只需要选择另一个分位数。我们展示了0.950.900.50分位数的结果,见图 15.14

图 15.14 – Grafana 中同一图表上的多个分位数

图 15.14 – Grafana 中同一图表上的多个分位数

图例可以自定义

注意,图例中的系列名称是可以自定义的。在每个查询的选项部分,你可以根据需要进行自定义。你甚至可以包含来自查询的动态值,例如指标标签。

最后,我们可以通过在右侧列中设置面板标题来给我们的图表命名。现在我们对图表感到满意,可以点击右上角的应用按钮,将其添加到我们的仪表板中,如图 15.15所示。

图 15.15 – Grafana 仪表板

图 15.15 – Grafana 仪表板

就这样!我们可以开始监控我们的应用程序了。你可以随意调整每个面板的大小和位置。你还可以设置想要查看的查询时间范围,甚至启用自动刷新功能,这样数据就能实时更新!别忘了点击保存按钮来保存你的仪表板。

我们可以使用完全相同的配置,构建一个类似的图表,用于监控执行 Dramatiq 任务所需的时间,这要感谢名为dramatiq_message_duration_milliseconds_bucket的指标。注意,这个指标是以毫秒为单位表示的,而不是秒,所以在选择图表单位时需要特别小心。我们在这里看到了 Prometheus 指标命名约定的一个优点!

添加柱状图

Grafana 提供了许多不同类型的图表。例如,我们可以将骰子投掷指标绘制成柱状图,其中每根柱子表示某一面出现的次数。让我们来试试:添加一个新面板并选择app_dice_rolls_total指标。你会看到类似图 15.6所示的内容。

图 15.16 – Grafana 中计数器指标的默认表示方式(柱状图)

图 15.16 – Grafana 中计数器指标的默认表示方式(柱状图)

确实,我们为每个面都有一个柱子,但有一个奇怪的地方:每个时间点都有一根柱子。这是理解 Prometheus 指标和 PromQL 的关键:所有指标都作为时间序列存储。这使我们能够回溯时间,查看指标随时间的演变。

然而,对于某些表示方式,像这里显示的这种,实际上并不具备很高的洞察力。对于这种情况,最好是展示我们选择的时间范围内的最新值。我们可以通过将指标面板中的类型设置为即时来实现这一点。我们会看到现在我们有一个单一的图表,显示一个时间点的数据,如图 15.17所示。

图 15.17 – 在 Grafana 中将计数器指标配置为即时类型

图 15.17 – 在 Grafana 中将计数器指标配置为即时类型

这样已经更好了,但我们可以更进一步。通常,我们希望 x 轴显示面孔标签,而不是时间点。首先,让我们用 {{face}} 自定义图例。现在图例将只显示 face 标签。

现在,我们将数据转换,使得 x 轴为 face 标签。点击 Transform 标签。你会看到一系列可以在可视化之前应用到数据的函数。在我们这里,我们将选择 Reduce。这个函数的作用是取每个序列,从中提取一个特定的值并将其绘制在 x 轴上。默认情况下,Grafana 会取最大值 Max,但也有其他选项,如 LastMeanStdDev。在这种情况下,它们没有区别,因为我们已经查询了即时值。

就是这样!我们的图表现在显示了我们看到面孔的次数。这就是我们在第 15.6 图中展示的内容。

总结

恭喜!现在你可以在 Grafana 中报告指标,并构建自己的仪表盘来监控你的数据科学应用程序。随着时间的推移,如果你发现一些盲点,不要犹豫添加新的指标或完善你的仪表盘:目标是能够一目了然地监控每个重要部分,从而快速采取纠正措施。这些指标也可以用来推动你工作的演进:通过监控你的机器学习模型的性能和准确性,你可以跟踪你所做的改动的效果,看看自己是否走在正确的道路上。

本书的内容和我们的 FastAPI 之旅到此结束。我们真诚希望你喜欢本书,并且在这段旅程中学到了很多。我们覆盖了许多主题,有时只是稍微触及表面,但现在你应该已经准备好使用 FastAPI 构建自己的项目,并提供智能数据科学算法。一定要查看我们在旅程中提到的所有外部资源,它们将为你提供掌握这些技能所需的所有见解。

近年来,Python 在数据科学社区中获得了极大的关注,尽管 FastAPI 框架仍然非常年轻,但它已经是一个改变游戏规则的工具,并且已经看到了前所未有的采用率。它很可能会成为未来几年许多数据科学系统的核心……而当你读完这本书时,你可能就是这些系统背后的开发者之一。干杯!