用Locust对FastAPI ML API进行性能测试的方法

MLOps知识已经成为一个机器学习工程师可以拥有的主要技能之一。然而，将机器学习模型成功投入生产并不是一件容易的事。除了对数据科学的理解，它还需要广泛的软件开发和DevOps能力。简而言之，为了提高你作为机器学习工程师的价值，你不仅要了解如何将各种机器学习和深度学习模型应用于特定的问题，还要了解如何测试、验证和部署它们。

拥有能够将机器学习模型投入生产的人，已经成为任何企业的一大优势。当涉及到将机器学习模型投入生产时，最后的问题之一是验证为该模型服务的API是否具有良好的性能。在这篇文章中，我们主要讨论如何用FastAPI训练和部署模型，以及如何用Locust测试这个API的性能。

在这篇文章中，我们涵盖了：

0.先决条件

1.负载测试

2.训练和部署机器学习模型

3.使用Locust进行性能测试

先决条件

要成功运行本教程中的例子，需要安装Python或更高版本。 Python 3.6 或更高版本需要安装。最简单的方法是使用 Anaconda 分发。它带有本教程所需的所有其他库，如Pandas、NumPy、SciKit Learn等。为了训练模型，我们使用SciKit Learn。

使用下面的命令来安装FastAPI和它的所有依赖项。

这包括Uvicorn，一个运行你的代码的ASGI服务器。如果你对其他的ASGI服务器（如Hypercorn） 更满意，那也是可以的，你可以在本教程中使用它。

最后，要安装本教程中使用的Locust测试框架，请使用此命令。

数据集

我们在本文中使用的数据是来自 PalmerPenguins数据集.这个数据集最近被介绍为著名的Iris数据集的替代品。它是由Kristen Gorman博士和南极洲LTER的Palmer站创建的。你可以获得这个数据集 这里，或通过Kaggle。

这个数据集本质上由两个数据集组成，每个数据集都包含344只企鹅的数据。就像在Iris数据集中，有3种不同的企鹅来自帕尔默群岛的3个岛屿。此外，这些数据集还包含每个物种的culmen尺寸。喙是鸟类喙的上脊。在简化的企鹅数据中， culmen的长度和深度被重新命名为culmen_length_mm和culmen_depth_mm两个变量。

1.负载测试

负载测试是一种非功能性的软件测试，在这种测试中，软件应用程序的性能在一定的负载下被评估。它决定了当许多人同时访问它时，软件应用程序是如何运行的，或者当一个API有大量的请求时，它的表现如何。负载测试用于识别性能瓶颈，确保应用程序的稳定性和平稳运行。

在将机器学习API发布到生产中之前，负载测试可以确定以下问题：

• 每个事务的响应时间
• 系统组件在不同负载下的性能
• 数据库组件在不同负载下的性能
• 网络延迟
• 软件设计方面的问题
• 网络服务器、应用服务器、数据库服务器和其他服务器配置困难。
• 硬件方面的问题

负载测试将表明系统是否需要进行微调，或者是否需要对硬件和软件进行修改以提高性能。

2.训练和部署机器学习模型

在本教程中，我们使用RandomForest来对PalmerPenguins数据集进行分类。我们并没有深入研究这个算法的功能以及你应该如何训练这个算法。本章只是简单地介绍了所有的概念，因为重点是性能测试。如果你想了解更多，可以在这里下载代码并查看 我们的指南.

2.1 训练分类模型

正如我们提到的分类例子，我们使用PalmerPenguins数据集。然而，由于我们想做二元分类，我们需要做一些准备工作。首先，我们加载数据集，删除我们在本文中不使用的特征，并删除一个类，因为我们进行二元分类。

data = pd.read_csv('./data/penguins_size.csv')
data.head()

2.1.1 准备数据

我们删除所有被标记为 "Chinstrap"类的样本和我们不想使用的特征（我们只使用culmen_length_mm 和culmen_depth_mm）。

data = data.dropna()
data = data.drop(['sex', 'island', 'flipper_length_mm', 'body_mass_g'], axis=1)
data = data[data['species'] != 'Chinstrap']

然后，我们将输入数据分开，并将其缩放。

X = data.drop(['species'], axis=1)
X = X.values
ss = StandardScaler()
X = ss.fit_transform(X)

之后，我们提取输出值并将其标记为-1和1，因为SVM算法需要这样做。

