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

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, PrivateAttr, model_validator

12

13 from srunx.exceptions import WorkflowValidationError

14 from srunx.logging import get_logger

15

16 logger = get_logger(__name__)

17

18

19 def _get_config_defaults():

20 """Get configuration defaults, with lazy import to avoid circular dependencies."""

21 try:

22 from srunx.config import get_config

23

24 return get_config()

25 except ImportError:

26 # Fallback if config module is not available

27 return None

28

29

30 def _default_nodes():

31 """Get default nodes from config."""

32 config = _get_config_defaults()

33 return config.resources.nodes if config else 1

34

35

36 def _default_gpus_per_node():

37 """Get default GPUs per node from config."""

38 config = _get_config_defaults()

39 return config.resources.gpus_per_node if config else 0

40

41

42 def _default_ntasks_per_node():

43 """Get default ntasks per node from config."""

44 config = _get_config_defaults()

45 return config.resources.ntasks_per_node if config else 1

46

47

48 def _default_cpus_per_task():

49 """Get default CPUs per task from config."""

50 config = _get_config_defaults()

51 return config.resources.cpus_per_task if config else 1

52

53

54 def _default_memory_per_node():

55 """Get default memory per node from config."""

56 config = _get_config_defaults()

57 return config.resources.memory_per_node if config else None

58

59

60 def _default_time_limit():

61 """Get default time limit from config."""

62 config = _get_config_defaults()

63 return config.resources.time_limit if config else None

64

65

66 def _default_nodelist():

67 """Get default nodelist from config."""

68 config = _get_config_defaults()

69 return config.resources.nodelist if config else None

70

71

72 def _default_partition():

73 """Get default partition from config."""

74 config = _get_config_defaults()

75 return config.resources.partition if config else None

76

77

78 def _default_conda():

79 """Get default conda environment from config."""

80 config = _get_config_defaults()

81 return config.environment.conda if config else None

82

83

84 def _default_venv():

85 """Get default venv path from config."""

86 config = _get_config_defaults()

87 return config.environment.venv if config else None

88

89

90 def _default_sqsh():

91 """Get default sqsh path from config."""

92 config = _get_config_defaults()

93 return config.environment.sqsh if config else None

94

95

96 def _default_env_vars():

97 """Get default environment variables from config."""

98 config = _get_config_defaults()

99 return config.environment.env_vars if config else {}

100

101

102 def _default_log_dir():

103 """Get default log directory from config."""

104 config = _get_config_defaults()

105 return config.log_dir if config else os.getenv("SLURM_LOG_DIR", "logs")

106

107

108 def _default_work_dir():

109 """Get default work directory from config."""

110 config = _get_config_defaults()

111 return config.work_dir if config else None

112

113

114 class JobStatus(Enum):

115 """Job status enumeration for both SLURM jobs and workflow jobs."""

116

117 UNKNOWN = "UNKNOWN"

118 PENDING = "PENDING"

119 RUNNING = "RUNNING"

120 COMPLETED = "COMPLETED"

121 FAILED = "FAILED"

122 CANCELLED = "CANCELLED"

123 TIMEOUT = "TIMEOUT"

124

125

126 class JobResource(BaseModel):

127 """SLURM resource allocation requirements."""

128

129 nodes: int = Field(

130 default_factory=_default_nodes, ge=1, description="Number of compute nodes"

131 )

132 gpus_per_node: int = Field(

133 default_factory=_default_gpus_per_node,

134 ge=0,

135 description="Number of GPUs per node",

136 )

137 ntasks_per_node: int = Field(

138 default_factory=_default_ntasks_per_node,

139 ge=1,

140 description="Number of jobs per node",

141 )

142 cpus_per_task: int = Field(

143 default_factory=_default_cpus_per_task,

144 ge=1,

145 description="Number of CPUs per task",

146 )

147 memory_per_node: str | None = Field(

148 default_factory=_default_memory_per_node,

149 description="Memory per node (e.g., '32GB')",

150 )

151 time_limit: str | None = Field(

152 default_factory=_default_time_limit, description="Time limit (e.g., '1:00:00')"

153 )

154 nodelist: str | None = Field(

155 default_factory=_default_nodelist,

156 description="Specific nodes to use (e.g., 'node001,node002')",

157 )

158 partition: str | None = Field(

159 default_factory=_default_partition,

160 description="SLURM partition to use (e.g., 'gpu', 'cpu')",

161 )

