From 23d2f2371e63d764950a061e3b8d87fb08d2c771 Mon Sep 17 00:00:00 2001 From: ethanglaser Date: Tue, 7 Apr 2026 15:16:08 -0700 Subject: [PATCH] Fix prefetch-dataset hang on GPU --- sklbench/runner/implementation.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/sklbench/runner/implementation.py b/sklbench/runner/implementation.py index cac0bba4..04ec6c0a 100644 --- a/sklbench/runner/implementation.py +++ b/sklbench/runner/implementation.py @@ -17,7 +17,7 @@ import argparse import json -from multiprocessing import Pool +from multiprocessing import get_context from typing import Dict, List, Tuple, Union from psutil import cpu_count @@ -102,7 +102,7 @@ def run_benchmarks(args: argparse.Namespace) -> int: logger.debug(f"Unique dataset names to load:\n{list(dataset_cases.keys())}") n_proc = min([16, cpu_count(), n_datasets]) logger.info(f"Prefetching {n_datasets} datasets with {n_proc} processes") - with Pool(n_proc) as pool: + with get_context("spawn").Pool(n_proc) as pool: pool.map(load_data_with_cleanup, dataset_cases.values()) # run bench_cases