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

from __future__ import print_function
import argparse
import logging
import signal
7
import os
8
import nrm.messaging
9
import uuid
Swann Perarnau's avatar
Swann Perarnau committed
10

11 12
RPC_MSG = nrm.messaging.MSGTYPES['up_rpc_req']
logger = logging.getLogger('nrm')
13

Swann Perarnau's avatar
Swann Perarnau committed
14 15 16 17 18 19

class CommandLineInterface(object):

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

    def __init__(self):
20
        pass
Swann Perarnau's avatar
Swann Perarnau committed
21

22
    def do_signal(self, signum, stackframe):
23
        logger.info("received signal %d, exiting", signum)
24
        exit(1)
Swann Perarnau's avatar
Swann Perarnau committed
25 26

    def setup(self):
27 28 29 30
        # upstream RPC port
        upstream_client_port = 3456
        upstream_client_param = "tcp://localhost:%d" % (upstream_client_port)
        self.client = nrm.messaging.UpstreamRPCClient(upstream_client_param)
Swann Perarnau's avatar
Swann Perarnau committed
31 32 33 34

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

35
        self.client.wait_connected()
Swann Perarnau's avatar
Swann Perarnau committed
36

37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52
    def do_listen(self, argv):
        """ Connect to the NRM and listen for pub/sub messages."""
        upstream_pub_port = 2345
        upstream_pub_param = "tcp://localhost:%d" % (upstream_pub_port)
        self.pub_client = nrm.messaging.UpstreamPubClient(upstream_pub_param)
        self.pub_client.wait_connected()

        while(True):
            msg = self.pub_client.recvmsg()
            if argv.uuid:
                uuid = getattr(msg, 'container_uuid', None)
                if argv.container == uuid:
                    logger.info("pub message", msg)
            else:
                logger.info("pub message", msg)

Swann Perarnau's avatar
Swann Perarnau committed
53
    def do_run(self, argv):
54 55 56
        """ Connect to the NRM and ask to spawn a container and run a command
        in it.

57
        The NRM should reply for container info."""
58 59 60 61

        # 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.
62
        environ = os.environ
63
        container_uuid = argv.ucontainername or str(uuid.uuid4())
64 65
        command = {'api': 'up_rpc_req',
                   'type': 'run',
66
                   'manifest': argv.manifest,
67
                   'path': argv.command,
68
                   'args': argv.args,
69
                   'environ': dict(environ),
70
                   'container_uuid': container_uuid,
71
                   }
72
        msg = RPC_MSG['run'](**command)
73 74 75 76 77
        # command fsm
        state = 'init'
        outeof = False
        erreof = False
        exitmsg = None
78
        self.client.sendmsg(msg)
79 80 81 82

        # the first message tells us if we started a container or not
        msg = self.client.recvmsg()
        assert msg.api == 'up_rpc_rep'
83
        assert msg.type == 'process_start'
84
        state = 'started'
85
        while(True):
86 87
            msg = self.client.recvmsg()
            assert msg.api == 'up_rpc_rep'
88
            assert msg.type in ['stdout', 'stderr', 'exit', 'process_exit']
89

90
            if msg.type == 'stdout':
91 92 93 94 95 96 97 98 99
                logger.info("container msg: %r", msg)
                if msg.payload == 'eof':
                    outeof = True
            elif msg.type == 'stderr':
                logger.info("container msg: %r", msg)
                if msg.payload == 'eof':
                    erreof = True
            elif msg.type == 'process_exit':
                logger.info("process ended: %r", msg)
100 101 102 103
                state = 'exiting'
                exitmsg = msg
            else:
                logger.error("unexpected message: %r", msg)
104 105
            if outeof and erreof and state == 'exiting':
                state = 'exit'
106
                logger.info("command ended: %r", exitmsg)
107
                break
Swann Perarnau's avatar
Swann Perarnau committed
108

109 110 111 112
    def do_list(self, argv):
        """Connect to the NRM and ask to list the containers present on the
        system.

113 114
        The NRM should respond to us with one message listing all
        containers."""
115

116 117 118 119 120 121 122 123
        command = {'api': 'up_rpc_req',
                   'type': 'list'}
        msg = RPC_MSG['list'](**command)
        self.client.sendmsg(msg)
        msg = self.client.recvmsg()
        assert msg.api == 'up_rpc_rep'
        assert msg.type == 'list'
        logger.info("list response: %r", msg)
124

Swann Perarnau's avatar
Swann Perarnau committed
125 126 127
    def do_kill(self, argv):
        """Connect to the NRM and ask to kill a container by uuid.

128 129
        The NRM should respond to us with a message containing the exit status
        of the top process of the container."""
Swann Perarnau's avatar
Swann Perarnau committed
130

131 132 133
        command = {'api': 'up_rpc_req',
                   'type': 'kill',
                   'container_uuid': argv.uuid
Swann Perarnau's avatar
Swann Perarnau committed
134
                   }
135 136 137 138 139 140
        msg = RPC_MSG['kill'](**command)
        self.client.sendmsg(msg)
        msg = self.client.recvmsg()
        assert msg.api == 'up_rpc_rep'
        assert msg.type == 'exit'
        logger.info("container exit: %r", msg)
Swann Perarnau's avatar
Swann Perarnau committed
141

Swann Perarnau's avatar
Swann Perarnau committed
142 143 144
    def do_setpower(self, argv):
        """ Connect to the NRM and ask to change the power limit.

145
        The NRM should answer with an acknowledgment."""
Swann Perarnau's avatar
Swann Perarnau committed
146

147 148 149 150 151
        # 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
152 153 154
        command = {'api': 'up_rpc_req',
                   'type': 'setpower',
                   'limit': str(argv.limit),
155
                   }
156 157 158 159 160 161
        msg = RPC_MSG['setpower'](**command)
        self.client.sendmsg(msg)
        msg = self.client.recvmsg()
        assert msg.api == 'up_rpc_rep'
        assert msg.type == 'getpower'
        logger.info("command received by the daemon: %r", msg)
Swann Perarnau's avatar
Swann Perarnau committed
162 163 164 165 166 167 168 169 170 171

    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")
172 173 174 175
        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)
176 177 178
        parser_run.add_argument("-u", "--ucontainername", help="""user-specified
                                name for container used to attach proceses""",
                                nargs='?', const=None, default=None)
Swann Perarnau's avatar
Swann Perarnau committed
179 180
        parser_run.set_defaults(func=self.do_run)

Swann Perarnau's avatar
Swann Perarnau committed
181 182 183 184 185
        # 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)

186 187 188 189
        # list containers
        parser_list = subparsers.add_parser("list")
        parser_list.set_defaults(func=self.do_list)

190 191 192 193 194 195 196
        # listen
        parser_listen = subparsers.add_parser("listen")
        parser_listen.add_argument("-u", "--uuid",
                                   help="container uuid to listen for",
                                   default=None)
        parser_listen.set_defaults(func=self.do_listen)

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

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

        self.setup()
        args.func(args)


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