cmd 7.04 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

    def do_run(self, argv):
38 39 40
        """ Connect to the NRM and ask to spawn a container and run a command
        in it.

41
        The NRM should reply for container info."""
42 43 44 45

        # 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.
46
        environ = os.environ
47 48
        command = {'api': 'up_rpc_req',
                   'type': 'run',
49
                   'manifest': argv.manifest,
50
                   'path': argv.command,
51
                   'args': argv.args,
52
                   'environ': dict(environ),
53
                   'container_uuid': argv.ucontainername or str(uuid.uuid4()),
54
                   }
55
        msg = RPC_MSG['run'](**command)
56 57 58 59 60
        # command fsm
        state = 'init'
        outeof = False
        erreof = False
        exitmsg = None
61
        self.client.sendmsg(msg)
62 63 64 65 66 67 68 69 70 71 72

        # the first message tells us if we started a container or not
        msg = self.client.recvmsg()
        assert msg.api == 'up_rpc_rep'
        assert msg.type in ['start', 'process_start']
        new_container = False
        if msg.type == 'start':
            new_container = True
            msg = self.client.recvmsg()
            assert msg.type == 'process_start'
        state = 'started'
73
        while(True):
74 75
            msg = self.client.recvmsg()
            assert msg.api == 'up_rpc_rep'
76
            assert msg.type in ['stdout', 'stderr', 'exit', 'process_exit']
77

78
            if msg.type == 'stdout':
79 80 81 82 83 84 85 86 87
                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)
88 89
                if not new_container:
                    state = 'exiting'
90 91 92 93 94 95
            elif msg.type == 'exit':
                if state == 'started':
                    state = 'exiting'
                    exitmsg = msg
                else:
                    logger.info("unexpected exit message: %r", msg)
96 97 98 99
            if outeof and erreof and state == 'exiting':
                state = 'exit'
                logger.info("container ended: %r", exitmsg)
                break
Swann Perarnau's avatar
Swann Perarnau committed
100

101 102 103 104
    def do_list(self, argv):
        """Connect to the NRM and ask to list the containers present on the
        system.

105 106
        The NRM should respond to us with one message listing all
        containers."""
107

108 109 110 111 112 113 114 115
        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)
116

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

120 121
        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
122

123 124 125
        command = {'api': 'up_rpc_req',
                   'type': 'kill',
                   'container_uuid': argv.uuid
Swann Perarnau's avatar
Swann Perarnau committed
126
                   }
127 128 129 130 131 132
        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
133

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

137
        The NRM should answer with an acknowledgment."""
Swann Perarnau's avatar
Swann Perarnau committed
138

139 140 141 142 143
        # 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
144 145 146
        command = {'api': 'up_rpc_req',
                   'type': 'setpower',
                   'limit': str(argv.limit),
147
                   }
148 149 150 151 152 153
        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
154 155 156 157 158 159 160 161 162 163

    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")
164 165 166 167
        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)
168 169 170
        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
171 172
        parser_run.set_defaults(func=self.do_run)

Swann Perarnau's avatar
Swann Perarnau committed
173 174 175 176 177
        # 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)

178 179 180 181
        # list containers
        parser_list = subparsers.add_parser("list")
        parser_list.set_defaults(func=self.do_list)

Swann Perarnau's avatar
Swann Perarnau committed
182 183 184 185 186
        # setpowerlimit
        parser_setpower = subparsers.add_parser("setpower")
        parser_setpower.add_argument("-f", "--follow",
                                     help="listen for power changes",
                                     action='store_true')
187
        parser_setpower.add_argument("limit",
Swann Perarnau's avatar
Swann Perarnau committed
188 189 190
                                     help="set new power limit",
                                     type=float)
        parser_setpower.set_defaults(func=self.do_setpower)
191

Swann Perarnau's avatar
Swann Perarnau committed
192 193
        args = parser.parse_args()
        if args.verbose:
194
            logger.setLevel(logging.DEBUG)
Swann Perarnau's avatar
Swann Perarnau committed
195 196 197 198 199 200 201 202 203

        self.setup()
        args.func(args)


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