162

163

164 class JobEnvironment(BaseModel):

165 """Job environment configuration."""

166

167 conda: str | None = Field(

168 default_factory=_default_conda, description="Conda environment name"

169 )

170 venv: str | None = Field(

171 default_factory=_default_venv, description="Virtual environment path"

172 )

173 sqsh: str | None = Field(

174 default_factory=_default_sqsh, description="SquashFS image path"

175 )

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

177 default_factory=_default_env_vars, description="Environment variables"

178 )

179

180 @model_validator(mode="after")

181 def validate_environment(self) -> Self:

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

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

184 if non_none_count != 1:

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

186 return self

187

188

189 class BaseJob(BaseModel):

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

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

192 depends_on: list[str] = Field(

193 default_factory=list, description="Task dependencies for workflow execution"

194 )

195

196 _status: JobStatus = PrivateAttr(default=JobStatus.PENDING)

197

198 @property

199 def status(self) -> JobStatus:

200 """

201 Accessing ``job.status`` always triggers a lightweight refresh

202 (only if we have a ``job_id`` and the status isn't terminal).

203 """

204 if self.job_id is not None and self._status not in {

205 JobStatus.COMPLETED,

206 JobStatus.FAILED,

207 JobStatus.CANCELLED,

208 JobStatus.TIMEOUT,

209 }:

210 self.refresh()

211 return self._status

212

213 @status.setter

214 def status(self, value: JobStatus) -> None:

215 self._status = value

216

217 def refresh(self, retries: int = 3) -> Self:

218 """Query sacct and update ``_status`` in-place."""

219 if self.job_id is None:

220 return self

221

222 for retry in range(retries):

223 try:

224 result = subprocess.run(

225 [

226 "sacct",

227 "-j",

228 str(self.job_id),

229 "--format",

230 "JobID,State",

231 "--noheader",

232 "--parsable2",

233 ],

234 capture_output=True,

235 text=True,

236 check=True,

237 )

238 except subprocess.CalledProcessError as e:

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

240 raise

241

242 line = result.stdout.strip().split("\n")[0] if result.stdout.strip() else ""

243 if not line:

244 if retry < retries - 1:

245 time.sleep(1)

246 continue

247 self._status = JobStatus.UNKNOWN

248 return self

249 break

250

251 _, state = line.split("|", 1)

252 self._status = JobStatus(state)

253 return self

254

255 def dependencies_satisfied(self, completed_job_names: list[str]) -> bool:

256 """All dependencies are completed & this job is still pending."""

257 return self.status == JobStatus.PENDING and all(

258 dep in completed_job_names for dep in self.depends_on

259 )

260

261

262 class Job(BaseJob):

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

264

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

266 resources: JobResource = Field(

267 default_factory=JobResource, description="Resource requirements"

268 )

269 environment: JobEnvironment = Field(

270 default_factory=JobEnvironment, description="Environment setup"

271 )

272 log_dir: str = Field(

273 default_factory=_default_log_dir,

274 description="Directory for log files",

275 )

276 work_dir: str = Field(

277 default_factory=lambda: _default_work_dir() or os.getcwd(),

278 description="Working directory",

279 )

280

281

282 class ShellJob(BaseJob):

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

284

285

286 type JobType = BaseJob | Job | ShellJob

287 type RunnableJobType = Job | ShellJob

288

289

290 class Workflow:

291 """Represents a workflow containing multiple jobs with dependencies."""

292

293 def __init__(self, name: str, jobs: list[RunnableJobType] | None = None) -> None:

294 if jobs is None:

295 jobs = []

296

297 self.name = name

298 self.jobs = jobs

299

300 def add(self, job: RunnableJobType) -> None:

301 # Check if job already exists

302 if job.depends_on:

303 for dep in job.depends_on:

304 if dep not in self.jobs:

305 raise WorkflowValidationError(

306 f"Job '{job.name}' depends on unknown job '{dep}'"

307 )

308 self.jobs.append(job)

309

310 def remove(self, job: RunnableJobType) -> None:

311 self.jobs.remove(job)

312

313 def get(self, name: str) -> RunnableJobType | None:

314 """Get a job by name."""

315 for job in self.jobs:

316 if job.name == name:

317 return job.refresh()

318 return None

319

320 def get_dependencies(self, job_name: str) -> list[str]:

321 """Get dependencies for a specific job."""

322 job = self.get(job_name)

