add tests for parallel magics

tests revealed some changes that should be made:

%px also prints remote stdout in blocking mode,
and actually raises the RemoteError (not just display)
if there is one

IPython/extensions/parallelmagic.py

@@ -1,276 +1,290 b''
 #!/usr/bin/env python
 # encoding: utf-8
 
 """Magic command interface for interactive parallel work."""
 
 #-----------------------------------------------------------------------------
 #  Copyright (C) 2008-2009  The IPython Development Team
 #
 #  Distributed under the terms of the BSD License.  The full license is in
 #  the file COPYING, distributed as part of this software.
 #-----------------------------------------------------------------------------
 
 #-----------------------------------------------------------------------------
 # Imports
 #-----------------------------------------------------------------------------
 
 import ast
 import re
 
 from IPython.core.plugin import Plugin
 from IPython.utils.traitlets import Bool, Any, Instance
 from IPython.testing import decorators as testdec
 
 #-----------------------------------------------------------------------------
 # Definitions of magic functions for use with IPython
 #-----------------------------------------------------------------------------
 
 
 NO_ACTIVE_VIEW = """
 Use activate() on a DirectView object to activate it for magics.
 """
 
 
 class ParalleMagic(Plugin):
     """A component to manage the %result, %px and %autopx magics."""
 
-    active_view = Any()
+    active_view = Instance('IPython.parallel.client.view.DirectView')
     verbose = Bool(False, config=True)
     shell = Instance('IPython.core.interactiveshell.InteractiveShellABC')
 
     def __init__(self, shell=None, config=None):
         super(ParalleMagic, self).__init__(shell=shell, config=config)
         self._define_magics()
         # A flag showing if autopx is activated or not
         self.autopx = False
 
     def _define_magics(self):
         """Define the magic functions."""
         self.shell.define_magic('result', self.magic_result)
         self.shell.define_magic('px', self.magic_px)
         self.shell.define_magic('autopx', self.magic_autopx)
 
     @testdec.skip_doctest
     def magic_result(self, ipself, parameter_s=''):
         """Print the result of command i on all engines..
 
         To use this a :class:`DirectView` instance must be created 
         and then activated by calling its :meth:`activate` method.
 
         Then you can do the following::
 
             In [23]: %result
             Out[23]: 
             <Results List>
             [0] In [6]: a = 10
             [1] In [6]: a = 10
     
             In [22]: %result 6
             Out[22]: 
             <Results List>
             [0] In [6]: a = 10
             [1] In [6]: a = 10
         """
         if self.active_view is None:
             print NO_ACTIVE_VIEW
             return
 
         try:
             index = int(parameter_s)
         except:
             index = None
         result = self.active_view.get_result(index)
         return result
 
     @testdec.skip_doctest
     def magic_px(self, ipself, parameter_s=''):
         """Executes the given python command in parallel.
 
         To use this a :class:`DirectView` instance must be created 
         and then activated by calling its :meth:`activate` method.
         
         Then you can do the following::
 
             In [24]: %px a = 5
             Parallel execution on engines: all
             Out[24]: 
             <Results List>
             [0] In [7]: a = 5
             [1] In [7]: a = 5
         """
 
         if self.active_view is None:
             print NO_ACTIVE_VIEW
             return
         print "Parallel execution on engines: %s" % self.active_view.targets
