Usa 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 la compatibilidad de Dataflow con las GPU. Si deseas obtener más información sobre el flujo de trabajo de los desarrolladores para compilar canalizaciones mediante GPU, consulta Desarrolla con GPU.

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 notebook 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.

Instala controladores de GPU

Indica a Dataflow que instale los controladores NVIDIA en los trabajadores. Para ello, agrega install-nvidia-driver a la opción worker_accelerator. Cuando se especifica 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 se ejecuta en /usr/local/nvidia/.

La versión del controlador de GPU depende de la versión de Container-Optimized OS que usa Dataflow en la actualidad. 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.

Puedes personalizar la imagen del contenedor si proporcionas 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 worker_harness_container_image. Si usas Apache Beam 2.30.0 o una versión posterior, puedes usar un nombre de opción más corto sdk_container_image para simplificar. Para obtener más información, consulta Usa contenedores personalizados.

Enfoque 1. 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 de la siguiente manera:

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 se usan 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 del catálogo de imágenes de contenedor de CUDA que no incluye el intérprete de Python. Puedes ajustar el ejemplo para instalar la versión deseada 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 desired 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 mediante el administrador de paquetes de SO. En este caso, puedes instalar un 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 desired 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" ]

Enfoque 2. Usa imágenes de contenedor de Apache Beam

Puedes configurar una imagen de contenedor para el uso de GPU sin usar imágenes preconfiguradas. No se recomienda este enfoque, a menos que las imágenes preconfiguradas no funcionen en tu caso. Para configurar tu propia imagen de contenedor, se deben 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 mediante la configuración de 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 los Requisitos de software y las configuraciones de compilación probadas.

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

Dataflow te permite configurar el tipo y la cantidad de GPU que se conectarán a los trabajadores de Dataflow mediante el parámetro worker_accelerator. Puedes seleccionar el tipo y la cantidad de GPU en función de tu caso práctico y cómo planeas usar las GPU en tu canalización.

Los siguientes tipos de GPU son compatibles con Dataflow:

  • NVIDIA® A100
  • NVIDIA® Tesla® T4
  • NVIDIA® Tesla® P4
  • NVIDIA® Tesla® V100
  • NVIDIA® Tesla® P100
  • NVIDIA® Tesla® K80

Para obtener información más detallada sobre cada tipo de GPU, incluidos los datos de rendimiento, lee el Gráfico de comparación de GPU.

Ejecuta tu trabajo con GPU

Para ejecutar un trabajo de Dataflow con GPU, usa el siguiente comando:

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: un extremo regional, por ejemplo, us-central1
  • IMAGE: la ruta de acceso de Container 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

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

Si usas TensorFlow, considera configurar los trabajadores para que usen un solo proceso mediante la configuración de una opción de canalización --experiments=no_use_multiple_sdk_containers o usar trabajadores con una CPU virtual. Si n1-standard-1 no proporciona suficiente memoria, puedes considerar un tipo personalizado de máquina, como n1-custom-1-NUMBER_OF_MB o n1-custom-1-NUMBER_OF_MB-ext para tener memoria extendida. Para obtener más información, consulta Paralelismo de trabajadores y GPU.

Verifica tu trabajo de Dataflow

Para confirmar que el trabajo use VM de trabajador con GPU, 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. Pega el ID de trabajo en el mensaje Buscar productos y recursos.
    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 que 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', ...]

Ve el uso de GPU

Para ver el uso de GPU en las VM 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. Especifica Dataflow Job como Tipo de recurso y GPU utilization o GPU memory utilization como métrica, según la métrica que desees supervisar.

Para obtener más información, lee la guía Explorador de métricas.

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, sigue estos pasos de solución de problemas.

Los trabajadores no inician

Si tu trabajo está atascado y los trabajadores de Dataflow nunca comienzan a procesar datos, es probable que tengas un problema relacionado con el uso de un contenedor personalizado con Dataflow. Para obtener más detalles, consulta la guía de solución de problemas de contenedores personalizados.

Si eres usuario de Python, verifica que se cumplan las siguientes condiciones:

  • La versión secundaria de intérprete de Python en tu imagen de contenedor es la misma versión que usas cuando inicias tu canalización. Si no hay coincidencia, es posible que veas errores como SystemError: unknown opcode con un seguimiento de pila que involucra apache_beam/internal/pickler.py.
  • Si usas el SDK 2.29.0 de Apache Beam o una versión anterior, se debe poder acceder a pip en la imagen en /usr/local/bin/pip.

Te recomendamos que reduzcas las personalizaciones a una configuración funcional mínima la primera vez que uses una imagen personalizada. Usa las imágenes de contenedor personalizadas de muestra que se proporcionan en los ejemplos de esta página. Asegúrate de que puedas ejecutar una canalización de Dataflow simple con esta imagen de contenedor sin solicitar GPUs. Luego, itera en la solución.

