Tutorial MLflow: MLOps de forma sencilla

Sin duda alguna, crear modelos es importante. Sin embargo, a medida que las empresas adoptan y ponen modelos en producción, cada vez es más importante saber gestionar esos modelos en producción. Esto es conocido como MLOps y, si bien hay muchas herramientas y formas de hacerlo, una muy usada es MLflow. Así pues, el objetivo de este tutorial de MLFlow es enseñarte cómo ponerlo en producción, así como el funcionamiento básico de MLFlow, tanto a nivel teórico como práctico. ¿Te suena bien? ¡Pues vamos a ello!

Introducción a MLflow

MLflow es una herramienta de Open Source para gestionar el ciclo de vida de los modelos de machine learning. Para ello, cuenta con varios aspectos principales:

  • Tracking: registra los resultados y parámetros de los modelos para poder compararlos.
  • Projects: paquetiza el código de tal forma que sea reproducible.
  • Models: permite gestionar el versionado de modelos, así como poner en producción modelos de ML a modo de endpoint. Este último se trata de un aspecto muy interesante, puesto que incluye integracciones para hacer el deploy del modelo tanto en Azure ML como en AWS SageMaker. Además, permite la exportación del modelo como fichero de Apache Spark.

Al usar MLflow no necesariamente se deben utilizar todas sus capacidades. En mi opinión, el apartado Projects no es tan potente. Además, decir que no es la única forma de hacer MLOps, aunque es una herramienta muy interesante en muchos casos.

Como ves, MLflow abarca muchas de las cuestiones clave en el desarrollo de modelos. En cualquier caso, el primer paso siempre será el mismo: instalar MLFlow en un servidor para que cualquier persona del equipo lo pueda usar. Veamos cómo hacerlo.

Cómo instalar MLflow

Para poner en producción un servidor con MLflow vamos a necesitar dos cuestiones:

  1. Una máquina virtual. En esta máquina virtual instalaremos MLflow, podremos ver la UI de MLflow, servirá modelos, etc.
  2. Una base de datos. La base de datos es el lugar donde MLflow guardará los metadatos del tracking de parámetros. Además, no puede ser cualquier tipo de base de datos, sino que debe ser alguna de las siguientes: MySQL, SQLite o PostgreSQL.
  3. Un lugar donde guardar artefactos: este es el lugar donde guardaremos los modelos. En este caso, podemos usar la propia máquina virtual, aunque lo más típico suele ser usar un Datalake como S3 o Cloud Storage, por ejemplo. En nuestro caso, al estar en el entorno de Google Cloud, usaremos Cloud Storage, ya que nos facilitará el trabajo.

Así pues, sigamos con el tutorial de MLflow, viendo cómo instalarlo en una máquina virtual:

Cómo instalar MLflow en una máquina virtual

Para poder hacer la instalación de MLflow primero debemos tener una máquina virtual corriendo. Esto se puede hacer sobre cualquier plataforma Cloud o servicio propio, aunque en mi caso yo lo haré en GCP, puesto que ofrece una máquina virtual e2-micro gratuita al mes.

En el caso de GCP debemos ir a Compute Engine y crear una instancia nueva, como aparece en la siguiente imagen:

Crear Máquina Virtual Compute Engine

Una vez la tengas creada, accede a la máquina virtual mediante SSH haciendo clic en el texto donde aparece SSH. Una vez hayas hecho clic, se abrirá una ventana. Perfecto, ya estás dentro de la máquina virtual. Ahora lo primero de todo tenemos que instalar Miniconda. Para ello ejecutamos el siguiente código:

# Instalamos conda
curl -O -J https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh
bash Miniconda3-latest-Linux-x86_64.sh

Tendrás que dar al enter hasta llegar al final del documento y finalmente aceptar los términos. Por últimos debes activar conda:

source .bashrc

Por último, tenemos que instalar MLflow. Para ello hay dos opciones:

  • Opción 1: Instalar MLlflow con todos los paquetes que puede llegar a utilizar, como Sklearn, Azure, etc. De esta forma, te evitas tener que andar instalando varios paquetes. Para ello, debes ejecutar el siguiente comando:
