From 543a3cd093d8e7986c540a011a02a4f5e5900a55 Mon Sep 17 00:00:00 2001
From: Douglas Blank <doug.blank@gmail.com>
Date: Fri, 12 Sep 2014 17:47:34 -0400
Subject: [PATCH] parallel magic: added draft of %pmap (with hack); parallel
 kernels now a subkernel; magickernel.do_function_direct prototype

---
 jupyter_kernel/magickernel.py           |  7 +++++++
 jupyter_kernel/magics/parallel_magic.py | 19 +++++++++++++++++++
 2 files changed, 26 insertions(+)

diff --git a/jupyter_kernel/magickernel.py b/jupyter_kernel/magickernel.py
index f6e3a0e9..9c656f13 100644
--- a/jupyter_kernel/magickernel.py
+++ b/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
diff --git a/jupyter_kernel/magics/parallel_magic.py b/jupyter_kernel/magics/parallel_magic.py
index 0a54b4bc..f5a502f4 100644
--- a/jupyter_kernel/magics/parallel_magic.py
+++ b/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