如何在Foundry代码仓库中使用Broadcast及在Pandas UDF中访问模型

获取Spark Context的方法 #

在Foundry的Python代码仓库里，你可以通过活跃的SparkSession实例拿到Spark Context，代码如下：

from pyspark.sql import SparkSession

spark = SparkSession.getActiveSession()
sc = spark.sparkContext
# 之后就可以用sc广播模型
model_bcast = sc.broadcast(model)

其他传入模型的方式 #

除了广播，还有两种常用方案：

1. 序列化模型后作为参数传入UDF #

如果你的模型支持序列化（比如用pickle），可以把模型序列化后，通过functools.partial绑定到UDF上，示例代码：

import pickle
from functools import partial
from pyspark.sql.functions import pandas_udf

# 序列化模型
model_bytes = pickle.dumps(model)

def predict_func(model_bytes, df):
    # 反序列化模型
    model = pickle.loads(model_bytes)
    # 执行预测逻辑
    df['prediction'] = model.predict(df['feature'])
    return df

# 绑定模型参数到函数
bound_predict_func = partial(predict_func, model_bytes)
# 注册为Pandas UDF
predict_udf = pandas_udf(bound_predict_func, returnType="struct<feature:double,prediction:double>")

注意：这种方式适合小型模型，大模型的话不如广播高效——广播是每个节点只存一份模型，而这种方式每个任务都会复制一份。

2. 从Foundry Model Registry直接加载 #

如果模型已经上传到Foundry的Model Registry，可以在UDF内部直接加载，无需提前处理广播或序列化，示例代码：

from foundry.ml.models import Model
from pyspark.sql.functions import pandas_udf

def predict_func(df):
    # 从Model Registry加载指定路径的模型
    model = Model.get_by_path("/path/to/model/in/registry").load()
    df['prediction'] = model.predict(df['feature'])
    return df

predict_udf = pandas_udf(predict_func, returnType="struct<feature:double,prediction:double>")

这种方式由Foundry负责模型的分发管理，省去了手动处理模型传递的麻烦。

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