[MLflow] Artifact Store

 

Artifact Store 란 MLflow 에서 학습된 모델을 저장하는 Model Registry로써 이용하기 위한 스토리지 (storage) 서버다.

Artifact Store 를 이용하면 기본적인 파일 시스템 보다 체계적으로 관리 할 수 있으며 외부에 있는 스토리지 서버도 사용 할 수 있다는 장점이 있다.

 

이 포스트에서는 MinIO를 사용해서 Artifact Store을 구축해보겠다.

 

MinIO를 사용하는 이유는 다음과 같다.

• MinIO 는 S3 를 대체할 수 있는 오픈 소스 고성능 개체 스토리지다.
• AWS S3 의 API 와도 호환되어 SDK 도 동일하게 사용 할 수 있다.
• MLflow 에서는 AWS S3 를 모델을 저장하기 위한 스토리지로 사용하도록 권장하고 있기 때문에 MinIO 를 사용한다.
• 실습에서 AWS credential 을 통해 MinIO 대신 AWS S3 를 사용해도 같은 결과를 얻을 수 있다.

 

이 포스트에서 만들 Model Registry의 구성은 다음과 같다.

 

각 서버는 Docker Container 서버로 구성되어 있고 docker compose로 관리된다.

[Docker] 도커 컴포즈(Docker Compose)

 

먼저 DB Server에 대하여 정의한다.

version: "1"

services:
  mlflow-db:
    image: postgres:14.0 # PostgreSQL 이미지
    container_name: mlflow-db # 컨테이너 이름
    environment:
      POSTGRES_USER: mlflowuser # DB 접근을 위한 사용자 이름
      POSTGRES_PASSWORD: qwer1234 # DB 접근을 위한 사용 패스워드
      POSTGRES_DB: mlflowdatabase # DB 이름
    healthcheck: # DB 서버가 잘 띄워졌는지 상태를 확인하기 위해 상태를 체크
      test: ["CMD", "pg_isready","-q","-U","mlflowuser","-d","mlflowdatabase"]
      interval: 10s
      timeout: 5s
      retries: 5

 

다음은 artifact store인 MinIO Server를 정의한다.

mlflow-artifact-store:
    image: minio/minio:RELEASE.2024-01-18T22-51-28Z # MinIO 이미지
    container_name: mlflow-artifact-store # 컨테이너 이름
    ports: # 실행 포트
      - 9000:9000
      - 9001:9001
    environment:
      MINIO_ROOT_USER: minio # MinIO에 접근하기 위한 사용자 이름
      MINIO_ROOT_PASSWORD: qwer1234 # MinIO에 접근하기 위한 비밀번호
    command: server /data/minio --console-address :9001 # MinIO 서버 실행 명령어
    # --console-address를 통해 9001 포트로 MinIO에 접근할 수 있도록 주소를 열어줌
    healthcheck: # 상태체크
      test: ["CMD", "mc", "ready", "local"]
      interval: 5s
      timeout: 5s
      retries: 5

 

이제 MLflow server를 만들어야하는데, 이를 위해 따로 Dockerfile을 만들어야 한다.

FROM arm64v8/python:3.9-slim