y = data['species']
spicies = {'Adelie': -1, 'Gentoo': 1}
y = [spicies[item] for item in y]
y = np.array(y) 

另外，我们删除了数据集中的第182个样本。为什么？好吧，这个样本处于类之间，它有点扰乱了我们试图做出的点，所以我们删除它并将数据分成训练和测试数据集。

# Remove sample that is too close
X = np.delete(X, 182, axis=0)
y = np.delete(y, 182, axis=0)

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=33)

就这样，我们的数据集已经准备好了，下面是我们的数据集的样子 情节它。

plt.figure(figsize=(11, 5))
plt.scatter(X[y == 0][:, 0], X[y == 0][:, 1], color='orange', label='Adelie')
plt.scatter(X[y == 1][:, 0], X[y == 1][:, 1], color='gray', label='Gentoo')
plt.legend();

2.1.2 训练随机森林

Sci-Kit学习库提供了一个伟大的随机森林算法的实现。下面是我们如何使用RandomForestClassifier。

rf_classifier = RandomForestClassifier(n_estimators=11, max_leaf_nodes=16, n_jobs=-1)
rf_classifier.fit(X_train, y_train)

rf_preditions = rf_classifier.predict(X_test)
print(accuracy_score(rf_preditions, y_test))

让我们来看看随机森林的分类图。

好的，我们的随机森林算法根据训练数据为两个类别创建了两个区域。看来，它做得很好。

2.2.用FastAPI部署机器学习模型

使用机器学习的系统通常使用微服务架构构建。从机器学习的角度来看，这意味着有一个服务执行模型的训练，另一个服务使用它。这样一来，我们就有了明确的关注点分离。这两个微服务之间的通信是通过一些存储模型的文件完成的。类似这样的事情。

2.2.1 将模型保存到一个文件中

下面是我们在使用Sci-Kit Learn时如何将模型保存到文件中。我们需要在训练文件的末尾添加这个。

from joblib import dump

Save model into file.
dump(model, ' ./static/model/classifier.joblib')

2.2.2 使用FastAPI的休息API

让我们建立一个API，用FastAPI加载和服务这个模型。首先，我们建立一个负责加载训练后的模型的cl组件。通过构造函数，它接收用户选择的算法的信息，并从那里开始处理它。下面是它的样子。

from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score
from penguin_sample import PenguinSample
from sklearn.preprocessing import StandardScaler
from joblib import load
import numpy as np

class ModelLoader():
    '''Class for loading models into memory.'''

    def __init__(self, algoritm):
        self.scaledData = False
        self.model = load('./models/random_forest.joblib')

        self.scaler = StandardScaler()

    def prepare_sample(self, raw_sample: PenguinSample):
        '''Prepare new sample so it can be processed by the model.'''
        island = self._island_map[raw_sample.island]
        sex = self._sex_map[raw_sample.sex]

        sample = [raw_sample.culmenLength, raw_sample.culmenDepth, raw_sample.flipperLength, raw_sample.bodyMass, island, sex]
        sample = np.array([np.asarray(sample)]).reshape(-1, 1)

        if(self.scaledData):
            self.scaler.fit_transform(sample)

        return sample.reshape(1, -1)

    def predict(self, data: PenguinSample):
        '''Predict class of the new sample.'''
        prepared_sample = self.prepare_sample(data)
        return self.model.predict(prepared_sample)

最重要的文件是main.py文件。这个文件将所有其他部分放在一起，并使用FastAPI构建REST API。下面是它的样子。

from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from model_loader import ModelLoader
from train_parameters import TrainParameters
from penguin_sample import PenguinSample

origins = [
    "http://localhost:8000",
    "http://127.0.0.1:8000/predict",
    "http://127.0.0.1:8000/load",
    "http://localhost:4200"
]

