使用 FastAPI 构建生成式 AI 服务——在AI工作负载中实现并发

在本章中，您将学习：

• 多线程、并行处理和异步编程在提升AI应用程序性能和可扩展性方面的作用和好处
• 线程池和事件循环在FastAPI服务器中的作用
• 如何在处理请求时避免阻塞服务器
• 使用异步编程与外部系统（如数据库、AI模型和网页内容）进行交互，通过构建网络爬虫和RAG模块
• 模型服务策略和LLM内存优化策略，以减少内存绑定的阻塞操作
• 处理长期运行的AI推理任务的策略

在本章中，您将进一步了解异步编程在提升GenAI服务的性能和可扩展性方面的作用和好处。作为其中的一部分，您将学习如何管理并发用户交互，并与外部系统（如数据库）进行交互，实施RAG，读取网页以丰富模型提示的上下文。您将掌握有效处理I/O绑定和CPU绑定操作的技巧，特别是在处理外部服务或处理长期推理任务时。

我们还将深入探讨高效处理长期生成AI推理任务的策略，包括使用FastAPI事件循环执行后台任务。

为多个用户优化GenAI服务

AI工作负载是计算密集型操作，可能会影响您的GenAI服务无法处理多个同时请求。在大多数生产环境中，多个用户将同时使用您的应用程序。因此，您的服务需要能够并发处理请求，以便多个重叠的任务能够同时执行。然而，如果您正在与GenAI模型和外部系统（如数据库、文件系统或互联网）进行交互，则会有一些操作可能会阻塞其他任务在服务器上执行。会中断程序执行流程的长期运行操作被称为阻塞操作。

这些阻塞操作可以分为两类：

• 输入/输出（I/O）绑定
  这是因为数据输入/输出操作需要等待，这些操作可能来自用户、文件、数据库、网络等。例如，读取或写入磁盘上的文件、发起网络请求和API调用、从数据库发送或接收数据、等待用户输入等。
• 计算绑定
  这是由于CPU或GPU上的计算密集型操作需要等待。计算绑定的程序通过执行高强度计算推动CPU或GPU核心到达其极限，通常会阻塞它们执行其他任务。例如，数据处理、AI模型推理或训练、3D对象渲染、运行仿真等。

您可以采用以下几种策略来服务多个用户：

• 系统优化
  针对I/O绑定任务，例如从数据库获取数据、与磁盘上的文件进行交互、发起网络请求或读取网页。
• 模型优化
  针对内存和计算绑定任务，如模型加载和推理。
• 队列系统
  用于处理长期推理任务，以避免响应延迟。

在本节中，我们将更详细地探讨每种策略。为了帮助加深您的理解，我们还将一起实现几个利用上述策略的功能：

• 构建一个网页爬虫，用于批量抓取和解析聊天中粘贴的HTTP URL，以便您可以让LLM了解网页内容。
• 在您的服务中添加一个检索增强生成（RAG）模块，使用自托管的向量数据库（如qdrant），这样您就可以上传文档并通过LLM服务与文档进行交互。
• 添加一个批量图像生成系统，以便您可以将图像生成工作负载作为后台任务运行。

在向您展示如何构建上述功能之前，我们应更深入地探讨并发和并行的主题，因为理解这两个概念将帮助您为自己的用例识别正确的策略。

并发是指服务在同一时间处理多个请求或任务的能力，而不是一个接一个地完成。在并发操作过程中，多个任务的时间线可能会重叠，并且可能会在不同的时间开始和结束。

在Python中，您可以通过在单个线程上切换任务（通过异步编程）或通过不同线程之间切换任务（通过多线程）来实现并发，即使只有一个CPU核心。

时间分片

时间分片是多线程和异步编程中的一种调度机制，在这种机制中，进程将CPU时间分配给多个任务，从而产生并发执行的假象。

在Python中，由于全局解释器锁（GIL）的存在，CPU时间在任何时刻只能分配给一个任务。Python的GIL允许只有一个线程（即Python进程中的执行流）控制Python解释器来执行代码。这意味着在任何时刻，只有一个线程可以处于执行状态。

GIL最初是为了简化Python的开发，便于Python进程中线程间的内存管理，并优先考虑单线程程序的性能而实现的。它的加入还确保了进程内的线程安全，防止了多线程共享资源时可能导致的数据竞争问题。

然而，Python引入GIL的同时也意味着，在单个Python进程中，线程的真正并行执行是不可能的。当一个任务在等待I/O操作完成时，CPU可以快速切换到另一个任务，以避免阻塞其他操作。暂停的任务会保存其状态，并且一旦I/O操作完成，可以恢复执行。

在多个核心的情况下，您还可以实现并发的一个子集，称为并行性，其中任务被分配给多个独立的工作进程（通过多进程），每个工作进程在其自己的隔离资源和独立进程中同时执行任务。

注意
尽管有计划在不久的将来移除Python的GIL，但在本文撰写时，多个线程无法同时处理任务。因此，即使在单核上，通过并发也可以产生并行的假象，尽管只有一个进程在执行所有工作。单个进程只能通过切换活动线程来进行多任务处理，以最小化I/O阻塞操作的等待时间。

要实现真正的并行性，您必须使用多个工作进程（在多进程中）。

尽管并发和并行性有许多相似之处，但它们并不是完全相同的概念。它们之间的主要区别在于，并发通过交替执行任务来帮助您管理多个任务，这对于I/O绑定任务非常有用。而并行性则涉及同时执行多个任务，通常在多核机器上进行，这对于CPU绑定任务更为有用。

您可以通过线程或异步编程等方法来实现并发（即在单核机器上进行时间分片，任务交替执行，给人以同时执行的假象）。

图5-1展示了并发与并行之间的关系。

在大多数可扩展系统中，您可以看到并发和并行性同时存在。

假设您正在访问一个快餐店并下单。在一个并发系统中，您会看到餐厅老板在做汉堡的同时接受订单，不时切换任务，有效地通过任务切换实现多任务处理。在一个并行系统中，您会看到多个员工同时接受订单，并且有几个人同时在做汉堡。在这种情况下，不同的工作人员同时处理每个任务。

在单线程进程中，如果没有任何多线程或异步编程，进程必须等待阻塞操作完成后才能开始新任务。如果没有多进程在多个核心上实现并行性，计算密集型操作可能会阻止应用程序开始其他任务。

图5-2展示了非并发执行、没有并行性的并发执行（单核）和具有并行性的并发执行（多核）之间的区别。

图5-2中展示的三种Python执行模型如下：

• 没有并发（同步）
  单个进程（在一个核心上）按顺序执行任务。
• 并发且无并行
  单个进程中的多个线程（在一个核心上）并发处理任务，但由于Python的GIL，它们无法并行执行。
• 并发且并行
  多个进程在多个核心上并行执行任务，充分利用多核处理器以达到最大效率。

在多进程中，每个进程都有自己的内存空间和资源，用来独立完成任务，彼此之间相互隔离。这种隔离性使得进程更加稳定——因为如果一个进程崩溃，它不会影响其他进程——但与共享同一内存空间的线程相比，进程间通信变得更加复杂，如图5-3所示。

分布式工作负载通常使用一个管理进程来协调这些进程的执行和协作，以避免数据损坏和重复工作的等问题。一个很好的多进程示例是，当您使用负载均衡器管理流量到多个容器时，每个容器运行您的应用程序实例。

多线程和异步编程都能减少I/O任务的等待时间，因为处理器在等待I/O时可以做其他工作。然而，它们对需要大量计算的任务（如AI推理）没有帮助，因为进程忙于计算某些结果。因此，为了将大型自托管的GenAI模型提供给多个用户，您应该通过多进程扩展服务，或者使用算法模型优化（通过专门的模型推理服务器，如vLLM）。

当您处理慢模型时，第一反应可能是通过在单台机器上创建多个FastAPI服务实例（多进程）来采用并行性，以并行处理请求。

不幸的是，多个在不同进程中运行的工作进程无法访问共享内存空间。因此，您无法在FastAPI的应用程序实例之间共享内存中加载的资源——例如GenAI模型。遗憾的是，您的模型新实例将需要重新加载，这将大幅消耗硬件资源。这是因为FastAPI是一个通用的Web服务器，无法原生优化GenAI模型的服务。

解决方案不是单纯的并行性，而是采用外部模型服务策略，正如第3章所讨论的那样。

在依赖第三方AI提供商API（例如OpenAI API）时，您唯一可以将AI推理工作负载视为I/O绑定而非计算绑定的情况。在这种情况下，您将计算密集型任务通过网络请求卸载给模型提供者。

在您的侧面，AI推理工作负载通过网络请求变为I/O绑定，从而可以通过时间分片使用并发。第三方提供商需要担心如何扩展其服务，以处理跨硬件资源的计算绑定模型推理。

您可以通过专门的服务器，如vLLM、Ray Serve或NVIDIA Triton，将更大的GenAI模型（如LLM）的服务和推理外包。

在本章的后续部分，我将详细介绍这些服务器如何在模型推理过程中最大化计算密集型操作的推理效率，同时最小化模型在数据生成过程中的内存占用。

为了帮助您消化到目前为止讨论的内容，请查看表5-1中的并发策略比较表，以了解何时以及为什么使用每种策略。

表5-1. 并发策略比较

策略	特点	挑战	使用场景
无并发（同步）	简单、可读、易于调试的代码	单核和单线程	可能存在较长的等待时间，取决于I/O或CPU阻塞操作的执行阻止
异步I/O（异步）	单核和单线程	通过事件循环在Python进程内管理多任务	线程安全，因为Python进程管理任务
多线程	单核，但同一进程内有多个线程	线程共享数据和资源	实现代码较异步I/O简单
多进程	在多个CPU核心上运行多个进程	每个进程分配一个CPU核心和独立资源	可以使用像Celery这样的工具通过管理进程分配工作

现在我们已经探讨了各种并发策略，让我们继续通过异步编程来增强您的服务， efficiently 管理I/O绑定操作。稍后我们将重点优化计算绑定任务，特别是通过专门的服务器进行模型推理。

通过异步编程优化I/O任务

在本节中，我们将探讨如何使用异步编程来防止在AI工作负载中，I/O绑定任务阻塞主服务器进程。您还将了解asyncio框架，它使您能够在Python中编写异步应用程序。

同步与异步（Async）执行

什么是异步应用程序？为了回答这个问题，让我们比较同步和异步程序。

