JupyterLab GPU 代码执行 MCP 服务器

jlab-mcp

A Model Context Protocol (MCP) server that enables Claude Code to execute Python code on GPU compute nodes via JupyterLab running on a SLURM cluster.

Inspired by and adapted from goodfire-ai/scribe, which provides notebook-based code execution for Claude. This project adapts that approach for HPC/SLURM environments where GPU resources are allocated via job schedulers.

Architecture

Claude Code
    ↕ stdio
MCP Server
    ↕ HTTP/WebSocket
JupyterLab (SLURM compute node or local subprocess)   ← one server, many kernels
    ↕
IPython Kernels (GPU access)

JupyterLab runs either on a SLURM compute node (HPC clusters) or as a local subprocess (laptops/workstations). The server process is decoupled from the MCP server and keeps running across Claude Code sessions. It can be started two ways:

• From Claude (recommended): the agent calls the start_server MCP tool, which asks you once whether to run on SLURM or locally (and, for SLURM, the job walltime and resources), launches jlab-mcp start in the background (bootstrapping a bare project with uv init/uv add if needed), then monitors startup with wait_for_server and tells you when it's ready.
• Manually: run jlab-mcp start in a separate terminal.

All sessions create separate kernels on this shared server. Each project directory gets its own JupyterLab instance — the status file is scoped by a hash of the working directory where the server was started.

Local Mode

On machines without SLURM (laptops, workstations), jlab-mcp runs JupyterLab as a local subprocess.

How the mode is chosen:

• Via the start_server MCP tool: on first use in a project it asks you (slurm or local) and saves the choice next to the project's status file (~/.jlab-mcp/servers/{name}-{hash}/run-mode). Pass mode explicitly to change it later.
• Via the CLI (jlab-mcp start): auto-detected — if sbatch is on PATH, SLURM mode; otherwise local mode.
• The JLAB_MCP_RUN_MODE environment variable overrides both:

export JLAB_MCP_RUN_MODE=local   # force local mode
export JLAB_MCP_RUN_MODE=slurm   # force SLURM mode

In local mode, jlab-mcp start runs in the foreground — press Ctrl+C to stop. The status file uses the same format as SLURM mode, so the MCP server works identically in both modes.

Setup

Zero-install (recommended)

The only prerequisites are Claude Code and uv. Drop this .mcp.json into any project directory — uvx fetches and runs jlab-mcp on demand (cached after the first run), no separate install step:

{
  "mcpServers": {
    "jlab-mcp": {
      "command": "uvx",
      "args": ["--from", "git+https://github.com/kdkyum/jlab-mcp.git", "jlab-mcp"]
    }
  }
}

Start Claude Code in that directory and approve the MCP server when prompted. The first launch clones and builds the package, so it can take a little longer; pin a tag (git+...@v1.0.2) for reproducibility.

Pre-installed alternative

For faster MCP startup (or offline login nodes), install the CLI once and reference it directly:

uv tool install git+https://github.com/kdkyum/jlab-mcp.git

The SLURM job activates .venv in the current working directory. If the project has no environment yet (e.g. a fresh directory with just .mcp.json), start_server bootstraps it automatically (uv init --bare + uv add jupyterlab ipykernel matplotlib numpy). To set it up manually, or to add GPU-enabled torch:

cd /shared/fs/my-project
uv venv
uv pip install jupyterlab ipykernel matplotlib numpy
uv pip install torch --index-url https://download.pytorch.org/whl/cu126   # NVIDIA GPUs
# AMD GPUs (e.g. MI300A): use the ROCm wheels instead
# uv pip install torch --index-url https://download.pytorch.org/whl/rocm6.3

Usage

Option A: Let Claude manage the server

Just start Claude Code in your project directory and ask it to run something in a notebook. When no server is running, the agent calls start_server (asking you SLURM vs local on first use), waits in the background via wait_for_server, and reports when JupyterLab is ready. To stop the server, ask Claude or run jlab-mcp stop.

Option B: Manual CLI

1. Start the compute node

In a separate terminal, start the SLURM job:

jlab-mcp start              # uses default time limit (4h)
jlab-mcp start 24:00:00     # 24 hour time limit
jlab-mcp start 1-00:00:00   # 1 day

This submits the job and waits until JupyterLab is ready:

SLURM job 24215408 submitted, waiting in queue...
Job running on ravg1011, JupyterLab starting...
JupyterLab ready at http://ravg1011:18432

2. Use Claude Code

In another terminal, start Claude Code. The MCP server connects to the running JupyterLab automatically.

3. Stop when done

jlab-mcp stop

CLI Commands

All commands accept --debug to enable verbose logging (status file reads, SLURM parameters, health check attempts, connection file paths) on stderr.

The SLURM job survives Claude Code restarts. You only need to run jlab-mcp start once per work session.

Configuration

All settings are configurable via environment variables. No values are hardcoded for a specific cluster.

Example: Cluster with A100 GPUs and CUDA module

export JLAB_MCP_SLURM_PARTITION=gpu1
export JLAB_MCP_SLURM_GRES=gpu:a100:1
export JLAB_MCP_SLURM_CPUS=18
export JLAB_MCP_SLURM_MEM=125000
export JLAB_MCP_SLURM_TIME=1-00:00:00
export JLAB_MCP_SLURM_MODULES="cuda/12.6"

Claude Code Integration

Use the zero-install .mcp.json from Setup, or — with the CLI pre-installed — reference the binary directly in ~/.claude.json or a project .mcp.json:

{
  "mcpServers": {
    "jlab-mcp": {
      "command": "jlab-mcp"
    }
  }
}

