Como usar GPUs

Nesta página, explicamos como executar um pipeline do Apache Beam no Dataflow com GPUs. Jobs que usam GPUs geram cobranças conforme especificado na página de preços do Dataflow.

Para mais informações sobre o uso de GPUs com o Dataflow, leia Suporte do Dataflow para GPUs.

Como usar os notebooks do Apache Beam

Se você já tiver um pipeline que queira executar com GPUs no Dataflow, pule esta seção.

Os notebooks do Apache Beam oferecem uma maneira conveniente de prototipar e desenvolver iterativamente o pipeline com GPUs sem configurar um ambiente de desenvolvimento. Para começar, leia o guia Como desenvolver com os notebooks do Apache Beam, inicie uma instância de notebooks do Apache Beam e siga o exemplo Usar GPUs com Apache Beam.

Como provisionar a cota da GPU

Os dispositivos da GPU estão sujeitos à disponibilidade de cota do projeto do Google Cloud. Solicite a cota de GPU na região de sua escolha.

Como instalar drivers de GPU

É necessário instruir o Dataflow a instalar drivers NVIDIA nos workers. Para isso, adicione install-nvidia-driver à opção worker_accelerator. Quando a opção install-nvidia-driver é especificada, o Dataflow instala drivers NVIDIA nos workers do Dataflow usando o utilitário cos-extensions fornecido pelo Container-Optimized OS. Ao especificar install-nvidia-driver, os usuários concordam em aceitar o contrato de licença da NVIDIA.

Binários e bibliotecas fornecidos pelo instalador de driver da NVIDIA são montados no contêiner que executa o código do usuário do pipeline em /usr/local/nvidia/.

A versão do driver da GPU depende da versão do Container-Optimized OS usada atualmente pelo Dataflow.

Como criar uma imagem de contêiner personalizada

Para interagir com as GPUs, talvez seja necessário usar software NVIDIA adicional, como bibliotecas aceleradas por GPU e o Kit de ferramentas do CUDA (links em inglês). É necessário fornecer essas bibliotecas no contêiner do Docker que está executando o código do usuário.

É possível personalizar a imagem de contêiner fornecendo uma imagem que atenda ao contrato de imagem de contêiner do SDK do Apache Beam e que tenha as bibliotecas de GPU necessárias.

Para fornecer uma imagem de contêiner personalizada, é necessário usar o Dataflow Runner v2 e fornecer a imagem do contêiner usando a opção worker_harness_container_image do pipeline. Se você estiver usando o Apache Beam 2.30.0 ou posterior, poderá usar um nome de opção mais curto sdk_container_image para simplificar. Para mais informações, consulte Como usar contêineres personalizados.

Abordagem 1. Como usar uma imagem existente configurada para uso da GPU

É possível criar uma imagem do Docker que atenda ao contrato de contêiner do SDK do Apache Beam a partir de uma imagem de base pré-configurada para uso da GPU. Por exemplo, as imagens do Docker do TensorFlow e as imagens de contêiner da NVIDIA são pré-configuradas para uso da GPU.

Veja abaixo um exemplo de Dockerfile criado com base na imagem do Docker do TensorFlow com o Python 3.6:

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 --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 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" ]

Ao usar as imagens do Docker do TensorFlow, use o TensorFlow 2.5.0 ou posterior. Imagens anteriores do Docker do TensorFlow instalam o pacote tensorflow-gpu em vez do pacote tensorflow. A distinção não é importante após o lançamento do TensorFlow 2.1.0, mas vários pacotes downstream, como tfx, exigem o pacote tensorflow.

Os contêineres grandes atrasam a inicialização do worker. Isso pode ocorrer ao usar contêineres como Deep Learning Containers.

Como instalar uma versão específica do Python

Se você tiver requisitos rígidos para a versão do Python, poderá criar sua imagem de uma imagem de base NVIDIA que tenha bibliotecas de GPU necessárias e instalar o interpretador do Python.

O exemplo a seguir demonstra a seleção de uma imagem NVIDIA do catálogo de imagens do contêiner CUDA que não inclui o interpretador de Python. É possível ajustar o exemplo para instalar a versão desejada do Python 3 e do pip. O exemplo usa o TensorFlow, portanto, ao escolher uma imagem, garantimos que as versões CUDA e cuDNN da imagem base atendam aos requisitos da versão do TensorFlow.

Uma amostra do Dockerfile se parece com o seguinte:

# 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 \
    # 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" ]

Em algumas distribuições do SO, pode ser difícil instalar versões específicas do Python usando o gerenciador de pacotes do SO. Nesse caso, você pode instalar o interpretador do Python com ferramentas como o Miniconda ou o pyenv.

Uma amostra do Dockerfile se parece com o seguinte:

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 \
    && sh 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 \
    # 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.
    # This can be omitted in 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" ]

