将Flask移植到FastAPI用于ML模型服务的教程

尼耳多
256 阅读7分钟

我看到新的数据科学家进入工业界的共同困难之一是让他们的机器学习模型与生产系统互动，也就是说，从他们的研究环境（通常在Jupyter笔记本或类似的环境中开发）跳到更广泛的代码生态系统中。把模型放在REST API后面，可以确保我们可以在飞行中向消费者提供预测，而不受计算资源、规模、甚至为生产生态系统的其他部分选择的语言的影响。

Flask是一个非常流行的框架，用于在Python中构建REST API，近一半的Python编码人员报告在2019年使用它。特别是，Flask对于服务于ML模型非常有用，与其他框架的 "包含电池 "的一体化功能相比，Flask的简单性和灵活性更值得期待，因为它更适合于一般的网络开发。然而，自从Flask问世以来，Python的性能、类型注解行为和异步编程接口有了一些发展，使我们能够建立一个更快、更健壮、更灵敏的API--在这里，我们将学习如何迁移到较新的FastAPI框架以利用这些进展。

这篇文章将假设我们对Flask和HTTP请求有一定的了解，尽管我们会先在Flask中完成应用的构建。这篇文章的完整docker化代码可以在这里找到。为了本文的目的，我们将涵盖。

  • 设置FastAPI应用，并使用uvicorn ，运行异步服务器。gunicorn
  • 路线规范和端点
  • 数据验证
  • 异步端点

我们不会涉及的东西。

  • 单元测试你的API - 但你可以使用starletteTestClient 来设计测试，其行为与Flask的非常相似。
  • 部署，因为这是非常具体的组织--但该应用很容易被Docker化，并通过gunicorn ，在生产就绪的服务器上运行，我们已经在完整的代码示例中包含了一个配置示例。

一个简单的Flask API

首先，我们需要一个模型。上面链接的演示代码提供了一个脚本，在scikit-learn （使用rec.sport.hockey 、sci.space 和talk.politics.misc 新闻组）中提供的 "20个新闻组 "数据集的一个子集上训练一个简单的Naive Bayes分类器，但请随意在此基础上构建和pickle 你自己的文本分类器--任何具有合理API的ML包在这个框架中都可以正常工作。

让我们从一个简单的、单文件的Flask应用开始，为这个模型服务。首先是标题内容。

import os
import pickle

from flask import Flask, jsonify, request
import numpy as np


app = Flask(__name__)

with open(os.getenv("MODEL_PATH"), "rb") as rf:
    clf = pickle.load(rf)


@app.route("/healthcheck", methods=["GET"])
def healthcheck():
    msg = (
        "this sentence is already halfway over, "
        "and still hasn't said anything at all"
    )
    return jsonify({"message": msg})

我们在其中设置了Flaskapp 对象，并在应用启动时将模型加载到内存中--实际的预测方法将引用这个模型对象作为一个闭包，避免每次模型调用时不必要的文件I/O。(我们将暂时搁置寻找比pickle 更好的模型加载范式，因为这将取决于所涉及的部署方案和模型的具体情况。)

包括一个 "健康检查 "端点，返回一个简单的 "我很好 "的状态，而不需要计算GET ，这对于部署监控和确保你自己的开发环境被正确设置都很有用。上述内容应该在你的环境中产生一个可以工作的（尽管有点傻）API--通过在你的shell中运行来测试它

$ export FLASK_APP=<path to your app file>.py
$ flask run

