MLFlow

현재 회사에서 MLOps 엔지니어로 근무하고 있습니다. 기존에는 GPU 병렬 컴퓨팅을 공부했고, 딥러닝에 대한 이론적 지식이 있으며 백엔드 엔지니어링을 공부한 배경을 바탕으로 MLOps 엔지니어 포지션을 제안받았기 때문입니다. MLOps 엔지니어는 아래와 같이 크게 세가지 일을 수행한다고 생각합니다.

 

1. 데이터 파이프라인 구축

2. 머신러닝 파이프라인 구축

3. 머신러닝 모델 서빙

 

이 번 포스팅에서는 머신러닝 파이프라인 구축 카테고리에서 머신러닝 모델 관리에 관한 오픈 소스를 리뷰할 까 합니다. 제목과 같이 MLFlow를 리뷰하려고 하고, MLFlow의 여러 기능 중에서 tracking과 registry 기능에 대해 리뷰하려고 합니다.

 

MLFlow는 클라이언트-서버 구조를 갖고 있습니다. 따라서 로컬머신에서 구동하여도 되고, 따로 중앙화된 서버를 구축하여도 됩니다. 개인이 사용할 경우에는 로컬머신을 사용하여도 무방하지만, 회사나 팀에서 사용할 경우에는 중앙화된 서버를 두고 모델 성능 확인이나 모델 공유를 손쉽게 하는 것을 추천하는 것 같습니다. 따라서 본 포스팅에서는 원격 서버를 사용하여 메트릭을 기록하고 모델을 저장하는 방법에 대해 작성하겠습니다.

 

MLflow client

mlflow의 기본 세팅은 로컬 서버를 사용하도록 되어있습니다. 따라서 이를 실제 사용할 서버 주소로 변경해야하며 이는 간단하게 mlflow.set_tracking_uri([서버주소]) 메소드를 사용함으로써 변경할 수 있습니다.

또 mlflow.get_tracking_uri() 를 사용하여 현재 tracking하고 있는 uri를 확인할 수 있습니다. 또 mlflow.set_experiment('test-experiment')를 사용하여 현재 실험 이름을 지정할 수 있습니다.

 

mlflow.set_tracking_uri('서버주소')
mlflow.set_experiment('test-experiment')

 

MLFlow Tracking

tracking은 모델의 parameter(모델의 파라미터), metric(loss 등 모델의 성능 지표)를 저장할 수 있습니다. 따라서 이를 통해 모델의 버전 및 성능을 저장하고, 성능 비교를 손쉽게 할 수 있을 것이라 생각합니다.

이는 mflow.log_params, mlflow.log_metrics 메소드를 통해서 활용이 가능합니다.

테스트를 위해서 fashion MNIST 데이터를 분류하는 분류기를 사용했습니다. MLP로만 구성되어 있으며, 20 epoch을 수행했습니다. 지금 코드에서는 각 epoch마다 metric을 기록했으나, 실제 개발환경에서는 학습 후에 metric을 기록하는 것이 좋아보입니다.

def test(dataloader, model, loss_fn):
    size = len(dataloader.dataset)
    num_batches = len(dataloader)
    model.eval()

    test_loss, correct = 0, 0

    with torch.no_grad():
        for X, y in dataloader:
            X, y = X.to(device), y.to(device)
            pred = model(X)
            test_loss += loss_fn(pred, y).item()
            correct += (pred.argmax(1) == y).type(torch.float).sum().item()
    test_loss /= num_batches
    correct /= size
    print(f"Test Error: \n Accuracy: {(100*correct):>0.1f}%, Avg loss: {test_loss:>8f} \n")
    mlflow.log_metrics({'accuracy': 100*correct, 'loss': test_loss})

코드를 실행히키면 metric이 수집되고, 수집된 metric은 mlflow 서버에 접속해서 확인할 수 있습니다. 

 

 

MLFlow Registry

MLFlow의 또다른 사용 이유 중 하나인, MLflow의 Registry 기능입니다. MLflow의 Registry는 중앙화된 모델 저장소로 모델의 단계(stage; production, staging 등) 및 버전등을 통한 관리가 가능합니다. MLFlow의 Registry에 저장하는 방법은 총 세가지입니다.

1. mlflow.<model_flavor>.log_model()을 사용하는 방법입니다. 아래는 sklearn패키지를 사용한 코드로 MLflow 공식 홈페이지에서 가져온 예제입니다.

from random import random, randint
from sklearn.ensemble import RandomForestRegressor

import mlflow
import mlflow.sklearn

with mlflow.start_run(run_name="YOUR_RUN_NAME") as run:
    params = {"n_estimators": 5, "random_state": 42}
    sk_learn_rfr = RandomForestRegressor(**params)

    # Log parameters and metrics using the MLflow APIs
    mlflow.log_params(params)
    mlflow.log_param("param_1", randint(0, 100))
    mlflow.log_metrics({"metric_1": random(), "metric_2": random() + 1})

    # Log the sklearn model and register as version 1
    mlflow.sklearn.log_model(
        sk_model=sk_learn_rfr,
        artifact_path="sklearn-model",
        registered_model_name="sk-learn-random-forest-reg-model"
    )

