In this blog post, we will explore how to implement the Jupyter Visual Studio Code Extension API to monitor the kernel state effectively. The Jupyter Extension API is fairly limited in its scope, but it provides a window into the Jupyter kernel that is very convenient to plug into for other Visual Studio Code extensions.

Prerequisites

Before we begin, ensure you have the following installed:

1. Visual Studio Code
2. Python
3. Jupyter Notebook Kernel
4. Jupyter Extension for VS Code

Setting Up Your VS Code Environment

1. Install the Jupyter Extension: Open VS Code, go to the Extensions view by clicking the square icon in the sidebar or pressing Ctrl+Shift+X, and search for "Jupyter". Install the Jupyter extension provided by Microsoft.

2. Create a Jupyter Notebook: Open the Command Palette (Ctrl+Shift+P), type Jupyter: Create New Blank Notebook, and select it. This will create a new Jupyter notebook file (.ipynb).

Accessing the Jupyter Extension API

The Jupyter extension for VS Code exposes an API that allows extensions to interact with Jupyter notebooks. We can use this API to monitor the kernel state.

Step 1: Create a VS Code Extension

First, create a new directory for your extension and navigate into it:

mkdir vscode-jupyter-kernel-monitor
cd vscode-jupyter-kernel-monitor

Next, initialize a new VS Code extension using yo (Yeoman) generator. If you don't have yo installed, you can install it using npm:

npm install -g yo generator-code
yo code

Follow the prompts to set up your extension. The following steps will be adding code to src/extension.ts

Step 2: Watch For New Open Notebooks

The Jupyter Extension API allows you to connect to running kernels, but you must provide the API with a Notebook URI. We need to write a function to watch for new Jupyter Notebooks opened in the editor so we can grab their URI and monitor it with the Jupyter kernel API.

async function notebookWatcher(context) {
    // Get Jupyter Extension public API implementation and make sure the extension is installed
    const jupyterExt = vscode.extensions.getExtension('ms-toolsai.jupyter');
    if (jupyterExt) {
        // Wait for the extension to go through its activation function
        const api = await jupyterExt.activate();
       
        let existingDocs = [];
        // Iterate forever and continue checking for new notebook documents in the workspace
        while(true) {
            for (const document of vscode.workspace.notebookDocuments) {
                if (!existingDocs.includes(document.uri)) {
                    // If we haven’t seen this notebook yet, pass it to a handler function
                    handleNotebookKernel(api, document.uri, context);
                    existingDocs.push(document.uri);
                }
            }
            await sleep(1000);
        }  
    }
 }

Step 3: Handle A Notebook’s Kernel State

Next we write a function to handle the kernel monitoring for each notebook in the workspace.

async function handleNotebookKernel(api, uri, context) {
    let kernelFound = false;
    let kernel = undefined;
    // Wait until a kernel is running. This requires the user to run the kernel, having the notebook document open in the workspace is not enough
    while(!kernelFound) {
        // Get the kernel for the notebook document URI passed to the function
        kernel = await api.kernels.getKernel(uri);
        if (kernel !== undefined) {
            kernelFound = true;
        }
        await sleep(1000);
    }
 
 
    // Notebook is open and the kernel is running. Add logic to handle the kernel’s status change
    context.subscriptions.push(kernel.onDidChangeStatus((e) => {
       vscode.window.showInformationMessage(e);
    }));
 }

Step 4: Tying Everything Together

Check out this GitHub Gist for the complete code for this example extension:

https://gist.github.com/2514millerj/7acb92d862f0ca40b9a57b66cad18db7

Example Use Case

At MeerkatIO we implement a similar process to monitor the state of execution of Jupyter Notebook cells in order to trigger personal notifications when a cell’s state transitions from “busy” to “idle”. If the cell execution data matches the MeerkatIO extension settings, a notification will be triggered by the MeerkatIO extension.

This is just one example, let us know how you plan to use the Jupyter Extension API!

Happy Hacking,
Justin
MeerkatIO.com

---

Sign up for MeerkatIO today to reclaim your free time and boost your productivity. With MeerkatIO, you can enjoy the peace of mind that comes from knowing you're not missing a beat while your code is running. Start your journey to a more efficient work-from-home experience today!

Back to Blog