Implementando 🤗 ViT en Kubernetes con TF Serving

'Implementing 🤗 ViT in Kubernetes with TF Serving'

En el post anterior, mostramos cómo implementar un modelo Vision Transformer (ViT) de 🤗 Transformers localmente con TensorFlow Serving. Cubrimos temas como la preprocesamiento y postprocesamiento de incrustaciones dentro del modelo Vision Transformer, manejo de solicitudes gRPC y más.

Aunque las implementaciones locales son un excelente punto de partida para construir algo útil, necesitarías realizar implementaciones que puedan atender a muchos usuarios en proyectos reales. En este post, aprenderás cómo escalar la implementación local del post anterior con Docker y Kubernetes. Por lo tanto, asumimos cierta familiaridad con Docker y Kubernetes.

Este post se basa en el post anterior, por lo que recomendamos encarecidamente leerlo primero. Puedes encontrar todo el código discutido a lo largo de este post en este repositorio.

El flujo de trabajo básico para escalar una implementación como la nuestra incluye los siguientes pasos:

• Contenerizar la lógica de la aplicación: La lógica de la aplicación implica un modelo servido que puede manejar solicitudes y devolver predicciones. Para la contenerización, Docker es la opción estándar de la industria.

  • Una Introducción Suave a la Multiplicación de Matrices de 8 bits pa...
  • Desplegando 🤗 ViT en Vertex AI
  • Preentrena BERT con Hugging Face Transformers y Habana Gaudi

• Implementar el contenedor Docker: Aquí tienes varias opciones. La opción más utilizada es implementar el contenedor Docker en un clúster de Kubernetes. Kubernetes proporciona numerosas características amigables para la implementación (por ejemplo, escalado automático y seguridad). Puedes utilizar una solución como Minikube para gestionar clústeres de Kubernetes localmente o una solución sin servidor como Elastic Kubernetes Service (EKS).

Es posible que te preguntes por qué utilizar una configuración explícita como esta en la era de Sagemaker, Vertex AI, que proporciona características específicas de implementación de ML desde el principio. Es justo pensarlo.

El flujo de trabajo anterior es ampliamente adoptado en la industria y muchas organizaciones se benefician de él. Ya ha sido probado en batalla durante muchos años. También te permite tener un control más granular de tus implementaciones al abstraer los aspectos no triviales.

Este post utiliza Google Kubernetes Engine (GKE) para aprovisionar y gestionar un clúster de Kubernetes. Asumimos que ya tienes un proyecto de GCP habilitado para facturación si estás utilizando GKE. Además, ten en cuenta que necesitarás configurar la utilidad gcloud para realizar la implementación en GKE. Pero los conceptos discutidos en este post también se aplican si decides utilizar Minikube.

Nota: Los fragmentos de código mostrados en este post se pueden ejecutar en una terminal Unix siempre y cuando hayas configurado la utilidad gcloud junto con Docker y kubectl. Se encuentran disponibles más instrucciones en el repositorio adjunto.

El modelo de servicio puede manejar entradas de imágenes sin procesar como bytes y es capaz de preprocesamiento y postprocesamiento.

En esta sección, verás cómo contenerizar ese modelo utilizando la imagen base de TensorFlow Serving. TensorFlow Serving consume modelos en formato SavedModel. Recuerda cómo obtuviste dicho SavedModel en el post anterior. Asumimos que tienes el SavedModel comprimido en formato tar.gz. Puedes obtenerlo aquí por si acaso. Luego, el SavedModel debe colocarse en la estructura de directorios especial de <NOMBRE_MODELO>/<VERSION>/<SavedModel>. De esta manera, TensorFlow Serving gestiona simultáneamente múltiples implementaciones de modelos con diferentes versiones.

Preparando la imagen de Docker

El script de shell a continuación coloca el SavedModel en hf-vit/1 dentro del directorio principal models. Copiarás todo lo que hay dentro de él al preparar la imagen de Docker. En este ejemplo, solo hay un modelo, pero este enfoque es más generalizable.

$ MODEL_TAR=model.tar.gz
$ MODEL_NAME=hf-vit
$ MODEL_VERSION=1
$ MODEL_PATH=models/$MODEL_NAME/$MODEL_VERSION

$ mkdir -p $MODEL_PATH
$ tar -xvf $MODEL_TAR --directory $MODEL_PATH