-        result = self.active_view.execute(parameter_s)
-        return result
+        result = self.active_view.execute(parameter_s, block=False)
+        if self.active_view.block:
+            result.get()
+            self._maybe_display_output(result)
 
     @testdec.skip_doctest
     def magic_autopx(self, ipself, parameter_s=''):
         """Toggles auto parallel mode.
 
         To use this a :class:`DirectView` instance must be created 
         and then activated by calling its :meth:`activate` method. Once this
         is called, all commands typed at the command line are send to
         the engines to be executed in parallel. To control which engine
         are used, set the ``targets`` attributed of the multiengine client
         before entering ``%autopx`` mode.
 
         Then you can do the following::
 
             In [25]: %autopx
             %autopx to enabled
 
             In [26]: a = 10
-            <Results List>
-            [0] In [8]: a = 10
-            [1] In [8]: a = 10
+            Parallel execution on engines: [0,1,2,3]
+            In [27]: print a
+            Parallel execution on engines: [0,1,2,3]
+            [stdout:0] 10
+            [stdout:1] 10
+            [stdout:2] 10
+            [stdout:3] 10
 
 
             In [27]: %autopx
             %autopx disabled
         """
         if self.autopx:
             self._disable_autopx()
         else:
             self._enable_autopx()
 
     def _enable_autopx(self):
         """Enable %autopx mode by saving the original run_cell and installing 
         pxrun_cell.
         """
         if self.active_view is None:
             print NO_ACTIVE_VIEW
             return
 
         # override run_cell and run_code
         self._original_run_cell = self.shell.run_cell
         self.shell.run_cell = self.pxrun_cell
         self._original_run_code = self.shell.run_code
         self.shell.run_code = self.pxrun_code
 
         self.autopx = True
         print "%autopx enabled"
     
     def _disable_autopx(self):
         """Disable %autopx by restoring the original InteractiveShell.run_cell.
         """
         if self.autopx:
             self.shell.run_cell = self._original_run_cell
             self.shell.run_code = self._original_run_code
             self.autopx = False
             print "%autopx disabled"
 
     def _maybe_display_output(self, result):
         """Maybe display the output of a parallel result.
 
         If self.active_view.block is True, wait for the result
         and display the result.  Otherwise, this is a noop.
         """
-        if self.active_view.block:
-            try:
-                result.get()
-            except:
-                self.shell.showtraceback()
-                return True
-            else:
-                targets = self.active_view.targets
-                if isinstance(targets, int):
-                    targets = [targets]
-                if targets == 'all':
-                    targets = self.active_view.client.ids
-                stdout = [s.rstrip() for s in result.stdout]
-                if any(stdout):
-                    for i,eid in enumerate(targets):
-                        print '[stdout:%i]'%eid, stdout[i]
-        return False
+        targets = self.active_view.targets
+        if isinstance(targets, int):
+            targets = [targets]
+        if targets == 'all':
+            targets = self.active_view.client.ids
+        stdout = [s.rstrip() for s in result.stdout]
+        if any(stdout):
+            for i,eid in enumerate(targets):
+                print '[stdout:%i]'%eid, stdout[i]
 
 
     def pxrun_cell(self, cell, store_history=True):
         """drop-in replacement for InteractiveShell.run_cell.
         
         This executes code remotely, instead of in the local namespace.
 
         See InteractiveShell.run_cell for details.
         """
         ipself = self.shell
         raw_cell = cell
         with ipself.builtin_trap:
             cell = ipself.prefilter_manager.prefilter_lines(cell)
         
             # Store raw and processed history
             if store_history:
                 ipself.history_manager.store_inputs(ipself.execution_count, 
                                                   cell, raw_cell)
 
             # ipself.logger.log(cell, raw_cell)
         
             cell_name = ipself.compile.cache(cell, ipself.execution_count)
         
             try:
                 code_ast = ast.parse(cell, filename=cell_name)
             except (OverflowError, SyntaxError, ValueError, TypeError, MemoryError):
                 # Case 1
                 ipself.showsyntaxerror()
                 ipself.execution_count += 1
                 return None
             except NameError:
                 # ignore name errors, because we don't know the remote keys
                 pass
 
         if store_history:
             # Write output to the database. Does nothing unless
             # history output logging is enabled.
             ipself.history_manager.store_output(ipself.execution_count)
             # Each cell is a *single* input, regardless of how many lines it has
             ipself.execution_count += 1
         
         if re.search(r'get_ipython\(\)\.magic\(u?"%?autopx', cell):
             self._disable_autopx()
             return False
         else:
             try:
                 result = self.active_view.execute(cell, block=False)
             except:
                 ipself.showtraceback()
-                return False
+                return True
             else:
