MLflow Deployments 实战:基于 Databricks Serving 端点管理外部模型与 LLM 推理
2026/9/12 5:46:13 网站建设 项目流程

MLflow Deployments 实战:基于 Databricks Serving 端点管理外部模型与 LLM 推理

【免费下载链接】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

导读

本文围绕仓库 examples/deployments/README.md 所指向的示例展开,讲解如何通过 MLflow Deployments 的get_deploy_client("databricks")客户端,在 Databricks 上创建、更新、查询、推理与删除 Serving 端点,并通过external_model接入 OpenAI 等外部大模型。读完本文,你将掌握端到端的管理流程、端点配置字段的完整语义,以及这些 API 背后在 mlflow/deployments 模块中的真实实现原理。

MLflow Deployments 是什么

MLflow Deployments 是 MLflow 提供的一套统一的模型部署抽象。它通过一个高度一致、可互换的 Python API 与 CLI 接口,让开发者以同样的代码风格管理不同平台的模型部署,而不必关心每个平台自身的 REST 细节。

核心入口定义在 mlflow/deployments/interface.py 的get_deploy_client(target_uri=None)

  • 传入目标 URI(如"databricks")时,它会解析目标名并通过插件注册表返回对应的客户端实例;
  • 不传参数时,会尝试读取MLFLOW_DEPLOYMENTS_TARGET环境变量或get_deployments_target()中设置的目标;
  • 找不到目标时会打印提示日志并返回None

客户端实例是 mlflow/deployments/base.py 中BaseDeploymentClient抽象类的子类,统一暴露了两类 API:

  • Deployment(部署)级 APIcreate_deploymentupdate_deploymentdelete_deploymentlist_deploymentsget_deploymentpredictpredict_streamexplain
  • Endpoint(端点)级 APIcreate_endpointupdate_endpointdelete_endpointlist_endpointsget_endpoint(base.py)。

其中端点级 API 在基类中默认抛出"未实现"异常,实际能力由各目标插件提供。这正是本文示例所展示的 Databricks 客户端的核心工作方式。

示例的整体脉络

仓库 examples/deployments 目录专门用于演示如何开始使用 MLflow Deployments,当前包含 Databricks 子目录。核心可运行脚本是 examples/deployments/databricks/databricks.py,它演示了一条完整的生命周期链路:

  1. 获取 Databricks 部署客户端;
  2. create_endpoint:创建一个代理 OpenAIgpt-4的外部模型端点,并附带 tags 与 rate_limits;
  3. update_endpoint:先更新served_entities,再单独更新rate_limits
  4. list_endpoints/get_endpoint:查询端点列表与单个端点详情;
  5. predict:向端点发送一轮 Chat 对话请求;
  6. finallydelete_endpoint:清理端点,保证示例可重复运行。

脚本头部以 docstring 形式给出了三步用法:

databricks secrets create-scope <scope> databricks secrets put-secret <scope> openai-api-key --string-value $OPENAI_API_KEY python examples/deployments/databricks.py --secret <scope>/openai-api-key

即先用 Databricks CLI 创建 Secret Scope 并存入 OpenAI API Key,再以--secret <scope>/openai-api-key的形式把密钥引用传给脚本,脚本会把该引用包装成{{secrets/scope/key}}形式写入外部模型配置,避免明文密钥出现在代码或日志中。

准备环境:认证与密钥

Databricks 认证

DatabricksDeploymentClient的类文档(mlflow/deployments/databricks/init.py)明确指出,需要先设置认证凭据:

export DATABRICKS_HOST=... export DATABRICKS_TOKEN=...

底层通过get_databricks_host_creds(self.target_uri)获取请求凭据(见 databricks/init.py),因此 Databricks 工作区地址与访问令牌是发起一切 API 调用的前提。除环境变量外,Databricks 官方还支持其他认证方式,可参考其 Dev Tools 文档。

外部模型密钥托管

示例脚本本身不接收任何明文 API Key,而是依赖 Databricks Secrets:

databricks secrets create-scope <scope> databricks secrets put-secret <scope> openai-api-key --string-value $OPENAI_API_KEY

随后脚本通过--secret参数接收形如<scope>/openai-api-key的引用,在构造配置时拼接为"{{" + args.secret + "}}",得到{{<scope>/openai-api-key}}。Databricks 在端点运行时解析该占位符并注入真实密钥。

创建端点:配置字段逐项拆解

databricks.py中的create_endpoint调用是全文信息量最大的一段(databricks.py):

client.create_endpoint( name=name, config={ "served_entities": [ { "name": "test", "external_model": { "name": "gpt-4", "provider": "openai", "task": "llm/v1/chat", "openai_config": { "openai_api_key": "{{" + args.secret + "}}", }, }, } ], "tags": [ {"key": "foo", "value": "bar"} ], "rate_limits": [ {"key": "user", "renewal_period": "minute", "calls": 5} ], }, )

