Machine Learning Operations: practicando con Sagemaker Pipelines.

Como ya os contamos, se denomina MLOps, al conjunto de buenas prácticas, reglas y normativas que permiten automatizar el ciclo de vida de los modelos construidos mediante Machine Learning a través de la creación de canalizaciones sobre el software, los datos y los modelos. La aplicación de este tipo de buenas prácticas permite la completa gobernabilidad de los modelos con el objetivo de analizar su funcionamiento y detectar posibles puntos de errores o de mejora de manera sencilla y ágil. Dentro del ecosistema de AWS existe un servicio denominado Sagemakerque ofrece los componentes básicos para permitir el diseño, desarrollo, despliegue y monitorización de los modelos. Entre los diferentes componentes, existe uno denominado Sagemaker Pipelines que permiten construir las canalizaciones necesarias para orquestar de manera sencilla y escalable las diferentes operaciones que componen el ciclo de vida de los modelos. En este post hablaremos de cómo crear canalizaciones de Machine Learning usando Sagemaker Pipelines y qué aspectos hay que tener en cuenta para su desarrollo.

Sagemaker Pipelines

Sagemaker es el servicio principal para la aplicación de técnicas de Machine Learning y Deep Learning en el ecosistema de AWS ya que proveía de los componentes de infraestructura necesarios para la ejecución de proyectos de Machine Learning (como, por ejemplo, la ejecución de procesos de entrenamiento o Jupyter notebook para la construcción de scripts), además de ofrecer su propio SDK para la gestión de los modelos. No obstante, no permitía automatizar el ciclo de vida de los modelos así como definir ciclos alternativos dependiendo de la complejidad que se le quiera aplicar al mismo, desde un ciclo que sólo implementa las operaciones de entrenamiento y despliegue, hasta uno que implementa todas las operaciones. Con el objetivo de ofrecer un sistema robusto y sencillo para la construcción del ciclo de vida de un modelo mediante la utilización de canalizaciones, Sagemaker ha incluido como parte de sus funcionalidades Sagemaker Pipelines.

Las canalizaciones (Pipelines) son secuencias de operaciones (procesos) del ciclo de vida de los modelos que se ejecutan de manera secuencial o paralela y que son definidas mediante un Grafo Acíclico Dirigido (Directed Acyclic Graph, DAG). Este grafo es utilizado para definir el orden de ejecución de las diferentes operaciones a lo largo de la canalización. Donde cada una de las operaciones de la canalización se corresponde con una ejecución de un proceso que es empaquetado normalmente mediante un contenedor Docker donde los diferentes tipos de operaciones pueden ser procesos de entrenamiento, de inferencia, de procesamiento, de explicabilidad, etc. A continuación, en la figura 1 se muestran los componentes básicos del ecosistema AWS que intervienen en la implementación del ciclo de vida de los modelos mediante Sagemaker Pipelines.

Además, una de las grandes ventajas que ofrece Sagemaker Pipelines frente a otros servicios de tipo MLOps es que ofrece una integración completa de manera automática con los diferentes servicios del ecosistema AWS, como, por ejemplo, S3 para crear volúmenes de datos utilizados en las operaciones de la canalización, ECR para acceder a imágenes Docker, CloudWatch para monitorizar los procesos de la canalización, etc.

Construyendo mi canalización

Para poder crear una canalización es necesario crear sus componentes básicos, es decir, las operaciones. Para poder implementar las diferentes operaciones de procesamiento, transformación, entrenamiento o inferencia dentro de una canalización es necesario crear una instancia de Estimator o Transformer (respectivamente) de Sagemaker y a continuación incluirla como una operación (step) dentro de la canalización de Sagemaker. Por ejemplo, si quisiéramos crear una canalización como la que se muestra en la figura 2.

Deberíamos crear tres operaciones: (1) entrenamiento, (2) registro y gestión del modelo; e (3) inferencia. Aunque, a la hora de implementar esta canalización sobre Sagemaker Pipelines es necesario implementar cuatro operaciones cuyo proceso de desarrollo se describe a continuación:

Paso 1 - Sesión y roles

Para poder crear, orquestar y definir procesos en Sagemaker es necesario utilizar un rol de Sagemaker con la policy AWSSagemakerFullAccess. Además, todos los objetos instanciados para la ejecución de un proceso deben estar definidos bajo la misma sesión de Sagemaker. Esta sesión se puede crear como se muestra en el siguiente fragmento de código.

import boto3
import pandas as pd
import numpy as np
import sagemaker
from sklearn.datasets import load_iris

sesion = sagemaker.Session()

Paso 2 - Entrenamiento del modelo

La operación más importante del ciclo de vida de los modelos es el entrenamiento, para poder crear una operación de entrenamiento en SageMaker Pipelines se debe crear un objeto de la clase SKLearn, que es de tipo Estimator, como se muestra en el siguiente fragmento de código:

from sagemaker.sklearn.estimator import SKLearn

