MLflow (MLflow) | KAITRUST AI 백과사전

📖 상세 설명

MLflow는 Databricks에서 개발한 오픈소스 ML 라이프사이클 관리 플랫폼입니다. 실험 추적(Tracking), 프로젝트 패키징(Projects), 모델 관리(Models), 모델 레지스트리(Registry) 4가지 핵심 컴포넌트를 제공하여 ML 워크플로우 전체를 관리합니다.

2018년 Matei Zaharia(Spark 창시자)가 이끄는 Databricks 팀이 발표했습니다. "ML 실험의 재현 불가능 문제"를 해결하기 위해 시작되었으며, 현재 GitHub 스타 18,000개 이상으로 가장 널리 사용되는 MLOps 도구 중 하나입니다.

MLflow Tracking은 파라미터, 메트릭, 아티팩트를 자동 기록하고 UI로 시각화합니다. Model Registry는 모델 버전 관리, 스테이지(Staging/Production) 전환, 승인 워크플로우를 제공합니다. MLflow는 PyTorch, TensorFlow, Scikit-learn 등 모든 주요 프레임워크와 호환됩니다.

실무에서 MLflow는 실험 관리의 표준 도구로 자리잡았습니다. 수백 개의 실험을 체계적으로 비교하고, 최적 모델을 원클릭으로 배포하며, REST API로 서빙까지 가능합니다. Databricks, AWS SageMaker, Azure ML 등과 통합되어 엔터프라이즈 환경에서도 활용됩니다.

💻 코드 예제

import mlflow
import mlflow.pytorch
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset

# MLflow 설정
mlflow.set_tracking_uri("sqlite:///mlflow.db")
mlflow.set_experiment("pytorch_classifier")

# 간단한 모델 정의
class SimpleNet(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        super().__init__()
        self.fc1 = nn.Linear(input_dim, hidden_dim)
        self.fc2 = nn.Linear(hidden_dim, output_dim)
        self.relu = nn.ReLU()

    def forward(self, x):
        return self.fc2(self.relu(self.fc1(x)))

# 하이퍼파라미터
config = {
    "input_dim": 20,
    "hidden_dim": 64,
    "output_dim": 2,
    "learning_rate": 0.001,
    "epochs": 10,
    "batch_size": 32
}

# MLflow autolog 활성화 (PyTorch 자동 추적)
mlflow.pytorch.autolog()

with mlflow.start_run(run_name="pytorch_simple_net"):
    # 하이퍼파라미터 로깅
    mlflow.log_params(config)

    # 모델 초기화
    model = SimpleNet(config["input_dim"], config["hidden_dim"], config["output_dim"])
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=config["learning_rate"])

    # 더미 데이터 생성
    X = torch.randn(1000, config["input_dim"])
    y = torch.randint(0, 2, (1000,))
    dataset = TensorDataset(X, y)
    loader = DataLoader(dataset, batch_size=config["batch_size"])

    # 학습 루프
    for epoch in range(config["epochs"]):
        total_loss = 0
        for batch_X, batch_y in loader:
            optimizer.zero_grad()
            outputs = model(batch_X)
            loss = criterion(outputs, batch_y)
            loss.backward()
            optimizer.step()
            total_loss += loss.item()

        avg_loss = total_loss / len(loader)
        mlflow.log_metric("train_loss", avg_loss, step=epoch)
        print(f"Epoch {epoch+1}: Loss = {avg_loss:.4f}")

    # 모델 저장
    mlflow.pytorch.log_model(model, "model")

    # 모델 레지스트리에 등록
    run_id = mlflow.active_run().info.run_id
    model_uri = f"runs:/{run_id}/model"
    mlflow.register_model(model_uri, "SimpleClassifier")

print("MLflow UI: http://localhost:5000")
# mlflow ui --port 5000 으로 실행

📊 성능 & 비용

2025년 1월 기준 MLflow 옵션별 비용입니다.

옵션	가격	특징
MLflow 오픈소스	무료 (Apache 2.0)	셀프호스팅, 전체 기능
Databricks MLflow	Databricks 요금 포함	$0.20/DBU~
AWS SageMaker MLflow	$0.07/시간~	관리형 서버
Azure ML MLflow	Azure ML 요금 포함	통합 환경

셀프호스팅 인프라 예상	비용	비고
EC2 (t3.medium)	$30/월	소규모 팀용 서버
RDS PostgreSQL	$15/월~	메타데이터 저장소
S3 저장소	$0.023/GB	아티팩트 저장

엔터프라이즈 기능 (Databricks): Unity Catalog 통합, Feature Store, Model Serving 자동 스케일링, 거버넌스.

🗣️ 실무에서 이렇게 말하세요

💬 회의에서

"지난 분기 동안 진행한 150개 실험이 모두 MLflow에 기록되어 있습니다. 베스트 모델은 Run ID abc123으로, F1 0.92를 달성했고 현재 Production 스테이지로 배포되어 있습니다."

💬 면접에서

"MLflow의 4가지 컴포넌트 중 Tracking은 실험 기록, Projects는 코드 패키징, Models는 모델 포맷 표준화, Registry는 버전 관리와 배포 승인을 담당합니다. autolog 기능으로 대부분의 프레임워크에서 자동 추적이 가능합니다."

💬 기술 토론에서

"MLflow 모델 서빙은 mlflow models serve 명령으로 바로 REST API 배포가 됩니다. 프로덕션 환경에서는 Seldon Core나 KServe와 연동하거나, MLflow의 Docker 빌드 기능을 써서 Kubernetes에 배포합니다."

⚠️ 흔한 실수 & 주의사항

❌

로컬 파일 시스템에만 저장

기본 설정으로 로컬에만 저장하면 팀 협업이 불가능하고 데이터 손실 위험이 있습니다.

✅

중앙 집중식 저장소 구성

PostgreSQL/MySQL을 백엔드로, S3/GCS를 아티팩트 저장소로 설정하세요. 팀 전체가 하나의 MLflow 서버를 공유해야 실험 결과를 효과적으로 비교할 수 있습니다.

🔗 관련 용어

📚 더 배우기

📄 MLflow 공식 문서 🎓 MLflow GitHub