Coverage for tests/test_transform.py: 96%

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 of transform.py.""" 

28 

29import dataclasses 

30import os 

31import shutil 

32import tempfile 

33import unittest 

34 

35from cqg_test_utils import make_test_clustered_quantum_graph 

36 

37from lsst.ctrl.bps import ( 

38 BPS_SEARCH_ORDER, 

39 BpsConfig, 

40 GenericWorkflow, 

41 GenericWorkflowExec, 

42 GenericWorkflowJob, 

43) 

44from lsst.ctrl.bps.transform import ( 

45 _enhance_command, 

46 _get_job_values, 

47 create_final_command, 

48 create_generic_workflow, 

49 create_generic_workflow_config, 

50) 

51 

52TESTDIR = os.path.abspath(os.path.dirname(__file__)) 

53 

54 

55class TestCreateGenericWorkflowConfig(unittest.TestCase): 

56 """Tests of create_generic_workflow_config.""" 

57 

58 def testCreate(self): 

59 """Test successful creation of the config.""" 

60 config = BpsConfig({"a": 1, "b": 2, "uniqProcName": "testCreate"}) 

61 wf_config = create_generic_workflow_config(config, "/test/create/prefix") 

62 self.assertIsInstance(wf_config, BpsConfig) 

63 for key in config: 

64 self.assertEqual(wf_config[key], config[key]) 

65 self.assertEqual(wf_config["workflowName"], "testCreate") 

66 self.assertEqual(wf_config["workflowPath"], "/test/create/prefix") 

67 

68 

69class TestCreateGenericWorkflow(unittest.TestCase): 

70 """Tests of create_generic_workflow.""" 

71 

72 def setUp(self): 

73 self.tmpdir = tempfile.mkdtemp(dir=TESTDIR) 

74 self.config = BpsConfig( 

75 { 

76 "runInit": True, 

77 "computeSite": "global", 

78 "runQuantumCommand": "gexe -q {qgraphFile} --qgraph-node-id {qgraphNodeId}", 

79 "clusterTemplate": "{D1}_{D2}", 

80 "cluster": { 

81 "cl1": {"pipetasks": "T1, T2", "dimensions": "D1, D2"}, 

82 "cl2": {"pipetasks": "T3, T4", "dimensions": "D1, D2"}, 

83 }, 

84 "cloud": { 

85 "cloud1": {"runQuantumCommand": "c1exe -q {qgraphFile} --qgraph-node-id {qgraphNodeId}"}, 

86 "cloud2": {"runQuantumCommand": "c2exe -q {qgraphFile} --qgraph-node-id {qgraphNodeId}"}, 

87 }, 

88 "site": { 

89 "site1": {"runQuantumCommand": "s1exe -q {qgraphFile} --qgraph-node-id {qgraphNodeId}"}, 

90 "site2": {"runQuantumCommand": "s2exe -q {qgraphFile} --qgraph-node-id {qgraphNodeId}"}, 

91 "global": {"runQuantumCommand": "s3exe -q {qgraphFile} --qgraph-node-id {qgraphNodeId}"}, 

92 }, 

93 # Needed because transform assumes they exist 

94 "whenSaveJobQgraph": "NEVER", 

95 "finalJob": {"whenRun": "ALWAYS", "command1": "/usr/bin/env"}, 

96 }, 

97 BPS_SEARCH_ORDER, 

98 ) 

99 _, self.cqg = make_test_clustered_quantum_graph(self.tmpdir) 

100 

101 def tearDown(self): 

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

103 

104 def testCreatingGenericWorkflowGlobal(self): 

105 """Test creating a GenericWorkflow with global settings.""" 

106 config = BpsConfig(self.config) 

107 config["computeCloud"] = "cloud1" 

108 config["computeSite"] = "site2" 

109 config["queue"] = "global_queue" 

110 print(config) 

111 workflow = create_generic_workflow(config, self.cqg, "test_gw", self.tmpdir) 

112 for jname in workflow: 

113 gwjob = workflow.get_job(jname) 

114 print(gwjob) 

115 self.assertEqual(gwjob.compute_site, "site2") 

116 self.assertEqual(gwjob.compute_cloud, "cloud1") 

117 self.assertEqual(gwjob.executable.src_uri, "s2exe") 

118 self.assertEqual(gwjob.queue, "global_queue") 

119 final = workflow.get_final() 

120 self.assertEqual(final.compute_site, "site2") 

121 self.assertEqual(final.compute_cloud, "cloud1") 

122 self.assertEqual(final.queue, "global_queue") 

123 

124 def testCreatingQuantumGraphMixed(self): 

125 """Test creating a GenericWorkflow with setting overrides.""" 

126 config = BpsConfig(self.config) 

127 config[".cluster.cl1.computeCloud"] = "cloud2" 

128 config[".cluster.cl1.computeSite"] = "notthere" 

129 config[".cluster.cl2.computeSite"] = "site1" 

130 config[".finalJob.queue"] = "special_final_queue" 

131 config[".finalJob.computeSite"] = "special_site" 

132 config[".finalJob.computeCloud"] = "special_cloud" 

133 workflow = create_generic_workflow(config, self.cqg, "test_gw", self.tmpdir) 

134 self.assertEqual(len(workflow) - 1, len(self.cqg)) # Don't count pipetaskInit 

135 for jname in workflow: 

136 gwjob = workflow.get_job(jname) 

137 print(gwjob) 

138 if jname.startswith("cl1"): 138 ↛ 139line 138 didn't jump to line 139 because the condition on line 138 was never true

139 self.assertEqual(gwjob.compute_site, "notthere") 

140 self.assertEqual(gwjob.compute_cloud, "cloud2") 

141 self.assertEqual(gwjob.executable.src_uri, "c2exe") 

142 elif jname.startswith("cl2"): 142 ↛ 143line 142 didn't jump to line 143 because the condition on line 142 was never true

143 self.assertEqual(gwjob.compute_site, "site1") 

144 self.assertIsNone(gwjob.compute_cloud) 

145 self.assertEqual(gwjob.executable.src_uri, "s1exe") 

146 elif jname.startswith("pipetask"): 

147 self.assertEqual(gwjob.compute_site, "global") 

148 self.assertIsNone(gwjob.compute_cloud) 

149 self.assertEqual(gwjob.executable.src_uri, "s3exe") 

150 final = workflow.get_final() 

151 self.assertEqual(final.compute_site, "special_site") 

152 self.assertEqual(final.compute_cloud, "special_cloud") 

153 self.assertEqual(final.queue, "special_final_queue") 

154 

155 

156class TestGetJobValues(unittest.TestCase): 

157 """Tests of _get_job_values.""" 

158 

159 def setUp(self): 

160 self.default_job = GenericWorkflowJob("default_job", "default_label") 

161 

162 def testGettingDefaults(self): 

163 """Test retrieving default values.""" 

164 config = BpsConfig({}) 

165 job_values = _get_job_values(config, {}, None) 

166 self.assertTrue( 

167 all( 

168 getattr(self.default_job, field.name) == job_values[field.name] 

169 for field in dataclasses.fields(self.default_job) 

170 ) 

171 ) 

172 

173 def testEnablingMemoryScaling(self): 

174 """Test enabling the memory scaling mechanism.""" 

175 config = BpsConfig({"memoryMultiplier": 2.0}) 

176 job_values = _get_job_values(config, {}, None) 

177 self.assertAlmostEqual(job_values["memory_multiplier"], 2.0) 

178 self.assertEqual(job_values["number_of_retries"], 5) 

179 

180 def testDisablingMemoryScaling(self): 

181 """Test disabling the memory scaling mechanism.""" 

182 config = BpsConfig({"memoryMultiplier": 0.5}) 

183 job_values = _get_job_values(config, {}, None) 

184 self.assertIsNone(job_values["memory_multiplier"]) 

185 

186 def testRetrievingCmdLine(self): 

187 """Test retrieving the command line.""" 

188 cmd_line_key = "runQuantum" 

189 config = BpsConfig({cmd_line_key: "/path/to/foo bar.txt"}) 

190 job_values = _get_job_values(config, {}, cmd_line_key) 

191 self.assertEqual(job_values["executable"].name, "foo") 

192 self.assertEqual(job_values["executable"].src_uri, "/path/to/foo") 

193 self.assertEqual(job_values["arguments"], "bar.txt") 

194 

195 def testEnvironment(self): 

196 config = BpsConfig( 

197 { 

198 "var1": "two", 

199 "environment": {"TEST_INT": 1, "TEST_BOOL": False, "TEST_SPACES": "one {var1} three"}, 

200 } 

201 ) 

202 job_values = _get_job_values(config, {}, None) 

203 truth = {"TEST_INT": "1", "TEST_BOOL": "False", "TEST_SPACES": "one two three"} 

204 self.assertEqual(truth, job_values["environment"]) 

205 

206 def testEnvironmentOptions(self): 

207 config = BpsConfig( 

208 { 

209 "var1": "two", 

210 "environment": {"TEST_INT": 1, "TEST_BOOL": False, "TEST_SPACES": "one {var1} three"}, 

211 "finalJob": {"requestMemory": 8096, "command1": "/usr/bin/env"}, 

212 } 

213 ) 

214 search_obj = config["finalJob"] 

215 search_opts = {"replaceVars": False, "searchobj": search_obj} 

216 job_values = _get_job_values(config, search_opts, None) 

217 truth = {"TEST_INT": "1", "TEST_BOOL": "False", "TEST_SPACES": "one two three"} 

218 self.assertEqual(truth, job_values["environment"]) 

219 self.assertEqual(search_opts["replaceVars"], False) 

220 self.assertEqual(search_opts["searchobj"]["requestMemory"], 8096) 

221 self.assertEqual(job_values["request_memory"], 8096) 

222 

223 def testVarsInEnvironment(self): 

224 config = BpsConfig( 

225 { 

226 "var1": "two", 

227 "environment": {"TEST_INT": 1, "TEST_BOOL": False, "TEST_SPACES": "one {var1} <ENV:var3>"}, 

228 } 

229 ) 

230 job_values = _get_job_values(config, {"replaceVars": True}, None) 

231 truth = {"TEST_INT": "1", "TEST_BOOL": "False", "TEST_SPACES": "one two <ENV:var3>"} 

232 self.assertEqual(truth, job_values["environment"]) 

233 

234 

235class TestCreateFinalCommand(unittest.TestCase): 

236 """Tests for the create_final_command function.""" 

237 

238 def setUp(self): 

239 self.tmpdir = tempfile.TemporaryDirectory() 

240 self.script_beginning = [ 

241 "#!/bin/bash\n", 

242 "\n", 

243 "set -e\n", 

244 "set -x\n", 

245 "qgraphFile=$1\n", 

246 "butlerConfig=$2\n", 

247 ] 

248 

249 def tearDown(self): 

250 self.tmpdir.cleanup() 

251 

252 def testSingleCommand(self): 

253 """Test with single final job command.""" 

254 config_butler = f"{self.tmpdir.name}/test_repo" 

255 config = BpsConfig( 

256 { 

257 "var1": "42a", 

258 "var2": "42b", 

259 "var3": "42c", 

260 "butlerConfig": config_butler, 

261 "finalJob": {"command1": "/usr/bin/echo {var1} {qgraphFile} {var2} {butlerConfig} {var3}"}, 

262 } 

263 ) 

264 gwf_exec, args = create_final_command(config, self.tmpdir.name) 

265 self.assertEqual(args, f"<FILE:runQgraphFile> {config_butler}") 

266 final_script = f"{self.tmpdir.name}/final_job.bash" 

267 self.assertEqual(gwf_exec.src_uri, final_script) 

268 with open(final_script) as infh: 

269 lines = infh.readlines() 

270 self.assertEqual( 

271 lines, self.script_beginning + ["/usr/bin/echo 42a ${qgraphFile} 42b ${butlerConfig} 42c\n"] 

272 ) 

273 

274 def testMultipleCommands(self): 

275 config_butler = f"{self.tmpdir.name}/test_repo" 

276 config = BpsConfig( 

277 { 

278 "var1": "42a", 

279 "var2": "42b", 

280 "var3": "42c", 

281 "butlerConfig": config_butler, 

282 "finalJob": { 

283 "command1": "/usr/bin/echo {var1} {qgraphFile} {var2} {butlerConfig} {var3}", 

284 "command2": "/usr/bin/uptime", 

285 }, 

286 } 

287 ) 

288 gwf_exec, args = create_final_command(config, self.tmpdir.name) 

289 self.assertEqual(args, f"<FILE:runQgraphFile> {config_butler}") 

290 final_script = f"{self.tmpdir.name}/final_job.bash" 

291 self.assertEqual(gwf_exec.src_uri, final_script) 

292 with open(final_script) as infh: 

293 lines = infh.readlines() 

294 self.assertEqual( 

295 lines, 

296 self.script_beginning 

297 + ["/usr/bin/echo 42a ${qgraphFile} 42b ${butlerConfig} 42c\n", "/usr/bin/uptime\n"], 

298 ) 

299 

300 def testZeroCommands(self): 

301 config_butler = f"{self.tmpdir.name}/test_repo" 

302 config = BpsConfig( 

303 { 

304 "var1": "42a", 

305 "var2": "42b", 

306 "var3": "42c", 

307 "butlerConfig": config_butler, 

308 "finalJob": { 

309 "cmd1": "/usr/bin/echo {var1} {qgraphFile} {var2} {butlerConfig} {var3}", 

310 "cmd2": "/usr/bin/uptime", 

311 }, 

312 } 

313 ) 

314 with self.assertRaisesRegex(RuntimeError, "finalJob.whenRun"): 

315 _, _ = create_final_command(config, self.tmpdir.name) 

316 

317 def testWhiteSpaceOnlyCommand(self): 

318 config_butler = f"{self.tmpdir.name}/test_repo" 

319 config = BpsConfig( 

320 { 

321 "butlerConfig": config_butler, 

322 "finalJob": {"command1": "", "command2": "\t \n"}, 

323 } 

324 ) 

325 with self.assertRaisesRegex(RuntimeError, "finalJob.whenRun"): 

326 _, _ = create_final_command(config, self.tmpdir.name) 

327 

328 def testSkipCommandUsingWhiteSpace(self): 

329 config_butler = f"{self.tmpdir.name}/test_repo" 

330 config = BpsConfig( 

331 { 

332 "var1": "42a", 

333 "var2": "42b", 

334 "var3": "42c", 

335 "butlerConfig": config_butler, 

336 "finalJob": { 

337 "command1": "/usr/bin/echo {var1} {qgraphFile} {var2} {butlerConfig} {var3}", 

338 "command2": "", # test skipping a command (i.e., overriding a default) 

339 "command3": "/usr/bin/uptime", 

340 }, 

341 } 

342 ) 

343 gwf_exec, args = create_final_command(config, self.tmpdir.name) 

344 self.assertEqual(args, f"<FILE:runQgraphFile> {config_butler}") 

345 final_script = f"{self.tmpdir.name}/final_job.bash" 

346 self.assertEqual(gwf_exec.src_uri, final_script) 

347 with open(final_script) as infh: 

348 lines = infh.readlines() 

349 self.assertEqual( 

350 lines, 

351 self.script_beginning 

352 + ["/usr/bin/echo 42a ${qgraphFile} 42b ${butlerConfig} 42c\n", "\n", "/usr/bin/uptime\n"], 

353 ) 

354 

355 

356class TestEnhanceCommand(unittest.TestCase): 

357 """Tests of _enhance_command function.""" 

358 

359 def setUp(self): 

360 self.gw_exec = GenericWorkflowExec("test_exec", "/dummy/dir/pipetask") 

361 self.config = BpsConfig( 

362 { 

363 # "profile": {}, 

364 "bpsUseShared": True, 

365 "whenSaveJobQgraph": "NEVER", 

366 "useLazyCommands": True, 

367 # "memoryLimit": 32768, 

368 "defOpts": "--long-log --log-file {submitPath}/{jobName}.{wmsAttemptNum}.json", 

369 "submitPath": "/the/path", 

370 } 

371 ) 

372 self.cached_vals = { 

373 "label1": { 

374 "profile": {}, 

375 "bpsUseShared": True, 

376 "whenSaveJobQgraph": "NEVER", 

377 "useLazyCommands": True, 

378 "memoryLimit": 32768, 

379 "key1": "val1", 

380 } 

381 } 

382 

383 def testAttemptNum(self): 

384 # test both in arguments as well as in variables in arguments 

385 gwjob = GenericWorkflowJob("job1", "label1", executable=self.gw_exec) 

386 gw = GenericWorkflow("test1") 

387 gw.add_job(gwjob) 

388 

389 first_args = "{defOpts} run-qbb repo test.qg --summary {submitPath}/{jobName}-summary." 

390 gwjob.arguments = first_args + "{wmsAttemptNum}.json" 

391 

392 new_arguments = first_args + "<WMS:attemptNum>.json" 

393 new_opts = "--long-log --log-file /the/path/job1.<WMS:attemptNum>.json" 

394 

395 _enhance_command(self.config, gw, gwjob, {}) 

396 

397 self.assertEqual(gwjob.arguments, new_arguments) 

398 self.assertEqual(gwjob.cmdvals["defOpts"], new_opts) 

399 

400 def testKeyCachedCmdVal(self): 

401 gwjob = GenericWorkflowJob("job1", "label1", executable=self.gw_exec) 

402 gw = GenericWorkflow("test1") 

403 gw.add_job(gwjob) 

404 gwjob.arguments = "run-qbb repo test.qg -x {key1}" 

405 self.assertNotIn("key1", gwjob.cmdvals) 

406 _enhance_command(self.config, gw, gwjob, self.cached_vals) 

407 self.assertEqual(gwjob.cmdvals["key1"], "val1") 

408 

409 def testS3Argument(self): 

410 """Make sure s3 double slashes are not getting removed.""" 

411 gwjob = GenericWorkflowJob("job1", "label1", executable=self.gw_exec) 

412 gw = GenericWorkflow("test1") 

413 gw.add_job(gwjob) 

414 s3 = "s3://user1@rubin-place-users/butler-pipeline1-processing.yaml" 

415 gwjob.arguments = s3 

416 _enhance_command(self.config, gw, gwjob, {}) 

417 self.assertEqual(gwjob.arguments, s3) 

418 

419 

420if __name__ == "__main__": 

421 unittest.main()