Add WORKER_POOL environmental option to facilitate pysawsim.manager benchmarking.

[sawsim.git] / pysawsim / manager / mpi.py

diff --git a/pysawsim/manager/mpi.py b/pysawsim/manager/mpi.py
index d99f4156dd13464f5c54130717fee922e8182f2a..a398c4f90ce40464fe3e8f05c8beec6df2bc0899 100644 (file)
--- a/pysawsim/manager/mpi.py
+++ b/pysawsim/manager/mpi.py
@@ -21,7 +21,7 @@
 
 mpi4py_ is a Python wrapper around MPI.
 
-.. _mpi4py: http://mpi4py.scipy.org/http://mpi4py.scipy.org/vb>
+.. _mpi4py: http://mpi4py.scipy.org/
 
 The MPIManager data flow is a bit complicated, so I've made a
 diagram::
@@ -49,18 +49,21 @@ There is also a `free_queue` running from `receive-thread` to
 nodes are free (and therefore ready to receive new jobs).
 """
 
+import os
 from Queue import Queue, Empty
 import sys
 from threading import Thread
 
 try:
     from mpi4py import MPI
+    _ENABLED = True
+    _DISABLING_ERROR = None
     if MPI.COMM_WORLD.Get_rank() == 0:
         _SKIP = ''
     else:
         _SKIP = '  # doctest: +SKIP'
-except ImportError, MPI_error:
-    MPI = None
+except ImportError, _DISABLING_ERROR:
+    _ENABLED = False
     _SKIP = '  # doctest: +SKIP'
 
 from .. import log
@@ -74,13 +77,14 @@ RECEIVE_TAG = 101
 
 
 def MPI_worker_death():
-    if MPI == None:
+    if _ENABLED != True:
         return
     if MPI.COMM_WORLD.Get_rank() != 0:
         sys.exit(0)
 
 def _manager_check():
-    assert MPI != None, MPI_error
+    if _ENABLED == False:
+        raise _DISABLING_ERROR
     rank = MPI.COMM_WORLD.Get_rank()
     assert rank == 0, (
         'process %d should have been killed by an MPI_worker_death() call'
@@ -205,9 +209,10 @@ class MPIManager (ThreadManager):
                 'w.run()',
                 'w.teardown()',
                 ])
-        if worker_pool == None:
-            worker_pool = MPI.COMM_WORLD.Get_size()
-        comm = MPI.COMM_SELF.Spawn(
+        if worker_pool is None:
+            worker_pool = int(os.environ.get('WORKER_POOL',
+                                             MPI.COMM_WORLD.Get_size()))
+        comm = MPI.COMM_SELF.Spawn(  # locks with mpich2 if no mpd running
             sys.executable, args=['-c', spawn_script], maxprocs=worker_pool)
         rank = comm.Get_rank()
         assert rank == 0, rank