from glimpse import util
from IPython.parallel import Client
from IPython.parallel.error import CompositeError, NoEnginesRegistered

def KillWorkers(client):
  """Kill all active workers."""
  client.shutdown(restart = True, block = True, hub = False)

def RestartWorkers(client):
  """Kill and relaunch all active workers."""
  # The ipython engine creates a file in the worker run directory to communicate the restart signal to the wrapper script.
  # This is a hackish stop-gap measure until client.shutdown responds to the 'restart' argument.
  dview = client.direct_view()
  def restart():
    import os, socket
    fname = os.path.join('run', 'glimpse-cluster', socket.getfqdn(), 'workers', 'restart')
    with file(fname, 'a'):
      os.utime(fname, None)
  dview.apply_sync(restart)
  client.shutdown(restart = False, block = True, hub = False)

def PingWorkers(client):
  """Determine the set of active workers."""
  def ping():
    import os, socket
    return socket.getfqdn(), os.getpid()
  dview = client.direct_view()
  try:
    results = dview.apply_sync(ping)
    print "\n".join(" ".join(map(str, r)) for r in results)
  except NoEnginesRegistered:
    print "--None--"

def main():
  methods = map(eval, ("KillWorkers", "RestartWorkers", "PingWorkers"))
  try:
    opts, args = util.GetOptions("v")
    for opt, arg in opts:
      if opt == '-v':
        import logging
        logging.getLogger().setLevel(logging.INFO)
    if len(args) < 1:
      raise util.UsageException
    method = eval(args[0])
    client = Client()
#    dview = client.direct_view()
    method(client, *args[1:])
  except util.UsageException, e:
    method_info = [ "  %s -- %s" % (m.func_name, m.__doc__.splitlines()[0])
        for m in methods ]
    util.Usage("[options] CMD [ARGS]\n"
        "  -v        Be verbose with logging\n"
        "CMDs include:\n" + "\n".join(method_info),
        e)

if __name__ == "__main__":
  main()

@@ -1265,16 +1265,33 @@
 kers%22))%0A
+  profile = None%0A
   try:%0A 
@@ -1323,16 +1323,18 @@
 ptions(%22
+p:
 v%22)%0A    
@@ -1370,16 +1370,62 @@
 f opt ==
+ '-p':%0A        profile = arg%0A      elif opt ==
  '-v':%0A 
@@ -1598,16 +1598,33 @@
  Client(
+profile = profile
 )%0A#    d
@@ -1853,24 +1853,72 @@
 D %5BARGS%5D%5Cn%22%0A
+        %22  -p PROF   Set the IPython profile%5Cn%22%0A
         %22  -

Add optional profile argument to ipython_cluster main script.