当任务按顺序执行，每个任务在开始前等待前一个任务完成时，应用程序被认为是同步的。对于不经常运行且处理时间只有几秒钟的应用程序，同步代码很少会导致问题，并且可以使实现变得更快且更简单。然而，如果您需要并发，并希望最大化服务在每个核心上的效率，那么您的服务应当在不等待阻塞操作完成的情况下进行多任务处理。这时，异步并发的实现就能发挥作用。

让我们看几个同步和异步函数的例子，以理解异步代码能带来多大的性能提升。在这两个示例中，我将使用sleep来模拟I/O阻塞操作，但您可以想象在实际场景中会执行其他I/O任务。

示例5-1展示了一个同步代码示例，它通过阻塞的time.sleep()函数模拟I/O阻塞操作。

示例5-1. 同步执行

import time

def task():
    print("Start of sync task")
    time.sleep(5) 
    print("After 5 seconds of sleep")

start = time.time()
for _ in range(3): 
    task()
duration = time.time() - start
print(f"\nProcess completed in: {duration} seconds")
"""
Start of sync task
After 5 seconds of sleep
Start of sync task
After 5 seconds of sleep
Start of sync task
After 5 seconds of sleep

Process completed in: 15.014271020889282 seconds
"""

• 使用sleep()来模拟I/O阻塞操作，如发送网络请求。
• 连续调用task()三次。该循环模拟了一个接一个发送多个网络请求。
• 在示例5-1中调用task()三次需要15秒钟完成，因为Python在等待阻塞操作sleep()完成时被阻塞。

要在Python中开发异步程序，您可以使用asyncio包，这是Python 3.5及更高版本标准库的一部分。使用asyncio，异步代码看起来与顺序同步代码类似，但通过添加async和await关键字来执行非阻塞的I/O操作。

示例5-2展示了如何使用async和await关键字与asyncio一起运行示例5-1的异步版本。

示例5-2. 异步执行

import time
import asyncio

async def task(): 
    print("Start of async task")
    await asyncio.sleep(5) 
    print("Task resumed after 5 seconds")

async def spawn_tasks():
    await asyncio.gather(task(), task(), task()) 

start = time.time()
asyncio.run(spawn_tasks()) 
duration = time.time() - start

print(f"\nProcess completed in: {duration} seconds")
"""
Start of async task
Start of async task
Start of async task
Task resumed after 5 seconds
Task resumed after 5 seconds
Task resumed after 5 seconds

Process completed in: 5.0057971477508545 seconds 
"""

• 实现一个任务协程，在阻塞操作时将控制权交给事件循环。
• 非阻塞的五秒睡眠会通知事件循环，在等待时执行另一个任务。
• 使用asyncio.create_task生成任务实例，使用asyncio.gather将任务链（或收集）并并发执行。
• 创建一个事件循环，使用asyncio.run方法调度异步任务。

执行时间是同步示例的1/3，因为这次Python进程并没有被阻塞。

在运行示例5-2后，您会发现task()函数是并发调用的三次。另一方面，示例5-1中的代码是顺序调用task()函数三次。异步函数在asyncio的事件循环内运行，事件循环负责在不等待的情况下执行代码。

在任何异步代码中，await关键字将I/O阻塞操作标记为非阻塞（即它们可以在不阻塞主进程的情况下执行）。通过识别阻塞操作，Python可以在等待阻塞操作完成时去做其他事情。

示例5-3展示了如何使用async和await关键字来声明和运行异步函数。

示例5-3. 如何使用async和await关键字

import asyncio

async def main():
    print("Before sleeping")
    await asyncio.sleep(3) 
    print("After sleeping for 3 seconds")

asyncio.run(main()) 

"""
Before sleeping
After sleeping for 3 seconds 
"""

• 通过等待asyncio.sleep()来模拟一个非阻塞的I/O操作，这样Python可以在等待时去做其他事情。
• 需要在asyncio.run()内部调用main()，因为它是一个异步函数。如果没有这样做，它不会被执行，而是返回一个协程对象。稍后我会详细介绍协程。

如果您运行代码，第二条语句会在第一次输出后3秒打印。在这个例子中，由于没有其他操作，Python会在等待时处于空闲状态，直到sleep操作完成。

在示例5-3中，我使用sleep来模拟I/O阻塞操作，如发起网络请求。

警告
您只能在用async def声明的函数内使用await关键字。如果在异步函数外使用await，Python会引发SyntaxError。另一个常见的陷阱是在异步函数内使用非异步的阻塞代码，这会无意中阻止Python在等待时执行其他任务。

现在，您已经理解了在异步程序中，为了防止主进程被阻塞，Python会在遇到阻塞操作时在函数之间切换。您可能会想知道：

• Python如何利用asyncio暂停和恢复函数？
• Python的asyncio使用什么机制在不忘记挂起的函数的情况下从一个函数切换到另一个？
• 函数如何在不中断其状态的情况下暂停或恢复？

为了解答这些问题，让我们深入探讨asyncio中的底层机制，理解这些问题的答案将帮助您在服务中调试异步代码。

asyncio的核心是一个称为事件循环的一级对象，它负责高效处理I/O事件、系统事件和应用程序上下文的变化。

图5-4展示了asyncio事件循环在Python中如何进行任务编排。

事件循环可以比作一个 while True 循环，它监视由Python进程中的协程函数发出的事件或消息，并在等待I/O阻塞操作完成时切换函数执行。这个编排机制允许其他函数异步执行，而不会中断。

协程函数

协程函数是专门设计的函数，它们会将控制权交还给调用者，而不会丢失它们的状态（即它们可以在稍后暂停并恢复）。将控制权交还给事件循环的机制依赖于协程。

暂停和恢复操作而不丢失状态的能力使得协程与生成器函数类似。事实上，在Python 3.5之前，协程函数没有原生支持时，您可以使用生成器来实现协程。与生成器类似，简单地调用协程函数并不会执行它。

要运行协程，您需要使用 asyncio.run() 或在Python 3.5及更高版本中通过 await 关键字来等待它们。

除了协程，asyncio 包还实现了其他并发原语，如期货（futures）、信号量（semaphores）和锁（locks）。

使用模型提供商API进行异步编程

到目前为止，我展示的三个示例都可以视为异步编程的“Hello World”示例。现在，让我们看看一个与构建GenAI服务相关的实际场景，在这个场景中，您需要使用模型提供商的API——例如OpenAI、Anthropic或Mistral——因为自己提供LLM服务可能会更加昂贵。

此外，如果您通过在短时间内发送多个请求来压力测试您在第3章中创建的生成端点，您会注意到每个请求在处理之前都有很长的等待时间。这是因为您将模型预加载并托管在与服务器运行的相同Python进程和CPU核心上。当您发送第一个请求时，整个服务器会被阻塞，直到推理工作负载完成。由于在推理过程中CPU正在尽最大努力工作，因此推理/生成过程是一个计算密集型的阻塞操作。然而，这并非必须如此。

当您使用提供商的API时，您不再需要担心计算密集型的AI工作负载，因为它们会变成I/O绑定，您将计算密集型的工作负载卸载给提供商。因此，了解如何利用异步编程与模型提供商的API并发交互是非常有意义的。

好消息是，API的所有者通常会发布同步和异步客户端以及软件开发工具包（SDK），以减少与其端点交互所需的工作量。

警告
如果您需要向其他外部服务发送请求、从数据库中获取数据或从文件中提取内容，那么您将向进程中添加其他I/O阻塞任务。如果不利用异步编程，这些阻塞任务可能会迫使服务器一直等待。

然而，任何同步代码都可以通过使用进程池或线程池执行器将其转换为异步代码，以避免在事件循环中运行任务。相反，您可以在单独的进程或线程中运行异步任务，从而避免阻塞事件循环。

您还可以通过查看库文档或源代码中是否提到了async或await关键字来验证任何异步支持。如果没有，您可以尝试测试该工具是否可以在异步函数中使用，并且当您在其上使用await时不会引发TypeError。

如果某个工具（如数据库库）只有同步实现，那么您不能使用该工具实现异步性。解决方案是将该工具切换为异步等效工具，这样您就可以使用async和await关键字来使用它们。

示例5-4中，您将通过同步和异步的OpenAI客户端与OpenAI GPT-3.5 API进行交互，以了解两者之间的性能差异。

注意
您需要安装openai库：

$ pip install openai

示例5-4. 比较同步和异步的OpenAI客户端

import os
from fastapi import FastAPI, Body
from openai import OpenAI, AsyncOpenAI

app = FastAPI()

sync_client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY"))
async_client = AsyncOpenAI(api_key=os.environ.get("OPENAI_API_KEY"))

@app.post("/sync")
def sync_generate_text(prompt: str = Body(...)):
    completion = sync_client.chat.completions.create(
        messages=[
            {
                "role": "user",
                "content": prompt,
            }
        ],
        model="gpt-3.5-turbo",
    )
    return completion.choices[0].message.content

@app.post("/async")
async def async_generate_text(prompt: str = Body(...)):
    completion = await async_client.chat.completions.create(
        messages=[
            {
                "role": "user",
                "content": prompt,
            }
        ],
        model="gpt-3.5-turbo",
    )
    return completion.choices[0].message.content

同步和异步客户端之间的区别在于，使用异步版本时，FastAPI可以在不等待OpenAI API对前一个用户输入的响应的情况下并行处理用户输入。

管理速率限制

请注意，对外部API的并发请求需要进行限制，以保持在速率限制范围内。

一种合理的策略是实现指数退避（exponential backoff），即您的服务根据指数延迟曲线来延迟访问外部API。收到的外部API的速率限制错误越多，您的服务在恢复请求之前的等待时间就会越长。

您可以申请提高等待限制，或者要求用户在再次尝试之前等待。

像stamina这样的第三方Python包可以帮助实现与外部API的速率限制策略。

通过利用异步代码，您可以大大提高吞吐量，并扩展到更大规模的并发请求。然而，编写异步（async）代码时，您必须小心。

以下是使用异步代码时可能遇到的一些常见问题和陷阱：

• 由于并发任务的非线性执行流程，理解和调试错误可能变得更加复杂。
• 一些库（如aiohttp）需要嵌套的异步上下文管理器来正确实现。这很快就会变得混乱。
• 混合使用异步和同步代码可能会抵消任何性能优势，例如，如果您忘记在函数中标记async和await关键字。
• 不使用异步兼容的工具和库也可能会消除任何性能优势；例如，使用requests库而不是aiohttp进行异步API调用。
• 忘记在任何异步函数中等待协程，或者等待非协程对象，会导致意外行为。所有的async关键字后面都必须跟await。
• 不正确地管理资源（例如打开的API/数据库连接或文件缓冲区）可能会导致内存泄漏，进而冻结计算机。如果在异步代码中不限制并发操作的数量，也可能会导致内存泄漏。
• 您还可能遇到并发和竞争条件问题，其中线程安全原则被违反，导致资源死锁，进而导致数据损坏。

