[Python] Cannot run PyVista/VTK inside a Huggingface multiprocessing map()

The following code crashes, with a forking error.

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}


# breaks if num_proc>1
ds = datasets.Dataset.from_list(data).map(render_point, num_proc=2)


The documentation on https://huggingface.co/docs/datasets/main/en/process#multiprocessing implies that the map() call is a fork, and it seems pyvista cannot work with forking. (not sure here, but that is what it seems like). However, adding

import multiprocess
multiprocess.set_start_method("spawn")


did not solve things, as it only generated the error TypeError: fork_exec() takes exactly 23 arguments (21 given).

How can I reap benefits of huggingface dataset parallellism with pyvista?

Continue reading...