src/nepi/execution/scheduler.py

   1 #
   2 #    NEPI, a framework to manage network experiments
   3 #    Copyright (C) 2013 INRIA
   4 #
   5 #    This program is free software: you can redistribute it and/or modify
   6 #    it under the terms of the GNU General Public License as published by
   7 #    the Free Software Foundation, either version 3 of the License, or
   8 #    (at your option) any later version.
   9 #
  10 #    This program is distributed in the hope that it will be useful,
  11 #    but WITHOUT ANY WARRANTY; without even the implied warranty of
  12 #    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  13 #    GNU General Public License for more details.
  14 #
  15 #    You should have received a copy of the GNU General Public License
  16 #    along with this program.  If not, see <http://www.gnu.org/licenses/>.
  17 #
  18 # Author: Alina Quereilhac <alina.quereilhac@inria.fr>
  19
  20 import itertools
  21 import heapq
  22
  23 class TaskStatus:
  24     NEW = 0
  25     DONE = 1
  26     ERROR = 2
  27
  28
  29 class Task(object):
  30     """ This class is to define a task, that is represented by an id,
  31     an execution time 'timestamp' and an action 'callback """
  32
  33     def __init__(self, timestamp, callback):
  34         self.id = None
  35         self.timestamp = timestamp
  36         self.callback = callback
  37         self.result = None
  38         self.status = TaskStatus.NEW
  39
  40 class HeapScheduler(object):
  41     """ Create a Heap Scheduler.
  42
  43     .. note::
  44
  45         This class is thread safe.
  46         All calls to C Extensions are made atomic by the GIL in the CPython implementation.
  47         heapq.heappush, heapq.heappop, and list access are therefore thread-safe.
  48
  49     """
  50
  51     def __init__(self):
  52         super(HeapScheduler, self).__init__()
  53         self._queue = list()
  54         self._valid = set()
  55         self._idgen = itertools.count(1)
  56
  57     def schedule(self, task):
  58         """ Add the task 'task' in the heap of the scheduler
  59
  60         :param task: task that need to be schedule
  61         :type task: task
  62         """
  63         if task.id == None:
  64             task.id = self._idgen.next()
  65         entry = (task.timestamp, task.id, task)
  66         self._valid.add(task.id)
  67         heapq.heappush(self._queue, entry)
  68         return task
  69
  70     def remove(self, tid):
  71         """ Remove a task form the heap
  72
  73         :param tid: Id of the task that need to be removed
  74         :type tid: int
  75         """
  76         try:
  77             self._valid.remove(tid)
  78         except:
  79             pass
  80
  81     def next(self):
  82         """ Get the next task in the scheduler
  83
  84         """
  85         while self._queue:
  86             try:
  87                 timestamp, tid, task = heapq.heappop(self._queue)
  88                 if tid in self._valid:
  89                     self.remove(tid)
  90                     return task
  91             except IndexError:
  92                 # heap empty
  93                 pass
  94         return None
  95