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

이 논문은 Databricks 작업 실행 시간 예측을 위해 LLM 에이전트와 검색 증강 생성 (RAG) 을 활용한 LeJOT-AutoML 프레임워크를 제안하여, 기존 수동 피처 엔지니어링의 한계를 극복하고 20~30 분 내에 200 개 이상의 동적 피처를 생성함으로써 19.01% 의 비용 절감 효과를 달성했다고 요약할 수 있습니다.

핵심

🍳 비유: 레스토랑의 '요리 시간 예측' 문제

대형 레스토랑 (Databricks) 이 매일 수천 개의 주문 (데이터 작업) 을 받습니다. 각 주문을 처리하는 데 걸리는 시간과 비용은 요리사 (컴퓨터 서버) 의 능력, 재료의 양, 그리고 주방의 혼잡도에 따라 달라집니다.

기존의 문제점 (수동 방식):
이전에는 주방장 (엔지니어) 이 직접 과거 기록을 보며 "아, 이 요리법은 보통 30 분 걸리겠지"라고 직접 계산하고 메모했습니다.

• 단점 1: 요리사가 갑자기 바뀌거나, 재료가 갑자기 많아지면 (데이터가 변하면) 계산이 틀립니다.
• 단점 2: 주방의 숨겨진 문제 (예: 특정 반죽이 너무 끈적해서 시간이 걸림) 는 기록에 안 나오기 때문에 놓치기 쉽습니다.
• 단점 3: 새로운 메뉴가 나오면 주방장이 다시 수기로 계산하는 데 한 달이나 걸립니다.

🤖 해결책: LeJOT-AutoML (AI 요리사 팀)

이 논문은 **LLM(거대 언어 모델) 이 주도하는 '자동 요리 팀'**을 제안합니다. 이 팀은 세 명의 AI 요리사 (에이전트) 로 구성되어 있습니다.

1. 메뉴 분석가 (Feature Analyzer Agent)

• 역할: "오늘 어떤 재료가 들어갈까? 어떤 문제가 생길까?"를 고민합니다.
• 비유: 이 요리사는 레시피 책 (지식 데이터베이스) 을 뒤적이며, "아, 이 요리는 재료가 섞일 때 (Shuffle) 시간이 더 걸릴 거야"라고 새로운 예측 포인트를 찾아냅니다.
• 특징: 사람이 생각하지 못했던 '숨은 신호' (예: 데이터가 한쪽으로 쏠리는 현상) 를 찾아냅니다.

2. 재료 수집가 (Feature Extraction Agent)

• 역할: 분석가가 찾은 포인트를 실제로 확인합니다.
• 비유: 이 요리사는 주방의 모든 구석 (로그, 데이터베이스, 작업 기록) 을 돌아다니며 실제 데이터를 가져옵니다.
• 안전 장치: "이 재료가 아직 안 익었으면 (미래 데이터) 쓸 수 없어!"라는 규칙을 지키기 위해, 데이터 유출 방지 검사를 거칩니다. (과거의 결과만 보고 미래를 예측해야 하니까요.)

3. 맛 평가자 (Feature Evaluation Agent)

• 역할: "이 재료가 정말 맛을 좋게 할까?"를 테스트합니다.
• 비유: 가져온 재료로 요리를 해보고, "이건 필요 없어"라고 버리거나 "이걸 더 넣으면 더 정확해져"라고 수정을 요청합니다. 이 과정이 몇 번 반복되면 요리가 완벽해집니다.

---

🚀 이 시스템이 가져온 변화

이 AI 팀이 도입된 후 어떤 일이 일어났을까요?

1. 시간 단축 (한 달 → 30 분):

   • 예전에는 새로운 예측 모델을 만들려면 엔지니어가 한 달을 꼬박 썼습니다.
   • 이제는 AI 가 20~30 분 만에 200 개 이상의 새로운 예측 요소 (재료) 를 찾아내고 모델을 완성합니다.

