如何在Huggingface数据集多进程map中运行PyVista/VTK？

问题描述 #

运行含PyVista渲染逻辑的Huggingface datasets并行map代码时，触发fork机制兼容性崩溃，报错：

objc[81151]: +[NSResponder initialize] may have been in progress in another thread when fork() was called. We cannot safely call it or ignore it in the fork() child process. Crashing instead. Set a breakpoint on objc_initializeAfterForkError to debug.

涉事代码：

import numpy as np
import pyvista as pv
import datasets

data = [{"coords": np.random.rand(5, 3)} for _ in range(3)]


def render_point(example):
  plotter = pv.Plotter(off_screen=True)
  cloud = pv.PolyData(example["coords"])
  plotter.add_mesh(cloud)
  img = plotter.screenshot(return_img=True)
  return {"image": img}


# 当num_proc>1时崩溃
ds = datasets.Dataset.from_list(data).map(render_point, num_proc=2)

尝试全局设置spawn启动方法：

import multiprocess
multiprocess.set_start_method("spawn")

但触发新错误：

TypeError: fork_exec() takes exactly 23 arguments (21 given)

解决方案 #

方案1：通过map参数指定spawn多进程上下文 #

无需全局修改启动方法，直接在map调用中通过mp_context参数指定spawn上下文，这是datasets官方推荐的安全方式：

import numpy as np
import pyvista as pv
import datasets
import multiprocessing

data = [{"coords": np.random.rand(5, 3)} for _ in range(3)]


def render_point(example):
  plotter = pv.Plotter(off_screen=True)
  cloud = pv.PolyData(example["coords"])
  plotter.add_mesh(cloud)
  img = plotter.screenshot(return_img=True)
  return {"image": img}


ds = datasets.Dataset.from_list(data).map(
    render_point, 
    num_proc=2,
    mp_context=multiprocessing.get_context("spawn")
)

方案2：优化子进程内PyVista资源管理 #

确保PyVista的初始化和资源释放都在子进程内完成，避免主进程残留状态影响：

import numpy as np
import pyvista as pv
import datasets

data = [{"coords": np.random.rand(5, 3)} for _ in range(3)]


def render_point(example):
    # 无GUI环境下启动虚拟显示（可选）
    pv.start_xvfb()
    plotter = pv.Plotter(off_screen=True)
    cloud = pv.PolyData(example["coords"])
    plotter.add_mesh(cloud)
    img = plotter.screenshot(return_img=True)
    # 显式关闭Plotter释放资源
    plotter.close()
    return {"image": img}


ds = datasets.Dataset.from_list(data).map(
    render_point, 
    num_proc=2,
    mp_context="spawn"  # 直接传入字符串"spawn"也可生效
)

关键说明 #

• spawn模式会创建全新的Python进程，而非复用主进程内存，彻底避免PyVista与fork机制的初始化冲突。
• 使用mp_context参数比全局设置set_start_method更安全，不会干扰其他库的多进程逻辑。
• 若在服务器等无GUI环境运行，pv.start_xvfb()是必要的前置操作，确保PyVista能完成离线渲染。

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