From 6c83c858baa5a8d3ae6d53825212ea62aaeb1c37 Mon Sep 17 00:00:00 2001 From: Giulia Baldini Date: Mon, 22 Jul 2024 16:30:18 +0200 Subject: [PATCH] Allow user to give number of processes --- config/default.yaml | 1 + exact_kmeans/ilp.py | 22 +++++++++++++++------- 2 files changed, 16 insertions(+), 7 deletions(-) diff --git a/config/default.yaml b/config/default.yaml index 8ceb28b..4a6c3a9 100644 --- a/config/default.yaml +++ b/config/default.yaml @@ -1,5 +1,6 @@ ilp_version: v4 +num_processes: 4 branch_priorities: true replace_min: true branching_ilp: true diff --git a/exact_kmeans/ilp.py b/exact_kmeans/ilp.py index ad61917..5a44cd6 100755 --- a/exact_kmeans/ilp.py +++ b/exact_kmeans/ilp.py @@ -60,13 +60,6 @@ def __init__( self._k = self.k + self._v self.kmeans_iterations = kmeans_iterations - self.num_processes = multiprocessing.cpu_count() - - logger.info( - f"Using {self.num_processes} processes for " - f"multiprocessing out of {multiprocessing.cpu_count()} total processes." - ) - self.changed_model_params = {} self.changed_bound_model_params = {} with Path(config_file).open("r") as f: @@ -111,6 +104,21 @@ def __init__( ) self.ilp_version += "-fill-sizes" + self.num_processes = self.config.get( + "num_processes", multiprocessing.cpu_count() + ) + if isinstance(self.num_processes, int): + self.num_processes = min(self.num_processes, multiprocessing.cpu_count()) + elif isinstance(self.num_processes, float): + self.num_processes = int(self.num_processes * multiprocessing.cpu_count()) + else: + raise ValueError("num_processes must be an integer or a float.") + + logger.info( + f"Using {self.num_processes} processes for " + f"multiprocessing out of {multiprocessing.cpu_count()} total processes." + ) + self.cache_current_run_path = cache_current_run_path if load_existing_run_path is not None and load_existing_run_path.exists():