各字段含义如下:

字段作用示例值
served_entities端点承载的模型实体列表,一个端点可托管多个实体单个 OpenAIgpt-4
served_entities[].name实体名称,用于端点内部区分"test"
served_entities[].external_model.name外部模型名称"gpt-4"
served_entities[].external_model.provider模型供应商标识"openai"
served_entities[].external_model.task模型任务类型,决定请求/响应协议"llm/v1/chat"(Chat 补全)
external_model.openai_config供应商专属配置,如密钥引用{"openai_api_key": "{{scope/key}}"}
tags端点的键值标签,用于标记、筛选与运维{"key": "foo", "value": "bar"}
rate_limits速率限制规则:key限定维度(如user/endpoint),renewal_period为重置周期(如minute),calls为周期内允许的调用次数每分钟 5 次

create_endpoint 的两种调用形态与兼容逻辑

从源码看,create_endpoint支持新旧两种调用形态(databricks/init.py):

  • 新形态config里直接包含完整 API 载荷,nameroute_optimized也放在 config 字典中(如{"name": ..., "config": {...}, "route_optimized": True});
  • 旧形态(已弃用)nameconfigroute_optimized作为独立参数传入,形如create_endpoint(name, config, route_optimized=False),示例脚本采用的就是这种形态。

源码会对新旧形态做冲突校验:若name参数与 config 内的name不一致会抛出MlflowException;旧形态下tagsrate_limits会被从 config 中剥离、与route_optimized一起平铺进 payload 顶层。从结构看,Databricks 端点的最终请求会被组装为{"name": ..., "config": {"served_entities": ..., "route_optimized": ...}, "tags": ..., "rate_limits": ...}并 POST 到端点创建接口。建议新代码直接采用全量 payload 的新形态。

更新端点:config、tags、rate_limits 的专项操作

示例的try块前半部分演示了两种更新方式(databricks.py):

# 更新 served_entities client.update_endpoint( endpoint=name, config={"served_entities": [ ...同样的 gpt-4 实体... ]}, ) # 更新 rate_limits client.update_endpoint( endpoint=name, config={"rate_limits": [{"key": "user", "renewal_period": "minute", "calls": 10}]}, )

update_endpoint在 databricks/init.py 中被标记为已弃用:当 config 仅含rate_limits时,它走PUT /serving-endpoints/<endpoint>/rate-limits;其余情况走PUT /serving-endpoints/<endpoint>/config。源码文档明确建议改用更精确的四个专项方法:

  • update_endpoint_config(endpoint, config):更新served_entities等端点配置(PUT config 路由);
  • update_endpoint_tags(endpoint, config):增删标签(PATCH tags 路由,config 形如{"add_tags": [{"key": "project", "value": "test"}]});
  • update_endpoint_rate_limits(endpoint, config):更新速率限制(PUT rate-limits 路由,config 形如{"rate_limits": [{"calls": 10, "key": "endpoint", "renewal_period": "minute"}]});
  • update_endpoint_ai_gateway(endpoint, config):更新 AI Gateway 配置,如启用用量追踪与推理表落库(PUT ai-gateway 路由)。

这些方法均返回DatabricksEndpoint对象(一种字典风格对象,可像endpoint.name一样用属性访问,见 databricks/init.py)。

查询端点:list 与 get

print(client.list_endpoints()[:3]) # 端点列表前 3 个 print(client.get_endpoint(endpoint=name)) # 单个端点详情
  • list_endpoints()发 GET 请求并从响应中取endpoints字段返回列表(databricks/init.py);
  • get_endpoint(endpoint)发 GET 请求到serving-endpoints/<endpoint>路径(databricks/init.py)。

返回的端点描述包含namecreatorcreation_timestamplast_updated_timestampstateconfigtagsid等字段,其中state.ready可反映端点的就绪状态。

发起推理:predict 与流式预测

predict:普通 Chat 请求

client.predict( endpoint=name, inputs={ "messages": [{"role": "user", "content": "Hello!"}], "max_tokens": 128, }, )

predict的实现(databricks/init.py)把请求 POST 到/<endpoint>/invocations路由(注意此处前缀为/,即完整路径为/serving-endpoints/<endpoint>/invocations)。请求体遵循 OpenAI 兼容的 Chat 补全协议:messages是对话消息数组,max_tokens限制生成长度。返回的响应包含idobjectcreatedmodelchoices(含message.rolemessage.contentfinish_reason)以及usage(token 用量统计)等字段。

predict_stream:流式响应