sklearn_estimator = SKLearn(
    source_dir='./src',
    entry_point='train.py',
    framework_version='0.23-1',
    instance_type="ml.c4.xlarge",
    role=role,
    sagemaker_session=sesion,
    hyperparameters={
        "min_leaf_nodes": 3,
        "n_estimators": 10,
        "target": "Species"
    }
)

Esta clase permite construir un proceso que se construirá en un contenedor que incluye las diferentes librerías de ScikitLearn, donde los elementos principales a configurar son:

• source_dir que se corresponde con la localización dentro del contenedor de nuestros scripts python.
• entry_point que se corresponde con nuestro código de entrenamiento (train.py) del modelo mediante Scikit-Learn.
• instance_type que se corresponde con el tipo de instancia de tipo EC2 donde se desplegará el contenedor.

Sagemaker provee además de otros tipos de clases de tipo Estimator preparados para la utilización de otros frameworks de ML (Tensorflow, PYTorch...) cuya información detallada se puede consultar en el siguiente enlace. La clase Estimator de Sagemaker permite utilizar configuraciones predefinidas para la construcción de los contenedores que ejecutarán procesos, es decir, que no será necesario definir las librerías que deben ser instaladas durante la fase de creación del contenedor.

Una vez hemos creado el Estimator es necesario encapsular el proceso de entrenamiento como una operación de la canalización, es decir, un step de Sagemaker Pipelines, mediante la creación de un objeto de la clase TrainingStep como se puede observar en el siguiente fragmento de código.

from sagemaker.inputs import TrainingInput
from sagemaker.workflow.steps import TrainingStep

train_step = TrainingStep(
    name="TrainStep",
    estimator=sklearn_estimator,
    inputs={
        "train": TrainingInput(s3_data=s3_data_uri, content_type="text/csv"),
    },
)

Esto objeto, además de convertir el proceso de entrenamiento como una operación (step) de la canalización, también permite definir volúmenes a archivos de S3 de entrada y salida para lectura y escritura de datos por medio de los objetos “TrainingInput”. Aunque a la hora de incluirlos es necesario seguir una serie de consideraciones:

• Todos los volúmenes que se definan por medio de los objetos TrainingInput, ProcessingInput, ProcessingOutput tienen que apuntar a un ruta dentro de la ruta principal “/opt/ml/”.
• El almacenamiento del modelo dentro del código de entrenamiento se tiene que hacer en la ruta “/opt/ml/model”, ya que los procesos de entrenamiento creados por medio del objeto TrainingStep, comprimen el contenido que se guarda en ese directorio y lo suben a S3 dentro de los logs del proceso.
• Del mismo modo, los procesos de inferencia creados posteriormente con un TransformStep crean un volumen del modelo entrenado a “/opt/ml/model” y descomprimen el contenido de esta ruta automáticamente (ya que el proceso de entrenamiento ha comprimido el modelo antes de subirlo).

Paso 3 - Registro y gestión del modelo

Además de los procesos de entrenamiento y despliegue de los modelos es necesario incluir operaciones relacionadas con la creación y el registro de los modelos. Sagemaker Pipelines ofrece operadores específicos para el registro de los modelos generados por los procesos de entrenamiento que, además, permite añadir metadatos propios de las características de la ejecución o resultados del entrenamiento. Este tipo de información es útil para poder gobernar el modelo y poder auditar los resultados de las diferentes versiones y/o parámetros de la operación de entrenamiento. Para ello existen dos tipos de operaciones (steps) dentro de Sagemaker Pipelines:

• CreateModelStep:

Dado un proceso de entrenamiento con TrainingStep que ha generado un artefacto de salida que se corresponde con un modelo es necesario ejecutar un proceso de identificación para que Sagemaker lo reconozca como un modelo y lo pueda utilizar durante la operación de inferencia, para lo cual se debe crear un objeto de tipo SKLearnModel que es de tipo Model. Para ello es necesario ejecutar el fragmento de código, presentado a continuación, para identificar el artefacto resultante como un modelo.

from sagemaker.inputs import CreateModelInput
from sagemaker.workflow.steps import CreateModelStep
from sagemaker.sklearn.model import SKLearnModel

model = SKLearnModel(
    model_data=step_train.properties.ModelArtifacts.S3ModelArtifacts,
    source_dir='./src',
    entry_point='train.py',
    framework_version='0.23-1',
    role=role,
    sagemaker_session=sesion,
)

Para la creación del modelo en la canalización es necesario encapsular el objeto SKLearnModel dentro de un objeto CreateModelStep.

create_model_step = CreateModelStep(
    name="CreateModel",
    model=model,
    inputs = CreateModelInput(instance_type="ml.m5.large")
)

• RegisterModel:

Dado un modelo ya creado e identificado es necesario registrarlo con el objetivo de que pueda ser utilizado de manera correcta incluyendo toda la información relativa a su creación y ejecución (como por ejemplo, si pertenece a un grupo nuevo de modelos o es una versión nueva del mismo, sobre qué tipo de máquinas se ejecuta, qué métricas de entrenamiento han sido utilizadas, etc.). A continuación, se presenta un fragmento de código para realizar el registro de un modelo:

