Conteinerizar e executar o código de treinamento localmente

Use o comando gcloud ai custom-jobs local-run para criar uma imagem de contêiner do Docker com base no código de treinamento e executá-la como um contêiner no computador local. Esse recurso oferece vários benefícios:

  • É possível criar uma imagem de contêiner com conhecimento mínimo do Docker. Você não precisa escrever seu próprio Dockerfile. Depois, envie essa imagem para o Artifact Registry e use-a para o treinamento de contêineres personalizados.

    Para casos de uso avançados, ainda convém gravar seu próprio Dockerfile.

  • A imagem do contêiner pode executar um aplicativo de treinamento em Python ou um script Bash.

    É possível usar um script Bash para executar um código de treinamento escrito em outra linguagem de programação, desde que você também especifique uma imagem de contêiner base compatível com a outra linguagem.

  • A execução local de um contêiner executa o código de treinamento de maneira semelhante à execução no Vertex AI.

    Executar o código localmente pode ajudar a depurar problemas com o código antes de executar o treinamento personalizado no Vertex AI.

Antes de começar

  1. Configure o ambiente de desenvolvimento do Vertex AI.

  2. Instale o Docker Engine.

  3. Se você estiver usando o Linux, configure o Docker para executá-lo sem sudo.

    O comando local-run requer essa configuração para usar o Docker.

Usar o comando local-run

Execute o comando a seguir para criar uma imagem de contêiner com base no código de treinamento e executar um contêiner localmente:

gcloud ai custom-jobs local-run \
  --executor-image-uri=BASE_IMAGE_URI \
  --local-package-path=WORKING_DIRECTORY \
  --script=SCRIPT_PATH \
  --output-image-uri=OUTPUT_IMAGE_NAME

Substitua:

  • BASE_IMAGE_URI: o URI de uma imagem do Docker a ser usada como base do contêiner. Escolha uma imagem de base que inclua as dependências necessárias do código de treinamento.

    É possível usar o URI para uma imagem de container de treinamento pré-criado ou qualquer outro valor que seja válido para uma instrução FROM do Dockerfile, por exemplo, uma imagem do Docker disponível publicamente ou uma imagem do Docker no Artifact Registry a que você tem acesso.

  • WORKING_DIRECTORY: o diretório de nível mais baixo no sistema de arquivos que contém todo o código de treinamento e as dependências locais que você precisa usar para o treinamento.

    Por padrão, o comando copia apenas o diretório pai do arquivo especificado pela sinalização --script (consulte o item de lista a seguir) para a imagem do Docker resultante. A imagem do Docker não inclui necessariamente todos os arquivos em WORKING_DIRECTORY. Para personalizar quais arquivos serão incluídos, consulte a seção deste documento sobre como incluir dependências.

    Se você omitir a sinalização --local-package-path (e esse marcador), o comando local-run usará o diretório de trabalho atual para esse valor.

  • SCRIPT_PATH: o caminho, relativo a WORKING_DIRECTORY no seu sistema de arquivos local, para o script que é o ponto de entrada do seu código de treinamento. Pode ser um script Python (terminado em .py) ou um script Bash.

    Por exemplo, se você quiser executar /hello-world/trainer/task.py e WORKING_DIRECTORY for /hello-world, use trainer/task.py para esse valor.

    Se você especificar um script Python, sua imagem de base precisará ter o Python instalado. Caso seja especificado um script bash, sua imagem de base precisará ter o Bash instalado. Todos os contêineres de treinamento pré-criados e muitas outras imagens do Docker disponíveis publicamente incluem essas duas dependências.

    Use --python-module em vez de --script

    Se você omitir a sinalização --script (e SCRIPT_PATH), use a sinalização --python-module para especificar o nome de um módulo Python em WORKING_DIRECTORY para ser executado como o ponto de entrada para o treinamento. Por exemplo, em vez de --script=trainer/task.py, é possível especificar --python-module=trainer.task.

    Nesse caso, o contêiner do Docker resultante carrega o código como um módulo, e não como um script. Essa opção provavelmente será usada se o script de ponto de entrada importar outros módulos Python em WORKING_DIRECTORY.

  • OUTPUT_IMAGE_NAME: um nome para a imagem do Docker resultante criada pelo comando. É possível usar qualquer valor aceito pela sinalização -t da docker build.

    Se você planeja enviar posteriormente a imagem para o Artifact Registry, convém usar um nome de imagem que atenda aos requisitos do Artifact Registry. Se você planeja enviar a imagem posteriormente para o Container Registry, use um nome de imagem que atenda aos requisitos do Container Registry. Você também pode marcar a imagem com nomes adicionais posteriormente.

    Se você omitir a sinalização --output-image-uri (e esse marcador), o comando local-run marcará a imagem com um nome baseado no horário atual e no nome do arquivo de script do ponto de entrada.