pip install mlflow[extras]
  • Opción 2: Instalar simplemente MLflow sin ningún extra. En cuyo caso, tendremos que instalar los extras que usemos nosotros mismos. Para ello, simplemente debemos ejecutar el siguiente comando:pip install mlflow

En mi caso, para facilitar el tutorial de MLflow lo instalaré con los extras, es decir, la opción 1.

Ahora, vamos a crear una regla en el firewall. Esto es necesario para poder acceder a la UI de MLflow, así como para poner modelos en producción. Para ello, tenemos que ir al apartado Firewall dentro de VPC Networw (enlace). En la parte superior podremos crear una regla, en la que tendremos que definir los siguientes puntos:

  • Targets (destino): Specified target tags
  • Target tags (etiquetas de destino): mlflow (o el nombre de la etiqueta que le hayas dado al crear la VM).
  • Source filter (filtros de origen): Rangos de IPV4
  • Source IP ranges (Rangos de IPv4 de origen): 0.0.0.0/0
  • Specified protocols and ports tcp (Protocolos y puertos especificados tcp): 8080,8081,1234.

Si más adelante vas a desplegar un modelo en un puerto (por ejemplo el puerto 1234) recuerda añadirlo a la regla para poder acceder al mismo.

La regla debería parecerse a la siguiente:

Regla de Firewall para VM

Una vez tengamos la regla creada, volvemos la consola de la VM y ejecutamos el siguiente comando:

mlflow ui -p 8080 -h 0.0.0.0

Ahora si vamos a la IP externa de nuestra máquina virtual (se indica en el UI de Compute Eninge y tiene la siguiente forma: XX.XXX.XXX.XXX donde x son números) y ponemos el puerto 8080 podemos acceder a la UI de MLflow. La url quedará algo como esto: 12.345.678.912:8080.

Perfecto, ya tenemos MLflow instalado en la máquina virtual. Ahora, veámos cómo crear una base de datos para poder usarlo como repositorio de los logs de modelos.

Cómo conectar una base de datos a MLflow

La base de datos es necesaria para guardar información como los parámetros usados en el modelo o las métricas obtenidas por el modelo.

Lo primero de todo para conectar la base de datos a MLflow es, obviamente, crear la base de datos. Como hemos dicho, crearé la base de datos dentro de la propia máquina virtual Así pues, para ello ejecutaré el siguiente código:

sudo apt-get install postgresql postgresql-contrib postgresql-server-dev-all -y

Ahora que tenemos la base de datos, nos conectamos a ella:

sudo -u postgres psql

Ahora que estamos conectados tendremos que crear la base de datos donde se guardará la información.

CREATE DATABASE mlflow;

Además, tendremos que crear un usuario y contraseña para poder acceder a la base de datos. Además, tendremos que darle permisos al usuario:

CREATE USER mlflow WITH ENCRYPTED PASSWORD 'your_password';
GRANT ALL PRIVILEGES ON DATABASE mlflow TO mlflow;

Por último, tenemos que instalar las librerías necesarias para que internamente MLflow se pueda conectar a esta base de datos para leer y escribir datos. Al ser Postgres tendremos que ejecutar el siguiente comando:

sudo apt install gcc -y 
pip install psycopg2

Con esto ya tenemos la base de datos creada. Ahora, podemos seguir con el tutorial de MLflow viendo cómo crear el repositorio de artefactos.

Cómo añadir el repositorio de artefactos a MLflow

El repositorio de artefactos es donde se guardarán todos los ficheros que no sean métrcias o parámetros, como modelos entrenados, datos, imágenes, o cualquier otro fichero que queramos guardar.

Tal como he comentado previamente en el tutorial de MLflow, hay dos formas de añadir un repositorio de artefactos: o bien usar la propia máquina virtual o usar un datalake (S3, Cloud Storage, etc.). En nuestro caso usaremos Cloud Storage.

Para usar Cloud Storage con MLflow tenemos que seguir tres pasos:

  1. Crear el bucket en Cloud Storage.
  2. Dar acceso a la cuenta de servicio de Compute Engine.
  3. Permitir acceso a Cloud Storage tanto a nivel de servidor como de cliente (el cliente necesita acceso para poder escribir artefactos, mientras que el servidor necesita para leerlos y utilizarlos).

