MLflow AI Gateway 集成 Hugging Face Text Generation Inference(TGI)实战指南
2026/9/13 0:09:23 网站建设 项目流程

MLflow AI Gateway 集成 Hugging Face Text Generation Inference(TGI)实战指南

【免费下载链接】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 中通过 AI Gateway(Deployments 服务)接入自建的 Hugging Face Text Generation Inference(TGI)服务器,将本地部署的开源 LLM 以统一llm/v1/completions接口暴露给上层应用。你将学会从 Docker 部署 TGI、编写 gateway 配置、启动服务到调用补齐的端到端流程,并理解 gateway 内部对 TGI 参数(temperature、max_tokens、details 等)的转换与约束原理。

TGI 是什么:为开源 LLM 推理而生的工具链

Hugging Face Text Generation Inference(TGI)是一套专门用于高效部署与托管大语言模型(LLM)的完整工具链,对 Llama、Falcon、StarCoder、BLOOM、GPT-Neo 等主流开源模型提供开箱即用的优化支持。TGI 内置的核心优化能力包括:

  • 一键启动器(Simple launcher):通过单条命令即可拉起大多数主流 LLM 的推理服务;
  • 张量并行(Tensor Parallelism):在多张 GPU 上并行切分模型权重,加速推理;
  • Safetensors 权重加载:安全、高效地加载模型权重文件;
  • 优化的 transformers 推理内核:在主流模型架构上使用 Flash Attention 与 Paged Attention 加速注意力计算。

需要特别注意的是,TGI 仅对精选模型列表使用自定义 CUDA 内核做推理优化。若你的模型不在列表中、或属于自建自定义模型,仍可尝试启动服务,但由于未针对 TGI 优化,性能不保证。若希望关闭自定义内核,可在docker run命令末尾追加--disable-custom-kernels参数。

环境准备:硬件要求与 NVIDIA Container Toolkit 安装

注意:本示例在 Linux(Debian 11)+ NVIDIA A100 GPU 环境下测试通过。

将 MLflow AI Gateway 与 TGI 对接的第一步,是先把 Hugging Face 模型部署到 TGI 服务器上。推荐使用官方 Docker 容器ghcr.io/huggingface/text-generation-inference:1.1.1启动 TGI,容器内含运行所需的全部依赖(库、二进制文件与配置)。

启动服务器前,需确认机器硬件满足要求:TGI 优化模型兼容NVIDIA A100、A10G、T4GPU。使用其他 GPU 硬件虽仍能获得性能提升,但 Flash Attention、Paged Attention 等操作不会执行。若你的机器没有 GPU 或 CUDA 支持,可去掉--gpus all参数并加上--disable-custom-kernels;但请注意 CPU 并非 TGI 的目标平台,如此选择会显著影响性能。

安装 NVIDIA Container Toolkit

NVIDIA Container Toolkit 是运行 GPU 加速容器的前置条件。首先添加软件源并刷新索引:

curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | sudo gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg \ && curl -s -L https://nvidia.github.io/libnvidia-container/stable/deb/nvidia-container-toolkit.list | \ sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' | \ sudo tee /etc/apt/sources.list.d/nvidia-container-toolkit.list \ && \ sudo apt-get update

然后安装 toolkit:

sudo apt-get install -y nvidia-container-toolkit

启动 TGI 服务器并验证

安装完成后,执行以下 Docker 命令在本地8000端口启动 TGI 服务器,并加载tiiuae/falcon-7b-instruct模型:

model=tiiuae/falcon-7b-instruct volume=$PWD/data # share a volume with the Docker container to avoid downloading weights every run docker run --gpus all --shm-size 1g -p 8000:80 -v $volume:/data ghcr.io/huggingface/text-generation-inference:1.1.1 --model-id $model

命令要点说明:

  • --gpus all:将全部 GPU 暴露给容器(无 GPU 环境请去掉);
  • --shm-size 1g:设置共享内存为 1GB,避免推理进程因共享内存不足而崩溃;
  • -p 8000:80:将容器的 80 端口映射到宿主机 8000 端口,即 TGI 对外服务端口;
  • -v $volume:/data:将宿主机$PWD/data目录挂载为容器/data,模型权重只需首次下载,后续运行直接复用,避免每次重复下载;
  • --model-id $model:指定要加载的 Hugging Face 模型 ID。

