File indexing completed on 2026-08-12 08:24:56
0001 """Slurm-backed scheduler for command jobs."""
0002
0003 from __future__ import annotations
0004
0005 import json
0006 import logging
0007 import shlex
0008 import subprocess
0009 import time
0010 import uuid
0011 from pathlib import Path
0012 from typing import Any, Dict, List, Optional
0013
0014 from aid2e.schedulers.Slurm.config import SlurmRunnerConfig
0015 from aid2e.schedulers.base import BaseScheduler, JobStatus, StageExecutionResult, StageStatus
0016
0017
0018 TERMINAL_STATUSES = {"completed", "failed", "cancelled"}
0019
0020 STATE_MAP = {
0021 "PENDING": "queued",
0022 "CONFIGURING": "queued",
0023 "RUNNING": "running",
0024 "COMPLETING": "running",
0025 "COMPLETED": "completed",
0026 "FAILED": "failed",
0027 "TIMEOUT": "failed",
0028 "OUT_OF_MEMORY": "failed",
0029 "NODE_FAIL": "failed",
0030 "PREEMPTED": "failed",
0031 "BOOT_FAIL": "failed",
0032 "DEADLINE": "failed",
0033 "REVOKED": "failed",
0034 "CANCELLED": "cancelled",
0035 }
0036
0037 RESOURCE_KEYS = (
0038 "partition",
0039 "account",
0040 "qos",
0041 "time",
0042 "nodes",
0043 "ntasks",
0044 "cpus_per_task",
0045 "mem",
0046 "gres",
0047 "constraint",
0048 )
0049
0050
0051 class SlurmScheduler(BaseScheduler):
0052 """Execute workflow stage jobs on Slurm using generated batch scripts."""
0053
0054 def __init__(self, config: Optional[SlurmRunnerConfig] = None) -> None:
0055 super().__init__(config)
0056 self.config = config or SlurmRunnerConfig()
0057 self.logger = logging.getLogger("SlurmScheduler")
0058 self.stages: Dict[str, Dict[str, Any]] = {}
0059 self.jobs: Dict[str, Dict[str, Any]] = {}
0060
0061 def run_stage(
0062 self,
0063 stage_name: str,
0064 job_definitions: List[Dict[str, Any]],
0065 parallelism_policy: Optional[Dict[str, Any]] = None,
0066 working_dir: Optional[str] = None,
0067 ) -> StageExecutionResult:
0068 """Submit a stage and block until it reaches a terminal state."""
0069
0070 stage_id = self.submit_stage(stage_name, job_definitions, parallelism_policy, working_dir)
0071 poll_interval = int((parallelism_policy or {}).get("poll_interval", self.config.poll_interval))
0072
0073 while True:
0074 status = self.check_stage_status(stage_id)
0075 if status.status in TERMINAL_STATUSES:
0076 return self.get_stage_results(stage_id)
0077 time.sleep(poll_interval)
0078
0079 def submit_stage(
0080 self,
0081 stage_name: str,
0082 job_definitions: List[Dict[str, Any]],
0083 parallelism_policy: Optional[Dict[str, Any]] = None,
0084 working_dir: Optional[str] = None,
0085 ) -> str:
0086 """Create scripts, submit them to Slurm, and return a stage id."""
0087
0088 stage_id = uuid.uuid4().hex
0089 stage_root = self._resolve_submit_root(working_dir) / f"{stage_name}_{stage_id}"
0090 stage_root.mkdir(parents=True, exist_ok=True)
0091
0092 stage_state: Dict[str, Any] = {
0093 "stage_id": stage_id,
0094 "stage_name": stage_name,
0095 "status": "queued",
0096 "created_at": time.time(),
0097 "parallelism_policy": dict(parallelism_policy or {}),
0098 "working_dir": str(stage_root),
0099 "job_ids": [],
0100 "result": None,
0101 }
0102 self.stages[stage_id] = stage_state
0103
0104 for index, job_def in enumerate(job_definitions):
0105 self._validate_job_definition(job_def)
0106 job_name = job_def.get("name", f"job_{index}")
0107 job_id = job_def.get("job_id") or f"{stage_name}_{job_name}_{index}"
0108 job_root = stage_root / job_id
0109 job_root.mkdir(parents=True, exist_ok=True)
0110
0111 script_path = job_root / "job.sbatch"
0112 runtime_dir = self._resolve_runtime_dir(job_def, job_root)
0113 output_dir = self._resolve_output_dir(job_def, job_root)
0114 runtime_dir.mkdir(parents=True, exist_ok=True)
0115 output_dir.mkdir(parents=True, exist_ok=True)
0116 stdout_path = output_dir / "stdout.log" if self.config.capture_stdout else None
0117 stderr_path = output_dir / "stderr.log" if self.config.capture_stderr else None
0118
0119 script_text = self._build_batch_script(
0120 job_id=job_id,
0121 job_name=job_name,
0122 job_def=job_def,
0123 runtime_dir=runtime_dir,
0124 stdout_path=stdout_path,
0125 stderr_path=stderr_path,
0126 )
0127 script_path.write_text(script_text, encoding="utf-8")
0128
0129 slurm_job_id = self._submit_script(script_path, stage_root)
0130 job_state = {
0131 "stage_id": stage_id,
0132 "job_id": job_id,
0133 "job_name": job_name,
0134 "job_root": str(job_root),
0135 "runtime_dir": str(runtime_dir),
0136 "output_dir": str(output_dir),
0137 "script_path": str(script_path),
0138 "stdout_path": str(stdout_path) if stdout_path else None,
0139 "stderr_path": str(stderr_path) if stderr_path else None,
0140 "slurm_job_id": slurm_job_id,
0141 "status": "queued",
0142 "raw_state": "PENDING",
0143 "return_code": None,
0144 "command": job_def.get("command", ""),
0145 "outputs": list(job_def.get("outputs", [])),
0146 "artifact_cache": None,
0147 "last_update": time.time(),
0148 }
0149 self.jobs[job_id] = job_state
0150 stage_state["job_ids"].append(job_id)
0151
0152 self._refresh_stage_status(stage_id)
0153 return stage_id
0154
0155 def check_stage_status(self, stage_id: str) -> StageStatus:
0156 """Poll Slurm and return a stage-level status summary."""
0157
0158 if stage_id not in self.stages:
0159 raise KeyError(f"Unknown stage_id: {stage_id}")
0160
0161 self._refresh_stage_status(stage_id)
0162 stage_state = self.stages[stage_id]
0163 job_statuses = [self._job_status_from_state(self.jobs[job_id]) for job_id in stage_state["job_ids"]]
0164 completed_jobs = sum(1 for status in job_statuses if status.status in TERMINAL_STATUSES)
0165 total_jobs = len(job_statuses)
0166 progress = float(completed_jobs) / float(total_jobs) if total_jobs else 1.0
0167
0168 return StageStatus(
0169 stage_id=stage_id,
0170 status=stage_state["status"],
0171 completed_jobs=completed_jobs,
0172 total_jobs=total_jobs,
0173 progress=progress,
0174 job_statuses=job_statuses,
0175 )
0176
0177 def get_stage_results(self, stage_id: str) -> StageExecutionResult:
0178 """Return final stage results after collecting logs and artifacts."""
0179
0180 if stage_id not in self.stages:
0181 raise KeyError(f"Unknown stage_id: {stage_id}")
0182
0183 self._refresh_stage_status(stage_id)
0184 stage_state = self.stages[stage_id]
0185 if stage_state["status"] not in TERMINAL_STATUSES:
0186 raise RuntimeError(f"Stage {stage_id} is not yet finished (status={stage_state['status']})")
0187
0188 cached = stage_state.get("result")
0189 if cached is not None:
0190 return cached
0191
0192 job_statuses: List[JobStatus] = []
0193 artifacts: Dict[str, Any] = {}
0194 stage_success = True
0195
0196 for job_id in stage_state["job_ids"]:
0197 job_state = self.jobs[job_id]
0198 job_outputs = self._collect_job_outputs(job_state)
0199 stdout = self._read_text(job_state.get("stdout_path"))
0200 stderr = self._read_text(job_state.get("stderr_path"))
0201 status = self._job_status_from_state(
0202 job_state,
0203 stdout=stdout,
0204 stderr=stderr,
0205 outputs=job_outputs,
0206 )
0207 job_statuses.append(status)
0208 if job_outputs:
0209 artifacts.update(job_outputs)
0210 if status.status != "completed":
0211 stage_success = False
0212
0213 result = StageExecutionResult(
0214 stage_name=stage_state["stage_name"],
0215 job_statuses=job_statuses,
0216 artifacts=artifacts,
0217 success=stage_success,
0218 error_message=None if stage_success else f"Some jobs failed in stage '{stage_state['stage_name']}'",
0219 )
0220 stage_state["result"] = result
0221 return result
0222
0223 def check_status(self, job_id: str) -> JobStatus:
0224 """Return the most recent cached job status."""
0225
0226 if job_id not in self.jobs:
0227 raise KeyError(f"Unknown job_id: {job_id}")
0228 self._refresh_job_state(self.jobs[job_id])
0229 return self._job_status_from_state(self.jobs[job_id])
0230
0231 def cancel_job(self, job_id: str) -> bool:
0232 """Request job cancellation through scancel."""
0233
0234 if job_id not in self.jobs:
0235 self.logger.warning("Cannot cancel unknown job %s", job_id)
0236 return False
0237
0238 job_state = self.jobs[job_id]
0239 slurm_job_id = job_state["slurm_job_id"]
0240 proc = subprocess.run(
0241 ["scancel", str(slurm_job_id)],
0242 capture_output=True,
0243 text=True,
0244 check=False,
0245 )
0246 if proc.returncode != 0:
0247 self.logger.warning("scancel failed for job %s (%s): %s", job_id, slurm_job_id, proc.stderr.strip())
0248 return False
0249
0250 job_state["status"] = "cancelled"
0251 job_state["raw_state"] = "CANCELLED"
0252 job_state["return_code"] = -1
0253 self._refresh_stage_status(job_state["stage_id"])
0254 return True
0255
0256 def _validate_job_definition(self, job_def: Dict[str, Any]) -> None:
0257 if job_def.get("function") is not None or job_def.get("params") is not None:
0258 raise ValueError("SlurmScheduler v1 supports command jobs only")
0259 command = str(job_def.get("command", "")).strip()
0260 if not command:
0261 raise ValueError("Job definition must include a non-empty 'command'")
0262
0263 def _resolve_submit_root(self, working_dir: Optional[str]) -> Path:
0264 submit_root = self.config.submit_working_dir or working_dir or str(Path.cwd())
0265 return Path(submit_root).expanduser().resolve()
0266
0267 def _resolve_runtime_dir(self, job_def: Dict[str, Any], default_job_root: Path) -> Path:
0268 if self.config.runtime_working_dir:
0269 return Path(self.config.runtime_working_dir).expanduser().resolve()
0270
0271 payload = job_def.get("payload") or {}
0272 execution_dir = payload.get("execution_dir")
0273 if execution_dir:
0274 return Path(str(execution_dir)).expanduser().resolve()
0275
0276 return default_job_root.resolve()
0277
0278 def _resolve_output_dir(self, job_def: Dict[str, Any], default_job_root: Path) -> Path:
0279 payload = job_def.get("payload") or {}
0280 output_dir = payload.get("output_dir")
0281 if output_dir:
0282 return Path(str(output_dir)).expanduser().resolve()
0283 return default_job_root.resolve()
0284
0285 def _resolve_job_resources(self, job_def: Dict[str, Any]) -> Dict[str, Any]:
0286 resolved: Dict[str, Any] = {key: getattr(self.config, key) for key in RESOURCE_KEYS}
0287 for key, value in (job_def.get("resources") or {}).items():
0288 if key in RESOURCE_KEYS and value is not None:
0289 resolved[key] = value
0290 return resolved
0291
0292 def _build_batch_script(
0293 self,
0294 job_id: str,
0295 job_name: str,
0296 job_def: Dict[str, Any],
0297 runtime_dir: Path,
0298 stdout_path: Optional[Path],
0299 stderr_path: Optional[Path],
0300 ) -> str:
0301 resources = self._resolve_job_resources(job_def)
0302 slurm_job_name = f"{self.config.job_name_prefix}_{job_name}"
0303 lines = ["#!/bin/bash", f"#SBATCH --job-name={slurm_job_name}"]
0304
0305 if stdout_path is not None:
0306 lines.append(f"#SBATCH --output={stdout_path}")
0307 if stderr_path is not None:
0308 lines.append(f"#SBATCH --error={stderr_path}")
0309
0310 directive_map = {
0311 "partition": "--partition",
0312 "account": "--account",
0313 "qos": "--qos",
0314 "time": "--time",
0315 "nodes": "--nodes",
0316 "ntasks": "--ntasks",
0317 "cpus_per_task": "--cpus-per-task",
0318 "mem": "--mem",
0319 "gres": "--gres",
0320 "constraint": "--constraint",
0321 }
0322 for key, flag in directive_map.items():
0323 value = resources.get(key)
0324 if value is not None:
0325 lines.append(f"#SBATCH {flag}={value}")
0326
0327 lines.extend(["", "set -euo pipefail"])
0328 lines.extend(self.config.setup_commands)
0329 lines.append(f"cd {shlex.quote(str(runtime_dir))}")
0330 lines.append(str(job_def["command"]))
0331 lines.append("")
0332 return "\n".join(lines)
0333
0334 def _submit_script(self, script_path: Path, submit_root: Path) -> str:
0335 command = ["sbatch", "--parsable", *self.config.sbatch_extra_args, str(script_path)]
0336 proc = subprocess.run(
0337 command,
0338 cwd=str(submit_root),
0339 capture_output=True,
0340 text=True,
0341 check=False,
0342 )
0343 if proc.returncode != 0:
0344 raise RuntimeError(
0345 f"sbatch failed for {script_path}: {proc.stderr.strip() or proc.stdout.strip()}"
0346 )
0347
0348 stdout = proc.stdout.strip()
0349 job_id = stdout.split(";", 1)[0].strip()
0350 if not job_id:
0351 raise RuntimeError(f"Could not parse Slurm job id from sbatch output: {stdout!r}")
0352 return job_id
0353
0354 def _refresh_stage_status(self, stage_id: str) -> None:
0355 stage_state = self.stages[stage_id]
0356 statuses = []
0357 for job_id in stage_state["job_ids"]:
0358 self._refresh_job_state(self.jobs[job_id])
0359 statuses.append(self.jobs[job_id]["status"])
0360
0361 if not statuses:
0362 stage_state["status"] = "completed"
0363 elif all(status == "completed" for status in statuses):
0364 stage_state["status"] = "completed"
0365 elif any(status == "failed" for status in statuses):
0366 stage_state["status"] = "failed"
0367 elif any(status == "cancelled" for status in statuses):
0368 stage_state["status"] = "cancelled"
0369 elif any(status == "running" for status in statuses):
0370 stage_state["status"] = "running"
0371 else:
0372 stage_state["status"] = "queued"
0373
0374 def _refresh_job_state(self, job_state: Dict[str, Any]) -> None:
0375 if job_state["status"] in TERMINAL_STATUSES:
0376 return
0377
0378 slurm_job_id = job_state["slurm_job_id"]
0379 active_state = self._query_squeue_state(slurm_job_id)
0380 if active_state is not None:
0381 job_state["raw_state"] = active_state
0382 job_state["status"] = self._normalize_state(active_state)
0383 job_state["last_update"] = time.time()
0384 return
0385
0386 account_state = self._query_sacct_state(slurm_job_id)
0387 if account_state is None:
0388 job_state["last_update"] = time.time()
0389 return
0390
0391 raw_state = account_state["state"]
0392 job_state["raw_state"] = raw_state
0393 job_state["status"] = self._normalize_state(raw_state)
0394 job_state["return_code"] = self._parse_exit_code(account_state.get("exit_code"), job_state["status"])
0395 job_state["last_update"] = time.time()
0396
0397 def _query_squeue_state(self, slurm_job_id: str) -> Optional[str]:
0398 proc = subprocess.run(
0399 ["squeue", "-h", "-j", str(slurm_job_id), "--format=%i|%T"],
0400 capture_output=True,
0401 text=True,
0402 check=False,
0403 )
0404 if proc.returncode != 0:
0405 self.logger.debug("squeue failed for %s: %s", slurm_job_id, proc.stderr.strip())
0406 return None
0407
0408 for line in proc.stdout.splitlines():
0409 line = line.strip()
0410 if not line:
0411 continue
0412 job_id, state = (line.split("|", 1) + [""])[:2]
0413 if job_id.strip() == str(slurm_job_id):
0414 return state.strip()
0415 return None
0416
0417 def _query_sacct_state(self, slurm_job_id: str) -> Optional[Dict[str, str]]:
0418 proc = subprocess.run(
0419 ["sacct", "-n", "-P", "-j", str(slurm_job_id), "--format=JobIDRaw,State,ExitCode"],
0420 capture_output=True,
0421 text=True,
0422 check=False,
0423 )
0424 if proc.returncode != 0:
0425 self.logger.debug("sacct failed for %s: %s", slurm_job_id, proc.stderr.strip())
0426 return None
0427
0428 best_match: Optional[Dict[str, str]] = None
0429 for line in proc.stdout.splitlines():
0430 line = line.strip()
0431 if not line:
0432 continue
0433 parts = line.split("|")
0434 if len(parts) < 3:
0435 continue
0436 job_id_raw, state, exit_code = parts[0].strip(), parts[1].strip(), parts[2].strip()
0437 if job_id_raw == str(slurm_job_id):
0438 best_match = {"state": state, "exit_code": exit_code}
0439 break
0440 return best_match
0441
0442 def _normalize_state(self, raw_state: Optional[str]) -> str:
0443 if not raw_state:
0444 return "unknown"
0445 token = raw_state.strip().upper().split()[0].rstrip("+")
0446 return STATE_MAP.get(token, "unknown")
0447
0448 def _parse_exit_code(self, exit_code: Optional[str], normalized_status: str) -> Optional[int]:
0449 if normalized_status == "completed":
0450 return 0
0451 if not exit_code:
0452 return -1 if normalized_status in {"failed", "cancelled"} else None
0453 token = exit_code.split(":", 1)[0].strip()
0454 try:
0455 return int(token)
0456 except ValueError:
0457 return -1 if normalized_status in {"failed", "cancelled"} else None
0458
0459 def _job_status_from_state(
0460 self,
0461 job_state: Dict[str, Any],
0462 stdout: Optional[str] = None,
0463 stderr: Optional[str] = None,
0464 outputs: Optional[Dict[str, Any]] = None,
0465 ) -> JobStatus:
0466 return JobStatus(
0467 job_id=job_state["job_id"],
0468 status=job_state["status"],
0469 return_code=job_state.get("return_code"),
0470 stdout=stdout,
0471 stderr=stderr,
0472 outputs=outputs,
0473 metrics={"slurm_job_id": job_state["slurm_job_id"], "raw_state": job_state.get("raw_state")},
0474 )
0475
0476 def _collect_job_outputs(self, job_state: Dict[str, Any]) -> Dict[str, Any]:
0477 cached = job_state.get("artifact_cache")
0478 if cached is not None:
0479 return cached
0480
0481 output_dir = Path(job_state.get("output_dir") or job_state["runtime_dir"])
0482 collected: Dict[str, Any] = {}
0483 for output_spec in job_state.get("outputs", []):
0484 output_path = self._get_output_spec_value(output_spec, "path")
0485 if not output_path:
0486 continue
0487
0488 path_obj = Path(str(output_path))
0489 full_path = path_obj if path_obj.is_absolute() else output_dir / path_obj
0490 if not full_path.exists():
0491 self.logger.warning("Expected output artifact missing for %s: %s", job_state["job_id"], full_path)
0492 continue
0493
0494 fmt = str(self._get_output_spec_value(output_spec, "format") or "").lower()
0495 if fmt == "json":
0496 with full_path.open("r", encoding="utf-8") as handle:
0497 collected.update(json.load(handle))
0498 else:
0499 collected[str(output_path)] = full_path.read_text(encoding="utf-8")
0500
0501 job_state["artifact_cache"] = collected
0502 return collected
0503
0504 def _get_output_spec_value(self, output_spec: Any, key: str) -> Any:
0505 if isinstance(output_spec, dict):
0506 return output_spec.get(key)
0507 return getattr(output_spec, key, None)
0508
0509 def _read_text(self, path_str: Optional[str]) -> Optional[str]:
0510 if not path_str:
0511 return None
0512 path = Path(path_str)
0513 if not path.exists():
0514 return None
0515 return path.read_text(encoding="utf-8")