Así pues, vamos por pasos:

Cómo crear un Bucket en Cloud Storage para MLflow

Crear un Bucket es muy sencillo. Simplemente tenemos que ir a crear un bucket en Cloud Storage y darle un nombre y asignar el tipo de almacenamiento que queramos.

Una vez indiquemos el bucket, si vamos a Configuración –> URI de gsutil veremos la URI de nuestro bucket, la cual tendrá una forma como la siguiente: gs://bucket_name

Perfecto, este punto hecho. Sigamos con el tutorial de MLflow viendo cómo dar acceso a la cuenta de servicio.

Cómo dar acceso a la cuenta de servicio de Compute Engine

Lo primero de todo tenemos que ir al apartado usuarios dentro de IAM (enlace). Una vez ahí, copiamos el correo electrónico de la cuenta de servicio de compute, el cual tiene la siguiente estructura xxxxxxxxxxxxx-compute@developer.gserviceaccount.com.

Una vez la tenemos copiada, volvemos a nuestro bucket y, esta vez, vamos al apartado Permisos.En la parte inferior veremos el botón agregar. Clicamos y pegamos la cuenta de servicio que hemos copiado. En el apartado funciones añadiremos la función Storage Object Admin. De esta forma la cuenta de servicio podrá crear y acceder a los artefactos.

Perfecto, ahora solo hace falta que tanto servidor como cliente tengan acceso a este bucket. Veámos cómo hacerlo:

Permitir acceso a Cloud Storage al servidor y al cliente

Como la máquina virtual está de lado del servidor y en el mismo proyecto que el bucket, no hace falta hacer nada. Respecto al cliente, es decir, nuestro ordenador, para poder acceder a Cloud Storage primero tenemos que descargarnos las claves de la cuenta de servicio de Compute Engine.Para ello, volvemos al apartado de Cuentas de Servicio dentro de IAM (enlace) y en los tres puntos hacemos clic en «Crear Claves». Hecho esto, se nos descargará un fichero .json con nuestras claves de servicio.

Volvemos a nuestro entorno de desarollo e instalamos la siguiente librería para poder acceder a Cloud Storage:

pip install google-cloud-storage

Por último, simplemente hay que llevar las claves de servicio a la carpeta donde tengamos nuestro proyecto de ML y ejecutar el siguiente código:

from google.cloud import storage
credentials = 'service_account.json
client = storage.Client(creds = credentials)

¡Perfecto! Ya tenemos nuestro Cloud Storage montado y listo para poder ser utilizado por nuestro servidor de MLflow. Ahora, veámos cómo lanzar el servidor. ¡Vamos a ello!

Cómo lanzar el servidor con MLflow

Para lanzar el servidor con MLflow tenemos que tener 3 cuestiones en mente:

  1. Host y puerto: lo mismo que hemos definido en el Host y puerto a la hora de crear la regla del Firewall, es decir, host 0.0.0.0 y puerto 8080.
  2. La URI de la base de datos. Como siempre, el formato será el siguiente: postgresql://<user>:<password>@<host>/<database>

En mi caso, la URI de la base de datos es la siguiente:

postgresql://mlflow:your_password@localhost/mlflow
  1. Ruta del repositorio de artefactos. Como hemos visto anteriormente, esta ruta tiene la siguiente forma: gs://bucket_name.

Con estos tres puntos, podemos lanzar el servidor de MLflow con el siguiente código:

mlflow server --backend-store-uri <db_URI> --default-artifact-root gs://bucket_name -h <host> -p <port>

¡Perfecto! Ya tenemos MLflow instalado en la máquina virtual. Ahora sigamos con el tutorial de MLflow viendo cómo funciona cada apartado. ¡Vamos allá!

Cómo usar MLflow Tracking

Conceptos importantes sobre MLflow Tracking