这份清单并不详尽，正如您所见，使用异步编程存在许多陷阱。因此，我建议首先编写同步程序，理解代码的基本流程和逻辑，再处理迁移到异步实现的复杂性。

FastAPI中的事件循环和线程池

在底层，FastAPI可以处理异步和同步的阻塞操作。它通过在其线程池中运行同步处理程序来实现这一点，这样阻塞操作就不会阻止事件循环执行任务。

正如我在第2章中提到的，FastAPI通过Starlette运行在ASGI Web框架上。如果没有ASGI，服务器将有效地以同步方式运行，您必须等待每个进程完成后才能处理下一个请求。然而，使用ASGI，FastAPI服务器通过多线程（通过线程池）和异步编程（通过事件循环）支持并发处理，从而并行处理多个请求，同时避免主服务器进程被阻塞。

FastAPI通过在应用启动时实例化一组线程来设置线程池，从而减少线程创建的运行时负担。然后，它将后台任务和同步工作负载委派给线程池，以防止事件循环因同步处理程序中的阻塞操作而被阻塞。事件循环也被称为主FastAPI服务器线程，负责协调请求的处理。

正如我提到的，事件循环是基于asyncio构建的每个应用程序的核心组件，包括实现并发的FastAPI。事件循环运行异步任务和回调，包括执行网络I/O操作和运行子进程。在FastAPI中，事件循环还负责协调异步处理请求。

如果可能，您应在事件循环上运行处理程序（通过异步编程），因为这比在线程池上运行（通过多线程）更高效。这是因为线程池中的每个线程在执行任何代码之前必须获取GIL，而这需要一定的计算工作。

假设多个并发用户同时使用FastAPI服务的同步和异步OpenAI GPT-3.5处理程序（端点），如示例5-4所示。FastAPI将异步处理程序请求运行在事件循环上，因为该处理程序使用非阻塞的异步OpenAI客户端。另一方面，FastAPI必须将同步处理程序请求委派给线程池，以保护事件循环免受阻塞。由于将请求委派给线程并在线程池中切换线程需要更多的工作，因此同步请求的完成时间要比异步请求晚。

注意
请记住，所有这些工作——处理同步和异步处理程序请求——都在同一个FastAPI Python进程的单个CPU核心上运行。这是为了最大程度地减少等待OpenAI API响应时的CPU空闲时间。

性能差异如图5-5所示。

图5-5显示了在I/O绑定工作负载下，异步实现更快，如果您需要并发，异步编程应该是您的首选方法。然而，即使FastAPI需要与同步的OpenAI客户端一起工作，它仍然能够很好地服务多个并发请求。FastAPI仅仅是将同步的API调用发送到线程池中的线程内，以实现某种形式的并发。这也是为什么FastAPI官方文档告诉您，您不必太担心将处理程序函数声明为async def或def。

然而，请记住，当您使用async def声明处理程序时，FastAPI假设您仅执行非阻塞操作。当您违反这一假设，并在异步路由中执行阻塞操作时，事件循环将被阻塞，直到阻塞操作完成之前，它将无法继续执行任务。

阻塞主服务器

如果您在定义函数时使用async关键字，请确保您在函数内部也使用了await关键字，并且函数内部使用的所有包依赖项都不是同步的。

避免将路由处理程序函数声明为async，如果它们的实现是同步的。否则，对受影响的路由处理程序的请求将阻塞主服务器，导致服务器在等待阻塞操作完成时无法处理其他请求。不管阻塞操作是I/O绑定的还是计算绑定的，都会发生这种情况。因此，如果您不小心，对数据库或AI模型的任何调用仍然可能导致阻塞。

这是一个容易犯的错误。例如，您可能会在您声明为async的处理程序中使用同步依赖项，如示例5-5所示。

示例5-5. 在FastAPI中错误实现异步处理程序

import os
from fastapi import FastAPI
from openai import AsyncOpenAI, OpenAI

app = FastAPI()

@app.get("/block")
async def block_server_controller():
    completion = sync_client.chat.completions.create(...) 
    return completion.choices[0].message.content

@app.get("/slow")
def slow_text_generator():
    completion = sync_client.chat.completions.create(...) 
    return completion.choices[0].message.content

@app.get("/fast")
async def fast_text_generator():
    completion = await async_client.chat.completions.create(...) 
    return completion.choices[0].message.content

• I/O阻塞操作：获取ChatGPT API的响应。由于路由处理程序被标记为async，FastAPI相信我们不会运行阻塞操作，但实际上我们正在运行，这将阻塞事件循环（主服务器线程）。其他请求现在会被阻塞，直到当前请求处理完毕。
• 同步路由处理程序：执行阻塞操作，但没有利用异步特性。同步请求被交给线程池在后台运行，这样主服务器就不会被阻塞。
• 异步路由：非阻塞操作。请求不会阻塞主线程，因此不需要交给线程池处理。结果，FastAPI事件循环可以使用异步OpenAI客户端更快地处理请求。

现在，您应该更有信心在FastAPI服务中实现需要执行I/O绑定任务的新功能。

为了帮助加深对I/O并发概念的理解，在接下来的几节中，您将构建几个新的功能，利用并发将它们集成到您的FastAPI服务中。这些功能包括：

• 与网站对话：构建并集成一个网页爬虫模块，允许您通过提供HTTP URL，向自托管的LLM询问网站内容。
• 与文档对话：构建并集成一个RAG模块，将文档处理到向量数据库中。向量数据库以支持高效相似性搜索的方式存储数据。然后，您可以使用语义搜索来与上传的文档进行交互。

这两个项目将使您获得与外部系统（如网站、数据库和文件系统）异步交互的实践经验。

项目：与网站对话（网页爬虫）

公司通常会将一系列内部网页托管为HTML页面，用于手册、流程和其他文档。对于较长的页面，您的用户可能希望在提问时提供URL，并期望您的LLM能够抓取并读取内容。这时，内置网页爬虫就非常有用。

构建一个自托管LLM的网页爬虫有多种方法。根据您的用例，您可以结合以下方法：

• 获取网页作为HTML，并将原始HTML（或内部文本内容）传递给LLM，将内容解析为您所需的格式。
• 使用像BeautifulSoup和Scrapy这样的网页爬虫框架，在抓取网页后解析内容。
• 使用无头浏览器，如Selenium和Microsoft Playwright，动态导航网页中的节点并解析内容。无头浏览器非常适合导航单页应用程序（SPA）。

警告
您或您的用户应避免使用LLM驱动的网页爬虫工具用于非法目的。在从URL提取内容之前，确保您已经获得许可：

• 审查每个网站的使用条款，特别是是否提到网页爬虫。
• 尽可能使用API。
• 如果不确定，直接向网站所有者请求许可。

对于这个小项目，我们将仅抓取并将HTML页面的原始内部文本提供给LLM，因为实现一个生产级的爬虫可能会成为一本独立的书。

构建一个简单的异步爬虫的过程如下：

• 开发一个函数，使用正则表达式（regex）匹配用户向LLM提交的URL模式。
• 如果找到了URL，循环遍历提供的URL列表，异步抓取页面。我们将使用一个异步HTTP库——aiohttp，而不是requests，因为requests只能进行同步网络请求。
• 开发一个解析函数，从抓取的HTML中提取文本内容。
• 将解析后的页面内容与原始用户提示一起传递给LLM。

示例5-6展示了如何实现上述步骤。

注意
您需要安装一些额外的依赖项来运行此示例：

$ pip install beautifulsoup lxml aiohttp

示例5-6. 构建一个异步网页爬虫

# scraper.py

import asyncio
import re

import aiohttp
from bs4 import BeautifulSoup
from loguru import logger

def extract_urls(text: str) -> list[str]:
    url_pattern = r"(?P<url>https?://[^\s]+)" 
    urls = re.findall(url_pattern, text) 
    return urls

def parse_inner_text(html_string: str) -> str:
    soup = BeautifulSoup(html_string, "lxml")
    if content := soup.find("div", id="bodyContent"): 
        return content.get_text()
    logger.warning("Could not parse the HTML content")
    return ""

async def fetch(session: aiohttp.ClientSession, url: str) -> str:
    async with session.get(url) as response: 
        html_string = await response.text()
        return parse_inner_text(html_string)

async def fetch_all(urls: list[str]) -> str:
    async with aiohttp.ClientSession() as session: 
        results = await asyncio.gather(
            *[fetch(session, url) for url in urls], return_exceptions=True
        )
    success_results = [result for result in results if isinstance(result, str)]
    if len(results) != len(success_results): 
        logger.warning("Some URLs could not be fetched")
    return " ".join(success_results)

• 使用简单的正则表达式模式来捕获URLs，并将它们放入一个名为url的命名组中，匹配http和https协议。为了简单起见，这个模式匹配的是更宽松定义的URL，并不验证域名或路径的结构，也不考虑查询字符串或URL中的锚点。
• 在文本中找到所有不重叠的正则匹配项。
• 使用BeautifulSoup包解析HTML字符串。在维基百科页面中，文章内容嵌套在id为bodyContent的div容器内，因此解析逻辑假设仅传递维基百科的URL。您可以根据需要更改此逻辑，或者直接使用soup.getText()来抓取HTML中的任何文本内容。然而，请记住，如果直接解析原始HTML，解析的内容会有很多噪音，这可能会混淆LLM。
• 给定一个aiohttp会话和一个URL，执行异步的GET请求。在这个上下文管理器内，异步地等待响应。
• 给定一个URL列表，创建一个客户端会话异步上下文管理器，以异步执行多个抓取操作。由于fetch()是一个协程函数（即使用await关键字），fetch_all()将需要在asyncio.gather()中运行多个fetch()协程，以便在事件循环上进行异步执行。
• 检查所有URL是否成功抓取，如果没有，发出警告。

现在，您已经拥有了实现网页爬虫功能所需的实用爬虫函数。

接下来，升级文本到文本的处理程序，通过依赖关系以异步方式使用爬虫函数，如示例5-7所示。

示例5-7. 将网页爬虫功能作为依赖注入到FastAPI LLM处理程序

# dependencies.py

from fastapi import Body
from loguru import logger

from schemas import TextModelRequest
from scraper import extract_urls, fetch_all

