Coverage for python/lsst/ctrl/pool/test/demoTask.py: 41%

42 statements  

« prev     ^ index     » next       coverage.py v6.4.1, created at 2022-06-28 09:40 +0000

1import math 

2import collections 

3import operator 

4from lsst.ctrl.pool.parallel import BatchPoolTask 

5from lsst.ctrl.pool.pool import Pool 

6from lsst.pipe.base import ArgumentParser 

7from lsst.pex.config import Config 

8 

9__all__ = ["DemoTask", ] 

10 

11 

12class DemoTask(BatchPoolTask): 

13 """Task for demonstrating the BatchPoolTask functionality""" 

14 ConfigClass = Config 

15 _DefaultName = "demo" 

16 

17 @classmethod 

18 def _makeArgumentParser(cls, *args, **kwargs): 

19 kwargs.pop('doBatch', False) # Unused 

20 parser = ArgumentParser(name="demo", *args, **kwargs) 

21 parser.add_id_argument("--id", datasetType="raw", level="visit", 

22 help="data ID, e.g. --id visit=12345") 

23 return parser 

24 

25 @classmethod 

26 def batchWallTime(cls, time, parsedCmd, numCores): 

27 """Return walltime request for batch job 

28 

29 Subclasses should override if the walltime should be calculated 

30 differently (e.g., addition of some serial time). 

31 

32 @param time: Requested time per iteration 

33 @param parsedCmd: Results of argument parsing 

34 @param numCores: Number of cores 

35 """ 

36 numTargets = [sum(1 for ccdRef in visitRef.subItems("ccd") if ccdRef.datasetExists("raw")) for 

37 visitRef in parsedCmd.id.refList] 

38 return time*sum(math.ceil(tt/numCores) for tt in numTargets) 

39 

40 def runDataRef(self, visitRef): 

41 """Main entry-point 

42 

43 Only the master node runs this method. It will dispatch jobs to the 

44 slave nodes. 

45 """ 

46 pool = Pool("test") 

47 

48 # Less overhead to transfer the butler once rather than in each dataRef 

49 dataIdList = dict([(ccdRef.get("ccdExposureId"), ccdRef.dataId) 

50 for ccdRef in visitRef.subItems("ccd") if ccdRef.datasetExists("raw")]) 

51 dataIdList = collections.OrderedDict(sorted(dataIdList.items())) 

52 

53 with self.logOperation("master"): 

54 total = pool.reduce(operator.add, self.run, list(dataIdList.values()), 

55 butler=visitRef.getButler()) 

56 self.log.info("Total number of pixels read: %d" % (total,)) 

57 

58 def run(self, cache, dataId, butler=None): 

59 """Read image and return number of pixels 

60 

61 Only the slave nodes run this method. 

62 """ 

63 assert butler is not None 

64 with self.logOperation("read %s" % (dataId,)): 

65 raw = butler.get("raw", dataId, immediate=True) 

66 dims = raw.getDimensions() 

67 num = dims.getX()*dims.getY() 

68 self.log.info("Read %d pixels for %s" % (num, dataId,)) 

69 return num 

70 

71 def _getConfigName(self): 

72 return None 

73 

74 def _getMetadataName(self): 

75 return None 

76 

77 def _getEupsVersionsName(self): 

78 return None