##// END OF EJS Templates
add `%%px --local` for local execution...
add `%%px --local` for local execution So that a cell can be run on engines *as well as* the local namespace. Prompted by SO Question [14184621](http://stackoverflow.com/questions/14184621).

File last commit:

r8996:c2fa6e5a
r8996:c2fa6e5a
Show More
magics.py
430 lines | 13.1 KiB | text/x-python | PythonLexer
# encoding: utf-8
"""
=============
parallelmagic
=============
Magic command interface for interactive parallel work.
Usage
=====
``%autopx``
{AUTOPX_DOC}
``%px``
{PX_DOC}
``%pxresult``
{RESULT_DOC}
``%pxconfig``
{CONFIG_DOC}
"""
#-----------------------------------------------------------------------------
# Copyright (C) 2008 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.error import UsageError
from IPython.core.magic import Magics
from IPython.core import magic_arguments
from IPython.testing.skipdoctest import skip_doctest
#-----------------------------------------------------------------------------
# Definitions of magic functions for use with IPython
#-----------------------------------------------------------------------------
NO_LAST_RESULT = "%pxresult recalls last %px result, which has not yet been used."
def exec_args(f):
"""decorator for adding block/targets args for execution
applied to %pxconfig and %%px
"""
args = [
magic_arguments.argument('-b', '--block', action="store_const",
const=True, dest='block',
help="use blocking (sync) execution",
),
magic_arguments.argument('-a', '--noblock', action="store_const",
const=False, dest='block',
help="use non-blocking (async) execution",
),
magic_arguments.argument('-t', '--targets', type=str,
help="specify the targets on which to execute",
),
magic_arguments.argument('--local', action="store_const",
const=True, dest="local",
help="also execute the cell in the local namespace",
),
magic_arguments.argument('--verbose', action="store_const",
const=True, dest="set_verbose",
help="print a message at each execution",
),
magic_arguments.argument('--no-verbose', action="store_const",
const=False, dest="set_verbose",
help="don't print any messages",
),
]
for a in args:
f = a(f)
return f
def output_args(f):
"""decorator for output-formatting args
applied to %pxresult and %%px
"""
args = [
magic_arguments.argument('-r', action="store_const", dest='groupby',
const='order',
help="collate outputs in order (same as group-outputs=order)"
),
magic_arguments.argument('-e', action="store_const", dest='groupby',
const='engine',
help="group outputs by engine (same as group-outputs=engine)"
),
magic_arguments.argument('--group-outputs', dest='groupby', type=str,
choices=['engine', 'order', 'type'], default='type',
help="""Group the outputs in a particular way.
Choices are:
type: group outputs of all engines by type (stdout, stderr, displaypub, etc.).
engine: display all output for each engine together.
order: like type, but individual displaypub output from each engine is collated.
For example, if multiple plots are generated by each engine, the first
figure of each engine will be displayed, then the second of each, etc.
"""
),
magic_arguments.argument('-o', '--out', dest='save_name', type=str,
help="""store the AsyncResult object for this computation
in the global namespace under this name.
"""
),
]
for a in args:
f = a(f)
return f
class ParallelMagics(Magics):
"""A set of magics useful when controlling a parallel IPython cluster.
"""
# magic-related
magics = None
registered = True
# suffix for magics
suffix = ''
# A flag showing if autopx is activated or not
_autopx = False
# the current view used by the magics:
view = None
# last result cache for %pxresult
last_result = None
# verbose flag
verbose = False
def __init__(self, shell, view, suffix=''):
self.view = view
self.suffix = suffix
# register magics
self.magics = dict(cell={},line={})
line_magics = self.magics['line']
px = 'px' + suffix
if not suffix:
# keep %result for legacy compatibility
line_magics['result'] = self.result
line_magics['pxresult' + suffix] = self.result
line_magics[px] = self.px
line_magics['pxconfig' + suffix] = self.pxconfig
line_magics['auto' + px] = self.autopx
self.magics['cell'][px] = self.cell_px
super(ParallelMagics, self).__init__(shell=shell)
def _eval_target_str(self, ts):
if ':' in ts:
targets = eval("self.view.client.ids[%s]" % ts)
elif 'all' in ts:
targets = 'all'
else:
targets = eval(ts)
return targets
@magic_arguments.magic_arguments()
@exec_args
def pxconfig(self, line):
"""configure default targets/blocking for %px magics"""
args = magic_arguments.parse_argstring(self.pxconfig, line)
if args.targets:
self.view.targets = self._eval_target_str(args.targets)
if args.block is not None:
self.view.block = args.block
if args.set_verbose is not None:
self.verbose = args.set_verbose
@magic_arguments.magic_arguments()
@output_args
@skip_doctest
def result(self, line=''):
"""Print the result of the last asynchronous %px command.
This lets you recall the results of %px computations after
asynchronous submission (block=False).
Examples
--------
::
In [23]: %px os.getpid()
Async parallel execution on engine(s): all
In [24]: %pxresult
Out[8:10]: 60920
Out[9:10]: 60921
Out[10:10]: 60922
Out[11:10]: 60923
"""
args = magic_arguments.parse_argstring(self.result, line)
if self.last_result is None:
raise UsageError(NO_LAST_RESULT)
self.last_result.get()
self.last_result.display_outputs(groupby=args.groupby)
@skip_doctest
def px(self, line=''):
"""Executes the given python command in parallel.
Examples
--------
::
In [24]: %px a = os.getpid()
Parallel execution on engine(s): all
In [25]: %px print a
[stdout:0] 1234
[stdout:1] 1235
[stdout:2] 1236
[stdout:3] 1237
"""
return self.parallel_execute(line)
def parallel_execute(self, cell, block=None, groupby='type', save_name=None):
"""implementation used by %px and %%parallel"""
# defaults:
block = self.view.block if block is None else block
base = "Parallel" if block else "Async parallel"
targets = self.view.targets
if isinstance(targets, list) and len(targets) > 10:
str_targets = str(targets[:4])[:-1] + ', ..., ' + str(targets[-4:])[1:]
else:
str_targets = str(targets)
if self.verbose:
print base + " execution on engine(s): %s" % str_targets
result = self.view.execute(cell, silent=False, block=False)
self.last_result = result
if save_name:
self.shell.user_ns[save_name] = result
if block:
result.get()
result.display_outputs(groupby)
else:
# return AsyncResult only on non-blocking submission
return result
@magic_arguments.magic_arguments()
@exec_args
@output_args
@skip_doctest
def cell_px(self, line='', cell=None):
"""Executes the cell in parallel.
Examples
--------
::
In [24]: %%px --noblock
....: a = os.getpid()
Async parallel execution on engine(s): all
In [25]: %%px
....: print a
[stdout:0] 1234
[stdout:1] 1235
[stdout:2] 1236
[stdout:3] 1237
"""
args = magic_arguments.parse_argstring(self.cell_px, line)
if args.targets:
save_targets = self.view.targets
self.view.targets = self._eval_target_str(args.targets)
if args.local:
self.shell.run_cell(cell)
try:
return self.parallel_execute(cell, block=args.block,
groupby=args.groupby,
save_name=args.save_name,
)
finally:
if args.targets:
self.view.targets = save_targets
@skip_doctest
def autopx(self, line=''):
"""Toggles auto parallel mode.
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, the ``targets`` attribute of the view before
entering ``%autopx`` mode.
Then you can do the following::
In [25]: %autopx
%autopx to enabled
In [26]: a = 10
Parallel execution on engine(s): [0,1,2,3]
In [27]: print a
Parallel execution on engine(s): [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.
"""
# override run_cell
self._original_run_cell = self.shell.run_cell
self.shell.run_cell = self.pxrun_cell
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._autopx = False
print "%autopx disabled"
def pxrun_cell(self, raw_cell, store_history=False, silent=False):
"""drop-in replacement for InteractiveShell.run_cell.
This executes code remotely, instead of in the local namespace.
See InteractiveShell.run_cell for details.
"""
if (not raw_cell) or raw_cell.isspace():
return
ipself = self.shell
with ipself.builtin_trap:
cell = ipself.prefilter_manager.prefilter_lines(raw_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:
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.view.execute(cell, silent=False, block=False)
except:
ipself.showtraceback()
return True
else:
if self.view.block:
try:
result.get()
except:
self.shell.showtraceback()
return True
else:
with ipself.builtin_trap:
result.display_outputs()
return False
__doc__ = __doc__.format(
AUTOPX_DOC = ' '*8 + ParallelMagics.autopx.__doc__,
PX_DOC = ' '*8 + ParallelMagics.px.__doc__,
RESULT_DOC = ' '*8 + ParallelMagics.result.__doc__,
CONFIG_DOC = ' '*8 + ParallelMagics.pxconfig.__doc__,
)