Add support for kwargs to elfi.set_client (#271)

CHANGELOG.rst

@@ -1,6 +1,10 @@
 Changelog
 =========
 
+dev
+---
+- Add support for kwargs to elfi.set_client
+
 0.7.1 (2018-04-11)
 ------------------
 - Implemented model selection (elfi.compare_models). See API documentation.

elfi/client.py

@@ -28,13 +28,22 @@ def get_client():
     return _client
 
 
-def set_client(client=None):
-    """Set the current ELFI client instance."""
+def set_client(client=None, **kwargs):
+    """Set the current ELFI client instance.
+
+    Parameters
+    ----------
+    client : ClientBase or str
+        Instance of a client from ClientBase,
+        or a string from ['native', 'multiprocessing', 'ipyparallel'].
+        If string, the respective constructor is called with `kwargs`.
+
+    """
     global _client
 
     if isinstance(client, str):
         m = importlib.import_module('elfi.clients.{}'.format(client))
-        client = m.Client()
+        client = m.Client(**kwargs)
 
     _client = client
 

elfi/clients/ipyparallel.py

@@ -25,7 +25,7 @@ class Client(elfi.client.ClientBase):
     http://ipyparallel.readthedocs.io
     """
 
-    def __init__(self, ipp_client=None):
+    def __init__(self, ipp_client=None, **kwargs):
         """Create an ipyparallel client for ELFI.
 
         Parameters
@@ -34,7 +34,7 @@ def __init__(self, ipp_client=None):
             Use this ipyparallel client with ELFI.
 
         """
-        self.ipp_client = ipp_client or ipp.Client()
+        self.ipp_client = ipp_client or ipp.Client(**kwargs)
         self.view = self.ipp_client.load_balanced_view()
 
         self.tasks = {}

elfi/clients/multiprocessing.py

@@ -18,7 +18,7 @@ def set_as_default():
 class Client(elfi.client.ClientBase):
     """Client based on Python's built-in multiprocessing module."""
 
-    def __init__(self, num_processes=None):
+    def __init__(self, num_processes=None, **kwargs):
         """Create a multiprocessing client.
 
         Parameters
@@ -27,7 +27,8 @@ def __init__(self, num_processes=None):
             Number of worker processes to use. Defaults to os.cpu_count().
 
         """
-        self.pool = multiprocessing.Pool(processes=num_processes)
+        num_processes = num_processes or kwargs.pop('processes', None)
+        self.pool = multiprocessing.Pool(processes=num_processes, **kwargs)
 
         self.tasks = {}
         self._id_counter = itertools.count()

elfi/clients/native.py

@@ -20,7 +20,7 @@ class Client(elfi.client.ClientBase):
     Responsible for sending computational graphs to be executed in an Executor
     """
 
-    def __init__(self):
+    def __init__(self, **kwargs):
         """Create a native client."""
         self.tasks = {}
         self._ids = itertools.count()

tests/unit/test_client.py

@@ -29,4 +29,15 @@ def test_batch_handler(simple_model):
     assert not np.array_equal(out0['k2'], out1['k2'])
 
 
+def test_multiprocessing_kwargs(simple_model):
+    m = simple_model
+    num_proc = 2
+    elfi.set_client('multiprocessing', num_processes=num_proc)
+    rej = elfi.Rejection(m['k1'])
+    assert rej.client.num_cores == num_proc
+
+    elfi.set_client('multiprocessing', processes=num_proc)
+    rej = elfi.Rejection(m['k1'])
+    assert rej.client.num_cores == num_proc
+
 # TODO: add testing that client is cleared from tasks after they are retrieved