async def get_urls_content(body: TextModelRequest = Body(...)) -> str: 
    urls = extract_urls(body.prompt)
    if urls:
        try:
            urls_content = await fetch_all(urls)
            return urls_content
        except Exception as e:
            logger.warning(f"Failed to fetch one or several URls - Error: {e}")
    return ""

# main.py

from fastapi import Body, Depends, Request
from dependencies import construct_prompt
from schemas import TextModelResponse

@app.post("/generate/text", response_model_exclude_defaults=True) 
async def serve_text_to_text_controller(
    request: Request,
    body: TextModelRequest = Body(...),
    urls_content: str = Depends(get_urls_content) 
) -> TextModelResponse:
    ... # 其余的控制器逻辑
    prompt = body.prompt + " " + urls_content
    output = generate_text(models["text"], prompt, body.temperature)
    return TextModelResponse(content=output, ip=request.client.host)

• 实现一个get_urls_content的FastAPI依赖项，该依赖项从请求体中获取用户提示，查找所有URL。然后它返回所有URL的内容作为长字符串。该依赖项内置了异常处理，以处理任何I/O错误，返回空字符串并在服务器上记录警告。
• 在FastAPI中使用aiohttp时，您不需要自己管理事件循环，因为FastAPI作为一个异步框架会处理事件循环。您可以将端点定义为异步函数，并在处理程序或通过依赖关系（如本例所示）中使用aiohttp进行异步HTTP请求。
• 通过FastAPI的Depends类将get_urls_content依赖项的结果注入到处理程序中。使用FastAPI依赖项保持控制器逻辑简洁、清晰且易于阅读。

现在，运行Streamlit客户端并在浏览器中尝试您全新的功能。图5-6展示了我的实验结果。

恭喜！ 您已经学会了如何构建一个简单的非阻塞网页爬虫，并与您自己的LLM配合使用。在这个小项目中，您利用了re包来匹配用户提示中的URL模式，然后使用aiohttp库异步地并发抓取多个网页。接着，您使用BeautifulSoup包解析了维基百科文章的内容，通过抓取HTML字符串中id为bodyContent的div容器内的文本内容来实现。对于其他网站或公司内部网页，您可以随时调整解析逻辑以进行适当的解析。最后，您将整个爬虫逻辑封装在一个FastAPI依赖项中，并内置了异常处理机制，以便在升级文本模型处理程序时利用依赖注入。

请记住，您的爬虫无法处理具有动态布局的复杂页面，这些页面是由服务器渲染的。在这种情况下，您可以在网页爬虫中添加一个无头浏览器来导航动态页面。

此外，抓取外部网站的内容将是一个挑战，因为大多数网站可能会实施反爬虫保护措施，如IP封锁或验证码，这些都是常见的威慑手段。维护与外部网站的数据质量和一致性也是一个持续的挑战，因为您可能需要定期更新爬虫脚本，以确保提取的内容准确可靠。

现在，您应该更有信心构建需要通过异步网络请求与网页交互的GenAI服务。

接下来，我们将探讨其他I/O异步交互，例如与数据库和文件系统的交互，通过构建“与文档对话”功能来实现。

该功能允许用户通过Streamlit界面将文档上传到您的服务。上传的文档内容随后会被提取、处理，并保存到数据库中。随后，在用户与LLM的交互过程中，一个异步检索系统从数据库中检索语义相关的内容，并将其用来增强提供给LLM的上下文。

这个过程被称为RAG，我们接下来将为您的LLM构建一个RAG模块。

项目：与文档对话（RAG）

在这个项目中，我们将构建一个RAG模块到您的GenAI服务中，以便您获得与外部系统（如数据库和文件系统）异步交互的实践经验。

您可能好奇RAG模块的目的和必要性。RAG（检索增强生成）仅仅是一种通过自定义数据源增强LLM提示上下文的技术，适用于知识密集型任务。它是一种有效的技术，可以将LLM的回答基于数据中的事实，而无需复杂和昂贵的LLM微调。

组织希望使用RAG与他们自己的LLM进行互动，因为这使得员工可以通过LLM与庞大的内部知识库进行交互。通过RAG，企业期望能够让任何需要信息的人在需要时，随时访问和利用内部知识库、系统和程序来回答问题。对于任何企业而言，访问公司的信息体系有望提高生产力，减少寻找信息的时间和成本，并提升利润。

然而，LLM容易生成不符合用户指示的回答。换句话说，LLM可能会产生幻觉（hallucinations），即生成一些不基于事实或现实的数据或信息。

这些幻觉的产生是因为模型依赖于它所训练数据中的模式，而不是直接访问外部、最新和真实的数据。LLM可能会呈现自信但错误或荒谬的回答、编造的故事或没有事实依据的声明。

因此，对于更复杂和知识密集型的任务，您会希望LLM访问外部知识源来完成任务。这可以实现更高的事实一致性，并提高生成回答的可靠性。图5-7展示了整个过程。

在这个项目中，您将为您的LLM服务构建一个简单的RAG模块，使得用户可以上传并与他们的文档进行对话。

注意
关于RAG（检索增强生成）有很多内容需要了解。这足以填满几本教科书，每天都有新的论文发布，介绍新的技术和算法。我建议您查看其他关于LLM的出版物，学习RAG过程和高级RAG技术。

RAG的流程包括以下几个阶段：

1. 从文件系统中提取文档，将文本内容以块的形式加载到内存中。
2. 对文本内容进行转换，包括清理、拆分和准备，将它们传递到嵌入模型中，以生成表示文本块语义的嵌入向量。
3. 将嵌入向量与元数据（例如源文件和文本块）一起存储到向量存储中，如Qdrant。
4. 通过对用户查询进行语义搜索，从向量存储中检索语义相关的嵌入向量。检索的向量的元数据中的原始文本块将用于增强（即增强初始提示的上下文）提供给LLM的初始提示。
5. 生成LLM响应，将查询和检索到的文本块（即上下文）传递给LLM以获取响应。

您可以在图5-8中看到完整的流程。

您可以将图5-8中显示的流程构建到现有的服务中。图5-9展示了启用RAG的“与文档对话”服务的系统架构。

图5-9概述了用户通过Streamlit界面上传的文档是如何存储并被提取以进行处理和存储到数据库中，供以后检索并增强LLM提示的。

在实现图5-9中的RAG系统之前，第一步是在Streamlit客户端和后端API中添加文件上传功能。

通过使用FastAPI的UploadFile类，您可以接受用户分块上传的文档，并将其保存到文件系统或任何其他文件存储解决方案，如Blob存储。需要注意的重要事项是，这个I/O操作是通过异步编程非阻塞的，FastAPI的UploadFile类支持这一点。

提示
由于用户可能上传大型文档，FastAPI的UploadFile类支持分块存储上传的文档，每次上传一块。这将防止您的服务的内存被堵塞。您还需要通过禁止用户上传超过某个大小的文档来保护您的服务。

示例5-8展示了如何实现一个异步文件上传功能。

提示
您需要安装aiofiles包来异步上传文件，并安装python-multipart来接收来自HTML表单的上传文件：

$ pip install aiofiles python-multipart

示例5-8. 实现异步文件上传端点

# upload.py

import os
import aiofiles
from aiofiles.os import makedirs
from fastapi import UploadFile

DEFAULT_CHUNK_SIZE = 1024 * 1024 * 50  # 50 megabytes

async def save_file(file: UploadFile) -> str:
    await makedirs("uploads", exist_ok=True)
    filepath = os.path.join("uploads", file.filename)
    async with aiofiles.open(filepath, "wb") as f:
        while chunk := await file.read(DEFAULT_CHUNK_SIZE):
            await f.write(chunk)
    return filepath

# main.py

from fastapi import FastAPI, HTTPException, status, File
from typing import Annotated
from upload import save_file

@app.post("/upload")
async def file_upload_controller(
    file: Annotated[UploadFile, File(description="Uploaded PDF documents")]
):
    if file.content_type != "application/pdf":
        raise HTTPException(
            detail=f"Only uploading PDF documents are supported",
            status_code=status.HTTP_400_BAD_REQUEST,
        )
    try:
        await save_file(file)
    except Exception as e:
        raise HTTPException(
            detail=f"An error occurred while saving file - Error: {e}",
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
        )
    return {"filename": file.filename, "message": "File uploaded successfully"}

# client.py

import requests
import streamlit as st

st.write("Upload a file to FastAPI")
file = st.file_uploader("Choose a file", type=["pdf"])

if st.button("Submit"):
    if file is not None:
        files = {"file": (file.name, file, file.type)}
        response = requests.post("http://localhost:8000/upload", files=files)
        st.write(response.text)
    else:
        st.write("No file uploaded.")

现在，您应该能够通过Streamlit UI上传文件，如图5-10所示。

实现了上传功能后，您现在可以将注意力转向构建RAG模块。图5-11展示了详细的管道，打开了图5-9中的数据转换组件。

如图5-11所示，您需要异步地从硬盘中获取存储的文件，并通过数据转换管道，然后通过异步数据库客户端存储它们。

数据转换管道包括以下部分：

• 提取器
  提取PDF内容并将其存储为文本文件，返回硬盘。
• 加载器
  异步地将文本文件分块加载到内存中。
• 清理器
  从文本块中去除多余的空白或格式化字符。
• 嵌入器
  使用预训练的自托管嵌入模型将文本转换为嵌入向量。

嵌入

嵌入是高维向量，捕捉文本中单词和短语之间的语义含义和关系。语义相似的文本之间的距离较小。

使用嵌入模型，您可以将文本和图像编码成嵌入表示，作为连续向量空间中的点。这些转换后的数据点随后可以用于多种下游应用，如RAG中的信息检索以及聚类和分类任务。

有关嵌入向量的更多信息，请参阅第3章。

一旦用户通过示例5-8中的过程将PDF文件上传到您的服务器文件系统，您可以通过pypdf库立即将它们转换为文本文件。由于没有异步库来加载二进制PDF文件，您需要先将它们转换为文本文件。

示例5-9展示了如何加载PDF，提取并处理其内容，然后将其作为文本文件存储。

注意
您将需要安装几个依赖项来运行以下示例：

$ pip install qdrant_client aiofiles pypdf loguru

示例5-9. RAG PDF到文本提取器

# rag/extractor.py

from pypdf import PdfReader

def pdf_text_extractor(filepath: str) -> None:
    content = ""
    pdf_reader = PdfReader(filepath, strict=True) 
    for page in pdf_reader.pages:
        page_text = page.extract_text()
        if page_text:
            content += f"{page_text}\n\n" 
    with open(filepath.replace("pdf", "txt"), "w", encoding="utf-8") as file: 
        file.write(content)

