[Autotuner] Rebenchmark pool configs on owner workers

stack-info: PR: #2295, branch: choijon5/stack/52

Author: choijon5

helion/autotuner/benchmark_provider.py

@@ -1215,6 +1215,14 @@ def rebenchmark(
         if self._precompile_args_path is None:
             return None
 
+        if self._pool_manager is not None:
+            return self.benchmark_isolated(
+                fns,
+                warmup=25,
+                rep=100,
+                desc=desc,
+            )
+
         fn_specs: list[SerializedCompiledFunction] = []
         for fn in fns:
             fn_spec = self._serialize_fn_for_worker(cast("CompiledConfig", fn))
@@ -1298,9 +1306,7 @@ def benchmark_isolated(
                 if match_unrecoverable_runtime_error(e):
                     self.log.warning(f"{desc} sticky CUDA error skipped: {e}")
                 else:
-                    self.log.debug(
-                        f"{desc} subprocess raised: {type(e).__name__}: {e}"
-                    )
+                    self.log.debug(f"{desc} subprocess raised: {type(e).__name__}: {e}")
                 self._autotune_metrics.num_compile_failures += 1
                 timing = inf
             timings.append(float(timing))

test/test_benchmark_worker.py

@@ -27,7 +27,6 @@
 from helion.autotuner.base_search import PopulationBasedSearch
 from helion.autotuner.base_search import PopulationMember
 from helion.autotuner.benchmark_job import BenchmarkJob
-from helion.autotuner.benchmark_job import RebenchmarkJob
 from helion.autotuner.benchmark_job import _load_args
 from helion.autotuner.benchmark_pool import PoolBenchmarkManager
 from helion.autotuner.benchmark_provider import LocalBenchmarkProvider
@@ -306,27 +305,23 @@ def fake_fn_b() -> None:
             self.assertEqual(job.warmup, 25)
             self.assertEqual(job.rep, 100)
 
-    def test_rebenchmark_uses_worker_pool(self) -> None:
-        # Full-effort rebenchmarking should run on the worker that precompiled the config.
+    def test_rebenchmark_uses_owner_isolated_worker_pool(self) -> None:
+        # Pool rebenchmarking should isolate each candidate on its owner worker.
         class FakePoolManager:
             def __init__(self) -> None:
-                self.worker_index: int | None = None
-                self.job: object | None = None
-                self.timeout: float | None = None
+                self.calls: list[tuple[int, object, float]] = []
 
-            def worker_index_for_fn(self, _fn: object) -> int:
-                return 3
+            def worker_index_for_fn(self, fn: object) -> int:
+                return 3 if fn is fake_fn_a else 4
 
             def run_job_on_worker(
                 self,
                 worker_index: int,
                 job: object,
                 timeout: float,
-            ) -> list[float]:
-                self.worker_index = worker_index
-                self.job = job
-                self.timeout = timeout
-                return [1.0, 2.0]
+            ) -> float:
+                self.calls.append((worker_index, job, timeout))
+                return float(worker_index)
 
         class FakeLog:
             def warning(self, *_args: object, **_kwargs: object) -> None:
@@ -335,7 +330,10 @@ def warning(self, *_args: object, **_kwargs: object) -> None:
             def debug(self, *_args: object, **_kwargs: object) -> None:
                 pass
 
-        def fake_fn() -> None:
+        def fake_fn_a() -> None:
+            pass
+
+        def fake_fn_b() -> None:
             pass
 
         pool = FakePoolManager()
@@ -358,15 +356,12 @@ def fake_fn() -> None:
             module_name=None,
         )
 
-        result = provider.rebenchmark([fake_fn, fake_fn], repeat=7, desc="verify")
+        result = provider.rebenchmark([fake_fn_a, fake_fn_b], repeat=7, desc="verify")
 
-        self.assertEqual(result, [1.0, 2.0])
-        self.assertEqual(pool.worker_index, 3)
-        self.assertIsInstance(pool.job, RebenchmarkJob)
-        assert isinstance(pool.job, RebenchmarkJob)
-        self.assertEqual(pool.job.repeat, 7)
-        self.assertEqual(len(pool.job.fn_specs), 2)
-        self.assertEqual(pool.timeout, 20.0)
+        self.assertEqual(result, [3.0, 4.0])
+        self.assertEqual([call[0] for call in pool.calls], [3, 4])
+        self.assertTrue(all(isinstance(call[1], BenchmarkJob) for call in pool.calls))
+        self.assertTrue(all(call[2] == 10.0 for call in pool.calls))
 
     def test_population_rebenchmark_uses_provider_timings(self) -> None:
         # BaseSearch should use provider rebenchmark timings when available.