Guardar el tracking de nuestros modelos en nuestro servidor de MLflow es muy sencillo. Simplemente tenemos que:

  1. Conectarnos a nuestro servidor remoto con MLflow. Para hacer la conexión simplemente debemos usar el método set_tracking_uri e indicar la IP de nuestro servidor y el puerto. Ejemplo:
import mlflow
mlflow.set_tracking_uri('http://xx.xxx.xxx.xxx:8080')

Con esto, ya tendremos hecha la conexión con nuestro servidor de MLflow

2. Opcionalmente, podemos crear un experimento donde incluir los parámetros, si es que no lo tenemos ya. Cada proyecto diferente debería tener su propio experimento. Para ello, usaremos el método create_experiment:

experiment_name = "experiment_iris"

if not mlflow.get_experiment_by_name(experiment_name):
    mlflow.create_experiment(name=experiment_name)

3. A la hora de entrenar el modelo, pasar logs de los parámetris, métricas, artefactos y el modelo. Existe los siguientes datos que podemos incluir en MLflow:

  • Parámetros del modelo: indica parámetros del modelo utilizado. Se loguea usando el método log_param.
  • Métricas: se refiere a métricas de rendimiento, tales como el RMSE, accuracy, AUC, etc. Se loguea usando el método log_metric.
  • Artefactos: permite incluir ficheros y/o carpetas. El uso típico es incluir datos de entrenamiento, imágenes del entrenamiento, etc. Los artefactos se loguean usando el método log_artifact.
  • Modelos: permite incluir modelos. Los modelos se loguean usando el método log_model. Además, las librerías Sklearn, Tensforlow, Keras, Gluon, XGBoost, LightGBM, Statsmodels, Spark, Fastai y Pytorch están habilitadas para autologing. Es decir, que simplemente podemos usar el método autolog y MLflow automáticamente logueará los datos que vayamos generando. Puedes aprender más sobre el autlog aquí.

Importante: para poder loguear los parámetros primero debemos indicar a MLflow que «escuche». Esto podemos hacerlo de dos formas:

  1. Usar el método star_run junto con with para evitar cerrar el proceso. Ejemplo:
with mlflow.start_run():
   mlflow.log_param('max_depth', max_depth)
   #...

2. Usar los métodos `start_run` y `end_run`:

mlflow.start_run() 
mlflow.log_param('max_depth', max_depth) 
#...
mlflow.end_run()

Dicho esto, veámos cómo funciona con un ejemplo muy sencillo:

Ejemplo de MLflow Tracking

Como ejemplo vamos a crear un modelo que prediga el tipo de flor del dataset Iris. Para ello,crearemos un modelo de Random Forest con Sklearn aplicando Grid Search y Cross Validation y loguearemos los parámetros que mejor funcionan.

Asimismo, es importante recordar dos cuestiones clave:

  1. Conectarnos a nuestro Data Lake, en mi caso, Cloud Storage, para que podamos guardar ahí los artefactos.
  2. Hacer el setup de MLflow mediante set_tracking_uri y, opcionalmente, crear y definir el experimento.

Además, loguearemos los datos de entrenamiento y el modelo resultante.

Si no conoces alguna de las funciones de Sklearn que utilizo o te gustaría profundizar en esta librería, te recomiendo que leas este tutorial de Sklearn.

# Cargo las librerías
import importlib
import mlflow
importlib.reload(mlflow)
import numpy as np
import mlflow 
import mlflow.sklearn
from sklearn.datasets import load_iris
from sklearn.metrics import precision_score, accuracy_score, recall_score
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import GridSearchCV, train_test_split

# Conecto con Cloud Storage
service_account = 'credentials.json'
from google.cloud import storage
client = storage.Client(creds = service_account)

# Si no existe, creo el experimento
experiment_name = "experiment_iris"

if not mlflow.get_experiment_by_name(experiment_name):
    mlflow.create_experiment(name=experiment_name) 

experiment = mlflow.get_experiment_by_name(experiment_name)

# Setup de MLflow
mlflow.set_tracking_uri('http://XX.XXX.XXX.XXX:8080/')

# Cargo los datos
data = load_iris()

# Hago split entre train y test
x_train, x_test, y_train, y_test = train_test_split(
    data['data'],
    data['target'],
    test_size= 0.2,
    random_state= 1234
    )

