# encoding: utf-8 """ ============= parallelmagic ============= Magic command interface for interactive parallel work. Usage ===== ``%autopx`` {AUTOPX_DOC} ``%px`` {PX_DOC} ``%pxresult`` {RESULT_DOC} ``%pxconfig`` {CONFIG_DOC} """ from __future__ import print_function #----------------------------------------------------------------------------- # 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 running local, don't block until after local has run block = False if args.local else args.block try: ar = self.parallel_execute(cell, block=block, groupby=args.groupby, save_name=args.save_name, ) finally: if args.targets: self.view.targets = save_targets # run locally after submitting remote block = self.view.block if args.block is None else args.block if args.local: self.shell.run_cell(cell) # now apply blocking behavor to remote execution if block: ar.get() ar.display_outputs(args.groupby) if not block: return ar @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__, )