2. 정확도 향상 (숨은 신호 포착):

   • 사람이 만든 모델은 '서버 크기' 같은 눈에 보이는 것만 봤다면, AI 는 '데이터가 섞일 때 생기는 마찰' 같은 눈에 안 보이는 복잡한 신호까지 찾아냅니다.
   • (참고: 아직은 사람이 만든 모델이 하드웨어 변화에 더 강하지만, AI 는 빠르게 따라잡고 있습니다.)

3. 비용 절감 (19% 절약):

   • 정확한 예측 덕분에, 레스토랑은 불필요하게 큰 주방 (비싼 서버) 을 빌리지 않아도 됩니다.
   • 실제 실험에서 약 19% 의 비용을 아꼈습니다.

💡 핵심 요약

이 논문은 **"복잡한 데이터 작업을 예측할 때, 사람이 일일이 계산할 필요 없이 AI 가 스스로 지식을 찾아내고, 실험하고, 수정하게 하라"**는 것입니다.

마치 스마트한 주방장이 매일 새로운 레시피를 스스로 개발하고, 실수를 바로잡으며, 레스토랑의 비용을 자동으로 최적화하는 것과 같습니다. 덕분에 기업은 더 빠르고, 더 저렴하게, 더 똑똑하게 데이터를 처리할 수 있게 되었습니다.

기술 요약

1. 문제 정의 (Problem Statement)

Databricks 기반의 작업 오케스트레이션 시스템 (LeJOT) 은 지연 시간 (latency) 과 의존성 제약 조건을 만족하면서 저비용 컴퓨팅 구성을 선택하여 클라우드 비용을 절감합니다. 이를 위해서는 이질적인 인스턴스 유형과 비정상적인 (non-stationary) 런타임 환경 하에서 작업 실행 시간을 정확하게 예측하는 것이 필수적입니다.

기존의 파이프라인은 다음과 같은 한계를 가지고 있습니다:

• 정적 수동 피처 엔지니어링: 수동으로 설계된 정적 피처는 파티션 프루닝 (partition pruning), 데이터 스웨 (skew), 셔플 증폭 (shuffle amplification) 등 런타임에 발생하는 동적 신호를 포착하지 못합니다.
• 데이터 단편화: 예측에 필요한 신호는 로그, 메타데이터, 작업 스크립트 등 다양한 소스에 흩어져 있어 통합이 어렵습니다.
• 높은 엔지니어링 오버헤드: 수동 피처 설계와 모델 업데이트 주기가 길어 (수주~수개월), 작업 부하의 변화 (drift) 에 신속하게 대응하지 못합니다.

2. 방법론 (Methodology)

저자들은 LeJOT-AutoML을 제안했습니다. 이는 대규모 언어 모델 (LLM) 에이전트와 모델 컨텍스트 프로토콜 (MCP) 도구 체인을 결합하여 ML 라이프사이클 전반을 자동화하는 에이전트 기반 프레임워크입니다.

핵심 아키텍처 및 구성 요소

1. 피처 분석 에이전트 (Feature Analyzer Agent, FAA):
   • RAG(검색 증강 생성) 를 통해 도메인 지식 베이스 (Spark SQL 관행, 플랫폼 지식 등) 를 참조합니다.
   • 실행 로그, 메타데이터, 작업 스크립트 등을 분석하여 후보 피처 템플릿을 제안합니다.
2. 피처 추출 에이전트 (Feature Extraction Agent, FExA):
   • MCP 도구 체인 (로그 파서, 메타데이터 쿼리, 읽기 전용 SQL 샌드박스) 을 활용하여 제안된 피처를 실제로 추출하고 값을 구체화 (materialize) 합니다.
   • 정적 정보 (스크립트, 메타데이터) 와 런타임 정보 (로그 기반 신호) 를 모두 결합합니다.
3. 피처 평가 에이전트 (Feature Evaluation Agent, FEvA):
   • 피처의 품질 (커버리지, 편향성), 유용성 (중요도, 중복성), 모델 성능을 평가합니다.
   • 피드백을 통해 파이프라인을 반복적으로 개선합니다.
4. 안전 게이트 (Safety Gates):
   • 코드 완성도 체크: 생성된 코드가 실행 가능한지 확인합니다.
   • 데이터 누출 (Data Leakage) 체크: 피처 추출 시 레이블 (실행 시간) 정보가 유출되지 않도록 검증합니다.

