MLflow PySpark ML 自动日志记录（Autologging）实战指南：从 LogisticRegression 到 OneVsRest

MLflow PySpark ML 自动日志记录Autologging实战指南从 LogisticRegression 到 OneVsRest【免费下载链接】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/mlflowMLflow 的 PySpark ML 自动日志记录mlflow.pyspark.ml.autolog()能够在无需手动埋点的前提下自动捕获 Spark ML Estimator 的训练参数、指标、数据集信息与模型工件是构建可追踪、可复现的 Spark 机器学习流水线的关键能力。本文以仓库 examples/pyspark_ml_autologging 目录下的两个完整示例为主线结合 mlflow/pyspark/ml/init.py 的底层实现与 log_model_allowlist.txt 白名单机制逐行讲解如何跑通 PySpark ML 自动日志记录并深入剖析其参数体系、记录内容与实现原理读完后你将能够独立配置和使用这一能力。示例目录概览该目录是 MLflow 仓库中专门用于演示 PySpark ML 自动日志记录的官方示例包含两个文件文件说明logistic_regression.py训练一个LogisticRegression逻辑回归模型one_vs_rest.py训练一个OneVsRest一对多多分类模型两个示例均基于内置的 Iris鸢尾花数据集展示了 PySpark ML 自动日志记录最基本的两种使用形态直接训练单一模型以及通过元估计器meta-estimator包装基础分类器。示例只需几行代码即可运行是理解 autologging 行为的最佳起点。基础示例LogisticRegression 模型的自动记录logistic_regression.py 完整代码如下我们逐段拆解其工作流程from pyspark.ml.classification import LogisticRegression from pyspark.ml.feature import VectorAssembler from pyspark.sql import SparkSession from sklearn.datasets import load_iris import mlflow with SparkSession.builder.getOrCreate() as spark: df load_iris(as_frameTrue).frame.rename(columns{target: label}) df spark.createDataFrame(df) df VectorAssembler(inputColsdf.columns[:-1], outputColfeatures).transform(df) train, test df.randomSplit([0.8, 0.2]) mlflow.pyspark.ml.autolog() lor LogisticRegression(maxIter5) with mlflow.start_run(): lorModel lor.fit(train) pred lorModel.transform(test) pred.select(lorModel.getPredictionCol()).show(10)1. 数据准备用 sklearn 生成 DataFrame示例借助sklearn.datasets.load_iris加载经典鸢尾花数据集将其重命名为 ML 约定俗成的label目标列再通过spark.createDataFrame转换为 Spark DataFrame。随后使用 PySpark ML 的VectorAssembler把前 4 个特征列组装成向量类型的features列——这是 Spark ML 流水线中标准的特征向量化步骤VectorAssembler(inputColsdf.columns[:-1], outputColfeatures)表示除最后一列标签外的所有列作为输入。最后通过randomSplit([0.8, 0.2])将数据划分为 80% 训练集与 20% 测试集。2. 开启自动日志记录mlflow.pyspark.ml.autolog()这是整个示例的核心一行。它通过safe_patch机制mlflow/pyspark/ml/init.py对pyspark.ml.base.Estimator.fit、pyspark.ml.base.Model.transform和pyspark.ml.evaluation.Evaluator.evaluate三个方法打补丁从而在训练、预测、评估三个阶段自动收集信息调用Estimator.fit时记录参数、标签、训练数据集与模型工件调用Model.transform后若紧接着调用Evaluator.evaluate则自动捕获评估指标即训练后指标post-training metrics。注意与其它框架不同PySpark 的 autolog 在调用autolog()时会立即激活立即对pyspark与pyspark.ml模块执行setup_autologging而不是等待模块被导入后才触发这一特殊处理可见 mlflow/tracking/fluent.py。3. 训练并查看自动记录结果with mlflow.start_run(): lorModel lor.fit(train)在mlflow.start_run()的上下文内调用fit()后MLflow 会自动把以下内容写入当前 Run参数Parameters来自estimator.params的全部配置项例如maxIter5、regParam、elasticNetParam、tol、family等均以param_namevalue形式记录标签Tags估算器类名如LogisticRegression以及完全限定类名如pyspark.ml.classification.LogisticRegression由_get_estimator_info_tags生成mlflow/pyspark/ml/init.py数据集Dataset训练集以SparkDataset形式通过mlflow.log_input(dataset, train)记录mlflow/pyspark/ml/init.py形成数据血缘模型工件Artifact以mlflow.sparkflavor 记录的 MLflow Model落盘到model目录mlflow/pyspark/ml/init.py。4. 预测并打印结果pred lorModel.transform(test) pred.select(lorModel.getPredictionCol()).show(10)lorModel.transform(test)在测试集上产生包含prediction列的 DataFrame前 10 行预测结果被打印到控制台。这一行在开启log_post_training_metrics默认开启后还会被patched_transform捕获mlflow/pyspark/ml/init.py为后续Evaluator.evaluate的指标关联做准备。进阶示例OneVsRest 元估计器的自动记录one_vs_rest.py 与基础示例的结构几乎一致唯一区别在于用OneVsRest包装了逻辑回归from pyspark.ml.classification import LogisticRegression, OneVsRest from pyspark.ml.feature import VectorAssembler from pyspark.sql import SparkSession from sklearn.datasets import load_iris import mlflow with SparkSession.builder.getOrCreate() as spark: df load_iris(as_frameTrue).frame.rename(columns{target: label}) df spark.createDataFrame(df) df VectorAssembler(inputColsdf.columns[:-1], outputColfeatures).transform(df) train, test df.randomSplit([0.8, 0.2]) mlflow.pyspark.ml.autolog() lor LogisticRegression(maxIter5) ovr OneVsRest(classifierlor) with mlflow.start_run(): ovrModel ovr.fit(train) pred ovrModel.transform(test) pred.select(ovrModel.getPredictionCol()).show(10)OneVsRest是一种经典的一对多多分类元估计器对 K 个类别分别训练 K 个二分类器这里是 K 个LogisticRegression预测时取置信度最高的分类器输出。这个示例展示了 autologging 对**元估计器meta-estimator**的特殊支持不重复记录子估计器OneVsRest.fit()内部会多次调用子分类器的fit()。autologging 通过_SparkTrainingSession会话机制mlflow/pyspark/ml/init.py跟踪训练会话allow_childrenFalse确保仅在最外层fit()调用时记录一次内部子估计器的fit()调用全部被跳过避免产生大量噪音 Run层次结构工件对于OneVsRest、Pipeline、CrossValidator、TrainValidationSplit这类元估计器autologging 会额外生成并记录estimator_info.json工件其中hierarchy条目完整描述了元估计器的嵌套结构mlflow/pyspark/ml/init.py。从源码_should_log_hierarchymlflow/pyspark/ml/init.py可见Pipeline、OneVsRest、参数搜索估计器以及包含 stage 类型参数的估计器都会触发层次记录白名单递归校验_should_log_modelmlflow/pyspark/ml/init.py对OneVsRestModel会递归检查其内部模型spark_model.models[0]是否在白名单中只有被包装的模型也允许记录时整个OneVsRestModel才会被自动记录。autolog() 完整参数详解mlflow.pyspark.ml.autolog 的函数签名为def autolog( log_modelsTrue, log_datasetsTrue, disableFalse, exclusiveFalse, disable_for_unsupported_versionsFalse, silentFalse, log_post_training_metricsTrue, registered_model_nameNone, log_input_examplesFalse, log_model_signaturesTrue, log_model_allowlistNone, extra_tagsNone, ):各参数含义与源码依据如下源自 autolog 的 docstring参数默认值作用log_modelsTrue若训练出的模型在白名单中则作为 MLflow 模型工件记录False则不记录模型。注意内置白名单会排除部分体积过大的模型如 ALS可通过自定义白名单文件调整log_datasetsTrue是否将训练/评估数据集信息写入 MLflow Tracking形成数据血缘disableFalseTrue时禁用整个 PySpark ML autologging 集成exclusiveFalseTrue时自动记录内容不写入用户创建的 fluent RunFalse时写入当前活跃 Run可能由用户创建disable_for_unsupported_versionsFalseTrue时对未经测试或与当前 MLflow 客户端不兼容的 PySpark 版本禁用 autologgingsilentFalseTrue时抑制 autologging 过程中的全部事件日志与警告log_post_training_metricsTrue是否记录训练后指标即Evaluator.evaluate的结果详见下文训练后指标小节registered_model_nameNone若指定每次训练出的模型都会注册为该名称下的新模型版本模型不存在则自动创建log_input_examplesFalseTrue时从训练集采集输入示例并随模型工件记录log_model_signaturesTrue是否采集并记录模型输入/输出的 ModelSignature注意当前仅支持标量 Spark 数据类型若输入输出含pyspark.ml.linalg.Vector等非标量类型则不会记录签名log_model_allowlistNone覆盖默认模型记录白名单优先级高于 Spark 配置spark.mlflow.pysparkml.autolog.logModelAllowlistFileextra_tagsNone为 autologging 创建的每个托管 Run 附加的标签字典实际使用中常见的组合包括# 关闭模型记录仅保留参数与指标追踪 mlflow.pyspark.ml.autolog(log_modelsFalse) # 自动注册模型版本 mlflow.pyspark.ml.autolog(registered_model_nameiris_logistic_regression) # 记录输入示例与模型签名 mlflow.pyspark.ml.autolog(log_input_examplesTrue, log_model_signaturesTrue) # 关闭训练后指标捕获 mlflow.pyspark.ml.autolog(log_post_training_metricsFalse)autologging 记录的四类信息根据 autolog 源码 docstring开启后每次Estimator.fit调用会记录以下四类内容1. 参数Parameters通过estimator.params获取全部超参数。特殊规则若某个参数值本身也是一个 Estimator例如OneVsRest.classifier则被包装估计器的参数也会被记录嵌套参数键格式为{estimator_uid}.{param_name}。以OneVsRest为例LogisticRegression(maxIter5)的参数会以类似classifier_maxIter的键展开记录源码见_log_estimator_params与_get_param_map的实现 mlflow/pyspark/ml/init.py。2. 标签Tags记录两个标签估计器类名如LogisticRegression与完全限定类名如pyspark.ml.classification.LogisticRegression由_get_estimator_info_tagsmlflow/pyspark/ml/init.py生成。3. 训练后指标Post-training Metrics这是 PySpark ML autologging 的特色能力调用Evaluator.evaluate后MLflow 会自动把评估结果记录为对应 Run 的指标。指标键格式为{metric_name}[-{call_index}]_{dataset_name}metric_nameEvaluator.getMetricName()返回的指标名如accuracy、f1call_index同一评估指标被多次调用时从 2 开始追加序号首次调用无序号dataset_name调用model.transform时传入的预测输入数据集变量的变量名若数据集是未赋值的中间表达式则记为unknown_dataset同名变量后续出现会追加索引。实现上由_AutologgingMetricsManager全局管理器mlflow/pyspark/ml/init.py维护模型 → Run与预测结果 → 数据集名的映射patched_transform注册预测结果mlflow/pyspark/ml/init.pypatched_evaluate在评估时反查 Run 并写入指标mlflow/pyspark/ml/init.py同时每次 evaluate 调用还会把评估数据集以eval上下文的标签记录到 Tracking并输出metric_info.json工件包含评估器类名与参数。限制对预测结果做二次变换后派生的对象MLflow 无法再关联到原始 Run。4. 工件Artifactsmodel以mlflow.sparkflavor 记录的已拟合模型mlflow.spark.log_model写入见 mlflow/pyspark/ml/init.pyestimator_info.json估计器层次结构描述包含hierarchy条目对嵌套的 Pipeline 各阶段、被包装的模型与评估器等做了展开mlflow/pyspark/ml/init.pymetric_info.json训练后指标对应的评估器信息类名与参数best_parameters.json与search_results.csv仅参数搜索类元估计器CrossValidator、TrainValidationSplit会额外记录——前者为搜索到的最优参数后者为两列params、metric的搜索结果表best_model参数搜索模型还会在best_model目录下额外记录最优模型mlflow/pyspark/ml/init.py。参数搜索与流水线的自动记录除OneVsRest外autologging 对CrossValidator、TrainValidationSplit等参数搜索元估计器还有专属支持子 Run 记录_create_child_runs_for_parameter_searchmlflow/pyspark/ml/init.py会为每一组被探索的参数创建子 Run记录该组参数的评估指标estimator_info.json扩展额外记录tuning_parameter_map_list调参使用的全部参数映射与tuned_estimator_parameter_map被调优估计器的参数映射最优参数以best_前缀的参数形式如best_maxIter写入 Run 参数并输出best_parameters.json工件。对于Pipelineautologging 记录整个流水线的层次结构与最终PipelineModel并且流水线内部的 featurizer 在作为流水线阶段参与fit()时也会被纳入记录单独调用 featurizer 的fit()则不会被记录见 mlflow/pyspark/ml/init.py。模型记录白名单机制autologging 并不会记录任意模型——只有白名单内的模型才会被自动落盘这是出于存储与性能的考虑部分模型可能非常大。白名单机制由 log_model_allowlist.txt 与 mlflow/pyspark/ml/init.py 共同实现内置白名单包含以下类别完整内容见 log_model_allowlist.txt分类模型LogisticRegressionModel、LinearSVCModel、DecisionTreeClassificationModel、GBTClassificationModel、RandomForestClassificationModel、NaiveBayesModel聚类模型BisectingKMeansModel、KMeansModel、GaussianMixtureModel回归模型LinearRegressionModel、DecisionTreeRegressionModel、GBTRegressionModel、GeneralizedLinearRegressionModel、RandomForestRegressionModel、AFTSurvivalRegressionModel特征化模型StringIndexerModel、VectorIndexerModel、StandardScalerModel、MinMaxScalerModel、OneHotEncoderModel、CountVectorizerModel、IDFModel等约 18 个复合与流水线模型OneVsRestModel、PipelineModel超参调优模型CrossValidatorModel、TrainValidationSplitModelSynapseML 系列通过通配符synapse.ml.cognitive.*、synapse.ml.lightgbm.*等。自定义白名单的两种方式Spark 配置在创建 SparkSession 前设置spark ( SparkSession.builder .config(spark.mlflow.pysparkml.autolog.logModelAllowlistFile, /path/to/my_allowlist.txt) .getOrCreate() )白名单文件为换行分隔的完全限定类名列表读取失败时自动回退到内置白名单log_model_allowlist参数在autolog()中直接传入集合优先级最高mlflow.pyspark.ml.autolog(log_model_allowlist{pyspark.ml.classification.LogisticRegressionModel})白名单匹配支持尾部通配符一个*如synapse.ml.lightgbm.*未命中白名单的模型不会自动记录但可通过mlflow.spark.log_model()显式记录警告信息见_get_warning_msg_for_skip_log_modelmlflow/pyspark/ml/init.py。底层实现补丁与训练会话从 mlflow/pyspark/ml/init.py 可以看到autolog()的核心是三个safe_patch注册safe_patch(AUTOLOGGING_INTEGRATION_NAME, Estimator, fit, patched_fit, manage_runTrue, extra_tagsextra_tags) if log_post_training_metrics: safe_patch(..., Model, transform, patched_transform, manage_runFalse) safe_patch(..., Evaluator, evaluate, patched_evaluate, manage_runFalse)关键执行流程在fit_mlflowmlflow/pyspark/ml/init.pyfeaturizer 的直接fit()调用被跳过pyspark.ml.feature.*前缀判断传入params为列表时跳过对应fitMultiple返回模型迭代器的场景训练 DataFrame 被persist(StorageLevel.MEMORY_AND_DISK)持久化训练完成后unpersist()避免大数据集被重复计算先执行_log_pretraining_metadata参数、标签、层次结构、数据集再调用原始fit()最后执行_log_posttraining_metadata训练后指标、模型工件、参数搜索结果。patched_fit通过_SparkTrainingSession(estimatorself, allow_childrenFalse)建立训练会话会话确保元估计器内部子估计器的fit()不会被重复记录同时用_AUTOLOGGING_METRICS_MANAGER.disable_log_post_training_metrics()避免训练期间误捕获指标。fit()返回的模型随后被register_model注册供后续transform/evaluate关联。测试验证与运行前置条件仓库对 PySpark ML autologging 有完善的测试覆盖核心测试位于 tests/spark/autologging/ml/test_pyspark_ml_autologging.py其中test_basic_estimator、test_meta_estimator_fit、test_pipeline、test_param_search_estimator分别验证了基础估计器、元估计器、流水线与参数搜索场景的记录行为test_models_in_allowlist_exist校验白名单中的模型类真实存在另有 tests/spark/autologging/ml/test_pyspark_ml_autologging_custom_allowlist.py 专门验证自定义白名单与通配符匹配逻辑如test_should_log_model_with_wildcards_in_allowlist。运行示例需满足Spark 3.0 或以上autolog源码 docstring 明确要求 This API requires Spark 3.0 or above安装mlflow、pyspark与scikit-learn用于加载 Iris 数据集运行方式python examples/pyspark_ml_autologging/logistic_regression.py python examples/pyspark_ml_autologging/one_vs_rest.py运行前建议设置 MLflow Tracking URI如mlflow.set_tracking_uri(http://host:port)或本地./mlruns运行结束后可通过mlflow ui查看自动记录的 Run、参数、指标与模型工件。已知限制与注意事项综合 autolog 文档 与源码实现使用时有以下几点限制需要注意模型签名仅支持标量类型log_model_signaturesTrue时若输入/输出包含pyspark.ml.linalg.Vector等非标量 Spark 数据类型模型签名不会被记录相关逻辑见_get_columns_with_unsupported_data_typemlflow/pyspark/ml/init.py大模型可能不被自动记录出于性能与存储考虑内置白名单排除了部分体积较大的模型如 ALSPipelineModel 与超参调优元估计器模型的 autologging 在部分场景尚未完全支持训练后指标的关联局限对预测结果做进一步变换后得到的派生数据集MLflow 无法自动关联回原始 Run非线程安全autolog()本身不是线程安全的应在程序入口一次性开启with mlflow.start_run()与exclusive的关系exclusiveFalse默认时自动记录内容会写入用户创建的活跃 Run若想将自动记录与手动记录隔离可设置exclusiveTrue。小结通过 logistic_regression.py 与 one_vs_rest.py 两个示例可以快速掌握 MLflow PySpark ML autologging 的完整用法一行mlflow.pyspark.ml.autolog()即可自动捕获训练参数、估计器标签、训练后评估指标、数据集血缘与白名单内的模型工件OneVsRest示例则展示了元估计器的层次记录与子估计器去重机制。结合log_model_allowlist、registered_model_name、log_input_examples等参数开发者可以在几乎零侵入的前提下为 Spark ML 训练流程建立起完整的实验追踪与模型管理闭环。【免费下载链接】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），仅供参考