refactor: move max_worker_retries as a common arg in the base orchestrator

Author: shangkunwang01

MaxKernel/auto_search/algorithms/parallel_search.py

@@ -16,18 +16,13 @@ def __init__(
     self,
     num_parallel_runs: int = 2,
     strategies: Optional[List[str]] = None,
-    max_worker_retries: int = 1,
     agent_config: Optional[dict] = None,
     **kwargs,
   ):
     if num_parallel_runs <= 0:
       raise ValueError(
         f"num_parallel_runs must be a positive integer, got {num_parallel_runs}."
       )
-    if max_worker_retries < 1:
-      raise ValueError(
-        f"max_worker_retries must be at least 1, got {max_worker_retries}."
-      )
 
     self.num_parallel_runs = num_parallel_runs
     if strategies is not None:
@@ -43,7 +38,6 @@ def __init__(
       self.strategies = [""] * num_parallel_runs
 
     self.remaining_strategies = list(self.strategies)
-    self.max_worker_retries = max_worker_retries
     self.agent_config = agent_config
     self.worker = ADKSessionWorker()
 

MaxKernel/auto_search/orchestrator.py

@@ -17,7 +17,13 @@ def __init__(
     reference_code: str,
     graph_db_path: Optional[str] = None,
     max_concurrency: int = 2,
+    max_worker_retries: int = 1,
   ):
+    if max_worker_retries < 1:
+      raise ValueError(
+        f"max_worker_retries must be at least 1, got {max_worker_retries}."
+      )
+    self.max_worker_retries = max_worker_retries
     # Resolve graph db path and run directory
     if not graph_db_path:
       workdir = os.environ.get("WORKDIR", os.getcwd())
@@ -152,8 +158,6 @@ def _post_step_hook(self) -> None:
     """
     pass
 
-  # --- Concrete Flow ---
-
   @abc.abstractmethod
   async def _execute_expansions(self, tasks: Any) -> List[Node]:
     """Executes expansion tasks to evaluate new kernels.
@@ -163,6 +167,7 @@ async def _execute_expansions(self, tasks: Any) -> List[Node]:
     """
     pass
 
+  # --- Workflow ---
   async def run(self) -> Node:
     logger.info(
       f"Starting search orchestration for problem: {self.graph.problem_id}"