A continuación, mostramos cómo está estructurado el directorio models en nuestro caso:

$ find /models
/models
/models/hf-vit
/models/hf-vit/1
/models/hf-vit/1/keras_metadata.pb
/models/hf-vit/1/variables
/models/hf-vit/1/variables/variables.index
/models/hf-vit/1/variables/variables.data-00000-of-00001
/models/hf-vit/1/assets
/models/hf-vit/1/saved_model.pb

La imagen personalizada de TensorFlow Serving debe construirse sobre la base de la imagen base. Hay varios enfoques para esto, pero lo harás ejecutando un contenedor Docker como se ilustra en el documento oficial. Comenzamos ejecutando la imagen tensorflow/serving en modo de segundo plano, luego se copia el directorio completo models al contenedor en ejecución de la siguiente manera.

$ docker run -d --name serving_base tensorflow/serving
$ docker cp models/ serving_base:/models/

Utilizamos la imagen oficial de Docker de TensorFlow Serving como base, pero también puedes utilizar las que hayas construido desde el código fuente.

Nota: TensorFlow Serving se beneficia de optimizaciones de hardware que aprovechan conjuntos de instrucciones como AVX512. Estos conjuntos de instrucciones pueden acelerar la inferencia de modelos de aprendizaje profundo. Por lo tanto, si conoces el hardware en el que se implementará el modelo, a menudo es beneficioso obtener una versión optimizada de la imagen de TensorFlow Serving y utilizarla en todo momento.

Ahora que el contenedor en ejecución tiene todos los archivos necesarios en la estructura de directorios adecuada, debemos crear una nueva imagen de Docker que incluya estos cambios. Esto se puede hacer con el siguiente comando docker commit, y tendrás una nueva imagen de Docker llamada $NEW_IMAGE. Una cosa importante a tener en cuenta es que debes establecer la variable de entorno MODEL_NAME en el nombre del modelo, que en este caso es hf-vit. Esto le indica a TensorFlow Serving qué modelo implementar.

$ NEW_IMAGE=tfserving:$MODEL_NAME

$ docker commit \ 
    --change "ENV MODEL_NAME $MODEL_NAME" \ 
    serving_base $NEW_IMAGE

Ejecutando la imagen de Docker localmente

Por último, puedes ejecutar la imagen de Docker recién creada localmente para comprobar si funciona correctamente. A continuación, se muestra la salida del comando docker run. Dado que la salida es detallada, la hemos reducido para centrarnos en lo importante. Además, cabe señalar que abre los puertos 8500 y 8501 para los puntos finales gRPC y HTTP/REST, respectivamente.

$ docker run -p 8500:8500 -p 8501:8501 -t $NEW_IMAGE &


---------SALIDA---------
(Re-)adding model: hf-vit
Successfully reserved resources to load servable {name: hf-vit version: 1}
Approving load for servable version {name: hf-vit version: 1}
Loading servable version {name: hf-vit version: 1}
Reading SavedModel from: /models/hf-vit/1
Reading SavedModel debug info (if present) from: /models/hf-vit/1
Successfully loaded servable version {name: hf-vit version: 1}
Running gRPC ModelServer at 0.0.0.0:8500 ...
Exporting HTTP/REST API at:localhost:8501 ...

Subiendo la imagen de Docker

El último paso aquí es subir la imagen de Docker a un repositorio de imágenes. Utilizarás Google Container Registry (GCR) para este propósito. Las siguientes líneas de código pueden hacer esto por ti:

$ GCP_PROJECT_ID=<GCP_PROJECT_ID>
$ GCP_IMAGE=gcr.io/$GCP_PROJECT_ID/$NEW_IMAGE

$ gcloud auth configure-docker
$ docker tag $NEW_IMAGE $GCP_IMAGE
$ docker push $GCP_IMAGE

Dado que estamos utilizando GCR, debes agregar el prefijo de la etiqueta de la imagen de Docker (nota también otros formatos) con gcr.io/<GCP_PROJECT_ID>. Con la imagen de Docker preparada y subida a GCR, ahora puedes proceder a implementarla en un clúster de Kubernetes.

La implementación en un clúster de Kubernetes requiere lo siguiente:

