Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
import asyncio
import cProfile
import os
import pstats
import sys
from threading import Thread
import time
from jupyterhub.app import JupyterHub
from jupyterhub import orm
from jupyterhub.spawner import Spawner
from traitlets import Integer
from traitlets.config import Config
class DummySpawner(Spawner):
_stopped = True
def get_state(self):
if not self._stopped:
return {'stopped': False}
def load_state(self, state):
self._stopped = state.get('stopped', True)
async def start(self):
self._stopped = False
return 'http://127.0.0.1:9999'
async def stop(self):
self._stopped = True
async def poll(self):
if self._stopped:
return 0
else:
return None
cfg = Config()
cfg.JupyterHub.authenticator_class = 'github'
cfg.JupyterHub.spawner_class = DummySpawner
if len(sys.argv) >= 2:
n = int(sys.argv[1])
else:
n = 1000
cfg.JupyterHub.db_url = f'sqlite:///jupyterhub-user-test-{n}.db'
class GenerateUsers(JupyterHub):
num_users = Integer(100)
running_users = Integer(10)
async def generate_users(self):
offset = self.db.query(orm.User).count()
print(f"already have {offset} users")
for i in range(offset, self.num_users):
name = 'user-%i' % i
user = orm.User(name=name)
self.db.add(user)
self.users.add(user)
if offset > self.num_users:
too_many = offset - self.num_users
print(f"deleting {too_many} users")
to_delete = self.db.query(orm.User).limit(too_many)
for user in to_delete:
self.db.delete(user)
self.db.commit()
for orm_user in self.db.query(orm.User).limit(self.running_users):
user = self.users[orm_user]
spawner = user.spawner
status = await spawner.poll()
if status is not None:
await user.spawn()
async def init_spawners(self):
return 0
class HubStopper(JupyterHub):
async def start(self):
return
async def generate_users(n):
app = GenerateUsers(num_users=n, config=cfg)
await app.initialize()
await app.generate_users()
app.clear_instance()
async def profile_startup():
hub = HubStopper.instance(config=cfg)
await hub.initialize()
def listener(port):
import asyncio
asyncio.set_event_loop(asyncio.new_event_loop())
import tornado.web
import tornado.httpserver
import tornado.ioloop
class EchoHandler(tornado.web.RequestHandler):
def get(self):
self.write(self.request.path)
application = tornado.web.Application([('.*', EchoHandler)])
http_server = tornado.httpserver.HTTPServer(application)
http_server.listen(port)
tornado.ioloop.IOLoop.current().start()
if __name__ == '__main__':
t = Thread(target=listener, args=(9999,), daemon=True)
t.start()
tic = time.perf_counter()
asyncio.run(generate_users(n))
toc = time.perf_counter()
print(f"Created {n} users in {toc-tic:.3f} seconds")
p = cProfile.Profile()
tic = time.perf_counter()
p.runcall(asyncio.run, profile_startup())
toc = time.perf_counter()
print(f"Loaded {n} users in {toc-tic:.3f} seconds")
stats = pstats.Stats(p)
stats.sort_stats('cumtime')
stats.print_stats(32)
p.dump_stats(f'run-{n}.prof')
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.