cmd 5.9 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
#!/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__)

18 19 20
    def do_signal(self, signum, stackframe):
        self.logger.info("received signal %d, exiting", signum)
        exit(1)
Swann Perarnau's avatar
Swann Perarnau committed
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 53

    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):
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:
78 79
                    self.logger.info("container response: %r", msg)
                    break
Swann Perarnau's avatar
Swann Perarnau committed
80

81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100
    def do_list(self, argv):
        """Connect to the NRM and ask to list the containers present on the
        system.

        The NRM should respond to us on the pub socket with one message listing
        all containers."""

        command = {'command': 'list',
                   }

        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['event'] == 'list':
                    self.logger.info("list response: %r", msg)
                    break

Swann Perarnau's avatar
Swann Perarnau committed
101 102 103 104 105
    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."""

106 107 108 109 110 111 112 113
        # 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
114 115

        while(True):
116 117
            self.upstream_pub_socket.send_json(command)
            msg = self.upstream_sub_socket.recv_json()
Swann Perarnau's avatar
Swann Perarnau committed
118
            self.logger.info("new message: %r", msg)
119 120 121 122 123
            # 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
124 125 126 127 128 129 130 131 132 133

    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")
134 135 136 137
        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
138 139
        parser_run.set_defaults(func=self.do_run)

140 141 142 143
        # list containers
        parser_list = subparsers.add_parser("list")
        parser_list.set_defaults(func=self.do_list)

Swann Perarnau's avatar
Swann Perarnau committed
144 145 146 147 148
        # setpowerlimit
        parser_setpower = subparsers.add_parser("setpower")
        parser_setpower.add_argument("-f", "--follow",
                                     help="listen for power changes",
                                     action='store_true')
149
        parser_setpower.add_argument("limit",
Swann Perarnau's avatar
Swann Perarnau committed
150 151 152
                                     help="set new power limit",
                                     type=float)
        parser_setpower.set_defaults(func=self.do_setpower)
153

Swann Perarnau's avatar
Swann Perarnau committed
154 155 156 157 158 159 160 161 162 163 164 165
        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()