File indexing completed on 2026-08-12 08:24:56
0001 """Scheduler configuration models.
0002
0003 Unified scheduler configuration for AID2E workflow execution.
0004 Runner-specific parameters are validated via the scheduler registry.
0005
0006 Supports multiple execution backends:
0007 - JobLibRunner: Local parallel execution using joblib
0008 - SlurmRunner: HPC cluster execution via SLURM
0009 - PanDAiDDSRunner: Distributed execution via PanDA iDDS
0010
0011 Project: AID2E v0.0.0 - AI assisted Detector Design for EIC
0012 Homepage: https://aid2e.github.io/aid2e
0013 Repository: https://github.com/aid2e/AID2E-framework.git
0014 """
0015
0016 import json
0017 from pathlib import Path
0018 from typing import Dict, Optional, Any, Literal
0019 import yaml
0020 from pydantic import BaseModel, Field, ConfigDict
0021 from .scheduler_registry import get
0022
0023
0024 class SchedulerConfiguration(BaseModel):
0025 """Complete scheduler/runner configuration.
0026
0027 Specifies which scheduler backend to use and its parameters.
0028 Runner-specific configuration is validated via the scheduler registry.
0029
0030 Attributes:
0031 runner_type: Type of runner/scheduler to use (JobLibRunner, SlurmRunner, PanDAiDDSRunner).
0032 parameters: Runner-specific parameters as free-form dict.
0033 max_retries: Global maximum retries for failed jobs.
0034 output_location: Base directory for scheduler output files.
0035 monitor_interval: Monitoring interval in seconds for job status checks.
0036
0037 Example:
0038 >>> config = SchedulerConfiguration(
0039 ... runner_type="JobLibRunner",
0040 ... parameters={"n_jobs": -1, "backend": "threading"},
0041 ... output_location="./output"
0042 ... )
0043 """
0044
0045 model_config = ConfigDict(extra="forbid")
0046
0047 runner_type: Literal["JobLibRunner", "SlurmRunner", "PanDAiDDSRunner"] = Field(
0048 default="JobLibRunner",
0049 description="Type of runner/scheduler to use"
0050 )
0051
0052 parameters: Dict[str, Any] = Field(
0053 default_factory=dict,
0054 description="Runner-specific parameters (validated by scheduler registry)"
0055 )
0056
0057 max_retries: int = Field(
0058 default=3,
0059 ge=0,
0060 description="Global maximum retries for failed jobs"
0061 )
0062 output_location: str = Field(
0063 default="./scheduler_output",
0064 description="Base directory for scheduler output files"
0065 )
0066 monitor_interval: int = Field(
0067 default=30,
0068 ge=1,
0069 description="Monitoring interval in seconds for job status checks"
0070 )
0071
0072 def parse_runner_params(self) -> Optional[BaseModel]:
0073 """Parse and validate runner-specific parameters via registry.
0074
0075 Looks up the registered config model for this runner_type and
0076 validates the parameters dict against it.
0077
0078 Returns:
0079 Validated runner-specific config model instance, or None if
0080 runner type not found in registry.
0081
0082 Raises:
0083 ValidationError: If parameters don't match the runner's schema.
0084
0085 Example:
0086 >>> config = SchedulerConfiguration(
0087 ... runner_type="JobLibRunner",
0088 ... parameters={"n_jobs": 4}
0089 ... )
0090 >>> joblib_config = config.parse_runner_params()
0091 >>> joblib_config.n_jobs
0092 4
0093 """
0094 Model = get(self.runner_type)
0095 if Model:
0096 return Model(**self.parameters)
0097 return None
0098
0099
0100 class SchedulerConfigLoader:
0101 """Loader for scheduler YAML/CONFIG files.
0102
0103 Parses files following the scheduler schema:
0104
0105 scheduler:
0106 runner_type: "JobLibRunner"
0107 parameters:
0108 n_jobs: 4
0109
0110 Notes:
0111 Use `SchedulerConfigLoader.load()` to load from a file path or
0112 `SchedulerConfigLoader.from_dict()` to construct from an in-memory dictionary.
0113 """
0114
0115 @staticmethod
0116 def _build_from_scheduler_dict(
0117 scheduler_payload: Dict[str, Any],
0118 base_dir: Optional[Path] = None,
0119 ) -> SchedulerConfiguration:
0120 """Build SchedulerConfiguration from an inner ``scheduler`` mapping."""
0121 if not isinstance(scheduler_payload, dict):
0122 raise ValueError("Invalid scheduler definition: expected a mapping")
0123 scheduler_payload = dict(scheduler_payload)
0124 required_keys = ["runner_type", "parameters"]
0125 missing = [key for key in required_keys if key not in scheduler_payload]
0126 if missing:
0127 raise ValueError("Invalid scheduler definition, missing keys: " + ", ".join(missing))
0128 config = SchedulerConfiguration(**scheduler_payload)
0129 parameters = dict(config.parameters)
0130 if config.runner_type == "SlurmRunner":
0131 template_file = parameters.get("template_file")
0132 if template_file is None:
0133 if not parameters:
0134 raise ValueError("Invalid SlurmRunner scheduler parameters, provide inline definitions or template_file")
0135 else:
0136 template_path = Path(template_file).expanduser()
0137 if base_dir and not template_path.is_absolute():
0138 template_path = (base_dir / template_path).resolve()
0139 if not template_path.exists():
0140 raise FileNotFoundError(f"Slurm template file not found: {template_file}")
0141
0142 with open(template_path, "r") as f:
0143 template_data = json.load(f)
0144 if not isinstance(template_data, dict):
0145 raise ValueError("Invalid Slurm template file: expected a JSON object")
0146
0147 inline_parameters = dict(parameters)
0148 inline_parameters.pop("template_file")
0149 if not template_data and not inline_parameters:
0150 raise ValueError("Invalid Slurm template file: expected scheduler parameters")
0151 scheduler_payload["parameters"] = {
0152 **template_data,
0153 **inline_parameters,
0154 }
0155 config = SchedulerConfiguration(**scheduler_payload)
0156 parameters = dict(config.parameters)
0157 if config.runner_type == "PanDAiDDSRunner":
0158 if not parameters:
0159 raise ValueError("Invalid PanDAiDDSRunner scheduler parameters, provide PanDA definitions")
0160
0161 Model = get(config.runner_type)
0162 if Model is None:
0163 raise ValueError(f"No scheduler config model registered for {config.runner_type}")
0164 unknown_keys = sorted(set(parameters) - set(Model.model_fields))
0165 if unknown_keys:
0166 raise ValueError(
0167 f"Invalid {config.runner_type} scheduler parameters, unknown keys: "
0168 + ", ".join(unknown_keys)
0169 )
0170 Model(**parameters)
0171
0172 return config
0173
0174 @staticmethod
0175 def load(file_path: str) -> SchedulerConfiguration:
0176 """Load a scheduler configuration from a YAML file."""
0177 path = Path(file_path)
0178 if not path.exists():
0179 raise FileNotFoundError(f"Scheduler file not found: {file_path}")
0180
0181 with open(path, "r") as f:
0182 data = yaml.safe_load(f) or {}
0183
0184 if "scheduler" not in data or not isinstance(data["scheduler"], dict):
0185 raise ValueError("Invalid scheduler file: missing 'scheduler' section")
0186
0187 return SchedulerConfigLoader._build_from_scheduler_dict(
0188 data["scheduler"],
0189 base_dir=path.parent,
0190 )
0191
0192 @staticmethod
0193 def from_dict(
0194 scheduler_payload: Dict[str, Any],
0195 base_dir: Optional[str] = None,
0196 ) -> SchedulerConfiguration:
0197 """Construct SchedulerConfiguration from an inner scheduler mapping."""
0198 return SchedulerConfigLoader._build_from_scheduler_dict(
0199 scheduler_payload,
0200 base_dir=Path(base_dir) if base_dir else None,
0201 )