from sagemaker.workflow.step_collections import RegisterModel

register_step = RegisterModel(
    name="RegisterModel",
    model=model,
    content_types=["text/csv"],
    response_types=["text/csv"],
    inference_instances=["ml.t2.medium", "ml.m5.xlarge"],
    transform_instances=["ml.m5.xlarge"],
    model_package_group_name="SklearnIris",
    approval_status="Approved",
)

Paso 4 - Inferencia del modelo

Para poder crear un proceso de inferencia es necesario crear un objeto de tipo Transformer. A diferencia del objeto Estimator, en el que teníamos que pasarle como argumento el script de entrenamiento, el objeto Transformer no requiere del desarrollo de ningún script para su ejecución. Simplemente usando el modelo que se ha creado con el objeto CreateModelStep de tipo SKLearnModel es capaz de cargar el modelo y ejecutarlo automáticamente, como se puede observar en el siguiente fragmento de código:

from sagemaker.transformer import Transformer

transformer = Transformer(
    model_name=step_create_model.properties.ModelName,
    instance_type="ml.m5.xlarge",
    accept="text/csv",
    instance_count=1,
    output_path=output_data
)

Los argumentos básicos a configurar son:

• model_name: En este argumento le tenemos que indicar el path de S3 al artefacto del modelo generado por un proceso de entrenamiento que hemos visto en el Paso 1. Como aún no se ha ejecutado el proceso de entrenamiento, este path aún no existe, pero gracias al SDK de Sagemaker Pipelines, podemos referenciarlo con la propiedad “ModelName” del objeto CreateModelStep definido en el paso anterior.
• accept: En este argumento tenemos que especificar el tipo de formato que van a tener los datos de entrada. En nuestro caso va a ser un csv, como los datos de entrenamiento.
• output_path: Este argumento indica la ruta (path) de S3 donde se quiere que se guarden los resultados de la predicción.

Una vez hemos creado el objeto del tipo Transformer también es necesario encapsular el proceso de inferencia como una operación de la canalización, es decir, un step de Sagemaker Pipelines, mediante la creación de un objeto de la clase TransformStep como se puede observar en el siguiente fragmento de código.

from sagemaker.inputs import TransformInput
from sagemaker.workflow.steps import TransformStep

transform_step = TransformStep(
    name="TransformStep", 
    transformer=transformer,
    inputs=TransformInput(data=s3_path, content_type="text/csv")
)

De manera análoga a como hemos visto en el caso del entrenamiento, por medio de un objeto de la clase TransformInput podemos definir el dataset sobre el que queremos hacer la inferencia. Es importante tener en cuenta que este dataset deberá tener las mismas columnas que el usado para el entrenamiento salvo la variable de target.

Paso 5 - Construcción de la canalización

Una vez que se han creado las diferentes operaciones de la canalización es posible construir la canalización mediante la creación de un objeto de tipo Pipeline, al que se agregan las diferentes operaciones que hemos creado previamente como se muestra en el siguiente fragmento de código:

from sagemaker.session import Session
from sagemaker.workflow.pipeline import Pipeline

pipeline = Pipeline(
    name='pipeline',
    steps=[
        train_step,
        create_model_step,
        register_step,
        transform_step
    ],
    sagemaker_session=Session(),
)

pipeline.upsert(role_arn=args.role_arn)
pipeline.start()

De esta forma hemos creado una canalización formada por las operaciones de entrenamiento, registro y despliegue del modelo. Esta canalización (pipeline) tiene una estructura completamente secuencial, no obstante, se pueden definir pipelines más complejas, con diferentes caminos usando ConditionalStep, que permite evaluar que siguiente paso de la canalización usar en función de diferentes métricas, como el accuracy del modelo. Además las canalizaciones desarrolladas, así como sus metadatos se puede ver por medio de la herramienta Sagemaker Studio como se muestra en la siguiente imagen.

Conclusiones

Sagemaker Pipelines es uno de los nuevos servicios de AWS relacionados con Inteligencia Artificial que permite diseñar, desarrollar y ejecutar canalizaciones de Machine Learning sobre Sagemaker. Esto no solo supone una simplificación de los procesos de despliegue, sino que también ofrece gobierno sobre los modelos registrando y guardando los artefactos generados en cada proceso en S3 de manera automática.

A pesar de su simplicidad, Sagemaker no simplifica el proceso de migración de proyectos previos de Machine Learning, ya que los pasos de la canalización usan objetos a alto nivel de Sagemaker, como los Estimator, que normalmente implican modificaciones en el código para adaptarlo al modelo de funcionamiento de SageMaker Pipelines. La solución en estos casos pasa por construir tu propia imagen Docker con todas las dependencias y el código de proyecto (como ya veremos en futuros artículos del blog).