File indexing completed on 2026-11-09 09:29:52
0001
0002 """Submit ten Slurm jobs together, monitor them, and collect per-job outputs."""
0003
0004 from __future__ import annotations
0005
0006 import json
0007 import sys
0008 import time
0009 from collections import Counter
0010 from datetime import UTC, datetime
0011 from pathlib import Path
0012 from typing import Any, Dict, List
0013
0014
0015 REPO_ROOT = Path(__file__).resolve().parents[3]
0016 SRC_ROOT = REPO_ROOT / "src"
0017 for candidate in (str(REPO_ROOT), str(SRC_ROOT)):
0018 if candidate not in sys.path:
0019 sys.path.insert(0, candidate)
0020
0021 from aid2e.optimizers.base import Trial
0022 from aid2e.schedulers.Slurm import SlurmRunnerConfig, SlurmScheduler
0023
0024
0025 DEFAULT_SETUP_COMMANDS = [
0026 "module load miniforge3/24.9.2-0",
0027 "conda activate env_AID2E",
0028 ]
0029
0030
0031 def build_dtlz2_point(index: int) -> Dict[str, float]:
0032 return {
0033 "DTLZ2_variables.x1": 0.10 + 0.05 * index,
0034 "DTLZ2_variables.x2": 0.65 - 0.02 * index,
0035 "DTLZ2_variables.x3": 0.35 + 0.01 * index,
0036 "DTLZ2_variables.x4": 0.58 - 0.01 * index,
0037 "DTLZ2_variables.x5": 0.42 + 0.01 * index,
0038 **{f"DTLZ2_variables.x{i}": 0.5 for i in range(6, 11)},
0039 }
0040
0041
0042 def make_scheduler(run_dir: Path) -> SlurmScheduler:
0043 return SlurmScheduler(
0044 config=SlurmRunnerConfig(
0045 job_name_prefix="aid2e_multi10",
0046 setup_commands=list(DEFAULT_SETUP_COMMANDS),
0047 submit_working_dir=str(run_dir),
0048 runtime_working_dir=str(REPO_ROOT),
0049 poll_interval=2,
0050 )
0051 )
0052
0053
0054 def main() -> int:
0055 run_dir = REPO_ROOT / "experimental_tests" / "output" / (
0056 f"slurm_ten_simultaneous_jobs_{datetime.now(UTC).strftime('%Y%m%d_%H%M%S')}"
0057 )
0058 run_dir.mkdir(parents=True, exist_ok=True)
0059
0060 scheduler = make_scheduler(run_dir)
0061 job_definitions: List[Dict[str, Any]] = []
0062 expected_outputs: Dict[str, Path] = {}
0063
0064 for index in range(10):
0065 design_point = build_dtlz2_point(index)
0066 job_name = f"dtlz2_design_{index}"
0067 design_path = run_dir / "designs" / f"{job_name}.json"
0068 output_path = run_dir / "results" / f"{job_name}.json"
0069 Trial(index=index, parameters=design_point, status="pending").save_to_json(design_path)
0070
0071 job_definitions.append(
0072 {
0073 "name": job_name,
0074 "command": (
0075 "python examples/dtlz2/dtlz2_evaluator.py "
0076 f"--design {design_path} "
0077 f"--output {output_path} "
0078 f"--label {job_name} "
0079 "--sleep-sec 0.6 "
0080 "--repeat 3"
0081 ),
0082 "payload": {"execution_dir": str(run_dir / "runtime" / job_name)},
0083 "outputs": [{"path": str(output_path), "format": "json"}],
0084 "resources": {"time": "00:05:00"},
0085 }
0086 )
0087 expected_outputs[job_name] = output_path
0088
0089 stage_id = scheduler.submit_stage(
0090 "ten_dtlz2_designs",
0091 job_definitions,
0092 parallelism_policy={"poll_interval": 1},
0093 working_dir=str(run_dir / "scheduler"),
0094 )
0095
0096 history: List[Dict[str, Any]] = []
0097 while True:
0098 stage_status = scheduler.check_stage_status(stage_id)
0099 counts = Counter(job.status for job in stage_status.job_statuses or [])
0100 snapshot = {
0101 "timestamp_utc": datetime.now(UTC).isoformat(),
0102 "stage_status": stage_status.status,
0103 "completed_jobs": stage_status.completed_jobs,
0104 "total_jobs": stage_status.total_jobs,
0105 "status_counts": dict(counts),
0106 }
0107 history.append(snapshot)
0108 print(f"[ten-jobs] {snapshot}")
0109 if stage_status.status in {"completed", "failed", "cancelled"}:
0110 break
0111 time.sleep(1)
0112
0113 result = scheduler.get_stage_results(stage_id)
0114 if not result.success:
0115 raise RuntimeError(f"Expected all 10 jobs to succeed, got: {result.error_message}")
0116 if len(result.job_statuses) != 10:
0117 raise RuntimeError(f"Expected 10 job statuses, got {len(result.job_statuses)}")
0118 if any(job.status != "completed" for job in result.job_statuses):
0119 raise RuntimeError("Expected all ten jobs to complete successfully")
0120
0121 per_job_outputs = {}
0122 for job_status in result.job_statuses:
0123 if not job_status.outputs or "f1" not in job_status.outputs or "f2" not in job_status.outputs:
0124 raise RuntimeError(f"Missing objective outputs for job {job_status.job_id}: {job_status.outputs}")
0125 per_job_outputs[job_status.job_id] = job_status.outputs
0126
0127 for job_name, output_path in expected_outputs.items():
0128 if not output_path.exists():
0129 raise RuntimeError(f"Expected output file missing for {job_name}: {output_path}")
0130
0131 summary = {
0132 "stage_id": stage_id,
0133 "history": history,
0134 "jobs": [
0135 {
0136 "job_id": job_status.job_id,
0137 "status": job_status.status,
0138 "return_code": job_status.return_code,
0139 "slurm_job_id": (job_status.metrics or {}).get("slurm_job_id"),
0140 "outputs": job_status.outputs,
0141 }
0142 for job_status in result.job_statuses
0143 ],
0144 "design_files": sorted(str(path) for path in (run_dir / "designs").glob("*.json")),
0145 "result_files": sorted(str(path) for path in (run_dir / "results").glob("*.json")),
0146 }
0147 summary_path = run_dir / "run_slurm_ten_simultaneous_jobs.summary.json"
0148 summary_path.write_text(json.dumps(summary, indent=2), encoding="utf-8")
0149 print(json.dumps(summary, indent=2))
0150 return 0
0151
0152
0153 if __name__ == "__main__":
0154 raise SystemExit(main())