(或使用示例代码中的docker化调用）来测试。

预测端点

接下来，让我们添加一些实际的预测功能。

@app.route("/predict", methods=["POST"])
def predict():
    samples = request.get_json()["samples"]
    data = np.array([sample["text"] for sample in samples])

    probas = clf.predict_proba(data)
    predictions = probas.argmax(axis=1)

    return jsonify(
        {
            "predictions": (
                np.tile(clf.classes_, (len(predictions), 1))[
                    np.arange(len(predictions)), predictions
                ].tolist()
            ),
            "probabilities": probas[np.arange(len(predictions)), predictions].tolist()
        }
    )

这是我们用于服务预测的典型设计模式--客户POST 输入数据到端点，并以固定的格式接收预测结果，模型本身通过闭包调用。这就要求有一个形式的请求体

{
  "samples": [
    {
      "text": "this is a text about hockey"
    }
  ]
}

因此，我们将传递给它一个samples ，每个数组包含我们期望的文本输入数据的键值集（对于具有多个不同特征的模型，将每个特征作为一个键值对传递）。这确保我们以我们想要的形式明确地访问数据，与pandas 和其他表格数据工具很好地配合，并反映了我们可能期望从生产数据源中获得的单个数据点的形状。我们的predict() 方法将其解压为一个数组，通过scikit-learn 分类器管道对象，并返回预测的标签（映射到类名）和其相应的概率。请注意，我们可以简单地调用分类器的predict 方法，而不是将predict_proba 的概率输出映射到类标签上，但是一些简单的numpy 代数通常会比第二次运行推理快得多。

@app.route("/predict/<label>", methods=["POST"])
def predict_label(label):
    samples = request.get_json()["samples"]
    data = np.array([sample["text"] for sample in samples])

    probas = clf.predict_proba(data)
    target_idx = clf.classes_.tolist().index(label)

    return jsonify({"label": label, "probabilities": probas[:, target_idx].tolist()})

同样，我们可以要求每个样本的特定标签的概率--这种类型的参数化很容易通过路径本身完成。就这样，我们拥有了它!

这个API可以用来提供预测服务（自己试试！）......但如果你试图把它投入生产，它完全可能会倒下。值得注意的是，它缺乏任何数据验证--传入请求中的任何缺失、错误命名或错误输入的信息都会导致未处理的异常，从Flask ，返回一个500错误和一个毫无帮助的响应（以及一些在调试模式下相当吓人的HTML），同时可能向你的监控系统发出警报并唤醒你的开发团队。

通常，Flask 中的错误处理最终都是以脆弱的、纠结的 try-catches 和受保护的dict 访问结束。更好的方法是使用一个包，如 pydantic或 marshmallow来实现更多的程序化数据验证。幸运的是，FastAPI ，包括开箱即用的pydantic 验证。

进入FastAPI

FastAPI是一个现代的Python网络框架，旨在。

  • 提供一个轻量级的微框架，有一个直观的、类似Flask的路由系统
  • 利用Python 3.6以上版本的类型注解支持进行数据验证和编辑支持
  • 利用对Python异步支持的改进和ASGI规范的发展，使异步API更加容易。
  • 使用OpenAPI和JSON Schema自动生成有用的API文档

在引擎盖下，FastAPI使用pydantic进行数据验证，并在其网络工具中使用starlette，使其与Flask等框架相比快得令人发指，并使其性能与Node或Go中的高速网络API相当。

幸运的是，由于FastAPI明确借鉴了Flask的路由规范，所以过渡到使用它是相当快的--让我们开始把我们的服务应用程序的功能移植过来。对于相应的头文件内容。

import os
import pickle

from fastapi import FastAPI
import numpy as np

app = FastAPI()

with open(os.getenv("MODEL_PATH"), "rb") as rf:
    clf = pickle.load(rf)


@app.get("/healthcheck")
def healthcheck():
    msg = (
        "this sentence is already halfway over, "
        "and still hasn't said anything at all"
    )
    return {"message": msg}

到目前为止很容易--我们只需要做一些微小的语义变化。

  • 实例化Flask(__name__) 顶层对象变成实例化FastAPI() 对象
  • 路由仍然用装饰器来指定，但HTTP方法在装饰器本身而不是一个参数中--@app.route(..., methods=["GET"]) 变成了@app.get(...)

我们可以像flask的应用启动一样调用它，尽管它不包括一个内置的web服务器--相反，我们将直接启动uvicorn服务器。

$ uvicorn --reload <path to app file>:app

这给了我们一个简单的单工设置（或者，同样，只需使用示例代码中的dockerized调用）。

数据验证

对于指定预测端点，我们确实要对我们的思维引入一个重大改变。在上面提到的简单验证中，像try-catches和保护dicts这样的东西，本质上是试图用我们的数据拍下我们不想要的东西。如果我们只是简单地指定我们想要的东西，而让应用程序来处理其他的事情，会怎么样呢？这正是像marshmallow 或pydantic 这样的验证工具所做的--对于FastAPI中的pydantic ，我们只需使用Python的新（3.6+）类型注释指定模式，并将其作为参数传递给路由函数。由于Python的类型注解，FastAPI知道如何进行验证，这意味着我们只需要非常自然地指定我们对输入的期望，其余的就由它来完成。对于我们的predict() 端点，这看起来像

from typing import List

from pydantic import BaseModel


class TextSample(BaseModel):
    text: str


class RequestBody(BaseModel):
    samples: List[Sample]

我们简单地指定预期的数据形状（在这里使用基本的Python类型注释，尽管pydantic 支持扩展的类型检查，如字符串/电子邮件验证或数组尺寸和规范）到pydantic.BaseModel 的一个子类中。虽然pydantic 将其类型检查功能建立在这个类中，我们仍然可以像普通的Python对象一样使用它 - 子类化或添加额外的功能都是预期的。例如，我们可以在这个类中建立将样本的内容解包成一个数组的功能，以便进行推理，例如

class RequestBody(BaseModel):
    samples: List[Sample]

    def to_array(self):
        return [sample.text for sample in self.samples]

取代上面使用的列表理解，并保证数组格式的正确性。在端点中，我们将输入数据作为一个函数参数传递。

@app.post("/predict")
def predict(body: RequestBody):
    data = np.array(body.to_array())

    probas = clf.predict_proba(data)
    predictions = probas.argmax(axis=1)

    return {
        "predictions": (
            np.tile(clf.classes_, (len(predictions), 1))[
                np.arange(len(predictions)), predictions
            ].tolist()
        ),
        "probabilities": probas[np.arange(len(predictions)), predictions].tolist(),
    }

FastAPI将智能地处理这个问题，首先按名称查找路由参数，然后将请求体（用于POST 请求）或查询参数（用于GET ）打包到函数参数中。然后可以通过典型的属性/方法语法来访问产生的数据字段或方法。如果是畸形的输入数据，pydantic 将引发验证错误--FastAPI内部处理这个问题，返回422错误代码，其中包含有关错误的有用信息的JSON体。

枚举值和路径参数

我们也可以在数据验证中使用枚举值--例如，在predict_label() 端点中，我们可以用以下方式处理有效目标名称

from enum import Enum

class ResponseValues(str, Enum):
    hockey = "rec.sport.hockey"
    space = "sci.space"
    politics = "talk.politics.misc"

传递这个来验证路径参数将干净地处理坏目标名称的错误，这将扼杀试图在clf.classes_ 中找到相应值的索引。 在端点中，我们然后有

@app.post("/predict/{label}")
def predict_label(label: ResponseValues, body: RequestBody):
    data = np.array(body.to_array())

    probas = clf.predict_proba(data)
    target_idx = clf.classes_.tolist().index(label.value)

    return {"label": label.value, "probabilities": probas[:, target_idx].tolist()}

响应模型和文档

到目前为止，我们只在输入端添加了数据验证，但FastAPI允许我们也为输出端声明一个模式--我们定义了

class ResponseBody(BaseModel):
    predictions: List[str]
    probabilities: List[float]


class LabelResponseBody(BaseModel):
    label: str
    probabilities: List[float]

并将路由装饰器替换为

@app.post("/predict", response_model=ResponseBody)
...
@app.post("/predict/{label}", response_model=LabelResponseBody)
...

为我们的输出添加数据验证似乎很奇怪--毕竟，如果我们返回的东西不符合模式，这表明我们的代码中有更深的问题。然而，我们可以使用这些模式来限制向外部用户返回的确切数据（例如，通过从内部消息中删除敏感字段）。对我们来说更重要的是，这些模式会被纳入FastAPI的自动生成的文档中--当API运行时，点击{api url}/docs 或{api url}/redoc 端点会加载OpenAPI生成的文档，详细说明可用的端点及其输入和输出的预期结构（这些结构来自我们的模式）。我们甚至可以给API和它的端点添加注释。- 单独的路由可以被标记以进行分类（对API版本管理很有用）--路由函数的文档字符串将被拉入API文档以进行端点描述--FastAPI 对象本身可以接受标题、描述和版本等关键字参数，这些参数将被填充到文档中

异步端点

我们仍然没有触及FastAPI最强大的方面之一--它对异步代码的干净处理。坦率地说，这在数据科学家和ML工程师中并不罕见--如此多的模型训练和推理是受处理器限制的，以至于异步代码并没有出现那么多，相比之下（例如）Web开发中的异步代码要普遍得多。对于这个问题的深入研究，请查看Python自己的异步文档，尽管我实际上发现FastAPI自己的解释对于掌握并发和并行代码之间的差异要直观得多。

简而言之，异步执行（或并发，如果你愿意的话）的想法是将你的进程向外部资源发射工作，比如向外部 API 或数据存储的请求。在同步代码中，进程阻塞，直到工作完成 - 对于一个缓慢的请求，这意味着整个进程闲置，直到它收到一个返回值。异步执行允许进程切换上下文，并在不相关的事情上工作，直到它被告知所请求的工作已经完成，这时它又恢复了。(相比之下，并行代码的执行会有多条工作线都在独立地执行可能阻塞的代码）。

在机器学习中，执行通常是受处理器约束的--也就是说，工作始终是在一个或多个CPU核心上完全订阅处理能力。在这种情况下，异步执行并不是特别有帮助，因为实际上并不存在工作可以被处理器放下以等待结果，同时执行其他东西的情况（相反，有大量的场景可以并行计算ML工作）。然而，我们仍然可以设想这样的情况：我们的模型服务API可能在等待外部资源，而不是自己做计算的重任。例如，假设我们的API必须根据其计算结果从数据库或内存缓存中请求信息，或者我们的API是一个轻量级的中间人，在将工作移交给运行在GPU实例上的单独的tensorflow-serving API之前执行验证或预处理--正确处理异步处理可以使我们的工作以很小的代价获得显著的性能提升。

从历史上看，Python中的异步工作是不容易的（尽管自Python 3.4以来，其API已经迅速改善），特别是Flask。基本上，Flask（在大多数WSGI服务器上）默认是阻塞的--由对特定端点的请求触发的工作将完全保留服务器，直到该请求完成。相反，Flask（或者说，运行它的WSGI服务器，如gunicorn 或uWSGI ）通过并行运行应用程序的多个工作者实例来实现扩展，这样，当一个工作者忙碌时，请求可以被农场化。在一个工作者中，异步工作可以被包裹在一个阻塞调用中（路由函数本身仍然是阻塞的）、线程化（在较新的Flask版本中），或者被养殖到一个队列管理器中，如Celery - 但没有一个一致的故事，路由可以干净地处理异步请求，无需额外的工具。

FastAPI的异步性

相比之下，FastAPI从一开始就被设计为异步运行--由于其底层的starlette ASGI框架，路由函数默认为在异步事件循环中运行。有了一个好的ASGI服务器（FastAPI被设计成与uvicorn 耦合，在uvloop 上运行），这可以让我们的性能与Go或Node中的快速异步网络服务器相当，而不会失去Python更广泛的机器学习生态系统的好处。

与使用线程或Celery队列来实现Flask中的异步执行相比，在FastAPI中异步运行一个端点是非常简单的--我们只需将路由函数声明为异步的（用async def ），就可以开始了！如果路由函数不是传统意义上的异步，我们甚至可以这样做--也就是说，我们没有任何可等待的调用（比如端点正在针对ML模型运行推理）。事实上，除非端点是专门执行一个阻塞的IO操作（例如到数据库），否则最好用async def （因为阻塞函数实际上被打到外部线程池，然后无论如何都要等待）来声明该函数。

对于我们上面的ML预测函数，我们可以用async def 来声明端点，尽管这并没有对我们的代码做出任何有趣的改变。但是，如果我们需要做一些真正的异步的事情，比如从外部API请求（和等待）资源呢？不幸的是，我们在Python中的传统requests 包是阻塞的，所以我们不能用它来做HTTP的异步请求--相反，我们将使用优秀的 aiohttp包中的请求功能。

异步请求

首先，我们需要设置一个客户端会话--这将保持一个持久的池子运行，以等待请求，而不是为每个请求创建一个新的会话（这实际上是requests ，如果像典型的requests.get ，requests.post ，等等那样调用）。我们将把它放在应用程序的顶层，这样任何路由函数都可以通过闭包调用它。

app = FastAPI()
...
client_session = aiohttp.ClientSession()

我们还需要确保这个会话正确关闭--幸运的是，FastAPI给了我们一个简单的装饰器来声明这些操作。

@app.on_event("shutdown")
async def cleanup():
    await client_session.close()

当FastAPI应用程序关闭时，这将执行在函数中调用的任何东西（这里等待aiohttp 客户端会话的干净关闭）。对于外部请求，我们在路由函数中封装了一个可等待的调用。

@app.get("/cat-facts", response_model=TextSample)
async def cat_facts():
    url = "https://cat-fact.herokuapp.com/facts/random"
    async with client_session.get(url) as resp:
        response = await resp.json()

    return response

将请求放在一个异步上下文块中，然后等待一个可解析的响应。在这里，我们利用我们的响应模型来限制返回值--"猫的事实 "API的响应会返回很多关于事实的额外元数据，但我们只想返回事实文本。与其在返回前摆弄响应，我们可以简单地重用我们现有的TextSample 模式，将其打包到响应中，并相信pydantic 来处理过滤问题，因此我们的响应看起来像

{
    "text": "In an average year, cat owners in the United States spend over $2 billion on cat food."
}

就这样了!我们可以将这种结构用于我们可能需要的对外部资源的任何异步调用，比如从数据存储中检索数据或向运行在GPU资源上的tensorflow-serving ，发射推理作业。

总结

在这篇文章中，我们已经走过了一个常见的、简单的布局，将你的机器学习模型立在REST API后面，以一种可消费的方式实现飞行预测，应该与各种代码环境干净地对接。虽然Flask是这些任务的一个极其常见的框架，但我们可以通过迁移到较新的FastAPI框架来利用Python的类型检查和异步支持的改进--幸运的是，从Flask移植到FastAPI是直接的

avatar
文章
阅读
粉丝