File indexing completed on 2026-08-12 08:24:56
0001 """Pydantic configuration model for the Slurm scheduler."""
0002
0003 from typing import List, Optional
0004
0005 from pydantic import BaseModel, Field
0006
0007 from aid2e.utilities.configurations.scheduler_registry import register as register_runner_config
0008
0009
0010 class SlurmRunnerConfig(BaseModel):
0011 """Configuration for Slurm command-job execution."""
0012
0013 partition: Optional[str] = Field(default=None, description="Slurm partition to submit into.")
0014 account: Optional[str] = Field(default=None, description="Slurm account to charge.")
0015 qos: Optional[str] = Field(default=None, description="Slurm QoS value.")
0016 time: Optional[str] = Field(default=None, description="Walltime limit in Slurm format, e.g. 00:10:00.")
0017 nodes: int = Field(default=1, ge=1, description="Number of nodes requested.")
0018 ntasks: int = Field(default=1, ge=1, description="Number of tasks requested.")
0019 cpus_per_task: Optional[int] = Field(default=None, ge=1, description="CPUs per task.")
0020 mem: Optional[str] = Field(default=None, description="Memory request, e.g. 4G.")
0021 gres: Optional[str] = Field(default=None, description="Generic resources request, e.g. gpu:1.")
0022 constraint: Optional[str] = Field(default=None, description="Optional Slurm node constraint.")
0023 job_name_prefix: str = Field(default="aid2e", description="Prefix for generated Slurm job names.")
0024 setup_commands: List[str] = Field(
0025 default_factory=list,
0026 description="Commands emitted verbatim before the workload command.",
0027 )
0028 submit_working_dir: Optional[str] = Field(
0029 default=None,
0030 description="Directory where sbatch is invoked and batch scripts are written.",
0031 )
0032 runtime_working_dir: Optional[str] = Field(
0033 default=None,
0034 description="Directory the batch script should cd into before running the workload.",
0035 )
0036 poll_interval: int = Field(default=5, ge=1, description="Default stage polling interval in seconds.")
0037 sacct_poll_interval: int = Field(
0038 default=10,
0039 ge=1,
0040 description="Reserved poll interval for sacct checks in seconds.",
0041 )
0042 sbatch_extra_args: List[str] = Field(
0043 default_factory=list,
0044 description="Extra arguments appended to the sbatch command line.",
0045 )
0046 capture_stdout: bool = Field(default=True, description="Capture stdout into a scheduler-owned log file.")
0047 capture_stderr: bool = Field(default=True, description="Capture stderr into a scheduler-owned log file.")
0048
0049
0050 register_runner_config("SlurmRunner", SlurmRunnerConfig)
