[Autotuner] Add isolated worker benchmarking

stack-info: PR: #2294, branch: choijon5/stack/51

Author: choijon5

helion/autotuner/benchmark_provider.py

@@ -299,6 +299,22 @@ def rebenchmark(
         """
         return None
 
+    def benchmark_isolated(
+        self,
+        fns: list[Callable[..., object]],
+        *,
+        warmup: int,
+        rep: int,
+        desc: str = "Benchmarking",
+    ) -> list[float] | None:
+        """Optionally benchmark callables independently via the provider.
+
+        Returning ``None`` means the caller should use its default in-process
+        path. Returning a list gives one timing per callable; isolated failures
+        should be represented as ``inf`` timings.
+        """
+        return None
+
     @abc.abstractmethod
     def setup(self) -> None:
         """Prepare resources needed before benchmarking begins (e.g. tmpdir)."""
@@ -1234,3 +1250,58 @@ def rebenchmark(
                 self.log.debug(f"{desc} subprocess raised: {type(e).__name__}: {e}")
         self._autotune_metrics.num_compile_failures += len(fn_specs)
         return [inf] * len(fn_specs)
+
+    def benchmark_isolated(
+        self,
+        fns: list[Callable[..., object]],
+        *,
+        warmup: int,
+        rep: int,
+        desc: str = "Benchmarking",
+    ) -> list[float] | None:
+        if not self._subprocess_benchmark_enabled():
+            return None
+        if self.settings.autotune_benchmark_fn is not None:
+            return None
+        if self._precompile_args_path is None:
+            return None
+
+        timings: list[float] = []
+        timeout = float(self.settings.autotune_benchmark_timeout)
+        for fn in fns:
+            fn_spec = self._serialize_fn_for_worker(cast("CompiledConfig", fn))
+            if fn_spec is None:
+                return None
+            job = BenchmarkJob(
+                fn_spec=fn_spec,
+                args_path=self._precompile_args_path,
+                warmup=warmup,
+                rep=rep,
+            )
+            worker_index = (
+                self._pool_manager.worker_index_for_fn(fn)
+                if self._pool_manager is not None
+                else 0
+            )
+            try:
+                timing = self._run_subprocess_job(
+                    job,
+                    timeout,
+                    worker_index=worker_index,
+                )
+            except BenchmarkSubprocessError as e:
+                self.log.warning(f"{desc} subprocess failed: {e}")
+                self._autotune_metrics.num_compile_failures += 1
+                timing = inf
+            except Exception as e:
+                e.__traceback__ = None
+                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._autotune_metrics.num_compile_failures += 1
+                timing = inf
+            timings.append(float(timing))
+        return timings

test/test_benchmark_worker.py

@@ -26,6 +26,7 @@
 from helion._testing import skipIfXPU
 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
@@ -238,6 +239,73 @@ def test_pool_mode_reports_disabled_worker_reason(self) -> None:
         assert reason is not None
         self.assertIn("disabled", reason)
 
+    def test_benchmark_isolated_routes_each_fn_to_owner_worker(self) -> None:
+        # Isolated benchmarking should route each config to its precompile owner.
+        class FakePoolManager:
+            def __init__(self) -> None:
+                self.calls: list[tuple[int, object, float]] = []
+
+            def worker_index_for_fn(self, fn: object) -> int:
+                return 2 if fn is fake_fn_a else 1
+
+            def run_job_on_worker(
+                self,
+                worker_index: int,
+                job: object,
+                timeout: float,
+            ) -> float:
+                self.calls.append((worker_index, job, timeout))
+                return float(worker_index)
+
+        class FakeLog:
+            def warning(self, *_args: object, **_kwargs: object) -> None:
+                pass
+
+            def debug(self, *_args: object, **_kwargs: object) -> None:
+                pass
+
+        def fake_fn_a() -> None:
+            pass
+
+        def fake_fn_b() -> None:
+            pass
+
+        pool = FakePoolManager()
+        provider = cast("Any", LocalBenchmarkProvider.__new__(LocalBenchmarkProvider))
+        provider.settings = Settings(
+            autotune_precompile="pool",
+            autotune_benchmark_timeout=10,
+        )
+        provider.config_spec = SimpleNamespace(backend=None)
+        provider.mutated_arg_indices = []
+        provider._precompile_args_path = "args.pt"
+        provider._pool_manager = pool
+        provider._benchmark_worker = None
+        provider._autotune_metrics = SimpleNamespace(num_compile_failures=0)
+        provider.log = FakeLog()
+        provider._serialize_fn_for_worker = lambda _fn: SerializedCompiledFunction(
+            function_name="fake_fn",
+            source_code="def fake_fn(): pass",
+            filename=None,
+            module_name=None,
+        )
+
+        result = provider.benchmark_isolated(
+            [fake_fn_a, fake_fn_b],
+            warmup=25,
+            rep=100,
+            desc="verify",
+        )
+
+        self.assertEqual(result, [2.0, 1.0])
+        self.assertEqual([call[0] for call in pool.calls], [2, 1])
+        self.assertTrue(all(isinstance(call[1], BenchmarkJob) for call in pool.calls))
+        self.assertTrue(all(call[2] == 10.0 for call in pool.calls))
+        for _, job, _ in pool.calls:
+            assert isinstance(job, BenchmarkJob)
+            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.
         class FakePoolManager: