如何在Spark/Databricks集群全局设置Python库常量？

根因说明 #

Spark采用分布式进程架构，driver节点与各worker节点的Python运行环境完全隔离，在笔记本顶层（driver侧执行的代码块）修改的Python库常量、全局变量，不会自动同步到worker进程。PIL.Image.MAX_IMAGE_PIXELS是进程级配置，每个Python进程加载PIL模块时会读取库内置的默认值，因此仅在driver侧修改该参数无法覆盖worker侧的配置，才会出现worker节点仍触发DecompressionBombError的问题。

可行方案 #

方案1：在pandas UDF内部完成配置（最推荐，无集群运维成本） #

所有UDF内的代码会在worker进程中执行，只要把PIL的参数配置放在UDF函数内部、且在导入fastai等依赖PIL的库之前执行，就能保证每个worker进程运行推理前配置生效。
示例代码：

import pandas as pd
from pyspark.sql.functions import pandas_udf

@pandas_udf("array<float>")
def image_batch_infer(img_path_col: pd.Series) -> pd.Series:
    # 必须在UDF内部导入PIL并设置参数，且顺序早于fastai导入
    from PIL import Image
    Image.MAX_IMAGE_PIXELS = None
    # 再导入fastai、加载模型执行推理
    from fastai.vision.all import load_learner
    learner = load_learner("/dbfs/path/to/your/model.pkl")
    
    pred_results = []
    for img_path in img_path_col:
        pred = learner.predict(img_path)[2].tolist()
        pred_results.append(pred)
    return pd.Series(pred_results)

# 调用UDF执行推理
df = df.withColumn("pred", image_batch_infer("img_path"))

注意：不要把PIL导入、参数设置、fastai导入、模型加载的代码写在UDF函数外部，这类代码仅会在driver节点执行，不会同步到worker。

方案2：配置集群全局init脚本（适合多任务统一配置场景） #

如果集群上有多个任务都需要关闭PIL的大图像检测，可以通过集群init脚本直接修改所有节点PIL库的默认配置，一次配置全集群生效。
操作步骤：

• 新建shell脚本，内容如下，用于替换PIL库源码中MAX_IMAGE_PIXELS的默认值：

#!/bin/bash
# 匹配所有Python环境下的PIL Image.py文件，替换默认配置
find / -name "Image.py" -path "*/site-packages/PIL/*" -exec sed -i 's/^MAX_IMAGE_PIXELS = .*/MAX_IMAGE_PIXELS = None/' {} \;

• 将脚本上传到Databricks工作区，配置为目标集群的全局init脚本
• 重启集群后，所有节点上的Python进程加载PIL时，默认就会关闭DecompressionBomb检测，无需在代码中重复设置。

注意事项 #

只要代码执行顺序上，Image.MAX_IMAGE_PIXELS = None的设置晚于PIL（或fastai等间接导入PIL的库）的导入，配置就会失效，必须保证参数设置在所有PIL相关导入操作之前。

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