Coverage for pySDC/projects/GPU/configs/base_config.py: 85%

1from pySDC.core.convergence_controller import ConvergenceController 

2import pickle 

3import numpy as np 

4 

5 

6def get_config(args): 

7 name = args['config'] 

8 if name[:2] == 'GS': 

9 from pySDC.projects.GPU.configs.GS_configs import get_config as _get_config 

10 elif name[:5] == 'RBC3D': 

11 from pySDC.projects.RayleighBenard.RBC3D_configs import get_config as _get_config 

12 elif name[:3] == 'RBC': 

13 from pySDC.projects.GPU.configs.RBC_configs import get_config as _get_config 

14 else: 

15 raise NotImplementedError(f'There is no configuration called {name!r}!') 

16 

17 return _get_config(args) 

18 

19 

20def get_comms(n_procs_list, comm_world=None, _comm=None, _tot_rank=0, _rank=None, useGPU=False): 

21 from mpi4py import MPI 

22 

23 comm_world = MPI.COMM_WORLD if comm_world is None else comm_world 

24 _comm = comm_world if _comm is None else _comm 

25 _rank = comm_world.rank if _rank is None else _rank 

26 

27 if len(n_procs_list) > 0: 

28 color = _tot_rank + _rank // n_procs_list[0] 

29 new_comm = comm_world.Split(color) 

30 

31 assert new_comm.size == n_procs_list[0] 

32 

33 if useGPU: 

34 import cupy_backends 

35 

36 try: 

37 import cupy 

38 from pySDC.helpers.NCCL_communicator import NCCLComm 

39 

40 new_comm = NCCLComm(new_comm) 

41 except ( 

42 ImportError, 

43 cupy_backends.cuda.api.runtime.CUDARuntimeError, 

44 cupy_backends.cuda.libs.nccl.NcclError, 

45 ): 

46 print('Warning: Communicator is MPI instead of NCCL in spite of using GPUs!') 

47 

48 return [new_comm] + get_comms( 

49 n_procs_list[1:], 

50 comm_world, 

51 _comm=new_comm, 

52 _tot_rank=_tot_rank + _comm.size * new_comm.rank, 

53 _rank=_comm.rank // new_comm.size, 

54 useGPU=useGPU, 

55 ) 

56 else: 

57 return [] 

58 

59 

60class Config(object): 

61 sweeper_type = None 

62 Tend = None 

63 base_path = './' 

64 logging_time_increment = 0.5 

65 

66 def __init__(self, args, comm_world=None): 

67 from mpi4py import MPI 

68 

69 self.args = args 

70 self.comm_world = MPI.COMM_WORLD if comm_world is None else comm_world 

71 self.n_procs_list = args["procs"] 

72 if args['mode'] in ['run', 'benchmark']: 

73 distribution = args.get('distribution', 'time_first') 

74 if distribution in ['space_first', 'space_major']: 

75 self.comms = get_comms( 

76 n_procs_list=self.n_procs_list[::-1], useGPU=args['useGPU'], comm_world=self.comm_world 

77 )[::-1] 

78 elif distribution in ['time_first', 'time_major']: 

79 self.comms = get_comms( 

80 n_procs_list=self.n_procs_list, useGPU=args['useGPU'], comm_world=self.comm_world 

81 ) 

82 else: 

83 raise NotImplementedError 

84 else: 

85 self.comms = [MPI.COMM_SELF, MPI.COMM_SELF, MPI.COMM_SELF] 

86 self.ranks = [me.rank for me in self.comms] 

87 

88 def get_file_name(self): 

89 res = self.args['res'] 

90 return f'{self.base_path}/data/{type(self).__name__}-res{res}.pySDC' 

91 

92 def get_hook(self, hook_class, key=None, **attributes): 

