Descubriendo MLOps con AWS, Docker y Terraform.

Con el creciente desarrollo de modelos de Machine Learning, que aportan soluciones a múltiples problemas, se hace más necesario que nunca tener un proceso definido para simplificar el despliegue en producción de dichos modelos y lograr que sean realmente útiles. En este post te explicamos un enfoque básico de cómo lograrlo usando AWS, Docker y Terraform.

Pero, antes de entrar en materia, es importante tener claro a qué hace referencia el término MLOps y qué beneficios ofrece.

¿Qué es MLOps?

En esencia, MLOps es un conjunto de buenas prácticas para permitir que los modelos de Machine Learning (ML), que muchas veces se quedan como meros experimentos, se desplieguen en producción de manera ágil, consistente, constante y automática, permitiendo que aporten soluciones a los problemas por los cuales fueron concebidos, fomentando la cooperación y comunicación entre los equipos de Data Scientists y TI.

¿Qué beneficios ofrece MLOps?

• Agilidad: Reduce la fricción entre equipos y mejora la colaboración entre Data Scientist y el área de TI, haciendo más rápido el despliegue de un modelo.
• Automatización: Crea un ciclo de vida para el modelo de ML, permitiendo el despliegue automático, constante, consistente y rápido de modelos.
• Gobierno: Mejora el seguimiento del modelo, el control de versiones, el monitoreo y la administración.
• Estandarización: MLOps se ha convertido en la práctica estándar que ayuda a los equipos de Data Scientists y TI, a administrar todo el ciclo de vida de modelos ML.

Dentro de un proceso de MLOps existen múltiples pasos que cubren el ciclo de vida de un modelo ML, desde la extracción de datos, validación de los mismos, ingeniería de características, pasando por el entrenamiento, validación y despliegue de modelos, entre otros más que puedes ver aquí.

En este post nos vamos a centrar en los pasos relativos al entrenamiento, validación y despliegue de modelos y cómo lograr su automatización en la nube de AWS.

Ahora que ya tenemos claro en qué consiste MLOps y el alcance de este post, ¡¡podemos ponernos manos a la obra!!

Definición del proceso MLOps

El objetivo de este post es explicar cómo automatizar el entrenamiento, generación, despliegue y validación de un modelo ML básico basado en el típico ejemplo de la clasificación de la flor Iris.

En el siguiente diagrama observamos los servicios de AWS y demás tecnologías y, sus interacciones, que van a permitir generar el proceso básico de MLOps sobre el modelo ML citado anteriormente.

Cada uno de los servicios y tecnologías del diagrama anterior abarcaría un post completo. Por ello, vamos a explicar a grosso modo cada uno de ellos y cómo interactúan entre sí.

Servicios de AWS

A continuación, os contamos cada uno de los servicios de AWS que se van a utilizar y la interacción entre los mismos dentro del proceso de MLOps.

