Fix prefetch datasets

Author: olegkkruglov

sklbench/runner/implementation.py

@@ -102,15 +102,15 @@ def run_benchmarks(args: argparse.Namespace) -> int:
             data_info = get_bench_case_value(bench_case, "data")
             if "dataset_sequence" in data_info:
                 for dataset_info in data_info["dataset_sequence"]:
-                    dataset_cases[get_dataset_name(dataset_info)] = dataset_info
+                    dataset_cases[get_dataset_name(dataset_info)] = (bench_case, dataset_info)
             else:
-                dataset_cases[get_dataset_name(data_info)] = data_info
+                dataset_cases[get_dataset_name(data_info)] = (bench_case, data_info)
 
         logger.debug(f"Unique dataset names to load:\n{list(dataset_cases.keys())}")
         n_proc = min([16, cpu_count(), len(dataset_cases)])
         logger.info(f"Prefetching datasets with {n_proc} processes")
         with Pool(n_proc) as pool:
-            pool.map(load_dataset, dataset_cases.values())
+            pool.starmap(load_dataset, dataset_cases.values())
 
     # run bench_cases
     return_code, result = call_benchmarks(