model_registry.py

  1"""
  2MLflow Model Registry Example
  3=============================
  4
  5MLflow Model Registry를 사용한 모델 버전 관리 예제입니다.
  6
  7실행 방법:
  8    # 먼저 tracking_example.py를 실행하여 모델을 학습/저장한 후
  9    python model_registry.py
 10"""
 11
 12import mlflow
 13from mlflow.tracking import MlflowClient
 14from sklearn.datasets import load_iris
 15from sklearn.model_selection import train_test_split
 16from sklearn.ensemble import RandomForestClassifier
 17from sklearn.metrics import accuracy_score
 18import os
 19
 20# MLflow 설정
 21TRACKING_URI = os.environ.get("MLFLOW_TRACKING_URI", "http://localhost:5000")
 22MODEL_NAME = "iris-classifier"
 23
 24
 25def setup():
 26    """MLflow 설정"""
 27    mlflow.set_tracking_uri(TRACKING_URI)
 28    mlflow.set_experiment("model-registry-demo")
 29    return MlflowClient()
 30
 31
 32def train_and_register_model(client, version_tag: str):
 33    """모델 학습 및 레지스트리 등록"""
 34    # 데이터 준비
 35    iris = load_iris()
 36    X_train, X_test, y_train, y_test = train_test_split(
 37        iris.data, iris.target, test_size=0.2, random_state=42
 38    )
 39
 40    with mlflow.start_run(run_name=f"training-{version_tag}") as run:
 41        # 모델 학습
 42        model = RandomForestClassifier(n_estimators=100, random_state=42)
 43        model.fit(X_train, y_train)
 44
 45        # 평가
 46        accuracy = accuracy_score(y_test, model.predict(X_test))
 47        mlflow.log_metric("accuracy", accuracy)
 48
 49        # 모델 저장 및 등록
 50        mlflow.sklearn.log_model(
 51            model,
 52            "model",
 53            registered_model_name=MODEL_NAME
 54        )
 55
 56        print(f"\n모델 등록 완료: {MODEL_NAME}")
 57        print(f"  Run ID: {run.info.run_id}")
 58        print(f"  Accuracy: {accuracy:.4f}")
 59
 60        return run.info.run_id
 61
 62
 63def get_model_versions(client):
 64    """등록된 모델 버전 조회"""
 65    print(f"\n{'='*50}")
 66    print(f"모델 '{MODEL_NAME}' 버전 목록:")
 67    print("="*50)
 68
 69    try:
 70        versions = client.search_model_versions(f"name='{MODEL_NAME}'")
 71        for v in versions:
 72            print(f"\n버전 {v.version}:")
 73            print(f"  상태: {v.current_stage}")
 74            print(f"  Run ID: {v.run_id}")
 75            print(f"  생성일: {v.creation_timestamp}")
 76            if v.description:
 77                print(f"  설명: {v.description}")
 78        return versions
 79    except Exception as e:
 80        print(f"모델을 찾을 수 없습니다: {e}")
 81        return []
 82
 83
 84def transition_to_staging(client, version: str):
 85    """모델을 Staging으로 전환"""
 86    client.transition_model_version_stage(
 87        name=MODEL_NAME,
 88        version=version,
 89        stage="Staging",
 90        archive_existing_versions=False
 91    )
 92    print(f"\n모델 v{version}을 Staging으로 전환했습니다.")
 93
 94
 95def transition_to_production(client, version: str):
 96    """모델을 Production으로 전환"""
 97    client.transition_model_version_stage(
 98        name=MODEL_NAME,
 99        version=version,
100        stage="Production",
101        archive_existing_versions=True
102    )
103    print(f"\n모델 v{version}을 Production으로 전환했습니다.")
104
105
106def update_model_description(client, version: str, description: str):
107    """모델 설명 업데이트"""
108    client.update_model_version(
109        name=MODEL_NAME,
110        version=version,
111        description=description
112    )
113    print(f"\n모델 v{version} 설명을 업데이트했습니다.")
114
115
116def add_model_tag(client, version: str, key: str, value: str):
117    """모델 태그 추가"""
118    client.set_model_version_tag(
119        name=MODEL_NAME,
120        version=version,
121        key=key,
122        value=value
123    )
124    print(f"\n모델 v{version}에 태그 추가: {key}={value}")
125
126
127def load_model_by_stage(stage: str):
128    """스테이지별 모델 로드"""
129    try:
130        model = mlflow.sklearn.load_model(f"models:/{MODEL_NAME}/{stage}")
131        print(f"\n{stage} 모델 로드 성공!")
132        return model
133    except Exception as e:
134        print(f"\n{stage} 모델 로드 실패: {e}")
135        return None
136
137
138def demo_workflow(client):
139    """전체 워크플로우 데모"""
140    print("\n" + "="*60)
141    print("MLflow Model Registry 워크플로우 데모")
142    print("="*60)
143
144    # 1. 첫 번째 모델 등록
145    print("\n[1] 첫 번째 모델 학습 및 등록...")
146    train_and_register_model(client, "v1")
147
148    # 2. 버전 조회
149    versions = get_model_versions(client)
150    if not versions:
151        return
152
153    latest_version = max(v.version for v in versions)
154
155    # 3. 설명 추가
156    print("\n[2] 모델 설명 추가...")
157    update_model_description(
158        client, latest_version,
159        "Initial model trained on Iris dataset with Random Forest"
160    )
161
162    # 4. 태그 추가
163    print("\n[3] 모델 태그 추가...")
164    add_model_tag(client, latest_version, "validated", "true")
165    add_model_tag(client, latest_version, "dataset", "iris")
166
167    # 5. Staging 전환
168    print("\n[4] Staging으로 전환...")
169    transition_to_staging(client, latest_version)
170
171    # 6. 두 번째 모델 등록
172    print("\n[5] 두 번째 모델 학습 및 등록 (개선 버전)...")
173    train_and_register_model(client, "v2")
174
175    # 7. 버전 재조회
176    versions = get_model_versions(client)
177    new_latest = max(v.version for v in versions)
178
179    # 8. 새 버전을 Staging으로
180    print("\n[6] 새 버전을 Staging으로...")
181    transition_to_staging(client, new_latest)
182
183    # 9. Production 승격
184    print("\n[7] Production으로 승격...")
185    transition_to_production(client, new_latest)
186
187    # 10. 최종 상태 확인
188    print("\n[8] 최종 모델 상태:")
189    get_model_versions(client)
190
191    # 11. Production 모델 로드 테스트
192    print("\n[9] Production 모델 로드 테스트...")
193    model = load_model_by_stage("Production")
194    if model:
195        # 간단한 예측 테스트
196        iris = load_iris()
197        sample = iris.data[:3]
198        predictions = model.predict(sample)
199        print(f"  샘플 예측 결과: {predictions}")
200        print(f"  실제 레이블: {iris.target[:3]}")
201
202
203def main():
204    """메인 함수"""
205    client = setup()
206
207    print("\nMLflow Model Registry 예제")
208    print("="*50)
209    print("\n옵션:")
210    print("1. 새 모델 학습 및 등록")
211    print("2. 등록된 모델 조회")
212    print("3. 전체 워크플로우 데모")
213
214    choice = input("\n선택 (1/2/3): ").strip()
215
216    if choice == "1":
217        train_and_register_model(client, "manual")
218    elif choice == "2":
219        get_model_versions(client)
220    elif choice == "3":
221        demo_workflow(client)
222    else:
223        print("잘못된 선택입니다.")
224
225
226if __name__ == "__main__":
227    main()