-                return self._maybe_display_output(result)
+                if self.active_view.block:
+                    try:
+                        result.get()
+                    except:
+                        self.shell.showtraceback()
+                        return True
+                    else:
+                        self._maybe_display_output(result)
+                return False
 
     def pxrun_code(self, code_obj, post_execute=True):
         """drop-in replacement for InteractiveShell.run_code.
 
         This executes code remotely, instead of in the local namespace.
 
         See InteractiveShell.run_code for details.
         """
         ipself = self.shell
         # check code object for the autopx magic
         if 'get_ipython' in code_obj.co_names and 'magic' in code_obj.co_names and \
             any( [ isinstance(c, basestring) and 'autopx' in c for c in code_obj.co_consts ]):
             self._disable_autopx()
             return False
         else:
             try:
                 result = self.active_view.execute(code_obj, block=False)
             except:
                 ipself.showtraceback()
-                return False
+                return True
             else:
-                return self._maybe_display_output(result)
+                if self.active_view.block:
+                    try:
+                        result.get()
+                    except:
+                        self.shell.showtraceback()
+                        return True
+                    else:
+                        self._maybe_display_output(result)
+                return False
 
 
 
 
 _loaded = False
 
 
 def load_ipython_extension(ip):
     """Load the extension in IPython."""
     global _loaded
     if not _loaded:
         plugin = ParalleMagic(shell=ip, config=ip.config)
         ip.plugin_manager.register_plugin('parallelmagic', plugin)
         _loaded = True
 

IPython/parallel/tests/test_view.py

@@ -1,302 +1,413 b''
 """test View objects"""
 #-------------------------------------------------------------------------------
 #  Copyright (C) 2011  The IPython Development Team
 #
 #  Distributed under the terms of the BSD License.  The full license is in
 #  the file COPYING, distributed as part of this software.
 #-------------------------------------------------------------------------------
 
 #-------------------------------------------------------------------------------
 # Imports
 #-------------------------------------------------------------------------------
 
+import sys
 import time
 from tempfile import mktemp