• Provisionar un clúster de Kubernetes, que se realiza con Google Kubernetes Engine (GKE) en esta publicación. Sin embargo, también puedes utilizar otras plataformas y herramientas como EKS o Minikube.

• Conectarse al clúster de Kubernetes para realizar una implementación.

• Escribir manifiestos YAML.

• Realizar la implementación con los manifiestos utilizando la herramienta kubectl.

Repasemos cada uno de estos pasos.

Provisión de un clúster de Kubernetes en GKE

Puedes usar un script de shell como este para esto (disponible aquí):

$ GKE_CLUSTER_NAME=tfs-cluster
$ GKE_CLUSTER_ZONE=us-central1-a
$ NUM_NODES=2
$ MACHINE_TYPE=n1-standard-8

$ gcloud container clusters create $GKE_CLUSTER_NAME \
    --zone=$GKE_CLUSTER_ZONE \
    --machine-type=$MACHINE_TYPE \
    --num-nodes=$NUM_NODES

GCP ofrece una variedad de tipos de máquinas para configurar la implementación de la manera que desees. Te recomendamos que consultes la documentación para obtener más información al respecto.

Una vez que se provisiona el clúster, debes conectarte a él para realizar la implementación. Dado que se utiliza GKE aquí, también necesitas autenticarte. Puedes usar un script de shell como este para hacer ambas cosas:

$ GCP_PROJECT_ID=<GCP_PROJECT_ID>

$ export USE_GKE_GCLOUD_AUTH_PLUGIN=True

$ gcloud container clusters get-credentials $GKE_CLUSTER_NAME \
    --zone $GKE_CLUSTER_ZONE \
    --project $GCP_PROJECT_ID

El comando gcloud container clusters get-credentials se encarga tanto de conectarse al clúster como de la autenticación. Una vez hecho esto, estás listo para escribir los manifiestos.

Escribiendo manifiestos de Kubernetes

Los manifiestos de Kubernetes se escriben en archivos YAML. Si bien es posible utilizar un solo archivo de manifiesto para realizar la implementación, crear archivos de manifiesto separados a menudo es beneficioso para delegar la separación de responsabilidades. Es común utilizar tres archivos de manifiesto para lograr esto:

• deployment.yaml define el estado deseado de la implementación proporcionando el nombre de la imagen Docker, argumentos adicionales al ejecutar la imagen Docker, los puertos para abrir para accesos externos y los límites de recursos.

• service.yaml define las conexiones entre clientes externos y los Pods dentro del clúster de Kubernetes.

• hpa.yaml define reglas para escalar hacia arriba y hacia abajo el número de Pods que forman parte de la implementación, como el porcentaje de utilización de la CPU.

Puedes encontrar los manifiestos relevantes para esta publicación aquí. A continuación, presentamos una visión general pictórica de cómo se consumen estos manifiestos.

A continuación, revisamos las partes importantes de cada uno de estos manifiestos.

deployment.yaml:

apiVersion: apps/v1
kind: Deployment
metadata:
  labels:
    app: tfs-server
  name: tfs-server
...
spec:
  containers:
  - image: gcr.io/$GCP_PROJECT_ID/tfserving-hf-vit:latest
    name: tfs-k8s
    imagePullPolicy: Always
    args: ["--tensorflow_inter_op_parallelism=2", 
           "--tensorflow_intra_op_parallelism=8"] 
    ports:
    - containerPort: 8500
      name: grpc
    - containerPort: 8501
      name: restapi
    resources:
      limits:
        cpu: 800m
      requests:
        cpu: 800m
...

Puedes configurar los nombres como tfs-server, tfs-k8s de la forma que desees. Bajo containers, especificas la URI de la imagen Docker que utilizará la implementación. La utilización actual de recursos se monitorea estableciendo los límites permitidos de los resources para el contenedor. Esto le permite al escalador de pod horizontal (discutido más adelante) decidir escalar hacia arriba o hacia abajo el número de contenedores. requests.cpu es la cantidad mínima de recursos de CPU necesarios para que el contenedor funcione correctamente, establecido por los operadores. Aquí, 800m significa el 80% de todos los recursos de la CPU. Por lo tanto, el escalador de pod horizontal monitorea la utilización promedio de la CPU a partir de la suma de requests.cpu en todos los Pods para tomar decisiones de escalado.

