필사 모드: MLflow 완벽 가이드: 실험 추적부터 Model Registry, 프로덕션 배포까지

# pip 설치
pip install mlflow

# 추가 프레임워크 지원
pip install mlflow[extras]  # sklearn, tensorflow, pytorch 등

# 서버 시작 (로컬)
mlflow server --host 0.0.0.0 --port 5000

# PostgreSQL + S3 백엔드로 프로덕션 서버
mlflow server \
  --backend-store-uri postgresql://mlflow:password@localhost:5432/mlflow \
  --default-artifact-root s3://mlflow-artifacts/ \
  --host 0.0.0.0 --port 5000

# docker-compose.yml
services:
  mlflow:
    image: ghcr.io/mlflow/mlflow:v3.15.1
    ports:
      - '5000:5000'
    environment:
      - MLFLOW_BACKEND_STORE_URI=postgresql://mlflow:password@postgres:5432/mlflow
      - MLFLOW_DEFAULT_ARTIFACT_ROOT=s3://mlflow-artifacts/
      - AWS_ACCESS_KEY_ID=${AWS_ACCESS_KEY_ID}
      - AWS_SECRET_ACCESS_KEY=${AWS_SECRET_ACCESS_KEY}
    command: >
      mlflow server
      --backend-store-uri postgresql://mlflow:password@postgres:5432/mlflow
      --default-artifact-root s3://mlflow-artifacts/
      --host 0.0.0.0 --port 5000
    depends_on:
      - postgres

  postgres:
    image: postgres:16
    environment:
      POSTGRES_USER: mlflow
      POSTGRES_PASSWORD: password
      POSTGRES_DB: mlflow
    volumes:
      - pgdata:/var/lib/postgresql/data

volumes:
  pgdata:

# 인자 없이 — MLflow 3.7.0 이상이면 ./mlflow.db 가 자동 생성된다
mlflow server

# 저장 위치를 명시하고 싶을 때
mlflow server \
  --backend-store-uri sqlite:///mlflow.db \
  --artifacts-destination ./mlartifacts \
  --host 127.0.0.1 --port 5000

# 클라이언트 쪽 (스크립트를 실행하는 셸)
export MLFLOW_TRACKING_URI=http://127.0.0.1:5000
export MLFLOW_EXPERIMENT_NAME=iris-classification

import mlflow
import mlflow.sklearn
from sklearn.ensemble import RandomForestClassifier
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, f1_score, precision_score

# 트래킹 서버 설정
mlflow.set_tracking_uri("http://localhost:5000")

# 실험 생성/설정
mlflow.set_experiment("iris-classification")

# 데이터 준비
X, y = load_iris(return_X_y=True)
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# 실험 실행
with mlflow.start_run(run_name="rf-baseline"):
    # 파라미터 기록
    params = {
        "n_estimators": 100,
        "max_depth": 5,
        "min_samples_split": 2,
        "random_state": 42
    }
    mlflow.log_params(params)

    # 모델 학습
    model = RandomForestClassifier(**params)
    model.fit(X_train, y_train)

    # 예측 및 메트릭
    y_pred = model.predict(X_test)
    metrics = {
        "accuracy": accuracy_score(y_test, y_pred),
        "f1_macro": f1_score(y_test, y_pred, average="macro"),
        "precision_macro": precision_score(y_test, y_pred, average="macro")
    }
    mlflow.log_metrics(metrics)

    # 태그
    mlflow.set_tag("model_type", "random_forest")
    mlflow.set_tag("dataset", "iris")

    # 모델 저장 (MLflow 3부터 artifact_path= 대신 name=)
    mlflow.sklearn.log_model(
        model,
        name="model",
        registered_model_name="iris-classifier"
    )

    # 커스텀 아티팩트 (그래프, 보고서 등)
    import matplotlib.pyplot as plt
    from sklearn.metrics import confusion_matrix, ConfusionMatrixDisplay

    cm = confusion_matrix(y_test, y_pred)
    fig, ax = plt.subplots()
    ConfusionMatrixDisplay(cm).plot(ax=ax)
    fig.savefig("confusion_matrix.png")
    mlflow.log_artifact("confusion_matrix.png")

    print(f"Run ID: {mlflow.active_run().info.run_id}")
    print(f"Metrics: {metrics}")

