1""" Class for handling interaction with the PBS queuing system.""" 

2from ase.io import write 

3import os 

4from ase.io.trajectory import Trajectory 

5from subprocess import Popen, PIPE 

6import time 

7 

8 

9class PBSQueueRun: 

10 

11 """ Class for communicating with the commonly used PBS queing system 

12 at a computer cluster. 

13 

14 The user needs to supply a job file generator which takes 

15 as input a job name and the relative path to the traj 

16 file which is to be locally optimized. The function returns 

17 the job script as text. 

18 If the traj file is called f the job must write a file 

19 f[:-5] + '_done.traj' which is then read by this object. 

20 

21 Parameters: 

22 

23 data_connection: The DataConnection object. 

24 tmp_folder: Temporary folder for all calculations 

25 job_prefix: Prefix of the job submitted. This identifier is used 

26 to determine how many jobs are currently running. 

27 n_simul: The number of simultaneous jobs to keep in the queuing system. 

28 job_template_generator: The function generating the job file. 

29 This function should return the content of the job file as a 

30 string. 

31 qsub_command: The name of the qsub command (default qsub). 

32 qstat_command: The name of the qstat command (default qstat). 

33 """ 

34 

35 def __init__(self, data_connection, tmp_folder, job_prefix, 

36 n_simul, job_template_generator, 

37 qsub_command='qsub', qstat_command='qstat', 

38 find_neighbors=None, perform_parametrization=None): 

39 self.dc = data_connection 

40 self.job_prefix = job_prefix 

41 self.n_simul = n_simul 

42 self.job_template_generator = job_template_generator 

43 self.qsub_command = qsub_command 

44 self.qstat_command = qstat_command 

45 self.tmp_folder = tmp_folder 

46 self.find_neighbors = find_neighbors 

47 self.perform_parametrization = perform_parametrization 

48 self.__cleanup__() 

49 

50 def relax(self, a): 

51 """ Add a structure to the queue. This method does not fail 

52 if sufficient jobs are already running, but simply 

53 submits the job. """ 

54 self.__cleanup__() 

55 self.dc.mark_as_queued(a) 

56 if not os.path.isdir(self.tmp_folder): 

57 os.mkdir(self.tmp_folder) 

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

59 a.info['confid']) 

60 write(fname, a) 

61 job_name = '{0}_{1}'.format(self.job_prefix, a.info['confid']) 

62 fd = open('tmp_job_file.job', 'w') 

63 fd.write(self.job_template_generator(job_name, fname)) 

64 fd.close() 

65 os.system('{0} tmp_job_file.job'.format(self.qsub_command)) 

66 

67 def enough_jobs_running(self): 

68 """ Determines if sufficient jobs are running. """ 

69 return self.number_of_jobs_running() >= self.n_simul 

70 

71 def number_of_jobs_running(self): 

72 """ Determines how many jobs are running. The user 

73 should use this or the enough_jobs_running method 

74 to verify that a job needs to be started before 

75 calling the relax method.""" 

76 self.__cleanup__() 

77 p = Popen(['`which {0}` -u `whoami`'.format(self.qstat_command)], 

78 shell=True, stdin=PIPE, stdout=PIPE, stderr=PIPE, 

79 close_fds=True, universal_newlines=True) 

80 fout = p.stdout 

81 lines = fout.readlines() 

82 n_running = 0 

83 for line in lines: 

84 if line.find(self.job_prefix) != -1: 

85 n_running += 1 

86 return n_running 

87 

88 def __cleanup__(self): 

89 """ Tries to load in structures previously 

90 submitted to the queing system. """ 

91 confs = self.dc.get_all_candidates_in_queue() 

92 for c in confs: 

93 fdone = '{0}/cand{1}_done.traj'.format(self.tmp_folder, 

94 c) 

95 if os.path.isfile(fdone) and os.path.getsize(fdone) > 0: 

96 try: 

97 a = [] 

98 niter = 0 

99 while len(a) == 0 and niter < 5: 

100 t = Trajectory(fdone, 'r') 

101 a = [ats for ats in t] 

102 if len(a) == 0: 

103 time.sleep(1.) 

104 niter += 1 

105 if len(a) == 0: 

106 txt = 'Could not read candidate ' + \ 

107 '{0} from the filesystem'.format(c) 

108 raise IOError(txt) 

109 a = a[-1] 

110 a.info['confid'] = c 

111 self.dc.add_relaxed_step( 

112 a, 

113 find_neighbors=self.find_neighbors, 

114 perform_parametrization=self.perform_parametrization) 

115 except IOError as e: 

116 print(e)