Source Code for Module multi.mpi4py_processor

  1  ############################################################################### 
  2  #                                                                             # 
  3  # Copyright (C) 2007 Gary S Thompson                                          # 
  4  # Copyright (C) 2008,2010-2012 Edward d'Auvergne                              # 
  5  #                                                                             # 
  6  # This file is part of the program relax (http://www.nmr-relax.com).          # 
  7  #                                                                             # 
  8  # This program is free software: you can redistribute it and/or modify        # 
  9  # it under the terms of the GNU General Public License as published by        # 
 10  # the Free Software Foundation, either version 3 of the License, or           # 
 11  # (at your option) any later version.                                         # 
 12  #                                                                             # 
 13  # This program is distributed in the hope that it will be useful,             # 
 14  # but WITHOUT ANY WARRANTY; without even the implied warranty of              # 
 15  # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the               # 
 16  # GNU General Public License for more details.                                # 
 17  #                                                                             # 
 18  # You should have received a copy of the GNU General Public License           # 
 19  # along with this program.  If not, see <http://www.gnu.org/licenses/>.       # 
 20  #                                                                             # 
 21  ############################################################################### 
 22   
 23  # Module docstring. 
 24  """The MPI processor fabric via the mpi4py Python implementation.""" 
 25   
 26   
 27  # TODO: clone communicators & resize 
 28  # TODO: check exceptions on master 
 29   
 30  # Python module imports. 
 31  try: 
 32      from mpi4py import MPI 
 33  except ImportError: 
 34      MPI = None 
 35  import os 
 36  import sys 
 37   
 38  # relax module imports. 
 39  from multi.slave_commands import Exit_command 
 40  from multi.multi_processor_base import Multi_processor, Too_few_slaves_exception 
 41   
 42   
 43 -class Mpi4py_processor(Multi_processor): 
 44      """The mpi4py multi-processor class.""" 
 45   
 46 -    def __init__(self, processor_size, callback): 
 47          """Initialise the mpi4py processor.""" 
 48   
 49          mpi_processor_size = MPI.COMM_WORLD.size-1 
 50   
 51          if processor_size == -1: 
 52              processor_size = mpi_processor_size 
 53   
 54          # FIXME: needs better support in relax generates stack trace 
 55          if mpi_processor_size == 0: 
 56              raise Too_few_slaves_exception() 
 57   
 58          msg = 'warning: mpi4py_processor is using 1 masters and %d slave processors you requested %d slaves\n' 
 59          if processor_size != (mpi_processor_size): 
 60              print(msg % (mpi_processor_size, processor_size)) 
 61   
 62          super(Mpi4py_processor, self).__init__(processor_size=mpi_processor_size, callback=callback) 
 63   
 64          # Initialise a flag for determining if we are in the run() method or not. 
 65          self.in_main_loop = False 
 66   
 67   
 68 -    def _broadcast_command(self, command): 
 69          for i in range(1, MPI.COMM_WORLD.size): 
 70              if i != 0: 
 71                  MPI.COMM_WORLD.send(obj=command, dest=i) 
 72   
 73   
 74 -    def _ditch_all_results(self): 
 75          for i in range(1, MPI.COMM_WORLD.size): 
 76              if i != 0: 
 77                  while True: 
 78                      result = MPI.COMM_WORLD.recv(source=i) 
 79                      if result.completed: 
 80                          break 
 81   
 82   
 83 -    def abort(self): 
 84          MPI.COMM_WORLD.Abort() 
 85   
 86   
 87 -    def assert_on_master(self): 
 88          """Make sure that this is the master processor and not a slave. 
 89   
 90          @raises Exception:  If not on the master processor. 
 91          """ 
 92   
 93          # Check if this processor is a slave, and if so throw an exception. 
 94          if self.on_slave(): 
 95              msg = 'running on slave when expected master with MPI.rank == 0, rank was %d'% self.rank() 
 96              raise Exception(msg) 
 97   
 98   
 99 -    def exit(self, status=0): 
100          """Exit the mpi4py processor with the given status. 
101   
102          @keyword status:    The program exit status. 
103          @type status:       int 
104          """ 
105   
106          # Execution on the slave. 
107          if MPI.COMM_WORLD.rank != 0: 
108              # Catch sys.exit being called on an executing slave. 
109              if self.in_main_loop: 
110                  raise Exception('sys.exit unexpectedly called on slave!') 
111   
112              # Catch sys.exit 
113              else: 
114                  sys.stderr.write('\n') 
115                  sys.stderr.write('***********************************************\n') 
116                  sys.stderr.write('\n') 
117                  sys.stderr.write('warning sys.exit called before mpi4py main loop\n') 
118                  sys.stderr.write('\n') 
119                  sys.stderr.write('***********************************************\n') 
120                  sys.stderr.write('\n') 
121                  MPI.COMM_WORLD.Abort() 
122   
123          # Execution on the master. 
124          else: 
125              # Slave clean up. 
126              if MPI.Is_initialized() and not MPI.Is_finalized() and MPI.COMM_WORLD.rank == 0: 
127                  # Send the exit command to all slaves. 
128                  self._broadcast_command(Exit_command()) 
129   
130                  # Dump all results. 
131                  self._ditch_all_results() 
132   
133              # Exit the program with the given status. 
134              sys.exit(status) 
135   
136   
137 -    def get_intro_string(self): 
138          """Return the string to append to the end of the relax introduction string. 
139   
140          @return:    The string describing this Processor fabric. 
141          @rtype:     str 
142          """ 
143   
144          # Get the specific MPI version. 
145          version_info = MPI.Get_version() 
146   
147          # The vendor info. 
148          vendor = MPI.get_vendor() 
149          vendor_name = vendor[0] 
150          vendor_version = str(vendor[1][0]) 
151          for i in range(1, len(vendor[1])): 
152              vendor_version = vendor_version + '.%i' % vendor[1][i] 
153   
154          # Return the string. 
155          return "MPI %s.%s running via mpi4py with %i slave processors & 1 master.  Using %s %s." % (version_info[0], version_info[1], self.processor_size(), vendor_name, vendor_version) 
156   
157   
158 -    def get_name(self): 
159          return '%s-pid%s' % (MPI.Get_processor_name(), os.getpid()) 
160   
161   
162 -    def master_queue_command(self, command, dest): 
163          """Slave to master processor data transfer - send the result command from the slave. 
164   
165          @param command: The results command to send to the master. 
166          @type command:  Results_command instance 
167          @param dest:    The destination processor's rank. 
168          @type dest:     int 
169          """ 
170   
171          # Use a basic MPI send call to transfer the result command. 
172          MPI.COMM_WORLD.send(obj=command, dest=dest) 
173   
174   
175 -    def master_receive_result(self): 
176          """Slave to master processor data transfer - receive the result command from the slave. 
177   
178          This is invoked by the master processor. 
179   
180          @return:        The result command sent by the slave. 
181          @rtype:         Result_command instance 
182          """ 
183   
184          # Catch and return the result command. 
185          return MPI.COMM_WORLD.recv(source=MPI.ANY_SOURCE) 
186   
187   
188 -    def rank(self): 
189          return MPI.COMM_WORLD.rank 
190   
191   
192 -    def return_result_command(self, result_object): 
193          MPI.COMM_WORLD.send(obj=result_object, dest=0) 
194   
195   
196 -    def run(self): 
197          self.in_main_loop = True 
198          super(Mpi4py_processor, self).run() 
199          self.in_main_loop = False 
200   
201   
202 -    def slave_receive_commands(self): 
203          return MPI.COMM_WORLD.recv(source=0) 
204