Coverage for tests/test_batch_submit.py: 100%

1# This file is part of ctrl_bps. 

2# 

3# Developed for the LSST Data Management System. 

4# This product includes software developed by the LSST Project 

5# (https://www.lsst.org). 

6# See the COPYRIGHT file at the top-level directory of this distribution 

7# for details of code ownership. 

8# 

9# This software is dual licensed under the GNU General Public License and also 

10# under a 3-clause BSD license. Recipients may choose which of these licenses 

11# to use; please see the files gpl-3.0.txt and/or bsd_license.txt, 

12# respectively. If you choose the GPL option then the following text applies 

13# (but note that there is still no warranty even if you opt for BSD instead): 

14# 

15# This program is free software: you can redistribute it and/or modify 

16# it under the terms of the GNU General Public License as published by 

17# the Free Software Foundation, either version 3 of the License, or 

18# (at your option) any later version. 

19# 

20# This program is distributed in the hope that it will be useful, 

21# but WITHOUT ANY WARRANTY; without even the implied warranty of 

22# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 

23# GNU General Public License for more details. 

24# 

25# You should have received a copy of the GNU General Public License 

26# along with this program. If not, see <https://www.gnu.org/licenses/>. 

27"""Unit tests for batch_submit.py.""" 

28 

29import logging 

30import shutil 

31import tempfile 

32import unittest 

33from pathlib import Path 

34from unittest.mock import MagicMock, patch 

35 

36from lsst.ctrl.bps import BpsConfig, batch_submit 

37 

38 

39class TestCreateBatchStages(unittest.TestCase): 

40 """Tests for create_batch_stages function.""" 

41 

42 def setUp(self): 

43 self.common_config = { 

44 "bpsUseShared": True, 

45 "whenSaveJobQgraph": "NEVER", 

46 "useLazyCommands": True, 

47 "submitPath": "/the/path", 

48 } 

49 

50 def testMissingBuildCmd(self): 

51 """Missing buildQuantumGraph jobCommand""" 

52 config_info = dict(self.common_config) 

53 config_info.update({"uniqProcName": "uniq_proc_name"}) 

54 config = BpsConfig(config_info) 

55 with self.assertRaisesRegex( 

56 RuntimeError, "Missing executable for buildQuantumGraph. Double check submit yaml for jobCommand" 

57 ): 

58 _ = batch_submit.create_batch_stages(config, "not_used_prefix") 

59 

60 def testMissingPrepareCmd(self): 

61 """Missing preparePayloadWorkflow jobCommand""" 

62 config_info = dict(self.common_config) 

63 config_info.update( 

64 { 

65 "configFile": "not_used_configFile", 

66 "uniqProcName": "uniq_proc_name", 

67 "operator": "testuser", 

68 "payload": {"payloadName": "testPayload"}, 

69 "buildQuantumGraph": {"jobCommand": "${CTRL_BPS_DIR}/bin/bps batch-acquire {configFile}"}, 

70 } 

71 ) 

72 config = BpsConfig(config_info) 

73 with self.assertRaisesRegex( 

74 RuntimeError, 

75 "Missing executable for preparePayloadWorkflow. Double check submit yaml for jobCommand", 

76 ): 

77 _ = batch_submit.create_batch_stages(config, "not_used_prefix") 

78 

79 def testSuccess(self): 

80 # No saving of files 

81 config_info = dict(self.common_config) 

82 config_info.update( 

83 { 

84 "configFile": "not_used_configFile", 

85 "uniqProcName": "uniq_proc_name", 

86 "operator": "testuser", 

87 "payload": {"payloadName": "testPayload"}, 

88 "buildQuantumGraph": { 

89 "jobCommand": "${CTRL_BPS_DIR}/bin/bps batch-acquire {configFile}", 

90 "requestMemory": 16384, 

91 }, 

92 "preparePayloadWorkflow": { 

93 "jobCommand": "${CTRL_BPS_DIR}/bin/bps batch-prepare {configFile}", 

94 "requestMemory": 24576, 

95 }, 

96 } 

97 ) 

98 config = BpsConfig(config_info) 

99 

100 with tempfile.TemporaryDirectory() as tmpdir: 

101 gw, config = batch_submit.create_batch_stages(config, tmpdir) 

102 self.assertIn("buildQuantumGraph", gw) 

103 job = gw.get_job("buildQuantumGraph") 

104 self.assertIn("batch-acquire", job.arguments) 

105 self.assertEqual(job.request_memory, 16384) 

106 self.assertIn("preparePayloadWorkflow", gw) 

107 job = gw.get_job("preparePayloadWorkflow") 

