PySpark调用.collect()/.count()时触发Socket超时错误求助

根因分析 #

• UDF与PI系统交互的网络波动：自定义UDF（get_average、get_tagfreq_sync等）直接调用AVEVA PI接口时，分布式环境下Python Worker与PI服务器的网络连接不稳定，单个请求超时会触发socket.timeout。
• Python Worker资源瓶颈：Worker进程内存/CPU不足，导致UDF执行卡顿，间接引发Spark Driver与Worker间的心跳或数据传输超时。
• PI系统并发限流：Spark分布式执行时，大量Worker同时请求PI系统，触发PI侧的并发阈值限制，导致请求超时。

针对性解决方案 #

1. 给UDF添加重试与超时控制 #

在UDF的PI接口调用逻辑中增加本地重试，避免单次网络波动导致失败；同时限制UDF执行时长，防止Worker进程挂起：

import socket
from tenacity import retry, stop_after_attempt, wait_exponential, retry_if_exception_type
import timeout_decorator

# 带重试的UDF示例
@retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=2, max=10), retry=retry_if_exception_type(socket.timeout))
def get_average(pi_tag, start_time, end_time):
    # 原PI接口调用逻辑
    pass

# 带超时控制的UDF示例
@timeout_decorator.timeout(30)  # 30秒超时
def get_tag_type(pi_tag):
    # 原PI接口调用逻辑
    pass

2. 调整Spark Worker相关配置 #

修改Spark会话配置，强化Worker稳定性并控制并发量：

spark = SparkSession.builder \
    .appName("PI-Data-Processing") \
    .config("spark.executor.python.worker.memory", "2g")  # 增大Python Worker内存
    .config("spark.python.worker.reuse", "true")  # 复用Worker进程，减少初始化开销
    .config("spark.network.timeout", "300s")  # 延长整体网络超时
    .config("spark.executor.heartbeatInterval", "60s")  # 调大心跳间隔
    .config("spark.executor.instances", "4")  # 控制Executor数量，避免PI系统过载
    .getOrCreate()

3. 批量处理PI接口请求 #

避免UDF中单条调用PI接口，改为批量查询后广播映射：

# 收集所有待查询的PI标签
pi_tags = df.select("pi_tag").distinct().collect()
tag_list = [row["pi_tag"] for row in pi_tags]

# 批量调用PI接口获取数据（需实现批量查询函数）
tag_type_map = batch_get_tag_types(tag_list)

# 广播映射到所有Executor
tag_type_broadcast = spark.sparkContext.broadcast(tag_type_map)

# 用广播变量替代UDF查询
df = df.withColumn("tag_type", tag_type_broadcast.value[col("pi_tag")])

4. 排查PI系统侧限制 #

联系AVEVA PI管理员确认：

• PI Web API的并发请求上限
• 单个请求的超时设置
• 是否存在IP白名单或限流策略

测试优化建议 #

• 修改pytest命令，增加日志输出定位超时节点：

  pytest -s -k test_main.py --log-level=INFO
  

• 模拟分布式并发场景测试UDF：

  test_tags = ["TAG001", "TAG002", "TAG003", ...]
  # 用8个切片模拟8个并发Worker
  rdd = spark.sparkContext.parallelize(test_tags, numSlices=8)
  result = rdd.map(get_average).collect()
  

内容的提问来源于stack exchange，提问作者Michael C.