examples/schedulers/run_config_driven_slurm_design_json.py

0001 #!/usr/bin/env python
0002 """Run the config-driven Slurm design-JSON workflow example end to end."""
0003
0004 from __future__ import annotations
0005
0006 import json
0007 import sys
0008 from datetime import UTC, datetime
0009 from pathlib import Path
0010
0011
0012 REPO_ROOT = Path(__file__).resolve().parents[2]
0013 SRC_ROOT = REPO_ROOT / "src"
0014 for candidate in (str(REPO_ROOT), str(SRC_ROOT)):
0015     if candidate not in sys.path:
0016         sys.path.insert(0, candidate)
0017
0018 from aid2e.utilities import build_workflow_executor_from_config
0019 from aid2e.utilities.configurations import load_config
0020
0021
0022 def main() -> int:
0023     config_path = REPO_ROOT / "examples" / "schedulers" / "workflow_example_slurm_design_json.yml"
0024     run_dir = REPO_ROOT / "experimental_tests" / "output" / (
0025         f"slurm_design_json_example_{datetime.now(UTC).strftime('%Y%m%d_%H%M%S')}"
0026     )
0027     run_dir.mkdir(parents=True, exist_ok=True)
0028     (REPO_ROOT / "experimental_tests" / "output" / "slurm_v2_output").mkdir(parents=True, exist_ok=True)
0029     (REPO_ROOT / "experimental_tests" / "output" / "slurm_v2_work").mkdir(parents=True, exist_ok=True)
0030
0031     config = load_config(str(config_path))
0032     if config.scheduler is None or config.workflows is None:
0033         raise RuntimeError("Expected scheduler and workflows sections in the Slurm design-JSON example")
0034
0035     executor = build_workflow_executor_from_config(
0036         config.workflows,
0037         problem_cfg=config.problem,
0038         scheduler_cfg=config.scheduler,
0039         workflow_name="dtlz2_slurm_design_json_eval",
0040         base_output_dir=str(run_dir),
0041         log_level="INFO",
0042     )
0043
0044     design_point = {
0045         "DTLZ2_variables.x1": 0.22,
0046         "DTLZ2_variables.x2": 0.61,
0047         "DTLZ2_variables.x3": 0.41,
0048         "DTLZ2_variables.x4": 0.56,
0049         "DTLZ2_variables.x5": 0.44,
0050     }
0051     objectives = executor.execute(design_point)
0052     if not objectives or "f1" not in objectives or "f2" not in objectives:
0053         raise RuntimeError(f"Expected f1/f2 objectives, got: {objectives}")
0054
0055     design_files = sorted(str(path) for path in executor.work_dir.glob("*/*/design_point.json"))
0056     result_files = sorted(str(path) for path in executor.output_dir.glob("*/*/objectives.json"))
0057     stdout_logs = sorted(str(path) for path in executor.output_dir.glob("*/*/stdout.log"))
0058     stderr_logs = sorted(str(path) for path in executor.output_dir.glob("*/*/stderr.log"))
0059     if not design_files:
0060         raise RuntimeError("Expected the Slurm workflow to materialize at least one design_point.json file")
0061     if not result_files:
0062         raise RuntimeError("Expected the Slurm workflow to produce at least one objectives.json file")
0063
0064     summary = {
0065         "config_path": str(config_path),
0066         "design_point": design_point,
0067         "objectives": objectives,
0068         "work_dir": str(executor.work_dir),
0069         "output_dir": str(executor.output_dir),
0070         "design_files": design_files,
0071         "result_files": result_files,
0072         "stdout_logs": stdout_logs,
0073         "stderr_logs": stderr_logs,
0074         "global_xcom_keys": sorted(executor.global_xcom.keys()),
0075     }
0076     summary_path = run_dir / "run_config_driven_slurm_design_json.summary.json"
0077     summary_path.write_text(json.dumps(summary, indent=2), encoding="utf-8")
0078     print(json.dumps(summary, indent=2))
0079     return 0
0080
0081
0082 if __name__ == "__main__":
0083     raise SystemExit(main())