Spark 中利用 MLflow 加载机器学习模型并对 pandas-on-Spark DataFrame 进行预测

大数据数据分析批处理流处理机器学习图计算【免费下载链接】sparkApache Spark - A unified analytics engine for large-scale data processing项目地址https://gitcode.com/gh_mirrors/sp/spark点击查看免费下载导读本文围绕 Apache Spark 的 pandas-on-Spark即pyspark.pandas与 MLflow 的集成模块pyspark.pandas.mlflow展开系统讲解如何把任意实现了 MLflowpyfuncflavor 的模型scikit-learn、PyTorch、TensorFlow 等加载为统一的预测器并直接作用在分布式 pandas-on-Spark DataFrame 上完成大规模推理。读完本文你将掌握load_model与PythonModelWrapper.predict的完整用法、返回类型推断机制、模型与 DataFrame 列合并的注意事项以及该模块在源码层的实现原理可直接在自己的 MLflow 模型 Spark 分布式推理场景中落地。本文对应仓库中的 API 参考文档为 python/docs/source/reference/pyspark.pandas/ml.rst核心实现位于 python/pyspark/pandas/mlflow.py。模块定位让 MLflow 模型与 pandas-on-Spark 数据框无缝对接pyspark.pandas.mlflow是 pandas-on-Spark 的机器学习工具模块解决的核心问题是训练阶段用 MLflow 统一管理的模型如何在 Spark 分布式环境中直接对 pandas-on-Spark DataFrame 做批量预测。从 API 参考文档的定义看该模块的前提条件非常明确任意 MLflow 模型只要实现了 pyfunc flavor即可用于 pandas-on-Spark DataFrame。绝大多数主流框架scikit-learn、pytorch、tensorflow 等都满足这一要求。同时文档给出了一个关键约束使用本模块必须安装 MLflow 包The MLflow package must be installed in order to use this module。也就是说pyspark.pandas.mlflow本身是一个轻量适配层它不依赖任何特定框架而是统一依赖 MLflow 的pyfunc接口把任意框架训练出的模型翻译为可在 Spark 上分布式执行的 UDF。这种设计让模型注册、加载、推理的链路与具体框架解耦模型只要能在 MLflow 中保存和加载就能接入 pandas-on-Spark。模块对外只暴露两个符号见 python/pyspark/pandas/mlflow.py 的__all__PythonModelWrapper围绕 MLflow Python 对象模型的封装类作为 pandas-on-Spark 上的预测器load_model加载 MLflow 模型并返回PythonModelWrapper。核心 APIload_model 与 PythonModelWrapper.predictload_model加载模型入口load_model的函数签名与参数语义如下def load_model( model_uri: str, predict_type: Union[str, type, Dtype] infer ) - PythonModelWrapper:参数类型说明model_uristr指向模型的 URI支持 MLflow 支持的各种寻址方式如runs:/run_id/model、models:/name/version、本地文件系统路径等详见 MLflow 文档predict_typePython 基础类型 / numpy 基础类型 / Spark 类型 /infer调用模型predict时期望的返回类型指定infer时包装器会尝试依据模型类型自动推断返回类型返回值为PythonModelWrapper。该包装器遵循mlflow.pyfunc.PythonModel的接口约定。PythonModelWrapper.predict对两种数据框做预测def predict(self, data: Union[DataFrame, pd.DataFrame]) - Union[Series, pd.Series]:predict会根据输入类型走两条完全不同的执行路径输入为 pandas DataFrame直接调用底层pyfunc对象的predict返回的是底层模型的原生输出通常是 pandas Series 或 numpy 数组适合小批量、单机场景输入为 pandas-on-Spark DataFrame通过 Spark UDF 将模型映射到分布式数据上执行返回 pandas-on-Spark Series可在集群上对海量数据进行并行推理。如果传入其他类型则会抛出ValueError(unknown data type: ...)。源码视角模型封装与分布式推理是如何实现的PythonModelWrapper的实现体现了惰性加载 类型推断 UDF 包装三个关键设计全部代码见 python/pyspark/pandas/mlflow.py。1. 惰性加载三个底层对象lazy_property包装器内部通过lazy_property缓存三个对象首次访问时才真正初始化_return_type由predict_type提示转换成的 SparkDataType。逻辑上当提示为infer或为空时默认使用np.float64——对应默认的连续值预测场景。转换通过pyspark.pandas.typedef中的as_spark_type完成_model调用mlflow.pyfunc.load_model(model_uri...)得到底层模型对象_model_udf调用mlflow.pyfunc.spark_udf(spark, model_uri..., result_typeself._return_type)得到可直接挂载到 Spark 列上的 UDF其中spark取自pyspark.pandas.utils.default_session()。源码中还留有明确的 TODO 注释目前返回类型推断逻辑比较简单仅覆盖连续值预测这一默认场景后续可针对sklearn.Classifier应返回整数或类别以及 PyTorch / TensorFlow / Keras 模型依据输出类型做更智能的推断但作者认为这部分更适合放在 MLflow 侧而非此处完成。这意味着当前版本对分类模型的predict_type需要用户显式指定而不是完全依赖自动推断。2. 分布式路径struct spark_udf 内部结构替换对 pandas-on-Spark DataFrame 的预测是模块的核心价值所在其执行链路为s struct(*data.columns) # 将整行特征打包成一个 struct 列 return_col self._model_udf(s) # 对 struct 列应用 MLflow pyfunc UDF column_labels [(col,) for col in data._internal.spark_frame.select(return_col).columns] internal data._internal.copy( column_labelscolumn_labels, data_spark_columns[return_col], data_fieldsNone ) return first_series(DataFrame(internal)) # 取回第一列作为 pandas-on-Spark Series可以看到实现要点用pyspark.sql.functions.struct把 DataFrame 的所有特征列打包为一个 struct 列一次传给 UDF避免逐列序列化spark_udf生成的 UDF 在 Spark 引擎内分布式执行每个 executor 上的分区数据由 MLflow 加载的模型完成本地推理预测结果作为一个新的 Spark 列return_col通过data._internal.copy替换列标签与数据列构造出一个新的 pandas-on-Spark DataFrame再用first_series取出唯一的预测列为 Series。这种列级 UDF 内部元数据复制的方式让用户感受到的 API 与 pandas 原生体验一致而底层计算已经被 Spark 调度到集群上。3. 模块在 pandas-on-Spark 中的集成该模块与 pandas-on-Spark 生态有两处明显的集成痕迹在 python/pyspark/pandas/namespace.py 的依赖版本检测列表中mlflow与pyspark、pandas、numpy、pyarrow等并列说明它是该发行版预期可用的可选依赖之一在 python/pyspark/pandas/usage_logging/init.py 中mlflow模块与mlflow.PythonModelWrapper类被纳入使用日志统计范围以try/except ImportError包裹未安装 MLflow 时静默跳过。完整实战scikit-learn 模型 MLflow pandas-on-Spark 推理以下示例完整来自 load_model 的 docstring是模块自带的 doctest 级可运行示例可直接复制验证。第一步初始化 MLflow 环境from mlflow.tracking import MlflowClient, set_tracking_uri import mlflow.sklearn from tempfile import mkdtemp d mkdtemp(pandas_on_spark_mlflow) set_tracking_uri(fsqlite:///{d}/mlflow.db) # 使用本地 sqlite 作为 tracking 后端 client MlflowClient() exp_id mlflow.create_experiment(my_experiment) exp mlflow.set_experiment(my_experiment)第二步训练并记录 scikit-learn 线性回归模型假设我们要学习的目标函数为y log(2 x)以x1、x2为特征import pandas as pd import numpy as np from sklearn.linear_model import LinearRegression train pd.DataFrame({x1: np.arange(8), x2: np.arange(8)**2, y: np.log(2 np.arange(8))}) train_x train[[x1, x2]] train_y train[[y]] with mlflow.start_run(): lr LinearRegression() lr.fit(train_x, train_y) mlflow.sklearn.log_model(lr, model)第三步加载模型并对 pandas-on-Spark DataFrame 预测from pyspark.pandas.mlflow import load_model import pyspark.pandas as ps run_info client.search_runs(exp_id)[-1].info model load_model(runs:/{run_id}/model.format(run_idrun_info.run_id)) prediction_df ps.DataFrame({x1: [2.0], x2: [4.0]}) prediction_df[prediction] model.predict(prediction_df) print(prediction_df)输出不同环境浮点精度可能有细微差异x1 x2 prediction 0 2.0 4.0 1.355551第四步对 pandas DataFrame 预测单机路径同一个model对象也接受 pandas DataFrame返回底层 pyfunc 的原生输出model.predict(prediction_df[[x1, x2]].to_pandas()) # array([[1.35555142]])值得注意的是同一个PythonModelWrapper同时支持 pandas 与 pandas-on-Spark 两种输入这种双路径设计让用户在本地调试用小数据pandas与集群推理用大数据pandas-on-Spark之间切换时无需更换 API。关键注意事项预测结果列如何与原始 DataFrame 合并官方文档在Notes一节明确指出了当前版本的一个重要限制目前模型预测结果只能与现有 DataFrame 合并回去。其他列必须手动 join。即model.predict(features)返回的 Series 只能赋值回用于预测的那个特征子集 DataFrame不能直接赋值给包含额外列的父 DataFrame。例如以下代码会失败报错信息为数据框未对齐df ps.DataFrame({x1: [2.0], x2: [3.0], z: [-1]}) features df[[x1, x2]] y model.predict(features) features[y] y # 可用预测列拼回 features 自身 df[y] y # 会失败features 与 df 结构不对齐官方给出的当前 workaround 是使用.merge()以特征值为连接键把预测结果拼回原表features[y] y everything df.merge(features, on[x1, x2]) print(everything)输出x1 x2 z y 0 2.0 3.0 -1 1.376932在实际生产代码中若x1、x2并非唯一键merge可能造成行数膨胀建议结合实际键的唯一性选择连接列或先对 DataFrame 进行行号/自增键标记再 merge。predict_type 的选择策略predict_type是load_model中唯一需要用户决策的参数官方语义为传 Python 基础类型、numpy 基础类型或 Spark 类型包装器通过as_spark_type直接转换为 SparkDataType作为 UDF 的result_type传infer默认值包装器尝试自动推断但当前实现只覆盖默认的连续值预测场景即推断为np.float64。从源码可以得出的实践建议是回归模型 / 输出连续浮点的模型直接使用默认predict_typeinfer即可分类模型输出整数类别或概率鉴于 mlflow.py 中的 TODO 注释表明分类推断尚未实现建议显式传入对应的 Spark 类型如IntegerType、DoubleType以确保spark_udf的result_type与模型输出一致避免类型不匹配导致的运行期错误返回多列输出的模型当前实现只取 UDF 结果的第一列first_series多输出场景需要自行评估是否符合预期。适用场景与边界总结适合的使用方式团队已用 MLflow 统一管理模型tracking、registry希望在同一套 Spark 作业里完成读数据 → 分布式推理 → 结果入库模型来自任何支持 pyfunc flavor 的框架希望推理代码与框架无关需要在小数据pandas上快速验证、再平滑迁移到大数据的 pandas-on-Spark 推理。需要留意的边界必须先安装 MLflowpip install mlflow否则from pyspark.pandas.mlflow import load_model无法正常使用预测列合并受限跨列赋值需走.merge()返回类型推断以np.float64为默认兜底分类/多输出场景应显式声明predict_type该模块面向批量离线/近线推理非实时单条请求的低延迟服务场景。进一步探索模块 API 参考文档python/docs/source/reference/pyspark.pandas/ml.rst完整源码与内嵌 doctest 示例python/pyspark/pandas/mlflow.py类型转换工具as_spark_type定义python/pyspark/pandas/typedef/typehints.py依赖版本检测含 mlflowpython/pyspark/pandas/namespace.py使用日志统计集成python/pyspark/pandas/usage_logging/init.py赞分享大数据数据分析批处理流处理机器学习图计算【免费下载链接】sparkApache Spark - A unified analytics engine for large-scale data processing项目地址https://gitcode.com/gh_mirrors/sp/spark点击查看免费下载相关推荐使用Spark进行分布式机器学习spark-sklearn使用Spark进行分布式机器学习spark sklearn 项目介绍 spark sklearn 是一个旨在将Apache Spark与流行的scikit lHound会话管理实战暂停、断点恢复与长期迭代审计的完整教程Hound会话管理实战暂停、断点恢复与长期迭代审计的完整教程 Hound 是一款语言无关的 AI 代码审计工具它通过自主构建并不断完善的 自适应知识图谱 大数据数据分析批处理流处理机器学习图计算技术解构Wand-Enhancer如何重新定义游戏修改工具生态技术解构Wand Enhancer如何重新定义游戏修改工具生态 在游戏修改工具领域付费订阅模式长期占据主导地位用户常常面临功能限制与持续支出的矛盾。Wan大数据数据分析批处理流处理机器学习图计算上一篇REFramework项目中native布局重复问题的分析与解决下一篇5分钟上手直播中实时显示键盘和手柄输入的免费神器创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考