ML Notebook Repair Agent — LLM Agents for Crash Repair

This repository contains code, data, and analysis for the paper "Beyond Crash Resolution: Evaluating LLM Agents for Repairing Machine Learning Notebooks". We build and evaluate agentic LLM systems that automatically diagnose and repair failing Jupyter notebooks, utilizing mini-swe-agent as our underlying framework. The project implements agents and baselines, reproduces experiments on the JunoBench benchmark and our self-constructed extension, and includes analysis scripts used to produce the paper results.

Contents

• src/ — core agent and baseline implementations and utilities (see src/run_agent.py, src/run_baseline.py).
• JunoBench/ — benchmark notebooks used in experiments. We also extended the benchmark, which can be found on Zenodo.
• data_analysis/ — scripts for aggregating results, computing metrics, and producing plots.
• results_JunoBench/, results_new/ — Agent and baseline outputs and processed results from our experiments.
• config/ — YAML configs for running agents and baselines.
• run_all.sh — streamlined script for running experiments.
• tests/ — Unit and integration tests for the project (see TEST_SUITE.md).

Datasets used

• Primary benchmark: JunoBench, excluding NBSpecific and Torch 13.
• Supplementary JunoBench Extension: We also constructed 16 new notebook instances with the most recent Kaggle notebooks from 2026-01, as well as tests for each notebook instance (JunoBench and the 16 new instances).

If you need to add other datasets, place them in a new directory and update the appropriate config (source_path_parent).

Reproducing the experiments

Prerequisites

• Python 3.9+ and a POSIX-like shell for some helper scripts (on Windows use WSL or Git Bash).
• Docker
• uv (for Python dependency management)
• Git LFS (required to pull the benchmark dataset)

1. Environment Setup

Clone the repository and install dependencies:

# Install project dependencies
uv sync
source .venv/bin/activate

2. Pull the docker image required by notebooks in the datasets:

The implemented notebook environment for evaluating all notebooks in the datasets uses the docker environment from JunoBench.

docker pull yarinamomo/kaggle_python_env:latest

3. Configure environment variables and API keys:

Create a .env file for local configuration using the following template:

NOTEBOOK_AGENT_CONFIG_PATH=./config/agent.yaml
NOTEBOOK_AGENT_DEFAULTS_PATH=./config/defaults_local.yaml
NOTEBOOK_AGENT_LOG_LEVEL=INFO
OPENAI_API_KEY=your_key_here

Then, set up your keys for parallel execution:

• Edit api_keys.txt with your model credentials (one key per line). This file is specifically used for threading/parallel execution. Modify run_all.sh to ensure API_KEYS matches your configured text file.

4. Prepare the datasets:

JunoBench is included as a git submodule and uses Git LFS (Large File Storage) for the notebook data. Initialize and update it:

git submodule update --init --recursive
git lfs pull

To include the supplementary extension, download it from Zenodo and extract the new instances into your local benchmark folders. You will need to explicitly point the agent to these folders in your configuration files by changing the source_path_parent key in your config file (e.g., source_path_parent: JunoBench/benchmark/, see config/CONFIG.md).

5. Running experiments

• To run the full benchmark pipeline across all configurations:

./run_all.sh

• To run a single experiment or setting:

python main.py --config config/agent.yaml

Configuration Mechanics

The runner (main.py) relies on a layered configuration system that merges a defaults file with a run-specific configuration. For full details on all available configurations and extensions, see config/CONFIG.md.

Notes

• Long-running jobs will create per-run outputs under the configured folder in the config files, e.g. trajectories/[setting_name]/[model]/[run_id].

Results & analysis

Metrics reported in the paper can be reproduced using (main_run_data_analysis.py).

Acknowledgments

This project is built on top of the mini-swe-agent framework. We thank the authors and maintainers for providing a robust foundation for agentic software engineering.