vLLM 插件系统

vLLM 插件系统#

社区经常要求扩展 vLLM 的自定义功能。为了方便这一点，vLLM 包含一个插件系统，允许用户添加自定义功能而无需修改 vLLM 代码库。本文档解释了插件在 vLLM 中的工作原理以及如何为 vLLM 创建插件。

插件在 vLLM 中的工作原理#

插件是用户注册的、由 vLLM 执行的代码。考虑到 vLLM 的架构（参见架构概述），可能会涉及多个进程，尤其是在使用具有各种并行技术的分布式推理时。为了成功启用插件，vLLM 创建的每个进程都需要加载插件。这通过 load_general_plugins 函数在 vllm.plugins 模块中完成。对于 vLLM 创建的每个进程，此函数都会在进程开始任何工作之前被调用。

vLLM 如何发现插件#

vLLM 的插件系统使用标准的 Python entry_points 机制。此机制允许开发者在其 Python 包中注册函数，以供其他包使用。插件示例

# inside `setup.py` file
from setuptools import setup

setup(name='vllm_add_dummy_model',
      version='0.1',
      packages=['vllm_add_dummy_model'],
      entry_points={
          'vllm.general_plugins':
          ["register_dummy_model = vllm_add_dummy_model:register"]
      })

# inside `vllm_add_dummy_model.py` file
def register():
    from vllm import ModelRegistry

    if "MyLlava" not in ModelRegistry.get_supported_archs():
        ModelRegistry.register_model("MyLlava",
                                        "vllm_add_dummy_model.my_llava:MyLlava")

有关向您的包添加入口点的更多信息，请查看官方文档。

每个插件包含三个部分

  1. 插件组：入口点组的名称。vLLM 使用入口点组 vllm.general_plugins 注册通用插件。这是 setup.py 文件中 entry_points 的键。对于 vLLM 的通用插件，始终使用 vllm.general_plugins。

  2. 插件名称：插件的名称。这是 entry_points 字典中的字典值。在上面的示例中，插件名称是 register_dummy_model。插件可以通过其名称使用 VLLM_PLUGINS 环境变量进行过滤。要仅加载特定插件，请将 VLLM_PLUGINS 设置为插件名称。

  3. 插件值：要在插件系统中注册的函数的完全限定名称。在上面的示例中，插件值是 vllm_add_dummy_model:register，它指的是 vllm_add_dummy_model 模块中名为 register 的函数。

支持的插件类型#

  • 通用插件（组名称为 vllm.general_plugins）：这些插件的主要用例是将自定义的、树外的模型注册到 vLLM 中。这通过调用 ModelRegistry.register_model 在插件函数内部注册模型来完成。

  • 平台插件（组名称为 vllm.platform_plugins）：这些插件的主要用例是将自定义的、树外的平台注册到 vLLM 中。当当前环境中不支持该平台时，插件函数应返回 None，当支持该平台时，则返回平台类的完全限定名称。

编写插件的指南#

  • 可重入性：入口点中指定的函数应该是可重入的，这意味着它可以被多次调用而不会引起问题。这是必要的，因为该函数在某些进程中可能会被多次调用。

兼容性保证#

vLLM 保证了已记录插件的接口，例如 ModelRegistry.register_model，将始终可用于插件注册模型。但是，插件开发者有责任确保其插件与他们所针对的 vLLM 版本兼容。例如，"vllm_add_dummy_model.my_llava:MyLlava" 应该与插件目标 vLLM 版本兼容。模型的接口可能会在 vLLM 开发过程中发生变化。