# jlab-mcp

A [Model Context Protocol (MCP)](https://modelcontextprotocol.io/) 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](https://github.com/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:

```bash
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](https://docs.astral.sh/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:

```json
{
  "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:

```bash
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:

```bash
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:

```bash
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

```bash
jlab-mcp stop
```

### CLI Commands

| Command | Description |
|---|---|
| `jlab-mcp start [TIME] [--debug]` | Start JupyterLab and wait until ready. In SLURM mode, submits a job and polls until the server responds. In local mode, spawns a subprocess and blocks in the foreground. Optional TIME overrides `JLAB_MCP_SLURM_TIME` (e.g. `24:00:00`). Skips submission if an existing server is still running. |
| `jlab-mcp stop` | Stop JupyterLab. In SLURM mode, runs `scancel`. In local mode, sends SIGTERM to the subprocess. Removes the status file in both cases. |
| `jlab-mcp wait` | Poll the status file from another terminal until the server is ready (up to 10 min). Prints state transitions (`pending → starting → ready`). Useful for scripts or for monitoring `start` progress from a separate shell. |
| `jlab-mcp status` | Print server state, mode, hostname, port, and whether the process/job is alive. Lists active kernels with execution state and last activity time. Queries GPU memory and utilization via `nvidia-smi` on a temporary kernel. |
| `jlab-mcp` | Run MCP server (stdio transport, used by Claude Code — not run manually) |

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.

| Environment Variable | Default | Description |
|---|---|---|
| `JLAB_MCP_DIR` | `~/.jlab-mcp` | Base working directory |
| `JLAB_MCP_NOTEBOOK_DIR` | `./notebooks` | Notebook storage (relative to cwd) |
| `JLAB_MCP_SERVER_ROOT_DIR` | cwd | JupyterLab root directory (what the file browser sees) |
| `JLAB_MCP_LOG_DIR` | `~/.jlab-mcp/logs` | SLURM job logs |
| `JLAB_MCP_STATUS_DIR` | `~/.jlab-mcp/servers/{name}-{hash}` | Per-project status directory (auto-derived from cwd) |
| `JLAB_MCP_CONNECTION_DIR` | `~/.jlab-mcp/connections` | Connection info files |
| `JLAB_MCP_SLURM_PARTITION` | `gpu` | SLURM partition |
| `JLAB_MCP_SLURM_GRES` | `gpu:1` | SLURM generic resource |
| `JLAB_MCP_SLURM_CPUS` | `4` | CPUs per task |
| `JLAB_MCP_SLURM_MEM` | `32000` | Memory in MB |
| `JLAB_MCP_SLURM_TIME` | `4:00:00` | Wall clock time limit |
| `JLAB_MCP_SLURM_BIND_IP` | `0.0.0.0` | Address JupyterLab binds to on the compute node. Default listens on all interfaces and advertises the node's `$(hostname)`; a concrete IP binds that one interface and advertises that exact IP |
| `JLAB_MCP_SLURM_MODULES` | *(empty)* | Space-separated modules to load (e.g. `cuda/12.6`) |
| `JLAB_MCP_QUEUE_TIMEOUT` | `300` | Seconds `start` waits for the job to leave the queue. On timeout the job **stays queued** — rerun `jlab-mcp start` to resume waiting |
| `JLAB_MCP_READY_TIMEOUT` | `120` | Seconds to wait for JupyterLab once the job is running. On timeout the job is cancelled |
| `JLAB_MCP_WALLTIME_WARN_SECONDS` | `3600` | Remaining SLURM walltime below which `ping`/`start_server` include a save-your-work warning |
| `JLAB_MCP_PORT_MIN` | `18000` | Port range lower bound |
| `JLAB_MCP_PORT_MAX` | `19000` | Port range upper bound |
| `JLAB_MCP_RUN_MODE` | *(auto)* | `local` or `slurm` (auto-detects based on `sbatch` availability) |
| `JLAB_MCP_LOCAL_BIND_IP` | `0.0.0.0` | Address JupyterLab binds to in local mode. Default listens on all interfaces (UI reachable from other hosts / a container host); the same-host MCP server still connects over loopback. Set `127.0.0.1` to restrict JupyterLab to loopback only |

### Example: Cluster with A100 GPUs and CUDA module

```bash
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](#setup), or — with the CLI pre-installed — reference the binary directly in `~/.claude.json` or a project `.mcp.json`:

```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

| Tool | Description |
|---|---|
| `start_server` | Launch JupyterLab in the background through `jlab-mcp start`. Saves the selected local or SLURM mode and resource options per project |
| `wait_for_server` | Monitor server startup with MCP progress notifications |
| `create_notebook` | Create an empty notebook without starting a kernel; existing filenames receive a numeric suffix |
| `open_notebook` | Read a document without a kernel; optional `include_cells="summary"` or `"source"` |
| `start_new_notebook` | Create a notebook and its kernel, preserving every existing kernel |
| `start_notebook` | Attach to a notebook and report `kernel_reused`, `kernel_created`, kernel status, and cell count. Source requires `include_cells="source"` |
| `list_cells` | List code and Markdown cells with stable `cell_id`, index, `source_hash`, and output counts |
| `get_cell` | Read one cell by ID or index, with optional source and execution history |
| `add_code_cell` | Insert source without executing code or advancing the kernel execution count |
| `edit_cell` | Edit code without execution. Changed source clears current outputs; identical source preserves them |
| `add_markdown`, `edit_markdown` | Insert or edit Markdown without execution |
| `delete_cell` | Delete a cell by stable ID or index |
| `batch_edit` | Validate insert, edit, delete, and move operations before applying them in one notebook save |
| `start_run` | Start a background run and return a durable run ID. Supports cell ranges, explicit cell IDs, and idempotent `request_id` retries |
| `get_run_status` | Query started and completed cells, elapsed times, latest text, errors, and capture/save status |
| `get_run_output` | Retrieve captured output while a cell runs or after it finishes, including results for old source revisions |
| `cancel_run` | Explicitly interrupt a run without shutting down its kernel |
| `execute_code` | Insert and execute a cell in the foreground, with durable run history |
| `run_cell`, `run_cells` | Execute existing source in the foreground with the same run records and kernel reservations |
| `get_cell_output` | Read saved notebook output with optional output selection, image previews, original images, and metadata |
| `execute_scratch` | Execute diagnostic code on a temporary kernel without a notebook |
| `interrupt_kernel` | Send an interrupt to a session's kernel |
| `shutdown_session` | Explicitly stop one kernel; other kernels and the JupyterLab server remain available |
| `ping` | Check server reachability, list kernels and notebook associations, and report SLURM walltime |
| `check_resources` | Inspect CPU, memory, and GPU usage through a temporary kernel |

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.

```python
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.

| State | Meaning |
|---|---|
| `queued` | The run is recorded and its worker is starting |
| `running` | The worker is dispatching or monitoring cells |
| `cancelling` | Cancellation was requested; the worker is waiting for execution to finish |
| `completed` | Every selected cell completed without a reported execution error |
| `failed` | One or more cells failed; `error_kind` distinguishes Python exceptions from failures before delivery. `stop_on_error` controls whether later cells run |
| `cancelled` | The run was cancelled, with no remaining cells dispatched |
| `kernel_died` | Kernel death, restart, or removal was detected |
| `unknown` | Execution outcome cannot be established; the kernel remains reserved |

`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.

```python
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

```bash
# 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](https://github.com/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