• 使用pypdf库以strict=True打开PDF文件，这样任何读取错误都会记录到终端。请注意，pypdf库没有异步实现，因此该函数使用正常的def关键字声明。重要的是避免在异步函数中使用此函数，以避免阻塞事件循环，从而影响主服务器线程的执行。您将看到如何通过FastAPI的后台任务来帮助解决这个问题。
• 循环遍历PDF文档的每一页，提取并将所有文本内容追加到一个长字符串中。
• 将PDF文档的内容写入文本文件，以供后续处理。指定encoding="utf-8"以避免在Windows等平台上出现问题。

文本提取器将PDF文件转换为简单的文本文件，我们可以通过异步文件加载器将其分块流入内存。每个块随后可以被清理并嵌入到嵌入向量中，使用像jinaai/jina-embeddings-v2-base-en这样的开源嵌入模型，它可以从Hugging Face模型库中下载。

注意
我选择了Jina基础嵌入模型，因为它的性能与OpenAI的专有text-embedding-ada-002模型相匹配。

示例5-10展示了包括异步文本加载器、清理器和嵌入函数的RAG数据转换管道的实现。

示例5-10. RAG数据转换函数

# rag/transform.py

import re
from typing import Any, AsyncGenerator

import aiofiles
from transformers import AutoModel

DEFAULT_CHUNK_SIZE = 1024 * 1024 * 50  # 50 megabytes

embedder = AutoModel.from_pretrained(
    "jinaai/jina-embeddings-v2-base-en", trust_remote_code=True 
)

async def load(filepath: str) -> AsyncGenerator[str, Any]:
    async with aiofiles.open(filepath, "r", encoding="utf-8") as f: 
        while chunk := await f.read(DEFAULT_CHUNK_SIZE): 
            yield chunk 

def clean(text: str) -> str:
    t = text.replace("\n", " ")
    t = re.sub(r"\s+", " ", t)
    t = re.sub(r". ,", "", t)
    t = t.replace("..", ".")
    t = t.replace(". .", ".")
    cleaned_text = t.replace("\n", " ").strip()
    return cleaned_text 

def embed(text: str) -> list[float]:
    return embedder.encode(text).tolist() 

• 下载并使用开源的jina-embeddings-v2-base-en模型，将文本字符串转换为嵌入向量。设置trust_remote_code=True以下载模型权重和分词器配置。如果没有将此参数设置为True，下载的模型权重将使用随机值初始化，而不是经过训练的值。
• 使用aiofiles库打开对文件系统中一个文件的异步连接。
• 以分块方式异步加载文本文件的内容，以提高内存效率。
• 不返回块，而是通过yield它，这样load()函数就成了一个异步生成器。异步生成器可以通过async for循环进行迭代，这样其中的阻塞操作可以被等待，从而让事件循环开始或恢复其他任务。async for循环和普通for循环都按顺序遍历可迭代对象，但async for循环允许对异步迭代器进行迭代。
• 清理文本，去除任何多余的空格、逗号、点和换行符。
• 使用Jina嵌入模型将文本块转换为嵌入向量。

一旦数据处理成嵌入向量，您可以将其存储到向量数据库中。与传统的关系型数据库不同，向量数据库专门用于处理优化语义搜索的数据存储和检索操作，它比基于关键字的搜索能返回更好的结果，后者可能会返回不理想或不完整的结果。

执行语义搜索

语义搜索使用一种数学运算，称为余弦相似度，来计算用户查询的嵌入向量与数据库中存储的嵌入向量（即嵌入文档）之间的相似度，从而检索最相关的项。

余弦相似度计算执行两个向量之间的归一化点积运算，返回一个介于-1和1之间的数字。归一化确保结果在-1和1之间。

得分为1意味着两个向量完全对齐，具有相似的语义含义，而得分为-1则意味着它们是完全对立的，表示相反的含义。

得分为0表示两个向量之间没有语义相关性，暗示这些向量可以从搜索结果中排除。

从数据库中返回的结果是一组有序的嵌入向量，这些向量包含元数据，包括原始文本。您可以将这些文本片段直接注入到用户的提示中，在转发给LLM之前，通过相关上下文来增强提示内容。

从数据库中检索到的语义搜索结果按相似度得分降序排序。这个排序，称为上下文排名，非常重要，因为研究表明，LLM对出现在提示前部的词汇比后部的词汇更敏感。

这也有道理，因为我们人类在一次阅读中，理解和注意力通常更集中在文档的前面，而不是中间或末尾。这就是为什么我们在大报告中有类似“执行摘要”这样的部分，用来向读者传达报告中最重要的内容。

以下代码示例要求你在本地机器上运行一个 qdrant 向量数据库实例，用于 RAG 模块的操作。拥有本地数据库设置将让你亲身体验如何与生产级向量数据库异步协作。为了在容器中运行该数据库，你应该在机器上安装 Docker，然后拉取并运行 qdrant 向量数据库容器。如果你不熟悉 Docker，不用担心。你将在第12章中深入学习 Docker 和容器化相关内容。

$ docker pull qdrant/qdrant 
$ docker run -p 6333:6333 -p 6334:6334 \  
    -v $(pwd)/qdrant_storage:/qdrant/storage:z \ 
    qdrant/qdrant

从 Docker 注册表中的 qdrant 仓库下载 qdrant 向量数据库镜像。

运行 qdrant/qdrant 镜像，然后将容器的端口 6333 和 6334 映射到主机机器的相同端口。

将 qdrant 数据库的存储挂载到主机的文件系统中，并将其设置为项目的根目录。

由于数据库存储和检索是 I/O 操作，你应该使用异步数据库客户端。幸运的是，qdrant 提供了一个异步数据库客户端供你使用。

提示： 你可以使用其他向量数据库提供商，如 Weaviate、Elastic、Milvus、Pinecone、Chroma 等，来替代 Qdrant。每个提供商都有一套特性和限制，你需要根据自己的使用场景做出选择。如果选择其他数据库提供商，确保有可用的异步数据库客户端供你使用。

你不必为存储和检索数据编写多个函数，可以使用第二章提到的仓储模式。在仓储模式中，你可以抽象出低级的创建、读取、更新和删除数据库操作，并使用适合你的用例的默认值。

示例 5-11 展示了如何使用仓储模式实现 Qdrant 向量数据库的客户端设置。

示例 5-11. 使用仓储模式设置向量数据库客户端

# rag/repository.py

from loguru import logger
from qdrant_client import AsyncQdrantClient
from qdrant_client.http import models
from qdrant_client.http.models import ScoredPoint

class VectorRepository: 
    def __init__(self, host: str = "localhost", port: int = 6333) -> None:
        self.db_client = AsyncQdrantClient(host=host, port=port)

    async def create_collection(self, collection_name: str, size: int) -> bool: 
        vectors_config = models.VectorParams(
            size=size, distance=models.Distance.COSINE 
        )
        response = await self.db_client.get_collections()

        collection_exists = any(
            collection.name == collection_name
            for collection in response.collections
        )
        if collection_exists: 
            logger.debug(
                f"Collection {collection_name} already exists - recreating it"
            )
            await self.db_client.delete_collection(collection_name)
            return await self.db_client.create_collection(
                collection_name,
                vectors_config=vectors_config,
            )

        logger.debug(f"Creating collection {collection_name}")
        return await self.db_client.create_collection(
            collection_name=collection_name,
            vectors_config=models.VectorParams(
                size=size, distance=models.Distance.COSINE
            ),
        )

    async def delete_collection(self, name: str) -> bool:
        logger.debug(f"Deleting collection {name}")
        return await self.db_client.delete_collection(name)

    async def create(
        self,
        collection_name: str,
        embedding_vector: list[float],
        original_text: str,
        source: str,
    ) -> None:
        response = await self.db_client.count(collection_name=collection_name)
        logger.debug(
            f"Creating a new vector with ID {response.count} "
            f"inside the {collection_name}"
        )
        await self.db_client.upsert(
            collection_name=collection_name,
            points=[
                models.PointStruct(
                    id=response.count,
                    vector=embedding_vector,
                    payload={
                        "source": source,
                        "original_text": original_text,
                    },
                )
            ],
        )

    async def search(
        self,
        collection_name: str,
        query_vector: list[float],
        retrieval_limit: int,
        score_threshold: float, 
    ) -> list[ScoredPoint]:
        logger.debug(
            f"Searching for relevant items in the {collection_name} collection"
        )
        response = await self.db_client.query_points(
            collection_name=collection_name,
            query_vector=query_vector,
            limit=retrieval_limit,
            score_threshold=score_threshold,
        )
        return response.points

使用仓储模式通过异步客户端与向量数据库进行交互。在仓储模式中，通常你会实现创建、获取、更新和删除的方法。现在，我们将实现 create_collection、delete_collection、create 和 search 方法。

向量需要存储在集合中。集合是一个命名的点集，你可以在搜索时使用。集合类似于关系数据库中的表。

通过余弦相似度计算让数据库知道，集合中的任何向量应该通过余弦距离进行比较。

在创建新集合之前检查该集合是否已经存在。否则，重新创建该集合。

设置 retrieval_limit 和 score_threshold 来限制搜索结果中的项目数量。

现在，VectorRepository 类应该能让你更容易地与数据库交互。当存储向量嵌入时，你还将存储一些元数据，包括源文档的名称、文本在源中的位置以及原始提取的文本。RAG 系统依赖这些元数据来增强 LLM 提示并向用户显示源信息。

提示： 目前，将文本转换为嵌入向量是一个不可逆的过程。因此，你需要将创建嵌入的文本与嵌入向量一起作为元数据进行存储。

接下来，你可以扩展 VectorRepository 并创建 VectorService，以便将数据处理和存储管道链式化，示例 5-12 展示了该实现。

示例 5-12. 向量数据库服务

# rag/service.py

import os

from loguru import logger
from .repository import VectorRepository
from .transform import clean, embed, load


class VectorService(VectorRepository): 
    def __init__(self):
        super().__init__()

    async def store_file_content_in_db( 
        self,
        filepath: str,
        chunk_size: int = 512,
        collection_name: str = "knowledgebase",
        collection_size: int = 768,
    ) -> None:
        await self.create_collection(collection_name, collection_size)
        logger.debug(f"Inserting {filepath} content into database")
        async for chunk in load(filepath, chunk_size): 
            logger.debug(f"Inserting '{chunk[0:20]}...' into database")

            embedding_vector = embed(clean(chunk))
            filename = os.path.basename(filepath)
            await self.create(
                collection_name, embedding_vector, chunk, filename
            )


