Allow BenchmarkProblem to take in custom BenchmarkRunner (facebook#4948)

Summary:

`BenchmarkRunner` by default assumes all trials are completed in `poll_trial_status`. We need a custom runner (`FailureAwareBenchmarkRunner` from previous diff) to override `poll_trial_status` which stores trial status depending on if the trial data is NaN (should be abandoned) or is available (completed)

Differential Revision: D87091317

Author: Sunny Shen

ax/benchmark/benchmark.py

@@ -576,11 +576,17 @@ def run_optimization_with_orchestrator(
         include_status_quo=sq_arm is not None,
         logging_level=orchestrator_logging_level,
     )
-    runner = get_benchmark_runner(
-        problem=problem,
-        max_concurrency=orchestrator_options.max_pending_trials,
-        force_use_simulated_backend=method.early_stopping_strategy is not None,
-    )
+
+    # Use custom runner if provided on the problem, otherwise create standard runner
+    if problem.runner is not None:
+        runner = problem.runner
+    else:
+        runner = get_benchmark_runner(
+            problem=problem,
+            max_concurrency=orchestrator_options.max_pending_trials,
+            force_use_simulated_backend=method.early_stopping_strategy is not None,
+        )
+
     experiment = Experiment(
         name=f"{problem.name}|{method.name}_{int(time())}",
         search_space=problem.search_space,

ax/benchmark/benchmark_problem.py

@@ -10,6 +10,7 @@
 from dataclasses import dataclass, field
 
 from ax.benchmark.benchmark_metric import BenchmarkMapMetric, BenchmarkMetric
+from ax.benchmark.benchmark_runner import BenchmarkRunner
 from ax.benchmark.benchmark_step_runtime_function import TBenchmarkStepRuntimeFunction
 from ax.benchmark.benchmark_test_function import BenchmarkTestFunction
 from ax.benchmark.noise import GaussianNoise, Noise
@@ -107,6 +108,7 @@ class BenchmarkProblem(Base):
         dict[AuxiliaryExperimentPurpose, list[AuxiliaryExperiment]] | None
     ) = None
     tracking_metrics: list[Metric] | None = None
+    runner: BenchmarkRunner | None = None
 
     def __post_init__(self) -> None:
         # Handle backward compatibility for noise_std parameter