Messy Code Makes Managing ML Pipelines Difficult? Just Let LLMs Rewrite the Code!

TL;DR

The paper tackles regulatory and data-management gaps in ML pipelines caused by the absence of declarative, data-centric abstractions. It introduces Lester, an approach that leverages LLM-assisted rewriting to convert messy imperative data-science code into a declarative pipeline with fine-grained provenance and an incremental view maintenance capability, exemplified by $\mathcal{D}_{prep}$, $\mathbf{X}$, and $\mathbf{y}$. Key contributions include a formal ML-pipeline model with provenance polynomials, a practical LLM-based rewrite workflow, and a demonstration of low-latency updates for deployed pipelines, plus a small user study; the code and scenario are publicly available. The work demonstrates how automated code rewriting can enable provenance-aware, compliant ML deployments and reduces the manual burden of refactoring legacy data-science code, with broad applicability beyond ML pipelines.

Abstract

Machine learning (ML) applications that learn from data are increasingly used to automate impactful decisions. Unfortunately, these applications often fall short of adequately managing critical data and complying with upcoming regulations. A technical reason for the persistence of these issues is that the data pipelines in common ML libraries and cloud services lack fundamental declarative, data-centric abstractions. Recent research has shown how such abstractions enable techniques like provenance tracking and automatic inspection to help manage ML pipelines. Unfortunately, these approaches lack adoption in the real world because they require clean ML pipeline code written with declarative APIs, instead of the messy imperative Python code that data scientists typically write for data preparation. We argue that it is unrealistic to expect data scientists to change their established development practices. Instead, we propose to circumvent this "code abstraction gap" by leveraging the code generation capabilities of large language models (LLMs). Our idea is to rewrite messy data science code to a custom-tailored declarative pipeline abstraction, which we implement as a proof-of-concept in our prototype Lester. We detail its application for a challenging compliance management example involving "incremental view maintenance" of deployed ML pipelines. The code rewrites for our running example show the potential of LLMs to make messy data science code declarative, e.g., by identifying hand-coded joins in Python and turning them into joins on dataframes, or by generating declarative feature encoders from NumPy code.

Figures (4)

• Figure 1: LLM-assisted rewrite of messy data preparation code to declarative dataframe operations in Lester with relational semantics and fine-grained row and column provenance tracking.
• Figure 2: LLM-assisted rewrite of messy data feature encoding code to estimator/transformer operations in scikit-learn for Lester with fine-grained provenance tracking.
• Figure 3: -- Initial pipeline execution with artifact capture and provenance tracking; -- Provenance-driven low-latency removal of security-critical data across all pipeline artifacts (intermediate training data, feature matrix, ML model).
• Figure 4: Time (in logarithmic scale) to re-execute a pipeline from scratch versus the time to update the pipeline artifacts with IVM. Lester enables sub-second updates.