hazza/jupyterhub
1
0
Fork 0
mirror of https://github.com/jupyterhub/jupyterhub.git synced 2025-10-16 14:33:00 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #48 from minrk/resume-from-db

Browse Source 
enable resuming Hub state from database
This commit is contained in:
Min RK
2014-09-21 22:57:17 -07:00
parent cf64828d32 f120ae6154
commit ea04bdb30c
9 changed files with 269 additions and 76 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
jupyterhub/apihandlers/users.py
View File

@@ -133,9 +133,11 @@ class UserServerAPIHandler(BaseUserHandler):
def post(self, name):
user = self.find_user(name)
if user.spawner:
raise web.HTTPError(400, "%s's server is already running" % name)
else:
yield self.spawn_single_user(user)
state = yield user.spawner.poll()
if state is None:
raise web.HTTPError(400, "%s's server is already running" % name)
yield self.spawn_single_user(user)
self.set_status(201)
@gen.coroutine

252
jupyterhub/app.py
View File

@@ -25,11 +25,10 @@ from tornado.log import LogFormatter
from tornado import gen, web
from IPython.utils.traitlets import (
Unicode, Integer, Dict, TraitError, List, Bool, Bytes, Any,
DottedObjectName, Set,
Unicode, Integer, Dict, TraitError, List, Bool, Any,
Type, Set, Instance,
)
from IPython.config import Application, catch_config_error
from IPython.utils.importstring import import_item
here = os.path.dirname(__file__)
@@ -37,7 +36,7 @@ from . import handlers, apihandlers
from . import orm
from ._data import DATA_FILES_PATH
from .utils import url_path_join
from .utils import url_path_join, random_hex, TimeoutError
# classes for config
from .auth import Authenticator, PAMAuthenticator
@@ -57,10 +56,13 @@ aliases = {
}
flags = {
'debug': ({'Application' : {'log_level' : logging.DEBUG}},
'debug': ({'Application' : {'log_level': logging.DEBUG}},
"set log level to logging.DEBUG (maximize logging output)"),
'generate-config': ({'JupyterHubApp': {'generate_config' : True}},
"generate default config file")
'generate-config': ({'JupyterHubApp': {'generate_config': True}},
"generate default config file"),
'no-db': ({'JupyterHubApp': {'db_url': 'sqlite:///:memory:'}},
"disable persisting state database to disk"
),
}
@@ -150,10 +152,13 @@ class JupyterHubApp(Application):
"""
)
proxy_auth_token = Unicode(config=True,
help="The Proxy Auth token"
help="""The Proxy Auth token.
Loaded from the CONFIGPROXY_AUTH_TOKEN env variable by default.
"""
)
def _proxy_auth_token_default(self):
return orm.new_token()
return os.environ.get('CONFIGPROXY_AUTH_TOKEN', orm.new_token())
proxy_api_ip = Unicode('localhost', config=True,
help="The ip for the proxy API handlers"
@@ -190,11 +195,16 @@ class JupyterHubApp(Application):
if newnew != new:
self.hub_prefix = newnew
cookie_secret = Bytes(config=True)
cookie_secret = Unicode(config=True,
help="""The cookie secret to use to encrypt cookies.
Loaded from the JPY_COOKIE_SECRET env variable by default.
"""
)
def _cookie_secret_default(self):
return b'secret!'
return os.environ.get('JPY_COOKIE_SECRET', random_hex(64))
authenticator = DottedObjectName("jupyterhub.auth.PAMAuthenticator", config=True,
authenticator_class = Type("jupyterhub.auth.PAMAuthenticator", config=True,
help="""Class for authenticating users.
This should be a class with the following form:
@@ -208,16 +218,38 @@ class JupyterHubApp(Application):
and `data` is the POST form data from the login page.
"""
)
authenticator = Instance(Authenticator)
def _authenticator_default(self):
return self.authenticator_class(config=self.config)
# class for spawning single-user servers
spawner_class = DottedObjectName("jupyterhub.spawner.LocalProcessSpawner", config=True,
spawner_class = Type("jupyterhub.spawner.LocalProcessSpawner", config=True,
help="""The class to use for spawning single-user servers.
Should be a subclass of Spawner.
"""
)
db_url = Unicode('sqlite:///:memory:', config=True)
debug_db = Bool(False)
db_url = Unicode('sqlite:///jupyterhub.sqlite', config=True,
help="url for the database. e.g. `sqlite:///jupyterhub.sqlite`"
)
def _db_url_changed(self, name, old, new):
if '://' not in new:
# assume sqlite, if given as a plain filename
self.db_url = 'sqlite:///%s' % new
db_kwargs = Dict(config=True,
help="""Include any kwargs to pass to the database connection.
See sqlalchemy.create_engine for details.
"""
)
reset_db = Bool(False, config=True,
help="Purge and reset the database."
)
debug_db = Bool(False, config=True,
help="log all database transactions. This has A LOT of output"
)
db = Any()
admin_users = Set({getpass.getuser()}, config=True,
@@ -288,7 +320,6 @@ class JupyterHubApp(Application):
self.handlers = self.add_url_prefix(self.hub_prefix, h)
# some extra handlers, outside hub_prefix
self.handlers.extend([
(r"%s" % self.hub_prefix.rstrip('/'), web.RedirectHandler,
@@ -302,48 +333,139 @@ class JupyterHubApp(Application):
])
def init_db(self):
# TODO: load state from db for resume
# TODO: if not resuming, clear existing db contents
self.db = orm.new_session(self.db_url, echo=self.debug_db)
for name in self.admin_users:
user = orm.User(name=name, admin=True)
self.db.add(user)
self.db.commit()
"""Create the database connection"""
self.db = orm.new_session(self.db_url, reset=self.reset_db, echo=self.debug_db,
**self.db_kwargs
)
def init_hub(self):
"""Load the Hub config into the database"""
self.hub = orm.Hub(
server=orm.Server(
ip=self.hub_ip,
port=self.hub_port,
base_url=self.hub_prefix,
cookie_secret=self.cookie_secret,
cookie_name='jupyter-hub-token',
self.hub = self.db.query(orm.Hub).first()
if self.hub is None:
self.hub = orm.Hub(
server=orm.Server(
ip=self.hub_ip,
port=self.hub_port,
base_url=self.hub_prefix,
cookie_secret=self.cookie_secret,
cookie_name='jupyter-hub-token',
)
)
)
self.db.add(self.hub)
self.db.add(self.hub)
else:
server = self.hub.server
server.ip = self.hub_ip
server.port = self.hub_port
server.base_url = self.hub_prefix
self.db.commit()
def init_users(self):
"""Load users into and from the database"""
db = self.db
for name in self.admin_users:
# ensure anyone specified as admin in config is admin in db
user = orm.User.find(db, name)
if user is None:
user = orm.User(name=name, admin=True)
db.add(user)
else:
user.admin = True
# the admin_users config variable will never be used after this point.
# only the database values will be referenced.
whitelist = self.authenticator.whitelist
if not whitelist:
self.log.info("Not using whitelist. Any authenticated user will be allowed.")
# add whitelisted users to the db
for name in whitelist:
user = orm.User.find(db, name)
if user is None:
user = orm.User(name=name)
db.add(user)
if whitelist:
# fill the whitelist with any users loaded from the db,
# so we are consistent in both directions.
# This lets whitelist be used to set up initial list,
# but changes to the whitelist can occur in the database,
# and persist across sessions.
for user in db.query(orm.User):
whitelist.add(user.name)
# The whitelist set and the users in the db are now the same.
# From this point on, any user changes should be done simultaneously
# to the whitelist set and user db, unless the whitelist is empty (all users allowed).
db.commit()
# load any still-active spawners from JSON
run_sync = IOLoop().run_sync
user_summaries = ['']
def _user_summary(user):
parts = ['{0: >8}'.format(user.name)]
if user.admin:
parts.append('admin')
if user.server:
parts.append('running at %s' % user.server)
return ' '.join(parts)
for user in db.query(orm.User):
if not user.state:
user_summaries.append(_user_summary(user))
continue
self.log.debug("Loading state for %s from db", user.name)
spawner = self.spawner_class.fromJSON(user.state, user=user, hub=self.hub, config=self.config)
status = run_sync(spawner.poll)
if status is None:
self.log.info("User %s still running", user.name)
user.spawner = spawner
else:
self.log.warn("Failed to load state for %s, assuming server is not running.", user.name)
# not running, state is invalid
user.state = {}
user.server = None
user_summaries.append(_user_summary(user))
self.log.debug("Loaded users: %s", '\n'.join(user_summaries))
db.commit()
def init_proxy(self):
"""Load the Proxy config into the database"""
self.proxy = orm.Proxy(
public_server=orm.Server(
ip=self.ip,
port=self.port,
),
api_server=orm.Server(
ip=self.proxy_api_ip,
port=self.proxy_api_port,
base_url='/api/routes/'
),
auth_token = orm.new_token(),
)
self.db.add(self.proxy)
self.proxy = self.db.query(orm.Proxy).first()
if self.proxy is None:
self.proxy = orm.Proxy(
public_server=orm.Server(),
api_server=orm.Server(),
auth_token = self.proxy_auth_token,
)
self.db.add(self.proxy)
self.db.commit()
self.proxy.log = self.log
self.proxy.public_server.ip = self.ip
self.proxy.public_server.port = self.port
self.proxy.api_server.ip = self.proxy_api_ip
self.proxy.api_server.port = self.proxy_api_port
self.proxy.api_server.base_url = '/api/routes/'
if self.proxy.auth_token is None:
self.proxy.auth_token = self.proxy_auth_token
self.db.commit()
@gen.coroutine
def start_proxy(self):
"""Actually start the configurable-http-proxy"""
if self.proxy.public_server.is_up() and \
self.proxy.api_server.is_up():
self.log.warn("Proxy already running at: %s", self.proxy.public_server.url)
self.proxy_process = None
return
env = os.environ.copy()
env['CONFIGPROXY_AUTH_TOKEN'] = self.proxy.auth_token
cmd = [self.proxy_cmd,
@@ -385,7 +507,9 @@ class JupyterHubApp(Application):
def check_proxy(self):
if self.proxy_process.poll() is None:
return
self.log.error("Proxy stopped with exit code %i", self.proxy_process.poll())
self.log.error("Proxy stopped with exit code %r",
'unknown' if self.proxy_process is None else self.proxy_process.poll()
)
yield self.start_proxy()
self.log.info("Setting up routes on new proxy")
yield self.proxy.add_all_users()
@@ -407,10 +531,10 @@ class JupyterHubApp(Application):
proxy=self.proxy,
hub=self.hub,
admin_users=self.admin_users,
authenticator=import_item(self.authenticator)(config=self.config),
spawner_class=import_item(self.spawner_class),
authenticator=self.authenticator,
spawner_class=self.spawner_class,
base_url=base_url,
cookie_secret=self.cookie_secret,
cookie_secret=self.hub.server.cookie_secret,
login_url=url_path_join(self.hub.server.base_url, 'login'),
static_path=os.path.join(self.data_files_path, 'static'),
static_url_prefix=url_path_join(self.hub.server.base_url, 'static/'),
@@ -438,12 +562,13 @@ class JupyterHubApp(Application):
if self.generate_config:
return
self.load_config_file(self.config_file)
self.write_pid_file()
self.init_logging()
self.write_pid_file()
self.init_ports()
self.init_db()
self.init_hub()
self.init_proxy()
self.init_users()
self.init_handlers()
self.init_tornado_settings()
self.init_tornado_application()
@@ -456,15 +581,24 @@ class JupyterHubApp(Application):
futures = []
for user in self.db.query(orm.User):
if user.spawner is not None:
futures.append(user.spawner.stop())
futures.append(user.stop())
# clean up proxy while SUS are shutting down
self.log.info("Cleaning up proxy[%i]..." % self.proxy_process.pid)
self.proxy_process.terminate()
if self.proxy_process and self.proxy_process.poll() is None:
self.log.info("Cleaning up proxy[%i]...", self.proxy_process.pid)
try:
self.proxy_process.terminate()
except Exception as e:
self.log.error("Failed to terminate proxy process: %s", e)
# wait for the requests to stop finish:
for f in futures:
yield f
try:
yield f
except Exception as e:
self.log.error("Failed to stop user: %s", e)
self.db.commit()
if self.pid_file and os.path.exists(self.pid_file):
self.log.info("Cleaning up PID file %s", self.pid_file)
@@ -474,6 +608,7 @@ class JupyterHubApp(Application):
self.log.info("...done")
def write_config_file(self):
"""Write our default config to a .py config file"""
if os.path.exists(self.config_file) and not self.answer_yes:
answer = ''
def ask():
@@ -509,9 +644,14 @@ class JupyterHubApp(Application):
return
loop = IOLoop.current()
loop.add_callback(self.proxy.add_all_users)
pc = PeriodicCallback(self.check_proxy, self.proxy_check_interval)
pc.start()
if self.proxy_process:
# only check / restart the proxy if we started it in the first place.
# this means a restarted Hub cannot restart a Proxy that its
# predecessor started.
pc = PeriodicCallback(self.check_proxy, self.proxy_check_interval)
pc.start()
# start the webserver
http_server = tornado.httpserver.HTTPServer(self.tornado_application)

7
jupyterhub/handlers/base.py
View File

@@ -73,7 +73,7 @@ class BaseHandler(RequestHandler):
if not match:
return None
token = match.group(1)
orm_token = self.db.query(orm.APIToken).filter(orm.APIToken.token == token).first()
orm_token = orm.APIToken.find(self.db, token)
if orm_token is None:
return None
else:
@@ -83,8 +83,7 @@ class BaseHandler(RequestHandler):
"""get_current_user from a cookie token"""
token = self.get_cookie(self.hub.server.cookie_name, None)
if token:
cookie_token = self.db.query(orm.CookieToken).filter(
orm.CookieToken.token==token).first()
cookie_token = orm.CookieToken.find(self.db, token)
if cookie_token:
return cookie_token.user
else:
@@ -103,7 +102,7 @@ class BaseHandler(RequestHandler):
return None if no such user
"""
return self.db.query(orm.User).filter(orm.User.name==name).first()
return orm.User.find(self.db, name)
def user_from_username(self, username):
"""Get ORM User for username"""

7
jupyterhub/handlers/login.py
View File

@@ -51,7 +51,12 @@ class LoginHandler(BaseHandler):
authorized = yield self.authenticate(data)
if authorized:
user = self.user_from_username(username)
yield self.spawn_single_user(user)
already_running = False
if user.spawner:
status = yield user.spawner.poll()
already_running = (status == None)
if not already_running:
yield self.spawn_single_user(user)
self.set_login_cookie(user)
next_url = self.get_argument('next', default='') or self.hub.server.base_url
self.redirect(next_url)

31
jupyterhub/orm.py
View File

@@ -9,6 +9,7 @@ import socket
import uuid
from tornado import gen
from tornado.log import app_log
from tornado.httpclient import HTTPRequest, AsyncHTTPClient, HTTPError
from sqlalchemy.types import TypeDecorator, VARCHAR
@@ -71,7 +72,7 @@ class Server(Base):
ip = Column(Unicode, default=u'localhost')
port = Column(Integer, default=random_port)
base_url = Column(Unicode, default=u'/')
cookie_secret = Column(Binary, default=b'secret')
cookie_secret = Column(Unicode, default=u'')
cookie_name = Column(Unicode, default=u'cookie')
def __repr__(self):
@@ -103,12 +104,11 @@ class Server(Base):
socket.create_connection((self.ip or 'localhost', self.port))
except socket.error as e:
if e.errno == errno.ECONNREFUSED:
return True
return False
else:
raise
else:
return True
class Proxy(Base):
@@ -124,6 +124,7 @@ class Proxy(Base):
public_server = relationship(Server, primaryjoin=_public_server_id == Server.id)
_api_server_id = Column(Integer, ForeignKey('servers.id'))
api_server = relationship(Server, primaryjoin=_api_server_id == Server.id)
log = app_log
def __repr__(self):
if self.public_server:
@@ -136,6 +137,9 @@ class Proxy(Base):
@gen.coroutine
def add_user(self, user, client=None):
"""Add a user's server to the proxy table."""
self.log.info("Adding user %s to proxy %s => %s",
user.name, user.server.base_url, user.server.host,
)
client = client or AsyncHTTPClient()
req = HTTPRequest(url_path_join(
@@ -155,6 +159,7 @@ class Proxy(Base):
@gen.coroutine
def delete_user(self, user, client=None):
"""Remove a user's server to the proxy table."""
self.log.info("Removing user %s from proxy", user.name)
client = client or AsyncHTTPClient()
req = HTTPRequest(url_path_join(
self.api_server.url,
@@ -262,6 +267,14 @@ class User(Base):
"""Return a new cookie token"""
return self._new_token(CookieToken)
@classmethod
def find(cls, db, name):
"""Find a user by name.
Returns None if not found.
"""
return db.query(cls).filter(cls.name==name).first()
@gen.coroutine
def spawn(self, spawner_class, base_url='/', hub=None, config=None):
db = inspect(self).session
@@ -321,6 +334,14 @@ class Token(object):
u=self.user.name,
)
@classmethod
def find(cls, db, token):
"""Find a token object by value.
Returns None if not found.
"""
return db.query(cls).filter(cls.token==token).first()
class APIToken(Token, Base):
"""An API token"""
@@ -332,13 +353,15 @@ class CookieToken(Token, Base):
__tablename__ = 'cookie_tokens'
def new_session(url="sqlite:///:memory:", **kwargs):
def new_session(url="sqlite:///:memory:", reset=False, **kwargs):
"""Create a new session at url"""
kwargs.setdefault('connect_args', {'check_same_thread': False})
kwargs.setdefault('poolclass', StaticPool)
engine = create_engine(url, **kwargs)
Session = sessionmaker(bind=engine)
session = Session()
if reset:
Base.metadata.drop_all(engine)
Base.metadata.create_all(engine)
return session

9
jupyterhub/spawner.py
View File

@@ -60,7 +60,7 @@ class Spawner(LoggingConfigurable):
env = os.environ.copy()
for key in ['HOME', 'USER', 'USERNAME', 'LOGNAME', 'LNAME']:
env.pop(key, None)
self._env_key(env, 'COOKIE_SECRET', self.user.server.cookie_secret.decode('ascii'))
self._env_key(env, 'COOKIE_SECRET', self.user.server.cookie_secret)
self._env_key(env, 'API_TOKEN', self.api_token)
return env
@@ -103,7 +103,7 @@ class Spawner(LoggingConfigurable):
state: dict
a JSONable dict of state
"""
return {}
return dict(api_token=self.api_token)
def get_args(self):
"""Return the arguments to be passed after self.cmd"""
@@ -211,10 +211,13 @@ class LocalProcessSpawner(Spawner):
raise ValueError("This should be impossible")
def load_state(self, state):
super(LocalProcessSpawner, self).load_state(state)
self.pid = state['pid']
def get_state(self):
return dict(pid=self.pid)
state = super(LocalProcessSpawner, self).get_state()
state['pid'] = self.pid
return state
def sudo_cmd(self, user):
return ['sudo', '-u', user.name] + self.sudo_args

9
jupyterhub/tests/mocking.py
View File

@@ -56,11 +56,14 @@ class MockHubApp(JupyterHubApp):
def _ip_default(self):
return 'localhost'
def _authenticator_default(self):
return '%s.%s' % (__name__, 'MockPAMAuthenticator')
def _db_url_default(self):
return 'sqlite:///:memory:'
def _authenticator_class_default(self):
return MockPAMAuthenticator
def _spawner_class_default(self):
return '%s.%s' % (__name__, 'MockSpawner')
return MockSpawner
def _admin_users_default(self):
return {'admin'}

11
jupyterhub/tests/test_orm.py
View File

@@ -21,7 +21,7 @@ def test_server(db):
assert server.proto == 'http'
assert isinstance(server.port, int)
assert isinstance(server.cookie_name, unicode)
assert isinstance(server.cookie_secret, bytes)
assert isinstance(server.cookie_secret, unicode)
assert server.url == 'http://localhost:%i/' % server.port
@@ -71,6 +71,11 @@ def test_user(db):
assert user.server.ip == u'localhost'
assert user.state == {'pid': 4234}
found = orm.User.find(db, u'kaylee')
assert found.name == user.name
found = orm.User.find(db, u'badger')
assert found is None
def test_tokens(db):
user = orm.User(name=u'inara')
@@ -87,3 +92,7 @@ def test_tokens(db):
assert len(user.api_tokens) == 1
assert len(user.cookie_tokens) == 3
found = orm.CookieToken.find(db, token=token.token)
assert found.token == token.token
found = orm.APIToken.find(db, token.token)
assert found is None

11
jupyterhub/utils.py
View File

@@ -3,7 +3,9 @@
# Copyright (c) Jupyter Development Team.
# Distributed under the terms of the Modified BSD License.
import binascii
import errno
import os
import socket
from tornado import web, gen, ioloop
from tornado.log import app_log
@@ -11,7 +13,8 @@ from tornado.log import app_log
from IPython.html.utils import url_path_join
try:
TimeoutError
# make TimeoutError importable on Python >= 3.3
TimeoutError = TimeoutError
except NameError:
# python < 3.3
class TimeoutError(Exception):
@@ -25,6 +28,12 @@ def random_port():
sock.close()
return port
def random_hex(nbytes):
"""Return nbytes random bytes as a unicode hex string
It will have length nbytes * 2
"""
return binascii.hexlify(os.urandom(nbytes)).decode('ascii')
@gen.coroutine
def wait_for_server(ip, port, timeout=10):