smbinterp/bin/sheep.py

#!/usr/bin/env python

import sys
import os
import time

from   multiprocessing.managers import BaseManager

from optparse import OptionParser

import numpy as np

import interp.bootstrap
from   interp.grid.gmsh import ggrid
from   interp.tools import *

from interp.cluster import QueueManager, get_qs

if __name__ == '__main__':
  parser = OptionParser(usage = "usage: %s [options] <server> <gmsh file> <dimension> <baker|friendly>")

  parser.add_option("-v", "--verbose",
    action="store_true", dest="verbose", default=False,
    help="verbose flag (default: %default)")

  parser.add_option('-p', '--port',
    type="int", dest="port", default=6666,
    help="specify the port to use on the server (default: %default)")

  (options, args) = parser.parse_args()

  if len(args) != 4:
    parser.print_usage()
    sys.exit(1)

  server, input_file, dimension, exact_type = args
  dimension = int(dimension)

  if   dimension == 2 and exact_type == 'baker':
    exact = baker_exact_2D
  elif dimension == 2 and exact_type == 'friendly':
    exact = friendly_exact_2D
  elif dimension == 3 and exact_type == 'baker':
    exact = baker_exact_3D
  elif dimension == 3 and exact_type == 'friendly':
    exact = friendly_exact_3D


  myname = "%s-%d" % (os.uname()[1], os.getpid())
  if options.verbose:
    print "%s: started" % myname


  m = QueueManager(address=(server, options.port), authkey='asdf')
  m.connect()

  tasksq, resultsq, masterq, minionsq = get_qs(m)

  if options.verbose:
    print "%s: starting parse input file" % myname
  g   = ggrid(input_file, dimension)
  g.q = np.array([exact(x) for x in g.verts])
  if options.verbose:
    print "%s: done parsing input file" % myname


  while True:
    if options.verbose:
      print "%s: letting master know that I am ready" % myname
    masterq.put(myname)

    if options.verbose:
      print "%s: waiting for master to tell me to start" % myname
    action = minionsq.get()
    if options.verbose:
      print "%s: got message from master: %s" % (myname, action)

    if action in ('teardown', 'slay'):
      masterq.get()
      break

    while not tasksq.empty():
      i, order, extra, X = tasksq.get()
      try:
        s = time.time()
        a = g.run_baker(X, order = order, extra_points = extra)
        cur_qlin  = a['qlin' ]
        cur_error = a['error']
        cur_final = a['final']
        cur_exact = exact(X)
        e = time.time()
        duration = e-s
      except Exception as e:
        print >>sys.stderr, X, e
        cur_qlin  = 0.0
        cur_error = 0.0
        cur_final = 0.0
        cur_exact = 0.0
        duration  = 0

      resultsq.put((i,
          cur_qlin,
          cur_error,
          cur_final,
          cur_exact,
          duration,
        ))

  if options.verbose:
    print "%s: exiting" % myname