MLOps의 핵심 엔진, MLflow가 Databricks 플랫폼에 어떻게 통합되어 있는지 깊이 파헤쳐봅니다.
Databricks 내 MLflow 대시보드 — 실험 추적부터 모델 배포까지 한눈에
🔍 MLflow란 무엇인가?
MLflow는 머신러닝 라이프사이클 전반을 관리하기 위해 Databricks가 만든 오픈소스 플랫폼입니다. 2018년 Databricks가 처음 공개한 이후, 현재는 ML 실험 추적의 사실상 표준(de facto standard)으로 자리잡았습니다.
MLflow는 크게 네 가지 핵심 컴포넌트로 구성됩니다.
| 컴포넌트 | 역할 |
|---|---|
| MLflow Tracking | 실험 파라미터, 메트릭, 아티팩트 기록 |
| MLflow Projects | 재현 가능한 ML 코드 패키징 |
| MLflow Models | 다양한 배포 환경을 위한 모델 패키징 |
| MLflow Model Registry | 모델 버전 관리 및 스테이지 전환 |
⚙️ Databricks에서 MLflow는 어떻게 통합되어 있나?
Databricks는 MLflow를 플랫폼에 네이티브로 내장시켰습니다. 별도 설치나 설정 없이 노트북 환경에서 바로 import mlflow를 실행하면 곧바로 사용할 수 있습니다.
1. 완전 관리형(Managed) MLflow
일반적인 오픈소스 MLflow를 직접 운영하면 Tracking Server, Backend Store, Artifact Store를 모두 직접 구성해야 합니다. 하지만 Databricks에서는 이 모든 인프라를 플랫폼이 자동으로 관리합니다.
import mlflow
# Databricks에서는 이게 전부입니다.
# Tracking URI 설정이 필요 없어요!
with mlflow.start_run():
mlflow.log_param("learning_rate", 0.01)
mlflow.log_metric("accuracy", 0.95)
mlflow.sklearn.log_model(model, "model")Databricks 환경에서는 DATABRICKS_HOST와 DATABRICKS_TOKEN을 기반으로 Tracking URI가 자동 설정됩니다.
2. Unity Catalog와의 통합
MLflow Model Registry — 모델의 전체 라이프사이클을 관리
Databricks의 최신 아키텍처에서 MLflow Model Registry는 Unity Catalog와 깊이 통합되어 있습니다.
import mlflow
# Unity Catalog 기반 Model Registry 사용
mlflow.set_registry_uri("databricks-uc")
# 모델 등록
mlflow.register_model(
model_uri=f"runs:/{run_id}/model",
name="main.ml_project.my_classifier"
# ^^^^ ^^^^^^^^^^^ ^^^^^^^^^^^^
# catalog schema model name
)Unity Catalog를 통해 모델에 대해서도 테이블과 동일한 수준의 거버넌스가 가능해집니다. 접근 권한, 데이터 리니지, 태그 관리가 하나의 인터페이스에서 이루어집니다.
3. 실험(Experiment) 구조
Databricks에서 MLflow 실험은 두 가지 형태로 존재합니다.
노트북 실험(Notebook Experiment)
- 각 노트북에 자동으로 연결된 실험
- 노트북과 동일한 경로에 위치
- 노트북 삭제 시 함께 삭제됨
워크스페이스 실험(Workspace Experiment)
- 팀이 공유하는 독립적인 실험
- 여러 노트북, 여러 Job에서 로그 가능
- 프로덕션 ML 프로젝트에 권장
# 워크스페이스 실험 명시적 지정
mlflow.set_experiment("/Users/alice@company.com/fraud-detection-v2")
with mlflow.start_run(run_name="xgboost-baseline"):
# 학습 코드
...4. MLflow와 Databricks Jobs 연동
MLflow Tracking UI — 수백 번의 실험을 체계적으로 비교
Databricks Job으로 ML 파이프라인을 실행할 때, MLflow는 Job의 run context를 자동으로 인식합니다.
# Job 환경에서 자동으로 parent run이 설정됨
with mlflow.start_run() as parent_run:
# 데이터 전처리 단계
with mlflow.start_run(nested=True, run_name="preprocessing"):
mlflow.log_param("feature_count", 128)
# 모델 학습 단계
with mlflow.start_run(nested=True, run_name="training"):
mlflow.log_metric("train_loss", 0.042)
mlflow.log_metric("val_loss", 0.051)중첩 실행(Nested Runs)을 활용하면 복잡한 ML 파이프라인의 각 단계를 계층적으로 추적할 수 있습니다.
🚀 MLflow Model Serving (Mosaic AI)
Databricks는 최근 MLflow 모델 서빙 기능을 Mosaic AI Model Serving으로 통합했습니다. Model Registry에 등록된 모델을 몇 번의 클릭만으로 REST API 엔드포인트로 배포할 수 있습니다.
import requests
# 배포된 모델 엔드포인트 호출
response = requests.post(
url="https://<workspace>.azuredatabricks.net/serving-endpoints/my-model/invocations",
headers={"Authorization": f"Bearer {token}"},
json={"inputs": [[1.2, 3.4, 5.6, 7.8]]}
)
print(response.json())
# {'predictions': [0]}Auto-scaling, A/B 테스트, 모델 버전 트래픽 분산까지 UI에서 바로 설정 가능합니다.
📊 MLflow 아키텍처 전체 그림
MLflow 전체 아키텍처 — 학습부터 배포까지의 흐름
Databricks에서 MLflow의 전체 데이터 흐름을 정리하면 다음과 같습니다.
[노트북 / Job / 외부 환경]
|
| mlflow.log_*()
▼
[Databricks MLflow Tracking Server] ← 완전 관리형
|
┌────┴────┐
▼ ▼
[메타데이터] [아티팩트]
(파라미터, (모델 파일,
메트릭) 이미지, CSV)
| |
▼ ▼
[Delta Lake] [DBFS / S3 / ADLS]
|
▼
[MLflow Model Registry (Unity Catalog)]
|
▼
[Mosaic AI Model Serving]💡 Databricks MLflow 사용 시 팁
1. autolog 적극 활용하기
mlflow.autolog() # sklearn, xgboost, pytorch 등 자동 로깅2. 실험 비교를 위한 태그 전략
mlflow.set_tags({
"team": "fraud-detection",
"data_version": "v3.1",
"model_type": "ensemble"
})3. 커스텀 아티팩트 저장
# 혼동 행렬, SHAP 플롯 등 저장
mlflow.log_figure(fig, "confusion_matrix.png")
mlflow.log_dict(feature_importance, "feature_importance.json")🎯 정리
Databricks의 MLflow 통합은 단순한 "MLflow 설치"가 아닙니다. 인프라 관리를 없애고, Unity Catalog 거버넌스를 부여하며, Job과 서빙까지 연결되는 완전한 MLOps 플랫폼을 제공합니다.
오픈소스 MLflow의 유연성을 그대로 유지하면서도, 엔터프라이즈 환경에서 필요한 보안·거버넌스·확장성을 모두 갖춘 것이 Databricks MLflow의 핵심 가치입니다.
이 글이 도움이 되었다면 공유해주세요! 궁금한 점은 댓글로 남겨주세요 🙌
