Coverage for src/srunx/client.py: 87%

1"""SLURM client for job submission and management."""

2

3 import subprocess

4 import tempfile

5 import time

6 from importlib.resources import files

7 from pathlib import Path

8

9 from srunx.logging import get_logger

10 from srunx.models import (

11 BaseJob,

12 Job,

13 JobStatus,

14 ShellJob,

15 render_job_script,

16)

17 from srunx.utils import get_job_status

18

19 logger = get_logger(__name__)

20

21

22 class Slurm:

23 """Client for interacting with SLURM workload manager."""

24

25 def __init__(self, default_template: str | None = None):

26 """Initialize SLURM client.

27

28 Args:

29 default_template: Path to default job template.

30 """

31 self.default_template = default_template or self._get_default_template()

32

33 def run(

34 self, job: Job | ShellJob, template_path: str | None = None

35 ) -> Job | ShellJob:

36 """Submit a job to SLURM.

37

38 Args:

39 job: Job configuration.

40 template_path: Optional template path (uses default if not provided).

41

42 Returns:

43 Job instance with updated job_id and status.

44

45 Raises:

46 subprocess.CalledProcessError: If job submission fails.

47 """

48

49 if isinstance(job, Job):

50 template = template_path or self.default_template

51

52 with tempfile.TemporaryDirectory() as temp_dir:

53 script_path = render_job_script(template, job, temp_dir)

54 logger.debug(f"Generated SLURM script at: {script_path}")

55

56 # Handle container execution

57 sbatch_cmd = ["sbatch"]

58 if job.environment.sqsh:

59 sbatch_cmd.extend(["--sqsh", job.environment.sqsh])

60 logger.debug(f"Using sqsh: {job.environment.sqsh}")

61

62 sbatch_cmd.append(script_path)

63 logger.debug(f"Executing command: {' '.join(sbatch_cmd)}")

64

65 try:

66 result = subprocess.run(

67 sbatch_cmd,

68 capture_output=True,

69 text=True,

70 check=True,

71 )

72 except subprocess.CalledProcessError as e:

73 logger.error(f"Failed to submit job '{job.name}': {e}")

74 logger.error(f"Command: {' '.join(e.cmd)}")

75 logger.error(f"Return code: {e.returncode}")

76 logger.error(f"Stdout: {e.stdout}")

77 logger.error(f"Stderr: {e.stderr}")

78 raise

79

80 elif isinstance(job, ShellJob):

81 try:

82 result = subprocess.run(

83 ["sbatch", job.path],

84 capture_output=True,

85 text=True,

86 check=True,

87 )

88 except subprocess.CalledProcessError as e:

89 logger.error(f"Failed to submit job '{job.name}': {e}")

90 logger.error(f"Command: {' '.join(e.cmd)}")

91 logger.error(f"Return code: {e.returncode}")

92 logger.error(f"Stdout: {e.stdout}")

93 logger.error(f"Stderr: {e.stderr}")

94 raise

95

96 else:

97 raise ValueError("Either 'command' or 'file' must be set")

98

99 time.sleep(3)

100 job_id = int(result.stdout.split()[-1])

101 job.job_id = job_id

102 job.status = JobStatus.PENDING

103

104 logger.info(f"Successfully submitted job '{job.name}' with ID {job_id}")

105 return job

106

107 @staticmethod

108 def retrieve(job_id: int) -> BaseJob:

109 return get_job_status(job_id)

110

111 def cancel(self, job_id: int) -> None:

112 """Cancel a SLURM job.

113

114 Args:

115 job_id: SLURM job ID to cancel.

116

117 Raises:

118 subprocess.CalledProcessError: If job cancellation fails.

119 """

120 logger.info(f"Cancelling job {job_id}")

121

122 try:

123 subprocess.run(

124 ["scancel", str(job_id)],

125 check=True,

126 )

127 logger.info(f"Successfully cancelled job {job_id}")

128 except subprocess.CalledProcessError as e:

129 logger.error(f"Failed to cancel job {job_id}: {e}")

130 raise

131

132 def list(self, user: str | None = None) -> list[BaseJob]:

133 """List jobs for a user.

134

135 Args:

136 user: Username (defaults to current user).

137

138 Returns:

139 List of Job objects.

140 """

141 cmd = [

142 "squeue",

143 "--format",

144 "%.18i %.9P %.15j %.8u %.8T %.10M %.9l %.6D %R",

145 "--noheader",

146 ]

147 if user:

148 cmd.extend(["--user", user])

149

150 result = subprocess.run(cmd, capture_output=True, text=True, check=True)

151

152 jobs = []

153 for line in result.stdout.strip().split("\n"):

154 if not line.strip():

155 continue

156

157 parts = line.split()

158 if len(parts) >= 5:

159 job_id = int(parts[0])

160 job_name = parts[2]

161 status_str = parts[4]

162

163 try:

164 status = JobStatus(status_str)

165 except ValueError:

166 status = JobStatus.PENDING # Default for unknown status

167

168 job = BaseJob(

169 name=job_name,

170 job_id=job_id,

171 status=status,

172 )

173 jobs.append(job)

174

175 return jobs

176

177 def monitor(

178 self, job_obj_or_id: BaseJob | Job | ShellJob | int, poll_interval: int = 30

179 ) -> BaseJob | Job | ShellJob:

180 """Wait for a job to complete.

181

182 Args:

183 job_obj_or_id: Job object or job ID.

184 poll_interval: Polling interval in seconds.

185

186 Returns:

187 Completed job object.

188

189 Raises:

190 RuntimeError: If job fails.

191 """

192 if isinstance(job_obj_or_id, int):

193 job = self.retrieve(job_obj_or_id)

194 else:

195 job = job_obj_or_id

196

197 msg = f"Waiting for job {job.job_id} to complete (polling every {poll_interval}s)."

198 if isinstance(job, Job):

199 msg += f" Logging to {job.log_dir}/{job.name}_{job.job_id}.out"

200 logger.info(msg)

201

202 previous_status = None

203

204 while True:

205 job.refresh()

206

207 # Log status changes

208 if job.status != previous_status:

209 status_str = job.status.value if job.status else "Unknown"

210 logger.info(f"Job {job.job_id} status: {status_str}")

211 previous_status = job.status

212

213 match job.status:

214 case JobStatus.COMPLETED:

215 logger.info(f"Job {job.job_id} completed successfully")

216 return job

217 case JobStatus.FAILED:

218 err_msg = f"SLURM job {job.job_id} failed.\n"

219 if isinstance(job, Job):

220 log_file = Path(job.log_dir) / f"{job.name}_{job.job_id}.out"

221 if log_file.exists():

222 with open(log_file) as f:

223 err_msg += f.read()

224 err_msg += f"\nLog file: {log_file}"

225 else:

226 err_msg += f"Log file not found: {log_file}"

227 raise RuntimeError(err_msg)

228 case JobStatus.CANCELLED | JobStatus.TIMEOUT:

229 err_msg = (

230 f"SLURM job {job.job_id} was {job.status.value.lower()}.\n"

231 )

232 if isinstance(job, Job):

233 log_file = Path(job.log_dir) / f"{job.name}_{job.job_id}.out"

234 if log_file.exists():

235 with open(log_file) as f:

236 err_msg += f.read()

237 err_msg += f"\nLog file: {log_file}"

238 else:

239 err_msg += f"Log file not found: {log_file}"

240 raise RuntimeError(err_msg)

241 time.sleep(poll_interval)

242

243 def _get_default_template(self) -> str:

244 """Get the default job template path."""

245 from srunx import templates

246

247 return str(files(templates).joinpath("base.slurm.jinja"))

248

249

250# Convenience functions for backward compatibility

251 def submit_job(job: Job | ShellJob, template_path: str | None = None) -> Job | ShellJob:

252 """Submit a job to SLURM (convenience function)."""

253 client = Slurm()

254 return client.run(job, template_path)

255

256

257 def retrieve_job(job_id: int) -> BaseJob:

258 """Get job status (convenience function)."""

259 client = Slurm()

260 return client.retrieve(job_id)

261

262

263 def cancel_job(job_id: int) -> None:

264 """Cancel a job (convenience function)."""

265 client = Slurm()

266 client.cancel(job_id)