Hide keyboard shortcuts

Hot-keys on this page

r m x p   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

1# This file is part of ctrl_mpexec. 

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 program is free software: you can redistribute it and/or modify 

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

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

12# (at your option) any later version. 

13# 

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

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

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

17# GNU General Public License for more details. 

18# 

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

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

21 

22"""Simple unit test for cmdLineFwk module. 

23""" 

24 

25import argparse 

26import contextlib 

27import logging 

28import os 

29import pickle 

30import tempfile 

31import unittest 

32 

33from lsst.ctrl.mpexec.cmdLineFwk import CmdLineFwk 

34from lsst.ctrl.mpexec.cmdLineParser import (_ACTION_ADD_TASK, _ACTION_CONFIG, 

35 _ACTION_CONFIG_FILE, _ACTION_ADD_INSTRUMENT) 

36from lsst.daf.butler import Config, Quantum, Registry 

37from lsst.daf.butler.registry import RegistryConfig 

38from lsst.obs.base import Instrument 

39import lsst.pex.config as pexConfig 

40from lsst.pipe.base import (Pipeline, PipelineTask, PipelineTaskConfig, 

41 QuantumGraph, QuantumGraphTaskNodes, 

42 TaskDef, TaskFactory, PipelineTaskConnections) 

43import lsst.pipe.base.connectionTypes as cT 

44import lsst.utils.tests 

45from testUtil import (AddTask, AddTaskFactoryMock, makeSimpleQGraph) 

46 

47 

48logging.basicConfig(level=logging.INFO) 

49 

50# Have to monkey-patch Instrument.fromName() to not retrieve non-existing 

51# instrument from registry, these tests can run fine without actual instrument 

52# and implementing full mock for Instrument is too complicated. 

53Instrument.fromName = lambda name, reg: None 53 ↛ exitline 53 didn't run the lambda on line 53

54 

55 

56@contextlib.contextmanager 

57def makeTmpFile(contents=None): 

58 """Context manager for generating temporary file name. 

59 

60 Temporary file is deleted on exiting context. 

61 

62 Parameters 

63 ---------- 

64 contents : `bytes` 

65 Data to write into a file. 

66 """ 

67 fd, tmpname = tempfile.mkstemp() 

68 if contents: 

69 os.write(fd, contents) 

70 os.close(fd) 

71 yield tmpname 

72 with contextlib.suppress(OSError): 

73 os.remove(tmpname) 

74 

75 

76@contextlib.contextmanager 

77def makeSQLiteRegistry(): 

78 """Context manager to create new empty registry database. 

79 

80 Yields 

81 ------ 

82 config : `RegistryConfig` 

83 Registry configuration for initialized registry database. 

84 """ 

85 with makeTmpFile() as filename: 

86 uri = f"sqlite:///{filename}" 

87 config = RegistryConfig() 

88 config["db"] = uri 

89 Registry.fromConfig(config, create=True) 

90 yield config 

91 

92 

93class SimpleConnections(PipelineTaskConnections, dimensions=(), 

94 defaultTemplates={"template": "simple"}): 

95 schema = cT.InitInput(doc="Schema", 

96 name="{template}schema", 

97 storageClass="SourceCatalog") 

98 

99 

100class SimpleConfig(PipelineTaskConfig, pipelineConnections=SimpleConnections): 

101 field = pexConfig.Field(dtype=str, doc="arbitrary string") 

102 

103 def setDefaults(self): 

104 PipelineTaskConfig.setDefaults(self) 

105 

106 

107class TaskOne(PipelineTask): 

108 ConfigClass = SimpleConfig 

109 _DefaultName = "taskOne" 

110 

111 

112class TaskTwo(PipelineTask): 

113 ConfigClass = SimpleConfig 

114 _DefaultName = "taskTwo" 

115 

116 

117class TaskFactoryMock(TaskFactory): 

118 def loadTaskClass(self, taskName): 

119 if taskName == "TaskOne": 

120 return TaskOne, "TaskOne" 

121 elif taskName == "TaskTwo": 

122 return TaskTwo, "TaskTwo" 

123 

