Merge pull request #2963 from takluyver/shutdown-no-kernels

minrk · web-flow · commit a2f72da17f3d · 2017-11-20T09:49:20.000+01:00
Config option to shut down server after N seconds with no kernels
diff --git a/notebook/notebookapp.py b/notebook/notebookapp.py
@@ -104,7 +104,7 @@
 from jupyter_core.paths import jupyter_runtime_dir, jupyter_path
 from notebook._sysinfo import get_sys_info
 
-from ._tz import utcnow
+from ._tz import utcnow, utcfromtimestamp
 from .utils import url_path_join, check_pid, url_escape
 
 #-----------------------------------------------------------------------------
@@ -336,6 +336,26 @@ def init_handlers(self, settings):
         new_handlers.append((r'(.*)', Template404))
         return new_handlers
 
+    def last_activity(self):
+        """Get a UTC timestamp for when the server last did something.
+
+        Includes: API activity, kernel activity, kernel shutdown, and terminal
+        activity.
+        """
+        sources = [
+            self.settings['started'],
+            self.settings['kernel_manager'].last_kernel_activity,
+        ]
+        try:
+            sources.append(self.settings['api_last_activity'])
+        except KeyError:
+            pass
+        try:
+            sources.append(self.settings['terminal_last_activity'])
+        except KeyError:
+            pass
+        return max(sources)
+
 
 class NotebookPasswordApp(JupyterApp):
     """Set a password for the notebook server.
@@ -1139,6 +1159,16 @@ def _update_server_extensions(self, change):
     rate_limit_window = Float(3, config=True, help=_("""(sec) Time window used to 
         check the message and data rate limits."""))
 
+    shutdown_no_activity_timeout = Integer(0, config=True,
+        help=("Shut down the server after N seconds with no kernels or "
+              "terminals running and no activity. "
+              "This can be used together with culling idle kernels "
+              "(MappingKernelManager.cull_idle_timeout) to "
+              "shutdown the notebook server when it's not in use. This is not "
+              "precisely timed: it may shut down up to a minute later. "
+              "0 (the default) disables this automatic shutdown.")
+    )
+
     def parse_command_line(self, argv=None):
         super(NotebookApp, self).parse_command_line(argv)
 
@@ -1426,6 +1456,37 @@ def init_mime_overrides(self):
         # mimetype always needs to be text/css, so we override it here.
         mimetypes.add_type('text/css', '.css')
 
+
+    def shutdown_no_activity(self):
+        """Shutdown server on timeout when there are no kernels or terminals."""
+        km = self.kernel_manager
+        if len(km) != 0:
+            return   # Kernels still running
+
+        try:
+            term_mgr = self.web_app.settings['terminal_manager']
+        except KeyError:
+            pass  # Terminals not enabled
+        else:
+            if term_mgr.terminals:
+                return   # Terminals still running
+
+        seconds_since_active = \
+            (utcnow() - self.web_app.last_activity()).total_seconds()
+        self.log.debug("No activity for %d seconds.",
+                       seconds_since_active)
+        if seconds_since_active > self.shutdown_no_activity_timeout:
+            self.log.info("No kernels or terminals for %d seconds; shutting down.",
+                          seconds_since_active)
+            self.stop()
+
+    def init_shutdown_no_activity(self):
+        if self.shutdown_no_activity_timeout > 0:
+            self.log.info("Will shut down after %d seconds with no kernels or terminals.",
+                          self.shutdown_no_activity_timeout)
+            pc = ioloop.PeriodicCallback(self.shutdown_no_activity, 60000)
+            pc.start()
+
     @catch_config_error
     def initialize(self, argv=None):
         super(NotebookApp, self).initialize(argv)
@@ -1439,6 +1500,7 @@ def initialize(self, argv=None):
         self.init_signal()
         self.init_server_extensions()
         self.init_mime_overrides()
+        self.init_shutdown_no_activity()
 
     def cleanup_kernels(self):
         """Shutdown all kernels.