Además de la configuración específica de Kubernetes, puedes especificar opciones específicas de TensorFlow Serving en args. En este caso, tienes dos opciones:

• tensorflow_inter_op_parallelism, que establece el número de hilos para ejecutar operaciones independientes en paralelo. El valor recomendado para esto es 2.

• tensorflow_intra_op_parallelism, que establece el número de hilos para ejecutar operaciones individuales en paralelo. El valor recomendado es el número de núcleos físicos de la CPU de implementación.

Puedes obtener más información sobre estas opciones (y otras) y consejos sobre cómo ajustarlas para su implementación desde aquí y aquí.

service.yaml :

apiVersion: v1
kind: Service
metadata:
  labels:
    app: tfs-server
  name: tfs-server
spec:
  ports:
  - port: 8500
    protocol: TCP
    targetPort: 8500
    name: tf-serving-grpc
  - port: 8501
    protocol: TCP
    targetPort: 8501
    name: tf-serving-restapi
  selector:
    app: tfs-server
  type: LoadBalancer

Hemos configurado el tipo de servicio como ‘LoadBalancer’ para que los puntos de conexión sean accesibles externamente desde el clúster de Kubernetes. Se selecciona el despliegue ‘tfs-server’ para establecer conexiones con clientes externos a través de los puertos especificados. Abrimos dos puertos, ‘8500’ y ‘8501’, para conexiones gRPC y HTTP/REST respectivamente.

hpa.yaml :

apiVersion: autoscaling/v1
kind: HorizontalPodAutoscaler
metadata:
 name: tfs-server

spec:
 scaleTargetRef:
   apiVersion: apps/v1
   kind: Deployment
   name: tfs-server
 minReplicas: 1
 maxReplicas: 3
 targetCPUUtilizationPercentage: 80

HPA significa H orizontal P od A utoscaler. Establece los criterios para decidir cuándo escalar el número de Pods en el despliegue objetivo. Puedes obtener más información sobre el algoritmo de escalado automático utilizado internamente por Kubernetes aquí.

Aquí se especifica cómo Kubernetes debe manejar el escalado automático. En particular, se define el límite de réplicas dentro del cual debe realizar el escalado automático, minReplicas y maxReplicas, y la utilización objetivo de la CPU. targetCPUUtilizationPercentage es una métrica importante para el escalado automático. El siguiente enlace resume adecuadamente su significado (tomado de aquí):

La utilización de la CPU es el uso promedio de la CPU de todos los Pods en un despliegue en el último minuto, dividido por la CPU solicitada de este despliegue. Si la media de la utilización de la CPU de los Pods es mayor que el objetivo que has definido, se ajustarán tus réplicas.

Recuerda especificar resources en el manifiesto de despliegue. Al especificar los resources, el plano de control de Kubernetes comienza a supervisar las métricas, por lo que el targetCPUUtilization funciona. De lo contrario, HPA no conoce el estado actual del despliegue.

Puedes experimentar y establecer estos números según tus necesidades. Sin embargo, el escalado automático dependerá de la cuota disponible en GCP, ya que GKE utiliza internamente Google Compute Engine para gestionar estos recursos.

Realizando el despliegue

Una vez que los manifiestos estén listos, puedes aplicarlos al clúster de Kubernetes actualmente conectado con el comando kubectl apply.

$ kubectl apply -f deployment.yaml
$ kubectl apply -f service.yaml
$ kubectl apply -f hpa.yaml

Aunque usar kubectl está bien para aplicar cada uno de los manifiestos y realizar el despliegue, puede volverse complicado rápidamente si tienes muchos manifiestos diferentes. Aquí es donde puede ser útil una utilidad como Kustomize. Simplemente defines otra especificación llamada kustomization.yaml de la siguiente manera:

commonLabels:
  app: tfs-server
resources:
- deployment.yaml
- hpa.yaml
- service.yaml
apiVersion: kustomize.config.k8s.io/v1beta1
kind: Kustomization

Entonces, solo necesitas una línea para realizar el despliegue real:

$ kustomize build . | kubectl apply -f -

Las instrucciones completas están disponibles aquí. Una vez que se haya realizado el despliegue, puedes obtener la dirección IP del punto de conexión de la siguiente manera:

$ kubectl rollout status deployment/tfs-server
$ kubectl get svc tfs-server --watch

