from IPython.parallel import Client

class ClusterPool(object):

  def __init__(self, *args, **kwargs):
    self.client = Client(*args, **kwargs)
    self.lbview = self.client.load_balanced_view()
    self.chunksize = 1

  def _map(self, func, args, chunksize = None, block = False):
    if chunksize == None:
      chunksize = self.chunksize or 1
    return self.lbview.map(func, args, block = block, chunksize = chunksize, ordered = ordered)

  def map(self, func, args, chunksize = None):
    map = self.lbview.map
    return iter(map(func, args, chunksize = chunksize or self.chunksize, block = True, ordered = True))

  def imap(self, func, args, chunksize = None):
    map = self.lbview.map
    return iter(map(func, args, chunksize = chunksize or self.chunksize, block = False, ordered = True))

  def imap_unordered(self, func, args, chunksize = None):
    map = self.lbview.map
    return iter(map(func, args, chunksize = chunksize or self.chunksize, block = False, ordered = False))

from IPython.parallel import Client
from IPython.parallel.util import interactive

@interactive  # The direct view pushes to a particular namespace (i.e., not this module's NS). Thus, we must use this decorator for FUNC to be found in the correct namespace.
def func_int(*args):
  global FUNC, SHARED_DATA
  return FUNC(SHARED_DATA, *args)

class ClusterPool(object):

  def __init__(self, *args, **kwargs):
    self.client = Client(*args, **kwargs)
    self.lbview = self.client.load_balanced_view()
    self.chunksize = 1

  def map_with_shared_data(self, func, shared_data, args, chunksize = None):
    """Map a function over each in a set of arguments, also passing a constant shared variable to each invocation."""
    # no imap with shared data, since we couldn't guarantee the integrity of FUNC and SHARED_DATA
    self.dview.push(dict(FUNC = func, SHARED_DATA = shared_data), block = True)
    return self.lbview.map_sync(func_int, args, chunksize = chunksize or self.chunksize, ordered = True)

  def map(self, func, args, chunksize = None):
    map = self.lbview.map
    return iter(map(func, args, chunksize = chunksize or self.chunksize, block = True, ordered = True))

  def imap(self, func, args, chunksize = None):
    map = self.lbview.map
    return iter(map(func, args, chunksize = chunksize or self.chunksize, block = False, ordered = True))

  def imap_unordered(self, func, args, chunksize = None):
    map = self.lbview.map
    return iter(map(func, args, chunksize = chunksize or self.chunksize, block = False, ordered = False))

Add map_with_shared_data to IPython cluster pool, to reduce network traffic for large models.