# LeJOT-AutoML: LLM-Driven Feature Engineering for Job Execution Time Prediction in Databricks Cost Optimization

arXiv: https://arxiv.org/abs/2603.07897 · HTML: https://mlsanity.com/paper/2603.07897/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 12 of its 24 references that are in the map.
- First posted on arXiv 9 Mar 2026.

## Builds on (most cited first, 12 of 12)

- [Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks](https://mlsanity.com/paper/2005.11401/) (2020)
- [Qwen3 Technical Report](https://mlsanity.com/paper/2505.09388/) (2025)
- [ReAct: Synergizing Reasoning and Acting in Language Models](https://mlsanity.com/paper/2210.03629/) (2022)
- [XGBoost: A Scalable Tree Boosting System](https://mlsanity.com/paper/1603.02754/) (2016)
- [GPT-4o System Card](https://mlsanity.com/paper/2410.21276/) (2024)
- [Toolformer: Language Models Can Teach Themselves to Use Tools](https://mlsanity.com/paper/2302.04761/) (2023)
- [Auto-Sklearn 2.0: Hands-free AutoML via Meta-Learning](https://mlsanity.com/paper/2007.04074/) (2020)
- [A Survey of Controllable Text Generation using Transformer-based Pre-trained Language Models](https://mlsanity.com/paper/2201.05337/) (2022)
- [Large Language Models for Automated Data Science: Introducing CAAFE for Context-Aware Automated Feature Engineering](https://mlsanity.com/paper/2305.03403/) (2023)
- [Learned Cardinalities: Estimating Correlated Joins with Deep Learning](https://mlsanity.com/paper/1809.00677/) (2018)
- [SLAQ: Quality-Driven Scheduling for Distributed Machine Learning](https://mlsanity.com/paper/1802.04819/) (2018)
- [LeJOT: An Intelligent Job Cost Orchestration Solution for Databricks Platform](https://mlsanity.com/paper/2512.18266/) (2025)

## Abstract (arXiv)

Databricks job orchestration systems (e.g., LeJOT) reduce cloud costs by selecting low-priced compute configurations while meeting latency and dependency constraints. Accurate execution-time prediction under heterogeneous instance types and non-stationary runtime conditions is therefore critical. Existing pipelines rely on static, manually engineered features that under-capture runtime effects (e.g., partition pruning, data skew, and shuffle amplification), and predictive signals are scattered across logs, metadata, and job scripts-lengthening update cycles and increasing engineering overhead. We present LeJOT-AutoML, an agent-driven AutoML framework that embeds large language model agents throughout the ML lifecycle. LeJOT-AutoML combines retrieval-augmented generation over a domain knowledge base with a Model Context Protocol toolchain (log parsers, metadata queries, and a read-only SQL sandbox) to analyze job artifacts, synthesize and validate feature-extraction code via safety gates, and train/select predictors. This design materializes runtime-derived features that are difficult to obtain through static analysis alone. On enterprise Databricks workloads, LeJOT-AutoML generates over 200 features and reduces the feature-engineering and evaluation loop from weeks to 20-30 minutes, while maintaining competitive prediction accuracy. Integrated into the LeJOT pipeline, it enables automated continuous model updates and achieves 19.01% cost savings in our deployment setting through improved orchestration.