Automated Assertion Generation and Regression Testing for Machine Learning Notebooks

What is it about?

Jupyter Notebooks have become the de-facto choice for data scientists and machine learning (ML) engineers for prototyping and experimenting with ML pipelines. Notebooks provide a rich interactive interface with support for code, data, and visualization in one place. However, notebooks provide limited support for testing. As a result, during continuous development, many silent (non-crashing) regressions often go unnoticed and make notebooks unreliable and results hard to reproduce. To enable more systematic testing of ML notebooks, we introduce NBTestGen – the first automated assertion generation approach for ML notebooks. NBTestGen generates regression-based assertions that check properties related to data processing, model building, and model evaluation steps in a typical ML notebook. To support systematic integration of such assertions in notebooks, we introduce the first regression testing framework (called NBTest) that can be used as a Jupyter plugin and allows developers to write cell-scoped assertions in notebooks. Three key features of such assertions are that they are 1) cell-scoped: they are linked to specific notebook cells and execute only after those cells are executed, 2) non-intrusive: they do not block notebook execution (in a Jupyter session), so that development can continue when they fail, and 3) they integrate with pytest and CI pipelines, allowing developers to easily do regression testing of their notebooks. We evaluate NBTestGen on a corpus of 585 notebooks from the popular Kaggle platform. NBTestGen generates a total of 21542 assertions (36.82 on average per notebook). The generated assertions kill 72.21% of ML-specific mutations, while maintaining a high passrate of 100.00%. We also show that NBTestGen can detect 69.68% of historical regressions in 531 older versions of Kaggle notebooks. A popular ML library, SHAP, integrated NBTest into their CI. Further, we perform a user study with 17 ML developers that shows that such users find NBTest highly intuitive and useful.

Featured Image

Photo by Daniil Komov on Unsplash

Photo by Daniil Komov on Unsplash