Coverage for src/gridtk/easy.py: 0%

1# Copyright © 2022 Idiap Research Institute <contact@idiap.ch> 

2# 

3# SPDX-License-Identifier: GPL-3.0-or-later 

4"""Common arguments to grid jobs.""" 

5 

6import os 

7import sys 

8 

9from . import manager, tools 

10 

11 

12def add_arguments(parser): 

13 """Adds stock arguments to argparse parsers from scripts that submit grid 

14 jobs.""" 

15 

16 default_log_path = os.path.realpath("logs") 

17 

18 parser.add_argument( 

19 "--log-dir", 

20 metavar="LOG", 

21 type=str, 

22 dest="logdir", 

23 default=default_log_path, 

24 help='Base directory used for logging (defaults to "%(default)s")', 

25 ) 

26 

27 q_choices = ( 

28 "default", 

29 "all.q", 

30 "q_1day", 

31 "q1d", 

32 "q_1week", 

33 "q1w", 

34 "q_1month", 

35 "q1m", 

36 "q_1day_mth", 

37 "q1dm", 

38 "q_1week_mth", 

39 "q1wm", 

40 "q_gpu", 

41 "gpu", 

42 "q_long_gpu", 

43 "lgpu", 

44 "q_short_gpu", 

45 "sgpu", 

46 "vsgpu", 

47 ) 

48 

49 parser.add_argument( 

50 "--queue-name", 

51 metavar="QUEUE", 

52 type=str, 

53 dest="queue", 

54 default=q_choices[0], 

55 choices=q_choices, 

56 help="Queue for submission - one of " 

57 + "|".join(q_choices) 

58 + ' (defaults to "%(default)s")', 

59 ) 

60 

61 parser.add_argument( 

62 "--hostname", 

63 metavar="HOSTNAME", 

64 type=str, 

65 dest="hostname", 

66 default=None, 

67 help="If set, it asks the queue to use only a subset of the available nodes", 

68 ) 

69 parser.add_argument( 

70 "--memfree", 

71 metavar="MEMFREE", 

72 type=str, 

73 dest="memfree", 

74 default=None, 

75 help="Adds the '-l mem_free' argument to qsub", 

76 ) 

77 parser.add_argument( 

78 "--hvmem", 

79 metavar="HVMEM", 

80 type=str, 

81 dest="hvmem", 

82 default=None, 

83 help="Adds the '-l h_vmem' argument to qsub", 

84 ) 

85 parser.add_argument( 

86 "--pe-opt", 

87 metavar="PE_OPT", 

88 type=str, 

89 dest="pe_opt", 

90 default=None, 

91 help="Adds the '--pe ' argument to qsub", 

92 ) 

93 

94 parser.add_argument( 

95 "--no-cwd", 

96 default=True, 

97 action="store_false", 

98 dest="cwd", 

99 help="Do not change to the current directory when starting the grid job", 

100 ) 

101 

102 parser.add_argument( 

103 "--dry-run", 

104 default=False, 

105 action="store_true", 

106 dest="dryrun", 

107 help="Does not really submit anything, just print what would do instead", 

108 ) 

109 

110 parser.add_argument( 

111 "--job-database", 

112 default=None, 

113 dest="statefile", 

114 help="The path to the state file that will be created with the submissions (defaults to the parent directory of your logs directory)", 

115 ) 

116 

117 return parser 

118 

119 

120def create_manager(arguments): 

121 """A simple wrapper to JobManager() that places the statefile on the 

122 correct path by default.""" 

123 

124 if arguments.statefile is None: 

125 arguments.statefile = os.path.join( 

126 os.path.dirname(arguments.logdir), "submitted.db" 

127 ) 

128 

129 arguments.statefile = os.path.realpath(arguments.statefile) 

130 

131 return manager.JobManager(statefile=arguments.statefile) 

132 

133 

134class DryRunJob: 

135 """A simple wrapper for dry-run jobs that behaves like a normal job.""" 

136 

137 # distributed as jobs are "submitted" 

138 current_id = 0 

139 

140 def __init__( 

141 self, 

142 cmd, 

143 cwd, 

144 queue, 

145 hostname, 

146 memfree, 

147 hvmem, 

148 gpumem, 

149 pe_opt, 

150 stdout, 

151 stderr, 

152 name, 

153 array, 

154 deps, 

155 ): 

156 self.myid = DryRunJob.current_id 

157 DryRunJob.current_id += 1 

158 

159 self.cmd = cmd 

160 self.cwd = cwd 

161 self.queue = queue 

162 self.hostname = hostname 

163 self.memfree = memfree 

164 self.hvmem = hvmem 

165 self.gpumem = gpumem 

166 self.pe_opt = pe_opt 

167 self.stdout = stdout 

168 self.stderr = stderr 

169 self.name = name 

170 self.array = array 

171 self.deps = deps 

172 

173 def __str__(self): 

174 return """ 

175 id : %d 

176 command : %s 

177 cwd : %s 

178 queue : %s 

179 hostname : %s 

180 memfree : %s 

181 hvmem : %s 

182 gpumem : %s 

183 pe_opt : %s 

184 stdout : %s 

185 stderr : %s 

186 name : %s 

187 array : %s 

188 depends : %s""" % ( 

189 self.myid, 

190 self.cmd, 

191 self.cwd, 

192 self.queue, 

193 self.hostname, 

194 self.memfree, 

195 self.hvmem, 

196 self.gpumem, 

197 self.pe_opt, 

198 self.stdout, 

199 self.stderr, 

200 self.name, 

201 self.array, 

202 self.deps, 

203 ) 

204 

205 def id(self): 

206 return self.myid 

207 

208 

209def submit(jman, command, arguments, deps=[], array=None): 

210 """An easy submission option for grid-enabled scripts. 

211 

212 Create the log directories using random hash codes. Use the 

213 arguments as parsed by the main script. 

214 """ 

215 logdir = os.path.join(os.path.realpath(arguments.logdir), tools.random_logdir()) 

216 

217 jobname = os.path.splitext(os.path.basename(command[0]))[0] 

218 cmd = tools.make_shell(sys.executable, command) 

219 

220 if arguments.dryrun: 

221 return DryRunJob( 

222 cmd, 

223 cwd=arguments.cwd, 

224 queue=arguments.queue, 

225 hostname=arguments.hostname, 

226 memfree=arguments.memfree, 

227 hvmem=arguments.hvmem, 

228 gpumem=arguments.gpumem, 

229 pe_opt=arguments.pe_opt, 

230 stdout=logdir, 

231 stderr=logdir, 

232 name=jobname, 

233 deps=deps, 

234 array=array, 

235 ) 

236 

237 # really submit 

238 return jman.submit( 

239 cmd, 

240 cwd=arguments.cwd, 

241 queue=arguments.queue, 

242 hostname=arguments.hostname, 

243 memfree=arguments.memfree, 

244 hvmem=arguments.hvmem, 

245 gpumem=arguments.gpumem, 

246 pe_opt=arguments.pe_opt, 

247 stdout=logdir, 

248 stderr=logdir, 

249 name=jobname, 

250 deps=deps, 

251 array=array, 

252 )