From 8bc449e143ae83828755924e1f77ffb7fb428ff2 Mon Sep 17 00:00:00 2001 From: IlyasMoutawwakil Date: Mon, 6 Nov 2023 11:29:22 +0000 Subject: [PATCH] expose multiprocessing start_method --- optimum_benchmark/experiment.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/optimum_benchmark/experiment.py b/optimum_benchmark/experiment.py index 3fcbbdc2..f89a0e75 100644 --- a/optimum_benchmark/experiment.py +++ b/optimum_benchmark/experiment.py @@ -168,10 +168,10 @@ def run(experiment: DictConfig) -> None: raise e -def run_isolated(experiment: DictConfig) -> None: +def run_isolated(experiment: DictConfig, start_method: str = "spawn") -> None: # Set the multiprocessing start method if not already set - if multiprocessing.get_start_method(allow_none=True) is None: - multiprocessing.set_start_method("spawn") + if multiprocessing.get_start_method(allow_none=True) != start_method: + multiprocessing.set_start_method(start_method) # Spawn a new process p = multiprocessing.Process(target=run, args=(experiment,)) @@ -185,4 +185,4 @@ def main(experiment: DictConfig) -> None: LOGGER.warning("Skipping because results already exist in experiment directory.") return - run_isolated(experiment) + run_isolated(experiment, start_method="spawn")