Add server and profile options to the notebook() command

Jason Grout · Jason Grout · commit 1174ddc52ebf · 2012-05-22T10:05:00.000-05:00
The server now takes either 'twistd' or 'flask' as options.  The profile option saves profile data to a user-specified file.
diff --git a/sagenb/notebook/notebook_object.py b/sagenb/notebook/notebook_object.py
@@ -164,6 +164,14 @@ class NotebookObject:
           now, so if the machines are separate the server machine must
           NSF export ``/tmp``.
 
+        - ``server`` -- string ("twistd" (default) or "flask").  The server
+          to use to server content.
+
+        - ``profile`` -- True, False, or file prefix (default: False - no profiling),
+          If True, profiling is saved to a randomly-named file like `sagenb-*-profile*.stats`
+          in the $DOT_SAGE directory.  If a string, that string is used as a
+          prefix for the pstats data file.
+
         - ``ulimit`` -- string (initial default: None -- leave as is),
           if given and ``server_pool`` is also given, the worksheet
           processes are run with these constraints. See the ulimit
diff --git a/sagenb/notebook/run_notebook.py b/sagenb/notebook/run_notebook.py
@@ -36,6 +36,76 @@
 template_file = os.path.join(conf_path, 'cert.cfg')
 
 FLASK_NOTEBOOK_CONFIG = """
+####################################################################
+# WARNING -- Do not edit this file!   It is autogenerated each time
+# the notebook(...) command is executed.
+####################################################################
+
+import sagenb.notebook.misc
+sagenb.notebook.misc.DIR = %(cwd)r #We should really get rid of this!
+
+import signal, sys, random
+def save_notebook(notebook):
+    print "Quitting all running worksheets..."
+    notebook.quit()
+    print "Saving notebook..."
+    notebook.save()
+    print "Notebook cleanly saved."
+
+#########
+# Flask #
+#########
+import os, sys, random
+flask_dir = os.path.join(os.environ['SAGE_ROOT'], 'devel', 'sagenb', 'flask_version')
+sys.path.append(flask_dir)
+import base as flask_base
+opts={}
+startup_token = '{0:x}'.format(random.randint(0, 2**128))
+if %(login)s:
+    opts['startup_token'] = startup_token
+flask_app = flask_base.create_app(%(notebook_opts)s, **opts)
+sys.path.remove(flask_dir)
+
+if %(secure)s:
+    from OpenSSL import SSL
+    ssl_context = SSL.Context(SSL.SSLv23_METHOD)
+    ssl_context.use_privatekey_file(%(pkey)r)
+    ssl_context.use_certificate_file(%(cert)r)
+else:
+    ssl_context = None
+
+import logging
+logger=logging.getLogger('werkzeug')
+#logger.setLevel(logging.WARNING)
+#logger.setLevel(logging.INFO) # to see page requests
+logger.setLevel(logging.DEBUG)
+logger.addHandler(logging.StreamHandler())
+
+if %(secure)s:
+    # Monkey-patch werkzeug so that it works with pyOpenSSL and Python 2.7
+    # otherwise, we constantly get TypeError: shutdown() takes exactly 0 arguments (1 given)
+
+    # Monkey patching idiom: http://mail.python.org/pipermail/python-dev/2008-January/076194.html
+    def monkeypatch_method(cls):
+        def decorator(func):
+            setattr(cls, func.__name__, func)
+            return func
+        return decorator
+    from werkzeug import serving
+
+    @monkeypatch_method(serving.BaseWSGIServer)
+    def shutdown_request(self, request):
+        request.shutdown()
+
+%(open_page)s
+try:
+    flask_app.run(host=%(host)r, port=%(port)s, threaded=True,
+                  ssl_context=ssl_context, debug=False)
+finally:
+    save_notebook(flask_base.notebook)
+"""
+
+TWISTD_NOTEBOOK_CONFIG = """
 ####################################################################        
 # WARNING -- Do not edit this file!   It is autogenerated each time
 # the notebook(...) command is executed.
@@ -50,17 +120,10 @@
 
 from twisted.internet import reactor
 
-# Now set things up and start the notebook
-import sagenb.notebook.notebook
-sagenb.notebook.notebook.JSMATH=True
-import sagenb.notebook.notebook as notebook
-import sagenb.notebook.worksheet as worksheet
-
-import sagenb.notebook.misc as misc
-
-misc.DIR = %(cwd)r #We should really get rid of this!
+import sagenb.notebook.misc
+sagenb.notebook.misc.DIR = %(cwd)r #We should really get rid of this!
 
-import signal, sys, random
+import signal
 def save_notebook(notebook):
     from twisted.internet.error import ReactorNotRunning
     print "Quitting all running worksheets..."
@@ -84,7 +147,7 @@ def my_sigint(x, n):
 #########
 # Flask #
 #########
-import os
+import os, sys, random
 flask_dir = os.path.join(os.environ['SAGE_ROOT'], 'devel', 'sagenb', 'flask_version')
 sys.path.append(flask_dir)
 import base as flask_base
@@ -215,16 +278,16 @@ def notebook_setup(self=None):
            '--outfile %s' % (template_file, private_pem, public_pem)]
     print cmd[0]
     subprocess.call(cmd, shell=True)
-    
+
     # Set permissions on private cert
     os.chmod(private_pem, 0600)
 
     print "Successfully configured notebook."
 
-def notebook_twisted(self,
+def notebook_run(self,
              directory     = None,
              port          = 8080,
-             interface     = 'localhost',        
+             interface     = 'localhost',
              port_tries    = 50,
              secure        = False,
              reset         = False,
@@ -241,7 +304,10 @@ def notebook_twisted(self,
              start_path    = "",
              fork          = False,
              quiet         = False,
-             
+
+             server = "twistd",
+             profile = False,
+
              subnets = None,
              require_login = None,
              open_viewer = None,
@@ -285,10 +351,9 @@ def notebook_twisted(self,
     # if none use defaults
 
     nb = notebook.load_notebook(directory)
-    
+
     directory = nb._dir
-    conf = os.path.join(directory, 'twistedconf.tac')
-    
+
     if not quiet:
         print "The notebook files are stored in:", nb._dir
 
@@ -314,7 +379,7 @@ def notebook_twisted(self,
 
     if not nb.user_manager().user_exists('admin'):
         reset = True
-        
+
     if reset:
         passwd = get_admin_passwd()                
         if reset:
@@ -345,9 +410,52 @@ def notebook_twisted(self,
     nb.save()
     del nb
 
-    def run(port):
+    def run_flask():
+        """Run a flask (werkzeug) webserver."""
+        # TODO: Check to see if server is running already (PID file?)
+        conf = os.path.join(directory, 'run_flask')
+
+        notebook_opts = '"%s",interface="%s",port=%s,secure=%s' % (
+            os.path.abspath(directory), interface, port, secure)
+
+        if automatic_login:
+            start_path = "'/?startup_token=%s' % startup_token"
+            if interface:
+                hostname = interface
+            else:
+                hostname = 'localhost'
+            open_page = "from sagenb.misc.misc import open_page; open_page('%s', %s, %s, %s)" % (hostname, port, secure, start_path)
+        else:
+            open_page = ''
+
+        config = open(conf, 'w')
+
+        config.write(FLASK_NOTEBOOK_CONFIG%{'notebook_opts': notebook_opts,
+                                            'cwd':cwd,
+                                            'open_page': open_page, 'login': automatic_login,
+                                            'secure': secure, 'pkey': private_pem, 'cert': public_pem,
+                                            'host': interface, 'port': port})
+
+        config.close()
+
+        if profile:
+            import random
+            if isinstance(profile, basestring):
+                profilefile = profile+'%s.stats'%random.random()
+            else:
+                profilefile = 'sagenb-flask-profile-%s.stats'%random.random()
+            profilecmd = '-m cProfile -o %s'%profilefile
+        else:
+            profilecmd=''
+        cmd = 'python %s %s' % (profilecmd, conf)
+        return cmd
+        # end of inner function run_flask
+
+    def run_twistd():
+        """Run a twistd webserver."""
         # Is a server already running? Check if a Twistd PID exists in
         # the given directory.
+        conf = os.path.join(directory, 'twistedconf.tac')
         pidfile = os.path.join(directory, 'twistd.pid')
         if platformType != 'win32':
             from twisted.scripts._twistd_unix import checkPID
@@ -358,6 +466,7 @@ def run(port):
 
                 if str(e).startswith('Another twistd server is running,'):
                     print 'Another Sage Notebook server is running, PID %d.' % pid
+
                     old_interface, old_port, old_secure = get_old_settings(conf)
                     if automatic_login and old_port:
                         old_interface = old_interface or 'localhost'
@@ -367,20 +476,12 @@ def run(port):
 
                         from sagenb.misc.misc import open_page as browse_to
                         browse_to(old_interface, old_port, old_secure, '/')
-                        return
+                        return None
                     print '\nPlease either stop the old server or run the new server in a different directory.'
-                    return
+                    return None
 
         ## Create the config file
         if secure:
-            if (not os.path.exists(private_pem) or
-                not os.path.exists(public_pem)):
-                print "In order to use an SECURE encrypted notebook, you must first run notebook.setup()."
-                print "Now running notebook.setup()"
-                notebook_setup()
-            if (not os.path.exists(private_pem) or
-                not os.path.exists(public_pem)):
-                print "Failed to setup notebook.  Please try notebook.setup() again manually."
             strport = '%s:%s:interface=%s:privateKey=%s:certKey=%s'%(
                 protocol, port, interface, private_pem, public_pem)
         else:
@@ -401,35 +502,27 @@ def run(port):
 
         config = open(conf, 'w')
 
-        config.write(FLASK_NOTEBOOK_CONFIG%{'notebook_opts': notebook_opts,
+        config.write(TWISTD_NOTEBOOK_CONFIG%{'notebook_opts': notebook_opts,
                                             'cwd':cwd, 'strport': strport,
                                             'open_page': open_page, 'login': automatic_login})
 
 
-        config.close()                     
+        config.close()
 
         ## Start up twisted
-        cmd = 'twistd --pidfile="%s" -ny "%s"' % (pidfile, conf)
-        if not quiet:
-            print_open_msg('localhost' if not interface else interface,
-            port, secure=secure)
-        if secure and not quiet:
-            print "There is an admin account.  If you do not remember the password,"
-            print "quit the notebook and type notebook(reset=True)."
-
-        if fork:
-            import pexpect
-            return pexpect.spawn(cmd)
+        if profile:
+            import random
+            if isinstance(profile, basestring):
+                profilefile = profile+'%s.stats'%random.random()
+            else:
+                profilefile = 'sagenb-twistd-profile-%s.stats'%random.random()
+            profilecmd = '--profile=%s --profiler=cprofile --savestats'%profilefile
         else:
-            e = os.system(cmd)
-
-        os.chdir(cwd)
-        if e == 256:
-            raise socket.error
+            profilecmd=''
+        cmd = 'twistd %s --pidfile="%s" -ny "%s"' % (profilecmd, pidfile, conf)
+        return cmd
+        # end of inner function run_twistd
 
-        return True
-        # end of inner function run
-                     
     if interface != 'localhost' and not secure:
             print "*" * 70
             print "WARNING: Insecure notebook server listening on external interface."
@@ -440,7 +533,38 @@ def run(port):
     port = find_next_available_port(interface, port, port_tries)
     if automatic_login:
         "Automatic login isn't fully implemented.  You have to manually open your web browser to the above URL."
-    return run(port)
+    if secure:
+        if (not os.path.exists(private_pem) or
+            not os.path.exists(public_pem)):
+            print "In order to use an SECURE encrypted notebook, you must first run notebook.setup()."
+            print "Now running notebook.setup()"
+            notebook_setup()
+        if (not os.path.exists(private_pem) or
+            not os.path.exists(public_pem)):
+            print "Failed to setup notebook.  Please try notebook.setup() again manually."
+    if server=="flask":
+        cmd = run_flask()
+    elif server=="twistd":
+        cmd = run_twistd()
+    if cmd is None:
+        return
+    
+    if not quiet:
+        print_open_msg('localhost' if not interface else interface,
+        port, secure=secure)
+    if secure and not quiet:
+        print "There is an admin account.  If you do not remember the password,"
+        print "quit the notebook and type notebook(reset=True)."
+    print "Executing", cmd
+    if fork:
+        import pexpect
+        return pexpect.spawn(cmd)
+    else:
+        e = os.system(cmd)
+
+    os.chdir(cwd)
+    if e == 256:
+        raise socket.error
 
 def get_admin_passwd():
     print "\n" * 2