app = FastAPI()
app.add_middleware(
    CORSMiddleware,
    allow_origins=origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Initially load SVM Model.
app.model = ModelLoader('svm')

@app.post("/load")
async def load(params: TrainParameters):
    '''API Route for loading defined model into memory.'''
    print("Model Loading Started")
    app.model =  ModelLoader(params.model.lower())
    return True

@app.post("/predict")
async def predict(data:PenguinSample):
    '''API Route for making predictions based on the entered information.'''
    print("Predicting")
    spicies_map = {0: 'Adelie', 1: 'Chinstrap', 2: 'Gentoo'}
    species = app.model.predict(data)
    return spicies_map[species[0]]

predict方法从Web应用程序的predict标签中接收数据。它从ModelLoader中调用预测方法，并返回预测值。要运行这个应用程序，请执行以下命令。

uvicorn main:app --port 8000 

我们可以通过浏览器访问我们的API，在那里我们可以使用Swagger来测试/predict端点。

3.用Locust进行性能测试

Locust是一个性能测试工具，使用简单，可编写脚本，并且可扩展。你不需要被困在用户界面或有限的领域专业语言中，而是用标准的Python代码描述你的用户行为。因此，Locust具有无限的可扩展性和开发者友好性。

3.1 Locust的主要组成部分

从本质上讲，Locust测试是一个Python应用程序。这给了它很大的灵活性，而且它特别擅长处理复杂的用户流。有三个主要组件，我们用来创建测试：

• task- 这是一个描述符，我们用它来定义哪个方法是Locust服务器应该运行的测试方法。
• HttpUser - 这是一个核心或解决方案。这个类应该继承自我们的测试类。它为每个用户提供了一个客户端属性，它是HttpSession的一个实例，可以用来执行对负载测试目标系统的HTTP请求。
• between(m,n) - 使模拟用户在每个任务执行后等待m和n秒之间。

3.2 编写蝗虫测试

下面是蝗虫测试文件的样子。

import json
from locust import HttpUser, task, between
from penguin_sample import PenguinSample

class PerformanceTests(HttpUser):
    wait_time = between(1, 3)

    @task(1)
    def testFastApi(self):
        sample = PenguinSample(island="Torgersen", 
                               culmenLength=39.1, 
                               culmenDepth=18.7, 
                               flipperLength=181, 
                               bodyMass=3750, 
                               sex="MALE", 
                               species="Adelie")
        headers = {'Accept': 'application/json', 'Content-Type': 'application/json'}
        self.client.post("/predict", data=json.dumps(sample.dict()), headers=headers)

让我们把它分解一下。在必要的导入之后，我们创建了一个PerformanceTests类，它继承了locust模块中的HttpUser。这样，我们就能与我们实现的API进行交互。这个类有一个方法testFastApi，这是我们的测试任务。

这个任务创建了一个新的PenguinSample的样本，我们想得到它的预测结果。然后我们为我们的发布方法定义头信息。最后，我们使用客户端 来发送一个post方法。注意，我们只是定义端点，即"/predict*"，* 而不是API的完整地址。另外，我们将样本作为数据发送。

一般来说，这是一个简单的例子，只有一个任务。我们可以根据自己的需要定义更多的任务。

3.3 使用用户界面运行蝗虫测试

为了运行先前定义的文件，我们可以使用命令。

locust -f performance_tests.py

这个命令在我们的机器上运行一个Locust服务器进行性能测试。我们可以通过浏览器访问它，默认端口是8989。在这里，我们需要向Locust提供有关测试的信息。在屏幕上，我们定义我们想要创建多少个用户（即进程）。这些进程将向API发送请求。

此外，我们还需要定义这些进程的创建速度。最后，我们需要定义API的地址。注意，这里你不需要定义端点，只需要定义地址。下面是我的填写方法。

所以，Locust将创建100个进程，速度为每秒1个进程。这些进程中的每一个都将执行PerformanceTest类中定义的任务。一旦点击 "启动 蜂群"按钮。一旦运行，我们有几个屏幕，我们可以观察进度。第一个屏幕是统计数据。这个屏幕给我们提供了测试的总体情况。在这里，我们可以看到我们的API在测试过程中的响应速度，每秒的请求数（RPS）以及我们有多少次失败。

同样的事情，但完整的历史可以在图表标签中看到。这是一个更有趣的标签，你可以看到请求数是如何随着进程数（用户数）增长的，以及这些数值是如何影响API的速度的。

除此之外，还有 "失败 "和 "异常 "标签，你可以看到为什么你的请求会失败。在这个特定的案例中，我们没有任何错误，但这里有一个我最近测试的API的例子，在统计中出现了一些错误，我不得不在Failures标签中检查原因是什么。

你也可以在下载标签中以csv或html格式下载Locust报告。

总结

在这篇文章中，我们涵盖了从训练模型到在API中部署模型以及用Locust测试该API性能的完整过程。

谢谢您的阅读!