1  # -*- coding: utf-8 -*-
   2  """Manage background (threaded) jobs conveniently from an interactive shell.
   3  
   4  This module provides a BackgroundJobManager class.  This is the main class
   5  meant for public usage, it implements an object which can create and manage
   6  new background jobs.
   7  
   8  It also provides the actual job classes managed by these BackgroundJobManager
   9  objects, see their docstrings below.
  10  
  11  
  12  This system was inspired by discussions with B. Granger and the
  13  BackgroundCommand class described in the book Python Scripting for
  14  Computational Science, by H. P. Langtangen:
  15  
  16  http://folk.uio.no/hpl/scripting
  17  
  18  (although ultimately no code from this text was used, as IPython's system is a
  19  separate implementation).
  20  
  21  An example notebook is provided in our documentation illustrating interactive
  22  use of the system.
  23  """
  24  
  25  #*****************************************************************************
  26  #       Copyright (C) 2005-2006 Fernando Perez <fperez@colorado.edu>
  27  #
  28  #  Distributed under the terms of the BSD License.  The full license is in
  29  #  the file COPYING, distributed as part of this software.
  30  #*****************************************************************************
  31  
  32  # Code begins
  33  import sys
  34  import threading
  35  
  36  from IPython import get_ipython
  37  from IPython.core.ultratb import AutoFormattedTB
  38  from logging import error, debug
  39  
  40  
  41  class BackgroundJobManager(object):
  42      """Class to manage a pool of backgrounded threaded jobs.
  43  
  44      Below, we assume that 'jobs' is a BackgroundJobManager instance.
  45      
  46      Usage summary (see the method docstrings for details):
  47  
  48        jobs.new(...) -> start a new job
  49        
  50        jobs() or jobs.status() -> print status summary of all jobs
  51  
  52        jobs[N] -> returns job number N.
  53  
  54        foo = jobs[N].result -> assign to variable foo the result of job N
  55  
  56        jobs[N].traceback() -> print the traceback of dead job N
  57  
  58        jobs.remove(N) -> remove (finished) job N
  59  
  60        jobs.flush() -> remove all finished jobs
  61        
  62      As a convenience feature, BackgroundJobManager instances provide the
  63      utility result and traceback methods which retrieve the corresponding
  64      information from the jobs list:
  65  
  66        jobs.result(N) <--> jobs[N].result
  67        jobs.traceback(N) <--> jobs[N].traceback()
  68  
  69      While this appears minor, it allows you to use tab completion
  70      interactively on the job manager instance.
  71      """
  72  
  73      def __init__(self):
  74          # Lists for job management, accessed via a property to ensure they're
  75          # up to date.x
  76          self._running  = []
  77          self._completed = []
  78          self._dead = []
  79          # A dict of all jobs, so users can easily access any of them
  80          self.all = {}
  81          # For reporting
  82          self._comp_report = []
  83          self._dead_report = []
  84          # Store status codes locally for fast lookups
  85          self._s_created   = BackgroundJobBase.stat_created_c
  86          self._s_running   = BackgroundJobBase.stat_running_c
  87          self._s_completed = BackgroundJobBase.stat_completed_c
  88          self._s_dead      = BackgroundJobBase.stat_dead_c
  89          self._current_job_id = 0
  90  
  91      @property
  92      def running(self):
  93          self._update_status()
  94          return self._running
  95  
  96      @property
  97      def dead(self):
  98          self._update_status()
  99          return self._dead
 100  
 101      @property
 102      def completed(self):
 103          self._update_status()
 104          return self._completed
 105  
 106      def new(self, func_or_exp, *args, **kwargs):
 107          """Add a new background job and start it in a separate thread.
 108  
 109          There are two types of jobs which can be created:
 110  
 111          1. Jobs based on expressions which can be passed to an eval() call.
 112          The expression must be given as a string.  For example:
 113  
 114            job_manager.new('myfunc(x,y,z=1)'[,glob[,loc]])
 115  
 116          The given expression is passed to eval(), along with the optional
 117          global/local dicts provided.  If no dicts are given, they are
 118          extracted automatically from the caller's frame.
 119  
 120          A Python statement is NOT a valid eval() expression.  Basically, you
 121          can only use as an eval() argument something which can go on the right
 122          of an '=' sign and be assigned to a variable.
 123  
 124          For example,"print 'hello'" is not valid, but '2+3' is.
 125  
 126          2. Jobs given a function object, optionally passing additional
 127          positional arguments:
 128  
 129            job_manager.new(myfunc, x, y)
 130  
 131          The function is called with the given arguments.
 132  
 133          If you need to pass keyword arguments to your function, you must
 134          supply them as a dict named kw:
 135  
 136            job_manager.new(myfunc, x, y, kw=dict(z=1))
 137  
 138          The reason for this asymmetry is that the new() method needs to
 139          maintain access to its own keywords, and this prevents name collisions
 140          between arguments to new() and arguments to your own functions.
 141  
 142          In both cases, the result is stored in the job.result field of the
 143          background job object.
 144  
 145          You can set `daemon` attribute of the thread by giving the keyword
 146          argument `daemon`.
 147  
 148          Notes and caveats:
 149  
 150          1. All threads running share the same standard output.  Thus, if your
 151          background jobs generate output, it will come out on top of whatever
 152          you are currently writing.  For this reason, background jobs are best
 153          used with silent functions which simply return their output.
 154  
 155          2. Threads also all work within the same global namespace, and this
 156          system does not lock interactive variables.  So if you send job to the
 157          background which operates on a mutable object for a long time, and
 158          start modifying that same mutable object interactively (or in another
 159          backgrounded job), all sorts of bizarre behaviour will occur.
 160  
 161          3. If a background job is spending a lot of time inside a C extension
 162          module which does not release the Python Global Interpreter Lock
 163          (GIL), this will block the IPython prompt.  This is simply because the
 164          Python interpreter can only switch between threads at Python
 165          bytecodes.  While the execution is inside C code, the interpreter must
 166          simply wait unless the extension module releases the GIL.
 167  
 168          4. There is no way, due to limitations in the Python threads library,
 169          to kill a thread once it has started."""
 170          
 171          if callable(func_or_exp):
 172              kw  = kwargs.get('kw',{})
 173              job = BackgroundJobFunc(func_or_exp,*args,**kw)
 174          elif isinstance(func_or_exp, str):
 175              if not args:
 176                  frame = sys._getframe(1)
 177                  glob, loc = frame.f_globals, frame.f_locals
 178              elif len(args)==1:
 179                  glob = loc = args[0]
 180              elif len(args)==2:
 181                  glob,loc = args
 182              else:
 183                  raise ValueError(
 184                        'Expression jobs take at most 2 args (globals,locals)')
 185              job = BackgroundJobExpr(func_or_exp, glob, loc)
 186          else:
 187              raise TypeError('invalid args for new job')
 188  
 189          if kwargs.get('daemon', False):
 190              job.daemon = True
 191          job.num = self._current_job_id
 192          self._current_job_id += 1
 193          self.running.append(job)
 194          self.all[job.num] = job
 195          debug('Starting job # %s in a separate thread.' % job.num)
 196          job.start()
 197          return job
 198  
 199      def __getitem__(self, job_key):
 200          num = job_key if isinstance(job_key, int) else job_key.num
 201          return self.all[num]
 202  
 203      def __call__(self):
 204          """An alias to self.status(),
 205  
 206          This allows you to simply call a job manager instance much like the
 207          Unix `jobs` shell command."""
 208  
 209          return self.status()
 210  
 211      def _update_status(self):
 212          """Update the status of the job lists.
 213  
 214          This method moves finished jobs to one of two lists:
 215            - self.completed: jobs which completed successfully
 216            - self.dead: jobs which finished but died.
 217  
 218          It also copies those jobs to corresponding _report lists.  These lists
 219          are used to report jobs completed/dead since the last update, and are
 220          then cleared by the reporting function after each call."""
 221  
 222          # Status codes
 223          srun, scomp, sdead = self._s_running, self._s_completed, self._s_dead
 224          # State lists, use the actual lists b/c the public names are properties
 225          # that call this very function on access
 226          running, completed, dead = self._running, self._completed, self._dead
 227  
 228          # Now, update all state lists
 229          for num, job in enumerate(running):
 230              stat = job.stat_code
 231              if stat == srun:
 232                  continue
 233              elif stat == scomp:
 234                  completed.append(job)
 235                  self._comp_report.append(job)
 236                  running[num] = False
 237              elif stat == sdead:
 238                  dead.append(job)
 239                  self._dead_report.append(job)
 240                  running[num] = False
 241          # Remove dead/completed jobs from running list
 242          running[:] = filter(None, running)
 243  
 244      def _group_report(self,group,name):
 245          """Report summary for a given job group.
 246  
 247          Return True if the group had any elements."""
 248  
 249          if group:
 250              print('%s jobs:' % name)
 251              for job in group:
 252                  print('%s : %s' % (job.num,job))
 253              print()
 254              return True
 255  
 256      def _group_flush(self,group,name):
 257          """Flush a given job group
 258  
 259          Return True if the group had any elements."""
 260  
 261          njobs = len(group)
 262          if njobs:
 263              plural = {1:''}.setdefault(njobs,'s')
 264              print('Flushing %s %s job%s.' % (njobs,name,plural))
 265              group[:] = []
 266              return True
 267          
 268      def _status_new(self):
 269          """Print the status of newly finished jobs.
 270  
 271          Return True if any new jobs are reported.
 272  
 273          This call resets its own state every time, so it only reports jobs
 274          which have finished since the last time it was called."""
 275  
 276          self._update_status()
 277          new_comp = self._group_report(self._comp_report, 'Completed')
 278          new_dead = self._group_report(self._dead_report,
 279                                        'Dead, call jobs.traceback() for details')
 280          self._comp_report[:] = []
 281          self._dead_report[:] = []
 282          return new_comp or new_dead
 283                  
 284      def status(self,verbose=0):
 285          """Print a status of all jobs currently being managed."""
 286  
 287          self._update_status()
 288          self._group_report(self.running,'Running')
 289          self._group_report(self.completed,'Completed')
 290          self._group_report(self.dead,'Dead')
 291          # Also flush the report queues
 292          self._comp_report[:] = []
 293          self._dead_report[:] = []
 294  
 295      def remove(self,num):
 296          """Remove a finished (completed or dead) job."""
 297  
 298          try:
 299              job = self.all[num]
 300          except KeyError:
 301              error('Job #%s not found' % num)
 302          else:
 303              stat_code = job.stat_code
 304              if stat_code == self._s_running:
 305                  error('Job #%s is still running, it can not be removed.' % num)
 306                  return
 307              elif stat_code == self._s_completed:
 308                  self.completed.remove(job)
 309              elif stat_code == self._s_dead:
 310                  self.dead.remove(job)
 311  
 312      def flush(self):
 313          """Flush all finished jobs (completed and dead) from lists.
 314  
 315          Running jobs are never flushed.
 316  
 317          It first calls _status_new(), to update info. If any jobs have
 318          completed since the last _status_new() call, the flush operation
 319          aborts."""
 320  
 321          # Remove the finished jobs from the master dict
 322          alljobs = self.all
 323          for job in self.completed+self.dead:
 324              del(alljobs[job.num])
 325  
 326          # Now flush these lists completely
 327          fl_comp = self._group_flush(self.completed, 'Completed')
 328          fl_dead = self._group_flush(self.dead, 'Dead')
 329          if not (fl_comp or fl_dead):
 330              print('No jobs to flush.')
 331  
 332      def result(self,num):
 333          """result(N) -> return the result of job N."""
 334          try:
 335              return self.all[num].result
 336          except KeyError:
 337              error('Job #%s not found' % num)
 338  
 339      def _traceback(self, job):
 340          num = job if isinstance(job, int) else job.num
 341          try:
 342              self.all[num].traceback()
 343          except KeyError:
 344              error('Job #%s not found' % num)
 345  
 346      def traceback(self, job=None):
 347          if job is None:
 348              self._update_status()
 349              for deadjob in self.dead:
 350                  print("Traceback for: %r" % deadjob)
 351                  self._traceback(deadjob)
 352                  print()
 353          else:
 354              self._traceback(job)
 355  
 356  
 357  class BackgroundJobBase(threading.Thread):
 358      """Base class to build BackgroundJob classes.
 359  
 360      The derived classes must implement:
 361  
 362      - Their own __init__, since the one here raises NotImplementedError.  The
 363        derived constructor must call self._init() at the end, to provide common
 364        initialization.
 365  
 366      - A strform attribute used in calls to __str__.
 367  
 368      - A call() method, which will make the actual execution call and must
 369        return a value to be held in the 'result' field of the job object.
 370      """
 371  
 372      # Class constants for status, in string and as numerical codes (when
 373      # updating jobs lists, we don't want to do string comparisons).  This will
 374      # be done at every user prompt, so it has to be as fast as possible
 375      stat_created   = 'Created'; stat_created_c = 0
 376      stat_running   = 'Running'; stat_running_c = 1
 377      stat_completed = 'Completed'; stat_completed_c = 2
 378      stat_dead      = 'Dead (Exception), call jobs.traceback() for details'
 379      stat_dead_c = -1
 380  
 381      def __init__(self):
 382          """Must be implemented in subclasses.
 383  
 384          Subclasses must call :meth:`_init` for standard initialisation.
 385          """
 386          raise NotImplementedError("This class can not be instantiated directly.")
 387  
 388      def _init(self):
 389          """Common initialization for all BackgroundJob objects"""
 390          
 391          for attr in ['call','strform']:
 392              assert hasattr(self,attr), "Missing attribute <%s>" % attr
 393          
 394          # The num tag can be set by an external job manager
 395          self.num = None
 396        
 397          self.status    = BackgroundJobBase.stat_created
 398          self.stat_code = BackgroundJobBase.stat_created_c
 399          self.finished  = False
 400          self.result    = '<BackgroundJob has not completed>'
 401          
 402          # reuse the ipython traceback handler if we can get to it, otherwise
 403          # make a new one
 404          try:
 405              make_tb = get_ipython().InteractiveTB.text
 406          except:
 407              make_tb = AutoFormattedTB(mode = 'Context',
 408                                        color_scheme='NoColor',
 409                                        tb_offset = 1).text
 410          # Note that the actual API for text() requires the three args to be
 411          # passed in, so we wrap it in a simple lambda.
 412          self._make_tb = lambda : make_tb(None, None, None)
 413  
 414          # Hold a formatted traceback if one is generated.
 415          self._tb = None
 416          
 417          threading.Thread.__init__(self)
 418  
 419      def __str__(self):
 420          return self.strform
 421  
 422      def __repr__(self):
 423          return '<BackgroundJob #%d: %s>' % (self.num, self.strform)
 424  
 425      def traceback(self):
 426          print(self._tb)
 427          
 428      def run(self):
 429          try:
 430              self.status    = BackgroundJobBase.stat_running
 431              self.stat_code = BackgroundJobBase.stat_running_c
 432              self.result    = self.call()
 433          except:
 434              self.status    = BackgroundJobBase.stat_dead
 435              self.stat_code = BackgroundJobBase.stat_dead_c
 436              self.finished  = None
 437              self.result    = ('<BackgroundJob died, call jobs.traceback() for details>')
 438              self._tb       = self._make_tb()
 439          else:
 440              self.status    = BackgroundJobBase.stat_completed
 441              self.stat_code = BackgroundJobBase.stat_completed_c
 442              self.finished  = True
 443  
 444  
 445  class BackgroundJobExpr(BackgroundJobBase):
 446      """Evaluate an expression as a background job (uses a separate thread)."""
 447  
 448      def __init__(self, expression, glob=None, loc=None):
 449          """Create a new job from a string which can be fed to eval().
 450  
 451          global/locals dicts can be provided, which will be passed to the eval
 452          call."""
 453  
 454          # fail immediately if the given expression can't be compiled
 455          self.code = compile(expression,'<BackgroundJob compilation>','eval')
 456                  
 457          glob = {} if glob is None else glob
 458          loc = {} if loc is None else loc
 459          self.expression = self.strform = expression
 460          self.glob = glob
 461          self.loc = loc
 462          self._init()
 463          
 464      def call(self):
 465          return eval(self.code,self.glob,self.loc)
 466  
 467  
 468  class BackgroundJobFunc(BackgroundJobBase):
 469      """Run a function call as a background job (uses a separate thread)."""
 470  
 471      def __init__(self, func, *args, **kwargs):
 472          """Create a new job from a callable object.
 473  
 474          Any positional arguments and keyword args given to this constructor
 475          after the initial callable are passed directly to it."""
 476  
 477          if not callable(func):
 478              raise TypeError(
 479                  'first argument to BackgroundJobFunc must be callable')
 480          
 481          self.func = func
 482          self.args = args
 483          self.kwargs = kwargs
 484          # The string form will only include the function passed, because
 485          # generating string representations of the arguments is a potentially
 486          # _very_ expensive operation (e.g. with large arrays).
 487          self.strform = str(func)
 488          self._init()
 489  
 490      def call(self):
 491          return self.func(*self.args, **self.kwargs)
 492