MLflow Spark UDF 批量推理实战从 sklearn 模型到分布式环境复现【免费下载链接】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本篇技术指南围绕 examples/spark_udf/ 目录下的官方示例展开系统讲解mlflow.pyfunc.spark_udf这一核心 API 的使用方式——它允许你将已记录的 MLflow 模型包装成 Spark UDF直接对 Spark DataFrame 做分布式批量推理。你将掌握四类实战场景基础 sklearn 批量推理、datetime 类型输入、struct/array 复杂类型输入以及基于 Databricks Connect 的预构建模型环境env_managerconda与prebuilt_env_uri的配置与原理并深入源码理解其底层调用链。为什么需要mlflow.pyfunc.spark_udf在真实生产环境中模型训练完成后往往需要以**批量离线推理batch inference**的方式作用于海量数据。Spark 是这类场景的主流计算引擎但直接把训练好的模型塞进 Spark 并不简单模型对象通常驻留在 driver 进程内无法直接分发到各 executor训练时的 Python 依赖版本与推理环境可能存在差异导致结果不一致甚至报错输入特征可能是 datetime、struct、array 等复杂类型需要正确地在 PySpark 与 pandas 之间转换。mlflow.pyfunc.spark_udf正是为消除这些障碍而生它以 mlflow.pyfunc flavor 为统一接口将任意 MLflow 模型sklearn、PyTorch、LangChain 等包装成pandas_udf并可通过env_manager参数在 UDF 执行时重建训练时的 Python 环境保证训练什么环境推理什么环境。环境准备与先决条件官方示例基于 scikit-learn 构建模型运行前需要先安装依赖pip install scikit-learn此外运行示例还需要一个可用的 Spark 环境pyspark已安装本地模式即可如local[*]已安装 mlflow示例代码从源码目录或通过pip install mlflow获得。from pyspark.sql import SparkSession from sklearn import datasets from sklearn.neighbors import KNeighborsClassifier import mlflow from mlflow.models import infer_signature基础示例sklearn 模型 conda 环境复现spark_udf.py 是理解整个 API 的最小完整示例其完整流程如下with SparkSession.builder.getOrCreate() as spark: X, y datasets.load_iris(as_frameTrue, return_X_yTrue) model KNeighborsClassifier() model.fit(X, y) predictions model.predict(X) signature infer_signature(X, predictions) with mlflow.start_run(): model_info mlflow.sklearn.log_model(model, namemodel, signaturesignature) infer_spark_df spark.createDataFrame(X) pyfunc_udf mlflow.pyfunc.spark_udf(spark, model_info.model_uri, env_managerconda) result infer_spark_df.select(pyfunc_udf(*X.columns).alias(predictions)).toPandas() print(result)步骤拆解训练并记录模型用mlflow.sklearn.log_model记录模型同时通过mlflow.models.infer_signature推断并保存输入/输出 schema。签名signature非常重要——后续spark_udf在没有显式指定result_type时会依据模型输出 schema 自动推断 UDF 返回类型。构造推理 DataFramespark.createDataFrame(X)把训练特征转换为 Spark DataFrame这正是 UDF 的输入。创建 UDFmlflow.pyfunc.spark_udf(spark, model_info.model_uri, env_managerconda)。env_managerconda告诉 MLflow 在 UDF 执行时用模型记录时保存的 conda 环境包含训练时的精确依赖版本重建推理环境。执行批量推理infer_spark_df.select(pyfunc_udf(*X.columns).alias(predictions))对整张表逐行/逐批次执行推理最终toPandas()把结果拉回本地打印。为什么是env_managerconda从 mlflow/pyfunc/init.py 的源码spark_udf定义位于 L2042 起可以看到env_manager支持四种取值取值行为conda使用 Conda 重建训练时的软件环境推荐用于保证推理环境与训练一致virtualenv使用 virtualenv 重建训练环境默认选项之一源码 L2168-L2170uv使用 uv 重建训练环境L2171-L2172local直接使用当前 Python 环境推理可能与训练环境不同可能引发错误或无效预测L2175-L2177源码中针对local模式会打印明确警告L2352-L2362Callingspark_udf()withenv_managerlocaldoes not recreate the same environment that was used during training, which may lead to errors or inaccurate predictions. We recommend specifyingenv_managerconda...同时在非local模式下driver 端会调用pyfunc_backend.prepare_env(...)预创建环境L2440-L2443因此首次执行 UDF 时重建环境会消耗额外时间源码在 L2364-L2367 有相应日志提示。这个先准备、后分发的设计使得 Spark 各个 executor 上的 UDF 任务都能共享已还原的环境实现环境复现的分布式一致性。进阶场景一datetime 类型输入spark_udf_datetime.py 演示了当模型输入包含 datetime 类型列时如何处理其核心是一个 sklearn Pipelinedef extract_month(df): transformed df.assign(monthdf[timestamp].dt.month) return transformed month_extractor FunctionTransformer(extract_month, validateFalse) timestamp_remover ColumnTransformer( [(selector, passthrough, X.columns.drop(timestamp))], remainderdrop ) model Pipeline([ (month_extractor, month_extractor), (timestamp_remover, timestamp_remover), (knn, KNeighborsClassifier()), ]) model.fit(X, y) with mlflow.start_run(): model_info mlflow.sklearn.log_model( model, namemodel, signaturesignature, serialization_formatcloudpickle )该示例的关键点训练数据在 iris 特征之外追加了一列随机生成的timestampdatetime对象Pipeline 内先用FunctionTransformer提取月份再用ColumnTransformer去掉时间戳列保证送入 KNN 的仍是数值特征因为 Pipeline 包含自定义 Python 函数记录模型时显式指定serialization_formatcloudpickle推理侧同样用spark.createDataFrame构造含 datetime 列的 DataFrame并通过infer_spark_df._jdf.schema().treeString()打印 Spark 侧 schema直观观察类型映射。注意根据 mlflow/pyfunc/init.py 中spark_udf的文档注释L2076-L2078pyspark.sql.types.DateType类型的输入在 Spark 2.4 及更早版本上不受支持。使用 datetime 类型输入时请确保 Spark 版本满足要求。进阶场景二struct 与 array 复杂类型输入structs_and_arrays.py 演示了 Spark 中最常见的复杂类型——字符串、数组ArrayType、结构体StructType、结构体数组——如何作为 UDF 输入并显式指定返回类型。示例定义了一个自定义PythonModelclass MyModel(mlflow.pyfunc.PythonModel): def predict(self, context, model_input): return [str( | .join(map(str, row))) for _, row in model_input.iterrows()] df spark.createDataFrame( [(a, [0], {bool: True}, [{double: 0.1}])], schemaT.StructType([ T.StructField(str, T.StringType()), T.StructField(arr, T.ArrayType(T.IntegerType())), T.StructField(obj, T.StructType([T.StructField(bool, T.BooleanType())])), T.StructField(obj_arr, T.ArrayType(T.StructType([T.StructField(double, T.DoubleType())]))), ]), ) with mlflow.start_run(): model_info mlflow.pyfunc.log_model( namemodel, python_modelMyModel(), signaturemlflow.models.infer_signature(df), ) udf mlflow.pyfunc.spark_udf(sparkspark, model_urimodel_info.model_uri, result_typestring) df.withColumn(output, udf(str, arr, obj, obj_arr)).show()这个示例展示了两点核心机制任意PythonModel都可被包装成 UDF不需要是 sklearn 等特定 flavor只要实现了predict(context, model_input)即可result_type显式指定返回类型这里使用stringUDF 会将预测结果统一转换为字符串result_type支持的值包括int/long/float/double/string/boolean等原始类型以及ArrayType、struct...等复合类型源码 L2131-L2160 给出了完整清单与类型转换规则。从源码_predict_row_batchL2484-L2586可以看到内部处理逻辑当 UDF 收到多个列参数时会按模型的输入 schema 命名拼装 pandas DataFrame 再调用predict_fn(pdf, params)若模型输出为 DataFrame 且指定了result_type则会按目标类型过滤与转换列如数值列转 float32/float64、整型列转 int32/int64 等。进阶场景三Databricks Connect 与预构建模型环境prebuilt_env_urispark_udf_with_prebuilt_env.py 面向一个特殊场景在本地非 Databricks runtime使用 Databricks Connect 远程连接 Databricks 集群执行推理。此时每次推理都重建环境代价高昂MLflow 提供了预构建环境机制。from databricks.connect import DatabricksSession from databricks.sdk import WorkspaceClient ws WorkspaceClient() spark DatabricksSession.builder.remote( hostos.environ[DATABRICKS_HOST], tokenos.environ[DATABRICKS_TOKEN], cluster_idcluster id, ).getOrCreate() X, y datasets.load_iris(as_frameTrue, return_X_yTrue) model KNeighborsClassifier() model.fit(X, y) with mlflow.start_run(): model_info mlflow.sklearn.log_model(model, namemodel) model_uri model_info.model_uri # 在 Databricks runtime 中预构建模型环境 # model_env_uc_path mlflow.pyfunc.build_model_env(model_uri, /Volumes/...) model_env_uc_path dbfs:/Volumes/... infer_spark_df spark.createDataFrame(X) pyfunc_udf mlflow.pyfunc.spark_udf(spark, model_uri, prebuilt_env_urimodel_env_uc_path) result infer_spark_df.select(pyfunc_udf(*X.columns).alias(predictions)).toPandas() print(result)预构建环境的使用要点构建在 Databricks 集群中调用mlflow.pyfunc.build_model_env(model_uri, save_path)它会生成一个包含模型 Python 环境快照的tar.gz归档文件源码 mlflow/pyfunc/init.py L1949 起。save_path可以是本地目录、挂载的 DBFS 路径/dbfs/...或 UC Volume 挂载路径/Volumes/...默认env_managervirtualenv也支持uv。使用在spark_udf中传入prebuilt_env_uriMLflow 会直接使用预构建环境而跳过重建显著缩短 UDF 启动时间。路径可以是本地文件或dbfs:/...形式的 DBFS 路径。约束prebuilt_env_uri与env_manager互斥——设置前者时不得再设置后者源码 L2179-L2180、L2233-L2239该参数只能在 Databricks Serverless notebook REPL、Databricks Shared cluster notebook REPL 以及 Databricks Connect 客户端环境中使用源码 L2252-L2258使用 Databricks Connect 从自有机器连接时prebuilt_env_uri是必需参数源码 L2260-L2264否则会直接抛出RuntimeError远程 Databricks Connect 要求集群 runtime 版本 ≥ 15.4Databricks Serverless 要求 MLflow ≥ 2.18.0若 runtime 恰为 15.4 且为标准访问模式集群还需将集群配置spark.databricks.safespark.archive.artifact.unpack.disabled设为false源码 L2281-L2313build_model_env仅支持在 Databricks runtime 中执行且环境归档不可跨 Databricks runtime 版本或不同平台复用——切换集群或 runtime 后需要重新构建源码 L1960-L1975 的说明。其他适用注意事项Databricks Serverless 限制Serverless 的 UDF 沙箱总计容量约 1GB内存 本地磁盘且无 GPU。因此包含大权重或依赖 GPU 的深度学习模型不适合部署在 Databricks Serverless源码 L2090-L2096。Spark Connect 限制在连接非 Databricks Spark 集群的 Spark Connect 客户端环境中virtualenv、conda、uv等环境管理器不受支持只能使用local源码 L2328-L2336。参数透传params可向模型推理传入额外参数extra_env可向 UDF executor 注入环境变量例如通过MLFLOW_SCORING_SERVER_REQUEST_TIMEOUT覆盖评分服务超时详见源码 L2182-L2186 的文档注释。测试与验证依据仓库对spark_udf的覆盖相当完整相关测试文件包括tests/pyfunc/test_spark.pySpark 模式下spark_udf的核心行为测试类型转换、环境管理、参数传递等tests/pyfunc/test_spark_connect.pySpark Connect 模式下的适配测试tests/examples/test_examples.py对examples/目录下示例脚本的端到端冒烟测试。如果你希望验证本文示例在你的环境中可运行可以依次执行python spark_udf.py python spark_udf_datetime.py python structs_and_arrays.py第四个示例 spark_udf_with_prebuilt_env.py 需要 Databricks Connect 与远程集群配置需按DATABRICKS_HOST、DATABRICKS_TOKEN环境变量及集群 ID 准备环境。总结mlflow.pyfunc.spark_udf是连接MLflow 模型管理与Spark 分布式计算的关键桥梁。通过本文四个官方示例你可以掌握基础批量推理spark_udf(spark, model_uri, env_managerconda)一行完成环境复现与分布式推理复杂输入类型datetime、struct、array 均可正确映射到模型输入 schema生产级优化prebuilt_env_uribuild_model_env在 Databricks Connect 场景下跳过环境重建显著加速 UDF 启动底层原理从 mlflow/pyfunc/init.py 源码可确认其内部基于pandas_udf实现、按模型签名推断/校验返回类型、通过prepare_env在 driver 端预建环境并分发给 executors 的完整调用链。无论你运行的是本地 Spark 集群、自建 Spark Connect 还是 Databricks 托管环境上述 API 与示例都提供了从训练记录到批量推理的最小可复现路径。【免费下载链接】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），仅供参考 MLflow Spark UDF 批量推理实战：从 sklearn 模型到分布式环境复现
MLflow Spark UDF 批量推理实战从 sklearn 模型到分布式环境复现【免费下载链接】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本篇技术指南围绕 examples/spark_udf/ 目录下的官方示例展开系统讲解mlflow.pyfunc.spark_udf这一核心 API 的使用方式——它允许你将已记录的 MLflow 模型包装成 Spark UDF直接对 Spark DataFrame 做分布式批量推理。你将掌握四类实战场景基础 sklearn 批量推理、datetime 类型输入、struct/array 复杂类型输入以及基于 Databricks Connect 的预构建模型环境env_managerconda与prebuilt_env_uri的配置与原理并深入源码理解其底层调用链。为什么需要mlflow.pyfunc.spark_udf在真实生产环境中模型训练完成后往往需要以**批量离线推理batch inference**的方式作用于海量数据。Spark 是这类场景的主流计算引擎但直接把训练好的模型塞进 Spark 并不简单模型对象通常驻留在 driver 进程内无法直接分发到各 executor训练时的 Python 依赖版本与推理环境可能存在差异导致结果不一致甚至报错输入特征可能是 datetime、struct、array 等复杂类型需要正确地在 PySpark 与 pandas 之间转换。mlflow.pyfunc.spark_udf正是为消除这些障碍而生它以 mlflow.pyfunc flavor 为统一接口将任意 MLflow 模型sklearn、PyTorch、LangChain 等包装成pandas_udf并可通过env_manager参数在 UDF 执行时重建训练时的 Python 环境保证训练什么环境推理什么环境。环境准备与先决条件官方示例基于 scikit-learn 构建模型运行前需要先安装依赖pip install scikit-learn此外运行示例还需要一个可用的 Spark 环境pyspark已安装本地模式即可如local[*]已安装 mlflow示例代码从源码目录或通过pip install mlflow获得。from pyspark.sql import SparkSession from sklearn import datasets from sklearn.neighbors import KNeighborsClassifier import mlflow from mlflow.models import infer_signature基础示例sklearn 模型 conda 环境复现spark_udf.py 是理解整个 API 的最小完整示例其完整流程如下with SparkSession.builder.getOrCreate() as spark: X, y datasets.load_iris(as_frameTrue, return_X_yTrue) model KNeighborsClassifier() model.fit(X, y) predictions model.predict(X) signature infer_signature(X, predictions) with mlflow.start_run(): model_info mlflow.sklearn.log_model(model, namemodel, signaturesignature) infer_spark_df spark.createDataFrame(X) pyfunc_udf mlflow.pyfunc.spark_udf(spark, model_info.model_uri, env_managerconda) result infer_spark_df.select(pyfunc_udf(*X.columns).alias(predictions)).toPandas() print(result)步骤拆解训练并记录模型用mlflow.sklearn.log_model记录模型同时通过mlflow.models.infer_signature推断并保存输入/输出 schema。签名signature非常重要——后续spark_udf在没有显式指定result_type时会依据模型输出 schema 自动推断 UDF 返回类型。构造推理 DataFramespark.createDataFrame(X)把训练特征转换为 Spark DataFrame这正是 UDF 的输入。创建 UDFmlflow.pyfunc.spark_udf(spark, model_info.model_uri, env_managerconda)。env_managerconda告诉 MLflow 在 UDF 执行时用模型记录时保存的 conda 环境包含训练时的精确依赖版本重建推理环境。执行批量推理infer_spark_df.select(pyfunc_udf(*X.columns).alias(predictions))对整张表逐行/逐批次执行推理最终toPandas()把结果拉回本地打印。为什么是env_managerconda从 mlflow/pyfunc/init.py 的源码spark_udf定义位于 L2042 起可以看到env_manager支持四种取值取值行为conda使用 Conda 重建训练时的软件环境推荐用于保证推理环境与训练一致virtualenv使用 virtualenv 重建训练环境默认选项之一源码 L2168-L2170uv使用 uv 重建训练环境L2171-L2172local直接使用当前 Python 环境推理可能与训练环境不同可能引发错误或无效预测L2175-L2177源码中针对local模式会打印明确警告L2352-L2362Callingspark_udf()withenv_managerlocaldoes not recreate the same environment that was used during training, which may lead to errors or inaccurate predictions. We recommend specifyingenv_managerconda...同时在非local模式下driver 端会调用pyfunc_backend.prepare_env(...)预创建环境L2440-L2443因此首次执行 UDF 时重建环境会消耗额外时间源码在 L2364-L2367 有相应日志提示。这个先准备、后分发的设计使得 Spark 各个 executor 上的 UDF 任务都能共享已还原的环境实现环境复现的分布式一致性。进阶场景一datetime 类型输入spark_udf_datetime.py 演示了当模型输入包含 datetime 类型列时如何处理其核心是一个 sklearn Pipelinedef extract_month(df): transformed df.assign(monthdf[timestamp].dt.month) return transformed month_extractor FunctionTransformer(extract_month, validateFalse) timestamp_remover ColumnTransformer( [(selector, passthrough, X.columns.drop(timestamp))], remainderdrop ) model Pipeline([ (month_extractor, month_extractor), (timestamp_remover, timestamp_remover), (knn, KNeighborsClassifier()), ]) model.fit(X, y) with mlflow.start_run(): model_info mlflow.sklearn.log_model( model, namemodel, signaturesignature, serialization_formatcloudpickle )该示例的关键点训练数据在 iris 特征之外追加了一列随机生成的timestampdatetime对象Pipeline 内先用FunctionTransformer提取月份再用ColumnTransformer去掉时间戳列保证送入 KNN 的仍是数值特征因为 Pipeline 包含自定义 Python 函数记录模型时显式指定serialization_formatcloudpickle推理侧同样用spark.createDataFrame构造含 datetime 列的 DataFrame并通过infer_spark_df._jdf.schema().treeString()打印 Spark 侧 schema直观观察类型映射。注意根据 mlflow/pyfunc/init.py 中spark_udf的文档注释L2076-L2078pyspark.sql.types.DateType类型的输入在 Spark 2.4 及更早版本上不受支持。使用 datetime 类型输入时请确保 Spark 版本满足要求。进阶场景二struct 与 array 复杂类型输入structs_and_arrays.py 演示了 Spark 中最常见的复杂类型——字符串、数组ArrayType、结构体StructType、结构体数组——如何作为 UDF 输入并显式指定返回类型。示例定义了一个自定义PythonModelclass MyModel(mlflow.pyfunc.PythonModel): def predict(self, context, model_input): return [str( | .join(map(str, row))) for _, row in model_input.iterrows()] df spark.createDataFrame( [(a, [0], {bool: True}, [{double: 0.1}])], schemaT.StructType([ T.StructField(str, T.StringType()), T.StructField(arr, T.ArrayType(T.IntegerType())), T.StructField(obj, T.StructType([T.StructField(bool, T.BooleanType())])), T.StructField(obj_arr, T.ArrayType(T.StructType([T.StructField(double, T.DoubleType())]))), ]), ) with mlflow.start_run(): model_info mlflow.pyfunc.log_model( namemodel, python_modelMyModel(), signaturemlflow.models.infer_signature(df), ) udf mlflow.pyfunc.spark_udf(sparkspark, model_urimodel_info.model_uri, result_typestring) df.withColumn(output, udf(str, arr, obj, obj_arr)).show()这个示例展示了两点核心机制任意PythonModel都可被包装成 UDF不需要是 sklearn 等特定 flavor只要实现了predict(context, model_input)即可result_type显式指定返回类型这里使用stringUDF 会将预测结果统一转换为字符串result_type支持的值包括int/long/float/double/string/boolean等原始类型以及ArrayType、struct...等复合类型源码 L2131-L2160 给出了完整清单与类型转换规则。从源码_predict_row_batchL2484-L2586可以看到内部处理逻辑当 UDF 收到多个列参数时会按模型的输入 schema 命名拼装 pandas DataFrame 再调用predict_fn(pdf, params)若模型输出为 DataFrame 且指定了result_type则会按目标类型过滤与转换列如数值列转 float32/float64、整型列转 int32/int64 等。进阶场景三Databricks Connect 与预构建模型环境prebuilt_env_urispark_udf_with_prebuilt_env.py 面向一个特殊场景在本地非 Databricks runtime使用 Databricks Connect 远程连接 Databricks 集群执行推理。此时每次推理都重建环境代价高昂MLflow 提供了预构建环境机制。from databricks.connect import DatabricksSession from databricks.sdk import WorkspaceClient ws WorkspaceClient() spark DatabricksSession.builder.remote( hostos.environ[DATABRICKS_HOST], tokenos.environ[DATABRICKS_TOKEN], cluster_idcluster id, ).getOrCreate() X, y datasets.load_iris(as_frameTrue, return_X_yTrue) model KNeighborsClassifier() model.fit(X, y) with mlflow.start_run(): model_info mlflow.sklearn.log_model(model, namemodel) model_uri model_info.model_uri # 在 Databricks runtime 中预构建模型环境 # model_env_uc_path mlflow.pyfunc.build_model_env(model_uri, /Volumes/...) model_env_uc_path dbfs:/Volumes/... infer_spark_df spark.createDataFrame(X) pyfunc_udf mlflow.pyfunc.spark_udf(spark, model_uri, prebuilt_env_urimodel_env_uc_path) result infer_spark_df.select(pyfunc_udf(*X.columns).alias(predictions)).toPandas() print(result)预构建环境的使用要点构建在 Databricks 集群中调用mlflow.pyfunc.build_model_env(model_uri, save_path)它会生成一个包含模型 Python 环境快照的tar.gz归档文件源码 mlflow/pyfunc/init.py L1949 起。save_path可以是本地目录、挂载的 DBFS 路径/dbfs/...或 UC Volume 挂载路径/Volumes/...默认env_managervirtualenv也支持uv。使用在spark_udf中传入prebuilt_env_uriMLflow 会直接使用预构建环境而跳过重建显著缩短 UDF 启动时间。路径可以是本地文件或dbfs:/...形式的 DBFS 路径。约束prebuilt_env_uri与env_manager互斥——设置前者时不得再设置后者源码 L2179-L2180、L2233-L2239该参数只能在 Databricks Serverless notebook REPL、Databricks Shared cluster notebook REPL 以及 Databricks Connect 客户端环境中使用源码 L2252-L2258使用 Databricks Connect 从自有机器连接时prebuilt_env_uri是必需参数源码 L2260-L2264否则会直接抛出RuntimeError远程 Databricks Connect 要求集群 runtime 版本 ≥ 15.4Databricks Serverless 要求 MLflow ≥ 2.18.0若 runtime 恰为 15.4 且为标准访问模式集群还需将集群配置spark.databricks.safespark.archive.artifact.unpack.disabled设为false源码 L2281-L2313build_model_env仅支持在 Databricks runtime 中执行且环境归档不可跨 Databricks runtime 版本或不同平台复用——切换集群或 runtime 后需要重新构建源码 L1960-L1975 的说明。其他适用注意事项Databricks Serverless 限制Serverless 的 UDF 沙箱总计容量约 1GB内存 本地磁盘且无 GPU。因此包含大权重或依赖 GPU 的深度学习模型不适合部署在 Databricks Serverless源码 L2090-L2096。Spark Connect 限制在连接非 Databricks Spark 集群的 Spark Connect 客户端环境中virtualenv、conda、uv等环境管理器不受支持只能使用local源码 L2328-L2336。参数透传params可向模型推理传入额外参数extra_env可向 UDF executor 注入环境变量例如通过MLFLOW_SCORING_SERVER_REQUEST_TIMEOUT覆盖评分服务超时详见源码 L2182-L2186 的文档注释。测试与验证依据仓库对spark_udf的覆盖相当完整相关测试文件包括tests/pyfunc/test_spark.pySpark 模式下spark_udf的核心行为测试类型转换、环境管理、参数传递等tests/pyfunc/test_spark_connect.pySpark Connect 模式下的适配测试tests/examples/test_examples.py对examples/目录下示例脚本的端到端冒烟测试。如果你希望验证本文示例在你的环境中可运行可以依次执行python spark_udf.py python spark_udf_datetime.py python structs_and_arrays.py第四个示例 spark_udf_with_prebuilt_env.py 需要 Databricks Connect 与远程集群配置需按DATABRICKS_HOST、DATABRICKS_TOKEN环境变量及集群 ID 准备环境。总结mlflow.pyfunc.spark_udf是连接MLflow 模型管理与Spark 分布式计算的关键桥梁。通过本文四个官方示例你可以掌握基础批量推理spark_udf(spark, model_uri, env_managerconda)一行完成环境复现与分布式推理复杂输入类型datetime、struct、array 均可正确映射到模型输入 schema生产级优化prebuilt_env_uribuild_model_env在 Databricks Connect 场景下跳过环境重建显著加速 UDF 启动底层原理从 mlflow/pyfunc/init.py 源码可确认其内部基于pandas_udf实现、按模型签名推断/校验返回类型、通过prepare_env在 driver 端预建环境并分发给 executors 的完整调用链。无论你运行的是本地 Spark 集群、自建 Spark Connect 还是 Databricks 托管环境上述 API 与示例都提供了从训练记录到批量推理的最小可复现路径。【免费下载链接】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），仅供参考