Merge lp:~doanac/lava-celery/scheduler-monitor-support into lp:lava-celery

[Andy Doan]

This updates the celery module a few ways:

1) improve run_command_task to take files as arguments
2) use fork/exec for run_command task
3) add a new celery-schedulermonitor command

[Michael Hudson-Doyle (mwhudson)]

Yay for all that Unix fd yoga.

I think you should be a bunch more anal about errors in the child process: don't call sys.exit() but rather os._exit() and probably do it in a finally: block whose try: you enter first thing in the pid == 0 branch (going through the finally: clause in the run_command twice would be pretty confusing I think!).

You add celery-schedulermonitor to commands.__all__, not celery_schedulermonitor!

There's something about the way files are handled that makes me uneasy -- the code on the worker side doesn't do anything about directories or permissions, so I could imagine that if we started to write the json into /tmp/blah/foo.json rather than directly in /tmp/ things might break in unexpected ways. But well, maybe that just needs an XXX.

Otherwise, all looks ok!

review: Approve

[lp:~doanac/lava-celery/scheduler-monitor-support]

25. By Andy Doan on 2012-07-11

retval is now handled in the _exec_command function

26. By Andy Doan on 2012-07-11

make _exec_command a little safer

This adds some exception handling per mwhudson's review comments to help
ensure we always properly exit the child process.

27. By Andy Doan on 2012-07-11

fix celery-schedulermonitor command entry

mwhudson caught this in the code review

[Andy Doan (doanac)]

I just added commits 25,26, and 27 per your review comments. Let me know if they seem sane, and I'll proceed with ther merge

[Michael Hudson-Doyle (mwhudson)]

Looks good! Sorry for not replying sooner, I think Canonical mail was broken for a while yesterday.

=== modified file 'lava/celery/commands.py'
--- lava/celery/commands.py	2012-06-26 05:05:01 +0000
+++ lava/celery/commands.py	2012-07-11 03:19:19 +0000
@@ -15,7 +15,8 @@
 from lava.tool.command import Command, CommandGroup
 
 
-__all__ = ["celeryd", "celery", "hello_world", "run_remote"]
+__all__ = ["celeryd", "celery", "celery_schedulermonitor", "hello_world",
+            "run_remote"]
 
 class celeryd(Command):
     """
@@ -100,6 +101,7 @@
 
     @classmethod
     def register_arguments(cls, parser):
+        super(run_remote, cls).register_arguments(parser)
         parser.add_argument(
             '-s', '--silent',
             action="store_true",
@@ -195,7 +197,7 @@
             self.say("Unknown message: {0}", message)
             message.reject()
 
-    def invoke(self):
+    def invoke(self, arg0="lava", files=None):
         self._configure_celery()
         # Importing only after we call ._configure_celery()
         from celery.messaging import establish_connection
@@ -216,7 +218,7 @@
                     # Invoke a command with specified arguments.  NOTE: that we
                     # pass queue.name here which is conveniently also the
                     # routing key. See ._setup_plumbing() for an explanation.
-                    args=[queue.name, self.args.args],
+                    args=[queue.name, self.args.args, arg0, files],
                     # Routing key for the celery task workers. This routing
                     # key is different from queue.name. It is used by celery
                     # and celeryd to associate workers with tasks to perform.
@@ -247,3 +249,68 @@
             connection.close()
         self.logger.debug("Remote processing done, return code is %r", self.retval)
         return self.retval
+
+class celery_schedulermonitor(run_remote):
+    """
+    Launches a schedulermonitor task on a celery worker node.
+    """
+
+    @classmethod
+    def register_arguments(cls, parser):
+        super(celery_schedulermonitor, cls).register_arguments(parser)
+
+        group = parser.add_argument_group(title="Dispatcher Configuration")
+        group.add_argument('--settings-module',
+            default="lava_server.settings.debian",
+            help='The Django server settings module to use.')
+        group.add_argument('dispatcher', metavar='dispatcher',
+            help='The dispatcher to use like "lava-dispatch"')
+        group.add_argument('boardname', metavar='boardname',
+            help='The board to run the job on')
+        group.add_argument('json_file', metavar='json_file',
+            help='The job definition to run')
+
+    def _get_logfile_impl(self, logfile):
+        self.logfile = logfile.path
+
+    def _get_logfile(self):
+        os.environ["DJANGO_SETTINGS_MODULE"] = self.args.settings_module
+        from lava_scheduler_daemon.dbjobsource import DatabaseJobSource
+        from twisted.internet import reactor
+
+        source = DatabaseJobSource()
+
+        def run():
+            d = source.getLogFileForJobOnBoard(self.args.boardname)
+            d = d.addCallback(self._get_logfile_impl)
+            d.addCallback(lambda result: reactor.stop())
+
+        reactor.callWhenRunning(run)
+        reactor.run()
+        return self.logfile
+
+    def invoke(self):
+        logfile = self._get_logfile()
+
+        files = {}
+        with open(self.args.json_file, 'r') as f:
+            files[self.args.json_file] = f.read()
+
+        self._logfd = open(logfile, 'wb')
+
+        self.args.args = ['manage', 'schedulermonitor',
+            self.args.dispatcher, self.args.boardname, self.args.json_file]
+        super(celery_schedulermonitor, self).invoke('lava-server', files)
+
+    def _process_message(self, body, message):
+        self.logger.debug("Remote worker message: %r", message.payload)
+        if 'stdout' in message.payload:
+            self._logfd.write(message.payload['stdout'])
+            self._logfd.flush()
+        elif 'stderr' in message.payload:
+            self._logfd.write(message.payload['stderr'])
+            self._logfd.flush()
+        else:
+            if 'done' in message.payload:
+                self._logfd.close()
+            super(celery_schedulermonitor, self)._process_message(body, message)

=== modified file 'lava/celery/tasks.py'
--- lava/celery/tasks.py	2012-03-26 14:56:09 +0000
+++ lava/celery/tasks.py	2012-07-11 03:19:19 +0000
@@ -5,58 +5,91 @@
 Task definitions imported by celeryd when configured to use lava.celery.config
 """
 
