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

Het paper introduceert LeJOT-AutoML, een door LLM-agenten aangedreven AutoML-framework dat automatisch complexe runtime-kenmerken genereert voor het voorspellen van uitvoeringstijden in Databricks, waardoor de feature-engineeringcyclus van weken naar minuten wordt teruggebracht en de kosten met 19,01% worden verlaagd.

De kern

Stel je voor dat je een gigantische, complexe fabriek hebt (zoals Databricks) waar duizenden taken elke dag worden uitgevoerd. Sommige taken zijn simpel, andere zijn zwaar als een vrachtwagen. Het doel van het systeem LeJOT is om te beslissen: "Welke machine (computer) gebruiken we voor deze taak om het goedkoopst te zijn, zonder dat het te lang duurt?"

Het probleem? Als je de verkeerde machine kiest, is het te duur of te traag. Om dit te weten, moet je precies kunnen voorspellen hoe lang een taak duurt.

Hier komt LeJOT-AutoML in beeld. Dit is een slimme, nieuwe manier om die voorspellingen te maken, met behulp van kunstmatige intelligentie (LLM's). Hier is de uitleg in gewone taal, met een paar creatieve vergelijkingen:

1. Het Oude Probleem: De "Handgemaakte" Chef-kok

Vroeger deden mensen dit werk. Ze waren als chef-koks die een recept schreven voor het voorspellen van de tijd.

• Ze keken naar statische dingen: "Hoe groot is het bestand?" of "Hoeveel CPU heeft de machine?"
• Het nadeel: Dit is als een kookboek dat niet weet dat de oven warm is of dat de kip nat is. In de echte wereld verandert alles constant. Soms is een taak sneller omdat de data al gesorteerd is, of trager omdat de data scheef ligt (zoals een auto die vastzit in modder).
• Mensen moesten maanden werken om nieuwe regels te bedenken als de situatie veranderde. Het was traag, duur en vaak onnauwkeurig.

2. De Nieuwe Oplossing: Het "AI-Team"

LeJOT-AutoML is geen enkele robot, maar een team van slimme agenten die samenwerken. Denk aan een super-efficiënte keuken met drie specialisten:

• De Onderzoeker (Feature Analyzer Agent):
  Deze agent leest duizenden oude rapporten en kookboeken (de kennisbank). Hij zegt: "Hé, ik heb gezien dat als er veel data wordt verplaatst, het vaak trager gaat. Laten we dat gaan meten!" Hij bedenkt honderden nieuwe manieren om de taak te analyseren.
• De Uitvoerder (Feature Extraction Agent):
  Deze agent is de "handjes". Hij pakt de ideeën van de Onderzoeker en gaat ze direct testen in de fabriek. Hij kijkt in de logboeken, vraagt aan de database en simuleert kleine tests. Hij verzamelt de feiten: "Oké, deze taak heeft inderdaad veel data verplaatst, hier is het exacte getal."
• De Kwaliteitscontroleur (Feature Evaluation Agent):
  Deze agent is de strenge keurmeester. Hij kijkt of de nieuwe data nuttig is en of er geen fouten in zitten. Hij zorgt dat niemand "valsspelen" (bijvoorbeeld door de uitkomst al te kennen voordat de taak begint).

3. De "Veiligheidspoortjes"

Voordat de AI iets doet, moet het door twee poortjes:

1. De Code-Controle: Is het recept wel compleet? (Zitten er ingrediënten in die ontbreken?)
2. De Data-Lek-Controle: Kijkt de AI niet naar het antwoord voordat de taak is gedaan? (Dat zou valsspelen zijn).

4. Wat levert dit op? (De Resultaten)

• Snelheid: Waar mensen maanden nodig hadden om een nieuw voorspellingssysteem te bouwen, doet dit AI-team het in 20 tot 30 minuten. Het is alsof je van handmatig koken overschakelt op een robot die in een seconde een perfect maaltijd bereidt.
• Meer Inzichten: Het systeem vindt meer dan 200 nieuwe factoren om naar te kijken (zoals hoe de data zich gedraagt tijdens het draaien), terwijl mensen er maar ongeveer 40 vonden.
• Kostenbesparing: Omdat de voorspellingen beter zijn, kiest het systeem de juiste, goedkopere machines. In de praktijk heeft dit geleid tot 19% kostenbesparing. Dat is alsof je voor hetzelfde eten 19% minder betaalt.

Samenvattend

LeJOT-AutoML is als het vervangen van een ouderwetse, trage planner door een slim, zelflerend team van detectives. Ze kijken niet alleen naar wat er op papier staat, maar voelen ook hoe de machine echt aanvoelt tijdens het werk. Hierdoor wordt het bedrijf goedkoper, sneller en slimmer, zonder dat mensen urenlang hoeven te knutselen aan de code.

Technische samenvatting

1. Het Probleem

In Databricks-omgevingen is het nauwkeurig voorspellen van de uitvoeringstijd van taken (jobs) cruciaal voor het optimaliseren van cloudkosten. Het bestaande systeem LeJOT kiest goedkope compute-configuraties op basis van deze voorspellingen, mits latentie- en afhankelijkheidsconstraints worden gehaald.

De huidige uitdagingen zijn:

• Onvoldoende statische kenmerken: Bestaande pipelines vertrouwen op handmatig ontworpen, statische kenmerken (bijv. tabelgrootte, aantal rijen). Deze vangen dynamische runtime-effecten niet goed op, zoals partition pruning (waarbij minder data wordt gescand dan de totale tabelgrootte), data skew (onevenwichtige verdeling van taken) en shuffle amplification.
• Versnipperde data: De signalen die nodig zijn voor nauwkeurige voorspelling zijn verspreid over logs, metadata, job-scripts en configuratiegeschiedenis.
• Hoge engineering- overhead: Het handmatig aanpassen van kenmerken vereist diepe domeinkennis en is traag. Dit leidt tot verouderde modellen wanneer werklasten veranderen (drift), wat resulteert in inefficiënte resource-toewijzing en hogere kosten.

2. Methodologie: LeJOT-AutoML

LeJOT-AutoML is een agent-gedreven AutoML-framework dat Large Language Models (LLM's) integreert in de volledige ML-leefcyclus om een dynamisch, zelfverbeterend systeem te creëren.

Architectuur en Componenten:
Het systeem werkt in twee fasen: een trainingsfase (automatisch genereren en valideren) en een inference-fase (realtime voorspelling). De kern bestaat uit drie LLM-agenten die samenwerken via een Model Context Protocol (MCP) toolchain:

1. Feature Analyzer Agent (FAA):
   • Gebruikt Retrieval-Augmented Generation (RAG) om kennis op te halen uit een domein-database (Spark SQL-praktijken, platformkennis).
   • Analyseert job-artifacten (logs, scripts, metadata) en stelt een lijst met kandidaat-kenmerken op die corresponderen met waarneembare runtime-signalen.
2. Feature Extraction Agent (FExA):
   • Voert de plannen van de FAA uit via de MCP-toolchain.
   • Gebruikt tools zoals log-parsers, metadata-query's en een read-only SQL-sandbox om zowel statische als runtime-afgeleide kenmerken te materialiseren.
   • Voert normalisatie en data-kwaliteitscontroles uit.
3. Feature Evaluation Agent (FEvA):
   • Evalueert de kwaliteit van de kenmerken (dekking, stabiliteit) en de prestaties van het model.
   • Geeft feedback voor iteratieve verbetering van de pipeline.

Veiligheidsgates:
Voordat code wordt uitgevoerd, passeert deze twee strenge controles:

• Code-completion checker: Verifieert syntaxis en dat alle benodigde libraries beschikbaar zijn.
• Data-leakage checker: Zorgt ervoor dat kenmerken alleen gebaseerd zijn op informatie die vooraf bekend is (bijv. scripts, historische logs) en geen informatie lekt uit de uitvoeringstijd zelf (label leakage).

Mathematische Formulering:
Het systeem optimaliseert een kostenfunctie die een afweging maakt tussen voorspellingsnauwkeurigheid en de extractie-latentie (budget $B$), terwijl het voldoet aan veiligheidseisen ($gcc$ en $gdl$ predicates).

3. Belangrijkste Bijdragen

• LLM-gedreven AutoML-pipeline: Een end-to-end systeem dat agents gebruikt voor analyse, tool-uitvoering, kenmerkextractie, validatie en modelselectie, waardoor snelle hertraining mogelijk is.
• Agent-tool samenwerking via MCP: Combinatie van LLM-planning met tool-gebaseerde uitvoering om dynamische kenmerken te extraheren die met puur statische analyse ontoegankelijk zijn.
• Iteratieve evaluatie met veiligheid: Een feedbacklus met geautomatiseerde veiligheidschecks die de betrouwbaarheid garandeert en de pipeline continu verfijnt tot aan vooraf gedefinieerde criteria.

4. Resultaten

De prestaties zijn getest op enterprise Databricks-werklasten en vergeleken met handmatige feature engineering:

• Kenmerkdiversiteit: LeJOT-AutoML genereert >200 kenmerken (inclusief log-profielen, tijdreeksdata en driver-node geschiedenis), terwijl handmatige engineering slechts ~40 kenmerken produceert (voornamelijk gebaseerd op configuratiegeschiedenis).
• Snelheid: De cyclus voor feature-engineering en evaluatie is gereduceerd van weken naar 20–30 minuten.
• Voorspellingsnauwkeurigheid:
  • Handmatig: $R^2 = 0.91$, MAPE = 19.49%.
  • AutoML: $R^2 = 0.81$, MAPE = 20.13%.
  • Opmerking: Hoewel de handmatige methode iets nauwkeuriger is, levert AutoML een zeer concurrerende prestatie tegen een fractie van de ontwikkelkosten.
• Kostenefficiëntie: Geïntegreerd in de LeJOT-pipeline resulteert AutoML in 19.01% kostenbesparing door verbeterde orchestration. Handmatige methoden behalen 27.94% besparing, maar vereisen veel meer menselijke inspanning.
• Iteratieve verbetering: Gedurende drie iteraties verbeterde de AutoML-pipeline de $R^2$ van 0.61 naar 0.81 en verlaagde de MAE van 247.95 naar 145.64 seconden.

5. Betekenis en Conclusie

LeJOT-AutoML demonstreert dat LLM-gebaseerde agenten complexe, domeinspecifieke ML-taken kunnen automatiseren die traditioneel veel menselijke expertise vereisten.

• Schaalbaarheid: Het systeem maakt het mogelijk om continu te leren van drift in werklasten zonder lange engineering-cycli.
• Runtime-inzicht: Het slaagt erin om "verborgen" runtime-signalen (zoals shuffle-omvang na pruning) te vangen die statische metadata niet kan bieden.
• Praktische toepasbaarheid: Hoewel handmatige engineering momenteel nog iets beter generaliseert over verschillende hardware-configuraties (waarschijnlijk door betere toegang tot configuratiegeschiedenis in de huidige implementatie), biedt LeJOT-AutoML een schaalbaar, onderhoudsarm alternatief dat aanzienlijke kostenbesparingen realiseert.

Toekomstig werk richt zich op het verbeteren van de "resource awareness" door rijkere configuratie- en runtime-indicatoren op te nemen om de nauwkeurigheidskloof met handmatige methoden verder te dichten.