jupyterhub_config.py

import sys

c = get_config()  # noqa

# start the service
c.JupyterHub.services = [
    {"name": "list-notebooks", "command": [sys.executable, "list_notebooks.py"]}
]

# give the service the permissions it needs
c.JupyterHub.load_roles = [
    {
        "name": "list-notebooks",
        "scopes": [
            # list users
            "list:users",
            # read server status, URLs
            "read:servers",
            # access each server's API
            "access:servers",
        ],
        "services": ["list-notebooks"],
    }
]

# for demo purposes
c.JupyterHub.allowed_users = {f"test-{i}" for i in range(10)}
c.JupyterHub.spawner_class = "simple"
c.JupyterHub.authenticator_class = "dummy"
c.JupyterHub.ip = "127.0.0.1"

list_notebooks.py

"""
List all open notebooks in a JupyterHub instance

uses JupyterHub authentication to talk to individual server APIs
to list open notebooks.
"""
import asyncio
import os

import aiohttp

# token env injected when the service is started
token = os.environ["JUPYTERHUB_API_TOKEN"]
# hub api
hub_api = os.environ["JUPYTERHUB_API_URL"]
# the proxy's base URL, not available from env
public_url = "http://127.0.0.1:8000"


async def list_users(session, offset=0):
    """List users with running & ready servers"""
    async with session.get(
        f"{hub_api}/users", params={"offset": offset, "state": "ready"}
    ) as r:
        response = await r.json()
        for user in response["items"]:
            yield user
        # handle pagination
        next_offset = response["_pagination"]["next"]
        if next_offset:
            async for user in list_users(session, offset=next_offset):
                yield user


async def list_notebooks(session, user):
    """List open notebook sessions for a given user"""
    # iterate through servers
    for server_name, server_info in user["servers"].items():
        if server_info["ready"]:
            url = server_info["url"]
            # get open sessions
            async with session.get(f"{public_url}{url}api/sessions") as r:
                sessions = await r.json()
                # iterate over sessions
                for session in sessions:
                    # if it's a notebook session (could also be a console, etc.), print the path
                    # also has other fields like kernel info, last_activity, execution_state
                    if session["type"] == "notebook":
                        print(f"{user['name']}/{server_name}: {session['path']}")


async def main():
    pid = os.getpid()
    print(f"{pid=}")
    while True:
        async with aiohttp.ClientSession() as session:
            # add jupyterhub token to requests
            session.headers["Authorization"] = f"Bearer {token}"
            # opt-in to jupyterhub user list pagination
            session.headers["Accept"] = "application/jupyterhub-pagination+json"
            notebook_futures = []
            async for user in list_users(session):
                notebook_futures.append(list_notebooks(session, user))

            # run user requests concurrently
            await asyncio.gather(*notebook_futures)

        # collect info every minute
        await asyncio.sleep(60)


if __name__ == "__main__":
    asyncio.run(main())