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

Sebastian Schelter,Stefan Grafberger

2024-09-16

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.

Databases

What problem does this paper attempt to address?

The problem that this paper attempts to solve is that in machine learning (ML) applications, due to the fact that the code usually written by data scientists is rather chaotic and lacks declarative abstractions, it is difficult to manage and comply with upcoming regulations. Specifically, the paper points out that current machine - learning pipeline libraries and services lack fundamental data - centric abstractions, such as logical query plans in databases. These services often focus on flexibility and ease of deployment, and regard data as part of the workflow rather than a core element. Consequently, the responsibility for handling complex business requirements (such as regulatory compliance) falls on developers, which increases the difficulty of implementation.

To address this challenge, the paper proposes leveraging the capabilities of large - language models (LLMs) to rewrite chaotic data - science code and transform it into declarative pipeline abstractions. In this way, the manageability of the code and compliance with regulatory requirements can be improved. The paper introduces a prototype system named Lester, which is able to convert chaotic Python code into declarative data - frame operations with relational semantics and fine - grained row - and - column lineage tracking. The design of Lester aims to cover existing research on enhancing ML pipelines and demonstrates its potential in solving complex tasks such as compliance. For example, Lester can quickly identify and update data artifacts containing sensitive information by tracing the origin and usage of data, thereby meeting the need for urgently deleting leaked data.