vector_service = VectorService()

通过继承 VectorRepository 类创建 VectorService 类，以便你可以使用并扩展示例 5-11 中的常见数据库操作方法。

使用 store_file_content_in_db 服务方法异步加载、转换并以块的形式将原始文本文档存储到数据库中。

使用异步生成器 load() 异步加载来自文件的文本块。

创建 VectorService 的实例，以便在应用程序中导入并使用。

RAG 数据处理和存储管道的最后一步是将文本提取和存储逻辑作为后台任务在 file_upload_controller 中运行。实现如下所示，在处理完用户请求后，处理器将触发两个后台操作。

示例 5-13. 更新上传处理程序，将 PDF 文件内容处理并存储到向量数据库中

# main.py

from fastapi import (
    BackgroundTasks,
    FastAPI,
    File,
    UploadFile,
    status,
    HTTPException,
)
from typing import Annotated
from rag import pdf_text_extractor, vector_service

@app.post("/upload")
async def file_upload_controller(
    file: Annotated[UploadFile, File(description="A file read as UploadFile")],
    bg_text_processor: BackgroundTasks, 
):
    ... # Raise an HTTPException if data upload is not a PDF file
    try:
        filepath = await save_file(file)
        bg_text_processor.add_task(pdf_text_extractor, filepath) 
        bg_text_processor.add_task( 
            vector_service.store_file_content_in_db,
            filepath.replace("pdf", "txt"),
            512,
            "knowledgebase",
            768,
        )

    except Exception as e:
        raise HTTPException(
            detail=f"An error occurred while saving file - Error: {e}",
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
        )
    return {"filename": file.filename, "message": "File uploaded successfully"}

将 FastAPI 后台任务功能注入上传文件处理器，以便在后台处理文件上传。FastAPI 后台任务将在处理器发送响应给客户端后立即按顺序执行。

在返回响应后，后台运行 PDF 文本提取功能。由于 pdf_text_extractor 是同步函数，FastAPI 将在一个单独的线程池中运行此函数，以避免阻塞事件循环。

在 pdf_text_extractor 完成处理后，后台运行 vector_service.store_file_content_in_db 异步函数，按块加载文本内容，并将其存储在知识库向量集合中，集合接受大小为 768 的向量。

在构建完 RAG 数据存储管道之后，你现在可以专注于搜索与检索系统，它将允许你通过数据库中的知识增强 LLM 的用户提示。示例 5-14 将 RAG 搜索与检索操作与 LLM 处理程序集成，以便通过额外的上下文增强 LLM 的提示。

示例 5-14. RAG 与 LLM 服务端点的集成

# dependencies.py

from rag import vector_service
from rag.transform import embed
from schemas import TextModelRequest, TextModelResponse

async def get_rag_content(body: TextModelRequest = Body(...)) -> str: 
    rag_content = await vector_service.search( 
        "knowledgebase", embed(body.prompt), 3, 0.7
    )
    rag_content_str = "\n".join( 
        [c.payload["original_text"] for c in rag_content]
    )

    return rag_content_str


# main.py

... # other imports
from dependencies import get_rag_content, get_urls_content

@app.post("/generate/text", response_model_exclude_defaults=True)
async def serve_text_to_text_controller(
    request: Request,
    body: TextModelRequest = Body(...),
    urls_content: str = Depends(get_urls_content),
    rag_content: str = Depends(get_rag_content), 
) -> TextModelResponse:
    ... # Raise HTTPException for invalid models
    prompt = body.prompt + " " + urls_content + rag_content
    output = generate_text(models["text"], prompt, body.temperature)
    return TextModelResponse(content=output, ip=request.client.host)

为 LLM 服务端点创建 get_rag_content 依赖函数进行注入。该依赖函数可以访问请求体，从而获取用户的提示。

使用 vector_service 来搜索数据库，获取与用户提示相关的内容。当将提示传递给 vector_service.search 函数时，使用 embed 函数将用户提示转换为嵌入。仅在余弦相似度分数高于 0.7（即 70%）时，才检索最相关的三个项。

将最相关的三个项的文本负载合并为 rag_content_str 并返回。

将 get_rag_content 依赖函数的结果注入到 LLM 处理程序中，用来自向量数据库知识库的内容增强最终的 LLM 提示。现在，LLM 处理程序可以获取网页内容以及 RAG 向量数据库的内容。

如果你现在访问浏览器并上传一个 PDF 文档，你应该能够向 LLM 提问有关它的问题。图 5-12 显示了我通过上传这本书的原始样本并要求 LLM 描述我是谁来进行的实验。

注意
根据模型和输入的大小，你可能会观察到性能下降或遇到类似于令牌长度限制的问题。

恭喜！你现在已经拥有了一个完全可用的 RAG 系统，它由开源模型和向量数据库支持。

这个较长的项目作为一个动手教程，帮助你通过为 LLM 系统构建 RAG 模块，学习与异步编程、文件系统和向量数据库相关的概念。请注意，我们刚刚一起构建的 RAG 系统仍然存在许多局限性：

• 文本拆分可能会将单词一分为二，导致检索效果差和 LLM 产生混淆。
• 即使使用了增强的提示，LLM 仍然可能会产生幻觉和不一致的输出。
• 搜索和检索系统在某些情况下可能表现不佳。
• 增强后的提示可能超出 LLM 的上下文窗口。
• 从数据库中检索到的信息可能缺乏相关事实，原因可能是知识库过时或不完整、查询模糊或检索算法差。
• 检索到的上下文可能没有根据与用户查询的相关性进行排序。

你可以通过实现各种其他技术进一步改进 RAG 模块，这些技术我在本书中不会覆盖：

• 优化文本拆分、块大小、清理和嵌入操作。
• 使用 LLM 执行查询转换，通过提示压缩、链式处理、精炼和聚合等技术帮助检索和增强系统，以减少幻觉并提高 LLM 性能。
• 总结或分解大的增强提示，使用滑动窗口方法将上下文传递给模型。
• 改进检索算法，以处理模糊查询并实施用于不完整数据的回退机制。
• 使用最大边际相关性（MMR）等方法增强检索性能，通过更多样化的文档丰富增强过程。
• 实现其他高级 RAG 技术，如检索重新排序与过滤、层级数据库索引、RAG 融合、检索增强思维（RAT）等，以提高整体生成性能。

我将让你更详细地研究这些技术，并将它们作为额外的练习自行实现。

在下一节，我们将回顾优化 GenAI 服务的其他技术，以避免通过计算密集型操作（如模型推理）阻塞服务器。

优化内存和计算密集型 AI 推理任务的模型服务

到目前为止，我们已经讨论了优化服务中 I/O 密集型操作的方法。你学习了如何通过构建网页爬虫和 RAG 模块，利用异步编程与网页、数据库和文件进行交互。

通过使用异步工具和技术，在与网页、文件系统和数据库交互时，你的服务保持了响应性。然而，如果你是自托管模型，切换到异步编程技术并不会完全消除长时间的等待。这是因为瓶颈将是模型推理操作。

计算密集型操作

你可以通过在 GPU 上运行模型来加速推理，充分利用 GPU 的并行计算能力。现代 GPU 拥有惊人的计算能力，以每秒浮点运算次数（FLOPS）来衡量，现代 GPU 达到了 teraFLOPS（NVIDIA A100）或 petaflops（NVIDIA H100）的计算能力。然而，尽管现代 GPU 核心具有强大的计算能力和并行化能力，但在并发负载下，尤其是在处理更大的模型时，GPU 核心往往会处于低效状态。

当在 GPU 上自托管模型时，模型参数从磁盘加载到 RAM（I/O 密集型），然后通过 CPU 将其从 RAM 移动到 GPU 高带宽内存（内存密集型）。一旦模型参数加载到 GPU 内存中，便执行推理操作（计算密集型）。

与直觉相反，较大的 GenAI 模型，如 SDXL 和 LLM 的推理并不是 I/O 或计算密集型的，而是内存密集型的。这意味着将 1 MB 的数据加载到 GPU 计算核心所需的时间要长于这些计算核心处理 1 MB 数据的时间。为了最大化服务的并发性，你不可避免地需要将推理请求批处理，并尽可能地将最大的批量大小装入 GPU 高带宽内存。

因此，即使使用异步技术和最新的 GPU，你的服务器仍可能会在每个请求期间被阻塞，等待数十亿个模型参数加载到 GPU 高带宽内存中。为了避免服务器阻塞，你可以通过外部化模型服务，将内存密集型的模型服务操作与 FastAPI 服务器解耦，就像我们在第 3 章中提到的那样。

让我们来看一下如何将模型服务委托给另一个进程。

外部化模型服务

在外部化你的模型服务工作负载时，你有几个选择。你可以将模型托管在另一个 FastAPI 服务器上，或者使用专门的模型推理服务器。

专门的推理服务器只支持有限的 GenAI 模型架构。然而，如果你的模型架构受到支持，你将节省大量时间，而无需自己实现推理优化。例如，如果你需要自托管 LLM，LLM 服务框架可以为你执行多个推理优化，如批处理、张量并行化、量化、缓存、流式输出、GPU 内存管理等。

模型优化技术

你可以尝试几种模型压缩技术，以提高推理性能：

• 量化：压缩模型
• 剪枝：总结模型参数
• 蒸馏：创建比“教师”模型小得多的“学生”模型
• 微调小型模型：使其专门化，以适应你的使用场景

如果你正在服务基于 Transformer 的模型，可以使用以下技术进一步优化模型推理：

• 快速注意力（Fast Attention） ：优化 GPU 上的注意力图计算
• KV 缓存（KV Caching） ：利用内存缓存技术，通过重用计算过的注意力图结果来加速推理
• 分页注意力（Paged Attention） ：在注意力计算后优化 KV 缓存内存使用
• 请求批处理（Request Batching） ：包括简单批处理和连续批处理，以最大化 GPU 使用率

第 10 章将深入探讨更多的模型优化技术。

由于我们在本章中主要处理的是 LLM，我将向你展示如何集成 vLLM，这是一款开源 LLM 服务器，可以为你启动一个符合 OpenAI API 规范的 FastAPI 服务器。vLLM 还与流行的开源 Hugging Face 模型架构（包括 GPT、Llama、Gemma、Mistral、Falcon 等）无缝集成。

注意
在撰写本文时，你可以使用的其他 LLM 托管服务器包括 NVIDIA Triton 推理服务器、Ray Serve、Hugging Face 推理和 OpenLLM 等。
每种服务器都有其功能、优缺点以及支持的模型架构。我建议在将这些服务器应用于你的用例之前，先进行相关研究。