작동 원리

• 학습 단계: FAA 가 피처를 설계하고, FExA 가 도구를 호출하여 데이터를 추출한 후, 안전 게이트를 통과하면 모델을 학습시킵니다. FEvA 가 성능을 평가하고 피드백을 주어 과정을 반복합니다.
• 추론 단계: 새로운 작업이 들어오면 학습된 피처 템플릿을 재사용하여 실시간으로 피처를 추출하고, 예측된 실행 시간을 기반으로 LeJOT 이 최적의 비용 구성을 선택합니다.

3. 주요 기여 (Key Contributions)

1. LLM 기반 엔터프라이즈용 AutoML 파이프라인: 분석, 도구 호출, 피처 추출, 검증, 학습, 모델 선택 전 과정에 LLM 에이전트를 내장하여 빠른 재학습과 추론 시 피처 구체화를 가능하게 함.
2. MCP 를 통한 에이전트 - 도구 협업 피처 추출: LLM 의 계획 능력과 도구의 실행/검증 능력을 결합하여 정적 분석만으로는 접근 불가능한 동적 런타임 피처를 효율적으로 추출함.
3. 안전 게이트가 포함된 반복적 평가 루프: 코드 완성도 및 데이터 누출 검사를 통해 신뢰성을 확보하고, 피드백 기반의 평가 에이전트를 통해 사전 정의된 기준이 충족될 때까지 파이프라인을 지속적으로 정제함.

4. 실험 결과 (Results)

엔터프라이즈급 Databricks 워크로드를 대상으로 수동 피처 엔지니어링과 LeJOT-AutoML 을 비교 평가했습니다.

• 피처 다양성: AutoML 은 200 개 이상의 피처 (로그 프로파일링, 시계열 데이터, 노드 히스토리 등) 를 생성한 반면, 수동 방식은 약 40 개의 피처 (주로 노드 구성 히스토리 기반) 에 그쳤습니다.
• 예측 정확도:
  • 수동 방식이 R² 0.91 로 AutoML(R² 0.81) 보다 정확도가 약간 높았습니다. 이는 수동 방식이 하드웨어 구성 변경에 대한 풍부한 컨텍스트를 더 잘 반영하기 때문입니다.
  • 그러나 AutoML 은 MAPE 20.13%, MAE 123.29 로 경쟁력 있는 성능을 보였습니다.
• 개발 시간 및 효율성:
  • 수동 방식: 약 1 개월 소요.
  • AutoML: 3 회 반복 (20~30 분) 으로 완료. 피처 엔지니어링 및 평가 루프를 수주에서 수분 단위로 단축했습니다.
• 비용 절감 효과:
  • LeJOT 파이프라인에 통합된 AutoML 은 19.01% 의 비용 절감을 달성했습니다. (수동 ML 은 27.94% 절감이었으나, AutoML 은 유지보수 비용이 훨씬 낮음).
  • 인스턴스 유형 업그레이드 (Photon 등) 시 수동 모델의 예측 편차가 큰 반면, AutoML 은 더 안정적인 반응을 보였습니다.

5. 의의 및 결론 (Significance)

• 동적 신호 포착: 정적 분석으로는 불가능했던 런타임 신호 (셔플 증폭, 스웨 등) 를 LLM 에이전트가 도구를 통해 직접 추출함으로써 예측 정확도를 높였습니다.
• 지속 가능한 적응: 작업 부하의 변화 (drift) 에 따라 모델을 수주 단위가 아닌 분 단위로 업데이트할 수 있어, 동적인 클라우드 환경에서 모델의 노후화를 방지합니다.
• 실용성: 완벽한 정확도 (수동 방식 대비) 는 약간 떨어질 수 있으나, 개발 비용과 시간을 획기적으로 줄이면서도 상당한 비용 절감 효과를 제공하여 엔터프라이즈 환경에 실용적인 대안이 됩니다.

향후 과제: 리소스 인식 (Resource Awareness) 을 높이기 위해 구성 변경 로그 및 인스턴스 사양/가격 정보를 MCP 도구 인터페이스에 더 풍부하게 통합하는 방향으로 발전할 예정입니다.