cmd 4.92 KB
Newer Older
Swann Perarnau's avatar
Swann Perarnau committed
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52
#!/usr/bin/env python2

from __future__ import print_function
import argparse
import logging
import uuid
import signal
import zmq


class CommandLineInterface(object):

    """Implements a command line interface to the NRM."""

    def __init__(self):
        self.logger = logging.getLogger(__name__)

    def do_signal(self):
        pass

    def setup(self):
        # SUB port to the upstream API (connected to its PUB port)
        upstream_sub_port = 2345
        # PUB port to the upstream API (connected to its SUB port)
        upstream_pub_port = 3456

        self.context = zmq.Context()
        self.upstream_pub_socket = self.context.socket(zmq.PUB)
        self.upstream_sub_socket = self.context.socket(zmq.SUB)

        upstream_pub_param = "tcp://*:%d" % (upstream_pub_port)
        upstream_sub_param = "tcp://localhost:%d" % (upstream_sub_port)

        self.upstream_pub_socket.bind(upstream_pub_param)
        self.upstream_sub_socket.connect(upstream_sub_param)
        # we want to receive everything for now
        upstream_sub_filter = ""
        self.upstream_sub_socket.setsockopt(zmq.SUBSCRIBE, upstream_sub_filter)

        self.logger.info("upstream pub socket bound to: %s",
                         upstream_pub_param)
        self.logger.info("upstream sub socket connected to: %s",
                         upstream_sub_param)

        # take care of signals
        signal.signal(signal.SIGINT, self.do_signal)

        # create a uuid for this client instance
        self.uuid = str(uuid.uuid4())
        self.logger.info("client uuid: %r", self.uuid)

    def do_run(self, argv):
53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77
        """ Connect to the NRM and ask to spawn a container and run a command
        in it.

        The NRM should notify us on the pub socket of the container
        creation."""

        # build the command as a JSON dict containing enough info. We add to
        # the command a container uuid as a way to make sure that we can make
        # the command idempotent.
        containerid = str(uuid.uuid4())
        command = {'command': 'run',
                   'manifest': argv.manifest,
                   'file': argv.command,
                   'args': argv.args,
                   'uuid': containerid,
                   }

        while(True):
            self.upstream_pub_socket.send_json(command)
            msg = self.upstream_sub_socket.recv_json()
            self.logger.info("new message: %r", msg)
            # ignore other messages
            if isinstance(msg, dict) and msg.get('type') == 'container':
                if msg['uuid'] == containerid:
                    pass
Swann Perarnau's avatar
Swann Perarnau committed
78 79 80 81 82 83

    def do_setpower(self, argv):
        """ Connect to the NRM and ask to change the power limit.

        The NRM should answer on the pub socket with an acknowledgment."""

84 85 86 87 88 89 90 91
        # build the command as a JSON dict giving enough info. This is an
        # idempotent command, so we will repeat the command if we don't get a
        # timely answer.
        # TODO: check that the level makes a little bit of sense in the first
        # place
        command = {'command': 'setpower',
                   'limit': argv.limit,
                   }
Swann Perarnau's avatar
Swann Perarnau committed
92 93

        while(True):
94 95
            self.upstream_pub_socket.send_json(command)
            msg = self.upstream_sub_socket.recv_json()
Swann Perarnau's avatar
Swann Perarnau committed
96
            self.logger.info("new message: %r", msg)
97 98 99 100 101
            # ignore other messages
            if isinstance(msg, dict) and msg.get('type') == 'power':
                if msg['limit'] == argv.limit:
                    self.logger.info("command received by the daemon")
                    break
Swann Perarnau's avatar
Swann Perarnau committed
102 103 104 105 106 107 108 109 110 111

    def main(self):
        parser = argparse.ArgumentParser()
        parser.add_argument("-v", "--verbose",
                            help="verbose logging information",
                            action='store_true')
        subparsers = parser.add_subparsers()

        # run container
        parser_run = subparsers.add_parser("run")
112 113 114 115
        parser_run.add_argument("manifest", help="manifest file to apply")
        parser_run.add_argument("command", help="command to execute")
        parser_run.add_argument("args", help="command arguments",
                                nargs=argparse.REMAINDER)
Swann Perarnau's avatar
Swann Perarnau committed
116 117 118 119 120 121 122
        parser_run.set_defaults(func=self.do_run)

        # setpowerlimit
        parser_setpower = subparsers.add_parser("setpower")
        parser_setpower.add_argument("-f", "--follow",
                                     help="listen for power changes",
                                     action='store_true')
123
        parser_setpower.add_argument("limit",
Swann Perarnau's avatar
Swann Perarnau committed
124 125 126 127 128 129 130 131 132 133 134 135 136 137 138
                                     help="set new power limit",
                                     type=float)
        parser_setpower.set_defaults(func=self.do_setpower)
        args = parser.parse_args()
        if args.verbose:
            self.logger.setLevel(logging.DEBUG)

        self.setup()
        args.func(args)


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    cli = CommandLineInterface()
    cli.main()