你可以通过一个命令启动自己的 vLLM FastAPI 服务器，如示例 5-15 所示。要运行示例 5-15 中的代码，你需要先安装 vLLM：

$ pip install vllm

警告
在撰写本文时，vLLM 仅支持 Linux 平台（包括 WSL），并要求 NVIDIA 兼容的 GPU 来运行 CUDA 工具包依赖项。不幸的是，你无法在 Mac 或 Windows 机器上安装 vLLM 进行本地测试。
vLLM 旨在支持在 Linux 环境中运行的生产推理工作负载，在此环境下，服务器可以通过张量并行性将请求分配给多个 GPU 核心。它也支持通过 Ray Serve 依赖项进行分布式计算，以便在超出单一机器时扩展服务。
有关分布式推理和服务的更多详情，请参阅 vLLM 文档。

示例 5-15. 在配备 4 个 16GB NVIDIA T4 GPU 的 Linux 机器上启动 vLLM FastAPI OpenAI API 服务器，使用 TinyLlama 模型

$ python -m vllm.entrypoints.openai.api_server \ 
--model "TinyLlama/TinyLlama-1.1B-Chat-v1.0" \
--dtype float16 \ 
--tensor-parallel-size 4 \ 
--api-key "your_secret_token"

• 启动一个与 OpenAI 兼容的 API 服务器，使用 FastAPI 提供 TinyLlama 模型服务。
• 使用 float16 中等精度数据类型。float16 兼容 GPU 硬件，而 bfloat16 通常兼容 CPU 硬件。
• 利用 vLLM 的张量并行特性，在 4 个 GPU 上运行 API 服务器。
• 设置一个密钥令牌以进行基本身份验证，保护 LLM 服务器。这对安全的机器对机器通信很有用，例如，直接与当前的 FastAPI 服务进行通信。

示例 5-16. 用异步 API 调用替换当前服务中的模型服务逻辑，连接到新的 vLLM 服务器

# models.py

import os
import aiohttp
from loguru import logger

async def generate_text(prompt: str, temperature: float = 0.7) -> str:
    system_prompt = "You are an AI assistant"
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": prompt},
    ]
    data = {"temperature": temperature, "messages": messages}
    headers = {"Authorization": f"Bearer {os.environ.get('VLLM_API_KEY')}"}
try:
   async with aiohttp.ClientSession() as session: 
        response = await session.post(
            "http://localhost:8000/v1/chat", json=data, headers=headers
        )
        predictions = await response.json()
except Exception as e:
    logger.error(f"Failed to obtain predictions from vLLM - Error: {e}")
    return (
        "Failed to obtain predictions from vLLM - "
        "See server logs for more details"
    )
try:
    output = predictions["choices"][0]["message"]["content"] 
    logger.debug(f"Generated text: {output}")
    return output
except KeyError as e:
    logger.error(f"Failed to parse predictions from vLLM - Error: {e}")
    return (
        "Failed to parse predictions from vLLM - "
        "See server logs for more details"
    )

• 使用 aiohttp 创建一个异步会话，向 vLLM FastAPI 服务器发送 POST 请求。此逻辑替换了当前 FastAPI 服务器中的 Hugging Face 模型管道推理逻辑。
• 由于 vLLM 服务器与 OpenAI 兼容，你可以按照 OpenAI API 规范来访问输出内容。

示例 5-17. 移除 FastAPI 生命周期并更新文本生成处理程序为异步

# main.py

from fastapi import FastAPI, Request
from schemas import TextModelRequest, TextModelResponse
from models import generate_text

# Remove the asynccontextmanager to remove TinyLlama from FastAPI 
# @asynccontextmanager
# async def lifespan(app: FastAPI):
#     models["text"] = load_text_model()
#     yield
#     models.clear()

# Remove the `lifespan` argument from `FastAPI()`
app = FastAPI()

@app.post("/generate/text")
async def serve_text_to_text_controller(
    request: Request, body: TextModelRequest
) -> TextModelResponse: 
    ...  # controller logic
    output = await generate_text(body.prompt, body.temperature)
    return TextModelResponse(content=output, ip=request.client.host)

• 现在无需使用 FastAPI 生命周期，因为模型由外部的 vLLM FastAPI 服务器提供服务。
• 将 serve_text_to_text_controller 更新为异步路由处理程序，因为它现在执行的是向 vLLM 服务器发送 I/O 操作，而不是运行同步的计算密集型模型推理操作。

恭喜，你已经成功实现了 AI 推理工作负载的并发性。通过将 LLM 推理工作负载转移到另一台服务器，你在单机上实现了某种形式的多进程。现在，两个服务器分别运行在不同的核心上，LLM 服务器通过并行化将工作分配给多个 GPU 核心。这意味着你的主服务器现在可以处理多个传入请求，并执行比单次处理一个 LLM 推理操作更多的任务。

提示
请记住，迄今为止，你实现的并发性仅限于单一机器。
为了支持更多的并发用户，你可能需要更多的机器，配备 CPU 和 GPU 核心。在那时，像 Ray Serve 和 Kubernetes 这样的分布式计算框架可以帮助你使用并行化技术扩展和协调服务。
在集成 vLLM 之前，你会发现请求之间的等待时间很长，因为主服务器忙于运行推理操作。而通过 vLLM，现在延迟大大减少，LLM 服务的吞吐量也得到了显著提高。

延迟和吞吐量

在 LLM 的背景下，延迟指的是从向模型发出请求到接收到首次响应所花费的时间。这是衡量单个请求处理过程中延迟的标准。另一方面，吞吐量是 LLM 在给定时间框架内可以处理的请求数量，它表示服务器处理并发或顺序请求的能力。延迟通常以延迟秒数来衡量，而吞吐量则以每分钟处理的令牌数（TPM）来衡量。

作为开发者，你肯定希望你的服务拥有尽可能低的延迟和最高的吞吐量。然而，模型的大小和质量与这两个指标之间存在权衡关系。通常，拥有更多参数的 LLM 可以实现更高的质量，但也会增加延迟并减少吞吐量。

目前，研究人员正在使用模型压缩技术，如蒸馏、量化和剪枝，来保持语言模型小型化，同时在 AI 推理服务中保持高质量、高吞吐量和低延迟。

除了像量化这样的模型压缩机制，vLLM 还使用其他优化技术，包括连续请求批处理、缓存分区（分页注意力）、通过内存共享减少 GPU 内存占用，以及流式输出，以实现更小的延迟和更高的吞吐量。

让我们更详细地了解请求批处理和分页注意力机制，以理解如何进一步优化 LLM 推理。

请求批处理和连续批处理

正如我们在第 3 章中讨论的那样，LLM 以自回归的方式生成下一个令牌预测，如图 5-13 所示。

这意味着 LLM 必须在一个循环中执行多个推理迭代，以生成响应，每次迭代产生一个单独的输出令牌。随着每次迭代输出令牌的生成，输入序列会不断增长，新的序列会在下次迭代时传递给模型。一旦模型生成了序列结束令牌，生成循环就会停止。本质上，LLM 会生成一系列完成令牌，直到生成一个停止令牌或达到最大序列长度为止。

LLM 必须为序列中的每个令牌计算多个注意力图，以便逐步做出下一个令牌的预测。

幸运的是，GPU 可以并行化每次迭代的注意力图计算。正如你所了解的，这些注意力图捕捉了输入序列中每个令牌的意义和上下文，并且计算起来非常昂贵。因此，为了优化推理，LLM 使用键值（KV）缓存，将计算出的图存储在 GPU 内存中。

提示
注意力图公式基于给定的查询（Q）和键（K）计算一个值（V）：

Q = KV

这个计算必须为序列中的每个令牌进行，但幸运的是可以使用 GPU 上的大矩阵乘法操作对其进行向量化。

然而，在 GPU 内存中存储参数以便在迭代之间重用会消耗大量 GPU 内存。例如，一个 13B 参数的模型在每个令牌上消耗接近 1 MB 的状态内存，再加上所有 13B 模型参数。这意味着你可以存储在内存中供重用的令牌数量有限。

如果你使用的是高端 GPU，例如具有 40 GB 内存的 A100，你只能一次存储 14K 个令牌，而剩余的内存用于存储 26 GB 的模型参数。简而言之，GPU 内存的消耗随着基础模型的大小和令牌序列的长度而增加。

更糟糕的是，如果你需要通过批处理请求来同时服务多个用户，GPU 内存必须在多个 LLM 推理之间共享。因此，你有更少的内存来存储更长的序列，LLM 会被限制在较短的上下文窗口中。另一方面，如果你想保持较大的上下文窗口，那么你将无法处理更多的并发用户。例如，一个 2048 长度的序列意味着你的批量大小将限制为 7 个并发请求（或 7 个提示序列）。实际上，这是一个上限限制，并且没有为存储中间计算留出空间，这将进一步减少上述数字。

这一切意味着，LLM 无法充分利用 GPU 的可用资源。主要原因是 GPU 的大量内存带宽被加载模型参数所消耗，而不是用于处理输入。

减少服务负载的第一步是集成最有效的模型。通常，较小且压缩的模型可以完成你要求的任务，并且具有与较大模型类似的性能。

解决 GPU 未充分利用问题的另一个合适方案是实施请求批处理，其中模型以组的形式处理多个输入，从而减少每次请求加载模型参数的开销。这种方法更高效地利用了芯片的内存带宽，导致更高的计算利用率、更高的吞吐量，并且降低了 LLM 推理的成本。像 vLLM 这样的 LLM 推理服务器通过批处理、快速注意力、KV 缓存和分页注意力机制来最大化吞吐量。

你可以在图 5-14 中看到启用和不启用批处理时响应延迟和吞吐量的差异。

实现批处理有两种方式：

静态批处理

批量大小保持不变。

动态或连续批处理

批量大小根据需求决定。

在静态批处理中，我们等待预定数量的传入请求到达，然后再将它们批量处理。然而，由于请求可以在批次中的任何时候完成，我们实际上是在延迟每个请求的响应——并增加延迟——直到整个批次处理完毕。

处理批次并将新的请求添加到批次中时，释放 GPU 资源也可能是一个难题，因为这些新请求可能处于不同的完成状态。因此，由于批次内生成的序列长度不同，并且与该批次中最长的序列长度不匹配，GPU 资源会保持未充分利用。

图 5-15 说明了在 LLM 推理中静态批处理的情况。

在图 5-15 中，你会注意到白色的块代表了未充分利用的 GPU 计算时间。在整个批处理的处理时间线中，批次中的只有一个输入序列使 GPU 得到了充分利用。

