1""" Class for handling several simultaneous jobs. 

2The class has been tested on Niflheim-opteron4. 

3""" 

4from multiprocessing import Pool 

5import time 

6from ase.io import write, read 

7 

8 

9class MultiprocessingRun: 

10 """Class that allows for the simultaneous relaxation of 

11 several candidates on a cluster. Best used if each individual 

12 calculation is too small for using a queueing system. 

13 

14 Parameters: 

15  

16 data_connection: DataConnection object. 

17  

18 tmp_folder: Folder for temporary files. 

19  

20 n_simul: The number of simultaneous relaxations. 

21  

22 relax_function: The relaxation function. This needs to return 

23 the filename of the relaxed structure. 

24 """ 

25 def __init__(self, data_connection, relax_function, 

26 tmp_folder, n_simul=None): 

27 self.dc = data_connection 

28 self.pool = Pool(n_simul) 

29 self.relax_function = relax_function 

30 self.tmp_folder = tmp_folder 

31 self.results = [] 

32 

33 def relax(self, a): 

34 """Relax the atoms object a by submitting the relaxation 

35 to the pool of cpus.""" 

36 self.dc.mark_as_queued(a) 

37 fname = '{0}/cand{1}.traj'.format(self.tmp_folder, 

38 a.info['confid']) 

39 write(fname, a) 

40 self.results.append(self.pool.apply_async(self.relax_function, 

41 [fname])) 

42 self._cleanup() 

43 

44 def _cleanup(self): 

45 for r in self.results: 

46 if r.ready() and r.successful(): 

47 fname = r.get() 

48 a = read(fname) 

49 self.dc.add_relaxed_step(a) 

50 self.results.remove(r) 

51 

52 def finish_all(self): 

53 """Checks that all calculations are finished, if not 

54 wait and check again. Return when all are finished.""" 

55 while len(self.results) > 0: 

56 self._cleanup() 

57 time.sleep(2.)