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

📄 arXiv: 2603.07897v1 📥 PDF

作者: Lizhi Ma, Yi-Xiang Hu, Yihui Ren, Feng Wu, Xiang-Yang Li

分类: cs.LG

发布日期: 2026-03-09

💡 一句话要点

LeJOT-AutoML：基于LLM的特征工程，优化Databricks作业执行时间预测

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: AutoML 大型语言模型 特征工程 作业执行时间预测 Databricks 成本优化 云原生 资源调度

📋 核心要点

  1. 现有Databricks作业执行时间预测依赖手动特征工程，难以捕捉运行时动态变化，导致预测不准确和成本优化受限。
  2. LeJOT-AutoML利用LLM代理自动分析作业工件，生成运行时特征，并通过安全门验证，实现高效的特征工程和模型更新。
  3. 实验表明，LeJOT-AutoML能生成200+特征，将特征工程周期从数周缩短至20分钟，并在实际部署中节省19.01%的成本。

📝 摘要（中文）

Databricks作业编排系统（如LeJOT）通过选择低价计算配置来降低云成本，同时满足延迟和依赖约束。准确预测异构实例类型和非平稳运行时条件下的执行时间至关重要。现有流程依赖于静态、手动设计的特征，无法充分捕捉运行时效应（如分区修剪、数据倾斜和shuffle放大），且预测信号分散在日志、元数据和作业脚本中，延长了更新周期并增加了工程开销。我们提出了LeJOT-AutoML，一个由代理驱动的AutoML框架，将大型语言模型代理嵌入到ML生命周期中。LeJOT-AutoML结合了基于领域知识库的检索增强生成和模型上下文协议工具链（日志解析器、元数据查询和只读SQL沙箱），以分析作业工件，通过安全门合成和验证特征提取代码，并训练/选择预测器。这种设计实现了仅通过静态分析难以获得的运行时派生特征。在企业Databricks工作负载上，LeJOT-AutoML生成了超过200个特征，并将特征工程和评估循环从数周缩短到20-30分钟，同时保持了有竞争力的预测精度。集成到LeJOT流程中，它实现了自动连续模型更新，并在我们的部署环境中通过改进的编排实现了19.01%的成本节省。

🔬 方法详解

问题定义：论文旨在解决Databricks作业执行时间预测中，现有方法依赖静态特征、难以捕捉运行时效应的问题。现有方法需要大量人工特征工程，周期长、成本高，且难以适应非平稳的运行时环境，导致预测精度不足，无法有效优化云资源成本。

核心思路：论文的核心思路是利用大型语言模型（LLM）的强大理解和生成能力，构建一个自动化的特征工程框架。通过LLM分析作业相关的各种信息源（日志、元数据、脚本），自动提取运行时特征，并进行验证，从而克服传统方法的局限性。

技术框架：LeJOT-AutoML框架包含以下主要模块：1) 领域知识库：存储Databricks作业相关的知识，用于检索增强生成。2) 模型上下文协议工具链：包括日志解析器、元数据查询和只读SQL沙箱，用于访问和分析作业工件。3) LLM代理：负责分析作业工件，生成特征提取代码，并通过安全门进行验证。4) 训练/选择模块：使用生成的特征训练预测模型，并选择最佳模型。整体流程是，LLM代理利用领域知识库和模型上下文协议工具链分析作业，生成特征提取代码，经过安全验证后，用于训练预测模型。

关键创新：LeJOT-AutoML的关键创新在于利用LLM驱动的自动化特征工程。与传统的手动特征工程相比，LeJOT-AutoML能够自动分析作业相关的各种信息源，提取运行时特征，并进行验证，大大提高了特征工程的效率和质量。此外，该框架还采用了安全门机制，确保生成的特征提取代码的安全性。

关键设计：论文中提到使用了检索增强生成（Retrieval-Augmented Generation）来利用领域知识库，但没有详细说明检索和生成的具体算法。安全门（Safety Gates）的具体实现方式也未详细描述，但强调了其在保证特征提取代码安全性的作用。模型训练和选择的具体算法也未详细说明。

🖼️ 关键图片

fig_0

📊 实验亮点

LeJOT-AutoML在企业Databricks工作负载上生成了超过200个特征，并将特征工程和评估循环从数周缩短到20-30分钟。集成到LeJOT流程后，通过改进的编排实现了19.01%的成本节省。这些结果表明，LeJOT-AutoML能够显著提高特征工程的效率和预测精度，从而有效降低云资源成本。

🎯 应用场景

LeJOT-AutoML可应用于各种云原生数据处理平台的成本优化，通过更准确的作业执行时间预测，实现更高效的资源调度和成本控制。该研究对自动化特征工程和LLM在云资源管理领域的应用具有重要参考价值，未来可扩展到其他类型的云服务和工作负载。

📄 摘要（原文）

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.