108 self.assertIn("batch-prepare", job.arguments) 

109 self.assertEqual(job.request_memory, 24576) 

110 

111 # Check we didn't make any files 

112 self.assertEqual(list(Path(tmpdir).iterdir()), []) 

113 

114 def testSaving(self): 

115 config_info = dict(self.common_config) 

116 config_info.update( 

117 { 

118 "configFile": "not_used_configFile", 

119 "uniqProcName": "uniq_proc_name", 

120 "operator": "testuser", 

121 "payload": {"payloadName": "testPayload"}, 

122 "bpsPreCommandOpts": "--long-log --log-level=VERBOSE", 

123 "buildQuantumGraph": { 

124 "jobCommand": "${CTRL_BPS_DIR}/bin/bps batch-acquire {configFile}", 

125 "requestMemory": 16384, 

126 }, 

127 "preparePayloadWorkflow": { 

128 "jobCommand": "${CTRL_BPS_DIR}/bin/bps batch-prepare {configFile}", 

129 "requestMemory": 24576, 

130 }, 

131 "saveGenericWorkflow": True, 

132 } 

133 ) 

134 config = BpsConfig(config_info) 

135 with tempfile.TemporaryDirectory() as tmpdir: 

136 gw, config = batch_submit.create_batch_stages(config, tmpdir) 

137 self.assertTrue((Path(tmpdir) / "bps_stages_generic_workflow.pickle").exists()) 

138 

139 

140class TestBatchPayloadPrepare(unittest.TestCase): 

141 """Tests for batch_payload_prepare function.""" 

142 

143 def setUp(self): 

144 self.tmpdir = tempfile.mkdtemp() 

145 self.config_info = { 

146 "runQgraphFile": "run.qgraph", 

147 "uniqProcName": "uniq_proc_name", 

148 "computeSite": "site1", 

149 "qgraphFileTemplate": "template.qgraph", 

150 "bps_defined": {"submitPath": self.tmpdir}, 

151 } 

152 

153 def tearDown(self): 

154 shutil.rmtree(self.tmpdir, ignore_errors=True) 

155 

156 def _make_mocks(self, mock_cluster, mock_transform, mock_prepare): 

157 """Configure the standard set of dependency mocks.""" 

158 clustered_qgraph = MagicMock() 

159 clustered_qgraph.__len__.return_value = 3 

160 mock_cluster.return_value = clustered_qgraph 

161 

162 generic_workflow = MagicMock() 

163 generic_workflow.name = "test_workflow" 

164 generic_workflow.job_counts = {"label1": 5} 

165 gwfile = MagicMock() 

166 generic_workflow.get_file.return_value = gwfile 

167 generic_workflow_config = MagicMock() 

168 mock_transform.return_value = (generic_workflow, generic_workflow_config) 

169 

170 wms_workflow = MagicMock() 

171 mock_prepare.return_value = wms_workflow 

172 

173 return clustered_qgraph, generic_workflow, generic_workflow_config, gwfile, wms_workflow 

174 

175 @patch("lsst.ctrl.bps.batch_submit.prepare") 

176 @patch("lsst.ctrl.bps.batch_submit.transform") 

177 @patch("lsst.ctrl.bps.batch_submit.cluster_quanta") 

178 @patch("lsst.ctrl.bps.batch_submit.read_quantum_graph") 

179 def testSuccessBasic(self, mock_read, mock_cluster, mock_transform, mock_prepare): 

180 """Test success with all save flags off and no run temp space.""" 

181 _, generic_workflow, generic_workflow_config, gwfile, wms_workflow = self._make_mocks( 

182 mock_cluster, mock_transform, mock_prepare 

183 ) 

184 config = BpsConfig(self.config_info) 

185 

186 batch_submit.batch_payload_prepare(config, self.tmpdir) 

187 

188 mock_read.assert_called_once_with("run.qgraph") 

189 mock_cluster.assert_called_once() 

190 mock_transform.assert_called_once() 

191 mock_prepare.assert_called_once() 

192 # The runQgraphFile should be marked as not transferred by the WMS. 

193 self.assertFalse(gwfile.wms_transfer) 

194 # The payload workflow should be attached to the running workflow. 

195 wms_workflow.add_to_parent_workflow.assert_called_once_with(generic_workflow_config) 

196 # No files should be written with all save flags off. 

197 self.assertEqual(list(Path(self.tmpdir).iterdir()), []) 

198 

199 @patch("lsst.ctrl.bps.batch_submit.prepare") 

200 @patch("lsst.ctrl.bps.batch_submit.transform") 

