have the client manage the TPE

Author: KRRT7

codeflash/api/aiservice.py

@@ -35,9 +35,6 @@
     from codeflash.models.models import AIServiceCodeRepairRequest, AIServiceRefinerRequest
     from codeflash.result.explanation import Explanation
 
-multi_model_executor = concurrent.futures.ThreadPoolExecutor(max_workers=10, thread_name_prefix="multi_model")
-
-
 class AiServiceClient:
     def __init__(self) -> None:
         self.base_url = self.get_aiservice_base_url()
@@ -251,6 +248,7 @@ def optimize_python_code_multi_model(
         *,
         is_async: bool = False,
         sequence_offset: int = 0,
+        executor: concurrent.futures.ThreadPoolExecutor | None = None,
     ) -> tuple[list[OptimizedCandidate], int]:
         """Generate optimizations using multiple models in parallel."""
         logger.info("Generating optimized candidates…")
@@ -264,7 +262,7 @@ def optimize_python_code_multi_model(
                 call_trace_id = f"{base_trace_id[:-3]}0{call_index:02x}"
                 call_sequence = sequence_offset + call_index + 1
                 call_index += 1
-                future = multi_model_executor.submit(
+                future = executor.submit(
                     self.optimize_python_code,
                     source_code,
                     dependency_code,
@@ -299,6 +297,7 @@ def optimize_python_code_line_profiler_multi_model(
         model_distribution: list[tuple[str, int]],
         experiment_metadata: ExperimentMetadata | None = None,
         sequence_offset: int = 0,
+        executor: concurrent.futures.ThreadPoolExecutor | None = None,
     ) -> tuple[list[OptimizedCandidate], int]:
         """Generate line profiler optimizations using multiple models in parallel."""
         logger.info("Generating optimized candidates with line profiler…")
@@ -312,7 +311,7 @@ def optimize_python_code_line_profiler_multi_model(
                 call_trace_id = f"{base_trace_id[:-3]}1{call_index:02x}"
                 call_sequence = sequence_offset + call_index + 1
                 call_index += 1
-                future = multi_model_executor.submit(
+                future = executor.submit(
                     self.optimize_python_code_line_profiler,
                     source_code,
                     dependency_code,

codeflash/optimization/function_optimizer.py

@@ -961,6 +961,7 @@ def determine_best_candidate(
             if self.experiment_id
             else None,
             sequence_offset=self.optimize_calls_count,
+            executor=self.executor,
         )
 
         processor = CandidateProcessor(
@@ -1395,6 +1396,7 @@ def generate_optimizations(
             ExperimentMetadata(id=self.experiment_id, group="control") if run_experiment else None,
             is_async=self.function_to_optimize.is_async,
             sequence_offset=N_TESTS_TO_GENERATE_EFFECTIVE,
+            executor=self.executor,
         )
 
         future_references = self.executor.submit(
@@ -1419,6 +1421,7 @@ def generate_optimizations(
                 ExperimentMetadata(id=self.experiment_id, group="experiment"),
                 is_async=self.function_to_optimize.is_async,
                 sequence_offset=N_TESTS_TO_GENERATE_EFFECTIVE,
+                executor=self.executor,
             )
             futures.append(future_candidates_exp)