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 4f2b2f338fac934975bb018c48f0148c2562ed0f..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,6 +49,7 @@ 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
@@ -208,8 +209,9 @@ class MPIManager (ThreadManager):
                 'w.run()',
                 'w.teardown()',
                 ])
-        if worker_pool == None:
-            worker_pool = MPI.COMM_WORLD.Get_size()
+        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()