Add n_jobs param to BenchmarkConfig

GitOrigin-RevId: feaf8a226e10527cd1b778bdfa8b7fc5933224c6
gretelai · Apr 10, 2024 · c4181f9 · c4181f9
1 parent f175bb9
commit c4181f9
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/src/gretel_trainer/benchmark/core.py b/src/gretel_trainer/benchmark/core.py
@@ -46,12 +46,14 @@ def __init__(
         trainer: bool = False,
         working_dir: Optional[Union[str, Path]] = None,
         additional_report_scores: Optional[list[str]] = None,
+        n_jobs: int = 5,
     ):
         self.project_display_name = project_display_name or _default_name()
         self.working_dir = Path(working_dir or self.project_display_name)
         self.refresh_interval = refresh_interval
         self.trainer = trainer
         self.additional_report_scores = additional_report_scores or []
+        self.n_jobs = n_jobs
 
 
 class Timer:

diff --git a/src/gretel_trainer/benchmark/session.py b/src/gretel_trainer/benchmark/session.py
@@ -52,7 +52,7 @@ def __init__(
         self._gretel_executors: dict[RunKey, Executor] = {}
         self._custom_executors: dict[RunKey, Executor] = {}
         self._trainer_project_names: dict[str, str] = {}
-        self._thread_pool = ThreadPoolExecutor(5)
+        self._thread_pool = ThreadPoolExecutor(self._config.n_jobs)
         self._futures: dict[FutureKeyT, Future] = {}
 
         self._report_scores = {