Databricks部署MLlib ALS推荐模型，数据量增大致端点报错

背景 #

基于MLlib的ALS模型训练商品推荐模型，通过Databricks Model Serving部署服务。数据集result_df约4万行，结构如下：

Schema：

('row_id', 'bigint'),
('user_id', 'int'),
('product_id', 'int'),
('rating', 'double')

问题现象 #

使用4千行左右的小数据集训练时，模型端点可正常返回预测结果；但训练数据量超过1万行后，调用评分接口出现如下报错：

Encountered an unexpected error while evaluating the model. Verify that the input is compatible with the model for inference. Error 'An error occurred while calling o219.collectToPython.

org.apache.spark.SparkException: Job aborted due to stage failure: Task 0 in stage 11.0 failed 1 times, most recent failure: Lost task 0.0 in stage 11.0 (TID 11) (<driver_id> executor driver): TaskResultLost (result lost from block manager)
Driver stacktrace:
    at org.apache.spark.scheduler.DAGScheduler.failJobAndIndependentStages(DAGScheduler.scala:2785)
    ...（省略中间栈信息）
py4j.protocol.Py4JJavaError: An error occurred while calling o219.collectToPython.

补充说明 #

1. 排除内存问题：同一集群上用100万行数据集训练相同模型，端点可正常工作。
2. 排除输入兼容性问题：两种数据量下的训练、评分数据集结构完全一致，仅数据量不同。

相关代码 #

X = result_df.sample(fraction=FRAC) # FRAC取0.1或0.3，控制训练数据量

als = ALS(
  rank=100,
  maxIter=5,  
  userCol='user_id', 
  itemCol='product_id', 
  ratingCol='rating',
  implicitPrefs=True,
  coldStartStrategy="drop"
)

# 训练模型
model = als.fit(X)

mlflow.spark.log_model(model, artifact_path='model', registered_model_name=MODEL_NAME)

import os
import requests
import numpy as np
import pandas as pd
import json


def create_tf_serving_json(data):
    return {'inputs': {name: data[name].tolist() for name in data.keys()} if isinstance(data, dict) else data.tolist()}

def score_model(dataset):
    url = <模型服务端点URL>
    headers = {'Authorization': f'Bearer {os.environ.get("DATABRICKS_TOKEN")}', 
'Content-Type': 'application/json'}
    ds_dict = {'dataframe_split': dataset.to_dict(orient='split')} if isinstance(dataset, pd.DataFrame) else create_tf_serving_json(dataset)
    data_json = json.dumps(ds_dict, allow_nan=False)
    response = requests.request(method='POST', headers=headers, url=url, data=data_json)
    if response.status_code != 200:
        raise Exception(f'Request failed with status {response.status_code}, {response.text}')

    return response.json()

score_model(X.toPandas()) # 此调用触发上述报错

问题分析与解决方案 #

核心原因 #

报错中的TaskResultLost (result lost from block manager)指向：当训练数据量超过1万行时，ALS模型生成的用户/物品特征矩阵大小达到了Spark广播变量的默认阈值，导致BlockManager无法处理广播结果，进而任务失败。而100万行数据训练的模型能正常工作，是因为更大的数据集可能让特征矩阵的稀疏度更高，实际占用的广播资源反而未触发阈值。

解决方案 #

1. 调整Spark广播配置

在训练模型前或日志模型时，修改Spark的广播阈值参数，避免自动广播触发的限制：

# 训练前设置Spark配置
spark.conf.set("spark.sql.autoBroadcastJoinThreshold", "-1")  # 禁用自动广播阈值
spark.conf.set("spark.broadcast.blockSize", "128m")  # 增大广播块大小

2. 日志模型时绑定Spark配置

在mlflow.spark.log_model中指定服务端使用的Spark配置，确保部署后的环境生效：

mlflow.spark.log_model(
    model, 
    artifact_path='model', 
    registered_model_name=MODEL_NAME,
    spark_conf={
        "spark.sql.autoBroadcastJoinThreshold": "-1",
        "spark.broadcast.blockSize": "128m"
    }
)

3. 优化推理数据处理逻辑

避免一次性收集全量预测结果，改用分批处理减少单批次数据量：

# 修改评分函数，分批处理输入数据
def score_model_batch(dataset, batch_size=1000):
    url = <模型服务端点URL>
    headers = {'Authorization': f'Bearer {os.environ.get("DATABRICKS_TOKEN")}', 
'Content-Type': 'application/json'}
    
    results = []
    for i in range(0, len(dataset), batch_size):
        batch = dataset.iloc[i:i+batch_size]
        ds_dict = {'dataframe_split': batch.to_dict(orient='split')}
        data_json = json.dumps(ds_dict, allow_nan=False)
        response = requests.post(url, headers=headers, data=data_json)
        if response.status_code != 200:
            raise Exception(f'Batch {i//batch_size} failed: {response.status_code}, {response.text}')
        results.extend(response.json()['predictions'])
    
    return results

4. 检查Model Serving端点资源配置

确认端点的实例类型和资源配额，适当提升CPU/内存配置，避免推理时资源不足导致的结果丢失。

验证步骤 #

1. 应用上述配置后重新训练并部署模型。
2. 使用相同测试数据调用评分接口，观察是否仍报错。
3. 若问题持续，查看Spark日志中的BlockManager指标，确认广播变量实际大小是否仍超出阈值。

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