93 """ 

94 Get a subclass of a hook configured through class attributes, the same one for every call with the same key. 

95 

96 Setting the attributes on `hook_class` itself would configure every other run in the process. The subclass is 

97 kept because the restart counter set on it has to reach the hook that does the logging and `LogStats`. 

98 

99 Args: 

100 hook_class (pySDC.Hook): The hook to configure 

101 key: Distinguishes differently configured subclasses of the same hook 

102 attributes: Class attributes to set on the subclass 

103 

104 Returns: 

105 pySDC.Hook: The configured subclass 

106 """ 

107 hooks = self.__dict__.setdefault('_hooks', {}) 

108 if (hook_class, key) not in hooks: 

109 hooks[(hook_class, key)] = type(hook_class.__name__, (hook_class,), {}) 

110 hook = hooks[(hook_class, key)] 

111 for name, value in attributes.items(): 

112 setattr(hook, name, value) 

113 return hook 

114 

115 def get_LogToFile(self, *args, **kwargs): 

116 if self.comms[1].rank > 0: 

117 return None 

118 from pySDC.implementations.hooks.log_solution import LogToFile 

119 

120 return self.get_hook( 

121 LogToFile, 

122 filename=self.get_file_name(), 

123 time_increment=self.logging_time_increment, 

124 allow_overwriting=True, 

125 ) 

126 

127 def get_description(self, *args, MPIsweeper=False, useGPU=False, **kwargs): 

128 description = {} 

129 description['problem_class'] = None 

130 description['problem_params'] = {'useGPU': useGPU, 'comm': self.comms[2]} 

131 description['sweeper_class'] = self.get_sweeper(useMPI=MPIsweeper) 

132 description['sweeper_params'] = {'initial_guess': 'copy'} 

133 description['level_params'] = {} 

134 description['step_params'] = {} 

135 description['convergence_controllers'] = {} 

136 

137 hook = self.get_LogToFile() 

138 if hook: 

139 path = self.get_file_name()[:-6] 

140 description['convergence_controllers'][LogStats] = {'path': path, 'hook': hook} 

141 

142 if MPIsweeper: 

143 description['sweeper_params']['comm'] = self.comms[1] 

144 return description 

145 

146 def get_controller_params(self, *args, logger_level=15, **kwargs): 

147 from pySDC.implementations.hooks.log_work import LogWork 

148 from pySDC.implementations.hooks.log_step_size import LogStepSize 

149 from pySDC.implementations.hooks.log_restarts import LogRestarts 

150 

151 controller_params = {} 

152 controller_params['logger_level'] = logger_level if self.comm_world.rank == 0 else 40 

153 controller_params['hook_class'] = [LogWork, LogStepSize, LogRestarts] 

154 logToFile = self.get_LogToFile() 

155 if logToFile: 

156 controller_params['hook_class'] += [logToFile] 

157 controller_params['mssdc_jac'] = False 

158 return controller_params 

159 

160 def get_sweeper(self, useMPI): 

161 if useMPI and self.sweeper_type == 'IMEX': 

162 from pySDC.implementations.sweeper_classes.imex_1st_order_MPI import imex_1st_order_MPI as sweeper 

163 elif not useMPI and self.sweeper_type == 'IMEX': 

164 from pySDC.implementations.sweeper_classes.imex_1st_order import imex_1st_order as sweeper 

165 elif useMPI and self.sweeper_type == 'generic_implicit': 

166 from pySDC.implementations.sweeper_classes.generic_implicit_MPI import generic_implicit_MPI as sweeper 

167 elif not useMPI and self.sweeper_type == 'generic_implicit': 

168 from pySDC.implementations.sweeper_classes.generic_implicit import generic_implicit as sweeper 

169 else: 

170 raise NotImplementedError(f'Don\'t know the sweeper for {self.sweeper_type=}') 

171 

172 return sweeper 

173 

174 def prepare_caches(self, prob): 

175 pass 

176 

177 def get_path(self, *args, ranks=None, **kwargs): 

178 ranks = self.ranks if ranks is None else ranks 

179 return f'{type(self).__name__}{self.args_to_str()}-{ranks[0]}-{ranks[2]}' 

180 