Abordagem 2. Como usar imagens de contêiner do Apache Beam

É possível configurar uma imagem de contêiner para uso da GPU sem usar imagens pré-configuradas. Essa abordagem não é recomendada, a menos que as imagens pré-configuradas não funcionem para você. A configuração da própria imagem de contêiner requer a seleção de bibliotecas compatíveis e a configuração do ambiente de execução.

Uma amostra do Dockerfile se parece com o seguinte:

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

As bibliotecas de driver em /usr/local/nvidia/lib64 precisam ser descobertas no contêiner como bibliotecas compartilhadas configurando a variável de ambiente LD_LIBRARY_PATH.

Para usar o TensorFlow, é preciso escolher uma combinação compatível de versões do CUDA Toolkit e do cuDNN. Para ver mais detalhes, leia Requisitos de software e Configurações do build testados (links em inglês).

Como selecionar o tipo e o número de GPUs para workers do Dataflow

O Dataflow permite configurar o tipo e o número de GPUs a serem anexadas aos workers do Dataflow usando o parâmetro worker_accelerator. Selecione o tipo e número de GPUs com base no seu caso de uso e como você planeja utilizar as GPUs no pipeline.

Os seguintes tipos de GPU são compatíveis com o Dataflow:

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

Para informações mais detalhadas sobre cada tipo de GPU, incluindo dados de desempenho, leia o gráfico de comparação de GPUs.

Como executar um job com GPUs

Para executar um job do Dataflow com GPUs, use o seguinte 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" \
  --experiments "worker_accelerator=type:GPU_TYPE;count:GPU_COUNT;install-nvidia-driver" \
  --experiments "use_runner_v2"

Substitua:

  • PIPELINE: o arquivo de código-fonte do pipeline
  • PROJECT: o nome do projeto do Google Cloud.
  • BUCKET: o bucket do Cloud Storage.
  • REGION: um endpoint regional. Por exemplo, us-central1
  • IMAGE: o caminho do Container Registry da imagem do Docker
  • DISK_SIZE_GB: tamanho do disco de inicialização de cada VM de worker. Por exemplo, 50
  • GPU_TYPE: um tipo de GPU disponível. Por exemplo, nvidia-tesla-t4
  • GPU_COUNT: número de GPUs a serem anexadas a cada VM de worker. Por exemplo, 1

As considerações para executar um job do Dataflow com GPUs incluem:

Se você usa o TensorFlow, configure os workers para usar um único processo definindo uma opção de pipeline --experiments=no_use_multiple_sdk_containers ou usando workers com uma vCPU. Se o n1-standard-1 não fornecer memória suficiente, considere um tipo de máquina personalizado, como n1-custom-1-NUMBER_OF_MB ou n1-custom-1-NUMBER_OF_MB-ext, para memória estendida. Para mais informações, leia GPUs e paralelismo de workers.

Como verificar o job do Dataflow

Para confirmar que o job usa VMs de worker com GPUs, siga estas etapas:

  1. Verifique se os workers do Dataflow para o job foram iniciados.
  2. Durante a execução de um job, encontre uma VM de worker associada a ele.
    1. Cole o ID do job no prompt Pesquisar produtos e recursos.
    2. Selecione a instância de VM do Compute Engine associada ao job.

Também é possível encontrar uma lista de todas as instâncias em execução no console do Compute Engine.

  1. No Console do Google Cloud, acesse a página Instâncias de VM.

    Acessar instâncias de VM

  2. Clique em Detalhes da instância de VM.

  3. Verifique se a página de detalhes tem uma seção GPUs e se as GPUs estão anexadas.

Se o job não foi iniciado com GPUs, verifique se o experimento --worker_accelerator está configurado corretamente e visível na IU de monitoramento do Dataflow em experiments. A ordem dos tokens nos metadados do acelerador é importante.

Por exemplo, uma opção de pipeline "experimentos" na IU de monitoramento do Dataflow pode ter a seguinte aparência:

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

Como visualizar a utilização da GPU

Para ver a utilização da GPU nas VMs de worker, siga estas etapas:

  1. No Console do Google Cloud, acesse Monitoring ou clique no seguinte botão:

    Acessar o Monitoramento

  2. No painel de navegação do Monitoring, clique em Metrics Explorer.

  3. Especifique Dataflow Job como Tipo de recurso e GPU utilization ou GPU memory utilization como métrica, dependendo da métrica que você quer monitorar.

Para mais informações, leia o guia do Metrics Explorer.

Como usar GPUs com o Dataflow Prime

O Dataflow Prime permite solicitar aceleradores para uma etapa específica do pipeline. Para usar GPUs com o Dataflow Prime, não use a opção de pipeline --experiments=worker_accelerator. Em vez disso, solicite as GPUs com a dica de recurso accelerator. Para ver mais informações, consulte Como usar dicas de recursos.

