Coverage for src/srunx/models.py: 79%

1"""Data models for SLURM job management."""

2

3 import os

4 import subprocess

5 import time

6 from enum import Enum

7 from pathlib import Path

8 from typing import Self

9

10 import jinja2

11 from pydantic import BaseModel, Field, model_validator

12 from srunx.logging import get_logger

13

14 logger = get_logger(__name__)

15

16

17 class JobStatus(Enum):

18 """SLURM job status enumeration."""

19

20 PENDING = "PENDING"

21 RUNNING = "RUNNING"

22 COMPLETED = "COMPLETED"

23 FAILED = "FAILED"

24 CANCELLED = "CANCELLED"

25 TIMEOUT = "TIMEOUT"

26

27

28 class JobResource(BaseModel):

29 """SLURM resource allocation requirements."""

30

31 nodes: int = Field(default=1, ge=1, description="Number of compute nodes")

32 gpus_per_node: int = Field(default=0, ge=0, description="Number of GPUs per node")

33 ntasks_per_node: int = Field(

34 default=1, ge=1, description="Number of tasks per node"

35 )

36 cpus_per_task: int = Field(default=1, ge=1, description="Number of CPUs per task")

37 memory_per_node: str | None = Field(

38 default=None, description="Memory per node (e.g., '32GB')"

39 )

40 time_limit: str | None = Field(

41 default=None, description="Time limit (e.g., '1:00:00')"

42 )

43

44

45 class JobEnvironment(BaseModel):

46 """Job environment configuration."""

47

48 conda: str | None = Field(default=None, description="Conda environment name")

49 venv: str | None = Field(default=None, description="Virtual environment path")

50 sqsh: str | None = Field(default=None, description="SquashFS image path")

51 env_vars: dict[str, str] = Field(

52 default_factory=dict, description="Environment variables"

53 )

54

55 @model_validator(mode="after")

56 def validate_environment(self) -> Self:

57 envs = [self.conda, self.venv, self.sqsh]

58 non_none_count = sum(x is not None for x in envs)

59 if non_none_count != 1:

60 raise ValueError("Exactly one of 'conda', 'venv', or 'sqsh' must be set")

61 return self

62

63

64 class BaseJob(BaseModel):

65 name: str = Field(default="job", description="Job name")

66 job_id: int | None = Field(default=None, description="SLURM job ID")

67 status: JobStatus | None = Field(default=None, description="Current job status")

68 dependencies: list[str] = Field(

69 default_factory=list, description="Job dependencies"

70 )

71

72 def refresh(self, retries: int = 3) -> None:

73 for retry in range(retries):

74 try:

75 result = subprocess.run(

76 [

77 "sacct",

78 "-j",

79 str(self.job_id),

80 "--format",

81 "JobID,JobName,State",

82 "--noheader",

83 "--parsable2",

84 ],

85 capture_output=True,

86 text=True,

87 check=True,

88 )

89 except subprocess.CalledProcessError as e:

90 logger.error(f"Failed to query job {self.job_id} status: {e}")

91 raise

92

93 lines = result.stdout.strip().split("\n")

94 if not lines or not lines[0]:

95 error_msg = f"No job information found for job {self.job_id}"

96 if retry < retries - 1:

97 error_msg += f" Retrying {retry + 1} of {retries}..."

98 logger.error(error_msg)

99 time.sleep(1)

100 continue

101 raise ValueError(error_msg)

102 else:

103 break

104

105 # Parse the first line (main job entry)

106 job_data = lines[0].split("|")

107 if len(job_data) < 3:

108 error_msg = f"Cannot parse job data for job {self.job_id}"

109 logger.error(error_msg)

110 raise ValueError(error_msg)

111

112 status = job_data[2]

113 self.status = JobStatus(status)

114

115

116 class Job(BaseJob):

117 """Represents a SLURM job with complete configuration."""

118

119 command: list[str] = Field(description="Command to execute")

120 resources: JobResource = Field(

121 default_factory=JobResource, description="Resource requirements"

122 )

123 environment: JobEnvironment = Field(

124 default_factory=JobEnvironment, description="Environment setup"

125 )

126 log_dir: str = Field(

127 default=os.getenv("SLURM_LOG_DIR", "logs"),

128 description="Directory for log files",

129 )

130 work_dir: str = Field(default_factory=os.getcwd, description="Working directory")

131

132

133 class ShellJob(BaseJob):

134 path: str = Field(description="Shell script path to execute")

135

136

137 class WorkflowTask(BaseModel):

138 """Represents a single task in a workflow."""

139

140 name: str = Field(description="Task name")

141 job: BaseJob = Field(description="Job configuration")

142 depends_on: list[str] = Field(default_factory=list, description="Task dependencies")

143 async_execution: bool = Field(

144 default=False, description="Whether to run asynchronously"

145 )

146

147

148 class Workflow(BaseModel):

149 """Represents a workflow containing multiple tasks with dependencies."""

150

151 name: str = Field(description="Workflow name")

152 tasks: list[WorkflowTask] = Field(description="List of tasks in the workflow")

153

154 def get_task(self, name: str) -> WorkflowTask | None:

155 """Get a task by name."""

156 for task in self.tasks:

157 if task.name == name:

158 return task

159 return None

160

161 def get_task_dependencies(self, task_name: str) -> list[str]:

162 """Get dependencies for a specific task."""

163 task = self.get_task(task_name)

164 return task.depends_on if task else []

165

166

167 def render_job_script(

168 template_path: Path | str,

169 job: Job,

170 output_dir: Path | str,

171) -> str:

172 """Render a SLURM job script from a template.

173

174 Args:

175 template_path: Path to the Jinja template file.

176 job: Job configuration.

177 output_dir: Directory where the generated script will be saved.

178

179 Returns:

180 Path to the generated SLURM batch script.

181

182 Raises:

183 FileNotFoundError: If the template file does not exist.

184 jinja2.TemplateError: If template rendering fails.

185 """

186 template_file = Path(template_path)

187 if not template_file.is_file():

188 raise FileNotFoundError(f"Template file '{template_path}' not found")

189

190 with open(template_file, encoding="utf-8") as f:

191 template_content = f.read()

192

193 template = jinja2.Template(template_content, undefined=jinja2.StrictUndefined)

194

195 # Prepare template variables

196 template_vars = {

197 "job_name": job.name,

198 "command": " ".join(job.command or []),

199 "log_dir": job.log_dir,

200 "work_dir": job.work_dir,

201 "environment_setup": _build_environment_setup(job.environment),

202 **job.resources.model_dump(),

203 }

204

205 rendered_content = template.render(template_vars)

206

207 # Generate output file

208 output_path = Path(output_dir) / f"{job.name}.slurm"

209 with open(output_path, "w", encoding="utf-8") as f:

210 f.write(rendered_content)

211

212 return str(output_path)

213

214

215 def _build_environment_setup(environment: JobEnvironment) -> str:

216 """Build environment setup script."""

217 setup_lines = []

218

219 # Set environment variables

220 for key, value in environment.env_vars.items():

221 setup_lines.append(f"export {key}={value}")

222

223 # Activate environments

224 if environment.conda:

225 setup_lines.extend(["conda deactivate", f"conda activate {environment.conda}"])

226 elif environment.venv:

227 setup_lines.append(f"source {environment.venv}/bin/activate")

228 elif environment.sqsh:

229 setup_lines.extend(

230 [

231 f': "${{IMAGE:={environment.sqsh}}}"',

232 "declare -a CONTAINER_ARGS=(",

233 ' --container-image "$IMAGE"',

234 ")",

235 ]

236 )

237

238 return "\n".join(setup_lines)