对于需要逐 token 返回的场景,predict_stream(databricks/init.py)会在请求体中自动追加stream: True,以流式模式发起同样的 invocations 请求,然后逐行解析响应:

  • 每一行必须是data: <value>格式;
  • 遇到data: [DONE]时结束迭代;
  • 其余行 JSON 反序列化后逐个 yield,chunk 中通过choices[0].delta.content携带增量文本。

示例用法:

client = get_deploy_client("databricks") chunk_iter = client.predict_stream( endpoint="databricks-llama-2-70b-chat", inputs={ "messages": [{"role": "user", "content": "Hello!"}], "temperature": 0.0, "n": 1, "max_tokens": 500, }, ) for chunk in chunk_iter: print(chunk)

注意predict_stream对响应格式有严格要求:非data:前缀的行会抛出MlflowException,这正是 Databricks 流式协议的特征。

清理端点

finally: client.delete_endpoint(endpoint=name)

delete_endpoint发送 DELETE 请求到serving-endpoints/<endpoint>(databricks/init.py),删除操作是幂等的。示例把它放在finally中,无论中间哪一步失败都会执行清理,避免残留端点占用额度。

底层实现原理:请求如何打到 Databricks

所有端点操作最终都汇聚到_call_endpoint(databricks/init.py):

  • 路由统一拼接为{prefix}/serving-endpoints/{route},默认前缀/api/2.0predict例外地使用/前缀以得到/serving-endpoints/<endpoint>/invocations);
  • 请求携带自定义头X-Databricks-Endpoints-API-Client: Databricks Deployment Client,用于服务端识别客户端类型;
  • GET 请求参数走 query、其余方法走 JSON body;
  • 支持超时配置(MLFLOW_DEPLOYMENT_PREDICT_TIMEOUTMLFLOW_DEPLOYMENT_PREDICT_TOTAL_TIMEOUTMLFLOW_HTTP_REQUEST_TIMEOUT)与重试机制;
  • 重试码定义在 mlflow/deployments/constants.py:429、500、502、503——值得注意的是,它特意排除了超时类错误,因为对代理的上游供应商而言,长时间超时通常意味着查询参数或模型本身有问题,重试无意义。

此外需要特别说明一个设计差异:DatabricksDeploymentClient中的create_deploymentupdate_deploymentdelete_deploymentlist_deploymentsget_deployment全部直接抛出NotImplementedError(databricks/init.py)。也就是说,对 Databricks 目标而言,模型的托管以"端点 + served_entities"的形态存在,deployment 级 API 不适用,正确用法是本文展示的 endpoint 级 API。这与基类BaseDeploymentClient为自定义插件预留的 deployment 抽象并不冲突——其他目标(如 AWS SageMaker、RedisAI 等插件)仍然使用 deployment 级 API。

通过 CLI 完成相同操作

mlflow/deployments/cli.py 提供了mlflow deployments命令组,覆盖了 Python API 的等价能力,同样面向 Databricks 目标(-t databricks):

  • mlflow deployments create -t databricks --name <name> --model-uri <uri> [-C key=value]:创建部署;
  • mlflow deployments update -t databricks --name <name> [-m <model-uri>] [-C key=value]:更新部署;
  • mlflow deployments delete -t databricks --name <name>:删除部署;
  • mlflow deployments list / get -t databricks:列出/查看部署;
  • mlflow deployments create-endpoint -t databricks --endpoint <name> [-C key=value]:创建端点,-C支持key=value形式的自定义配置,重复配置项会报错;
  • mlflow deployments update-endpoint / delete-endpoint / list-endpoints / get-endpoint -t databricks:端点管理全套命令;
  • mlflow deployments predict -t databricks --endpoint <name> --input-path <file> [--output-path <file>]:从 JSON 或 CSV 文件读取输入并推理,--name--endpoint必须且只能指定一个;
  • mlflow deployments run-local:本地部署测试;
  • mlflow deployments help -t <target>:查看目标特有的 URI 格式与配置说明。

CLI 中的-C/--config选项通过_user_args_to_dict解析为字典(cli.py),支持值中包含=的情况(按第一个=分割),但同一 key 重复指定会直接报错。

小结

通过 examples/deployments/databricks/databricks.py 这一个脚本,可以完整看到 MLflow Deployments 在 Databricks 目标上的标准工作流:密钥托管 → 创建外部模型端点 → 更新配置/限流 → 查询 → 普通与流式推理 → 幂等清理。结合 mlflow/deployments 模块源码,还能进一步理解其统一客户端抽象、端点级 API 的设计差异、Databricks REST 路由的映射关系以及超时重试策略。这套"一套 API 管多家平台"的抽象,让模型部署与 LLM 推理的接入成本大幅降低,也是将外部大模型(OpenAI 等)安全、可控地接入 Databricks 工作区的推荐路径。

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

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

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

立即咨询