124 def makeTask(self, taskClass, config, overrides, butler): 

125 if config is None: 

126 config = taskClass.ConfigClass() 

127 if overrides: 

128 overrides.applyTo(config) 

129 return taskClass(config=config, butler=butler) 

130 

131 

132def _makeArgs(pipeline=None, qgraph=None, pipeline_actions=(), order_pipeline=False, 

133 save_pipeline="", save_qgraph="", save_single_quanta="", 

134 pipeline_dot="", qgraph_dot="", registryConfig=None): 

135 """Return parsed command line arguments. 

136 

137 Parameters 

138 ---------- 

139 pipeline : `str`, optional 

140 Name of the YAML file with pipeline. 

141 qgraph : `str`, optional 

142 Name of the pickle file with QGraph. 

143 pipeline_actions : itrable of `cmdLinePArser._PipelineAction`, optional 

144 order_pipeline : `bool` 

145 save_pipeline : `str` 

146 Name of the YAML file to store pipeline. 

147 save_qgraph : `str` 

148 Name of the pickle file to store QGraph. 

149 save_single_quanta : `str` 

150 Name of the pickle file pattern to store individual QGraph. 

151 pipeline_dot : `str` 

152 Name of the DOT file to write pipeline graph. 

153 qgraph_dot : `str` 

154 Name of the DOT file to write QGraph representation. 

155 """ 

156 args = argparse.Namespace() 

157 args.butler_config = Config() 

158 if registryConfig: 

159 args.butler_config["registry"] = registryConfig 

160 # The default datastore has a relocatable root, so we need to specify 

161 # some root here for it to use 

162 args.butler_config.configFile = "." 

163 args.pipeline = pipeline 

164 args.qgraph = qgraph 

165 args.pipeline_actions = pipeline_actions 

166 args.order_pipeline = order_pipeline 

167 args.save_pipeline = save_pipeline 

168 args.save_qgraph = save_qgraph 

169 args.save_single_quanta = save_single_quanta 

170 args.pipeline_dot = pipeline_dot 

171 args.qgraph_dot = qgraph_dot 

172 args.input = "" 

173 args.output = None 

174 args.output_run = None 

175 args.extend_run = False 

176 args.replace_run = False 

177 args.prune_replaced = False 

178 args.register_dataset_types = False 

179 args.skip_init_writes = False 

180 args.no_versions = False 

181 args.skip_existing = False 

182 args.init_only = False 

183 args.processes = 1 

184 args.profile = None 

185 args.enableLsstDebug = False 

186 args.graph_fixup = None 

187 args.timeout = None 

188 return args 

189 

190 

191def _makeQGraph(): 

192 """Make a trivial QuantumGraph with one quantum. 

193 

194 The only thing that we need to do with this quantum graph is to pickle 

195 it, the quanta in this graph are not usable for anything else. 

196 

197 Returns 

198 ------- 

199 qgraph : `~lsst.pipe.base.QuantumGraph` 

200 """ 

201 taskDef = TaskDef(taskName="taskOne", config=SimpleConfig()) 

202 quanta = [Quantum()] 

203 taskNodes = QuantumGraphTaskNodes(taskDef=taskDef, quanta=quanta, initInputs={}, initOutputs={}) 

204 qgraph = QuantumGraph([taskNodes]) 

205 return qgraph 

206 

207 

208class CmdLineFwkTestCase(unittest.TestCase): 

209 """A test case for CmdLineFwk 

210 """ 

211 

212 def testMakePipeline(self): 

213 """Tests for CmdLineFwk.makePipeline method 

214 """ 

215 fwk = CmdLineFwk() 

216 

217 # make empty pipeline 

218 args = _makeArgs() 

219 pipeline = fwk.makePipeline(args) 

220 self.assertIsInstance(pipeline, Pipeline) 

221 self.assertEqual(len(pipeline), 0) 

222 

223 # few tests with serialization 

224 with makeTmpFile() as tmpname: 

225 # make empty pipeline and store it in a file 

226 args = _makeArgs(save_pipeline=tmpname) 

227 pipeline = fwk.makePipeline(args) 