No env block is needed: on first use the agent surveys the cluster (sinfo) and asks you for walltime and resources, saving the choices per project. An env block with JLAB_MCP_* variables still works to pin cluster-specific defaults (e.g. JLAB_MCP_SLURM_MODULES for a CUDA module).

The MCP server uses the working directory to find .venv for the compute node. Claude Code launches from your project directory, so it picks up the right venv automatically.

MCP tools

Resource: jlab-mcp://server/status returns server information and active sessions.

Background runs

start_run snapshots the selected cell IDs and exact source before execution. Index ranges are inclusive; cell_ids executes IDs in the supplied order. Markdown is skipped. A second execution request on the same kernel returns state="busy" and active_run_id. This reservation covers the whole run, including foreground execution tools. It does not control requests sent directly from JupyterLab or other kernel clients.

run = start_run(session_id="SESSION", request_id="training-attempt-1")
get_run_status(run_id=run["run_id"], include_cells=True)
get_run_output(run_id=run["run_id"], cell_id="CELL")

A retry with the same request_id and execution parameters returns the existing run, even if its source has since changed. Reusing that identifier with different parameters returns a conflict. Request identifiers are unique within the project's run store. No automatic history expiration is applied.

started_code_cells counts cells acknowledged by the kernel. completed_code_cells counts cells whose execution reached the matching idle message, including cells that raised a Python exception. Neither count estimates training steps or loop iterations. current_cell.index is the index at the time of the source snapshot.

Status and output queries do not control execution. Cancelling a query or allowing it to time out leaves the run alone. cancel_run requests an interrupt; an unconfirmed interrupt remains uncertain. A foreground execution call retains its cancellation behavior and requests cancellation of its run. Per-cell timeout limits monitoring, not kernel computation; an expired timeout produces unknown and requires explicit cancellation before another run can use that kernel.

Durable capture and recovery

The project stores run records, source snapshots, received outputs, and session mappings in STATUS_DIR/runs/runs.sqlite3. SQLite commits captured outputs with synchronous=FULL. Image originals are retained. get_run_status reports notebook attachment through outputs_saved and durable capture through captured_outputs_saved; per-cell details distinguish notebook_saved, capture_saved, output_capture_complete, save_error, and capture_error.

Background workers live in the MCP process. They continue after a start response and across client reconnections while that process stays alive. If the MCP process exits, the kernel may continue the current cell, but the remaining cells are not dispatched. A new MCP process recovers the saved run and session mappings. Unfinished runs become unknown and are never restarted automatically. Captured messages remain retrievable; messages not received and committed before the crash cannot be promised recoverable.

A WebSocket reconnect listens for the existing request without resending code. Any connection gap marks output capture incomplete. Python exceptions, connection failures, and kernel death have separate status fields. An idle kernel after a disconnect does not establish whether a request finished or never arrived.

Source revisions and document edits

Editing and reading tools accept either session_id or notebook_path. Paths stay within JLAB_MCP_NOTEBOOK_DIR. create_notebook and open_notebook need no JupyterLab server. Cell edits, deletion, execution, and output retrieval accept stable cell IDs; integer indices remain available. expected_source_hash on edits and deletion returns a conflict with the current source instead of overwriting a changed cell. batch_edit supports the same check per operation and leaves the document unchanged if validation fails.

Each execution records its run ID, cell ID, exact source hash, execution count, and timestamps. Before saving results as current notebook output, the writer compares that hash with the cell's current source. Moving a cell preserves its output association. Editing its source during execution leaves the old result in run history and reports source_mismatch; it does not attach that result to the new source. get_run_output(include_source=true) retrieves the executed source alongside its result.

Identical-source edits preserve the file and outputs. Changed source clears current outputs. Earlier tool executions remain available by run ID. edit_cell(preserve_history=true) also copies the previous outputs and source into cell metadata, including outputs that predate run tracking; get_cell(include_history=true) reads that history.

Notebook writes use a lock shared by all managers for the same file within one MCP process. Revision checks compare the version read by the tool. They cannot lock out a separate JupyterLab process writing the file at the same instant.

Output and figure selection

Both output tools return a JSON metadata block followed by the selected content. metadata_only=true returns counts, output types, image dimensions, and execution provenance. output_index selects one output; start_output and end_output select an inclusive range. Negative indices count from the end. Indices refer to the output entries reported by that tool's metadata, not just images.

get_cell_output(session_id="SESSION", cell_id="CELL", metadata_only=True)
get_cell_output(
	session_id="SESSION", cell_id="CELL", output_index=4,
	include_images=True, image_resolution="original",
)

Images are omitted by default. image_resolution="preview" caps the longest dimension at max_image_dimension, default 2576. "original" returns the original captured image. Preview settings never alter saved output. max_chars limits returned text to its tail; zero disables truncation.

Kernel lifecycle

Creating notebook B preserves notebook A's kernel and variables. Opening a notebook reports whether it reused a recorded live kernel or created a fresh one. A failed liveness check leaves the existing session intact and does not create a replacement. Session mappings persist across MCP restarts. shutdown_session followed by start_notebook starts a fresh kernel for the same document.

Testing

# Unit tests (no SLURM needed)
uv run python -m pytest tests/ --ignore=tests/test_tools.py -v

# Isolated local integration tests, using a disposable JupyterLab server
JLAB_MCP_TEST_LOCAL=1 uv run python -m pytest tests/test_live_runs.py tests/test_tools.py -v --timeout=120

# Existing-server integration tests, after running `jlab-mcp start`
uv run python -m pytest tests/test_tools.py -v -s --timeout=600

Acknowledgments

This project is inspired by goodfire-ai/scribe, which provides MCP-based notebook code execution for Claude. The tool interface design, image resizing approach, and notebook management patterns are adapted from scribe for use on HPC/SLURM clusters.

License

MIT