import optuna
import mlflow

def objective(trial):
    params = {
        "n_estimators": trial.suggest_int("n_estimators", 50, 500),
        "max_depth": trial.suggest_int("max_depth", 2, 20),
        "min_samples_split": trial.suggest_int("min_samples_split", 2, 10),
        "min_samples_leaf": trial.suggest_int("min_samples_leaf", 1, 5),
    }

    with mlflow.start_run(nested=True, run_name=f"trial-{trial.number}"):
        mlflow.log_params(params)

        model = RandomForestClassifier(**params, random_state=42)
        model.fit(X_train, y_train)
        y_pred = model.predict(X_test)

        accuracy = accuracy_score(y_test, y_pred)
        mlflow.log_metric("accuracy", accuracy)

        return accuracy

# Optuna 스터디 실행
with mlflow.start_run(run_name="hyperparameter-tuning"):
    study = optuna.create_study(direction="maximize")
    study.optimize(objective, n_trials=50)

    # 최적 결과 기록
    mlflow.log_params(study.best_params)
    mlflow.log_metric("best_accuracy", study.best_value)
    mlflow.set_tag("best_trial", study.best_trial.number)

import torch
import torch.nn as nn
import mlflow.pytorch

class SimpleNet(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        super().__init__()
        self.fc1 = nn.Linear(input_dim, hidden_dim)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(hidden_dim, output_dim)

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

with mlflow.start_run(run_name="pytorch-model"):
    model = SimpleNet(4, 32, 3)
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
    criterion = nn.CrossEntropyLoss()

    mlflow.log_params({
        "hidden_dim": 32,
        "learning_rate": 0.001,
        "optimizer": "Adam",
        "epochs": 100
    })

    for epoch in range(100):
        # 학습 로직...
        loss = criterion(model(X_tensor), y_tensor)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        # 에폭별 메트릭 기록
        mlflow.log_metric("train_loss", loss.item(), step=epoch)

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

import mlflow
from mlflow.models import infer_signature

signature = infer_signature(model_input=X_train, model_output=model.predict(X_train))

# MLflow 3: start_run() 없이도 로깅되고, ModelInfo 를 돌려받는다
info = mlflow.sklearn.log_model(
    model,
    name="model",                 # artifact_path= 는 deprecated
    signature=signature,
    input_example=X_train[:2],
    registered_model_name="iris-classifier",
)

print(info.model_id)
loaded = mlflow.pyfunc.load_model(f"models:/{info.model_id}")

run_id  status  start_time  params.n_estimators  params.max_depth  metrics.accuracy  tags.model_type

from mlflow import MlflowClient

client = MlflowClient()

# 모델 등록 (log_model에서 registered_model_name 사용 시 자동 등록)
# 또는 수동 등록:
result = client.create_registered_model(
    name="iris-classifier",
    description="Iris 꽃 분류 모델"
)

# 특정 실행의 모델을 버전으로 등록
model_version = client.create_model_version(
    name="iris-classifier",
    source=f"runs:/{run_id}/model",
    run_id=run_id,
    description="RandomForest baseline v1"
)

print(f"Model Version: {model_version.version}")

# MLflow 2.x에서는 Alias 사용 (Stage는 deprecated)
client = MlflowClient()

# 프로덕션 alias 설정
client.set_registered_model_alias(
    name="iris-classifier",
    alias="champion",
    version=3
)

# 도전자 모델 설정
client.set_registered_model_alias(
    name="iris-classifier",
    alias="challenger",
    version=5
)

# Alias로 모델 로드
champion_model = mlflow.pyfunc.load_model("models:/iris-classifier@champion")
challenger_model = mlflow.pyfunc.load_model("models:/iris-classifier@challenger")

# A/B 테스트
champion_pred = champion_model.predict(X_test)
challenger_pred = challenger_model.predict(X_test)

print(f"Champion accuracy: {accuracy_score(y_test, champion_pred)}")
print(f"Challenger accuracy: {accuracy_score(y_test, challenger_pred)}")

# 모델 버전에 태그 추가
client.set_model_version_tag(
    name="iris-classifier",
    version=3,
    key="validation_status",
    value="approved"
)

client.set_model_version_tag(
    name="iris-classifier",
    version=3,
    key="approved_by",
    value="data-science-lead"
)

# 태그로 모델 검색
from mlflow import search_model_versions

approved_versions = search_model_versions(
    "name='iris-classifier' AND tag.validation_status='approved'"
)

# 로컬 REST API 서빙
mlflow models serve \
  -m "models:/iris-classifier@champion" \
  --port 8080 \
  --env-manager local

# 테스트 요청
curl -X POST http://localhost:8080/invocations \
  -H "Content-Type: application/json" \
  -d '{"inputs": [[5.1, 3.5, 1.4, 0.2]]}'

mlflow models serve \
  -m "models:/iris-classifier@champion" \
  --host 0.0.0.0 --port 8080 \
  --workers 4 \
  --env-manager local

from fastapi import FastAPI
import mlflow.pyfunc
import numpy as np

app = FastAPI()

# 모델 로드 (서버 시작 시 1회)
model = mlflow.pyfunc.load_model("models:/iris-classifier@champion")

@app.post("/predict")
async def predict(features: list[list[float]]):
    predictions = model.predict(np.array(features))
    return {
        "predictions": predictions.tolist(),
        "model_version": "champion"
    }

@app.get("/health")
async def health():
    return {"status": "healthy", "model": "iris-classifier@champion"}

# 실험 검색 (CLI)
mlflow runs list --experiment-id 1

# 메트릭 기반 검색
mlflow runs list \
  --experiment-id 1 \
  --filter "metrics.accuracy > 0.95" \
  --order-by "metrics.accuracy DESC"

import mlflow
import pandas as pd

# 실험의 모든 실행 조회
runs = mlflow.search_runs(
    experiment_ids=["1"],
    filter_string="metrics.accuracy > 0.9",
    order_by=["metrics.accuracy DESC"],
    max_results=10
)

# DataFrame으로 분석
print(runs[["run_id", "params.n_estimators", "params.max_depth", "metrics.accuracy"]])

# 최적 실행 찾기
best_run = runs.iloc[0]
print(f"Best run: {best_run.run_id}, Accuracy: {best_run['metrics.accuracy']}")

runs = mlflow.search_runs(
    experiment_ids=["1"],
    filter_string='metrics.accuracy > 0.72 AND metrics.loss <= 0.15',
    order_by=["metrics.accuracy DESC"],
)

# 파라미터는 문자열 — 따옴표가 없으면 원하는 결과가 나오지 않는다
exact = mlflow.search_runs(filter_string='params.n_estimators = "100"')

# OR 가 없으므로 두 번 조회해서 붙인다
import pandas as pd

merged = pd.concat([
    mlflow.search_runs(filter_string='tags.model_type = "random_forest"'),
    mlflow.search_runs(filter_string='tags.model_type = "xgboost"'),
]).drop_duplicates(subset="run_id")

□ 백엔드 스토어를 PostgreSQL/MySQL로 설정
□ 아티팩트 스토어를 S3/GCS/MinIO로 설정
□ 인증/권한 설정 (OIDC, Basic Auth)
□ 자동 실험 기록 (autolog) 설정
□ Model Registry alias 규칙 정립
□ CI/CD에서 모델 검증 자동화
□ 모델 서빙 헬스체크 설정
□ 실험 정리 정책 (오래된 실행 아카이브)

# autolog 가 Run 을 닫아 주는 경우와 아닌 경우
mlflow.autolog()

model.fit(X_train, y_train)          # 활성 Run 없음 -> 만들고 닫아 준다

with mlflow.start_run(run_name="manual"):
    model.fit(X_train, y_train)      # 여기에 기록되고, 블록을 나갈 때 닫힌다

run = mlflow.start_run(run_name="leaky")
model.fit(X_train, y_train)          # 기록은 되지만 닫히지 않는다
mlflow.end_run()                     # 직접 닫아야 한다