O comando cria uma imagem de contêiner do Docker com base na sua configuração. Depois de criar a imagem, o comando imprime a seguinte saída:

A training image is built.
Starting to run ...

O comando usa imediatamente essa imagem de contêiner para executar um contêiner no seu computador local. Quando o contêiner sai, o comando imprime a seguinte saída:

A local run is finished successfully using custom image: OUTPUT_IMAGE_NAME

Opções adicionais

As seções a seguir descrevem opções adicionais que podem ser usadas para personalizar o comportamento do comando local-run.

Instalar dependências

O código de treinamento pode depender de todas as dependências instaladas na imagem de base. Por exemplo, as imagens de contêiner de treinamento pré-criadas incluem muitas bibliotecas Python para machine learning, bem como quaisquer arquivos que você inclui na imagem do Docker criada pelo comando local-run.

Quando você especifica um script com a sinalização --script ou --python-module, o comando copia o diretório pai do script e os subdiretórios para a imagem do Docker. Por exemplo, se você especificar --local-package-path=/hello-world e --script=trainer/task.py, o comando copiará /hello-world/trainer/ para a imagem do Docker.

Também é possível incluir outras dependências do Python ou arquivos arbitrários do seu sistema de arquivos concluindo as etapas extras descritas em uma das seguintes seções:

Instalar dependências extras do Python

É possível incluir outras dependências do Python na imagem do Docker de várias maneiras:

Usar um arquivo requirements.txt

Se houver um arquivo chamado requirements.txt no diretório de trabalho, o comando local-run o tratará como um arquivo de requisitos de pip e o usará para instalar dependências do Python na imagem do Docker.

Usar um arquivo setup.py

Se houver um arquivo chamado setup.py no diretório de trabalho, o comando local-run tratará isso como um arquivo setup.py do Python, copiará o arquivo para imagem do Docker e executará pip install no diretório na imagem do Docker que contém esse arquivo.

É possível, por exemplo, adicionar um argumento install_requires a setup.py para instalar dependências do Python na imagem do Docker.

Especificar dependências individuais do PyPI

É possível usar a sinalização --requirements para instalar dependências específicas do PyPI na imagem do Docker. Exemplo:

gcloud ai custom-jobs local-run \
  --executor-image-uri=BASE_IMAGE_URI \
  --local-package-path=WORKING_DIRECTORY \
  --script=SCRIPT_PATH \
  --output-image-uri=OUTPUT_IMAGE_NAME \
  --requirements=REQUIREMENTS

Substitua REQUIREMENTS por uma lista separada por vírgulas de especificadores de requisito do Python.

Especificar outras dependências locais do Python

Use a sinalização --extra-packages para instalar dependências locais específicas do Python. Essas dependências do Python precisam estar no diretório de trabalho e cada dependência precisa estar em um formato compatível com pip install. Por exemplo, um arquivo de roda ou uma distribuição de origem Python.

Exemplo:

gcloud ai custom-jobs local-run \
  --executor-image-uri=BASE_IMAGE_URI \
  --local-package-path=WORKING_DIRECTORY \
  --script=SCRIPT_PATH \
  --output-image-uri=OUTPUT_IMAGE_NAME \
  --extra-packages=LOCAL_DEPENDENCIES

Substitua LOCAL_DEPENDENCIES por uma lista separada por vírgulas de caminhos de arquivo local, expressa em relação ao diretório de trabalho.

Incluir outros arquivos

Para copiar diretórios adicionais para a imagem do Docker (sem instalá-los como dependências do Python), use a sinalização --extra-dirs. Você pode especificar diretórios somente no diretório de trabalho. Exemplo:

gcloud ai custom-jobs local-run \
  --executor-image-uri=BASE_IMAGE_URI \
  --local-package-path=WORKING_DIRECTORY \
  --script=SCRIPT_PATH \
  --output-image-uri=OUTPUT_IMAGE_NAME \
  --extra-dirs=EXTRA_DIRECTORIES

Substitua EXTRA_DIRECTORIES por uma lista separada por vírgulas de diretórios locais, expressa em relação ao diretório de trabalho.

Argumentos do aplicativo de treinamento

