parallel magic: added draft of %pmap (with hack); parallel kernels no…

…w a subkernel; magickernel.do_function_direct prototype

jupyter_kernel/magickernel.py

@@ -118,6 +118,13 @@ def do_execute_file(self, filename):
         """
         self.Error("This language does not support \"%run filename\".")
 
+    def do_function_direct(self, function_name, args):
+        """
+        Call a function in the kernel language with args (as a single item).
+        """
+        #self.Error("This language does not support \"%pmap function args\".")
+        return (function_name, args)
+
     def restart_kernel(self):
         """Restart the kernel"""
         pass

jupyter_kernel/magics/parallel_magic.py

@@ -15,6 +15,7 @@ def __getitem__(self, item):
 class ParallelMagic(Magic):
     client = None
     view = None
+    view_load_balanced = None
     module_name = None
     class_name = None
     kernel_name = None
@@ -65,6 +66,7 @@ def line_parallel(self, module_name, class_name, kernel_name="default", ids=None
                 for item in ids_slice:
                     view.append(self.client[item])
                 self.view = view
+        self.view_load_balanced = self.client.load_balanced_view()
         self.module_name = module_name
         self.class_name = class_name
         self.kernel_name = kernel_name
@@ -74,6 +76,7 @@ def line_parallel(self, module_name, class_name, kernel_name="default", ids=None
 except:
     kernels = {}
 from %(module_name)s import %(class_name)s
+%(class_name)s.subkernel(get_ipython().parent)
 kernels['%(kernel_name)s'] = %(class_name)s()
 """ % {"module_name": module_name, 
        "class_name": class_name,
@@ -145,6 +148,22 @@ def cell_px(self, kernel_name=None, evaluate=False):
             kernel_name, self._clean_code(self.code))]
         self.evaluate = evaluate
 
+    def line_pmap(self, function_name, args, kernel_name=None):
+        """
+        %pmap FUNCTION [ARGS1, ARGS2, ...] - call FUNCTION
+        """
+        if kernel_name is None:
+            kernel_name = self.kernel_name
+
+        # FIXME: horrible hack
+        self.view.execute("""
+import os
+os.kernels = kernels
+""")
+        self.retval = self.view_load_balanced.map_async(
+            lambda arg, kname=kernel_name, fname=function_name: os.kernels[kname].do_function_direct(fname, arg),
+            eval(args))
+
     def post_process(self, retval):
         try:
             ## any will crash on numpy arrays