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

2026-03-09

AI Summary

This work proposes LeJOT-AutoML, the first AutoML framework integrating a large language model (LLM) agent with the Model Context Protocol to address the challenge of job runtime prediction for cloud cost optimization in Databricks. Existing approaches rely on static, handcrafted features and fail to capture dynamic runtime effects such as data skew and partition pruning, while also suffering from lengthy iteration cycles. LeJOT-AutoML overcomes these limitations by leveraging retrieval-augmented generation (RAG) in coordination with a log parser, metadata query engine, and a read-only SQL sandbox to automatically construct and validate runtime-aware features. The framework enables end-to-end automated feature engineering, generating over 200 features on enterprise-scale workloads within 20–30 minutes—reducing the feature engineering cycle from weeks to minutes. When integrated into production pipelines, it achieves a 19.01% reduction in cloud costs while maintaining high prediction accuracy.

Technology Category

• Machine Learning: Auto ML and Hyperparameter Tuning
• Search and Optimization: Learning to Search
• Planning, Routing, and Scheduling: Planning with Language Models

Application Category

• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems
• Semantics and Knowledge: Data modeling to support human-machine intelligence, including LLMs agents, intelligent system behavior, explanations, and user-friendly interactions
• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation

Abstract

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.

Problem

• Job Execution Time Prediction
• Feature Engineering
• Databricks
• Cost Optimization
• Runtime Effects

Innovation

• LLM-driven AutoML
• feature engineering
• execution time prediction
• retrieval-augmented generation
• Model Context Protocol