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

        # the first message tells us if we started a container or not
        msg = self.client.recvmsg()
        assert msg.api == 'up_rpc_rep'
67
        assert msg.type == 'process_start'
68
        state = 'started'
69
        while(True):
70 71
            msg = self.client.recvmsg()
            assert msg.api == 'up_rpc_rep'
72
            assert msg.type in ['stdout', 'stderr', 'exit', 'process_exit']
73

74
            if msg.type == 'stdout':
75 76 77 78 79 80 81 82 83
                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)
84 85 86 87
                state = 'exiting'
                exitmsg = msg
            else:
                logger.error("unexpected message: %r", msg)
88 89
            if outeof and erreof and state == 'exiting':
                state = 'exit'
90
                logger.info("command ended: %r", exitmsg)
91
                break
Swann Perarnau's avatar
Swann Perarnau committed
92

93 94 95 96
    def do_list(self, argv):
        """Connect to the NRM and ask to list the containers present on the
        system.

97 98
        The NRM should respond to us with one message listing all
        containers."""
99

100 101 102 103 104 105 106 107
        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)
108

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

112 113
        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
114

115 116 117
        command = {'api': 'up_rpc_req',
                   'type': 'kill',
                   'container_uuid': argv.uuid
Swann Perarnau's avatar
Swann Perarnau committed
118
                   }
119 120 121 122 123 124
        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
125

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

129
        The NRM should answer with an acknowledgment."""
Swann Perarnau's avatar
Swann Perarnau committed
130

131 132 133 134 135
        # 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
136 137 138
        command = {'api': 'up_rpc_req',
                   'type': 'setpower',
                   'limit': str(argv.limit),
139
                   }
140 141 142 143 144 145
        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
146 147 148 149 150 151 152 153 154 155

    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")
156 157 158 159
        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)
160 161 162
        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
163 164
        parser_run.set_defaults(func=self.do_run)

Swann Perarnau's avatar
Swann Perarnau committed
165 166 167 168 169
        # 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)

170 171 172 173
        # list containers
        parser_list = subparsers.add_parser("list")
        parser_list.set_defaults(func=self.do_list)

Swann Perarnau's avatar
Swann Perarnau committed
174 175 176 177 178
        # setpowerlimit
        parser_setpower = subparsers.add_parser("setpower")
        parser_setpower.add_argument("-f", "--follow",
                                     help="listen for power changes",
                                     action='store_true')
179
        parser_setpower.add_argument("limit",
Swann Perarnau's avatar
Swann Perarnau committed
180 181 182
                                     help="set new power limit",
                                     type=float)
        parser_setpower.set_defaults(func=self.do_setpower)
183

Swann Perarnau's avatar
Swann Perarnau committed
184 185
        args = parser.parse_args()
        if args.verbose:
186
            logger.setLevel(logging.DEBUG)
Swann Perarnau's avatar
Swann Perarnau committed
187 188 189 190 191 192 193 194 195

        self.setup()
        args.func(args)


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