+import os
+import select
 import sys
+import traceback
 
 from celery.messaging import establish_connection
 from celery.task import task
 from kombu.compat import Publisher
-from lava.tool.main import LavaDispatcher
 
 from lava.celery.queues import StreamExchange
 
 
 __all__ = ["run_command"]
 
-
-class CloudifiedDispatcher(LavaDispatcher):
-    """
-    Dispatcher suitable to run commands and stream their output back to the
-    originator via kombu.
-    """
-
-    def __init__(self, publisher):
-        super(CloudifiedDispatcher, self).__init__()
-        self.publisher = publisher
-
-    def say(self, command, message, *args, **kwargs):
-        text = message.format(*args, **kwargs)
-        self.publisher.send({
-            'command': command.get_name(),
-            'say': text
-        })
-
-
-class WriteProxy(object):
-    """
-    Simple proxy for File-like objects that redirects write() to a callback
-    """
-
-    def __init__(self, stream, write_cb):
-        self._stream = stream
-        self._write_cb = write_cb
-
-    def write(self, data):
-        self._write_cb(data)
-
-    def __getattr__(self, attr):
-        return getattr(self._stream, attr)
-
+def _exec_command(publisher, args, arg0):
+    """
+    A traditional fork/exec style action that prints stdout/stderr to the
+    publisher
+    """
+    (out_read, out_write) = os.pipe()
+    (err_read, err_write) = os.pipe()
+
+    pid = os.fork()
+    if pid == 0:
+        try:
+            os.close(out_read)
+            os.close(err_read)
+
+            os.dup2(out_write, 1)
+            os.dup2(err_write, 2)
+            args.insert(0, arg0)
+            os.execvp(args[0], args)
+            print >>sys.stderr, "lava celeryd error executing ",args
+        except:
+            traceback.print_exc()
+        finally:
+            os._exit(1)
+
+    os.close(out_write)
+    os.close(err_write)
+    out_read = os.fdopen(out_read)
+    err_read = os.fdopen(err_read)
+
+    readset = [out_read, err_read]
+    while readset:
+        rlist, wlist, xlist = select.select(readset, [], [])
+        if out_read in rlist:
+            data = os.read(out_read.fileno(), 1024)
+            if data == "":
+                out_read.close()
+                readset.remove(out_read)
+            publisher.send({'stdout': data})
+        if err_read in rlist:
+            data = os.read(err_read.fileno(), 1024)
+            if data == "":
+                err_read.close()
+                readset.remove(err_read)
+            publisher.send({'stderr': data})
+
+    (pid, retval) = os.waitpid(pid,0)
+    publisher.send({'done': retval})
+
+def _copy_files(files):
+    """
+    run command can take a dictionary of filname->content pairs that
+    the remote side of the command needs
+    """
+    if files is None:
+        return []
+    deletes = []
+    for k,v in files.items():
+        # if path exists, let's assume the celery worker and lava scheduler
+        # running on the same local system, so we don't need to create the file
+        if not os.path.exists(k):
+            with open(k, 'w') as f:
+                f.write(v)
+                deletes.insert(0, k)
+    return deletes
 
 @task(acks_late=True, ignore_result=True)
-def run_command(queue_name, args):
+def run_command(queue_name, args, arg0="lava", files=None):
     """
     Run a lava-tool command with the specified arguments
     """
+    print "command received:", args
     connection = establish_connection()
     channel = connection.channel()
     publisher = Publisher(
@@ -64,18 +97,17 @@
         exchange=StreamExchange,
         routing_key=queue_name,
         exchange_type="direct")
-    orig_stdout = sys.stdout
-    sys.stdout = WriteProxy(sys.stdout, lambda data: publisher.send({"stdout": data}))
-    orig_stderr = sys.stderr
-    sys.stderr = WriteProxy(sys.stderr, lambda data: publisher.send({"stderr": data}))
-    retval = None
+
     try:
-        retval = CloudifiedDispatcher(publisher).dispatch(args)
-    except SystemExit as exc:
-        retval = exc.code
+        files = _copy_files(files)
+        _exec_command(publisher, args, arg0)
+    except:
+        traceback.print_exc()
     finally:
-        publisher.send({'done': retval})
         publisher.close()
         connection.close()
-        sys.stdout = orig_stdout
-        sys.stderr = orig_stderr
+
+        for f in files:
+            os.unlink(f)
+
+    print "command complete for:",args

=== modified file 'setup.py'
--- setup.py	2012-06-26 05:05:01 +0000
+++ setup.py	2012-07-11 03:19:19 +0000
@@ -17,6 +17,7 @@
     [lava.commands]
     celery = lava.celery.commands:celery
     celeryd = lava.celery.commands:celeryd
+    celery-schedulermonitor- = lava.celery.commands:celery_schedulermonitor
     [lava.celery.commands]
     run-remote= lava.celery.commands:run_remote
     hello-world = lava.celery.commands:hello_world