Verifica que los trabajadores tengan suficiente espacio en el disco para descargar la imagen del contenedor. Ajusta el tamaño del disco si es necesario. Las imágenes grandes tardan más en descargarse, lo que aumenta el tiempo de inicio de los trabajadores.

El trabajo falla inmediatamente en el inicio

Si encuentras los errores ZONE_RESOURCE_POOL_EXHAUSTED o ZONE_RESOURCE_POOL_EXHAUSTED_WITH_DETAILS, puedes seguir estos pasos:

  • No especifiques la zona del trabajador para que Dataflow seleccione la zona óptima.

  • Inicia la canalización en una zona diferente o con un tipo de acelerador diferente.

El trabajo falla en el entorno de ejecución

Si el trabajo falla en el entorno de ejecución, revisa los errores de memoria insuficiente (OOM) en la máquina de trabajador y en la GPU. Los errores de OOM de GPU pueden manifestarse como errores cudaErrorMemoryAllocation out of memory en los registros de trabajador. Si usas TensorFlow, verifica que uses solo un proceso de TensorFlow para acceder a un dispositivo de GPU. Para obtener más información, consulta Paralelismo de trabajadores y GPU.

No hay uso de GPU

Si tu canalización se ejecuta de forma correcta, pero no se usan las GPU, verifica lo siguiente:

  • Las bibliotecas de NVIDIA instaladas en la imagen del contenedor coinciden con los requisitos del código de usuario de la canalización y las bibliotecas que usa.
  • Se puede acceder a las bibliotecas de NVIDIA en imágenes de contenedor como bibliotecas compartidas.

Si los dispositivos no están disponibles, puede que estés usando una configuración de software incompatible. Por ejemplo, si usas TensorFlow, verifica que tengas una combinación compatible de TensorFlow, la versión de cuDNN y la versión de CUDA Toolkit.

Para verificar la configuración de la imagen, considera ejecutar una canalización simple que verifique que las GPU estén disponibles y que los trabajadores puedan acceder a ellas.

Depura con una VM independiente

Mientras diseñas e iteras en una imagen de contenedor que funciona para ti, puede ser más rápido reducir el ciclo de reacción si pruebas tu imagen de contenedor en una VM independiente.

Puedes depurar tu contenedor personalizado en una VM independiente con GPU si creas una VM de Compute Engine que ejecute GPU en Container-Optimized OS, la instalación de controladores y el inicio de tu contenedor.

  1. Cree una instancia de VM:

    gcloud compute instances create INSTANCE_NAME \
  --project "PROJECT" \
  --image-family cos-stable \
  --image-project=cos-cloud  \
  --zone=us-central1-f \
  --accelerator type=nvidia-tesla-t4,count=1 \
  --maintenance-policy TERMINATE \
  --restart-on-failure  \
  --boot-disk-size=200G \
  --scopes=cloud-platform

  2. Usa ssh para conectarte a la VM:

    gcloud compute ssh INSTANCE_NAME --project "PROJECT"

  3. Instala los controladores de GPU. Después de conectarte a la VM mediante ssh, ejecuta los siguientes comandos en la VM:

    # Run these commands on the virtual machine
cos-extensions install gpu
sudo mount --bind /var/lib/nvidia /var/lib/nvidia
sudo mount -o remount,exec /var/lib/nvidia
/var/lib/nvidia/bin/nvidia-smi

  4. Inicia tu contenedor personalizado.

    Los contenedores del SDK de Apache Beam usan el punto de entrada /opt/apache/beam/boot. Para fines de depuración, puedes iniciar tu contenedor de forma manual con un punto de entrada diferente:

    docker-credential-gcr configure-docker
docker run --rm \
  -it \
  --entrypoint=/bin/bash \
  --volume /var/lib/nvidia/lib64:/usr/local/nvidia/lib64 \
  --volume /var/lib/nvidia/bin:/usr/local/nvidia/bin \
  --privileged \
  IMAGE

    Reemplaza IMAGE con la ruta de acceso de Container Registry para tu imagen de Docker.

  5. Luego, verifica que las bibliotecas de GPU instaladas en tu contenedor puedan acceder a los dispositivos de GPU.

    Si usas TensorFlow, puedes imprimir dispositivos disponibles en el intérprete de Python con lo siguiente:

    >>> import tensorflow as tf
>>> print(tf.config.list_physical_devices("GPU"))

    Si usas PyTorch, puedes inspeccionar los dispositivos disponibles en el intérprete de Python con lo siguiente:

    >>> import torch
>>> print(torch.cuda.is_available())
>>> print(torch.cuda.device_count())
>>> print(torch.cuda.get_device_name(0))

Para iterar en la canalización, puedes iniciarla en Direct Runner. También puedes iniciar canalizaciones en Dataflow Runner desde este entorno.

Para obtener información adicional, consulta:

¿Qué sigue?