Ejecuta una canalización con las GPU

En esta página, se explica cómo ejecutar una canalización de Apache Beam en Dataflow con GPU. Los trabajos que usan GPU generan cargos, como se especifica en la página de precios de Dataflow.

Para obtener más información sobre el uso de GPU con Dataflow, consulta Compatibilidad de Dataflow con las GPU. Si deseas obtener más información sobre el flujo de trabajo de los desarrolladores para compilar canalizaciones con GPU, consulta Acerca de las GPU con Dataflow.

Usa notebooks de Apache Beam

Si ya tienes una canalización que deseas ejecutar con GPU en Dataflow, puedes omitir esta sección.

Los notebooks de Apache Beam ofrecen una forma conveniente de prototipar y desarrollar de forma iterativa tu canalización con GPU sin configurar un entorno de desarrollo. Para comenzar, lee la guía Desarrolla con notebooks de Apache Beam, inicia una instancia de notebooks de Apache Beam y sigue el notebook de ejemplo Usa GPU con Apache Beam.

Aprovisiona la cuota de GPU

Los dispositivos de GPU están sujetos a la disponibilidad de cuota de tu proyecto de Google Cloud Platform. Solicita la cuota de GPU en la región que elijas.

Instalar los controladores de GPU.

Para instalar los controladores NVIDIA en los trabajadores de Dataflow, agrega install-nvidia-driver a la opción de servicio worker_accelerator.

Cuando especificas la opción install-nvidia-driver, Dataflow instala controladores NVIDIA en los trabajadores de Dataflow con la utilidad cos-extensions proporcionada por Container-Optimized OS. Cuando se especifica install-nvidia-driver, aceptas el contrato de licencia de NVIDIA.

Los objetos binarios y las bibliotecas que proporciona el instalador del controlador de NVIDIA se activan en el código de usuario de ejecución del contenedor que ejecuta el código de usuario de la canalización en /usr/local/nvidia/.

La versión del controlador de GPU depende de la versión de Container-Optimized OS que usa Dataflow. Para encontrar la versión del controlador de GPU de un trabajo de Dataflow determinado, en los registros de paso de Dataflow de tu trabajo, busca GPU driver.

Compila una imagen de contenedor personalizada

Para interactuar con las GPU, es posible que necesites software adicional de NVIDIA, como bibliotecas aceleradas por GPU y el kit de herramientas de CUDA. Proporciona estas bibliotecas en el contenedor de Docker que ejecuta el código de usuario.

Para personalizar la imagen del contenedor, proporciona una imagen que cumpla con el contrato de imagen de contenedor del SDK de Apache Beam y que tenga las bibliotecas de GPU necesarias.

Para proporcionar una imagen de contenedor personalizada, usa Dataflow Runner v2 y proporciona la imagen de contenedor con la opción de canalización sdk_container_image. Si usas la versión 2.29.0 de Apache Beam o una anterior, usa la opción de canalización worker_harness_container_image. Para obtener más información, consulta Usa contenedores personalizados.

Para compilar una imagen de contenedor personalizada, usa uno de los siguientes dos enfoques:

Usa una imagen existente configurada para el uso de GPU

Puedes compilar una imagen de Docker que cumpla con el contrato de contenedor del SDK de Apache Beam a partir de una imagen base existente que esté preconfigurada para el uso de GPU. Por ejemplo, las imágenes de Docker de TensorFlow y las imágenes de contenedor de NVIDIA están preconfiguradas para el uso de GPU.

Un Dockerfile de muestra que se basa en la imagen de Docker de TensorFlow con Python 3.6 se ve como el siguiente ejemplo:

ARG BASE=tensorflow/tensorflow:2.5.0-gpu
FROM $BASE

# Check that the chosen base image provides the expected version of Python interpreter.
ARG PY_VERSION=3.6
RUN [[ $PY_VERSION == `python -c 'import sys; print("%s.%s" % sys.version_info[0:2])'` ]] \
   || { echo "Could not find Python interpreter or Python version is different from ${PY_VERSION}"; exit 1; }

