MLOps·머신러닝 운영/MLflow를 활용한 머신러닝 실험 관리

MLflowClient 완전 정리 — MLflow를 코드로 제어하는 방법

Data Jun 2025. 10. 25. 10:26

1. MLflowClient란?

MLflowClient는 MLflow의 핵심 기능(Tracking, Registry 등)
코드로 직접 제어할 수 있는 Python 클래스입니다.

즉, “UI 대신 코드로 MLflow를 다루는 리모컨” 같은 존재죠.

🔹 주요 기능

  • 실험(Experiment) 생성 및 조회
  • 실행(Run) 생성 및 로그 기록
  • 모델(Model) 등록 및 버전 관리
  • 아티팩트(Artifact) 업로드 / 다운로드
  • 모델 스테이지(Stage) 전환 (Staging → Production)

요약

MlflowClient는 MLflow를 자동화하거나 외부 시스템에서 제어할 때 사용하는 저수준(Low-Level) API입니다.

 

 

 

2. MLflow의 두 가지 접근 방식

MLflow를 사용할 때는 크게 두 가지 방법이 있습니다.

 구분 고수준 API  저수준 API
대표 모듈 mlflow mlflow.tracking.MlflowClient
사용 목적 간단한 실험 기록 자동화, 백엔드 통합, 대규모 제어
주 사용 위치 모델 학습 코드 (train.py) 관리 스크립트 (manage_mlflow.py)
Run 컨텍스트 자동으로 관리(start_run) 명시적으로 run_id 지정
모델 스테이지 관리 불가능 가능
Tracking 서버 접근 제한적 완전 지원
코드 스타일 직관적, 선언형 명시적, 세밀한 제어

비유하자면

  • mlflow는 “노트북에서 실험할 때 쓰는 인터페이스”
  • MlflowClient는 “운영 서버나 파이프라인에서 모델을 자동으로 관리하는 엔진”

 

3. 기본 사용법과 예제

 1️⃣  MLflowClient 초기화

from mlflow.tracking import MlflowClient

# MLflowClient 객체 생성
client = MlflowClient()

이 객체는 MLflow Tracking 서버(또는 로컬 mlruns/ 디렉토리)에 직접 연결되어
모든 실험 및 모델을 코드 기반으로 조작할 수 있습니다.

 

 2️⃣ 실험 생성 및 조회

# 새로운 실험 생성
experiment_id = client.create_experiment("My_New_Experiment")

# 실험 정보 조회
experiment = client.get_experiment(experiment_id)
print(f"Experiment ID: {experiment.experiment_id}, Name: {experiment.name}")

UI에서 “새 실험 생성”을 클릭하는 대신,
create_experiment() 메서드 한 줄로 동일한 작업을 수행합니다.

 

 

4.  실행(Run) 생성과 메트릭 기록

import mlflow

# Run 시작
with mlflow.start_run(experiment_id=experiment_id) as run:
    run_id = run.info.run_id
    client.log_param(run_id, "learning_rate", 0.01)
    client.log_metric(run_id, "accuracy", 0.95)

# Run 정보 조회
retrieved_run = client.get_run(run_id)
print(f"Run ID: {retrieved_run.info.run_id}, Metrics: {retrieved_run.data.metrics}")

핵심 포인트

  • mlflow.start_run() 으로 실행 컨텍스트 생성
  • client.log_param, client.log_metric 으로 세밀한 로그 기록
  • get_run() 으로 결과 조회

이 방식은 단순한 실험 기록을 넘어서,
CI/CD 파이프라인에서도 동일한 로그 구조를 보장할 수 있습니다.

 

5. 아티팩트 관리

import os

artifact_path = "output.txt"
with open(artifact_path, "w") as f:
    f.write("Sample artifact content")

# 로그 및 다운로드
client.log_artifact(run_id, artifact_path)
artifacts = client.list_artifacts(run_id)
print(f"Artifacts: {[a.path for a in artifacts]}")
downloaded_path = client.download_artifacts(run_id, artifact_path)
print(f"Artifact downloaded to: {downloaded_path}")

모델 파일, 로그, 이미지 등 모든 산출물을 아티팩트(Artifact) 로 남기고 관리합니다.

 

6. 모델 등록 및 버전 관리

model_name = "my_model"
model_uri = f"runs:/{run_id}/model"

# 모델 등록 및 버전 생성
client.create_registered_model(model_name)
client.create_model_version(
    name=model_name,
    source=model_uri,
    run_id=run_id
)

# 등록된 모델 정보 조회
registered_model = client.get_registered_model(model_name)
print(f"Registered Model: {registered_model.name}")

학습 코드에서는 보통 mlflow.register_model()을 사용하지만,
MlflowClient는 등록뿐 아니라 버전 관리, 상태 전환까지 모두 제어할 수 있습니다.

 

7. 모델 스테이지 전환 (Production 배포)

client.transition_model_version_stage(
    name=model_name,
    version=1,
    stage="Production"
)

model_version = client.get_model_version(model_name, 1)
print(f"Model version {model_version.version} is now in {model_version.current_stage}")

모델을 자동으로 "Staging" → "Production" 으로 올리는 코드
➡️ MLOps 파이프라인에서 자동 배포 및 검증 워크플로우 구축 가능

 

8. 고수준 vs 저수준 — 실제 코드 비교

 1️⃣ 고수준 API (mlflow)

import mlflow

with mlflow.start_run():
    mlflow.log_param("alpha", 0.5)
    mlflow.log_metric("accuracy", 0.94)
    mlflow.sklearn.log_model(model, "model")

✔️ 간단하고 빠르지만
❌ 모델 버전 관리나 스테이지 전환은 불가능

 

 2️⃣  저수준 API (MlflowClient)

from mlflow.tracking import MlflowClient
client = MlflowClient()

client.log_param(run_id, "alpha", 0.5)
client.log_metric(run_id, "accuracy", 0.94)
client.create_registered_model("my_model")
client.transition_model_version_stage("my_model", version=1, stage="Production")

run_id 단위로 제어 가능
✔️ Registry 관리 및 자동화 파이프라인에 적합
➡️ 운영 단계에서 모델을 자동 관리할 때 필수

 

 

정리하면

 

상황 추천 API
모델 학습 시 실험 기록 mlflow
실험 결과 관리 및 모델 등록 자동화 MlflowClient
모델 스테이지 전환 및 배포 자동화 MlflowClient

mlflow는 실험을 “기록”하기 위한 도구,
MlflowClient는 MLflow 시스템을 “제어”하기 위한 도구입니다.

둘은 경쟁 관계가 아니라 역할이 다릅니다.

  • 개발자 / 연구자 → mlflow API로 빠르게 실험 기록
  • MLOps 엔지니어 / 시스템 자동화 담당자 → MlflowClient로 관리 및 배포 자동화

 

MLflowClient는 MLflow의 숨은 제어판이다.
단순한 로그 기록을 넘어, 모델의 생애주기 전체를 코드로 통제할 수 있게 만든다.