expose multiprocessing start_method

huggingface · Nov 6, 2023 · 8bc449e · 8bc449e
1 parent a860946
commit 8bc449e
Showing 1 changed file with 4 additions and 4 deletions.
diff --git 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")