---------OUTPUT---------
NAME        TYPE          CLUSTER-IP   EXTERNAL-IP  PORT(S)                        AGE
tfs-server  LoadBalancer  xxxxxxxxxx   xxxxxxxxxx   8500:30869/TCP,8501:31469/TCP  xxx

Toma nota de la IP externa cuando esté disponible.

¡Y eso resume todos los pasos que necesitas para implementar tu modelo en Kubernetes! Kubernetes proporciona de manera elegante abstracciones para aspectos complejos como el escalado automático y la gestión de clústeres, mientras te permite centrarte en los aspectos cruciales en los que debes enfocarte al implementar un modelo. Estos incluyen la utilización de recursos, la seguridad (no la cubrimos aquí), los principales indicadores de rendimiento como la latencia, etc.

Dado que obtuviste una IP externa para el endpoint, puedes utilizar la siguiente lista para probarlo:

import tensorflow as tf 
import json
import base64

image_path = tf.keras.utils.get_file(
    "image.jpg", "http://images.cocodataset.org/val2017/000000039769.jpg"
)
bytes_inputs = tf.io.read_file(image_path)
b64str = base64.urlsafe_b64encode(bytes_inputs.numpy()).decode("utf-8")
data = json.dumps(
    {"signature_name": "serving_default", "instances": [b64str]}
)

json_response = requests.post(
    "http://<ENDPOINT-IP>:8501/v1/models/hf-vit:predict", 
    headers={"content-type": "application/json"}, 
    data=data
)
print(json.loads(json_response.text))

---------SALIDA---------
{'predictions': [{'label': 'Gato egipcio', 'confidence': 0.896659195}]}

Si estás interesado en saber cómo se desempeñaría esta implementación si recibe más tráfico, te recomendamos que consultes este artículo . Consulta el repositorio correspondiente para obtener más información sobre cómo ejecutar pruebas de carga con Locust y visualizar los resultados.

TensorFlow Serving ofrece varias opciones para adaptar la implementación según el caso de uso de tu aplicación. A continuación, discutimos brevemente algunas de ellas.

enable_batching habilita la capacidad de inferencia por lotes que recopila las solicitudes entrantes durante un cierto intervalo de tiempo, las agrupa en un lote, realiza una inferencia por lotes y devuelve los resultados de cada solicitud a los clientes correspondientes. TensorFlow Serving proporciona un conjunto completo de opciones configurables (como max_batch_size, num_batch_threads) para adaptar tus necesidades de implementación. Puedes obtener más información sobre ellas aquí . El procesamiento por lotes es especialmente beneficioso para aplicaciones en las que no necesitas predicciones de un modelo de forma instantánea. En esos casos, normalmente recopilarías varias muestras para la predicción en lotes y luego enviarías esos lotes para la predicción. Afortunadamente, TensorFlow Serving puede configurar todo esto automáticamente cuando se habilitan sus capacidades de procesamiento por lotes.

enable_model_warmup realiza un precalentamiento de algunos de los componentes de TensorFlow que se instancian de forma perezosa con datos de entrada ficticios. De esta manera, puedes asegurarte de que todo esté cargado correctamente y de que no haya retrasos durante el tiempo real del servicio.

En esta publicación y el repositorio asociado, aprendiste cómo implementar el modelo Vision Transformer de 🤗 Transformers en un clúster de Kubernetes. Si estás haciendo esto por primera vez, los pasos pueden parecer un poco desafiantes, pero una vez que los comprendas, pronto se convertirán en un componente esencial de tu conjunto de herramientas. Si ya estabas familiarizado con este flujo de trabajo, esperamos que esta publicación aún te haya sido útil.

Aplicamos el mismo flujo de trabajo de implementación para una versión optimizada de ONNX del mismo modelo Vision Transformer. Para obtener más detalles, consulta este enlace . Los modelos optimizados de ONNX son especialmente beneficiosos si estás utilizando CPUs x86 para la implementación.

En la próxima publicación, te mostraremos cómo realizar estas implementaciones con significativamente menos código con Vertex AI, ¡más como model.deploy(autoscaling_config=...) y listo! Esperamos que estés tan emocionado como nosotros.

Gracias al equipo del Programa de Relaciones con Desarrolladores de ML en Google, que nos proporcionó créditos de GCP para realizar los experimentos.