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
  66         entry = (task.timestamp, task.id, task)
  67         self._valid.add(task.id)
  68         heapq.heappush(self._queue, entry)
  69         return task
  70
  71     def remove(self, tid):
  72         """ Remove a task form the heap
  73
  74         :param tid: Id of the task that need to be removed
  75         :type tid: int
  76         """
  77         try:
  78             self._valid.remove(tid)
  79         except:
  80             pass
  81
  82     def next(self):
  83         """ Get the next task in the scheduler
  84
  85         """
  86         while self._queue:
  87             try:
  88                 timestamp, tid, task = heapq.heappop(self._queue)
  89                 if tid in self._valid:
  90                     self.remove(tid)
  91                     return task
  92             except IndexError:
  93                 # heap empty
  94                 pass
  95         return None
  96