+from StringIO import StringIO
 
 import zmq
 
 from IPython import parallel  as pmod
 from IPython.parallel import error
 from IPython.parallel import AsyncResult, AsyncHubResult, AsyncMapResult
 from IPython.parallel import LoadBalancedView, DirectView
 from IPython.parallel.util import interactive
 
 from IPython.parallel.tests import add_engines
 
 from .clienttest import ClusterTestCase, segfault, wait, skip_without
 
 def setup():
     add_engines(3)
 
 class TestView(ClusterTestCase):
     
     def test_segfault_task(self):
         """test graceful handling of engine death (balanced)"""
         # self.add_engines(1)
         ar = self.client[-1].apply_async(segfault)
         self.assertRaisesRemote(error.EngineError, ar.get)
         eid = ar.engine_id
         while eid in self.client.ids:
             time.sleep(.01)
             self.client.spin()
     
     def test_segfault_mux(self):
         """test graceful handling of engine death (direct)"""
         # self.add_engines(1)
         eid = self.client.ids[-1]
         ar = self.client[eid].apply_async(segfault)
         self.assertRaisesRemote(error.EngineError, ar.get)
         eid = ar.engine_id
         while eid in self.client.ids:
             time.sleep(.01)
             self.client.spin()
     
     def test_push_pull(self):
         """test pushing and pulling"""
         data = dict(a=10, b=1.05, c=range(10), d={'e':(1,2),'f':'hi'})
         t = self.client.ids[-1]
         v = self.client[t]
         push = v.push
         pull = v.pull
         v.block=True
         nengines = len(self.client)
         push({'data':data})
         d = pull('data')
         self.assertEquals(d, data)
         self.client[:].push({'data':data})
         d = self.client[:].pull('data', block=True)
         self.assertEquals(d, nengines*[data])
         ar = push({'data':data}, block=False)
         self.assertTrue(isinstance(ar, AsyncResult))
         r = ar.get()
         ar = self.client[:].pull('data', block=False)
         self.assertTrue(isinstance(ar, AsyncResult))
         r = ar.get()
         self.assertEquals(r, nengines*[data])
         self.client[:].push(dict(a=10,b=20))
         r = self.client[:].pull(('a','b'), block=True)
         self.assertEquals(r, nengines*[[10,20]])
     
     def test_push_pull_function(self):
         "test pushing and pulling functions"
         def testf(x):
             return 2.0*x
         
         t = self.client.ids[-1]
         v = self.client[t]
         v.block=True
         push = v.push
         pull = v.pull
         execute = v.execute
         push({'testf':testf})
         r = pull('testf')
         self.assertEqual(r(1.0), testf(1.0))
         execute('r = testf(10)')
         r = pull('r')
         self.assertEquals(r, testf(10))
         ar = self.client[:].push({'testf':testf}, block=False)
         ar.get()
         ar = self.client[:].pull('testf', block=False)
         rlist = ar.get()
         for r in rlist:
             self.assertEqual(r(1.0), testf(1.0))
         execute("def g(x): return x*x")
         r = pull(('testf','g'))
         self.assertEquals((r[0](10),r[1](10)), (testf(10), 100))
     
     def test_push_function_globals(self):
         """test that pushed functions have access to globals"""
         @interactive
         def geta():
             return a
         # self.add_engines(1)
         v = self.client[-1]
         v.block=True
         v['f'] = geta
         self.assertRaisesRemote(NameError, v.execute, 'b=f()')
         v.execute('a=5')
         v.execute('b=f()')
         self.assertEquals(v['b'], 5)
     
     def test_push_function_defaults(self):
         """test that pushed functions preserve default args"""
         def echo(a=10):
             return a
         v = self.client[-1]
         v.block=True
         v['f'] = echo
         v.execute('b=f()')
         self.assertEquals(v['b'], 10)
 
     def test_get_result(self):
         """test getting results from the Hub."""
         c = pmod.Client(profile='iptest')
         # self.add_engines(1)
         t = c.ids[-1]
         v = c[t]
         v2 = self.client[t]
         ar = v.apply_async(wait, 1)
         # give the monitor time to notice the message
         time.sleep(.25)
         ahr = v2.get_result(ar.msg_ids)
         self.assertTrue(isinstance(ahr, AsyncHubResult))
         self.assertEquals(ahr.get(), ar.get())
         ar2 = v2.get_result(ar.msg_ids)
         self.assertFalse(isinstance(ar2, AsyncHubResult))
         c.spin()
         c.close()
     
     def test_run_newline(self):
         """test that run appends newline to files"""
         tmpfile = mktemp()
         with open(tmpfile, 'w') as f:
             f.write("""def g():
                 return 5
                 """)
         v = self.client[-1]
         v.run(tmpfile, block=True)
         self.assertEquals(v.apply_sync(lambda f: f(), pmod.Reference('g')), 5)
 
     def test_apply_tracked(self):
         """test tracking for apply"""
         # self.add_engines(1)
         t = self.client.ids[-1]
         v = self.client[t]
         v.block=False
         def echo(n=1024*1024, **kwargs):
             with v.temp_flags(**kwargs):
                 return v.apply(lambda x: x, 'x'*n)
         ar = echo(1, track=False)
         self.assertTrue(isinstance(ar._tracker, zmq.MessageTracker))
         self.assertTrue(ar.sent)
         ar = echo(track=True)
         self.assertTrue(isinstance(ar._tracker, zmq.MessageTracker))
         self.assertEquals(ar.sent, ar._tracker.done)
         ar._tracker.wait()
         self.assertTrue(ar.sent)
         
     def test_push_tracked(self):
         t = self.client.ids[-1]
         ns = dict(x='x'*1024*1024)
         v = self.client[t]
         ar = v.push(ns, block=False, track=False)
         self.assertTrue(isinstance(ar._tracker, zmq.MessageTracker))
         self.assertTrue(ar.sent)
         
         ar = v.push(ns, block=False, track=True)
         self.assertTrue(isinstance(ar._tracker, zmq.MessageTracker))
         self.assertEquals(ar.sent, ar._tracker.done)
         ar._tracker.wait()
         self.assertTrue(ar.sent)
         ar.get()
         
     def test_scatter_tracked(self):
         t = self.client.ids
         x='x'*1024*1024
         ar = self.client[t].scatter('x', x, block=False, track=False)
         self.assertTrue(isinstance(ar._tracker, zmq.MessageTracker))
         self.assertTrue(ar.sent)
         
         ar = self.client[t].scatter('x', x, block=False, track=True)
         self.assertTrue(isinstance(ar._tracker, zmq.MessageTracker))
         self.assertEquals(ar.sent, ar._tracker.done)
         ar._tracker.wait()
         self.assertTrue(ar.sent)
         ar.get()
     
     def test_remote_reference(self):
         v = self.client[-1]
         v['a'] = 123
         ra = pmod.Reference('a')
         b = v.apply_sync(lambda x: x, ra)
         self.assertEquals(b, 123)
 
 
     def test_scatter_gather(self):
         view = self.client[:]
         seq1 = range(16)
         view.scatter('a', seq1)
         seq2 = view.gather('a', block=True)
         self.assertEquals(seq2, seq1)
         self.assertRaisesRemote(NameError, view.gather, 'asdf', block=True)
     
     @skip_without('numpy')
     def test_scatter_gather_numpy(self):
         import numpy
         from numpy.testing.utils import assert_array_equal, assert_array_almost_equal
         view = self.client[:]
         a = numpy.arange(64)
         view.scatter('a', a)
         b = view.gather('a', block=True)
         assert_array_equal(b, a)
     
     def test_map(self):
         view = self.client[:]
         def f(x):
             return x**2
         data = range(16)
         r = view.map_sync(f, data)
         self.assertEquals(r, map(f, data))
     
     def test_scatterGatherNonblocking(self):
         data = range(16)
         view = self.client[:]
         view.scatter('a', data, block=False)
         ar = view.gather('a', block=False)
         self.assertEquals(ar.get(), data)
     
     @skip_without('numpy')
     def test_scatter_gather_numpy_nonblocking(self):
         import numpy
         from numpy.testing.utils import assert_array_equal, assert_array_almost_equal
         a = numpy.arange(64)
         view = self.client[:]
         ar = view.scatter('a', a, block=False)
         self.assertTrue(isinstance(ar, AsyncResult))
         amr = view.gather('a', block=False)
         self.assertTrue(isinstance(amr, AsyncMapResult))
         assert_array_equal(amr.get(), a)
 
     def test_execute(self):
         view = self.client[:]
         # self.client.debug=True
         execute = view.execute
         ar = execute('c=30', block=False)
         self.assertTrue(isinstance(ar, AsyncResult))
         ar = execute('d=[0,1,2]', block=False)
         self.client.wait(ar, 1)
         self.assertEquals(len(ar.get()), len(self.client))
         for c in view['c']:
             self.assertEquals(c, 30)
     
     def test_abort(self):
         view = self.client[-1]
         ar = view.execute('import time; time.sleep(0.25)', block=False)
         ar2 = view.apply_async(lambda : 2)
         ar3 = view.apply_async(lambda : 3)
         view.abort(ar2)
         view.abort(ar3.msg_ids)
         self.assertRaises(error.TaskAborted, ar2.get)
         self.assertRaises(error.TaskAborted, ar3.get)
         
     def test_temp_flags(self):
         view = self.client[-1]
         view.block=True
         with view.temp_flags(block=False):
             self.assertFalse(view.block)
         self.assertTrue(view.block)
     
     def test_importer(self):
         view = self.client[-1]
         view.clear(block=True)
         with view.importer:
             import re
         
         @interactive
         def findall(pat, s):
             # this globals() step isn't necessary in real code
             # only to prevent a closure in the test
             return globals()['re'].findall(pat, s)
         
         self.assertEquals(view.apply_sync(findall, '\w+', 'hello world'), 'hello world'.split())
     
