1"""A basic kernel monitor with autorestarting.
2
3This watches a kernel's state using KernelManager.is_alive and auto
4restarts the kernel if it dies.
5
6It is an incomplete base class, and must be subclassed.
7"""
8
9# Copyright (c) Jupyter Development Team.
10# Distributed under the terms of the Modified BSD License.
11
12from traitlets.config.configurable import LoggingConfigurable
13from traitlets import (
14    Instance, Float, Dict, Bool, Integer,
15)
16
17
18class KernelRestarter(LoggingConfigurable):
19    """Monitor and autorestart a kernel."""
20
21    kernel_manager = Instance('jupyter_client.KernelManager')
22
23    debug = Bool(False, config=True,
24        help="""Whether to include every poll event in debugging output.
25
26        Has to be set explicitly, because there will be *a lot* of output.
27        """
28    )
29
30    time_to_dead = Float(3.0, config=True,
31        help="""Kernel heartbeat interval in seconds."""
32    )
33
34    restart_limit = Integer(5, config=True,
35        help="""The number of consecutive autorestarts before the kernel is presumed dead."""
36    )
37
38    random_ports_until_alive = Bool(True, config=True,
39        help="""Whether to choose new random ports when restarting before the kernel is alive."""
40    )
41    _restarting = Bool(False)
42    _restart_count = Integer(0)
43    _initial_startup = Bool(True)
44
45    callbacks = Dict()
46    def _callbacks_default(self):
47        return dict(restart=[], dead=[])
48
49    def start(self):
50        """Start the polling of the kernel."""
51        raise NotImplementedError("Must be implemented in a subclass")
52
53    def stop(self):
54        """Stop the kernel polling."""
55        raise NotImplementedError("Must be implemented in a subclass")
56
57    def add_callback(self, f, event='restart'):
58        """register a callback to fire on a particular event
59
60        Possible values for event:
61
62          'restart' (default): kernel has died, and will be restarted.
63          'dead': restart has failed, kernel will be left dead.
64
65        """
66        self.callbacks[event].append(f)
67
68    def remove_callback(self, f, event='restart'):
69        """unregister a callback to fire on a particular event
70
71        Possible values for event:
72
73          'restart' (default): kernel has died, and will be restarted.
74          'dead': restart has failed, kernel will be left dead.
75
76        """
77        try:
78            self.callbacks[event].remove(f)
79        except ValueError:
80            pass
81
82    def _fire_callbacks(self, event):
83        """fire our callbacks for a particular event"""
84        for callback in self.callbacks[event]:
85            try:
86                callback()
87            except Exception as e:
88                self.log.error("KernelRestarter: %s callback %r failed", event, callback, exc_info=True)
89
90    def poll(self):
91        if self.debug:
92            self.log.debug('Polling kernel...')
93        if self.kernel_manager.shutting_down:
94            self.log.debug('Kernel shutdown in progress...')
95            return
96        if not self.kernel_manager.is_alive():
97            if self._restarting:
98                self._restart_count += 1
99            else:
100                self._restart_count = 1
101
102            if self._restart_count >= self.restart_limit:
103                self.log.warning("KernelRestarter: restart failed")
104                self._fire_callbacks('dead')
105                self._restarting = False
106                self._restart_count = 0
107                self.stop()
108            else:
109                newports = self.random_ports_until_alive and self._initial_startup
110                self.log.info('KernelRestarter: restarting kernel (%i/%i), %s random ports',
111                    self._restart_count,
112                    self.restart_limit,
113                    'new' if newports else 'keep'
114                )
115                self._fire_callbacks('restart')
116                self.kernel_manager.restart_kernel(now=True, newports=newports)
117                self._restarting = True
118        else:
119            if self._initial_startup:
120                self._initial_startup = False
121            if self._restarting:
122                self.log.debug("KernelRestarter: restart apparently succeeded")
123            self._restarting = False
124