cmd 7.9 KB
Newer Older
Swann Perarnau's avatar
Swann Perarnau committed
1 2 3 4 5 6 7 8 9
#!/usr/bin/env python2

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

10 11
logger = logging.getLogger('nrm-cmd')

Swann Perarnau's avatar
Swann Perarnau committed
12 13 14 15 16 17

class CommandLineInterface(object):

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

    def __init__(self):
18
        pass
Swann Perarnau's avatar
Swann Perarnau committed
19

20
    def do_signal(self, signum, stackframe):
21
        logger.info("received signal %d, exiting", signum)
22
        exit(1)
Swann Perarnau's avatar
Swann Perarnau committed
23 24 25 26 27 28 29 30 31 32 33

    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)

34
        upstream_pub_param = "tcp://localhost:%d" % (upstream_pub_port)
Swann Perarnau's avatar
Swann Perarnau committed
35 36
        upstream_sub_param = "tcp://localhost:%d" % (upstream_sub_port)

37
        self.upstream_pub_socket.connect(upstream_pub_param)
Swann Perarnau's avatar
Swann Perarnau committed
38 39 40 41 42
        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)

43 44
        logger.info("upstream pub socket bound to: %s", upstream_pub_param)
        logger.info("upstream sub socket connected to: %s", upstream_sub_param)
Swann Perarnau's avatar
Swann Perarnau committed
45 46 47 48 49 50

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

        # create a uuid for this client instance
        self.uuid = str(uuid.uuid4())
51
        logger.info("client uuid: %r", self.uuid)
Swann Perarnau's avatar
Swann Perarnau committed
52 53

    def do_run(self, argv):
54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69
        """ 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,
                   }
70 71 72 73 74 75
        # command fsm
        state = 'init'
        outeof = False
        erreof = False
        exitmsg = None
        self.upstream_pub_socket.send_json(command)
76 77 78 79
        while(True):
            msg = self.upstream_sub_socket.recv_json()
            if isinstance(msg, dict) and msg.get('type') == 'container':
                if msg['uuid'] == containerid:
80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104
                    if msg['event'] == 'start':
                        if state == 'init':
                            state = 'started'
                            logger.info("container started: %r", msg)
                        else:
                            logger.info("unexpected start message: %r", state)
                            exit(1)
                    elif msg['event'] == 'stdout':
                        logger.info("container msg: %r", msg)
                        if msg['payload'] == 'eof':
                            outeof = True
                    elif msg['event'] == 'stderr':
                        logger.info("container msg: %r", msg)
                        if msg['payload'] == 'eof':
                            erreof = True
                    elif msg['event'] == 'exit':
                        if state == 'started':
                            state = 'exiting'
                            exitmsg = msg
                        else:
                            logger.info("unexpected exit message: %r", msg)
            if outeof and erreof and state == 'exiting':
                state = 'exit'
                logger.info("container ended: %r", exitmsg)
                break
Swann Perarnau's avatar
Swann Perarnau committed
105

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

116
        self.upstream_pub_socket.send_json(command)
117 118
        while(True):
            msg = self.upstream_sub_socket.recv_json()
119
            logger.info("new message: %r", msg)
120 121 122
            # ignore other messages
            if isinstance(msg, dict) and msg.get('type') == 'container':
                if msg['event'] == 'list':
123
                    logger.info("list response: %r", msg)
124 125
                    break

Swann Perarnau's avatar
Swann Perarnau committed
126 127 128 129 130 131 132 133 134 135
    def do_kill(self, argv):
        """Connect to the NRM and ask to kill a container by uuid.

        The NRM should respond to us on the pub socket with a message
        containing the exit status of the top process of the container."""

        command = {'command': 'kill',
                   'uuid': argv.uuid
                   }

136
        self.upstream_pub_socket.send_json(command)
Swann Perarnau's avatar
Swann Perarnau committed
137 138
        while(True):
            msg = self.upstream_sub_socket.recv_json()
139
            logger.info("new message: %r", msg)
Swann Perarnau's avatar
Swann Perarnau committed
140 141 142
            # ignore other messages
            if isinstance(msg, dict) and msg.get('type') == 'container':
                if msg['event'] == 'exit' and msg['uuid'] == argv.uuid:
143
                    logger.info("container exit: %r", msg)
Swann Perarnau's avatar
Swann Perarnau committed
144 145
                    break

Swann Perarnau's avatar
Swann Perarnau committed
146 147 148 149 150
    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."""

151 152 153 154 155 156 157 158
        # 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
159

160
        self.upstream_pub_socket.send_json(command)
Swann Perarnau's avatar
Swann Perarnau committed
161
        while(True):
162
            msg = self.upstream_sub_socket.recv_json()
163
            logger.info("new message: %r", msg)
164 165 166
            # ignore other messages
            if isinstance(msg, dict) and msg.get('type') == 'power':
                if msg['limit'] == argv.limit:
167
                    logger.info("command received by the daemon")
168
                    break
Swann Perarnau's avatar
Swann Perarnau committed
169 170 171 172 173 174 175 176 177 178

    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")
179 180 181 182
        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
183 184
        parser_run.set_defaults(func=self.do_run)

Swann Perarnau's avatar
Swann Perarnau committed
185 186 187 188 189
        # kill container
        parser_kill = subparsers.add_parser("kill")
        parser_kill.add_argument("uuid", help="uuid of the container")
        parser_kill.set_defaults(func=self.do_kill)

190 191 192 193
        # list containers
        parser_list = subparsers.add_parser("list")
        parser_list.set_defaults(func=self.do_list)

Swann Perarnau's avatar
Swann Perarnau committed
194 195 196 197 198
        # setpowerlimit
        parser_setpower = subparsers.add_parser("setpower")
        parser_setpower.add_argument("-f", "--follow",
                                     help="listen for power changes",
                                     action='store_true')
199
        parser_setpower.add_argument("limit",
Swann Perarnau's avatar
Swann Perarnau committed
200 201 202
                                     help="set new power limit",
                                     type=float)
        parser_setpower.set_defaults(func=self.do_setpower)
203

Swann Perarnau's avatar
Swann Perarnau committed
204 205
        args = parser.parse_args()
        if args.verbose:
206
            logger.setLevel(logging.DEBUG)
Swann Perarnau's avatar
Swann Perarnau committed
207 208 209 210 211 212 213 214 215

        self.setup()
        args.func(args)


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