+    # parallel magic tests
+    
+    def test_magic_px_blocking(self):
+        ip = get_ipython()
+        v = self.client[-1]
+        v.activate()
+        v.block=True
+
+        ip.magic_px('a=5')
+        self.assertEquals(v['a'], 5)
+        ip.magic_px('a=10')
+        self.assertEquals(v['a'], 10)
+        sio = StringIO()
+        savestdout = sys.stdout
+        sys.stdout = sio
+        ip.magic_px('print a')
+        sys.stdout = savestdout
+        sio.read()
+        self.assertTrue('[stdout:%i]'%v.targets in sio.buf)
+        self.assertRaisesRemote(ZeroDivisionError, ip.magic_px, '1/0')
+
+    def test_magic_px_nonblocking(self):
+        ip = get_ipython()
+        v = self.client[-1]
+        v.activate()
+        v.block=False
+
+        ip.magic_px('a=5')
+        self.assertEquals(v['a'], 5)
+        ip.magic_px('a=10')
+        self.assertEquals(v['a'], 10)
+        sio = StringIO()
+        savestdout = sys.stdout
+        sys.stdout = sio
+        ip.magic_px('print a')
+        sys.stdout = savestdout
+        sio.read()
+        self.assertFalse('[stdout:%i]'%v.targets in sio.buf)
+        ip.magic_px('1/0')
+        ar = v.get_result(-1)
+        self.assertRaisesRemote(ZeroDivisionError, ar.get)
+    
+    def test_magic_autopx_blocking(self):
+        ip = get_ipython()
+        v = self.client[-1]
+        v.activate()
+        v.block=True
+
+        sio = StringIO()
+        savestdout = sys.stdout
+        sys.stdout = sio
+        ip.magic_autopx()
+        ip.run_cell('\n'.join(('a=5','b=10','c=0')))
+        ip.run_cell('print b')
+        ip.run_cell("b/c")
+        ip.run_code(compile('b*=2', '', 'single'))
+        ip.magic_autopx()
+        sys.stdout = savestdout
+        sio.read()
+        output = sio.buf.strip()
+        self.assertTrue(output.startswith('%autopx enabled'))
+        self.assertTrue(output.endswith('%autopx disabled'))
+        self.assertTrue('RemoteError: ZeroDivisionError' in output)
+        ar = v.get_result(-2)
+        self.assertEquals(v['a'], 5)
+        self.assertEquals(v['b'], 20)
+        self.assertRaisesRemote(ZeroDivisionError, ar.get)
+
+    def test_magic_autopx_nonblocking(self):
+        ip = get_ipython()
+        v = self.client[-1]
+        v.activate()
+        v.block=False
+
+        sio = StringIO()
+        savestdout = sys.stdout
+        sys.stdout = sio
+        ip.magic_autopx()
+        ip.run_cell('\n'.join(('a=5','b=10','c=0')))
+        ip.run_cell('print b')
+        ip.run_cell("b/c")
+        ip.run_code(compile('b*=2', '', 'single'))
+        ip.magic_autopx()
+        sys.stdout = savestdout
+        sio.read()
+        output = sio.buf.strip()
+        self.assertTrue(output.startswith('%autopx enabled'))
+        self.assertTrue(output.endswith('%autopx disabled'))
+        self.assertFalse('ZeroDivisionError' in output)
+        ar = v.get_result(-2)
+        self.assertEquals(v['a'], 5)
+        self.assertEquals(v['b'], 20)
+        self.assertRaisesRemote(ZeroDivisionError, ar.get)
+    
+    def test_magic_result(self):
+        ip = get_ipython()
+        v = self.client[-1]
+        v.activate()
+        v['a'] = 111
+        ra = v['a']
+        
+        ar = ip.magic_result()
+        self.assertEquals(ar.msg_ids, [v.history[-1]])
+        self.assertEquals(ar.get(), 111)
+        ar = ip.magic_result('-2')
+        self.assertEquals(ar.msg_ids, [v.history[-2]])
+        
+
+