[yt-svn] commit/yt: MatthewTurk: Adding a 'root' option to mpi_bcast, and adding some niceties to workgroups.

[Bitbucket]

1 new commit in yt:


https://bitbucket.org/yt_analysis/yt/changeset/271456ace2d6/
changeset:   271456ace2d6
branch:      yt
user:        MatthewTurk
date:        2012-05-22 23:29:06
summary:     Adding a 'root' option to mpi_bcast, and adding some niceties to workgroups.
affected #:  1 file

diff -r e4de6464d110fcd21a9c11e6bbae6157f811647f -r 271456ace2d6c25a1ee3363ef7ac28fb420ead16 yt/utilities/parallel_tools/parallel_analysis_interface.py
--- a/yt/utilities/parallel_tools/parallel_analysis_interface.py
+++ b/yt/utilities/parallel_tools/parallel_analysis_interface.py
@@ -288,7 +288,7 @@
         if size is None:
             size = len(self.available_ranks)
         if len(self.available_ranks) < size:
-            print 'Not enough resources available'
+            print 'Not enough resources available', size, self.available_ranks
             raise RuntimeError
         if ranks is None:
             ranks = [self.available_ranks.pop(0) for i in range(size)]
@@ -315,6 +315,26 @@
         for wg in self.workgroups:
             self.free_workgroup(wg)
 
+    @classmethod
+    def from_sizes(cls, sizes):
+        sizes = ensure_list(sizes)
+        pool = cls()
+        rank = pool.comm.rank
+        for i,size in enumerate(sizes):
+            if iterable(size):
+                size, name = size
+            else:
+                name = "workgroup_%02i" % i
+            pool.add_workgroup(size, name = name)
+        for wg in pool.workgroups:
+            if rank in wg.ranks: workgroup = wg
+        return pool, workgroup
+
+    def __getitem__(self, key):
+        for wg in self.workgroups:
+            if wg.name == key: return wg
+        raise KeyError(key)
+
 class ResultsStorage(object):
     slots = ['result', 'result_id']
     result = None
@@ -517,24 +537,24 @@
         raise NotImplementedError
 
     @parallel_passthrough
-    def mpi_bcast(self, data):
+    def mpi_bcast(self, data, root = 0):
         # The second check below makes sure that we know how to communicate
         # this type of array. Otherwise, we'll pickle it.
         if isinstance(data, na.ndarray) and \
                 get_mpi_type(data.dtype) is not None:
-            if self.comm.rank == 0:
+            if self.comm.rank == root:
                 info = (data.shape, data.dtype)
             else:
                 info = ()
-            info = self.comm.bcast(info, root=0)
-            if self.comm.rank != 0:
+            info = self.comm.bcast(info, root=root)
+            if self.comm.rank != root:
                 data = na.empty(info[0], dtype=info[1])
             mpi_type = get_mpi_type(info[1])
-            self.comm.Bcast([data, mpi_type], root = 0)
+            self.comm.Bcast([data, mpi_type], root = root)
             return data
         else:
             # Use pickled methods.
-            data = self.comm.bcast(data, root = 0)
+            data = self.comm.bcast(data, root = root)
             return data
 
     def preload(self, grids, fields, io_handler):

Repository URL: https://bitbucket.org/yt_analysis/yt/

--

This is a commit notification from bitbucket.org. You are receiving
this because you have the service enabled, addressing the recipient of
this email.