RUN pip install --upgrade pip \
    && pip install --no-cache-dir apache-beam[gcp]==2.29.0 \
    # Verify that there are no conflicting dependencies.
    && pip check

# Copy the Apache Beam worker dependencies from the Beam Python 3.6 SDK image.
COPY --from=apache/beam_python3.6_sdk:2.29.0 /opt/apache/beam /opt/apache/beam

# Apache Beam worker expects pip at /usr/local/bin/pip by default.
# Some images have pip in a different location. If necessary, make a symlink.
# This line can be omitted in Beam 2.30.0 and later versions.
RUN [[ `which pip` == "/usr/local/bin/pip" ]] || ln -s `which pip` /usr/local/bin/pip

# Set the entrypoint to Apache Beam SDK worker launcher.
ENTRYPOINT [ "/opt/apache/beam/boot" ]

Cuando uses las imágenes de Docker de TensorFlow, usa TensorFlow 2.5.0 o una versión posterior. Las imágenes de Docker de TensorFlow anteriores instalan el paquete tensorflow-gpu en lugar del paquete tensorflow. La distinción no es importante después de la versión 2.1.0 de TensorFlow, pero varios paquetes descendentes, como tfx, requieren el paquete tensorflow.

Los tamaños grandes de contenedores ralentizan el tiempo de inicio del trabajador. Este cambio en el rendimiento puede ocurrir cuando usas contenedores como contenedores de aprendizaje profundo.

Instala una versión específica de Python

Si tienes requisitos estrictos para la versión de Python, puedes compilar una imagen a partir de una imagen base de NVIDIA que tenga bibliotecas de GPU necesarias y, luego, instalar el intérprete de Python.

En el siguiente ejemplo, se muestra cómo seleccionar una imagen de NVIDIA que no incluye el intérprete de Python del catálogo de imágenes de contenedor de CUDA. Ajusta el ejemplo para instalar la versión necesaria de Python 3 y pip. En el ejemplo, se usa TensorFlow. Por lo tanto, cuando elijas una imagen, las versiones CUDA y cuDNN en la imagen base satisfacen los requisitos de la versión de TensorFlow.

Un Dockerfile de muestra se ve de la siguiente manera:

# Select an NVIDIA base image with needed GPU stack from https://ngc.nvidia.com/catalog/containers/nvidia:cuda

FROM nvidia/cuda:11.0.3-cudnn8-runtime-ubuntu20.04

