Merge branch '2881-node-manager'
[arvados.git] / services / nodemanager / arvnodeman / launcher.py
1 #!/usr/bin/env python
2
3 from __future__ import absolute_import, print_function
4
5 import argparse
6 import logging
7 import signal
8 import sys
9 import time
10
11 import daemon
12 import pykka
13
14 from . import config as nmconfig
15 from .computenode import \
16     ComputeNodeSetupActor, ComputeNodeShutdownActor, ComputeNodeUpdateActor, \
17     ShutdownTimer
18 from .daemon import NodeManagerDaemonActor
19 from .jobqueue import JobQueueMonitorActor, ServerCalculator
20 from .nodelist import ArvadosNodeListMonitorActor, CloudNodeListMonitorActor
21 from .timedcallback import TimedCallBackActor
22
23 node_daemon = None
24
25 def abort(msg, code=1):
26     print("arvados-node-manager: " + msg)
27     sys.exit(code)
28
29 def parse_cli(args):
30     parser = argparse.ArgumentParser(
31         prog='arvados-node-manager',
32         description="Dynamically allocate Arvados cloud compute nodes")
33     parser.add_argument(
34         '--foreground', action='store_true', default=False,
35         help="Run in the foreground.  Don't daemonize.")
36     parser.add_argument(
37         '--config', help="Path to configuration file")
38     return parser.parse_args(args)
39
40 def load_config(path):
41     if not path:
42         abort("No --config file specified", 2)
43     config = nmconfig.NodeManagerConfig()
44     try:
45         with open(path) as config_file:
46             config.readfp(config_file)
47     except (IOError, OSError) as error:
48         abort("Error reading configuration file {}: {}".format(path, error))
49     return config
50
51 def setup_logging(path, level, **sublevels):
52     handler = logging.FileHandler(path)
53     handler.setFormatter(logging.Formatter(
54             '%(asctime)s %(name)s[%(process)d] %(levelname)s: %(message)s',
55             '%Y-%m-%d %H:%M:%S'))
56     root_logger = logging.getLogger()
57     root_logger.addHandler(handler)
58     root_logger.setLevel(level)
59     for logger_name, sublevel in sublevels.iteritems():
60         sublogger = logging.getLogger(logger_name)
61         sublogger.setLevel(sublevel)
62
63 def launch_pollers(config):
64     cloud_client = config.new_cloud_client()
65     arvados_client = config.new_arvados_client()
66     cloud_size_list = config.node_sizes(cloud_client.list_sizes())
67     if not cloud_size_list:
68         abort("No valid node sizes configured")
69
70     server_calculator = ServerCalculator(
71         cloud_size_list, config.getint('Daemon', 'max_nodes'))
72     poll_time = config.getint('Daemon', 'poll_time')
73     max_poll_time = config.getint('Daemon', 'max_poll_time')
74
75     timer = TimedCallBackActor.start(poll_time / 10.0).proxy()
76     cloud_node_poller = CloudNodeListMonitorActor.start(
77         cloud_client, timer, poll_time, max_poll_time).proxy()
78     arvados_node_poller = ArvadosNodeListMonitorActor.start(
79         arvados_client, timer, poll_time, max_poll_time).proxy()
80     job_queue_poller = JobQueueMonitorActor.start(
81         config.new_arvados_client(), timer, server_calculator,
82         poll_time, max_poll_time).proxy()
83     return timer, cloud_node_poller, arvados_node_poller, job_queue_poller
84
85 _caught_signals = {}
86 def shutdown_signal(signal_code, frame):
87     current_count = _caught_signals.get(signal_code, 0)
88     _caught_signals[signal_code] = current_count + 1
89     if node_daemon is None:
90         pykka.ActorRegistry.stop_all()
91         sys.exit(-signal_code)
92     elif current_count == 0:
93         node_daemon.shutdown()
94     elif current_count == 1:
95         pykka.ActorRegistry.stop_all()
96     else:
97         sys.exit(-signal_code)
98
99 def main(args=None):
100     global node_daemon
101     args = parse_cli(args)
102     config = load_config(args.config)
103
104     if not args.foreground:
105         daemon.DaemonContext().open()
106     for sigcode in [signal.SIGINT, signal.SIGQUIT, signal.SIGTERM]:
107         signal.signal(sigcode, shutdown_signal)
108
109     setup_logging(config.get('Logging', 'file'), **config.log_levels())
110     timer, cloud_node_poller, arvados_node_poller, job_queue_poller = \
111         launch_pollers(config)
112     cloud_node_updater = ComputeNodeUpdateActor.start(
113         config.new_cloud_client).proxy()
114     node_daemon = NodeManagerDaemonActor.start(
115         job_queue_poller, arvados_node_poller, cloud_node_poller,
116         cloud_node_updater, timer,
117         config.new_arvados_client, config.new_cloud_client,
118         config.shutdown_windows(), config.getint('Daemon', 'max_nodes'),
119         config.getint('Daemon', 'poll_stale_after'),
120         config.getint('Daemon', 'node_stale_after')).proxy()
121
122     signal.pause()
123     daemon_stopped = node_daemon.actor_ref.actor_stopped.is_set
124     while not daemon_stopped():
125         time.sleep(1)
126     pykka.ActorRegistry.stop_all()
127
128
129 if __name__ == '__main__':
130     main()