201 @patch("lsst.ctrl.bps.batch_submit.cluster_quanta") 

202 @patch("lsst.ctrl.bps.batch_submit.read_quantum_graph") 

203 def testSaveClusteredQgraph(self, mock_read, mock_cluster, mock_transform, mock_prepare): 

204 """Test saving of the clustered quantum graph.""" 

205 clustered_qgraph, *_ = self._make_mocks(mock_cluster, mock_transform, mock_prepare) 

206 self.config_info["saveClusteredQgraph"] = True 

207 config = BpsConfig(self.config_info) 

208 

209 batch_submit.batch_payload_prepare(config, self.tmpdir) 

210 

211 clustered_qgraph.save.assert_called_once() 

212 self.assertIn("bps_clustered_qgraph.pickle", clustered_qgraph.save.call_args[0][0]) 

213 

214 @patch("lsst.ctrl.bps.batch_submit.prepare") 

215 @patch("lsst.ctrl.bps.batch_submit.transform") 

216 @patch("lsst.ctrl.bps.batch_submit.cluster_quanta") 

217 @patch("lsst.ctrl.bps.batch_submit.read_quantum_graph") 

218 def testSaveDotClustered(self, mock_read, mock_cluster, mock_transform, mock_prepare): 

219 """Test writing of the dot file.""" 

220 clustered_qgraph, *_ = self._make_mocks(mock_cluster, mock_transform, mock_prepare) 

221 self.config_info["saveDot"] = True 

222 config = BpsConfig(self.config_info) 

223 

224 batch_submit.batch_payload_prepare(config, self.tmpdir) 

225 

226 clustered_qgraph.draw.assert_called_once() 

227 self.assertIn("bps_clustered_qgraph.dot", clustered_qgraph.draw.call_args[0][0]) 

228 

229 @patch("lsst.ctrl.bps.batch_submit.prepare") 

230 @patch("lsst.ctrl.bps.batch_submit.transform") 

231 @patch("lsst.ctrl.bps.batch_submit.cluster_quanta") 

232 @patch("lsst.ctrl.bps.batch_submit.read_quantum_graph") 

233 def testSaveGenericWorkflow(self, mock_read, mock_cluster, mock_transform, mock_prepare): 

234 """Test writing of the GenericWorkflow to a file.""" 

235 _, generic_workflow, *_ = self._make_mocks(mock_cluster, mock_transform, mock_prepare) 

236 self.config_info["saveGenericWorkflow"] = True 

237 config = BpsConfig(self.config_info) 

238 

239 batch_submit.batch_payload_prepare(config, self.tmpdir) 

240 

241 generic_workflow.save.assert_called_once() 

242 self.assertTrue((Path(self.tmpdir) / "bps_generic_workflow.pickle").exists()) 

243 

244 @patch("lsst.ctrl.bps.batch_submit.prepare") 

245 @patch("lsst.ctrl.bps.batch_submit.transform") 

246 @patch("lsst.ctrl.bps.batch_submit.cluster_quanta") 

247 @patch("lsst.ctrl.bps.batch_submit.read_quantum_graph") 

248 def testSaveDotGeneric(self, mock_read, mock_cluster, mock_transform, mock_prepare): 

249 """Test saving the generic workflow dot file.""" 

250 _, generic_workflow, *_ = self._make_mocks(mock_cluster, mock_transform, mock_prepare) 

251 self.config_info["saveDot"] = True 

252 config = BpsConfig(self.config_info) 

253 

254 batch_submit.batch_payload_prepare(config, self.tmpdir) 

255 

256 generic_workflow.draw.assert_called_once() 

257 self.assertEqual(generic_workflow.draw.call_args[0][1], "dot") 

258 self.assertTrue((Path(self.tmpdir) / "bps_generic_workflow.dot").exists()) 

259 

260 @patch("lsst.ctrl.bps.batch_submit.prepare") 

261 @patch("lsst.ctrl.bps.batch_submit.transform") 

262 @patch("lsst.ctrl.bps.batch_submit.cluster_quanta") 

263 @patch("lsst.ctrl.bps.batch_submit.read_quantum_graph") 

264 def testUseRunTempSpaceFound(self, mock_read, mock_cluster, mock_transform, mock_prepare): 

265 """When run temp space is enabled and endpoint set, src_uri updates.""" 

266 _, _, _, gwfile, _ = self._make_mocks(mock_cluster, mock_transform, mock_prepare) 

267 self.config_info["bpsUseRunTempSpace"] = True 

268 self.config_info["fileDistributionEndpoint"] = "/run/temp/space" 

269 config = BpsConfig(self.config_info) 