TGI 启动后,可用如下 Python 脚本验证服务是否正常工作:

import requests headers = { "Content-Type": "application/json", } data = { 'inputs': 'What is Deep Learning?', 'parameters': { 'max_new_tokens': 20, }, } response = requests.post('http://127.0.0.1:8000/generate', headers=headers, json=data) print(response.json()) # {'generated_text': '\nDeep learning is a branch of machine learning that uses artificial neural networks to learn and make decisions.'}

该请求走的是 TGI 的原生/generate接口,返回体中的generated_text即模型生成结果。

编写 gateway 配置:新增 completions 端点

服务器就绪后,编辑 MLflow AI Gateway 的配置文件 examples/gateway/huggingface/config.yaml,将 TGI 服务器注册为新的端点:

endpoints: - name: completions endpoint_type: llm/v1/completions model: provider: "huggingface-text-generation-inference" name: falcon-7b-instruct config: hf_server_url: http://127.0.0.1:8080

字段含义与取值说明:

  • name:端点名称,也是后续client.predict(endpoint=...)调用时的唯一标识;
  • endpoint_type:端点语义类型,此处为llm/v1/completions(文本补全)。从 gateway 源码看,该 provider 仅实现补全路由,chat 与 embeddings 路由会返回 501 错误(详见下文"源码级原理");
  • model.provider:固定为字符串"huggingface-text-generation-inference",gateway 通过该值在 provider_registry.py 中注册对应的HFTextGenerationInferenceServerProvider实现;
  • model.name:模型展示名,会原样写入补全响应的model字段,建议设置为实际加载的模型 ID;
  • model.config.hf_server_url:TGI 服务器地址。注意示例脚本中原始 README 使用http://127.0.0.1:8000/generate这种带路径的形式,而仓库中的实际 config.yaml 使用http://127.0.0.1:8080这种不含路径的形式——两种写法 gateway 均能处理,因为底层会通过append_to_uri_pathgenerate路径拼接到 base URL 之后。

在 mlflow/gateway/config.py 中,HuggingFaceTextGenerationInferenceConfig的完整定义只有必填字段hf_server_url: str,即 TGI 服务器地址是唯一必填配置项,可见整个接入过程极其轻量。

启动 MLflow AI Gateway

配置文件就绪后,通过 CLI 启动 gateway 服务(端口 7000):

mlflow gateway start --config-path examples/gateway/huggingface/config.yaml --port 7000

该命令属于mlflow gateway命令组(见 mlflow/gateway/cli.py),--config-path指定 YAML 配置文件路径,--port指定监听端口。启动后,gateway 会读取配置、实例化各 provider 并对外提供 OpenAI 兼容的补全接口。

查询端点:用 Deployment Client 发起补全请求

仓库提供了完整的调用示例脚本 examples/gateway/huggingface/example.py,演示如何查询已部署的falcon-7b-instruct模型:

from mlflow.deployments import get_deploy_client def main(): client = get_deploy_client("http://localhost:7000") print(f"Hugging Face TGI endpoints: {client.list_endpoints()}\n") print( f"Hugging Face completions endpoint info: {client.get_endpoint(endpoint='completions')}\n" ) # Completions request response_completions = client.predict( endpoint="completions", inputs={ "prompt": ("What is Deep Learning?"), "temperature": 0.1, }, ) print(f"Hugging Face TGI response for completions: {response_completions}") if __name__ == "__main__": main()

执行python examples/gateway/huggingface/example.py后,脚本依次完成三件事:

  1. 列出端点list_endpoints()返回 gateway 当前注册的全部端点;
  2. 查看端点详情get_endpoint(endpoint='completions')返回指定端点的元信息;
  3. 发起补全predict()prompt+ 采样参数(如temperature)调用completions端点,并打印 TGI 的生成结果。

get_deploy_client正是文档中mlflow gateway start --config-path ... --port 7000这条部署方式对应的官方客户端入口(见 mlflow/deployments/mlflow/init.py 中对该命令的说明)。

