File indexing completed on 2026-11-09 09:29:52
0001 """Run the advanced inline Ax configuration from dtlz2_starter_kit.ipynb in debug mode."""
0002
0003 from __future__ import annotations
0004
0005 import argparse
0006 import json
0007 import logging
0008 import os
0009 import sys
0010 from datetime import datetime, timezone
0011 from pathlib import Path
0012 from pprint import pformat
0013 from typing import Any
0014
0015 REPO_ROOT = Path(__file__).resolve().parents[3]
0016 if str(REPO_ROOT) not in sys.path:
0017 sys.path.insert(0, str(REPO_ROOT))
0018
0019 from examples.dtlz2.dtlz2_evaluator import objective_payload
0020 from aid2e.utilities import build_optimizer_from_config
0021 from aid2e.utilities.configurations import load_config
0022 from aid2e.utilities.configurations.optimizer_config import OptimizerConfiguration
0023
0024
0025 DEFAULT_CONFIG_PATH = Path(__file__).resolve().parents[1] / "configurations" / "dtlz2_ax_joblib.yml"
0026 LOGGER = logging.getLogger("examples.optimizers.ax_advanced_inline_debug")
0027
0028
0029 def parse_args(argv: list[str]) -> argparse.Namespace:
0030 parser = argparse.ArgumentParser(
0031 description=(
0032 "Exercise the inline advanced Ax setup from dtlz2_starter_kit.ipynb "
0033 "with debug logging and JSON output."
0034 )
0035 )
0036 parser.add_argument(
0037 "--config",
0038 type=Path,
0039 default=DEFAULT_CONFIG_PATH,
0040 help="Path to the DTLZ2 optimizer config YAML.",
0041 )
0042 parser.add_argument(
0043 "--output-dir",
0044 type=Path,
0045 default=Path.cwd(),
0046 help="Directory for JSON artifacts. Defaults to the current working directory.",
0047 )
0048 parser.add_argument(
0049 "--debug",
0050 action=argparse.BooleanOptionalAction,
0051 default=True,
0052 help="Enable verbose debug logging.",
0053 )
0054 return parser.parse_args(argv)
0055
0056
0057 def configure_logging(debug: bool) -> None:
0058 level = logging.DEBUG if debug else logging.INFO
0059 logging.basicConfig(
0060 level=level,
0061 format="%(asctime)s %(levelname)-8s [%(name)s] %(message)s",
0062 force=True,
0063 )
0064 for logger_name in ("aid2e", "ax", "botorch", "gpytorch"):
0065 logging.getLogger(logger_name).setLevel(level)
0066
0067
0068 def dtlz2_objectives(parameters: dict[str, Any]) -> dict[str, float]:
0069 """Compute the two-objective DTLZ2 values for the example design space."""
0070 return objective_payload(design_point=parameters)
0071
0072
0073 def trial_to_dict(trial: Any) -> dict[str, Any]:
0074 return {
0075 "index": trial.index,
0076 "status": trial.status,
0077 "parameters": dict(trial.parameters),
0078 "metrics": dict(trial.metrics or {}),
0079 "metadata": dict(trial.metadata or {}),
0080 }
0081
0082
0083 def evaluate_candidates(
0084 optimizer: Any,
0085 candidates: list[dict[str, Any]],
0086 *,
0087 phase: str,
0088 ) -> list[dict[str, Any]]:
0089 """Evaluate locally and feed results back into the optimizer."""
0090 start_index = len(optimizer.get_trials()) - len(candidates)
0091 records: list[dict[str, Any]] = []
0092
0093 for offset, parameters in enumerate(candidates):
0094 trial_index = start_index + offset
0095 metrics = dtlz2_objectives(parameters)
0096 optimizer.update_with_results(
0097 trial_index=trial_index,
0098 parameters=parameters,
0099 metrics=metrics,
0100 )
0101 record = {
0102 "trial_index": trial_index,
0103 "phase": phase,
0104 "parameters": dict(parameters),
0105 "metrics": metrics,
0106 }
0107 LOGGER.debug("Recorded %s", pformat(record))
0108 records.append(record)
0109
0110 return records
0111
0112
0113 def build_inline_payload() -> dict[str, Any]:
0114 return {
0115 "name": "ax",
0116 "type": "bayesian",
0117 "parameters": {
0118 "initialization_strategy": "sobol",
0119 "generator": "BOTORCH_MODULAR",
0120 "generator_kwargs": {
0121 "surrogate_spec": {
0122 "model_configs": [
0123 {
0124 "botorch_model_class": "SaasFullyBayesianSingleTaskGP"
0125 }
0126 ]
0127 },
0128 "botorch_acqf_class": "qLogNoisyExpectedHypervolumeImprovement",
0129 },
0130 "objective_thresholds": {"f1": 1.0, "f2": 1.0},
0131 "n_initial_samples": 10,
0132 "n_iterations": 20,
0133 "batch_size": 4,
0134 "seed": 7,
0135 },
0136 }
0137
0138
0139 def main(argv: list[str]) -> int:
0140 args = parse_args(argv)
0141 configure_logging(args.debug)
0142
0143 config_path = args.config.resolve()
0144 output_dir = args.output_dir.resolve()
0145 output_dir.mkdir(parents=True, exist_ok=True)
0146
0147 LOGGER.info("Repository root: %s", REPO_ROOT)
0148 LOGGER.info("Config path: %s", config_path)
0149 LOGGER.info("Output directory: %s", output_dir)
0150 LOGGER.info("Debug logging enabled: %s", args.debug)
0151
0152 ax_full_config = load_config(str(config_path))
0153 ax_inline_optimizer_payload = build_inline_payload()
0154 ax_inline_optimizer_cfg = OptimizerConfiguration(**ax_inline_optimizer_payload)
0155 ax_inline_optimizer = build_optimizer_from_config(
0156 ax_full_config.problem,
0157 ax_inline_optimizer_cfg,
0158 )
0159
0160 optimizer_config = ax_inline_optimizer_cfg.parse_algorithm_params()
0161 if optimizer_config is None:
0162 raise RuntimeError("Failed to parse the inline Ax optimizer configuration.")
0163
0164 nodes = getattr(ax_inline_optimizer.generation_strategy, "nodes", None)
0165 generation_summary = {
0166 "name": ax_inline_optimizer.generation_strategy.name,
0167 "nodes": [node.name for node in nodes] if nodes else [],
0168 }
0169
0170 print("Inline Ax optimizer payload (SAAS surrogate + qLogNEHVI):")
0171 print(pformat(ax_inline_optimizer_payload))
0172 print("Generation strategy summary:")
0173 print(pformat(generation_summary))
0174
0175 records: list[dict[str, Any]] = []
0176 first_candidate: dict[str, Any] | None = None
0177
0178 remaining_init = optimizer_config.n_initial_samples
0179 init_round = 0
0180 while remaining_init > 0:
0181 init_round += 1
0182 current_batch = min(optimizer_config.batch_size, remaining_init)
0183 candidates = ax_inline_optimizer.suggest_candidates(n_candidates=current_batch)
0184 if first_candidate is None and candidates:
0185 first_candidate = dict(candidates[0])
0186 print("First Ax candidate from the inline config:")
0187 print(pformat(first_candidate))
0188 records.extend(
0189 evaluate_candidates(
0190 ax_inline_optimizer,
0191 candidates,
0192 phase=f"init-{init_round}",
0193 )
0194 )
0195 remaining_init -= current_batch
0196
0197 for iteration in range(optimizer_config.n_iterations):
0198 candidates = ax_inline_optimizer.suggest_candidates(
0199 n_candidates=optimizer_config.batch_size
0200 )
0201 records.extend(
0202 evaluate_candidates(
0203 ax_inline_optimizer,
0204 candidates,
0205 phase=f"iter-{iteration + 1}",
0206 )
0207 )
0208
0209 summary = {
0210 "n_trials": ax_inline_optimizer.get_optimization_results()["n_trials"],
0211 "pareto_points": len(ax_inline_optimizer.get_pareto_front()),
0212 "objective_names": [obj.name for obj in ax_full_config.problem.objectives],
0213 }
0214
0215 artifact = {
0216 "timestamp_utc": datetime.now(timezone.utc).isoformat(),
0217 "slurm_job_id": os.environ.get("SLURM_JOB_ID"),
0218 "cwd": str(Path.cwd()),
0219 "config_path": str(config_path),
0220 "output_dir": str(output_dir),
0221 "debug": args.debug,
0222 "inline_payload": ax_inline_optimizer_payload,
0223 "generation_strategy": generation_summary,
0224 "first_candidate": first_candidate,
0225 "records": records,
0226 "summary": summary,
0227 "trials": [trial_to_dict(trial) for trial in ax_inline_optimizer.get_trials()],
0228 }
0229
0230 job_suffix = os.environ.get("SLURM_JOB_ID", "local")
0231 artifact_path = output_dir / f"ax_advanced_inline_debug_{job_suffix}.json"
0232 artifact_path.write_text(json.dumps(artifact, indent=2), encoding="utf-8")
0233
0234 print("Optimization summary:")
0235 print(pformat(summary))
0236 print(f"Artifact written to: {artifact_path}")
0237 LOGGER.info("Artifact written to %s", artifact_path)
0238 return 0
0239
0240
0241 if __name__ == "__main__":
0242 raise SystemExit(main(sys.argv[1:]))
