Databricks Notebook中ThreadPool/进程池卡顿问题排查求助

我调用的单请求API响应较慢，但支持每秒最多70次并发请求。于是我把数据拆分成多个ID区间[(intervall_1_lowest_id,intervall_1_highest_id),....(intervall_n_lowest_id, intervall_n_highest_id)]，先单线程生成这些区间，再用ThreadPool批量下载数据并写入不同文件（无文件写入冲突）。但运行时总会出现99%任务完成后，剩余1-5个任务卡住的情况，线程看似静默失败，即使获取结果也无反馈。

运行环境 #

Databricks Notebook（11.3 LTS，含Apache Spark 3.3.0、Scala 2.12），Python 3.9.5。

复现代码（内存泄漏导致进程崩溃） #

我在14GB RAM的小集群上用以下代码复现了问题：运行后内存耗尽，第二个进程静默崩溃，但主进程未收到任何异常通知，池长时间无进展。

import multiprocessing
import logging
import time
import requests
import psutil

def execute_single_task(task):
    session =[]
    max_s = 1000000
    print(f"before:{multiprocessing.current_process().pid} {task}")
    while len(session) < max_s:
        # 模拟内存泄漏的示例代码
        session.append(requests.Session())

        if len(session) %1000 == 0:
            print(len(session))
            print(psutil.virtual_memory())

    print(f"after:{multiprocessing.current_process().pid} {task}")
    time.sleep(2)
    return True

def execute_all_tasks(tasks):
    PARALLELIZATION = 2
    print(PARALLELIZATION)
    
    with multiprocessing.Pool(processes=PARALLELIZATION) as pool:
         results = []
         print("with pool")
         for result in pool.imap_unordered(execute_single_task, tasks):
            if result:
                logging.info("Task succeded!")
            else:
                logging.info("Task failed!")
                raise Exception("The entire main process failed.")
        logging.info("All Tasks succeeded!")

if __name__ == "__main__":
    print(f"if name:{__name__}")
    tasks = [1,2]
    execute_all_tasks(tasks)

初始业务代码示例 #

from multiprocessing.pool import ThreadPool
import logging
from dataclasses import dataclass
import copy

@dataclass
class Task:
  id: int
  range: tuple[int,int]

def download_and_save(range):
    resp=requests.post(....)    
    # 怀疑这里是否有dbutils.fs的竞态问题
    dbutils.fs.put(some_path, resp.text)

def download_range(task):
    logging.info("some info here and there")
    download_and_save(task.range)
    return task

def determine_ranges()-> [Task]: # id, range
   # 生成区间逻辑
   return list_of_tasks

def main():
    logging.basicConfig(level=...,format=...)
    tasks = determine_ranges() # Task(id, range)
    remaining_tasks = copy.deepcopy(tasks) 
    with ThreadPool(processes=70) as pool:
        for result in pool.imap_unordered(download_range,tasks):
            print(f"Task with id: {task.id} done")
            if isinstance(result,Task):
               # 从剩余任务中移除
               #....
               logging.info(f"remaining: {remaining_tasks}")
            elif isinstance(result, Exception):
              logging.exception(result)
              raise result
            else: 
              raise ValueError(f"unexpected result: {type(result)}")

核心疑问 #

1. 如何排查这类任务卡住、进程静默崩溃的问题？
2. 为什么子进程崩溃后，主进程收不到任何异常通知？

一、子进程崩溃后主进程无通知的原因 #

Python的multiprocessing.Pool（包括ThreadPool，本质是线程池但底层复用类似通信机制）在子进程因内存耗尽被系统OOM Killer强制杀死时，主进程无法捕获异常：

• 子进程是被操作系统直接终止，而非正常抛出Python异常，因此没有可传递给主进程的异常信息。
• imap_unordered会持续等待子进程返回结果，但被杀死的进程不会返回任何数据，导致主进程阻塞在结果等待环节，表现为任务"卡住"。

二、排查方法 #

1. 监控系统级资源使用 #

• 查看Databricks集群监控面板的节点内存、CPU使用率，确认是否出现OOM（内存耗尽）触发的进程终止。
• 在任务逻辑中加入psutil实时打印内存状态（如复现代码中的psutil.virtual_memory()），定位内存飙升的时间点和触发任务。

2. 捕获子进程退出状态 #

• 替换imap_unordered为apply_async，配合成功/失败回调函数，同时监控子进程的退出码：

  def on_success(result):
      logging.info(f"Task {result.id} completed")
      # 更新剩余任务逻辑
  
  def on_error(error):
      logging.error(f"Task failed with error: {str(error)}")
  
  with ThreadPool(processes=70) as pool:
      for task in tasks:
          pool.apply_async(download_range, args=(task,), callback=on_success, error_callback=on_error)
      pool.close()
      pool.join()
  

• 对于进程池，可通过multiprocessing.active_children()检查存活进程数量，对比任务数判断是否有进程异常退出。

3. 增强日志粒度 #

• 在download_range和download_and_save中添加详细日志，包括任务ID、开始/结束时间、dbutils.fs.put的执行状态，确认是下载阶段还是写入阶段卡住。
• 给每个线程/进程的日志添加唯一标识（如进程ID、任务ID），方便定位具体故障任务。

4. 验证dbutils.fs.put的并发安全性 #

• 即使无文件写入冲突，Databricks的dbutils.fs.put在高并发下可能存在元数据同步延迟问题。可先将数据写入本地临时文件，再用dbutils.fs.cp上传，或改用Spark文件写入API替代。

5. 复现OOM场景 #

• 用你编写的内存泄漏模拟代码复现问题，确认OOM是导致进程静默崩溃的核心诱因。

三、解决方案 #

1. 限制单任务内存占用 #

• 优化下载逻辑，及时释放资源：关闭requests会话、采用流式写入避免内存缓存大响应：

  def download_and_save(range):
      with requests.post(..., stream=True) as resp:
          # 流式写入本地临时文件
          temp_path = f"/tmp/task_{range[0]}.tmp"
          with open(temp_path, "wb") as f:
              for chunk in resp.iter_content(chunk_size=8192):
                  f.write(chunk)
          # 上传到DBFS
          dbutils.fs.cp(f"file:{temp_path}", some_path)
          # 删除本地临时文件
          import os
          os.remove(temp_path)
  

2. 给任务添加超时机制 #

• 使用tenacity库给任务添加超时和重试逻辑，避免单个任务无限阻塞：

  from tenacity import retry, stop_after_attempt, wait_exponential, timeout
  
  @retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=2, max=10), timeout=timeout(60))
  def download_range(task):
      logging.info(f"Starting task {task.id}")
      download_and_save(task.range)
      return task
  

3. 改用Spark并行机制替代Python线程池 #

• 在Databricks环境中，利用Spark的分布式调度能力实现并行下载，比手动管理线程池更稳定：

  from pyspark.sql import SparkSession
  
  spark = SparkSession.builder.getOrCreate()
  # 将任务列表转为RDD，设置分片数匹配并发数
  tasks_rdd = spark.sparkContext.parallelize(tasks, numSlices=70)
  # 并行执行下载任务
  tasks_rdd.foreach(lambda task: download_range(task))
  

4. 监控并处理异常任务 #

• 自定义进度监控逻辑，定期检查任务完成情况，对长时间未完成的任务进行重试或标记为失败。

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