Coverage for tests / test_transform.py: 19%

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 BPS_SEARCH_ORDER, BpsConfig, GenericWorkflowJob 

38from lsst.ctrl.bps.transform import ( 

39 _get_job_values, 

40 create_final_command, 

41 create_generic_workflow, 

42 create_generic_workflow_config, 

43) 

44 

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

46 

47 

48class TestCreateGenericWorkflowConfig(unittest.TestCase): 

49 """Tests of create_generic_workflow_config.""" 

50 

51 def testCreate(self): 

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

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

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

55 self.assertIsInstance(wf_config, BpsConfig) 

56 for key in config: 

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

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

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

60 

61 

62class TestCreateGenericWorkflow(unittest.TestCase): 

63 """Tests of create_generic_workflow.""" 

64 

65 def setUp(self): 

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

67 self.config = BpsConfig( 

68 { 

69 "runInit": True, 

70 "computeSite": "global", 

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

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

73 "cluster": { 

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

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

76 }, 

77 "cloud": { 

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

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

80 }, 

81 "site": { 

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

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

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

85 }, 

86 # Needed because transform assumes they exist 

87 "whenSaveJobQgraph": "NEVER", 

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

89 }, 

90 BPS_SEARCH_ORDER, 

91 ) 

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

93 

94 def tearDown(self): 

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

96 

97 def testCreatingGenericWorkflowGlobal(self): 

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

99 config = BpsConfig(self.config) 

100 config["computeCloud"] = "cloud1" 

101 config["computeSite"] = "site2" 

102 config["queue"] = "global_queue" 

103 print(config) 

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

105 for jname in workflow: 

106 gwjob = workflow.get_job(jname) 

107 print(gwjob) 

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

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

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

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

112 final = workflow.get_final() 

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

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

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

116 

117 def testCreatingQuantumGraphMixed(self): 

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

119 config = BpsConfig(self.config) 

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

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

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

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

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

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

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

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

128 for jname in workflow: 

129 gwjob = workflow.get_job(jname) 

130 print(gwjob) 

131 if jname.startswith("cl1"): 

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

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

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

135 elif jname.startswith("cl2"): 

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

137 self.assertIsNone(gwjob.compute_cloud) 

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

139 elif jname.startswith("pipetask"): 

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

141 self.assertIsNone(gwjob.compute_cloud) 

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

143 final = workflow.get_final() 

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

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

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

147 

148 

149class TestGetJobValues(unittest.TestCase): 

150 """Tests of _get_job_values.""" 

151 

152 def setUp(self): 

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

154 

155 def testGettingDefaults(self): 

156 """Test retrieving default values.""" 

157 config = BpsConfig({}) 

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

159 self.assertTrue( 

160 all( 

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

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

163 ) 

164 ) 

165 

166 def testEnablingMemoryScaling(self): 

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

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

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

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

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

172 

173 def testDisablingMemoryScaling(self): 

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

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

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

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

178 

179 def testRetrievingCmdLine(self): 

180 """Test retrieving the command line.""" 

181 cmd_line_key = "runQuantum" 

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

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

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

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

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

187 

188 def testEnvironment(self): 

189 config = BpsConfig( 

190 { 

191 "var1": "two", 

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

193 } 

194 ) 

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

196 truth = BpsConfig({"TEST_INT": "1", "TEST_BOOL": "False", "TEST_SPACES": "one two three"}, {}, None) 

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

198 

199 def testEnvironmentOptions(self): 

200 config = BpsConfig( 

201 { 

202 "var1": "two", 

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

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

205 } 

206 ) 

207 search_obj = config["finalJob"] 

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

209 job_values = _get_job_values(config, search_opts, None) 

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

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

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

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

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

215 

216 

217class TestCreateFinalCommand(unittest.TestCase): 

218 """Tests for the create_final_command function.""" 

219 

220 def setUp(self): 

221 self.tmpdir = tempfile.TemporaryDirectory() 

222 self.script_beginning = [ 

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

224 "\n", 

225 "set -e\n", 

226 "set -x\n", 

227 "qgraphFile=$1\n", 

228 "butlerConfig=$2\n", 

229 ] 

230 

231 def tearDown(self): 

232 self.tmpdir.cleanup() 

233 

234 def testSingleCommand(self): 

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

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

237 config = BpsConfig( 

238 { 

239 "var1": "42a", 

240 "var2": "42b", 

241 "var3": "42c", 

242 "butlerConfig": config_butler, 

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

244 } 

245 ) 

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

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

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

249 self.assertEqual(gwf_exec.src_uri, final_script) 

250 with open(final_script) as infh: 

251 lines = infh.readlines() 

252 self.assertEqual( 

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

254 ) 

255 

256 def testMultipleCommands(self): 

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

258 config = BpsConfig( 

259 { 

260 "var1": "42a", 

261 "var2": "42b", 

262 "var3": "42c", 

263 "butlerConfig": config_butler, 

264 "finalJob": { 

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

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

267 }, 

268 } 

269 ) 

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

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

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

273 self.assertEqual(gwf_exec.src_uri, final_script) 

274 with open(final_script) as infh: 

275 lines = infh.readlines() 

276 self.assertEqual( 

277 lines, 

278 self.script_beginning 

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

280 ) 

281 

282 def testZeroCommands(self): 

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

284 config = BpsConfig( 

285 { 

286 "var1": "42a", 

287 "var2": "42b", 

288 "var3": "42c", 

289 "butlerConfig": config_butler, 

290 "finalJob": { 

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

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

293 }, 

294 } 

295 ) 

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

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

298 

299 def testWhiteSpaceOnlyCommand(self): 

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

301 config = BpsConfig( 

302 { 

303 "butlerConfig": config_butler, 

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

305 } 

306 ) 

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

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

309 

310 def testSkipCommandUsingWhiteSpace(self): 

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

312 config = BpsConfig( 

313 { 

314 "var1": "42a", 

315 "var2": "42b", 

316 "var3": "42c", 

317 "butlerConfig": config_butler, 

318 "finalJob": { 

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

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

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

322 }, 

323 } 

324 ) 

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

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

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

328 self.assertEqual(gwf_exec.src_uri, final_script) 

329 with open(final_script) as infh: 

330 lines = infh.readlines() 

331 self.assertEqual( 

332 lines, 

333 self.script_beginning 

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

335 ) 

336 

337 

338if __name__ == "__main__": 

339 unittest.main()