透传与参数转换:gateway 如何与 TGI 交互

当你向 MLflow Deployments 服务器发起请求时,请求体中的信息会被透传给 TGI,从而让你对 TGI 的生成输出拥有更多控制权。但需要注意:detailsdecoder_input_details这两个参数无法关闭,它们是 TGI 端点正常工作的必需项。

以 provider 核心实现 mlflow/gateway/providers/huggingface.py 为据,gateway 在转发前会做如下关键处理:

1. 参数名映射(max_tokensmax_new_tokens

TGI 使用max_new_tokens,而 OpenAI 兼容接口使用max_tokens。provider 通过rename_payload_keys完成映射(源码 huggingface.py);若用户直接传了max_new_tokens,gateway 会返回 422 错误,提示"请改用max_tokens",避免双重指定造成歧义。

2. temperature 缩放(0–2 映射到 0–100)

TGI 的 temperature 取值范围是 0–100,而 gateway 的补全接口范围为 0–2,因此 provider 将用户传入值乘以 50(源码 huggingface.py)。同时 TGI 不支持 0 温度,provider 会用max(scaled_temp, 1e-3)兜底,保证即使传 0 也能得到合法的极小值。

3.n参数约束:只能生成单候选

TGI 不支持一次生成多个候选序列,provider 会弹出n参数;若n != 1,返回 422 错误:"'n' must be '1' for the Text Generation Inference provider."(源码 huggingface.py)。

4. 强制注入detailsdecoder_input_details

provider 会无条件写入parameters["details"] = Trueparameters["decoder_input_details"] = True(源码 huggingface.py),因为响应解析依赖details中的generated_tokensfinish_reasonprefill等字段来计算 token 用量与结束原因。

5. 响应转换为 OpenAI 补全格式

最终请求体为{"inputs": prompt, "parameters": parameters},发送到hf_server_url/generate(路径通过 utils.py 的append_to_uri_path拼接)。返回后,provider 从resp["details"]["generated_tokens"]len(resp["details"]["prefill"])分别得到输出与输入 token 数,组装成text_completion结构的completions.ResponsePayload(源码 huggingface.py),使 TGI 响应对上层调用方完全透明。

6. 仅支持补全路由

provider 只实现completions路由。chatembeddings路由在调用时会分别抛出 501 错误:"The chat route is not implemented for Hugging Face Text Generation Inference models.""The embeddings route is not implemented..."。这意味着本 provider 仅适合纯文本补全场景,对话与向量化请改用其他 provider。

测试佐证:TGI 集成的行为契约

仓库中的单元测试 tests/gateway/providers/test_huggingface.py 将上述行为固化为可验证的契约:

  • test_completions验证了请求被正确转发到https://testserverurl.com/generate,且max_tokens: 1000被转换为max_new_tokens: 1000,同时强制携带details: Truedecoder_input_details: True
  • test_completions_temperature_is_scaled_correctly断言temperature: 0.5最终以0.5 * 50 = 25发送给 TGI,验证了 50 倍缩放逻辑;
  • test_completion_fails_with_multiple_candidates验证n != 1时返回 422;
  • test_chat_is_not_supported_for_tgitest_embeddings_are_not_supported_for_tgi分别验证 chat(501)与 embeddings(501)路由不可用。

这些测试与 mlflow/gateway/providers/huggingface.py 的实现一一对应,读者若需自行扩展 TGI 集成的能力边界(例如新增流式输出),可先阅读这两个文件。

小结

至此,一条完整的"自托管 TGI → MLflow AI Gateway → 应用客户端"链路已经打通:硬件与 NVIDIA Container Toolkit 就绪后,用官方 Docker 容器拉起 TGI 并验证/generate接口;在 config.yaml 中声明huggingface-text-generation-inferenceprovider 端点;用mlflow gateway start启动网关;最后通过get_deploy_client以 OpenAI 兼容语义完成补全调用。理解 provider 在参数映射、temperature 缩放、n约束与details注入方面的内部处理,能帮助你在实际项目中准确控制生成行为,并规避"参数被网关拦截/改写"带来的困惑。

【免费下载链接】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),仅供参考

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询