4. Databricks运行时与Spark配置调优

机器学习任务在 Databricks 上跑得慢、频繁 OOM 或成本偏高，很多时候并不是模型本身的问题，而是 Spark 运行时配置与 ML 负载特征不匹配。ML 任务与典型 ETL 的差异在于：它往往要反复迭代、缓存特征表、在 driver 上收集小结果、用 Python UDF 或 pandas UDF 做特征工程，并且 shuffle 的数据形态常常是"宽而稀疏"的向量或键值对。理解内存管理、序列化、Shuffle 分区和动态资源分配这四件事如何作用于 ML 任务，是调优的起点。

内存管理：统一内存模型下的执行与存储竞争

Spark 在 executor 内使用统一内存管理（Unified Memory Management）。堆内内存被划分为预留内存（默认 300MB）、用户内存（约 40%）和 Spark 内存（约 60%）。Spark 内存又分为执行内存（execution）和存储内存（storage），二者共享一个池，可互相借用：当没有缓存块时，执行内存可占用全部；当有缓存块时，执行内存可强制驱逐部分存储内存，但存储内存不能反过来抢占执行内存。

对 ML 任务而言，这个机制的直接后果是：如果你用 cache() 缓存了特征表，同时又在一个 stage 里做大规模聚合或 join，执行内存会驱逐缓存块，导致后续迭代重新计算。典型症状是"缓存了却还是慢"。判断依据可以看 Spark UI 的 Storage 页签：如果 Cached Partitions 的 Fraction Cached 长期低于 100%，说明缓存被反复驱逐。

调优方向有三条。第一，对需要反复读取的特征表，用 MEMORY_AND_DISK 或 MEMORY_AND_DISK_DESER 存储级别，避免驱逐后完全重算。第二，控制单分区大小，让每个 task 处理的数据量落在 100–200MB 区间，避免单个 task 撑爆执行内存。第三，对 driver 侧收集结果的操作（如 collect()、toPandas()）保持警惕，它们消耗的是 driver 内存，与 executor 内存是两套账。

# 示例：为迭代式特征表设置合适的存储级别
feature_df = spark.table("ml.features_v3").repartition(400, "user_id")
feature_df.persist(StorageLevel.MEMORY_AND_DISK)

# 查看缓存实际命中情况
spark.sql("SELECT * FROM ml.features_v3").cache().count()

序列化：Kryo 与 ML 向量对象的开销

Spark 默认使用 Java 序列化，对 ML 场景中的稠密/稀疏向量、嵌套结构体、Python 对象包装开销很大。切换到 Kryo 序列化通常能带来 2–10 倍的序列化体积下降和明显的 CPU 节省，代价是需要注册自定义类。

在 Databricks 上设置方式如下：

spark.conf.set("spark.serializer", "org.apache.spark.serializer.KryoSerializer")
spark.conf.set("spark.kryo.registrationRequired", "false")
spark.conf.set("spark.kryoserializer.buffer.max", "256m")

需要注意两点。其一，spark.kryoserializer.buffer.max 默认 64MB，当单个记录（例如一个高维稀疏向量）超过该值时，Kryo 会抛 BufferOverflowException，ML 场景中把上限提到 256MB 是常见做法。其二，如果开启 registrationRequired=true，未注册的类会直接报错，适合在稳定生产作业中强制暴露序列化问题；开发阶段可先设为 false。

另一个容易被忽略的点是 Python UDF。标准 Python UDF 需要在 JVM 与 Python 进程之间序列化每一行数据，开销极高。对特征工程，应优先使用 pandas UDF（@pandas_udf）或 Spark 内置函数，它们以 Arrow 批量传输，序列化成本大幅降低。

Shuffle 分区：并行度与数据倾斜的平衡

Shuffle 是 ML 任务中最昂贵的操作之一，常见于 join、groupBy、repartition 和特征聚合。分区数由 spark.sql.shuffle.partitions 控制，默认 200。这个默认值对中小数据集偏大（大量空 task 调度开销），对大数据集又偏小（单分区过大导致 spill）。

一个实用的估算方法是：目标分区数 ≈ shuffle 输出总字节数 ÷ 目标单分区大小（128–200MB）。例如 shuffle 写出 80GB，则分区数约 400–640。Databricks 的 Adaptive Query Execution（AQE）可以在运行时自动合并小分区、拆分倾斜分区，建议保持开启：

spark.conf.set("spark.sql.adaptive.enabled", "true")
spark.conf.set("spark.sql.adaptive.coalescePartitions.enabled", "true")
spark.conf.set("spark.sql.adaptive.skewJoin.enabled", "true")
spark.conf.set("spark.sql.shuffle.partitions", "400")

对 ML 任务，还要特别关注 key 的分布。以 user_id 为 key 做特征 join 时，如果少数活跃用户占据大量记录，就会出现长尾 task。AQE 的 skew join 能缓解，但更根本的做法是加盐（salting）或对热门 key 单独处理。判断依据看 Spark UI 的 Stage 页签：如果个别 task 的 Shuffle Read Size 是其他 task 的数十倍，且 Duration 明显偏长，就是倾斜信号。

动态资源分配：让集群随负载伸缩

动态资源分配（Dynamic Allocation）允许 executor 数量根据待处理 task 队列自动增减。对 ML 任务尤其有价值，因为特征工程阶段需要大量 executor，而模型训练或 driver 侧计算阶段几乎不需要。开启方式：

spark.conf.set("spark.dynamicAllocation.enabled", "true")
spark.conf.set("spark.dynamicAllocation.minExecutors", "2")
spark.conf.set("spark.dynamicAllocation.maxExecutors", "32")
spark.conf.set("spark.dynamicAllocation.initialExecutors", "4")
spark.conf.set("spark.dynamicAllocation.executorIdleTimeout", "60s")

关键参数是 executorIdleTimeout：默认 60s，对 ML 任务中"计算—等待—再计算"的间歇模式，可以适当调大（如 120–180s），避免 executor 刚释放又要重新申请，造成反复的容器启动开销。反之，如果作业是短时突发型，调小可以更快释放资源、降低成本。

需要提醒的是，动态资源分配与 shuffle 有耦合：executor 被移除后，其 shuffle 输出需要保留在外部 shuffle 服务上，否则会触发重算。Databricks 的 Photon 与 shuffle 服务已处理这一层，但在自建 Spark 集群上必须确认 spark.shuffle.service.enabled=true。

四类配置的协同与取舍

这四项配置并非独立。序列化方式影响 shuffle 写出的字节数，进而影响分区数估算；分区数影响单 task 内存压力，进而影响缓存驱逐频率；动态资源分配决定 executor 数量，进而影响并行度上限。一个常见的调优顺序是：先确认序列化（Kryo + 合理 buffer），再根据 shuffle 字节数设定分区数并开启 AQE，然后调整存储级别保护缓存，最后用动态资源分配控制成本。

下面给出一份可直接参考的配置清单，适用于中等规模（数十 GB 级特征表）的 ML 特征工程作业：