Runtime-Augmented LLMs for Crash Detection and Diagnosis in ML Notebooks

TL;DR

CRANE-LLM enhances large language models with runtime data from ML notebooks to accurately detect and diagnose crashes, improving reliability in ML development environments.

Contribution

This work introduces CRANE-LLM, a novel method that combines static code and runtime information to improve crash detection and diagnosis in ML notebooks.

Findings

Runtime information improves detection accuracy by 7-10 percentage points.

Diagnosis F1-score increases by 8-11 points with runtime data.

Performance gains vary across ML libraries and crash types.

Abstract

Jupyter notebooks are widely used for machine learning (ML) development due to their support for interactive and iterative experimentation. However, ML notebooks are highly prone to bugs, with crashes being among the most disruptive. Despite their practical importance, systematic methods for crash detection and diagnosis in ML notebooks remain largely unexplored. We present CRANE-LLM, a novel approach that augments large language models (LLMs) with structured runtime information extracted from the notebook kernel state to detect and diagnose crashes before executing a target cell. Given previously executed cells and a target cell, CRANE-LLM combines static code context with runtime information, including object types, tensor shapes, and data attributes, to predict whether the target cell will crash (detection) and explain the underlying cause (diagnosis). We evaluate CRANE-LLM on JunoBench, a benchmark of 222 ML notebooks comprising 111 pairs of crashing and corresponding non-crashing notebooks across multiple ML libraries and crash root causes. Across three state-of-the-art LLMs (Gemini, Qwen, and GPT-5), runtime information improves crash detection and diagnosis by 7-10 percentage points in accuracy and 8-11 in F1-score, with larger gains for diagnosis. Improvements vary across ML libraries, crash causes, and LLMs, and depends on the integration of complementary categories of runtime information.