RUN apt-get update && apt-get install -y \
    git \
    wget \
    && rm -rf /var/lib/apt/lists/*

RUN pip install -U pip &&\
    pip install mlflow psycopg2-binary boto3

RUN cd /tmp && \
    wget https://dl.min.io/client/mc/release/linux-arm64/mc && \
    chmod +x mc && \
    mv mc /usr/bin/mc

 

🔥 python과 minio client 정의는 호스트 컴퓨터의 OS 환경을 생각해야함(본 포스트에서는 M1 arm64이기 때문에 arm64로 설치)

 

 

• FROM arm64v8/python:3.9-slim :
  • Base 이미지를 Python 3.9가 포함된 이미지로 설정한다.

 

• RUN apt-get update && apt-get install -y \ ~ :
  • git , wget 을 설치한다.
  • git 은 MLflow 서버 내부 동작에, wget 은 MinIO Client 를 설치하기 위해 사용된다.

 

• RUN pip install -U pip && \ ~ :
  • MLflow 를 비롯해 PostgreSQL DB, AWS S3 에 관련된 Python 패키지를 설치한다.

 

• RUN cd /tmp && \ ~ :
  • 앞서 설치한 wget 을 활용하여 MinIO Client 를 설치한다.

 

이제 이 Dockerfile을 빌드하도록 docker-compose에 정의한다.

mlflow-server:
    build:
      context: .
      dockerfile: Dockerfile
    container_name: mlflow-server
    depends_on: # MLflow 서버가 띄워지기 전에, PostgreSQL DB, MinIO 서버를 먼저 띄우도록 한다.
      mlflow-db:
        condition: service_healthy
      mlflow-artifact-store:
        condition: service_healthy
    ports:
      - 5001:5000
    environment: 
      AWS_ACCESS_KEY_ID: minio # AWS S3의 credential 정보, MINIO_ROOT_USER와 동일
      AWS_SECRET_ACCESS_KEY: qwer1234 # AWS S3의 credential 정보, MINIO_ROOT_PASSWORD와 동일
      MLFLOW_S3_ENDPOINT_URL: http://mlflow-artifact-store:9000 # AWS S3의 즈소, MinIO의 주소와 같음
    command: 
      - /bin/sh
      - -c
      - |
        mc config host add mlflowminio http://mlflow-artifact-store:9000 minio miniostorage &&
        mc mb --ignore-existing mlflowminio/mlflow
        mlflow server \
        --backend-store-uri postgresql://mlflowuser:mlflowpassword@mlflow-db/mlflowdatabase \
        --default-artifact-root s3://mlflow/ \
        --host 0.0.0.0
      # MinIO 초기 킷 생성 후 MLflow 서버 실행
      # mc config ~ : MinIO Client 를 활용해 MinIO 서버에 호스트를 등록합니다.
      # mc mb ~ : 등록된 호스트를 통해 초기 버켓을 생성합니다.
      # mlflow server : MLflow 서버를 동작시킵니다.
      # --db-uri : 명시된 정보를 통해 PostgreSQL DB 와 연결합니다.
      # --default-artifact-root : 명시된 버켓을 통해 MinIO 의 초기 버켓과 연결합니다.

 

 

MLflow - localhost:5001

 

MinIO- localhost:9001

 

이제 환경이 다 구성되었으니 테스트를 해보자.

 

 

Save model

모델을 학습하고 저장해보자.

iris 데이터와 pytorch DNN을 사용해서 예측을 해보자.

import numpy as np
import torch
import torch.nn as nn
import pandas as pd
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split

# parser
import os
from argparse import ArgumentParser

# mlflow
import mlflow


def loadData() -> tuple:
    X, y = load_iris(return_X_y=True, as_frame=True)
    mapper = {0: 'setosa', 1: 'versicolor', 2: 'virginica'}
    y = pd.get_dummies(y.apply(lambda x: mapper[x])).astype('int')
    X_train, X_test, Y_train, Y_test = train_test_split(X, y, test_size=0.2)
    return X_train, X_test, Y_train, Y_test


class DNN(nn.Module):
    def __init__(self, inp, outp):
        super(DNN, self).__init__()
        self.fc1 = nn.Linear(inp, 64)
        self.fc2 = nn.Linear(64, 32)
        self.fc3 = nn.Linear(32, outp)
        self.relu = nn.ReLU()

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


def model_make(input_size: int, output_size: int):
    DEVICE = torch.device('mps')
    model = DNN(input_size, output_size)
    model.to(DEVICE)  # use mps
    opt = torch.optim.Adam(model.parameters(), lr=0.01)
    loss_f = nn.CrossEntropyLoss()
    return model, opt, loss_f, DEVICE


def get_accuracy(pred_arr, original_arr):
    pred_arr = pred_arr.cpu().detach().numpy()
    original_arr = original_arr.cpu().detach().numpy()
    final_pred = []
    final_origin = []
    for i in range(len(pred_arr)):
        final_pred.append(np.argmax(pred_arr[i]))
        final_origin.append(np.argmax(original_arr[i]))
    count = 0

    for i in range(len(original_arr)):
        if final_pred[i] == final_origin[i]:
            count += 1
    return round(count/len(final_pred), 4)


def fit(X_train, Y_train, X_test, Y_test, model, opt, loss_f, DEVICE, epochs=10):
    X_train = torch.tensor(X_train.values, dtype=torch.float32).to(DEVICE)
    Y_train = torch.tensor(Y_train.values, dtype=torch.float32).to(DEVICE)
    X_test = torch.tensor(X_test.values, dtype=torch.float32).to(DEVICE)
    Y_test = torch.tensor(Y_test.values, dtype=torch.float32).to(DEVICE)
    loss_ = []
    acc_ = []
    for epoch in range(1, epochs+1):
        opt.zero_grad()
        # forward feed
        y_pred = model(X_train)
        # loss
        loss = loss_f(y_pred, Y_train)
        print(f"epoch{epoch} {'-'*20} loss: {loss:.4f}")
        loss_.append(round(loss.cpu().detach().numpy().item(), 4))

        # backward
        loss.backward()
        # update weight
        opt.step()
        with torch.no_grad():
            acc = get_accuracy(y_pred, Y_train)
            acc_.append(acc)
    return model, loss_, acc_


if __name__ == "__main__":
    # mlflow set env
    # 값들은 따로 실제 사용시 secret 만들어서 저장하기...
    os.environ['MLFLOW_S3_ENDPOINT_URL'] = "http://127.0.0.1:9000"
    os.environ['MLFLOW_TRACKING_URI'] = "http://127.0.0.1:5001"
    os.environ['AWS_ACCESS_KEY_ID'] = "minio" # minio 접속 ID
    os.environ['AWS_SECRET_ACCESS_KEY'] = "qwer1234" # minio 접속 pw
    # parser
    parser = ArgumentParser()
    parser.add_argument("--model-name", dest="model_name",
                        type=str, default="Pytorch_DNN_Iris")
    args = parser.parse_args()
    mlflow.set_experiment("new-exp")
    # data
    X_train, X_test, Y_train, Y_test = loadData()
    model, opt, loss_f, DEVICE = model_make(X_train.shape[1], Y_train.shape[1])

    signature = mlflow.models.infer_signature(
        model_input=X_train, model_output=Y_train)
    input_sample = X_train.iloc[:10]

    mlflow.pytorch.autolog()
    # mlflow
    with mlflow.start_run(run_name="iris_pytorch_Linear"):
        model, loss, acc = fit(X_train, Y_train, X_test, Y_test,
                               model, opt, loss_f, DEVICE, epochs=100)
        for i in range(len(loss)):
            mlflow.log_metric("train_loss", loss[i], i+1)
            mlflow.log_metric("train_acc", acc[i], i+1)
        mlflow.pytorch.log_model(
            model, args.model_name, signature=signature,
            input_example=input_sample
        )

 

 

Load Model

import os
import mlflow
import pandas as pd
import torch
from sklearn.datasets import load_iris
if __name__ == "__main__":
    # setting
    os.environ['MLFLOW_S3_ENDPOINT_URL'] = "http://127.0.0.1:9000"
    os.environ['MLFLOW_TRACKING_URI'] = "http://127.0.0.1:5001"
    os.environ['AWS_ACCESS_KEY_ID'] = "minio"
    os.environ['AWS_SECRET_ACCESS_KEY'] = "qwer1234"

    model_name = "iris_pytorch"
    run_id = "59e8a439df2a41b893fba07d128515b7"
    model = mlflow.pytorch.load_model(f"runs:/{run_id}/{model_name}")
    # data
    X, y = load_iris(return_X_y=True, as_frame=True)
    mapper = {0: 'setosa', 1: 'versicolor', 2: 'virginica'}
    y = pd.get_dummies(y.apply(lambda x: mapper[x])).astype('int')
    DEVICE = torch.device('mps')
    X_train = torch.tensor(X.iloc[:20].values, dtype=torch.float32).to(DEVICE)
    with torch.no_grad():
        pred = model(X_train)
    print(pred)

 

mlflow를 사용해 모델을 불러와 사용한다.