How To Implement The Jupyter Visual Studio Code Extension API To Monitor Kernel State
Posted: May 16th, 2024
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.
![Jupyter Notebook Visual Studio Code Extension API Development Tutorial](/static/images/jupyter-vscode.png)
Prerequisites
Before we begin, ensure you have the following installed:
Setting Up Your VS Code Environment
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.
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!
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