t/moe/pipeline.py

   1 #!/usr/bin/env python
   2
   3 import sys
   4 import bisect
   5 import imp
   6 import moe
   7 import moe.log
   8
   9 class MoePipeError(moe.MoeError):
  10     """Failure of the MoePipeline."""
  11
  12 class MoeAbortPipeline(Exception):
  13
  14     def __init__(self, skip_to=999):
  15         self.skip_to = skip_to
  16
  17 class MoePipeline:
  18     """Moe pipeline."""
  19
  20     def __init__(self, name):
  21         self.pipe = []
  22         self.index = -1
  23         self.name = name
  24
  25     def insert(self, pri, name, fun):
  26         triple = (pri,name,fun)
  27         pos = bisect.bisect(self.pipe, triple)
  28         if pos <= self.index:
  29             raise MoePipeError, "Pipeline insert cannot alter the past"
  30         self.pipe.insert(pos, triple)
  31
  32     def dump(self, file=sys.stdout, prefix=""):
  33         file.write(">>> Pipeline %s\n" % self.name)
  34         for pri,name,fun in self.pipe:
  35             file.write("%s%03d %s\n" % (prefix,pri,name))
  36
  37     def run(self, *args):
  38         self.index = 0
  39         min_pri = -1
  40         while self.index < len(self.pipe):
  41             (pri,name,fun) = self.pipe[self.index]
  42             if pri >= min_pri:
  43                 moe.log.default.verbose(">> Running %s:%s\n" % (self.name,name))
  44                 try:
  45                     fun(*args)
  46                 except MoeAbortPipeline, err:
  47                     min_pri = err.skip_to
  48             else:
  49                 moe.log.default.verbose(">> Skipping %s:%s\n" % (self.name,name))
  50             self.index += 1
  51         self.index = -1
  52         moe.log.default.verbose(">> Pipeline %s finished\n" % self.name)
  53
  54     def add_hook(self, name):
  55         modname = "moe.hooks." + name
  56         moe.log.default.verbose(">> Loading hook %s\n" % name)
  57         if not sys.modules.has_key(modname):
  58             ## FIXME: Configuration variable for the hook directory?
  59             try:
  60                 fp, path, desc = imp.find_module(name, ["moe/hooks"])
  61             except ImportError:
  62                 raise MoePipeError, "Cannot find hook module " + modname
  63             try:
  64                 imp.load_module(modname, fp, path, desc)
  65             finally:
  66                 fp.close()
  67         sys.modules[modname].init(self)
  68
  69     def configure(self, names):
  70         for name in names.split():
  71             self.add_hook(name)