• Step Function: Servicio que permite coordinar múltiples servicios de AWS en un flujo de trabajo. Mediante este servicio se crea un flujo de trabajo que realiza el entrenamiento, generación del modelo y la creación de jobs de SageMaker para el entrenamiento y transformación batch del modelo ML, así como la generación del endpoint correspondiente, el cual se usa para realizar inferencias on-line. Por último, ejecuta una función lambda que hace una llamada al endpoint mencionado anteriormente para realizar una validación.
• Lambda: Servicio que permite ejecutar código sin necesidad de aprovisionar y administrar servidor (Serverless). Mediante la función lambda con código python, se realiza una llamada de prueba al endpoint del modelo generado en SageMaker.
• S3: Servicio de almacenamiento de objetos, donde se crea un bucket para guardar los datasets de entrenamiento, scoring y el modelo ML entrenado por SageMaker y que se usará en las predicciones. Además, se crea un bucket para almacenar los logs generados por CloudTrail, donde se ven los eventos generados sobre el bucket de S3 con los datos de entrada del modelo ML.
• CloudWatch Event Rule: Se crea una regla en CloudWatch para identificar cuándo se realiza un evento PutObject sobre el bucket de S3, donde se almacenan los datos de entrada para el entrenamiento de un modelo ML, de tal manera, que al ejecutar dicho evento se inicia la ejecución de la Step Function automáticamente.
• CloudTrail: Se crea un trail para registrar todos los eventos de la API sobre el bucket de S3 donde se almacenan los datos de entrada para el entrenamiento, permitiendo que CloudWatch identifique cuando se realiza el evento PutObject sobre dicho bucket y ejecute la regla asociada al evento mencionado, logrando que se ejecute la Step Function automáticamente.
• ECR: Repositorio de contenedores Docker con las imágenes de los modelos ML, y a los que SageMaker tendrá acceso para realizar el entrenamiento y despliegue de los mismos.
• SageMaker: Servicio que permite la creación, entrenamiento y despliegue de modelos ML de forma rápida. Específicamente, dentro de este servicio, se hará uso de los componentes: Model, Endpoint Configuration y Endpoint.
• Rol: Se debe crear un rol para dar permisos a SageMaker para que obtenga objetos almacenados en S3 y, que realice operaciones sobre ECR para obtener las imágenes de Docker con los modelos ML. Dicho rol también tendrá los permisos necesarios para ejecutar las acciones sobre todos los demás servicios de AWS indicados en este post: Lambda, Step Function, CloudTrail, CloudWatch, etc.
• Política: Se crea una política con los diferentes permisos mencionados anteriormente y que permita al rol ejecutar todas las acciones necesarias sobre los servicios de AWS.

Docker

Herramienta para la generación de contenedores en base a los modelos ML, permitiendo su portabilidad a cualquier plataforma On-premise o cloud.

A grandes rasgos, debemos generar nuestro contenedor Docker con nuestro modelo ML y una jerarquización de carpetas previamente configurada, tal como se puede ver a continuación.

Jerarquización

