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     def __init__(self):
  50         super(HeapScheduler, self).__init__()
  51         self._queue = list()
  52         self._valid = set()
  53         self._idgen = itertools.count(1)
  54
  55     def schedule(self, task):
  56         """ Add the task 'task' in the heap of the scheduler
  57
  58         :param task: task that need to be schedule
  59         :type task: task
  60         """
  61         if task.id == None:
  62             task.id = self._idgen.next()
  63         entry = (task.timestamp, task.id, task)
  64         self._valid.add(task.id)
  65         heapq.heappush(self._queue, entry)
  66         return task
  67
  68     def remove(self, tid):
  69         """ Remove a task form the heap
  70
  71         :param tid: Id of the task that need to be removed
  72         :type tid: int
  73         """
  74         try:
  75             self._valid.remove(tid)
  76         except:
  77             pass
  78
  79     def next(self):
  80         """ Get the next task in the scheduler
  81
  82         """
  83         while self._queue:
  84             try:
  85                 timestamp, tid, task = heapq.heappop(self._queue)
  86                 if tid in self._valid:
  87                     self.remove(tid)
  88                     return task
  89             except IndexError:
  90                 # heap empty
  91                 pass
  92         return None
  93