323 return job.depends_on if job else []

324

325 def show(self):

326 msg = f"""\

327{" PLAN ":=^80}

328 Workflow: {self.name}

329 Jobs: {len(self.jobs)}

330"""

331

332 def add_indent(indent: int, msg: str) -> str:

333 return " " * indent + msg

334

335 for job in self.jobs:

336 msg += add_indent(1, f"Job: {job.name}\n")

337 if isinstance(job, Job):

338 msg += add_indent(

339 2, f"{'Command:': <13} {' '.join(job.command or [])}\n"

340 )

341 msg += add_indent(

342 2,

343 f"{'Resources:': <13} {job.resources.nodes} nodes, {job.resources.gpus_per_node} GPUs/node\n",

344 )

345 if job.environment.conda:

346 msg += add_indent(

347 2, f"{'Conda env:': <13} {job.environment.conda}\n"

348 )

349 if job.environment.sqsh:

350 msg += add_indent(2, f"{'Sqsh:': <13} {job.environment.sqsh}\n")

351 if job.environment.venv:

352 msg += add_indent(2, f"{'Venv:': <13} {job.environment.venv}\n")

353 elif isinstance(job, ShellJob):

354 msg += add_indent(2, f"{'Path:': <13} {job.path}\n")

355 if job.depends_on:

356 msg += add_indent(

357 2, f"{'Dependencies:': <13} {', '.join(job.depends_on)}\n"

358 )

359

360 msg += f"{'=' * 80}\n"

361 print(msg)

362

363 def validate(self):

364 """Validate workflow job dependencies."""

365 job_names = {job.name for job in self.jobs}

366

367 if len(job_names) != len(self.jobs):

368 raise WorkflowValidationError("Duplicate job names found in workflow")

369

370 for job in self.jobs:

371 for dependency in job.depends_on:

372 if dependency not in job_names:

373 raise WorkflowValidationError(

374 f"Job '{job.name}' depends on unknown job '{dependency}'"

375 )

376

377 # Check for circular dependencies (simple check)

378 visited = set()

379 rec_stack = set()

380

381 def has_cycle(job_name: str) -> bool:

382 if job_name in rec_stack:

383 return True

384 if job_name in visited:

385 return False

386

387 visited.add(job_name)

388 rec_stack.add(job_name)

389

390 job = self.get(job_name)

391 if job:

392 for dependency in job.depends_on:

393 if has_cycle(dependency):

394 return True

395

396 rec_stack.remove(job_name)

397 return False

398

399 for job in self.jobs:

400 if has_cycle(job.name):

401 raise WorkflowValidationError(

402 f"Circular dependency detected involving job '{job.name}'"

403 )

404

405

406 def render_job_script(

407 template_path: Path | str,

408 job: Job,

409 output_dir: Path | str,

410 verbose: bool = False,

411) -> str:

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

413

414 Args:

415 template_path: Path to the Jinja template file.

416 job: Job configuration.

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

418 verbose: Whether to print the rendered content.

419

420 Returns:

421 Path to the generated SLURM batch script.

422

423 Raises:

424 FileNotFoundError: If the template file does not exist.

425 jinja2.TemplateError: If template rendering fails.

426 """

427 template_file = Path(template_path)

428 if not template_file.is_file():

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

430

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

432 template_content = f.read()

433

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

435

436 # Prepare template variables

437 template_vars = {

438 "job_name": job.name,

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

440 "log_dir": job.log_dir,

441 "work_dir": job.work_dir,

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

443 **job.resources.model_dump(),

444 }

445

446 rendered_content = template.render(template_vars)

447

448 if verbose:

449 print(rendered_content)

450

451 # Generate output file

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

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

454 f.write(rendered_content)

455

456 return str(output_path)

457

458

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

460 """Build environment setup script."""

461 setup_lines = []

462

463 # Set environment variables

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

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

466

467 # Activate environments

468 if environment.conda:

469 home_dir = Path.home()

470 setup_lines.extend(

471 [

472 f"source {str(home_dir)}/miniconda3/bin/activate",

473 "conda deactivate",

474 f"conda activate {environment.conda}",

475 ]

476 )

477 elif environment.venv:

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

479 elif environment.sqsh:

480 setup_lines.extend(

481 [

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

483 "declare -a CONTAINER_ARGS=(",

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

485 ")",

486 ]

487 )

488

489 return "\n".join(setup_lines)