class ParamSweep:
"""Grid search with top-k training ranking and validation-set selection.
Parameters
----------
evaluator :
Shared scenario evaluator for running simulations.
config :
Configuration for this tuner run, including the aggregation metric and
other hyperparameters.
base_overrides :
Config overrides that are applied to *every* grid point (e.g.
optimised spin-ups from a previous phase).
run_dir :
Root directory for this tuner run.
phase_name :
Label for the current phase (e.g. ``"autoscaler"`` or
``"routing"``). Used for directory and log naming.
"""
def __init__(
self,
evaluator: ScenarioEvaluator,
initial_config: dict[str, Any],
run_dir: Path,
phase_name: str,
slo_objective: SloObjective,
agg_method: str,
verbose_progress: bool = True,
) -> None:
self._evaluator = evaluator
self._verbose_progress = verbose_progress
self._config = initial_config
self._run_dir = run_dir
self._phase_name = phase_name
self._slo_objective = slo_objective
self._agg_method = agg_method
# ------------------------------------------------------------------
# Public API
# ------------------------------------------------------------------
def sweep(
self,
train_workload_configs: list[WorkloadConfig],
val_workload_configs: list[WorkloadConfig],
param_sweep_config: ParamSweepConfig,
) -> tuple[
dict[str, Any], AggregatedExecutionResults, AggregatedExecutionResults
]:
"""Run a parameter sweep and return the best config with its metrics.
Parameters
----------
train_workload_configs :
Workload configurations for training scenarios.
val_workload_configs :
Workload configurations for validation scenarios.
param_sweep_config :
Sweep configuration.
Returns
-------
The full config dict with the best parameter values applied.
"""
strategy = param_sweep_config.strategy
phase_dir = self._run_dir / self._phase_name
# Empty params → nothing to sweep; fall back to grid (evaluates
# the base config once).
if not param_sweep_config.params:
strategy = "grid"
dump_yaml(self._config, phase_dir / "initial_config.yml")
# ── Generate & evaluate candidates (strategy-specific) ─────
if strategy == "grid":
candidates, grid_results = self._sweep_grid(
train_workload_configs, param_sweep_config, phase_dir
)
elif strategy == "random":
candidates, grid_results = self._sweep_random(
train_workload_configs, param_sweep_config, phase_dir
)
elif strategy == "coordinate_descent":
candidates, grid_results = self._sweep_coordinate_descent(
train_workload_configs, param_sweep_config, phase_dir
)
elif strategy == "adaptive_batch":
candidates, grid_results = self._sweep_adaptive_batch(
train_workload_configs, param_sweep_config, phase_dir
)
else:
raise ValueError(f"Unknown sweep strategy: {strategy!r}")
# ── Rank training candidates and select top-k for validation ─
val_top_k: int = param_sweep_config.val_top_k
points = [
ViolationCost(r["train_violation_agg"], r["train_cost_agg"])
for r in grid_results
]
ranked_indices = self._slo_objective.rank_indices(points)
for rank, idx in enumerate(ranked_indices):
grid_results[idx]["train_rank"] = rank
top_k_indices = ranked_indices[:val_top_k]
console.print(
f"\n [cyan]Top-k validation:[/] {len(top_k_indices)} of "
f"{len(candidates)} points (k={val_top_k})"
)
# ── Validate top-k candidates ──────────────────────────────
console.print(f"\n[bold cyan]Validation sweep:[/]")
val_overrides = [candidates[idx] for idx in top_k_indices]
all_val_results = self._evaluator.evaluate_batch_from_overrides(
progress_bar_label=self._phase_name,
workload_configs=val_workload_configs,
base_config=self._config,
all_config_overrides=val_overrides,
out_dir=phase_dir / "val",
workload_first=False,
verbose_progress=self._verbose_progress,
)
for i, idx in enumerate(top_k_indices):
val_results = all_val_results[i]
val_agg = AggregatedExecutionResults.aggregate_from(
val_results, self._agg_method
)
val_primary = val_agg.primary_violation(
self._slo_objective.slo_metric
)
grid_results[idx]["val_primary_violation_agg"] = val_primary
grid_results[idx]["val_cost_agg"] = val_agg.cost
grid_results[idx]["val_metrics"] = val_agg
# ── Select best ───────────────────────────────────────────
best_idx = self._select_best(grid_results, top_k_indices)
best_params = candidates[best_idx]
# ── Rich summary table ─────────────────────────────────────
self._print_top_k_table(grid_results, top_k_indices, best_idx)
# ── Persist results ────────────────────────────────────────
self._write_sweep_results(phase_dir, grid_results, best_idx)
final_config = cfgu.copy_and_apply_overrides(self._config, best_params)
dump_yaml(final_config, phase_dir / "final_config.yml")
best_train_agg = grid_results[best_idx]["train_metrics"]
best_val_agg = grid_results[best_idx]["val_metrics"]
return final_config, best_train_agg, best_val_agg
# ------------------------------------------------------------------
# Strategy implementations
# ------------------------------------------------------------------
def _evaluate_candidates(
self,
train_workload_configs: list[WorkloadConfig],
candidates: list[dict[str, Any]],
out_dir: Path,
) -> list[dict[str, Any]]:
"""Evaluate *candidates* on training scenarios and return result dicts."""
all_train_results = self._evaluator.evaluate_batch_from_overrides(
progress_bar_label=self._phase_name,
workload_configs=train_workload_configs,
base_config=self._config,
all_config_overrides=candidates,
out_dir=out_dir,
workload_first=False,
verbose_progress=self._verbose_progress,
)
grid_results: list[dict[str, Any]] = []
for idx, candidate in enumerate(candidates):
train_agg = AggregatedExecutionResults.aggregate_from(
all_train_results[idx], self._agg_method
)
train_primary = train_agg.primary_violation(
self._slo_objective.slo_metric
)
grid_results.append(
{
"grid_point": idx,
"params": candidate,
"train_violation_agg": train_primary,
"train_cost_agg": train_agg.cost,
"train_metrics": train_agg,
"train_rank": None,
"val_primary_violation_agg": None,
"val_cost_agg": None,
"val_metrics": None,
}
)
return grid_results
def _sweep_grid(
self,
train_workload_configs: list[WorkloadConfig],
param_sweep_config: ParamSweepConfig,
phase_dir: Path,
) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
"""Exhaustive grid search (original strategy)."""
grid = build_grid(param_sweep_config.params)
self._print_preflight(
param_sweep_config, grid, len(train_workload_configs)
)
console.print(f"\n[bold cyan]Training sweep:[/]")
grid_results = self._evaluate_candidates(
train_workload_configs, grid, phase_dir / "train"
)
return grid, grid_results
def _sweep_random(
self,
train_workload_configs: list[WorkloadConfig],
param_sweep_config: ParamSweepConfig,
phase_dir: Path,
) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
"""Random search: sample *budget* configs from the grid."""
full_grid = build_grid(param_sweep_config.params)
budget = param_sweep_config.budget
seed = param_sweep_config.seed
if budget >= len(full_grid):
grid = full_grid
else:
rng = stdlib_random.Random(seed)
grid = rng.sample(full_grid, budget)
self._print_preflight(
param_sweep_config,
grid,
len(train_workload_configs),
strategy_label=f"Random (budget={budget}, seed={seed})",
)
console.print(f"\n[bold cyan]Training sweep:[/]")
grid_results = self._evaluate_candidates(
train_workload_configs, grid, phase_dir / "train"
)
return grid, grid_results
def _sweep_coordinate_descent(
self,
train_workload_configs: list[WorkloadConfig],
param_sweep_config: ParamSweepConfig,
phase_dir: Path,
) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
"""Coordinate descent: optimise one parameter at a time."""
max_cycles = param_sweep_config.max_cycles
starting_point = param_sweep_config.starting_point
# Default starting point: middle value for each parameter.
if starting_point is None:
starting_point = {
name: values[len(values) // 2]
for name, values in param_sweep_config.params.items()
}
current_best = dict(starting_point)
all_candidates: list[dict[str, Any]] = []
all_grid_results: list[dict[str, Any]] = []
evaluated_cache: dict[frozenset, int] = {}
def _config_key(cfg: dict[str, Any]) -> frozenset:
return frozenset(sorted(cfg.items()))
total_values = sum(len(v) for v in param_sweep_config.params.values())
console.print(
f"\n[bold cyan]Coordinate descent[/] "
f"max_cycles={max_cycles} | "
f"params={len(param_sweep_config.params)} | "
f"total values={total_values} | "
f"training scenarios={len(train_workload_configs)} | "
f"max evals={max_cycles * total_values * len(train_workload_configs)}"
)
console.print(f" Starting point: {current_best}")
for cycle in range(max_cycles):
console.print(f"\n [bold cyan]Cycle {cycle + 1}/{max_cycles}:[/]")
changed = False
for param_name, param_values in param_sweep_config.params.items():
# Candidate configs: vary only this param.
candidates_for_param = []
for val in param_values:
point = dict(current_best)
point[param_name] = val
candidates_for_param.append(point)
# Identify configs not yet evaluated.
new_candidates = [
c
for c in candidates_for_param
if _config_key(c) not in evaluated_cache
]
# Evaluate new candidates.
if new_candidates:
out_dir = (
phase_dir
/ "train"
/ f"cycle_{cycle}"
/ param_name.replace(".", "_")
)
batch_results = self._evaluate_candidates(
train_workload_configs, new_candidates, out_dir
)
for j, c in enumerate(new_candidates):
idx = len(all_candidates)
batch_results[j]["grid_point"] = idx
all_candidates.append(c)
all_grid_results.append(batch_results[j])
evaluated_cache[_config_key(c)] = idx
# Select best value for this parameter.
indices = [
evaluated_cache[_config_key(c)]
for c in candidates_for_param
]
cd_candidates = [
ViolationCost(
all_grid_results[i]["train_violation_agg"],
all_grid_results[i]["train_cost_agg"],
)
for i in indices
]
best_local = self._slo_objective.idx_of_best(cd_candidates)
best_val = param_values[best_local]
if best_val != current_best[param_name]:
console.print(
f" {param_name}: "
f"{current_best[param_name]} → {best_val}"
)
current_best[param_name] = best_val
changed = True
else:
console.print(
f" {param_name}: unchanged "
f"({current_best[param_name]})"
)
if not changed:
console.print(
f"\n [dim]Converged after {cycle + 1} cycle(s).[/dim]"
)
break
return all_candidates, all_grid_results
def _sweep_adaptive_batch(
self,
train_workload_configs: list[WorkloadConfig],
param_sweep_config: ParamSweepConfig,
phase_dir: Path,
) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
"""Adaptive batch search: LHS initialisation then Noisy CEM refinement.
Round 0 uses Latin Hypercube Sampling (McKay et al. 1979) for uniform
coverage. Each subsequent round applies the Cross-Entropy Method with
a diagonal Gaussian (Rubinstein 1999; De Boer et al. 2005) refit to
the top-B elite set. A variance floor follows Noisy CEM (Szita &
Lörincz 2006) to prevent premature distribution collapse.
"""
params = param_sweep_config.params
budget = param_sweep_config.budget
seed = param_sweep_config.seed
max_rounds = param_sweep_config.max_rounds
beam_size = param_sweep_config.beam_size
min_sigma_frac = param_sweep_config.min_sigma_fraction
explr_mult = param_sweep_config.exploration_multiplier
param_names = list(params.keys())
# Validate and extract per-parameter bounds and type flags.
# Type: if all listed values are Python int, treat as integer;
# otherwise continuous float.
lo: dict[str, float] = {}
hi: dict[str, float] = {}
is_int: dict[str, bool] = {}
for name, values in params.items():
distinct = set(values)
if len(distinct) < 2:
raise ValueError(
f"adaptive_batch requires at least 2 distinct values for "
f"{name!r}; got {sorted(distinct)}"
)
lo[name] = float(min(values))
hi[name] = float(max(values))
is_int[name] = all(isinstance(v, int) for v in values)
rng = stdlib_random.Random(seed)
def _clamp(name: str, v: float) -> Any:
v = max(lo[name], min(hi[name], v))
return int(round(v)) if is_int[name] else v
def _config_key(cfg: dict[str, Any]) -> frozenset:
return frozenset(sorted(cfg.items()))
def _std(vals: list[float]) -> float:
if len(vals) < 2:
return 0.0
mean = sum(vals) / len(vals)
return (sum((v - mean) ** 2 for v in vals) / (len(vals) - 1)) ** 0.5
all_candidates: list[dict[str, Any]] = []
all_grid_results: list[dict[str, Any]] = []
evaluated_cache: dict[frozenset, int] = {}
def _absorb(
new_cands: list[dict[str, Any]], batch: list[dict[str, Any]]
) -> None:
"""Append a completed evaluation batch into the accumulators."""
for j, c in enumerate(new_cands):
idx = len(all_candidates)
batch[j]["grid_point"] = idx
all_candidates.append(c)
all_grid_results.append(batch[j])
evaluated_cache[_config_key(c)] = idx
console.print(
f"\n[bold cyan]Adaptive batch (LHS + Noisy CEM)[/] "
f"budget={budget} | max_rounds={max_rounds} | "
f"beam_size={beam_size} | params={len(param_names)} | "
f"training scenarios={len(train_workload_configs)}"
)
# ── Round 0: Latin Hypercube Sampling ─────────────────────────
# Divide [lo, hi] into `budget` equal intervals per dimension;
# draw one sample per interval; shuffle columns independently.
lhs_cols: dict[str, list] = {}
for name in param_names:
raw = [
lo[name] + (i + rng.random()) / budget * (hi[name] - lo[name])
for i in range(budget)
]
rng.shuffle(raw)
lhs_cols[name] = [_clamp(name, v) for v in raw]
round0_candidates = [
{name: lhs_cols[name][i] for name in param_names}
for i in range(budget)
]
console.print(f"\n[bold cyan]Round 0 (LHS, {budget} samples):[/]")
_absorb(
round0_candidates,
self._evaluate_candidates(
train_workload_configs,
round0_candidates,
phase_dir / "train" / "round_0",
),
)
# ── Rounds 1..max_rounds: CEM with Noisy CEM variance floor ───
for rnd in range(1, max_rounds + 1):
# Re-rank all accumulated results and pick the elite set.
points = [
ViolationCost(r["train_violation_agg"], r["train_cost_agg"])
for r in all_grid_results
]
ranked = self._slo_objective.rank_indices(points)
elite = [
all_candidates[i]
for i in ranked[: min(beam_size, len(ranked))]
]
# CEM variance refit: std(elite) × multiplier, floored at
# min_sigma_frac × range (Noisy CEM floor).
sigma: dict[str, float] = {}
for name in param_names:
elite_vals = [float(c[name]) for c in elite]
floor = min_sigma_frac * (hi[name] - lo[name])
sigma[name] = max(_std(elite_vals) * explr_mult, floor)
# Sample `budget` candidates via Gaussian perturbation around
# elite centers (round-robin); retry up to 5× to avoid cache hits.
new_candidates: list[dict[str, Any]] = []
for k in range(budget):
center = elite[k % len(elite)]
for _ in range(5):
proposal = {
name: _clamp(
name, rng.gauss(float(center[name]), sigma[name])
)
for name in param_names
}
if _config_key(proposal) not in evaluated_cache:
new_candidates.append(proposal)
break
# Halt early if the parameter space is saturated.
if len(new_candidates) < math.ceil(budget / 2):
console.print(
f"\n [dim]Converged after {rnd} CEM round(s) "
f"({len(new_candidates)} unique < "
f"{math.ceil(budget / 2)} threshold).[/dim]"
)
break
console.print(
f"\n[bold cyan]Round {rnd} "
f"(CEM, {len(new_candidates)} candidates):[/]"
)
_absorb(
new_candidates,
self._evaluate_candidates(
train_workload_configs,
new_candidates,
phase_dir / "train" / f"round_{rnd}",
),
)
return all_candidates, all_grid_results
# ------------------------------------------------------------------
# Selection logic
# ------------------------------------------------------------------
def _select_best(
self,
grid_results: list[dict[str, Any]],
candidate_indices: list[int],
) -> int:
"""
Pick the best validated candidate.
"""
candidates = [
ViolationCost(
grid_results[i]["val_primary_violation_agg"],
grid_results[i]["val_cost_agg"],
)
for i in candidate_indices
]
best_local_idx = self._slo_objective.idx_of_best(candidates)
return candidate_indices[best_local_idx]
# ------------------------------------------------------------------
# Rich output
# ------------------------------------------------------------------
def _print_preflight(
self,
param_sweep_config: ParamSweepConfig,
grid: list[dict[str, Any]],
n_train: int,
strategy_label: str = "Grid",
) -> None:
table = Table(
title=f"Parameter Ranges — {strategy_label}", show_lines=True
)
table.add_column("Parameter", justify="left")
table.add_column("Values", justify="left")
table.add_column("Count", justify="right")
for name, values in param_sweep_config.params.items():
table.add_row(name, str(values), str(len(values)))
console.print(table)
console.print(
f" Grid size: [bold]{len(grid)}[/] combinations | "
f"Training scenarios: [bold]{n_train}[/] | "
f"Total evaluations: [bold]{len(grid) * n_train}[/]"
)
def _print_top_k_table(
self,
grid_results: list[dict[str, Any]],
top_k_indices: list[int],
best_idx: int,
) -> None:
table = Table(
title=f"{self._phase_name} — Top-K Validation Results",
show_lines=True,
)
slo_label = self._slo_objective.slo_metric
slo_thresh = self._slo_objective.slo_threshold
table.add_column("GP", justify="right")
table.add_column("Train Rank", justify="right")
table.add_column("Params", justify="left")
table.add_column(
f"Train {slo_label} (≤{slo_thresh:.2f})", justify="right"
)
table.add_column("Train Cost", justify="right")
table.add_column(
f"Val {slo_label} (≤{slo_thresh:.2f})", justify="right"
)
table.add_column("Val Cost", justify="right")
table.add_column("Best", justify="center")
for idx in top_k_indices:
r = grid_results[idx]
is_best = idx == best_idx
style = "bold green" if is_best else ""
val_v = (
f"{r['val_primary_violation_agg']:.4f}"
if r["val_primary_violation_agg"] is not None
else "—"
)
val_c = (
f"{r['val_cost_agg']:.4f}"
if r["val_cost_agg"] is not None
else "—"
)
table.add_row(
str(idx),
str(r["train_rank"]),
str(r["params"]),
f"{r['train_violation_agg']:.4f}",
f"{r['train_cost_agg']:.4f}",
val_v,
val_c,
"✓" if is_best else "",
style=style,
)
console.print(table)
# ------------------------------------------------------------------
# Persistence
# ------------------------------------------------------------------
@staticmethod
def _write_sweep_results(
phase_dir: Path,
grid_results: list[dict[str, Any]],
best_idx: int,
) -> None:
phase_dir.mkdir(parents=True, exist_ok=True)
# Serialize grid_results, converting AggregatedMetrics to dicts.
serializable = []
for r in grid_results:
entry = {
k: v
for k, v in r.items()
if k not in ("train_metrics", "val_metrics")
}
tm = r.get("train_metrics")
if tm is not None:
entry["train_violation_rate"] = tm.violation_rate
entry["train_violation_amount_s"] = tm.violation_amount_s
entry["train_violation_relative_mean"] = (
tm.violation_relative_mean
)
vm = r.get("val_metrics")
if vm is not None:
entry["val_violation_rate"] = vm.violation_rate
entry["val_violation_amount_s"] = vm.violation_amount_s
entry["val_violation_relative_mean"] = (
vm.violation_relative_mean
)
serializable.append(entry)
output = {
"best_grid_point": best_idx,
"best_params": grid_results[best_idx]["params"],
"grid_results": serializable,
}
with open(phase_dir / "sweep_results.json", "w") as f:
json.dump(output, f, indent=2, default=str)