MLflow Gateway 插件 Provider 实战通过 Python Entry Point 扩展自定义模型端点【免费下载链接】mlflowThe open source AI engineering platform for agents, LLMs, and ML models. MLflow enables teams of all sizes to debug, evaluate, monitor, and optimize production-quality AI applications while controlling costs and managing access to models and data.项目地址: https://gitcode.com/GitHub_Trending/ml/mlflow导读MLflow GatewayAI 网关内置了 OpenAI、Anthropic、Bedrock、Gemini 等二十余种大模型 Provider但当需要接入内部自研模型或尚未内置的第三方服务时就需要用到本文介绍的**插件 Providerplugin provider**机制通过 Python 包的 Entry Point 声明将自己实现的 Provider 类注册进网关从而像内置 Provider 一样在 YAML 配置中直接引用。本文以仓库中的examples/gateway/plugin完整示例为主线讲解如何编写插件包、如何用mlflow gateway start加载并启动、以及如何通过客户端调用自定义端点并深入源码揭示插件注册的底层原理。读完本文你将能够独立为 MLflow Gateway 编写、安装和部署自己的 Provider 插件。一、示例总览一个自定义 LLM Provider 的最小闭环examples/gateway/plugin/目录构成了一个完整的插件 Provider 示例其结构如下examples/gateway/plugin/ ├── README.md # 示例说明本文讲解的主体 ├── config.yaml # 网关端点配置文件声明名为 chat 的插件端点 ├── example.py # 客户端脚本通过 deployments client 调用插件端点 └── my-llm/ # 自定义 Provider 的 Python 包 ├── pyproject.toml # 包声明文件核心是 mlflow.gateway.providers entry point └── my_llm/ ├── __init__.py ├── config.py # 端点配置模型含 API Key 的环境变量解析 └── providers.py # Provider 实现类实现 chat 方法示例的运行闭环非常清晰实现一个名为my_llm的 Provider 包 → 通过 Entry Point 注册 → 在config.yaml中声明一个chat端点并指向该 Provider → 启动网关 → 用客户端调用chat端点拿到模型响应。整个链路用最少的代码量演示了插件机制的全部关键环节是理解 MLflow Gateway 扩展性的最佳入门样例。二、端点配置在 config.yaml 中声明插件 Provider示例的配置文件 examples/gateway/plugin/config.yaml 内容如下endpoints: - name: chat endpoint_type: llm/v1/chat model: provider: my_llm name: my-model-0.1.2 config: my_llm_api_key: $MY_LLM_API_KEY该配置只声明了一个端点chat各字段含义如下字段取值说明namechat端点的唯一名称也是后续客户端调用时使用的端点标识endpoint_typellm/v1/chat端点类型表明这是一个 OpenAI 兼容的 Chat Completion 风格端点model.providermy_llmProvider 名称。这里my_llm不是内置 Provider而是插件包通过 Entry Point 注册的名字网关会据此在注册表中查找对应的 Provider 类model.namemy-model-0.1.2模型名称会被透传给 Provider 实现如响应中的model字段model.config.my_llm_api_key$MY_LLM_API_KEYProvider 自定义的配置项$前缀表示从同名环境变量解析真实值注意model.config下的键与插件包中定义的配置模型字段一一对应——这里的my_llm_api_key正是 my_llm/config.py 中MyLLMConfig模型的字段。这种配置模型驱动的设计让每个 Provider 都能声明自己专属的配置项而网关负责统一解析与校验。三、编写插件 Provider 包3.1 包声明与 Entry Point 注册插件包my_llm的声明文件 examples/gateway/plugin/my-llm/pyproject.toml 是整个插件机制的枢纽[project] name my_llm version 1.0 [project.entry-points.mlflow.gateway.providers] my_llm my_llm.providers:MyLLMProvider [tool.setuptools.packages.find] include [my_llm*] namespaces false关键点在于[project.entry-points.mlflow.gateway.providers]声明了一个命名空间为mlflow.gateway.providers的 Entry Point键名my_llm即该 Provider 在配置文件中使用的名字与config.yaml的provider: my_llm对应值my_llm.providers:MyLLMProvider指向实现 Provider 类的模块路径[tool.setuptools.packages.find]保证包内的my_llm子模块会被正确打包。从源码看网关启动时会通过get_entry_points(mlflow.gateway.providers)收集该命名空间下所有已安装的 Entry Point并逐个加载注册见 mlflow/gateway/provider_registry.py。因此只要插件包被pip install安装进当前 Python 环境网关就能自动发现它无需修改任何 MLflow 源码。3.2 实现 Provider 类Provider 实现位于 examples/gateway/plugin/my-llm/my_llm/providers.py完整代码如下import time from mlflow.gateway.config import EndpointConfig from mlflow.gateway.providers import BaseProvider from mlflow.gateway.schemas import chat from my_llm.config import MyLLMConfig class MyLLMProvider(BaseProvider): NAME MyLLM CONFIG_TYPE MyLLMConfig def __init__(self, config: EndpointConfig) - None: super().__init__(config) if config.model.config is None or not isinstance(config.model.config, MyLLMConfig): raise TypeError(fUnexpected config type {config.model.config}) self.my_llm_config: MyLLMConfig config.model.config async def chat(self, payload: chat.RequestPayload) - chat.ResponsePayload: return chat.ResponsePayload( idid-123, createdint(time.time()), modelself.config.model.name, choices[ chat.Choice( index0, messagechat.ResponseMessage( roleassistant, contentThis is a response from MyLLMProvider ), ) ], usagechat.ChatUsage( prompt_tokens10, completion_tokens18, total_tokens28, ), )对照网关源码中的抽象基类 mlflow/gateway/providers/base.py标注为developer_stable即面向开发者稳定开放可以提炼出实现一个插件 Provider 需要遵守的契约继承BaseProvider它是所有网关 Provider 的抽象基类定义了统一的初始化流程与类型标注。声明CONFIG_TYPE指向继承自ConfigModel的配置模型。网关在构造 Provider 前会把model.config解析为该类型实例因此__init__中可用isinstance校验避免拿到意外的配置类型。实现与端点类型匹配的异步方法对于llm/v1/chat端点需要实现async def chat(self, payload: chat.RequestPayload) - chat.ResponsePayload。方法签名完全基于网关的请求/响应 schemamlflow.gateway.schemas.chat返回结构包含id、created、model、choices、usage等标准字段——这也是为什么插件端点可以无缝兼容 OpenAI 风格的客户端。端点类型与方法名的对应关系从配置结构可以推断不同endpoint_type如llm/v1/embeddings、llm/v1/completions会对应 Provider 上的不同方法embeddings、completions等。示例仅实现了chat方法与配置中声明的llm/v1/chat端点类型严格匹配。示例 Provider 直接构造了一个固定的响应对象idid-123、固定文案与 token 用量用于演示响应协议的正确组装方式。真实场景中开发者应当在该方法内调用自有模型服务的 API并将返回结果映射为chat.ResponsePayload。3.3 配置模型与环境变量解析配置模型 examples/gateway/plugin/my-llm/my_llm/config.py 展示了网关对$前缀环境变量的处理约定import os from pydantic import field_validator from mlflow.gateway.base_models import ConfigModel class MyLLMConfig(ConfigModel): my_llm_api_key: str field_validator(my_llm_api_key, modebefore) def validate_my_llm_api_key(cls, value): if value.startswith($): # This resolves the API key from an environment variable env_var_name value[1:] if env_var : os.environ.get(env_var_name): return env_var else: raise ValueError(fEnvironment variable {env_var_name!r} is not set) return value要点如下配置模型必须继承网关提供的ConfigModel位于mlflow.gateway.base_models这样网关才能将其接入统一的配置解析链路。字段my_llm_api_key: str与config.yaml中的config.my_llm_api_key一一对应。借助 pydantic 的field_validator(..., modebefore)在字段校验前处理$前缀若值以$开头则把$后的部分当作环境变量名去os.environ中取值环境变量未设置时抛出ValueError。因此启动网关前必须保证MY_LLM_API_KEY已注入或者直接把明文密钥写在配置文件中不推荐。这是插件示例自带的解析逻辑意味着环境变量解析是每个 Provider 自行约定的行为插件作者可以自由设计配置项的解析方式例如支持$VAR前缀、默认值兜底或加密解密等。四、安装插件并启动网关4.1 安装插件包进入示例目录后用可编辑模式安装插件包使 Entry Point 对当前 Python 环境可见pip install -e ./my-llm-eeditable模式会在开发期间同步源码改动方便迭代。安装完成后可验证 Entry Point 是否生效python -c from importlib.metadata import entry_points; print([ep for ep in entry_points(groupmlflow.gateway.providers)])应能看到名为my_llm、指向my_llm.providers:MyLLMProvider的 Entry Point。这正是网关源码_register_plugin_providers在启动时扫描的对象。4.2 启动网关服务在示例目录即config.yaml所在目录执行MY_LLM_API_KEYsome-api-key mlflow gateway start --config-path config.yaml --port 7000环境变量MY_LLM_API_KEY与配置中的$MY_LLM_API_KEY对应由MyLLMConfig的校验器解析为真实的密钥值示例 Provider 实际上不使用该值但配置解析链路会校验它必须存在--config-path config.yaml指定端点配置文件--port 7000指定网关监听端口客户端脚本example.py中的http://127.0.0.1:7000与此保持一致。启动过程中网关会初始化全局的provider_registry先注册全部内置 Provider再通过get_entry_points(mlflow.gateway.providers)扫描插件并注册见 mlflow/gateway/provider_registry.py。注册完成后config.yaml中的provider: my_llm才能在注册表中命中。4.3 清理环境示例结束后卸载插件包以恢复环境pip uninstall my_llm卸载后该 Entry Point 消失下次启动网关时my_llm将不再可被引用。五、客户端调用插件端点仓库提供了现成的调用脚本 examples/gateway/plugin/example.py演示了通过统一客户端访问插件端点的三种操作from mlflow.deployments import get_deploy_client def main(): client get_deploy_client(http://127.0.0.1:7000) print(fPlugin endpoints: {client.list_endpoints()}\n) print(fPlugin chat endpoint info: {client.get_endpoint(endpointchat)}\n) # Chat request response_chat client.predict( endpointchat, inputs{ messages: [ { role: user, content: Tell me a joke, } ] }, ) print(fPlugin response for chat: {response_chat}) if __name__ __main__: main()脚本执行后依次完成client.list_endpoints()列出网关已注册的全部端点应包含chatclient.get_endpoint(endpointchat)查询端点详情模型名称、Provider 等元数据client.predict(endpointchat, inputs{messages: [...]})发送一个 OpenAI 风格的对话请求。inputs中的messages结构rolecontent与llm/v1/chat端点类型对齐插件 Provider 的chat方法会收到该载荷并返回标准chat.ResponsePayload最终打印固定响应This is a response from MyLLMProvider及 token 用量。可见插件端点在客户端侧与内置 Provider 完全无差别get_deploy_client、list_endpoints、get_endpoint、predict是统一接口插件只是底层实现的不同。这一特性让上层应用可以透明地在内置 Provider 与自定义插件之间切换而无需改动调用代码。六、插件注册的源码级原理examples/gateway/plugin示例背后是网关的插件发现与注册机制实现集中在 mlflow/gateway/provider_registry.pydef _register_plugin_providers(registry: ProviderRegistry): providers get_entry_points(mlflow.gateway.providers) for p in providers: cls p.load() registry.register(p.name, cls)整个过程可以拆解为四个阶段扫描get_entry_points(mlflow.gateway.providers)遍历当前 Python 环境中所有已安装包收集命名空间为mlflow.gateway.providers的 Entry Point——这正是my_llm包在pyproject.toml中声明的位置加载p.load()动态导入并取得 Entry Point 指向的类对象此处为MyLLMProvider注册registry.register(p.name, cls)以 Entry Point 的键名my_llm为 key 存入ProviderRegistry的_providers字典。若键名与已注册 Provider 冲突会抛出MlflowException提示重复注册查询与准入配置解析时通过registry.get(name)查找 Provider 类并调用is_provider_allowed(name)依据环境变量MLFLOW_GATEWAY_ALLOWED_PROVIDERS做白名单校验见 mlflow/gateway/provider_registry.py——这意味着即使是插件 Provider也可以纳入统一的网关 Provider 访问策略管理。值得注意的两点插件注册与内置 Provider 注册_register_default_providers共用同一个provider_registry实例见 mlflow/gateway/provider_registry.py插件 Provider 与内置 Provider 在地位上完全对等由于注册发生在模块导入级模块顶层语句插件必须在网关进程启动前就已安装到运行环境中——这也是示例要求先pip install -e ./my-llm再启动服务的原因。七、常见问题与排查要点结合示例与源码整理实践中容易踩坑的几点问题现象可能原因排查建议启动时提示Provider my_llm not found插件包未安装或 Entry Point 命名空间写错确认执行过pip install -e ./my-llm并用importlib.metadata.entry_points验证 Entry Point 存在且 group 为mlflow.gateway.providers提示Provider my_llm is not allowedMLFLOW_GATEWAY_ALLOWED_PROVIDERS白名单未包含插件名在启动网关时将该环境变量配置为允许的值或取消该策略限制启动报配置解析错误MY_LLM_API_KEY环境变量未设置$前缀解析失败在启动命令前显式注入环境变量如MY_LLM_API_KEYxxx mlflow gateway start ...调用时报不支持的端点类型/方法配置的endpoint_type与 Provider 实现的方法不匹配确认 Provider 实现了与endpoint_type对应的方法如chat并保持CONFIG_TYPE与配置字段一致修改插件代码后行为未更新未以可编辑模式安装使用pip install -e ./my-llm安装修改后重启网关进程八、总结examples/gateway/plugin示例完整演示了 MLflow Gateway 插件 Provider 的开发闭环在pyproject.toml中声明mlflow.gateway.providersEntry Point → 实现继承自BaseProvider的 Provider 类含CONFIG_TYPE配置模型与对应端点的异步方法→ 安装插件包 → 在config.yaml中以provider: entry-point-name声明端点 → 启动网关并通过get_deploy_client调用。底层机制上网关启动时会自动扫描并注册所有已安装的插件 Entry Point与内置 Provider 共用注册表、同等接受MLFLOW_GATEWAY_ALLOWED_PROVIDERS策略管控。对于希望将内部模型服务、自研推理框架或未内置的第三方 API 接入 MLflow Gateway 的团队这一插件机制提供了低侵入、可维护的扩展路径——只写一个符合契约的 Python 包无需改动网关本体任何代码即可获得与官方 Provider 一致的端点管理与统一客户端调用体验。相关参考文件示例说明examples/gateway/plugin/README.md端点配置examples/gateway/plugin/config.yaml插件包声明examples/gateway/plugin/my-llm/pyproject.tomlProvider 实现examples/gateway/plugin/my-llm/my_llm/providers.py配置模型与密钥解析examples/gateway/plugin/my-llm/my_llm/config.py客户端示例examples/gateway/plugin/example.py插件注册源码mlflow/gateway/provider_registry.pyProvider 基类契约mlflow/gateway/providers/base.py【免费下载链接】mlflowThe open source AI engineering platform for agents, LLMs, and ML models. MLflow enables teams of all sizes to debug, evaluate, monitor, and optimize production-quality AI applications while controlling costs and managing access to models and data.项目地址: https://gitcode.com/GitHub_Trending/ml/mlflow创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
