looking-up-notebooks.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "cdd153ec-5baa-4fda-8b0c-2b2aed5a32ba",
   "metadata": {},
   "source": [
    "# Gett\n",
    "\n",
    "\n",
    "## Step 1: list kernels and sessions\n",
    "\n",
    "In the Jupyter server API, a `session` is a relationship between a document (or console) and a kernel.\n",
    "In practice, it is typical for each kernel to be associated with one document.\n",
    "\n",
    "**Caveats**\n",
    "\n",
    "- Kernels do not _need_ to be associated with a document,\n",
    "  so listing `kernels` directly may discover more kernels than the sessions API,\n",
    "  but this is rare.\n",
    "- Technically, one kernel can be associated with multiple documents,\n",
    "  but this is also exceedingly rare and not super relevant."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "bce61d80-c9d1-42cd-ab8d-636e1c3a97bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "import requests\n",
    "\n",
    "token = \"abc123\"  # get this somewhere, e.g. $JUPYTERHUB_API_TOKEN or via a JupyterHub service token\n",
    "\n",
    "s = requests.Session()\n",
    "s.headers = {\"Authorization\": \"token abc123\"}\n",
    "server_url = \"http://127.0.0.1:8888\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "f0fe7824-1a52-407e-a106-1101f9744059",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'id': '0aa93fcc-0001-4de8-bb31-9aa5c9e00ffe',\n",
       "  'path': 'notebook-a.ipynb',\n",
       "  'name': 'notebook-a.ipynb',\n",
       "  'type': 'notebook',\n",
       "  'kernel': {'id': '2efac3d3-4adb-4193-9d52-a68940660347',\n",
       "   'name': 'python3',\n",
       "   'last_activity': '2022-08-03T12:15:36.181996Z',\n",
       "   'execution_state': 'busy',\n",
       "   'connections': 1},\n",
       "  'notebook': {'path': 'notebook-a.ipynb', 'name': 'notebook-a.ipynb'}},\n",
       " {'id': 'ca73e455-7593-4633-aebd-b080c9bbcde1',\n",
       "  'path': 'other notebook.ipynb',\n",
       "  'name': 'other notebook.ipynb',\n",
       "  'type': 'notebook',\n",
       "  'kernel': {'id': '1482c156-1abc-44fd-8c35-74805a3c8e98',\n",
       "   'name': 'python3',\n",
       "   'last_activity': '2022-08-03T12:15:30.509980Z',\n",
       "   'execution_state': 'idle',\n",
       "   'connections': 2},\n",
       "  'notebook': {'path': 'other notebook.ipynb',\n",
       "   'name': 'other notebook.ipynb'}},\n",
       " {'id': '237f6734-4a16-46b8-ae82-bc65e678d96d',\n",
       "  'path': 'console-1-49ab0dc9-a394-4a25-a4e3-f34cc22a5c6b',\n",
       "  'name': 'Console 1',\n",
       "  'type': 'console',\n",
       "  'kernel': {'id': '0518a5da-38d3-4dcd-8e1c-ddb7fbc494b3',\n",
       "   'name': 'python3',\n",
       "   'last_activity': '2022-08-03T12:15:32.717393Z',\n",
       "   'execution_state': 'idle',\n",
       "   'connections': 3}}]"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "session_list = s.get(f\"{server_url}/api/sessions\").json()\n",
    "session_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "249ec59e-9e1a-46d4-be3a-0f048acecf48",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'id': '2efac3d3-4adb-4193-9d52-a68940660347',\n",
       "  'name': 'python3',\n",
       "  'last_activity': '2022-08-03T12:15:36.981129Z',\n",
       "  'execution_state': 'busy',\n",
       "  'connections': 1},\n",
       " {'id': '1482c156-1abc-44fd-8c35-74805a3c8e98',\n",
       "  'name': 'python3',\n",
       "  'last_activity': '2022-08-03T12:15:30.509980Z',\n",
       "  'execution_state': 'idle',\n",
       "  'connections': 2},\n",
       " {'id': '0518a5da-38d3-4dcd-8e1c-ddb7fbc494b3',\n",
       "  'name': 'python3',\n",
       "  'last_activity': '2022-08-03T12:15:32.717393Z',\n",
       "  'execution_state': 'idle',\n",
       "  'connections': 3}]"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "kernel_list = s.get(f\"{server_url}/api/kernels\").json()\n",
    "kernel_list"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "be13e5f9-b43e-40a2-b505-abae00cbb048",
   "metadata": {},
   "source": [
    "## Step 2: look up kernel PIDs\n",
    "\n",
    "Jupyter doesn't include PIDs in any part of the API because kernels need not be local to the Jupyter server\n",
    "(e.g. when they are run via a gateway server).\n",
    "However, they are often enough that the IPython kernel has a prototype message handler called a `usage_request`.\n",
    "We can use this to ask for information, including the PID of the process.\n",
    "\n",
    "**Note:** this step only works for recent versions of `ipykernel`. This is a nonstandard message, and not supported by other kernels.\n",
    "An alternative would be to use an `execute_request` to run `os.getpid()`. This, in turn, is Python-specific and would need to be implemented.\n",
    "\n",
    "\n",
    "To do this, we connect a websocket to the Jupyter server to send messages directly to the kernel:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "aaae6bcc-6eb5-474b-8026-9a9a3e949e9d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "from tornado.websocket import websocket_connect\n",
    "from jupyter_client.session import Session, json_packer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "804ddcc9-5301-4a3a-bc8d-2facfa782a65",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'hostname': 'heavy',\n",
       " 'pid': 94479,\n",
       " 'kernel_cpu': 1.1,\n",
       " 'kernel_memory': 112164864,\n",
       " 'host_cpu_percent': 8.3,\n",
       " 'cpu_count': 10,\n",
       " 'host_virtual_memory': {'total': 17179869184,\n",
       "  'available': 4527505408,\n",
       "  'percent': 73.6,\n",
       "  'used': 7586217984,\n",
       "  'free': 165265408,\n",
       "  'active': 4448583680,\n",
       "  'inactive': 4358438912,\n",
       "  'wired': 3137634304}}"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "client_session = Session()\n",
    "\n",
    "kernel_id = session_list[0][\"kernel\"][\"id\"]\n",
    "\n",
    "ws = await websocket_connect(f\"ws{server_url[4:]}/api/kernels/{kernel_id}/channels?token={token}\")\n",
    "msg = client_session.msg(\"usage_request\", content={})\n",
    "msg[\"channel\"] = \"control\"\n",
    "json_msg = json_packer(msg)\n",
    "await ws.write_message(json_msg)\n",
    "while True:\n",
    "    reply = json.loads(await ws.read_message())\n",
    "    if reply[\"channel\"] == \"control\":\n",
    "        break\n",
    "reply[\"content\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "42f9fd7a-04d7-410b-b1db-f069155eb8f3",
   "metadata": {},
   "source": [
    "We can use this information to produce a map of PID to notebook document / console"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "604df843-cdf5-4584-9ffc-725f2136c003",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "94479: notebook 'notebook-a.ipynb'\n",
      "94480: notebook 'other notebook.ipynb'\n",
      "94481: console 'Console 1'\n"
     ]
    }
   ],
   "source": [
    "async def get_kernel_pid(kernel_id):\n",
    "    client_session = Session()\n",
    "    ws = await websocket_connect(f\"ws{server_url[4:]}/api/kernels/{kernel_id}/channels?token={token}\")\n",
    "    msg = client_session.msg(\"usage_request\", content={})\n",
    "    msg[\"channel\"] = \"control\"\n",
    "    json_msg = json_packer(msg)\n",
    "    await ws.write_message(json_msg)\n",
    "    while True:\n",
    "        reply = json.loads(await ws.read_message())\n",
    "        if reply[\"channel\"] == \"control\":\n",
    "            break\n",
    "    ws.close()\n",
    "    return reply[\"content\"][\"pid\"]\n",
    "\n",
    "for session in session_list:\n",
    "    pid = await get_kernel_pid(session[\"kernel\"][\"id\"])\n",
    "    print(f\"{pid}: {session['type']} {session['name']!r}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7dd12e3d-6d18-4bd6-a1b7-195e60ca761b",
   "metadata": {},
   "source": [
    "## Alternative: using kernel ids from connection files\n",
    "\n",
    "In [the question](https://discourse.jupyter.org/t/how-i-could-get-the-notebook-name-having-the-kernel-json-runtime-key-or-os-process-id/15168/2),\n",
    "the runtime file is also mentioned as a valid input.\n",
    "This may be simpler and more general to use.\n",
    "\n",
    "By convention, Jupyter servers put the kernel id in the connection file as well. So a simpler approach that works for all kernel languages\n",
    "is to check for the kernel id in.\n",
    "\n",
    "You could know _exactly_ how we produce these files and parse them out, but perhaps simpler and more robust to internal detail changes is simple substring matching:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "e18565c8-6b1e-42b6-99c5-12b587529a38",
   "metadata": {},
   "outputs": [],
   "source": [
    "from jupyter_core.paths import jupyter_runtime_dir\n",
    "import glob"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "b05af96c-d744-4e3f-8c3b-8b857201d053",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['/Users/minrk/Library/Jupyter/runtime/kernel-2efac3d3-4adb-4193-9d52-a68940660347.json',\n",
       " '/Users/minrk/Library/Jupyter/runtime/kernel-1482c156-1abc-44fd-8c35-74805a3c8e98.json',\n",
       " '/Users/minrk/Library/Jupyter/runtime/kernel-0518a5da-38d3-4dcd-8e1c-ddb7fbc494b3.json']"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "connection_files = glob.glob(f\"{jupyter_runtime_dir()}/kernel-*.json\")\n",
    "connection_files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "a64b86ce-4f78-43b6-85a4-61c6df1990cf",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "kernel-2efac3d3-4adb-4193-9d52-a68940660347.json: notebook 'notebook-a.ipynb'\n",
      "kernel-1482c156-1abc-44fd-8c35-74805a3c8e98.json: notebook 'other notebook.ipynb'\n",
      "kernel-0518a5da-38d3-4dcd-8e1c-ddb7fbc494b3.json: console 'Console 1'\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "\n",
    "for session in session_list:\n",
    "    kernel_id = session[\"kernel\"][\"id\"]\n",
    "    for connection_file in connection_files:\n",
    "        fname = os.path.basename(connection_file)\n",
    "        if kernel_id in fname:\n",
    "            print(f\"{fname}: {session['type']} {session['name']!r}\")\n",
    "            break\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.4"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {},
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}