# Definimos el modelo
rf_class = RandomForestClassifier()

# Definimos el grid de hiperparámetros
grid = {
    'max_depth':[6,8,10], 
    'min_samples_split':[2,3,4,5],
    'min_samples_leaf':[2,3,4,5],
    'max_features': [2,3]
    }

# Hago el Grid Search
rf_class_grid = GridSearchCV(rf_class, grid, cv = 5) 
rf_class_grid_fit = rf_class_grid.fit(x_train, y_train)

print(f'Best parameters: {rf_class_grid_fit.best_params_}')
Best parameters: {'max_depth': 6, 'max_features': 2, 'min_samples_leaf': 2, 'min_samples_split': 2}

Ahora que tenemos el modelo entrenado, vamos a hacer el loging en MLflow:

# Ahora hago el loging de los parámetros
with mlflow.start_run(experiment_id = experiment.experiment_id):

    # Logueo los mejores resultados
    mlflow.log_params(rf_class_grid_fit.best_params_)

    # Obtengo las predicciones
    y_pred = rf_class_grid_fit.predict(x_test)

    # Calculo el acuraccy y el AUC
    accuracy = accuracy_score(y_test, y_pred)
    precision = precision_score(y_test, y_pred, average='weighted')
    recall = recall_score(y_test, y_pred, average='weighted')
    print(f'Accuracy: {accuracy}\nPrecision: {precision}\nRecall: {recall}')

    # Log de parámetros
    metrics ={
        'accuracy': accuracy,
        'precision': precision, 
        'recall': recall 
        }

    mlflow.log_metrics(metrics)


    # Log model & artifacts
    np.save('data/artifacts/x_train', x_train)
    mlflow.log_artifact('x_train.npy')

    mlflow.sklearn.log_model(rf_class_grid_fit, 'iris_rf_first_attempt')
Accuracy: 1.0
Precision: 1.0
Recall: 1.0

¡Perfecto! Ya tenemos nuestro modelo creado y se ha traqueado en MLflow. Ahora, si vamos a la UI veremos que el modelo tiene tres partes.

Por un lado, tenemos los parámetros de nuestro modelo, de tal forma que sabemos con qué parámetros se ha entrenado ese modelo en concreto:

Parámetros de modelo en MLflow

Asimismo, en el apartado metricas podemos ver también cuál han sido las métricas que hemos conseguido con ese modelo. En mi caso solo hemos incluido métricas finales, pero si hubieramos entrenado una red neuronal, por ejemplo, podríamos haber guardado otras métricas como el validation accuracy de cada iteración. En esos casos, desde la UI podremos ver la evolución de la métrica.

Métricas del modelo en MLflow

Por último, el apartado artifacts incluye tanto los artefactos, en mi caso, los datos, como el modelo

Artefactos en MLFlow

Como ves, es muy sencillo registrar modelos y parámetros en MLflow. Pero MLflow va mucho más allá. Sigamos con este tutorial de MLflow viendo, lo que para mi, es algo clave de este software: la puesta en producción de modelos. ¡Vamos a ello!

Puesta en producción de modelos con MLflow

Una vez tenemos un modelo subido a MLflow podemos ponerlo en producción de una forma muy sencilla a modo de API. Para ello, vamos a la pestaña «Artifacts» del modelo que queremos poner en producción. Ahí, podremos hacer clic en el botón «Register Model» lo cual nos abrirá una ventana como la siguiente:

Registrar modelo en MLflow

Simplemente debemos indicar el nombre del modelo que, en mi caso, es iris (el nombre simplemente sirva para saber qué modelo es y ver diferentes versiones del mismo).

Ahora de cara a obtener predicciones hay dos opciones:

  • Leer el artefacto del servidor de MLflow y lo usemos para hacer predicciones.
  • Publicar el artefacto como un endpoint.

Veamos cómo hacer cada caso:

Hacer predicciones de un modelo de MLflow

Para obtener las predicciones, simplemente tenemos que ir a al apartado de artifacts y nos aparecerá doscómo podemos obtener predicciones de nuestro modelo, ya sea usando Spark o Python. Así pues, simplemente copiamos ese código y lo ejecutamos dando unos datos reales que queramos predecir:

