From 9a77d1df995d41aa41233aa4e50c95b4144c7269 Mon Sep 17 00:00:00 2001 From: Chico Venancio <fvenancio@wikimedia.org> Date: Mon, 12 Mar 2018 22:03:14 +0000 Subject: [PATCH] test cull_idle_servers url encoding --- images/paws-hub/Dockerfile | 2 + images/paws-hub/cull_idle_servers.py | 151 +++++++++++++++++++++++++++ 2 files changed, 153 insertions(+) create mode 100755 images/paws-hub/cull_idle_servers.py diff --git a/images/paws-hub/Dockerfile b/images/paws-hub/Dockerfile index 99969caf..2b57a50e 100644 --- a/images/paws-hub/Dockerfile +++ b/images/paws-hub/Dockerfile @@ -9,8 +9,10 @@ RUN adduser --disabled-password \ --home ${HOME} \ --force-badname \ ${NB_USER} +COPY cull_idle_servers.py /usr/local/bin/cull_idle_servers.py RUN chown ${NB_USER}:${NB_USER} /srv/jupyterhub USER ${NB_USER} + CMD ["jupyterhub", "--config", "/srv/jupyterhub_config.py"] diff --git a/images/paws-hub/cull_idle_servers.py b/images/paws-hub/cull_idle_servers.py new file mode 100755 index 00000000..db0e4658 --- /dev/null +++ b/images/paws-hub/cull_idle_servers.py @@ -0,0 +1,151 @@ +#!/usr/bin/env python3 +# Imported from https://github.com/jupyterhub/jupyterhub/blob/0.8.0rc1/examples/cull-idle/cull_idle_servers.py +"""script to monitor and cull idle single-user servers + +Caveats: + +last_activity is not updated with high frequency, +so cull timeout should be greater than the sum of: + +- single-user websocket ping interval (default: 30s) +- JupyterHub.last_activity_interval (default: 5 minutes) + +You can run this as a service managed by JupyterHub with this in your config:: + + + c.JupyterHub.services = [ + { + 'name': 'cull-idle', + 'admin': True, + 'command': 'python cull_idle_servers.py --timeout=3600'.split(), + } + ] + +Or run it manually by generating an API token and storing it in `JUPYTERHUB_API_TOKEN`: + + export JUPYTERHUB_API_TOKEN=`jupyterhub token` + python cull_idle_servers.py [--timeout=900] [--url=http://127.0.0.1:8081/hub/api] +""" + +import datetime +import json +import os + +try: + from urllib.parse import quote +except ImportError: + from urllib import quote + +from dateutil.parser import parse as parse_date + +from tornado.gen import coroutine +from tornado.log import app_log +from tornado.httpclient import AsyncHTTPClient, HTTPRequest +from tornado.ioloop import IOLoop, PeriodicCallback +from tornado.options import define, options, parse_command_line + + +@coroutine +def cull_idle(url, api_token, timeout, cull_users=False): + """Shutdown idle single-user servers + + If cull_users, inactive *users* will be deleted as well. + """ + auth_header = { + 'Authorization': 'token %s' % api_token + } + req = HTTPRequest(url=url + '/users', + headers=auth_header, + ) + now = datetime.datetime.utcnow() + cull_limit = now - datetime.timedelta(seconds=timeout) + client = AsyncHTTPClient() + resp = yield client.fetch(req) + users = json.loads(resp.body.decode('utf8', 'replace')) + futures = [] + + @coroutine + def cull_one(user, last_activity): + """cull one user""" + + # shutdown server first. Hub doesn't allow deleting users with running servers. + if user['server']: + app_log.info("Culling server for %s (inactive since %s)", user['name'], last_activity) + req = HTTPRequest(url=url + '/users/%s/server' % quote(user['name']), + method='DELETE', + headers=auth_header, + ) + resp = yield client.fetch(req) + if resp.code == 202: + msg = "Server for {} is slow to stop.".format(user['name']) + if cull_users: + app_log.warning(msg + " Not culling user yet.") + # return here so we don't continue to cull the user + # which will fail if the server is still trying to shutdown + return + app_log.warning(msg) + if cull_users: + app_log.info("Culling user %s (inactive since %s)", user['name'], last_activity) + req = HTTPRequest(url=url + '/users/%s' % user['name'], + method='DELETE', + headers=auth_header, + ) + yield client.fetch(req) + + for user in users: + if not user['server'] and not cull_users: + # server not running and not culling users, nothing to do + continue + if not user['last_activity']: + continue + last_activity = parse_date(user['last_activity']) + if last_activity < cull_limit: + # user might be in a transition (e.g. starting or stopping) + # don't try to cull if this is happening + if user['pending']: + app_log.warning("Not culling user %s with pending %s", user['name'], user['pending']) + continue + futures.append((user['name'], cull_one(user, last_activity))) + else: + app_log.debug("Not culling %s (active since %s)", user['name'], last_activity) + + for (name, f) in futures: + try: + yield f + except Exception: + app_log.exception("Error culling %s", name) + else: + app_log.debug("Finished culling %s", name) + + +if __name__ == '__main__': + define('url', default=os.environ.get('JUPYTERHUB_API_URL'), help="The JupyterHub API URL") + define('timeout', default=600, help="The idle timeout (in seconds)") + define('cull_every', default=0, help="The interval (in seconds) for checking for idle servers to cull") + define('cull_users', default=False, + help="""Cull users in addition to servers. + This is for use in temporary-user cases such as tmpnb.""", + ) + + parse_command_line() + if not options.cull_every: + options.cull_every = options.timeout // 2 + api_token = os.environ['JUPYTERHUB_API_TOKEN'] + + try: + AsyncHTTPClient.configure("tornado.curl_httpclient.CurlAsyncHTTPClient") + except ImportError as e: + app_log.warning("Could not load pycurl: %s\npycurl is recommended if you have a large number of users.", e) + + loop = IOLoop.current() + cull = lambda : cull_idle(options.url, api_token, options.timeout, options.cull_users) + # schedule first cull immediately + # because PeriodicCallback doesn't start until the end of the first interval + loop.add_callback(cull) + # schedule periodic cull + pc = PeriodicCallback(cull, 1e3 * options.cull_every) + pc.start() + try: + loop.start() + except KeyboardInterrupt: + pass