Se o script de ponto de entrada de seu aplicativo de treinamento esperar argumentos de linha de comando, especifique-os ao executar o comando local-run. Esses argumentos não são salvos na imagem do Docker. Em vez disso, eles são passados como argumentos quando a imagem é executada como um contêiner.

Para transmitir argumentos ao script do ponto de entrada, passe o argumento -- seguido dos argumentos do script para o comando local-run depois de todas as outras sinalizações do comando.

Por exemplo, imagine um script que é executado localmente com o seguinte comando:

python /hello-world/trainer/task.py \
  --learning_rate=0.1 \
  --input_data=gs://BUCKET/small-dataset/

Quando você usa o comando local-run, é possível usar as seguintes sinalizações para executar o script no contêiner com os mesmos argumentos:

gcloud ai custom-jobs local-run \\
  --executor-image-uri=BASE_IMAGE_URI \
  --local-package-path=/hello-world \
  --script=/trainer/task.py \
  --output-image-uri=OUTPUT_IMAGE_NAME \
  -- \
  --learning_rate=0.1 \
  --input_data=gs://BUCKET/small-dataset/

Acelerar o treinamento de modelos com GPUs

Se você quiser implantar a imagem do Docker criada pelo comando local-run no Vertex AI e usar GPUs para treinamento, certifique-se de gravar o código de treinamento que aproveita as GPUs e usa uma imagem do Docker ativada para GPU para o valor da sinalização --executor-image-uri. Por exemplo, é possível usar uma das imagens de contêiner de treinamento pré-criadas que são compatíveis com GPUs.

Se o computador local executa o Linux e tem GPUs, também é possível configurar o comando local-run para usar suas GPUs quando executar um contêiner localmente. Isso é opcional, mas pode ser útil se você quiser testar como seu código de treinamento funciona com GPUs. Faça o seguinte:

  1. Instale o NVIDIA Container Toolkit (nvidia-docker) no seu computador local, caso ainda não tenha feito isso.

  2. Use a sinalização --gpu ao executar o comando local-run. Exemplo:

    gcloud ai custom-jobs local-run \
  --executor-image-uri=BASE_IMAGE_URI \
  --local-package-path=WORKING_DIRECTORY \
  --script=SCRIPT_PATH \
  --output-image-uri=OUTPUT_IMAGE_NAME \
  --gpu

Especificar uma conta de serviço personalizada

Por padrão, quando olocal-run executa o código de treinamento em um contêiner local, ele ativa as credenciais do Google Cloud disponíveis no ambiente local por meio do comando Application Default Credentials (ADC) no contêiner, para que o código de treinamento possa usar o ADC para autenticação com as mesmas credenciais. Em outras palavras, as credenciais disponíveis pelo ADC no shell local também estão disponíveis pelo ADC no seu código quando você executa o comando local-run.

É possível usar o comandogcloud auth application-default login para usar sua conta de usuário para o ADC ou definir uma variável de ambiente no shell para usar um serviço pelo ADC.

Se você quiser que o contêiner seja executado com credenciais do Google Cloud diferentes das disponíveis pelo ADC no shell local, faça o seguinte:

  1. Crie ou selecione uma conta de serviço com as permissões que você quer que o código de treinamento tenha acesso.

  2. Faça o download de uma chave de conta de serviço para essa conta de serviço no computador local.

  3. Ao executar o comando local-run, especifique a sinalização --service-account-key-file. Exemplo:

    gcloud ai custom-jobs local-run \
  --executor-image-uri=BASE_IMAGE_URI \
  --local-package-path=WORKING_DIRECTORY \
  --script=SCRIPT_PATH \
  --output-image-uri=OUTPUT_IMAGE_NAME \
  --service-account-key-file=KEY_PATH

    Substitua KEY_PATH pelo caminho para a chave da conta de serviço no seu sistema de arquivos local. Ele precisa ser absoluto ou relativo ao diretório de trabalho atual do shell, não relativo ao diretório especificado pelo flag --local-package-path.

No contêiner resultante, seu código de treinamento pode usar o ADC para autenticação com as credenciais de conta de serviço especificadas.

Comparação com o treinamento na Vertex AI

Quando você executa o treinamento personalizado na Vertex AI, a Vertex AI usa o Agente de serviço de código personalizado da Vertex AI do seu projeto por padrão para executar seu código. Também é possível anexar uma conta de serviço diferente para treinamento personalizado.

Ao usar o comando local-run, não é possível se autenticar como o agente de serviço do código personalizado da Vertex AI, mas é possível criar uma conta de serviço com permissões semelhantes e usá-la localmente.

A seguir