2. 파이썬 스크립트가 모두 실행된 후에 사용할 수 있는 방법입니다. 아래와 코드로 모델 저장이 가능합니다.

result = mlflow.register_model(
    "runs:/d16076a3ec534311817565e6527539c0/sklearn-model",
    "sk-learn-random-forest-reg"
)

3. 세번째는 모델 저장소를 create하고 model version을 create하면서 저장을 하는 방법입니다. 코드는 아래와 같습니다.

from mlflow.tracking import MlflowClient

client = MlflowClient()
client.create_registered_model("sk-learn-random-forest-reg-model")
result = client.create_model_version(
    name="sk-learn-random-forest-reg-model",
    source="mlruns/0/d16076a3ec534311817565e6527539c0/artifacts/sklearn-model",
    run_id="d16076a3ec534311817565e6527539c0"
)

모델 fetching은 아래와 같은 코드로 가능하다고 합니다. mlflow.<model_flavor>.load_model 메소드를 사용하변 됩니다.

import mlflow.pyfunc

model_name = "sk-learn-random-forest-reg-model"
model_version = 1

model = mlflow.pyfunc.load_model(
    model_uri=f"models:/{model_name}/{model_version}"
)

model.predict(data)

MLflow registry server관련

mlflow server --backend-store-uri sqlite:///mlflow.db --default-artifact-root /home/dev/mlflow/artifacts --host 0.0.0.0 --port 8000

registry 기능을 사용하려면 backend store URI를 DB 주소로 입력해야한다고 합니다. 현재는 따로 MLFLOW용 DB를 갖고있지 않기 때문에 일단은 sqlite를 사용하였습니다. 또 default-artifact-root를 지정하여야한다고 하여서 임의로 제 서버의 디렉토리 경로를 지정하였습니다. 이렇게 설정하는 경우에는 문제가 발생하는데, 서버의 디렉토리 경로가 아니라 클라이언트의 디렉토리에 저장되는 구조입니다. 따라서 S3나 hdfs 등의 주소를 적어야 실제 모델이 클라이언트에 저장되지 않고 서버에 저장될 것이라고 합니다.

 

따라서 위 명령어를 아래와 같이 변경하였습니다.

mlflow server --backend-store-uri sqlite:///mlflow.db --default-artifact-root s3://my-s3-bucket/mlflow/ --host 0.0.0.0 --port 8000

 

default-artifact-root에 s3 uri를 넣어서 모델이 s3에 저장되도록 하였습니다. 따라서 log_model 함수를 호출하면 로컬에 파일이 저장되는 것이 아니라, 주어진 s3 루트에 저장됩니다.

다만 이 경우에는 클라이언트(연구자)도 읽기/쓰기가 가능한 S3 키를 갖고 있어야 하고, 서버에서도 S3키가 필요합니다. 이 때에 환경변수에 AWS_ACCESS_KEY_ID와 AWS_SECRET_ACCESS_KEY를 등록하면 됩니다.

위와 같이 세팅 후 모델을 저장하면 자동으로 S3에 모델이 업로드되고, 이 파일 정보는 mlflow 서버에서 불러오게 됩니다. 최종적으로 웹사이트에서 정보를 확인할 수 있습니다.

 

모델 탭에 가면 description이나 stage를 설정할 수 있습니다. 또, 모델의 버전 간의 stage를 결정할 수 있으며, 어떤 모델이 production으로 올라가게 되면 기존의 production 모델은 자동으로 archive stage로 변경되게 됩니다.

Model Registry에서 모델 다운로드

모델 다운로드는 업로드보다 간단합니다. 먼저, 새로운 노트북 혹은 파이썬 파일을 열고, 아래와 같이 tracking_uri와 model_information을입력합니다. 그 후에 load_model을 통해서 모델을 불러오면 됩니다.

 

import mlflow

mlflow.set_tracking_uri(my_mlflow_server_uri)
model_uri = "runs:/f650e3e6a00f4b27a854b7d631fb031c/pytorch-test2"
loaded_model = mlflow.pyfunc.load_model(model_uri)
loaded_model.predict(input_image)

보안

mlflow는 오픈소스인 만큼 아이디 패스워드 설정이 따로 불가능한 것 같습니다… 사실 이것이 가장 큰 단점인 것 같습니다. 하지만 mlflow는 상용툴을 제외하고는 가장 파워풀한 오픈소스로 mlflow를 사용하는 것이 바람직해 보입니다.

따라서 mlflow의 보안 이슈를 어느정도 해결하기 위해서는 nginx 프록시를 두고 nginx에 접근할 때 ID/PASSWORD를 입력하게 하는 것이 어떨까 합니다.