181 def args_to_str(self, args=None): 

182 args = self.args if args is None else args 

183 name = '' 

184 

185 name = f'{name}-res_{args["res"]}' 

186 name = f'{name}-useGPU_{args["useGPU"]}' 

187 name = f'{name}-procs_{args["procs"][0]}_{args["procs"][1]}_{args["procs"][2]}' 

188 return name 

189 

190 def plot(self, P, idx, num_procs_list): 

191 raise NotImplementedError 

192 

193 def get_initial_condition(self, P, *args, restart_idx=0, **kwargs): 

194 if restart_idx == 0: 

195 return P.u_exact(t=0), 0 

196 else: 

197 

198 from pySDC.helpers.fieldsIO import FieldsIO 

199 

200 P.setUpFieldsIO() 

201 outfile = FieldsIO.fromFile(self.get_file_name()) 

202 

203 t0, solution = outfile.readField(restart_idx) 

204 solution = solution[: P.spectral.ncomponents, ...] 

205 

206 u0 = P.u_init 

207 

208 if P.spectral_space: 

209 u0[...] = P.transform(solution) 

210 else: 

211 u0[...] = solution 

212 

213 return u0, t0 

214 

215 LogToFile = self.get_LogToFile() 

216 file = LogToFile.load(restart_idx) 

217 LogToFile.counter = restart_idx 

218 u0 = P.u_init 

219 if hasattr(P, 'spectral_space'): 

220 if P.spectral_space: 

221 u0[...] = P.transform(file['u']) 

222 else: 

223 u0[...] = file['u'] 

224 else: 

225 u0[...] = file['u'] 

226 return u0, file['t'] 

227 

228 

229class LogStats(ConvergenceController): 

230 

231 def get_stats_path(self, index=0): 

232 return f'{self.params.path}_{index:06d}-stats.pickle' 

233 

234 def merge_all_stats(self, controller): 

235 hook = self.params.hook 

236 

237 stats = {} 

238 for i in range(hook.counter - 1): 

239 try: 

240 with open(self.get_stats_path(index=i), 'rb') as file: 

241 _stats = pickle.load(file) 

242 stats = {**stats, **_stats} 

243 except (FileNotFoundError, EOFError): 

244 print(f'Warning: No stats found at path {self.get_stats_path(index=i)}') 

245 

246 stats = {**stats, **controller.return_stats()} 

247 return stats 

248 

249 def reset_stats(self, controller): 

250 for hook in controller.hooks: 

251 hook.reset_stats() 

252 self.logger.debug('Reset stats') 

253 

254 def setup(self, controller, params, *args, **kwargs): 

255 params['control_order'] = 999 

256 if 'hook' not in params.keys(): 

257 from pySDC.implementations.hooks.log_solution import LogToFile 

258 

259 params['hook'] = LogToFile 

260 

261 self.counter = params['hook'].counter 

262 return super().setup(controller, params, *args, **kwargs) 

263 

264 def post_step_processing(self, controller, S, **kwargs): 

265 hook = self.params.hook 

266 

267 P = S.levels[0].prob 

268 

269 while self.counter < hook.counter: 

270 path = self.get_stats_path(index=hook.counter - 2) 

271 stats = controller.return_stats() 

272 store = True 

273 if hasattr(S.levels[0].sweep, 'comm') and S.levels[0].sweep.comm.rank > 0: 

274 store = False 

275 elif P.comm.rank > 0: 

276 store = False 

277 if store: 

278 with open(path, 'wb') as file: 

279 pickle.dump(stats, file) 

280 self.log(f'Stored stats in {path!r}', S) 

281 # print(stats) 

282 self.reset_stats(controller) 

283 self.counter = hook.counter 

284 

285 def post_run_processing(self, controller, S, **kwargs): 

286 self.post_step_processing(controller, S, **kwargs) 

287 

288 stats = self.merge_all_stats(controller) 

289 

290 def return_stats(): 

291 return stats 

292 

293 controller.return_stats = return_stats