logged_model = 'gs://mlflow_artifacts_bucket/artifacts/7/72c46af3d3f649569c4df0c7cdfeb263/artifacts/iris_rf_first_attempt'

# Load model as a PyFuncModel.
loaded_model = mlflow.pyfunc.load_model(logged_model)

# Predict on a Pandas DataFrame.
import pandas as pd
loaded_model.predict(pd.DataFrame(x_test))
array([1, 1, 2, 0, 1, 0, 0, 0, 1, 2, 1, 0, 2, 1, 0, 1, 2, 0, 2, 1, 1, 1,
       1, 1, 2, 0, 2, 1, 2, 0])

Esta forma de poner en producción el modelo requiere que el cliente que vaya a hacer las peticiones pueda interactuar con la API de MLflow, es decir, que se haga desde Python, R o Java. Si queremos hacer la predicción de forma sencilla con REST, yo optaría por el siguiente método.

Visto cómo poner un modelo de MLflow en producción, sigamos con nuestro tutorial de MLflow viendo cómo podemos publicar un modelo como un endpoint. ¡Vamos a ello!

Cómo publicar un modelo de MLflow como un endpoint

De cara a poner un modelo de MLflow en producción como endpoint podemos hacer dos cosas:

  1. Hacer el deploy en el propio servidor de MLflow.
  2. Hacer el deploy en una herramienta Cloud externa, como AWS SageMaker, Azure ML o Apache Spark UDF.
  3. Descargar el modelo como Docker y ponerlo en producción en cualquier otra herramienta.

En mi caso usaré la primera opción, ya que es la más sencilla y, para modelos medianos funciona bien. Sin embargo, para modelos con mucho más tráfico, recomendaría usar el segundo punto.

Dicho esto, para poner en producción un modelo de MLflow como endpoint, simplemente tenemos que copiar el URI del modelo, el cual se encuentra en la parte superior de Artifacts, al llado de Full Path y ejecutar el siguiente código:

mlflow models serve -m "<model_uri>" -p <port> -h <host> --no-conda

Es importante que el puerto del endpoint sea diferente al puerto en el que está la UI. En mi caso, he usado el puerto 1234.

Hecho esto, MLflow desplegará el modelo en el puerto que hemos indicado y ahora simplemente tendremos que hacer peticiones POST al endpoint invocations en el puerto indicado:

import requests
import pandas as pd

url = 'http://XX.XXX.XXX.XXX:1234/invocations'

headers = {'Content-type': 'application/json'}
data = pd.DataFrame(x_test).to_json(orient='split')

resp = requests.post(
    url,
    headers=headers,
    data = data
    )

resp.content
b'[1, 1, 2, 0, 1, 0, 0, 0, 1, 2, 1, 0, 2, 1, 0, 1, 2, 0, 2, 1, 1, 1, 1, 1, 2, 0, 2, 1, 2, 0]'

Como vemos, ahora tenemos el modelo en producción en modo de API y podemos acceder al mismo de una forma muy sencilla mediante peticiones HTTP.

Conclusión

Sin duda alguna MLflow es una herramienta increible de cara a poner modelos en producción y MLOps. Y es que es una muy buena herramienta tanto para registrar los modelos, sus parámetros, métricas y datos de entrada, pero también para poner estos modelos en producción.

Además, es una plataforma que es agnóstica al ecosistema, lo puedes implementar usando diferentes bases de datos, diferentes tipos de sistemas de ficheros y puedes poner los modelos en producción en el propio servidor o de forma muy fácil en AWS y Azure.

En mi opinión, lo más costoso de MLflow es ponerlo en producción y conseguir que todas las piezas encajen y funcionan. En cualquier caso, espero que este tutorial de MLflow te haya servido para aprender a poner MLflow en producción y para saber usar el apartado de Tracking y Deployment de modelos de MLflow.

Si es así y quieres estar al día de más contenido así, te animo a que te suscribas para estar al día de los nuevos posts que voy subiendo. ¡Nos vemos en el siguiente post!