diff --git a/notebook/services/api/handlers.py b/notebook/services/api/handlers.py
@@ -33,14 +33,11 @@ class APIStatusHandler(APIHandler):
     def get(self):
         # if started was missing, use unix epoch
         started = self.settings.get('started', utcfromtimestamp(0))
-        # if we've never seen API activity, use started date
-        api_last_activity = self.settings.get('api_last_activity', started)
         started = isoformat(started)
-        api_last_activity = isoformat(api_last_activity)
 
         kernels = yield gen.maybe_future(self.kernel_manager.list_kernels())
         total_connections = sum(k['connections'] for k in kernels)
-        last_activity = max(chain([api_last_activity], [k['last_activity'] for k in kernels]))
+        last_activity = isoformat(self.application.last_activity())
         model = {
             'started': started,
             'last_activity': last_activity,
diff --git a/notebook/services/kernels/kernelmanager.py b/notebook/services/kernels/kernelmanager.py
@@ -8,6 +8,7 @@
 # Distributed under the terms of the Modified BSD License.
 
 from collections import defaultdict
+from datetime import datetime, timedelta
 from functools import partial
 import os
 
@@ -17,14 +18,14 @@
 
 from jupyter_client.session import Session
 from jupyter_client.multikernelmanager import MultiKernelManager
-from traitlets import Any, Bool, Dict, List, Unicode, TraitError, Integer, default, validate
+from traitlets import (Any, Bool, Dict, List, Unicode, TraitError, Integer,
+       Instance, default, validate
+)
 
 from notebook.utils import to_os_path, exists
 from notebook._tz import utcnow, isoformat
 from ipython_genutils.py3compat import getcwd
 
-from datetime import timedelta
-
 
 class MappingKernelManager(MultiKernelManager):
     """A KernelManager that handles notebook mapping and HTTP error handling"""
@@ -98,6 +99,13 @@ def _update_root_dir(self, proposal):
     def _default_kernel_buffers(self):
         return defaultdict(lambda: {'buffer': [], 'session_key': '', 'channels': {}})
 
+    last_kernel_activity = Instance(datetime,
+        help="The last activity on any kernel, including shutting down a kernel")
+
+    def __init__(self, **kwargs):
+        super(MappingKernelManager, self).__init__(**kwargs)
+        self.last_kernel_activity = utcnow()
+
     #-------------------------------------------------------------------------
     # Methods for managing kernels and sessions
     #-------------------------------------------------------------------------
@@ -255,6 +263,7 @@ def shutdown_kernel(self, kernel_id, now=False):
         kernel._activity_stream.close()
         self.stop_buffering(kernel_id)
         self._kernel_connections.pop(kernel_id, None)
+        self.last_kernel_activity = utcnow()
         return super(MappingKernelManager, self).shutdown_kernel(kernel_id, now=now)
 
     def restart_kernel(self, kernel_id):
@@ -360,7 +369,7 @@ def start_watching_activity(self, kernel_id):
 
         def record_activity(msg_list):
             """Record an IOPub message arriving from a kernel"""
-            kernel.last_activity = utcnow()
+            self.last_kernel_activity = kernel.last_activity = utcnow()
 
             idents, fed_msg_list = session.feed_identities(msg_list)
             msg = session.deserialize(fed_msg_list)
diff --git a/notebook/terminal/handlers.py b/notebook/terminal/handlers.py
@@ -6,6 +6,7 @@
 
 from tornado import web
 import terminado
+from notebook._tz import utcnow
 from ..base.handlers import IPythonHandler
 from ..base.zmqhandlers import WebSocketMixin
 
@@ -31,4 +32,11 @@ def get(self, *args, **kwargs):
         if not self.get_current_user():
             raise web.HTTPError(403)
         return super(TermSocket, self).get(*args, **kwargs)
-    
+
+    def on_message(self, message):
+        super(TermSocket, self).on_message(message)
+        self.application.settings['terminal_last_activity'] = utcnow()
+
+    def write_message(self, message, binary=False):
+        super(TermSocket, self).write_message(message, binary=binary)
+        self.application.settings['terminal_last_activity'] = utcnow()
