pysawsim/manager/subproc.py

   1 # Copyright (C) 2010  W. Trevor King <wking@drexel.edu>
   2 #
   3 # This program is free software: you can redistribute it and/or modify
   4 # it under the terms of the GNU General Public License as published by
   5 # the Free Software Foundation, either version 3 of the License, or
   6 # (at your option) any later version.
   7 #
   8 # This program is distributed in the hope that it will be useful,
   9 # but WITHOUT ANY WARRANTY; without even the implied warranty of
  10 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  11 # GNU General Public License for more details.
  12 #
  13 # You should have received a copy of the GNU General Public License
  14 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
  15 #
  16 # The author may be contacted at <wking@drexel.edu> on the Internet, or
  17 # write to Trevor King, Drexel University, Physics Dept., 3141 Chestnut St.,
  18 # Philadelphia PA 19104, USA.
  19
  20 """Functions for running external commands on other hosts.
  21 """
  22
  23 from multiprocessing import Manager, Process, Queue, cpu_count
  24
  25 from .. import log
  26 from . import Job
  27 from .thread import ThreadManager, CLOSE_MESSAGE
  28
  29
  30 class WorkerProcess (Process):
  31     def __init__(self, spawn_queue, receive_queue, *args, **kwargs):
  32         super(WorkerProcess, self).__init__(*args, **kwargs)
  33         self.spawn_queue = spawn_queue
  34         self.receive_queue = receive_queue
  35
  36     def run(self):
  37         while True:
  38             msg = self.spawn_queue.get()
  39             if msg == CLOSE_MESSAGE:
  40                 log().debug('%s closing' % self.name)
  41                 break
  42             assert isinstance(msg, Job), msg
  43             log().debug('%s running job %s' % (self.name, msg))
  44             msg.run()
  45             self.receive_queue.put(msg)
  46
  47
  48 class SubprocessManager (ThreadManager):
  49     """Manage asynchronous `Job` execution via :mod:`subprocess`.
  50
  51     >>> from time import sleep
  52     >>> from math import sqrt
  53     >>> m = SubprocessManager()
  54     >>> group_A = []
  55     >>> for i in range(10):
  56     ...     t = max(0, 5-i)
  57     ...     group_A.append(m.async_invoke(Job(target=sleep, args=[t])))
  58     >>> group_B = []
  59     >>> for i in range(10):
  60     ...     group_B.append(m.async_invoke(Job(target=sqrt, args=[i],
  61     ...                 blocks_on=[j.id for j in group_A])))
  62     >>> jobs = m.wait(ids=[j.id for j in group_A[5:8]])
  63     >>> print sorted(jobs.values(), key=lambda j: j.id)
  64     [<Job 5>, <Job 6>, <Job 7>]
  65     >>> jobs = m.wait()
  66     >>> print sorted(jobs.values(), key=lambda j: j.id)
  67     ... # doctest: +NORMALIZE_WHITESPACE
  68     [<Job 0>, <Job 1>, <Job 2>, <Job 3>, <Job 4>, <Job 8>, <Job 9>, <Job 10>,
  69      <Job 11>, <Job 12>, <Job 13>, <Job 14>, <Job 15>, <Job 16>, <Job 17>,
  70      <Job 18>, <Job 19>]
  71     >>> m.teardown()
  72     """
  73     def __init__(self, worker_pool=None):
  74         super(SubprocessManager, self).__init__(worker_pool=worker_pool)
  75
  76     def _setup_queues(self):
  77         self._spawn_queue = Queue()
  78         self._receive_queue = Queue()
  79
  80     def _spawn_workers(self, worker_pool=None):
  81         if worker_pool == None:
  82             worker_pool = cpu_count() + 1
  83         self._manager = Manager()
  84         self._workers = []
  85         for i in range(worker_pool):
  86             worker = WorkerProcess(spawn_queue=self._spawn_queue,
  87                                    receive_queue=self._receive_queue,
  88                                    name='worker-%d' % i)
  89             log().debug('start %s' % worker.name)
  90             worker.start()
  91             self._workers.append(worker)
  92
  93     def _put_job_in_spawn_queue(self, job):
  94         """Place a job in the spawn queue."""
  95         self._spawn_queue.put(job)