Solução de problemas do job do Dataflow

Se você tiver problemas ao executar o job do Dataflow com GPUs, siga as etapas de solução de problemas abaixo que podem resolver seu problema.

Os workers não são iniciados

Se o job estiver travado e os workers do Dataflow nunca começarem a processar dados, é provável que você tenha um problema relacionado ao uso de um contêiner personalizado com o Dataflow. Para mais detalhes, leia o guia de solução de problemas de contêineres personalizados.

Se você for usuário do Python, verifique se as seguintes condições são atendidas:

  • A versão secundária do interpretador do Python na imagem do contêiner é a mesma que você usa ao iniciar o pipeline. No caso de incompatibilidade, é possível ver erros como SystemError: unknown opcode com um rastreamento de pilha que envolve apache_beam/internal/pickler.py.
  • Se você estiver usando o SDK 2.29.0 ou anterior do Apache Beam, pip precisará estar acessível na imagem em /usr/local/bin/pip.

Recomendamos que você reduza as personalizações a um mínimo de trabalho na primeira vez que usar uma imagem personalizada. Use as imagens de contêiner personalizadas de amostra fornecidas nos exemplos desta página. Verifique se é possível executar um pipeline simples do Dataflow com essa imagem de contêiner sem solicitar GPUs e iterar na solução.

Verifique se os workers têm espaço em disco suficiente para fazer o download da imagem do contêiner e ajuste o tamanho do disco, se necessário. Imagens grandes levam mais tempo para fazer o download, o que aumenta o tempo de inicialização do worker.

O job falha imediatamente na inicialização

Se você encontrar os erros ZONE_RESOURCE_POOL_EXHAUSTED ou ZONE_RESOURCE_POOL_EXHAUSTED_WITH_DETAILS, siga estas etapas:

  • Não especifique a zona do worker para que o Dataflow selecione a zona ideal para você.

  • Inicie o pipeline em uma zona diferente ou com um tipo de acelerador diferente.

O job falha no ambiente de execução

Se o job falhar em tempo de execução, exclua os erros de memória insuficiente na máquina worker e na GPU. Os erros de OOM da GPU podem se manifestar como erros cudaErrorMemoryAllocation out of memory nos registros do worker. Se você estiver usando o TensorFlow, verifique se é usado apenas um processo do TensorFlow para acessar um dispositivo de GPU. Para mais informações, leia GPUs e paralelismo de worker.

Nenhum uso da GPU

Se o pipeline for executado com sucesso, mas as GPUs não forem usadas, verifique se:

  • As bibliotecas NVIDIA instaladas na imagem do contêiner correspondem aos requisitos do código do usuário do pipeline e das bibliotecas que ele usa.
  • As bibliotecas NVIDIA instaladas em imagens de contêiner podem ser acessadas como bibliotecas compartilhadas.

Se os dispositivos não estiverem disponíveis, é possível que você esteja usando uma configuração de software incompatível. Por exemplo, se você estiver usando o TensorFlow, verifique se tem uma combinação compatível (em inglês) com as versões do TensorFlow, do cuDNN e do CUDA Toolkit.

Para verificar a configuração da imagem, execute um pipeline simples que apenas verifique se as GPUs estão disponíveis e acessíveis para os workers.

Depurar com uma VM independente

Enquanto você projeta e itera em uma imagem de contêiner que funciona para você, pode ser mais rápido reduzir o loop de feedback testando a imagem de contêiner em uma VM independente.

É possível depurar o contêiner personalizado em uma VM autônoma com GPUs. Basta criar uma VM do Compute Engine que executa GPUs no Container-Optimized OS, instalar drivers e iniciar o contêiner da seguinte maneira.

  1. Crie uma instância 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. Use ssh para se conectar à VM.

    gcloud compute ssh INSTANCE_NAME --project "PROJECT"
    
  3. Instale os drivers da GPU. Depois de se conectar à VM via ssh, execute os seguintes comandos na 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. Inicie seu contêiner personalizado.

    Os contêineres do SDK do Apache Beam usam o ponto de entrada /opt/apache/beam/boot. Para fins de depuração, é possível iniciar o contêiner manualmente com um ponto de entrada diferente, como mostrado abaixo:

    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
    

    Substitua IMAGE pelo caminho do Container Registry da imagem do Docker.

  5. Verifique se as bibliotecas de GPU instaladas no contêiner podem acessar os dispositivos da GPU.

    Se estiver usando o TensorFlow, é possível imprimir os dispositivos disponíveis no interpretador do Python com o seguinte:

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

    Se estiver usando o PyTorch, será possível inspecionar dispositivos disponíveis no interpretador do Python com o seguinte:

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

Para iterar no pipeline, inicie o pipeline no Direct Runner. Também é possível iniciar pipelines no Dataflow Runner a partir desse ambiente.

Para mais informações, consulte:

A seguir