AutoML-Pipeline: A RAG-Enhanced Code Generation Framework With Pre-Validation for Cloud-Native Machine Learning Workflows
Wenyu Zhao, Tingjie Chen, J. C. Yang, Lei Qiu
Abstract
The proliferation of cloud-native machine learning platforms has significantly accelerated model development and deployment cycles. However, constructing and maintaining heterogeneous pipeline code spanning multiple languages (Python, YAML, Spark SQL) and cloud-specific configurations remains labor-intensive and error-prone. Existing LLM-based code generation tools lack awareness of runtime constraints and historical execution patterns, frequently producing code with resource misconfigurations or dependency conflicts that fail upon deployment. To address these challenges, we propose <italic xmlns:mml="http://www.w3.org/1998/Math/MathML" xmlns:xlink="http://www.w3.org/1999/xlink">AutoML-Pipeline</i>, a closed-loop code generation framework that integrates Retrieval-Augmented Generation (RAG) with reinforcement learning feedback mechanisms. Our approach leverages a knowledge base constructed from successful pipeline execution logs to guide GPT-4 in generating deployment-ready code that adheres to platform-specific constraints. The key innovation lies in a novel <italic xmlns:mml="http://www.w3.org/1998/Math/MathML" xmlns:xlink="http://www.w3.org/1999/xlink">Pre-validation Agent</i> that employs simulated execution environments to predict resource consumption and detect dependency conflicts before actual deployment. This agent iteratively refines generated code through a feedback loop informed by predicted execution profiles and dependency graphs. We evaluate our framework on the CodeSearchNet dataset augmented with Azure ML pipeline specifications, demonstrating a 43.7% improvement in first-submission success rate and 31.2% reduction in resource over-provisioning compared to vanilla GPT-4 baselines. Ablation studies confirm that both the RAG retrieval mechanism and pre-validation agent contribute substantially to performance gains. Our work establishes a practical paradigm for integrating large language models with domain-specific runtime intelligence, with potential applications extending to other infrastructure-as-code generation tasks.