Azure Databricks中MLflow模型大小与加载速率限制查询

一、模型注册大小限制 #

• Azure Databricks 官方未单独限制MLflow模型注册的整体大小，模型的核心工件存储依赖于DBFS或Azure Blob Storage（工作区默认存储）。
• 单个DBFS文件默认上限为4.95GB，如果需要存储更大的模型工件，建议拆分文件或直接使用Azure Blob Storage存储（无单个文件硬限制，仅受存储账户配额约束）。
• 模型注册本身仅存储元数据，大小限制主要来自底层存储。工作区存储配额可在Azure门户的Databricks资源配置页调整，具体取决于你的订阅级别。

二、模型加载速率限制 #

• MLflow的load_model请求没有官方硬速率限制，但实际瓶颈来自以下方面：
  • 集群资源：批量预测时频繁调用load_model会占用集群CPU、内存和网络带宽，导致性能下降。
  • 存储IOPS：模型工件存储在DBFS或Blob Storage时，频繁读取会受存储的IOPS限制。
• 没有全局可配置的load_model速率限制开关，但可通过集群配置（如增加节点、更换高性能实例）或缓存策略缓解瓶颈。

三、多模型场景优化方案 #

针对你提到的两种方案的潜在瓶颈，给出实际优化方向：

方案1：按需加载子模型（元模型） #

• 核心优化：缓存已加载的子模型，避免重复调用load_model。可以用Python字典或LRU缓存实现，同一子组模型仅加载一次。
• 示例代码：

from functools import lru_cache
import mlflow

class MetaModel(mlflow.pyfunc.PythonModel):
    @lru_cache(maxsize=100)  # 根据子模型数量调整缓存容量
    def load_submodel(self, submodel_path):
        return mlflow.pyfunc.load_model(submodel_path)
    
    def predict(self, context, model_input):
        submodel_id = model_input["submodel_id"].iloc[0]
        submodel = self.load_submodel(f"models:/submodel_{submodel_id}/Production")
        return submodel.predict(model_input)

方案2：打包多模型为单个自定义模型 #

• 拆分大模型：将子模型按逻辑分组打包，注册多个“分组元模型”，避免单个模型体积过大。
• 延迟加载：在自定义模型中实现按需加载逻辑，仅在需要预测对应子组数据时才加载该子模型，避免启动时加载所有子模型占用内存。
• 存储替换：将大模型工件存储到Azure Blob Storage（而非DBFS），利用Blob的高IOPS和无单个文件硬限制特性，通过Databricks挂载点访问。

四、配置调整说明 #

• 存储配额调整：Azure Databricks工作区的存储配额可在Azure门户对应资源的配置页调整，或联系Azure支持提升配额。
• 集群优化：批量预测场景建议使用高性能计算集群（如DSv3系列实例），增加节点数或调整实例类型提升IO和计算能力。

内容的提问来源于stack exchange，提问作者jtolle