除了增加不必要的等待时间和未能充分利用 GPU 外，静态批处理的问题还在于，LLM 驱动的聊天服务的用户并不会提供固定长度的提示或期望固定长度的输出。生成输出的变化可能会导致 GPU 大量未被充分利用。

一种解决方案是避免假设固定的输入或输出序列，而是在处理批次时设置动态批量大小。在动态或连续批处理中，批量大小可以根据传入请求的序列长度和可用的 GPU 资源来设置。采用这种方法时，新的生成请求可以通过替换已完成的请求插入到批次中，从而比静态批处理更高效地利用 GPU。

图 5-16 展示了动态或连续批处理如何能够充分利用 GPU 资源。

在模型参数加载时，请求可以继续流入，LLM 推理服务器将其调度并插入到批处理中，以最大化 GPU 的使用。这种方法能提高吞吐量并减少延迟。

如果你正在构建一个 LLM 推理服务器，可能会希望将连续批处理机制集成到服务器中。然而，好消息是，vLLM 服务器已经内置了连续批处理功能，并通过其 FastAPI 服务器提供，所以你不必自己实现这些功能。此外，vLLM 还具备另一个重要的 GPU 优化特性，这使它与其他 LLM 推理框架有所不同：分页注意力。

分页注意力

高效的内存使用对于处理高吞吐量服务的系统，尤其是 LLM 来说，是一个关键挑战。为了加速推理，现代模型依赖 KV 缓存来存储和重用注意力图，这些图随着输入序列长度的增加而呈指数增长。

分页注意力是一种新颖的解决方案，旨在最小化这些 KV 缓存的内存需求，从而提高 LLM 的内存效率，并使其更适合在资源有限的设备上使用。在基于 Transformer 的 LLM 中，为了捕捉必要的上下文，每个输入令牌会生成注意力键和值张量。与其在每一步都重新计算这些张量，不如将它们作为 KV 缓存保存到 GPU 内存中，这样它就成了模型的内存。然而，KV 缓存可以增长到巨大的尺寸，例如，13B 参数模型的 KV 缓存可能达到 40 GB，这对高效存储和访问提出了重大挑战，特别是在硬件资源有限的情况下。

分页注意力引入了一种方法，将 KV 缓存分解为较小、更易管理的部分，这些部分被称为页面，每个页面包含一组令牌的 KV 向量。通过这种分段，分页注意力能够在注意力计算过程中高效加载和访问 KV 缓存。你可以将这种技术与操作系统如何管理虚拟内存进行比较，其中数据的逻辑安排与其物理存储是分离的。本质上，一个块表将逻辑块映射到物理块，从而允许在处理新令牌时动态分配内存。核心思想是通过利用逻辑块（而非物理块）来避免内存碎片，并使用映射表快速访问存储在分页物理内存中的数据。

你可以将分页注意力机制分解为几个步骤：

1. 分割 KV 缓存
   缓存被分割成固定大小的页面，每个页面包含一部分键值对。
2. 构建查找表
   创建一个表，将查询键映射到对应的页面，便于快速分配和检索。
3. 选择性加载
   在推理过程中，仅加载当前输入序列所需的页面，减少内存占用。
4. 注意力计算
   模型使用加载页面中的键值对来计算注意力。这个方法旨在通过解决内存瓶颈，使 LLM 更加易于使用，可能使其能够在更多设备上部署。

上述步骤使得 vLLM 服务器能够通过映射物理和逻辑内存块最大化内存使用效率，从而在生成过程中高效地存储和检索 KV 缓存。

在 Anyscale.com 发布的博客文章中，作者研究并比较了各种 LLM 服务框架在推理过程中的性能。作者得出结论，利用分页注意力和连续批处理机制在优化 GPU 内存使用方面非常强大，以至于 vLLM 服务器能够将延迟减少 4 倍，吞吐量提高至 23 倍。

在下一节中，我们将关注那些处理时间较长且计算密集型的 GenAI 工作负载。这通常适用于像 SDXL 这样的较大非 LLM 模型，在为多个用户执行批量推理（如批量图像生成）时可能会面临挑战。

管理长时间运行的 AI 推理任务

通过将模型托管在 FastAPI 事件循环之外的独立进程中，你可以将注意力转向那些需要长时间才能完成的阻塞操作。

在上一节中，你利用了像 vLLM 这样的专用框架来外部托管并优化 LLM 的推理工作负载。然而，你可能仍然会遇到一些需要显著时间来生成结果的模型。为了防止用户长时间等待，你应该管理那些生成模型并需要较长时间才能完成的任务。

一些 GenAI 模型，例如 Stable Diffusion XL，即使在 GPU 上运行，也可能需要几分钟才能生成结果。在大多数情况下，你可以要求用户等待直到生成过程完成。但是，如果多个用户同时使用同一个模型，服务器就必须排队处理这些请求。当用户与生成模型交互时，他们需要多次与模型互动，以引导模型生成他们想要的结果。这种使用模式会造成大量的请求积压，排队的用户可能需要等待很长时间才能看到结果。

如果能有一种方法在不让用户等待的情况下处理长时间运行的任务，那就太完美了。幸运的是，FastAPI 提供了一个解决此类问题的机制。

FastAPI 的后台任务是你可以利用的机制，在模型忙于处理请求时，你仍然可以响应用户。你在构建 RAG 模块时，已经简要了解了这个功能，当时后台任务正在将上传的 PDF 文档内容填充到向量数据库中。

通过使用后台任务，用户可以继续发送请求或继续他们的工作，而无需等待。你可以将结果保存到磁盘或数据库中以供稍后检索，或者提供一个轮询系统，允许客户端在模型处理请求时查询更新。另一个选择是创建客户端和服务器之间的实时连接，以便在结果可用时立即更新用户界面。所有这些解决方案都可以通过 FastAPI 的后台任务来实现。

示例 5-18. 使用后台任务处理长时间运行的模型推理（例如，批量生成图像）

# main.py

from fastapi import BackgroundTasks
import aiofiles

...

async def batch_generate_image(prompt: str, count: int) -> None:
    images = generate_images(prompt, count) 
    for i, image in enumerate(images):
        async with aiofiles.open(f"output_{i}.png", mode='wb') as f:
            await f.write(image) 

@app.get("/generate/image/background")
def serve_image_model_background_controller(
    background_tasks: BackgroundTasks, prompt: str, count: int 
):
    background_tasks.add_task(batch_generate_image, prompt, count) 
    return {"message": "Task is being processed in the background"} 

• 使用外部模型服务 API（如 Ray Serve）批量生成多个图像。
• 循环处理生成的图像，并使用 aiofiles 库异步将每个图像保存到磁盘。在生产环境中，你还可以将输出图像保存到云存储解决方案中，客户端可以直接从中获取。
• 启用控制器执行后台任务。
• 将 batch_generate_image 函数定义传递给 FastAPI 背景任务处理程序，并提供所需的参数。
• 在处理后台任务之前，返回一个通用的成功消息给客户端，以便用户不必等待。

在示例 5-18 中，你让 FastAPI 在后台运行推理操作（通过外部模型服务器 API），这样事件循环保持不被阻塞，可以处理其他传入的请求。你甚至可以在后台运行多个任务，例如批量生成图像（在不同进程中）和发送通知电子邮件。这些任务被添加到队列中并按顺序处理，不会阻塞用户。然后，你可以存储生成的图像并暴露一个额外的端点，供客户端轮询状态更新并检索推理结果。

警告

后台任务运行在相同的事件循环中。它们不会提供真正的并行性；它们只提供并发性。如果你在后台任务中运行像 AI 推理这样的 CPU 密集型操作，它将阻塞主事件循环，直到所有后台任务完成。类似地，使用异步后台任务时要小心。如果你没有等待阻塞的 I/O 操作，该任务将阻止主服务器响应其他请求，即使它是在后台运行。FastAPI 会在内部线程池中运行非异步后台任务。

尽管 FastAPI 的后台任务是处理简单批处理任务的绝佳工具，但它无法很好地扩展，且无法像专用工具那样处理异常或重试。其他 ML 服务框架，如 Ray Serve、BentoML 和 vLLM，可能通过提供如请求批处理等功能，在大规模部署时更好地处理模型服务。更复杂的工具，如 Celery（一个队列管理器）、Redis（一个缓存数据库）和 RabbitMQ（一个消息代理），也可以结合使用，以实现更强大和可靠的推理管道。

总结

本章探讨了在 AI 系统中应用并发的复杂方面。
你了解了并发和并行的概念，包括几种阻塞操作类型，这些操作会阻止你同时为用户提供服务。你发现了并发技术，如多线程、进程间并行和异步编程，以及它们在不同用例中的区别、相似性、优缺点。

接下来，你学习了线程池和事件循环，尤其是在 FastAPI 服务器环境中的应用，并理解了它们在并发处理请求中的作用。这包括理解如果你不小心声明路由处理程序，为什么服务器可能会被阻塞。

随后，你了解到如何实现异步编程来管理 I/O 阻塞操作。通过动手示例，你加深了对与数据库和网页内容进行异步交互的理解，构建了一个网页爬虫和一个 RAG 模块。

此外，你了解了为什么较大的 GenAI 模型可能会占用大量内存，并导致内存密集型的阻塞操作。在这部分内容中，你还学习了内存优化技术，如连续批处理和分页注意力，在为 LLM 提供服务时最小化与内存相关的瓶颈。

最后，你了解了如何处理长时间运行的 AI 推理过程，确保你的服务在长时间操作中保持响应。

通过本章的知识，你现在已经准备好将并发原理应用到自己的服务中，构建出具有弹性、可扩展性和高性能的 AI 应用程序。

能够同时处理多个用户是一个重要的里程碑。但你仍然可以进行更多优化，进一步提升 GenAI 服务的用户体验。你可以通过流式技术提供实时更新，逐步展示生成过程中的近实时结果给用户。这对在对话场景中可能具有更长生成时间的 LLM 特别有用。

接下来的章节将探讨 AI 流式工作负载，详细介绍如何使用实时通信技术，如服务器推送事件（SSE）和 WebSocket（WS）。你将学习这些技术之间的区别，并通过构建实时文本到文本、文本到语音、语音到文本的端点实现模型流式传输。

额外参考文献

• Kwon, W., et al. (2023). “Efficient Memory Management for Large Language Model Serving with PagedAttention”. arXiv preprint arXiv:2309.06180.
• Lewis, P., et al. (2022). “Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks”. arXiv preprint arXiv:2005.11401.