228 self.assertIsInstance(pipeline, Pipeline) 

229 

230 # read pipeline from a file 

231 args = _makeArgs(pipeline=tmpname) 

232 pipeline = fwk.makePipeline(args) 

233 self.assertIsInstance(pipeline, Pipeline) 

234 self.assertEqual(len(pipeline), 0) 

235 

236 # single task pipeline 

237 actions = [ 

238 _ACTION_ADD_TASK("TaskOne:task1") 

239 ] 

240 args = _makeArgs(pipeline_actions=actions) 

241 pipeline = fwk.makePipeline(args) 

242 self.assertIsInstance(pipeline, Pipeline) 

243 self.assertEqual(len(pipeline), 1) 

244 

245 # many task pipeline 

246 actions = [ 

247 _ACTION_ADD_TASK("TaskOne:task1a"), 

248 _ACTION_ADD_TASK("TaskTwo:task2"), 

249 _ACTION_ADD_TASK("TaskOne:task1b") 

250 ] 

251 args = _makeArgs(pipeline_actions=actions) 

252 pipeline = fwk.makePipeline(args) 

253 self.assertIsInstance(pipeline, Pipeline) 

254 self.assertEqual(len(pipeline), 3) 

255 

256 # single task pipeline with config overrides, cannot use TaskOne, need 

257 # something that can be imported with `doImport()` 

258 actions = [ 

259 _ACTION_ADD_TASK("testUtil.AddTask:task"), 

260 _ACTION_CONFIG("task:addend=100") 

261 ] 

262 args = _makeArgs(pipeline_actions=actions) 

263 pipeline = fwk.makePipeline(args) 

264 taskDefs = list(pipeline.toExpandedPipeline()) 

265 self.assertEqual(len(taskDefs), 1) 

266 self.assertEqual(taskDefs[0].config.addend, 100) 

267 

268 overrides = b"config.addend = 1000\n" 

269 with makeTmpFile(overrides) as tmpname: 

270 actions = [ 

271 _ACTION_ADD_TASK("testUtil.AddTask:task"), 

272 _ACTION_CONFIG_FILE("task:" + tmpname) 

273 ] 

274 args = _makeArgs(pipeline_actions=actions) 

275 pipeline = fwk.makePipeline(args) 

276 taskDefs = list(pipeline.toExpandedPipeline()) 

277 self.assertEqual(len(taskDefs), 1) 

278 self.assertEqual(taskDefs[0].config.addend, 1000) 

279 

280 # Check --instrument option, for now it only checks that it does not crash 

281 actions = [ 

282 _ACTION_ADD_TASK("testUtil.AddTask:task"), 

283 _ACTION_ADD_INSTRUMENT("Instrument") 

284 ] 

285 args = _makeArgs(pipeline_actions=actions) 

286 pipeline = fwk.makePipeline(args) 

287 

288 def testMakeGraphFromPickle(self): 

289 """Tests for CmdLineFwk.makeGraph method. 

290 

291 Only most trivial case is tested that does not do actual graph 

292 building. 

293 """ 

294 fwk = CmdLineFwk() 

295 

296 with makeTmpFile() as tmpname, makeSQLiteRegistry() as registryConfig: 

297 

298 # make non-empty graph and store it in a file 

299 qgraph = _makeQGraph() 

300 with open(tmpname, "wb") as pickleFile: 

301 qgraph.save(pickleFile) 

302 args = _makeArgs(qgraph=tmpname, registryConfig=registryConfig) 

303 qgraph = fwk.makeGraph(None, args) 

304 self.assertIsInstance(qgraph, QuantumGraph) 

305 self.assertEqual(len(qgraph), 1) 

306 

307 # pickle with wrong object type 

308 with open(tmpname, "wb") as pickleFile: 

309 pickle.dump({}, pickleFile) 

310 args = _makeArgs(qgraph=tmpname, registryConfig=registryConfig) 

311 with self.assertRaises(TypeError): 

312 fwk.makeGraph(None, args) 

313 

314 # reading empty graph from pickle should work but makeGraph() 

315 # will return None and make a warning 

316 qgraph = QuantumGraph() 