raiz/
|-- Dockerfile
|-- build_push_ecr_image.sh
|-- model_src/
|   |-- requirements.txt
|   `-- serve
|   `-- train
|-- test_dir/
|   |-- input/
|   |   |-- config/
|   |   |   |-- hyperparameters.json
|   |   |   `-- resourceConfig.json
|   |   `-- data/
|   |        `-- <training>/
|   |               `-- <input data>
|   |-- model/
|   |     `-- <model files>
    `-- output/
           `-- failure

• Dockerfile: Describe cómo debe realizarse el build del contenedor.
• build_push_ecr_image: script sh que utiliza el Dockerfile para construir las imágenes en un contenedor y luego realizar el push a ECR de AWS.
• Requirements.txt: Fichero con las dependencias del proyecto.
• Serve: Programa en python que se ejecuta cuando realizamos el host del modelo. Inicia una instancia flask definida dentro del propio fichero o en un fichero python externo.
• Train: se utiliza cuando el contenedor es iniciado en modo entrenamiento. Este fichero debe cambiarse por el algoritmo de entrenamiento que queramos desarrollar.

Podemos ver de forma más detalla la configuración necesaria en la documentación de AWS en este enlace: Sagemaker-scikit_bring_your_own.

Dockerfile

FROM python:3.6

# Get model name from args (we can have different prediction models in our stack and choose what model we want to deploy)
ARG MODEL_FOLDER_NAME=model_src
ARG DATA_INPUT=test_dir 
ARG SERVE=model_src/serve

# Set some environment variables. PYTHONUNBUFFERED keeps Python from buffering our standard
# output stream, which means that logs can be delivered to the user quickly. PYTHONDONTWRITEBYTECODE
# keeps Python from writing the .pyc files which are unnecessary in this case. We also update
# PATH so that the train and serve programs are found when the container is invoked. And set FLASK_APP to launch
# the scoring service
ENV PYTHONUNBUFFERED=TRUE
ENV PYTHONDONTWRITEBYTECODE=TRUE
ENV PROGRAM_PATH=/opt/program
ENV PATH=$PROGRAM_PATH:$PATH
ENV FLASK_APP=$PROGRAM_PATH/serve
ENV PATH=$DATA_INPUT:$PATH

# Set up the program in the image
COPY $MODEL_FOLDER_NAME $PROGRAM_PATH
WORKDIR $PROGRAM_PATH

#RUN apt-get -y update && apt-get install -y --no-install-recommends wget ca-certificates && rm -rf /var/lib/apt/lists/*

# Install libraries needed for training and scoring
RUN pip install -r requirements.txt

# Env vars needed for training
ENV PATH_PREFFIX=/opt/ml/
ENV TRAINING_ERROR_PATH=output/failure
ENV TRAINING_MODEL_PATH=model/model.pkl
ENV TRAINING_PARAMS_PATH=input/config/hyperparameters.json

COPY $DATA_INPUT $PATH_PREFFIX 

# Env var needed to be overwritten by the workflow with file path to train
ENV TRAINING_INPUT_DATA_PATH=input/data/training/iris_final.csv

Script SH - build and push

#!/usr/bin/env bash

docker_image=${1:-aws-docker-test}
docker_repositorie=${2:-aws-docker-test}
docker_path=${3:-../container/}
model_folder_name=${4:-model_src}
data_input=${5:-test_dir}

if [ "$docker_image" == "" ]
then
    echo "Usage: $0 <image-name>"
    exit 1
fi

chmod +x model_src/train
chmod +x model_src/serve

# construct the ECR name.
account=$(aws sts get-caller-identity --profile mlops-sagemaker --query Account --output text)
region=$(aws configure get region --profile mlops-sagemaker)
fullname="${account}.dkr.ecr.${region}.amazonaws.com/${docker_repositorie}:latest"

aws ecr describe-repositories --repository-names "${docker_repositorie}" --profile mlops-sagemaker > /dev/null 2>&1

# Check the error code, if it's non-zero then know we threw an error and no repo exists
if [ $? -ne 0 ]
then
    aws ecr create-repository --repository-name "${docker_repositorie}" --profile mlops-sagemaker > /dev/null
fi

sudo $(aws ecr get-login-password --region eu-west-1 | docker login --username AWS --password-stdin {user_name}.dkr.ecr.eu-west-1.amazonaws.com)

# Build the docker docker_image, tag it with the full name, and push it to ECR
docker build -t "${docker_image}" --build-arg MODEL_FOLDER_NAME="${model_folder_name}" --build-arg DATA_INPUT="${data_input}" ../container/

docker tag "${docker_image}" ${fullname}

docker push ${fullname

Terraform

Terraform es la herramienta con la que se define la Infrastructure as Code (IaC), es decir, la infraestructura sobre la cual se realizará el entrenamiento, validación y despliegue del modelo ML.

Dado que Terraform es Open Source, su uso es muy fácil y extendido, permitiendo la definición de la IaC con lenguaje de alto nivel HCL, además no está ligada a ningún proveedor Cloud. Para saber más sobre las cualidades de esta herramienta aquí tienes un post muy interesante.

Todos los servicios de AWS mencionados anteriormente se crean a partir de ficheros .tf, permitiendo que la creación de los recursos asociados a dichos servicios se pueda realizar en múltiples entornos de manera rápida y consistente, y manteniendo la trazabilidad de los cambios de la IaC en todo momento, a través del fichero tfstate, el cual se almacena en bucket de S3.

Vamos a explorar en qué consisten varios de los ficheros donde se definen los recursos a utilizar en AWS.

• Código HCL para crear la Step Function a partir de fichero json externo con la definición del flujo de trabajo, al que se le pueden pasar parámetros.

data "template_file" "sf_definition" {
  template = "${file(path_fichero_definicion_step_function)}"

  vars = {
    sagemaker_job_name       = "sagemaker_training_job_prueba"
    sagemaker_instance_count = 1
    sagemaker_volume_size_GB = 10
    sagemaker_s3_input_path  = "nombre_bucket_dataset_entrenamiento"
    sagemaker_s3_output_path = "nombre_bucket_resultados_entrenamiento"
    instance_training_type   = "ml.m5.4xlarge"
    role_arn                 = "arn_rol_con_permisos"
    ecr_image_repository     = "nombre_repositorio_ecr"
  }
}

resource "aws_sfn_state_machine" "create_state_machine" {
  name     = "sfn_state_machine_name"
  role_arn = "arn_rol_con_permisos"
  definition = data.template_file.sf_definition.rendered

  lifecycle {
    prevent_destroy = true
  }
}

• Código HCL para crear trail de CloudTrail.

resource "aws_cloudtrail" "trail_s3_event" {
  name           = "trail_s3_event"
  s3_bucket_name = "id_bucket_dataset_entrenamiento"

  include_global_service_events = false

  event_selector {
    read_write_type           = "All"
    include_management_events = true

    data_resource {
      type   = "AWS::S3::Object"
      values = ["arn_bucket_dataset_entrenamiento/training-data"]
    }
  }
}

• Código HCL para crear CloudWatch Events Rule.

resource "aws_cloudwatch_event_rule" "s3_event_rule" {
    name        = "event_rule_name"
    description = "Capture s3 event to trigger event target"

    event_pattern = <<PATTERN
  {
    "source": [
      "aws.s3"
    ],
    "detail-type": [
      "AWS API Call via CloudTrail"
    ],
    "detail": {
      "eventSource": [
        "s3.amazonaws.com"
      ],
      "eventName": [
        "PutObject"
      ],
      "requestParameters": {
        "bucketName": [
          "nombre_bucket_dataset_entrenamiento"
        ]
      }
    }
  }
  PATTERN
  }

  resource "aws_cloudwatch_event_target" "step_function_target" {
    role_arn  = "arn_rol_con_permisos"
    rule      = aws_cloudwatch_event_rule.s3_event_rule.name
    target_id = "SendToStepFunction"
    arn       = "arn_step_function"
  }

Una vez tengamos definidos todos los recursos de AWS necesarios para entrenar y desplegar un modelo ML, estos se pueden agrupar por módulos (tal como se puede ver en la imagen de abajo) permitiendo que puedan ser reutilizados para usarlos con múltiples modelos ML y/o crear recursos de manera consistente y repetible en los diferentes entornos donde se quiera desplegar un modelo ML.

Dichos módulos se deben “instanciar” en el módulo root llamado_ main.tf_, que será utilizado para crear todos los recursos en AWS mediante los comandos plan y apply.

Mediante el uso de workspaces es fácil desplegar la IaC en múltiples entornos sin necesidad de repetir código, ya que con una única definición en el módulo root, citado anteriormente, se puede lograr esto, manteniendo la trazabilidad y estado actual de la IaC en cada entorno por separado.

• Código HCL con las instancias a los módulos y la configuración del backend del fichero tfstate.

# Provider's definition 
provider "aws" {
  version = "~> 2.61"
  profile = "terraform_deploy"
  region  = "eu-west-1"
}

# Mandatory create previously backend's bucket-tfstate to stores tfstate file with real world resources deployed on AWS.
terraform {
  required_version = ">= 0.12"
  backend "s3" {
    bucket  = "bucket-tfstate"
    key     = "iac.tfstate"
    region  = "eu-west-1"
    encrypt = "true"
    profile = "terraform_deploy"
  }
}

#  Module for create Step Function and underlying SageMaker resources 
module "step_function_model_001" {
  source = "./../modules/step_functions"
}

#  Module for create CloudTrail resource 
module "cloud_trail_model_001" {
  source = "./../modules/cloud_trail"
}

#  Module for create CloudWatch Event Rule
module "cloudwatch_model_001" {
  source = "./../modules/cloudwatch"
}

Orquestación del proceso MLOps

Una vez que tenemos todas las piezas del puzzle, es decir, el código del modelo ML almacenado en el repositorio Git (p. ej. GitLab), el fichero Dockerfile para generar la imagen docker, la definición de la infraestructura de AWS necesaria para entrenar y desplegar el modelo, es momento de hacer que las piezas funcionen en conjunto.

Para lograr esto, vienen los ¡Pipelines de Jenkins al rescate!

Mediante la implementación de un pipeline compuesto por varias fases, como la generación de la imagen docker del modelo ML, hacer push de dicha imagen al repositorio ECR, generar la infraestructura definida en Terraform y validar que el endpoint del modelo se ha generado correctamente, estos pasos se pueden ejecutar automáticamente, respondiendo típicamente a un evento como un merge/pull request en la rama master del repositorio.

Un ejemplo del fichero Jenkinsfile escrito en groovy con la definición del pipeline podría ser el siguiente:

pipeline {
  agent any
  environment {
    TF_WORKSPACE = 'pro' //Sets the Terraform Workspace
    TF_IN_AUTOMATION = 'true'
  }
  stages {
    stage('Clone Repo') {
       steps {
          bat 'set'
          sh 'git clone --depth=1 <usuario@url_git_repo/nombre_repo.git>'
       }
    }
    stage('Build and Push Docker Image ML to ECR') {
       steps {
          sh 'cd mlops/cloud/docker/container && ./build_push_ecr_image.sh'
       }
    }
    stage('TF Plan Core IaC') {
       steps {
          sh 'cd mlops/IaC/environments/core && terraform init && terraform plan -var-file=core.tfvars'
       }
    }
    stage('TF Apply Core IaC') { //Deploy IaC core (role, policy, s3 bucket)
      steps {
          sh 'cd mlops/IaC/environments/core && terraform apply -auto-approve -var-file=core.tfvars'
      }
    }
    stage('TF Plan Remaining IaC') {
       steps {
         sh  'cd mlops/IaC/environments && terraform init && terraform plan -var-file=./${TF_WORKSPACE}/${TF_WORKSPACE}.tfvars'
       }
    }
    stage('TF Apply Remaining IaC') { //Deploy CloudTrail, CloudWatch, StepFunction...
      input{ //User must review and approve the IaC’s deploy
        message 'Apply Terraform?'
      }
      steps {
         sh 'cd mlops/IaC/environments && terraform apply -auto-approve -var-file=./${TF_WORKSPACE}/${TF_WORKSPACE}.tfvars'
      }
    }
  }
}

Ejecución del proceso MLOps

Ya tenemos el despliegue de todos los componentes necesarios para el entrenamiento de nuestro modelo, i.e., el modelo contenido y subido a un contenedor docker en AWS (ECR), para poder empezar la ejecución orquestada desde la Step Function.

Si Terraform nos sirve como orquestador para la creación y despliegue de la arquitectura deseada, Step Functions será nuestro orquestador de los servicios AWS que vamos a utilizar dentro de la propia plataforma.

Step Functions ejecutará de manera sincrona o asincrona los servicios para entrenar, guardar y desplegar el modelo, además de llamar a la función que realiza una predicción sobre el mismo.

Para que sea posible la comunicación entre los distintos servicios a través de la Step Function, debemos crear los roles con las políticas de acceso y ejecución (permisos mínimos) para: Sagemaker, Lambda, S3, etc.

El desencadenante de la ejecución en Step Functions es el evento PutObject de S3, que se realiza sobre el bucket donde se almacenará el fichero csv con los datos de entrada del entrenamiento del modelo ML. Cuando dicho fichero “iris_final.csv” se sube al bucket, se genera un evento PutObject, el cual es identificado y registrado por CloudTrail y, a continuación, CloudWatch ejecutará la regla que invoca a la Step Function en respuesta al evento mencionado sobre el bucket de S3. Aquí puedes consultar más información sobre cómo ejecutar Step Function en respuesta a un evento de S3.

Los pasos que se ejecutan en la Step Function se describen a continuación:

• Entrenamiento del modelo: Fase de entrenamiento del modelo a través de Sagemaker mediante un TrainingJob, el cual utiliza el contenedor ECR, donde se aloja la imagen Docker con nuestro modelo para realizar el entrenamiento. En este punto, si en el código hemos definido hiperparámetros y métricas, podemos configurar la llamada desde Step Functions para incluirlas (Train a Model with Amazon SageMaker). El recurso utilizado en la Step Function es sagemaker:createTrainingJob.

• Guardar el Modelo: Una vez culminado el entrenamiento, se procede a generar el Modelo en SageMaker, que lo guarda en el bucket de S3. Dicho modelo será utilizado por los jobs de transformación batch y online para realizar las inferencias en base a los datos de entrada. El recurso utilizado en la Step Function es sagemaker:createModel.

• Generar Job de transformación batch: La transformación batch nos permite realizar una única inferencia sobre el modelo generado en SageMaker. La predicción a través de la transformación batch se realiza con el modelo y los datos de entrada que queremos predecir. Dichos datos de entrada pueden ser transformados previamente o no. Por último, para ejecutar el trabajo debemos especificar el número y el tipo de instancias que vamos a utilizar. El recurso utilizado en Step Functions es sagemaker:createTransformJob. El resultado obtenido de la inferencia será guardado en un fichero plano en un bucket S3.

• Generar Endpoint para inferencia On-line: Endpoint que habilita la integración del modelo ML con aplicaciones externas mediante API, permitiendo que se realicen inferencias en tiempo real. Para la creación del Endpoint debemos anidar 2 pasos diferenciados: la creación de la configuración y la creación del endpoint propiamente dicho. Endpoint configuration, como su propio nombre indica, es el paso donde debemos especificar la configuración con la que vamos iniciar nuestro punto de enlace: modelo a utilizar, el número y tipo de instancias, y finalmente el variant name. Una vez tenemos la configuración podemos generar nuestro Endpoint con el nombre final que vamos a utilizar. Una vez levantado el punto de enlace podemos realizar peticiones al mismo. Los recursos utilizados en este paso son: sagemaker:createEndpointConfig y sagemaker:createEndpoint.

• Prueba del endpoint: Creamos una función Lambda que contiene una llamada en Python a la API generada en SageMaker. Esta lambda recogerá los parámetros de entrada, bien recogiendo los parámetros dentro del propio cuerpo de la Step Function (json), o bien como parámetros en la ejecución de la máquina de Lambda. En ambos casos, la función lambda debe recibir los datos de entrada para que el endpoint pueda tratar la llamada y devolver una predicción y, en ambos casos, la función debe estar creada antes de la ejecución, por ejemplo, creando la función mediante Terraform. El resultado de la predicción se aloja en un bucket S3.

Una vez finalizada la ejecución de todos los pasos definidos en la Step Function, el modelo ML ya está listo para ser utilizado.

Conclusiones

Adoptar la filosofía de MLOps es muy importante hoy en día. Así como DevOps ha ayudado a facilitar la entrega continua de software tradicional a los usuarios, MLOps hace los mismo con los modelos de Machine Learning pero manteniendo un seguimiento continuo de los mismos, a través de métricas, para identificar alguna degradación en la precisión de los resultados y actuar en consecuencia.

Hemos podido comprobar lo sencillo que puede llegar a ser desplegar un modelo en AWS. Con un poco de conocimiento de la plataforma podemos implementar una solución práctica que se adhiera a nuestros requerimientos. Además, con las plataformas cloud las limitaciones como servidores, capacidad, escalabilidad no son un problema; solamente debes pagar por lo que utilizas.

El proceso MLOps descrito en este post no es agnóstico a los servicios propios que utilizan las diversas plataformas en la nube; sin embargo, es una solución eficaz que sirve como esqueleto para la implementación y réplica de los mismos servicios e infraestructuras en diferentes áreas, modificando solamente una pequeña parte del script o de los parámetros que se deben embeber.

Existen herramientas como Mlflow, que aborda desde otro enfoque el uso de MLOps, permitiendo desplegar modelos en On-premise o en Cloud, sin necesidad de usar servicios como Sagemaker, comparar los resultados de varias ejecuciones del entrenamiento de un modelo y decidir cuál es mejor.

Hay que aprovechar al máximo los diversos servicios administrados que ofrece AWS y que hemos descrito en este post, ya que facilitan la implementación de un proceso de MLOps, incluso más amplio del que hemos intentado explicar. Solo hay que saber en qué consisten los servicios que se quieren usar y entender el modo de integrarlos entre sí y, una vez hecho esto, los resultados que se obtienen son muy buenos a la hora de poner en producción modelos ML para su uso en la solución de problemas.