MLflow의 Backend Store로 PostgreSQL 활용하기 / 2025.02.09

MLflow는 기본적으로 파일(/mlruns)에 실험 정보를 저장한다.

이번에는, MLflow의 Backend Store로 PostgreSQL을 사용하여 실험 정보를 DB에 기록하고 이를 쿼리문으로 조회해보자.

▶ PostgreSQL 데이터베이스 생성

MLflow와 연동하기 위한 DB를 생성하자.

#postgresql 접속
sudo -u postgres psql

# DB, USER 생성
CREATE DATABASE mlflow;
CREATE USER mlflow WITH ENCRYPTED PASSWORD 'mlflow';
GRANT ALL PRIVILEGES ON DATABASE mlflow TO mlflow;

---

▶ MLflow와 PostgreSQL 연동

PostgreSQL에 실험 정보를 저장할 수 있게, MLflow와 연동해야한다.

따라서 MLflow Server를 실행할 때, 아래와 같은 형식을 따른다.

굵은 글씨 부분은 사용자에 따라 변경해야할 부분이다.

 

backend-store-uri 형식: postgresql://<postgresql_user>:<postgresql_pwd>@localhost:5432/<DB>

web server port: 5001

mlflow server \
  --backend-store-uri postgresql://mlflow:mlflow@localhost:5432/mlflow \
  --default-artifact-root ./gnn/mlruns/ \
  --host 0.0.0.0 --port 5001

---

▶ 모델 학습 코드에 PostgreSQL 추가

MLflow의 tracking uri를 postgresql로 변경하면, 학습 정보를 기록할 수 있게 연동 된다.

MLFLOW_TRACKING_URI = "postgresql://mlflow:mlflow@localhost:5432/mlflow"
mlflow.set_tracking_uri(MLFLOW_TRACKING_URI)

 

postgresql를 사용한 전체 코드는 아래 github에서 확인할 수 있다.

https://github.com/Cow-Kite/MLOps_Practice/tree/main/Study

MLOps_Practice/Study at main · Cow-Kite/MLOps_Practice

---

▶ PostgreSQL에서 실험 정보 저장 확인

이제, 위에 코드를 실행하고 실제로 DB에 저장이 되었는지 확인해보자.

• 실험 목록 조회 

# 연동한 DB로 접속
\c mlflow;

SELECT * FROM experiments;

• Run 목록 조회

SELECT * FROM runs;

• 성능 지표 확인

# test_acc 성능 지표 확인
SELECT * FROM metrics WHERE key = 'test_acc' ORDER BY timestamp DESC;

• 하이퍼 파라미터 확인

# learning_rate 확인
SELECT * FROM params WHERE key = 'learning_rate';

---

▶ MLflow Web Server에서 실험 정보 저장 확인

DB에 실험 정보가 저장되고 이를 Web UI에서도 확인할 수 있다.

---

▶ 마치며...

오늘은 PostgreSQL을 mlflow의 backend store로 활용해보았다.

기존에 실험 정보가 mlruns에 파일 형태로 저장되고 이를 web ui로 확인할 수 있었지만,

비교해야할 실험 및 모델이 많아진다면, postgresql에 저장하여 쿼리문으로 조회하는게 더 간편하다고 생각한다.