270 

271 batch_submit.batch_payload_prepare(config, self.tmpdir) 

272 

273 self.assertEqual(gwfile.src_uri, str(Path("/run/temp/space") / "template.qgraph")) 

274 

275 @patch("lsst.ctrl.bps.batch_submit.prepare") 

276 @patch("lsst.ctrl.bps.batch_submit.transform") 

277 @patch("lsst.ctrl.bps.batch_submit.cluster_quanta") 

278 @patch("lsst.ctrl.bps.batch_submit.read_quantum_graph") 

279 def testUseRunTempSpaceMissingEndpoint(self, mock_read, mock_cluster, mock_transform, mock_prepare): 

280 """Run temp space enabled, missing endpoint should raise KeyError.""" 

281 self._make_mocks(mock_cluster, mock_transform, mock_prepare) 

282 self.config_info["bpsUseRunTempSpace"] = True 

283 config = BpsConfig(self.config_info) 

284 

285 with self.assertRaisesRegex(KeyError, "fileDistributionEndpoint"): 

286 batch_submit.batch_payload_prepare(config, self.tmpdir) 

287 

288 @patch("lsst.ctrl.bps.batch_submit.prepare") 

289 @patch("lsst.ctrl.bps.batch_submit.transform") 

290 @patch("lsst.ctrl.bps.batch_submit.cluster_quanta") 

291 @patch("lsst.ctrl.bps.batch_submit.read_quantum_graph") 

292 def testUseRunTempSpaceNotFound(self, mock_read, mock_cluster, mock_transform, mock_prepare): 

293 """When bpsUseRunTempSpace is absent, a debug message is logged.""" 

294 _, _, _, gwfile, _ = self._make_mocks(mock_cluster, mock_transform, mock_prepare) 

295 config = BpsConfig(self.config_info) 

296 

297 with self.assertLogs("lsst.ctrl.bps.batch_submit", level=logging.DEBUG) as cm: 

298 batch_submit.batch_payload_prepare(config, self.tmpdir) 

299 

300 self.assertTrue(any("missing bpsUseRunTempSpace" in msg for msg in cm.output)) 

301 

302 

303class TestBatchSubmit(unittest.TestCase): 

304 """Tests for batch_submit function.""" 

305 

306 def setUp(self): 

307 self.config_info = {"bps_defined": {"submitPath": "/the/path"}} 

308 

309 @patch("lsst.ctrl.bps.batch_submit._make_id_link") 

310 @patch("lsst.ctrl.bps.batch_submit.submit") 

311 @patch("lsst.ctrl.bps.batch_submit.prepare") 

312 @patch("lsst.ctrl.bps.batch_submit.create_batch_stages") 

313 def testSuccessSubmits(self, mock_create, mock_prepare, mock_submit, mock_make_id_link): 

314 """Without dryRun the control workflow is prepared and submitted.""" 

315 generic_workflow = MagicMock() 

316 config = BpsConfig(self.config_info) 

317 mock_create.return_value = (generic_workflow, config) 

318 wms_workflow = MagicMock() 

319 wms_workflow.run_id = "run123" 

320 mock_prepare.return_value = wms_workflow 

321 

322 result = batch_submit.batch_submit(config) 

323 

324 mock_create.assert_called_once() 

325 mock_prepare.assert_called_once() 

326 mock_submit.assert_called_once() 

327 mock_make_id_link.assert_called_once_with(config, "run123") 

328 self.assertIs(result, wms_workflow) 

329 

330 @patch("lsst.ctrl.bps.batch_submit._make_id_link") 

331 @patch("lsst.ctrl.bps.batch_submit.submit") 

332 @patch("lsst.ctrl.bps.batch_submit.prepare") 

333 @patch("lsst.ctrl.bps.batch_submit.create_batch_stages") 

334 def testDryRun(self, mock_create, mock_prepare, mock_submit, mock_make_id_link): 

335 """With dryRun the workflow is not submitted but still returned.""" 

336 generic_workflow = MagicMock() 

337 self.config_info["dryRun"] = True 

338 config = BpsConfig(self.config_info) 

339 mock_create.return_value = (generic_workflow, config) 

340 wms_workflow = MagicMock() 

341 wms_workflow.run_id = "run123" 

342 mock_prepare.return_value = wms_workflow 

343 

344 result = batch_submit.batch_submit(config) 

345 

346 mock_submit.assert_not_called() 

347 mock_make_id_link.assert_called_once_with(config, "run123") 

348 self.assertIs(result, wms_workflow) 

349 

350 

351if __name__ == "__main__": 

352 unittest.main()