RUN \
    # Add Deadsnakes repository that has a variety of Python packages for Ubuntu.
    # See: https://launchpad.net/~deadsnakes/+archive/ubuntu/ppa
    apt-key adv --keyserver keyserver.ubuntu.com --recv-keys F23C5A6CF475977595C89F51BA6932366A755776 \
    && echo "deb http://ppa.launchpad.net/deadsnakes/ppa/ubuntu focal main" >> /etc/apt/sources.list.d/custom.list \
    && echo "deb-src http://ppa.launchpad.net/deadsnakes/ppa/ubuntu focal main" >> /etc/apt/sources.list.d/custom.list \
    && apt-get update \
    && apt-get install -y curl \
        python3.8 \
        # With python3.8 package, distutils need to be installed separately.
        python3-distutils \
    && rm -rf /var/lib/apt/lists/* \
    && update-alternatives --install /usr/bin/python python /usr/bin/python3.8 10 \
    && curl https://bootstrap.pypa.io/get-pip.py | python \
    && pip install --upgrade pip \
    # Install Apache Beam and Python packages that will interact with GPUs.
    && pip install --no-cache-dir apache-beam[gcp]==2.29.0 tensorflow==2.4.0 \
    # Verify that there are no conflicting dependencies.
    && pip check

# Copy the Apache Beam worker dependencies from the Beam Python 3.8 SDK image.
COPY --from=apache/beam_python3.8_sdk:2.29.0 /opt/apache/beam /opt/apache/beam

# Set the entrypoint to Apache Beam SDK worker launcher.
ENTRYPOINT [ "/opt/apache/beam/boot" ]

En algunas distribuciones de SO, puede ser difícil instalar versiones específicas de Python con el administrador de paquetes de SO. En este caso, instala el intérprete de Python con herramientas como Miniconda o pyenv.

Un Dockerfile de muestra se ve de la siguiente manera:

FROM nvidia/cuda:11.0.3-cudnn8-runtime-ubuntu20.04

# The Python version of the Dockerfile must match the Python version you use
# to launch the Dataflow job.

ARG PYTHON_VERSION=3.8

# Update PATH so we find our new Conda and Python installations.
ENV PATH=/opt/python/bin:/opt/conda/bin:$PATH

RUN apt-get update \
    && apt-get install -y wget \
    && rm -rf /var/lib/apt/lists/* \
    # The NVIDIA image doesn't come with Python pre-installed.
    # We use Miniconda to install the Python version of our choice.
    && wget -q https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh \
    && bash Miniconda3-latest-Linux-x86_64.sh -b -p /opt/conda \
    && rm Miniconda3-latest-Linux-x86_64.sh \
    # Create a new Python environment with needed version, and install pip.
    && conda create -y -p /opt/python python=$PYTHON_VERSION pip \
    # Remove unused Conda packages, install necessary Python packages via pip
    # to avoid mixing packages from pip and Conda.
    && conda clean -y --all --force-pkgs-dirs \
    && pip install --upgrade pip \
    # Install Apache Beam and Python packages that will interact with GPUs.
    && pip install --no-cache-dir apache-beam[gcp]==2.29.0 tensorflow==2.4.0 \
    # Verify that there are no conflicting dependencies.
    && pip check \
    # Apache Beam worker expects pip at /usr/local/bin/pip by default.
    # You can omit this line when using Beam 2.30.0 and later versions.
    && ln -s $(which pip) /usr/local/bin/pip

# Copy the Apache Beam worker dependencies from the Apache Beam SDK for Python 3.8 image.
COPY --from=apache/beam_python3.8_sdk:2.29.0 /opt/apache/beam /opt/apache/beam

# Set the entrypoint to Apache Beam SDK worker launcher.
ENTRYPOINT [ "/opt/apache/beam/boot" ]

Usa una imagen de contenedor de Apache Beam

Puedes configurar una imagen de contenedor para el uso de GPU sin usar imágenes preconfiguradas. Este enfoque solo se recomienda cuando las imágenes preconfiguradas no funcionan en tu caso. Para configurar tu propia imagen de contenedor, debes seleccionar bibliotecas compatibles y configurar su entorno de ejecución.

Un Dockerfile de muestra se ve de la siguiente manera:

FROM apache/beam_python3.7_sdk:2.24.0
ENV INSTALLER_DIR="/tmp/installer_dir"

# The base image has TensorFlow 2.2.0, which requires CUDA 10.1 and cuDNN 7.6.
# You can download cuDNN from NVIDIA website
# https://developer.nvidia.com/cudnn
COPY cudnn-10.1-linux-x64-v7.6.0.64.tgz $INSTALLER_DIR/cudnn.tgz
RUN \
    # Download CUDA toolkit.
    wget -q -O $INSTALLER_DIR/cuda.run https://developer.download.nvidia.com/compute/cuda/10.1/Prod/local_installers/cuda_10.1.243_418.87.00_linux.run && \

    # Install CUDA toolkit. Print logs upon failure.
    sh $INSTALLER_DIR/cuda.run --toolkit --silent || (egrep '^\[ERROR\]' /var/log/cuda-installer.log && exit 1) && \
    # Install cuDNN.
    mkdir $INSTALLER_DIR/cudnn && \
    tar xvfz $INSTALLER_DIR/cudnn.tgz -C $INSTALLER_DIR/cudnn && \

    cp $INSTALLER_DIR/cudnn/cuda/include/cudnn*.h /usr/local/cuda/include && \
    cp $INSTALLER_DIR/cudnn/cuda/lib64/libcudnn* /usr/local/cuda/lib64 && \
    chmod a+r /usr/local/cuda/include/cudnn*.h /usr/local/cuda/lib64/libcudnn* && \
    rm -rf $INSTALLER_DIR

# A volume with GPU drivers will be mounted at runtime at /usr/local/nvidia.
ENV LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/nvidia/lib64:/usr/local/cuda/lib64

Las bibliotecas de controlador en /usr/local/nvidia/lib64 deben poder detectarse en el contenedor como bibliotecas compartidas. Para que las bibliotecas de controlador sean detectables, configura la variable de entorno LD_LIBRARY_PATH.

Si usas TensorFlow, debes elegir una combinación compatible de las herramientas de CUDA y las versiones de cuDNN. Para obtener más información, consulta Requisitos de software y Configuraciones de compilación probadas.

Selecciona el tipo y cantidad de GPU para los trabajadores de Dataflow

A fin de configurar el tipo y la cantidad de GPU para adjuntar a los trabajadores de Dataflow, usa la opción de servicio worker_accelerator. Selecciona el tipo y la cantidad de GPU en función de tu caso de uso y cómo planeas usar las GPU en tu canalización.

Para obtener una lista de los tipos de GPU compatibles con Dataflow, consulta Compatibilidad de Dataflow para GPU.

Ejecuta tu trabajo con GPUs

Las consideraciones para ejecutar un trabajo de Dataflow con las GPU incluyen las siguientes:

  • Debido a que los contenedores de GPU suelen ser grandes, para evitar quedarte sin espacio en el disco, haz lo siguiente:

  • Considera cuántos procesos usan simultáneamente la misma GPU en una VM de trabajador. Luego, decide si deseas limitar la GPU a un solo proceso o permitir que varios procesos usen la GPU.

    • Si un proceso del SDK de Apache Beam puede usar la mayor parte de la memoria de GPU disponible, por ejemplo, a través de la carga de un modelo grande en una GPU, es posible que desees configurar los trabajadores para que usen un solo proceso con la configuración de la opción de canalización --experiments=no_use_multiple_sdk_containers. Como alternativa, puedes usar trabajadores con una CPU virtual mediante un tipo personalizado de máquina, como n1-custom-1-NUMBER_OF_MB o n1-custom-1-NUMBER_OF_MB-ext, para la memoria extendida. Para obtener más información, consulta Usa un tipo de máquina con más memoria por CPU virtual.
    • Si varios procesos comparten la GPU, habilita el procesamiento simultáneo en una GPU compartida con el servicio de procesamiento múltiple de NVIDIA (MPS).

    Para obtener información general, consulta Paralelismo de trabajadores y GPU.

Para ejecutar un trabajo de Dataflow con GPU, usa el siguiente comando. Para usar el ajuste adecuado, en lugar de usar la opción de servicio worker_accelerator, usa la sugerencia de recursos accelerator..

Python

python PIPELINE \
  --runner "DataflowRunner" \
  --project "PROJECT" \
  --temp_location "gs://BUCKET/tmp" \
  --region "REGION" \
  --worker_harness_container_image "IMAGE" \
  --disk_size_gb "DISK_SIZE_GB" \
  --dataflow_service_options "worker_accelerator=type:GPU_TYPE;count:GPU_COUNT;install-nvidia-driver" \
  --experiments "use_runner_v2"

Reemplaza lo siguiente:

  • PIPELINE: tu archivo de código fuente de canalización
  • PROJECT: El nombre del proyecto de Google Cloud
  • BUCKET: el bucket de Cloud Storage
  • REGION: una región de Dataflow, por ejemplo, us-central1. Selecciona una "REGIÓN" que tenga zonas que admiten el GPU_TYPE. Dataflow asigna de forma automática a los trabajadores a una zona con GPU en esta región.
  • IMAGE: la ruta de acceso de Artifact Registry para tu imagen de Docker
  • DISK_SIZE_GB: el tamaño del disco de arranque para cada VM de trabajador, por ejemplo, 50
  • GPU_TYPE: un tipo de GPU disponible, por ejemplo, nvidia-tesla-t4
  • GPU_COUNT: cantidad de GPU que se conectará a cada VM de trabajador, por ejemplo, 1

Verifica tu trabajo de Dataflow

Para confirmar que el trabajo use VMs de trabajador con GPUs, sigue estos pasos:

  1. Verifica que los trabajadores de Dataflow para el trabajo se hayan iniciado.
  2. Mientras se ejecuta un trabajo, busca una VM de trabajador asociada con el trabajo.
    1. En el mensaje Buscar productos y recursos, pega el ID de trabajo.
    2. Selecciona la instancia de VM de Compute Engine asociada con el trabajo.

También puedes encontrar una lista de todas las instancias en ejecución en la consola de Compute Engine.

  1. En la consola de Google Cloud , ve a la página Instancias de VM.

    Ir a Instancias de VM

  2. Haz clic en Detalles de la instancia de VM.

  3. Verifica que la página de detalles tenga una sección GPU y que tus GPU estén conectadas.

Si tu trabajo no se inició con GPU, verifica que la opción de servicio worker_accelerator esté configurada correctamente y se vea en la interfaz de supervisión de Dataflow en dataflow_service_options. El orden de los tokens en los metadatos del acelerador es importante.

Por ejemplo, una opción de canalización dataflow_service_options en la interfaz de supervisión de Dataflow podría tener el siguiente aspecto:

['worker_accelerator=type:nvidia-tesla-t4;count:1;install-nvidia-driver', ...]

Visualiza el uso de GPU

Para ver el uso de GPU en las VMs de trabajador, sigue estos pasos:

  1. En la Google Cloud consola, ve a Monitoring o usa el siguiente botón:

    Ir a Monitoring

  2. En el panel de navegación de Monitoring, haz clic en Explorador de métricas.

  3. En Tipo de recurso, especifica Dataflow Job. Para la métrica, especifica GPU utilization o GPU memory utilization, según la métrica que desees supervisar.

Para obtener más información, consulta el Explorador de métricas.

Habilitar el servicio de procesamiento múltiple de NVIDIA

En las canalizaciones de Python que se ejecutan en trabajadores con más de una CPU virtual, puedes mejorar la simultaneidad para las operaciones de GPU mediante la habilitación del servicio de procesos múltiples (MPS) de NVIDIA. Si deseas obtener más información y los pasos para usar MPS, consulta Mejora el rendimiento en una GPU compartida mediante NVIDIA MPS.

Opcional: Configura un modelo de aprovisionamiento

Puedes mejorar la capacidad de obtener acceso a los recursos de GPU configurando un modelo de aprovisionamiento para tu canalización.

Dataflow admite los siguientes modelos de aprovisionamiento: estándar y de inicio flexible.

Aprovisionamiento estándar

El aprovisionamiento estándar es el modelo de aprovisionamiento predeterminado para todos los trabajos de Dataflow con GPUs. Las instancias que usan recursos de acelerador se crean de inmediato según la disponibilidad de recursos.

No tienes que configurar nada para usar el modelo de aprovisionamiento estándar.

Si las GPUs no están disponibles de inmediato en la zona o región en la que ejecutas tu trabajo de Dataflow, es posible que el trabajo no se inicie. Para obtener más información, consulta El trabajo falla inmediatamente en el inicio.

Aprovisionamiento de inicio flexible

Con el modelo de aprovisionamiento de inicio flexible, las instancias y los recursos del acelerador se programan para el aprovisionamiento y se completan según la disponibilidad de los recursos. Puedes usar el modelo de aprovisionamiento de inicio flexible para aumentar las probabilidades de obtener GPUs.

Para usar el modelo de aprovisionamiento de inicio flexible, agrega provisioning_model:FLEX_START a la opción de servicio worker_accelerator. Por ejemplo:

worker_accelerator=type:nvidia-tesla-t4;count:1;install-nvidia-driver:5xx;provisioning_model:FLEX_START

Los trabajos con inicio flexible habilitado se envían para su ejecución, pero solo se ejecutan cuando los recursos necesarios están disponibles. Para verificar que se haya habilitado el aprovisionamiento de inicio flexible, busca la siguiente entrada de registro en el registro job-message:

Se habilitó FLEX_START para el trabajo JOB_ID

Para determinar si el trabajo comenzó a ejecutarse, consulta el gráfico de ajuste de escala automático en la página de métricas del trabajo:

Gráfico de ajuste de escala automático que muestra el aprovisionamiento del trabajador después del inicio del trabajo con Inicio flexible habilitado

Los trabajos que comenzaron a ejecutarse mostrarán un recuento de trabajadores distinto de cero, mientras que los trabajos que están esperando recursos tendrán un recuento de trabajadores igual a cero.

Asistencia y limitaciones

  • El inicio flexible solo se admite en las canalizaciones por lotes. No se admiten las canalizaciones de transmisión.
  • Las VMs de trabajador aprovisionadas con el modelo de aprovisionamiento de inicio flexible tienen un tiempo de ejecución máximo de siete días. Después de este período, se interrumpen las VMs de trabajador con aceleradores. Dataflow intentará volver a aprovisionar los recursos. Si no se pueden volver a aprovisionar los recursos, la canalización fallará.
  • El inicio flexible intentará aprovisionar recursos hasta 1 hora después del envío del trabajo. Si no se pueden aprovisionar recursos después de 1 hora, el trabajo fallará.
  • El inicio flexible consume la cuota interrumpible. Si tu proyecto no tiene una cuota de recursos interrumpibles, se consumirá la cuota estándar. Para obtener más información, consulta Cuotas interrumpibles.
  • Si no se proporciona una configuración de zona de trabajador, Dataflow elegirá una sola zona para crear todos los recursos en función de la compatibilidad de hardware, la disponibilidad actual de recursos y cuotas, y las reservas coincidentes. Esta zona podría ser diferente de la zona en la que se encuentran los recursos de servicio para el trabajo.
  • No se admite el ajuste de escala automático horizontal. Para usar más de un trabajador, establece la opción de canalización --num_workers.
  • No se admiten las TPUs.
  • No se admite el ajuste adecuado.

Soluciona problemas de inicio flexible

Si tu trabajo falla 1 hora después del envío con el siguiente error:

The Dataflow job appears to be stuck because no worker activity has been seen in the last 1h.

Ve a la página Cuotas en la consola de Google Cloud para asegurarte de que tu proyecto tenga suficiente cuota de PREEMPTIBLE_GPU_TYPE_GPUS en la región configurada de tu trabajo.

Si había suficiente cuota en tu proyecto, flex-start no pudo aprovisionar recursos en 1 hora. Considera iniciar la canalización en una zona diferente o con un tipo de acelerador diferente.

Usa GPU con Dataflow Prime

Dataflow Prime te permite solicitar aceleradores para un paso específico de tu canalización. Para usar GPU con Dataflow Prime, no uses la opción de canalización --dataflow-service_options=worker_accelerator. En su lugar, solicita las GPU con la sugerencia de recursos de accelerator. Para obtener más información, consulta Usa sugerencias de recursos.

Soluciona problemas de tu trabajo de Dataflow

Si tienes problemas para ejecutar tu trabajo de Dataflow con GPU, consulta Soluciona problemas de tu trabajo de GPU de Dataflow.

¿Qué sigue?