File indexing completed on 2026-08-12 08:24:55
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[2]
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 }
0039
0040
0041 def make_scheduler(run_dir: Path) -> SlurmScheduler:
0042 return SlurmScheduler(
0043 config=SlurmRunnerConfig(
0044 job_name_prefix="aid2e_multi10",
0045 setup_commands=list(DEFAULT_SETUP_COMMANDS),
0046 submit_working_dir=str(run_dir),
0047 runtime_working_dir=str(REPO_ROOT),
0048 poll_interval=2,
0049 )
0050 )
0051
0052
0053 def main() -> int:
0054 run_dir = REPO_ROOT / "experimental_tests" / "output" / (
0055 f"slurm_ten_simultaneous_jobs_{datetime.now(UTC).strftime('%Y%m%d_%H%M%S')}"
0056 )
0057 run_dir.mkdir(parents=True, exist_ok=True)
0058
0059 scheduler = make_scheduler(run_dir)
0060 job_definitions: List[Dict[str, Any]] = []
0061 expected_outputs: Dict[str, Path] = {}
0062
0063 for index in range(10):
0064 design_point = build_dtlz2_point(index)
0065 job_name = f"dtlz2_design_{index}"
0066 design_path = run_dir / "designs" / f"{job_name}.json"
0067 output_path = run_dir / "results" / f"{job_name}.json"
0068 Trial(index=index, parameters=design_point, status="pending").save_to_json(design_path)
0069
0070 job_definitions.append(
0071 {
0072 "name": job_name,
0073 "command": (
0074 "python examples/schedulers/slurm_dtlz2_from_design_json.py "
0075 f"--design {design_path} "
0076 f"--output {output_path} "
0077 f"--label {job_name} "
0078 "--sleep-sec 0.6 "
0079 "--repeat 3"
0080 ),
0081 "payload": {"execution_dir": str(run_dir / "runtime" / job_name)},
0082 "outputs": [{"path": str(output_path), "format": "json"}],
0083 "resources": {"time": "00:05:00"},
0084 }
0085 )
0086 expected_outputs[job_name] = output_path
0087
0088 stage_id = scheduler.submit_stage(
0089 "ten_dtlz2_designs",
0090 job_definitions,
0091 parallelism_policy={"poll_interval": 1},
0092 working_dir=str(run_dir / "scheduler"),
0093 )
0094
0095 history: List[Dict[str, Any]] = []
0096 while True:
0097 stage_status = scheduler.check_stage_status(stage_id)
0098 counts = Counter(job.status for job in stage_status.job_statuses or [])
0099 snapshot = {
0100 "timestamp_utc": datetime.now(UTC).isoformat(),
0101 "stage_status": stage_status.status,
0102 "completed_jobs": stage_status.completed_jobs,
0103 "total_jobs": stage_status.total_jobs,
0104 "status_counts": dict(counts),
0105 }
0106 history.append(snapshot)
0107 print(f"[ten-jobs] {snapshot}")
0108 if stage_status.status in {"completed", "failed", "cancelled"}:
0109 break
0110 time.sleep(1)
0111
0112 result = scheduler.get_stage_results(stage_id)
0113 if not result.success:
0114 raise RuntimeError(f"Expected all 10 jobs to succeed, got: {result.error_message}")
0115 if len(result.job_statuses) != 10:
0116 raise RuntimeError(f"Expected 10 job statuses, got {len(result.job_statuses)}")
0117 if any(job.status != "completed" for job in result.job_statuses):
0118 raise RuntimeError("Expected all ten jobs to complete successfully")
0119
0120 per_job_outputs = {}
0121 for job_status in result.job_statuses:
0122 if not job_status.outputs or "f1" not in job_status.outputs or "f2" not in job_status.outputs:
0123 raise RuntimeError(f"Missing objective outputs for job {job_status.job_id}: {job_status.outputs}")
0124 per_job_outputs[job_status.job_id] = job_status.outputs
0125
0126 for job_name, output_path in expected_outputs.items():
0127 if not output_path.exists():
0128 raise RuntimeError(f"Expected output file missing for {job_name}: {output_path}")
0129
0130 summary = {
0131 "stage_id": stage_id,
0132 "history": history,
0133 "jobs": [
0134 {
0135 "job_id": job_status.job_id,
0136 "status": job_status.status,
0137 "return_code": job_status.return_code,
0138 "slurm_job_id": (job_status.metrics or {}).get("slurm_job_id"),
0139 "outputs": job_status.outputs,
0140 }
0141 for job_status in result.job_statuses
0142 ],
0143 "design_files": sorted(str(path) for path in (run_dir / "designs").glob("*.json")),
0144 "result_files": sorted(str(path) for path in (run_dir / "results").glob("*.json")),
0145 }
0146 summary_path = run_dir / "run_slurm_ten_simultaneous_jobs.summary.json"
0147 summary_path.write_text(json.dumps(summary, indent=2), encoding="utf-8")
0148 print(json.dumps(summary, indent=2))
0149 return 0
0150
0151
0152 if __name__ == "__main__":
0153 raise SystemExit(main())