317 with open(tmpname, "wb") as pickleFile: 

318 qgraph.save(pickleFile) 

319 args = _makeArgs(qgraph=tmpname, registryConfig=registryConfig) 

320 with self.assertWarnsRegex(UserWarning, "QuantumGraph is empty"): 

321 # this also tests that warning is generated for empty graph 

322 qgraph = fwk.makeGraph(None, args) 

323 self.assertIs(qgraph, None) 

324 

325 def testSimpleQGraph(self): 

326 """Test successfull execution of trivial quantum graph. 

327 """ 

328 

329 nQuanta = 5 

330 butler, qgraph = makeSimpleQGraph(nQuanta) 

331 

332 # should have one task and number of quanta 

333 self.assertEqual(len(qgraph), 1) 

334 self.assertEqual(len(list(qgraph.quanta())), nQuanta) 

335 

336 args = _makeArgs() 

337 fwk = CmdLineFwk() 

338 taskFactory = AddTaskFactoryMock() 

339 

340 # run whole thing 

341 AddTask.countExec = 0 

342 fwk.runPipeline(qgraph, taskFactory, args, butler=butler) 

343 self.assertEqual(AddTask.countExec, nQuanta) 

344 

345 def testSimpleQGraphSkipExisting(self): 

346 """Test continuing execution of trivial quantum graph with --skip-existing. 

347 """ 

348 

349 nQuanta = 5 

350 butler, qgraph = makeSimpleQGraph(nQuanta) 

351 

352 # should have one task and number of quanta 

353 self.assertEqual(len(qgraph), 1) 

354 self.assertEqual(len(list(qgraph.quanta())), nQuanta) 

355 

356 args = _makeArgs() 

357 fwk = CmdLineFwk() 

358 taskFactory = AddTaskFactoryMock() 

359 

360 # run whole thing 

361 AddTask.countExec = 0 

362 AddTask.stopAt = 3 

363 with self.assertRaises(RuntimeError): 

364 fwk.runPipeline(qgraph, taskFactory, args, butler=butler) 

365 self.assertEqual(AddTask.countExec, 3) 

366 

367 AddTask.stopAt = -1 

368 args.skip_existing = True 

369 args.no_versions = True 

370 fwk.runPipeline(qgraph, taskFactory, args, butler=butler) 

371 self.assertEqual(AddTask.countExec, nQuanta) 

372 

373 def testShowPipeline(self): 

374 """Test for --show options for pipeline. 

375 """ 

376 fwk = CmdLineFwk() 

377 

378 actions = [ 

379 _ACTION_ADD_TASK("testUtil.AddTask:task"), 

380 _ACTION_CONFIG("task:addend=100") 

381 ] 

382 args = _makeArgs(pipeline_actions=actions) 

383 pipeline = fwk.makePipeline(args) 

384 

385 args.show = ["pipeline"] 

386 fwk.showInfo(args, pipeline) 

387 args.show = ["config"] 

388 fwk.showInfo(args, pipeline) 

389 args.show = ["history=task::addend"] 

390 fwk.showInfo(args, pipeline) 

391 args.show = ["tasks"] 

392 fwk.showInfo(args, pipeline) 

393 

394 def testShowGraph(self): 

395 """Test for --show options for quantum graph. 

396 """ 

397 fwk = CmdLineFwk() 

398 

399 nQuanta = 2 

400 butler, qgraph = makeSimpleQGraph(nQuanta) 

401 

402 args = _makeArgs() 

403 args.show = ["graph"] 

404 fwk.showInfo(args, pipeline=None, graph=qgraph) 

405 # TODO: cannot test "workflow" option presently, it instanciates 

406 # butler from command line options and there is no way to pass butler 

407 # mock to that code. 

408 

409 

410class MyMemoryTestCase(lsst.utils.tests.MemoryTestCase): 

411 pass 

412 

413 

414def setup_module(module): 

415 lsst.utils.tests.init() 

416 

417 

418if __name__ == "__main__": 418 ↛ 419line 418 didn't jump to line 419, because the condition on line 418 was never true

419 lsst.utils.tests.init() 

420 unittest.main()