Ejecuta una canalización con 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 del proyecto de Google Cloud. 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 de 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 GPU

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 los recursos extendidos memoria. 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 tarea.
    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 consola de Google Cloud, 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.

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?