Criar e executar um job que usa TPUs

Neste documento, explicamos como criar e executar um job em lote que usa instâncias de máquina virtual (VM) com unidades de processamento de tensor (TPUs). As TPUs são circuitos integrados de aplicativo específico (ASICs) desenvolvidos pelo Google para acelerar as cargas de trabalho de machine learning (ML) e inteligência artificial (IA). Ao criar um job em lote, é possível especificar um tipo de máquina de TPU para acelerar cargas de trabalho de inferência, avaliação, ajuste e treinamento de IA e ML em contêineres ou baseadas em script.

Para saber mais sobre TPUs, consulte Sobre TPUs no Google Cloud e Tipos de máquina de TPU na documentação do Compute Engine. Como alternativa, é possível acelerar as cargas de trabalho de processamento de dados, IA e ML usando GPUs.

Antes de começar

  1. Se você nunca usou o Batch, consulte Introdução ao Batch e ative o Batch concluindo os pré-requisitos para projetos e usuários.

  2. Para receber as permissões necessárias para criar um job, peça ao administrador para conceder a você os seguintes papéis do IAM:

    Para mais informações sobre a concessão de papéis, consulte Gerenciar o acesso a projetos, pastas e organizações.

    Também é possível conseguir as permissões necessárias usando papéis personalizados ou outros papéis predefinidos.

Restrições

Quando um job usa TPUs, todas as restrições a seguir se aplicam, além das limitações do Compute Engine para TPUs.

  • Somente frações de TPU de host único: o Batch só oferece suporte a VMs de TPU independentes, chamadas de frações de TPU de host único. Frações de TPU de vários hosts, que consistem em várias VMs interconectadas, não são compatíveis.

    Um job pode ser executado em uma ou mais frações de TPU de host único. Por exemplo, se um job definir taskCount como 4 e taskCountPerNode como 1, o número de VMs em que o job é executado será 4 ou menos. Ao contrário das frações de TPU de vários hosts, as frações de TPU de host único operam como uma VM independente e não fornecem rede de interconexão entre chips (ICI) para comunicação de alta largura de banda e baixa latência.

    Saiba mais sobre frações de TPU de host único e de vários hosts na documentação do Compute Engine.

  • Versões de TPU e tipos de máquina compatíveis: o Batch só aceita frações de TPU de host único (VMs de TPU independentes) para as seguintes versões de TPU e tipos de máquina:

    • TPU7x (Ironwood): tpu7x-standard-4t

    • v6e (Trillium): ct6e-standard-1t, ct6e-standard-4t e ct6e-standard-8t

    • v5p: ct5p-hightpu-4t

    Para mais informações sobre cada tipo de máquina de TPU, como discos compatíveis, consulte Tipos de máquina de TPU na documentação do Compute Engine.

  • Opções de consumo compatíveis: o Batch não é compatível com VMs flex-start (modelo de provisionamento FLEX_START) para jobs que usam TPUs. Com exceção das VMs flex-start, é possível especificar qualquer opção de consumo compatível com a versão de TPU do job, de acordo com a disponibilidade de opções de consumo por versões de TPU na documentação do Compute Engine.

  • Local de zona única obrigatório: ao criar um job que usa TPUs, é necessário restringir o job a uma zona específica definindo os locais permitidos do job (campo allowedLocations[]) como uma zona que contenha sua versão de TPU ou, se aplicável, sua reserva.

  • Recursos não compatíveis: os jobs que usam TPUs não são compatíveis com os seguintes recursos:

Criar um job que use TPUs

Para criar um job que use TPUs, faça o seguinte:

  1. Planeje a configuração de um job que usa TPUs.
  2. Crie um job com a configuração que você identificou. Para exemplos de como criar um job que usa TPUs, consulte Criar um job que usa TPUs neste documento.

Planejar a configuração de um job que usa TPUs

Antes de criar um job que use TPUs, faça o seguinte:

  1. Escolha a versão da TPU: selecione a versão da TPU adequada para sua carga de trabalho. Para uma lista de versões de TPU por tipo de carga de trabalho, consulte Versões de TPU recomendadas por tipos de carga de trabalho.

    Selecione uma das versões de TPU compatíveis.

  2. Escolha uma opção de consumo de TPU: selecione uma opção de consumo que melhor se adapte à sua carga de trabalho, à duração dela e às suas necessidades de custo. Selecione uma das opções de consumo compatíveis. Os casos de uso e os requisitos de job para cada opção de consumo estão resumidos na tabela a seguir.

    Opções de consumo e os requisitos de trabalho delas

    Vinculado à reserva (reservas adiantadas no modo calendário ou reservas adiantadas para blocos de capacidade)

    • Caso de uso: recomendamos opções de consumo vinculadas à reserva para jobs que precisam de um nível muito alto de garantia de disponibilidade de recursos de TPU por um período específico. No Compute Engine, elas são descritas como reservas adiantadas no modo calendário (de 1 a 90 dias) e reservas adiantadas para blocos de capacidade (por 1 ano ou mais).

    • Detalhes: antes de criar um job, é necessário conseguir ou identificar uma reserva para a versão de TPU escolhida e o modelo de provisionamento RESERVATION_BOUND. Depois que uma reserva é criada, ela gera os custos das VMs especificadas pelo mesmo preço da execução das VMs até que você exclua a reserva. As VMs que estão consumindo uma reserva não geram custos separados, mas as reservas geram custos independentemente do consumo.

      Para mais informações, consulte Escolher um tipo de reserva na documentação do Compute Engine.

    O Batch consome reservas vinculadas a reservas para jobs que fazem o seguinte:

    • Defina o campo provisioningModel como RESERVATION_BOUND.
    • Segmentar especificamente uma reserva compatível. Defina o campo reservation como o nome curto (se a reserva estiver no mesmo projeto e zona) ou o caminho completo do recurso da reserva, do bloco de reserva ou do sub-bloco de reserva.

    Sob demanda

    • Caso de uso: o modelo on demand é rápido e simples de usar, por isso é a opção de consumo padrão para a maioria dos jobs. No entanto, não recomendamos o uso on demand se outra opção de consumo se adequar melhor às suas prioridades ou se você tiver problemas repetidos com a disponibilidade de recursos.

    • Detalhes: o acesso sob demanda geralmente é a maneira padrão de acessar as VMs do Compute Engine. Com o acesso on demand, você pode acessar recursos (se disponíveis) quando solicitado. Os recursos não são excluídos até que o job seja concluído ou o limite de tempo de execução seja atingido.

    O Batch usa o modo on demand para a maioria dos jobs por padrão. Se você quiser garantir que um job use o modo on demand, faça o seguinte:

    VMs spot

    • Caso de uso: recomendamos que você use VMs spot para reduzir custos de cargas de trabalho tolerantes a falhas.

    • Detalhes: se uma solicitação para criar um job que usa VMs spot for atendida, as VMs serão cobradas de acordo com os preços de VMs spot, que podem mudar com frequência e oferecer o maior desconto, de até 91% em relação aos preços sob demanda. No entanto, as VMs spot podem não estar sempre disponíveis e podem ser interrompidas a qualquer momento.

      Saiba mais sobre VMs spot e recomendações para jobs com VMs spot.

    O Batch usa VMs spot para jobs que fazem tudo o seguinte:

    • Defina o campo provisioningModel como SPOT.
    • (Opcional) Impedir o consumo de reservas. Não é possível usar VMs spot com reservas. Especificamente, o job pode definir o campo reservation como NO_RESERVATION. Para mais informações, consulte Criar e executar um job que não pode consumir VMs reservadas.

  3. Selecione uma zona em que as TPUs escolhidas estejam disponíveis: para usar uma versão de TPU, o campo allowedLocations[] de um job precisa ser definido como uma zona que contenha as TPUs escolhidas, da seguinte maneira:

    • Se você escolheu uma opção de consumo que usa reservas, use a zona em que a reserva está localizada.
    • Caso contrário, selecione uma zona que contenha sua versão de TPU de acordo com a disponibilidade de TPU.

    Em ambos os casos, a zona também precisa estar em uma das regiões com disponibilidade para o Batch.

  4. Verifique se o projeto tem cota de TPU suficiente: se você estiver criando um job que usa TPUs com VMs spot ou sob demanda, é preciso ter cota de TPU suficiente disponível na região que você quer usar. A criação de um job que usa TPUs que consomem uma reserva de TPU não requer nenhuma cota de TPU, já que a cota é usada quando a reserva é criada. Para uma lista de nomes de cotas de TPU, consulte Cota de TPU. Para instruções sobre como visualizar a cota, consulte Visualizar e gerenciar cotas.

  5. Escolha um tipo de máquina de TPU: selecione um tipo de máquina de TPU para a versão de TPU escolhida. Se você escolheu uma opção de consumo que usa reservas, use o tipo de máquina da sua reserva. Para uma lista de tipos de máquina disponíveis para cada versão de TPU, consulte Tipos de máquina de TPU.

    Devido às restrições, é necessário usar a topologia de host único para o tipo de máquina de TPU.

Criar um job que use TPUs

As seções a seguir explicam como criar um job de exemplo que usa TPUs para cada opção de consumo.

Opção de consumo Exemplos de instruções de trabalho
Vinculado à reserva (reservas adiantadas no modo calendário ou reservas adiantadas para blocos de capacidade)

Usar TPUs com opções de consumo vinculadas à reserva
VMs spot ou sob demanda

Usar TPUs com VMs spot ou on demand

Usar TPUs com VMs spot ou sob demanda

É possível criar um job que use TPUs com VMs spot ou sob demanda usando a CLI gcloud ou a API Batch.

gcloud

  1. Crie um arquivo de configuração JSON que especifique um tipo de máquina de TPU de host único compatível, defina o modelo de provisionamento como STANDARD ou SPOT e seja executado em um local que ofereça o tipo de máquina de TPU.

    Por exemplo, para criar um job de script básico, crie um arquivo JSON com o conteúdo a seguir:

    {
  "taskGroups": [
    {
      "taskSpec": {
        "runnables": [
          {
            "script": {
              "text": "echo Hello world from TPU task ${BATCH_TASK_INDEX}."
            }
          }
        ]
      },
      "taskCount": 3,
      "parallelism": 1
    }
  ],
  "allocationPolicy": {
    "instances": [
      {
        "policy": {
          "provisioningModel": "PROVISIONING_MODEL",
          "machineType": "MACHINE_TYPE"
        }
      }
    ],
    "location": {
      "allowedLocations": [
        "ALLOWED_LOCATIONS"
      ]
    }
  },
  "logsPolicy": {
    "destination": "CLOUD_LOGGING"
  }
}

    Substitua:

    • PROVISIONING_MODEL: o modelo de provisionamento da sua opção de consumo: STANDARD para VMs sob demanda ou SPOT para VMs spot.

    • MACHINE_TYPE: um tipo de máquina de TPU de host único compatível, por exemplo, ct5p-hightpu-4t.

    • ALLOWED_LOCATIONS: os locais permitidos para o job. Para um job que usa TPUs, é necessário definir o campo allowedLocations[] como uma zona que ofereça o tipo de máquina de TPU. Por exemplo, zones/us-central1-a oferece o tipo de máquina de TPU ct5p-hightpu-4t.

  2. Para criar e executar o job, use o comando gcloud batch jobs submit:

    gcloud batch jobs submit JOB_NAME \
    --location LOCATION \
    --config JSON_CONFIGURATION_FILE

    Substitua:

    • JOB_NAME: o nome do job.

    • LOCATION: o local do job.

    • JSON_CONFIGURATION_FILE: o caminho para um arquivo JSON com os detalhes de configuração do job.

API

Faça uma solicitação POST ao método jobs.create que especifica um tipo de máquina de TPU de host único compatível, define o modelo de provisionamento como STANDARD ou SPOT e é executado em um local que oferece o tipo de máquina de TPU.

Por exemplo, para criar um job de script básico, faça a seguinte solicitação:

POST https://batch.googleapis.com/v1/projects/PROJECT_ID/locations/LOCATION/jobs?job_id=JOB_NAME

{
  "taskGroups": [
    {
      "taskSpec": {
        "runnables": [
          {
            "script": {
              "text": "echo Hello world from TPU task ${BATCH_TASK_INDEX}."
            }
          }
        ]
      },
      "taskCount": 3,
      "parallelism": 1
    }
  ],
  "allocationPolicy": {
    "instances": [
      {
        "policy": {
          "provisioningModel": "PROVISIONING_MODEL",
          "machineType": "MACHINE_TYPE"
        }
      }
    ],
    "location": {
      "allowedLocations": [
        "ALLOWED_LOCATIONS"
      ]
    }
  },
  "logsPolicy": {
    "destination": "CLOUD_LOGGING"
  }
}

Substitua:

  • PROJECT_ID: o ID do projeto.

  • LOCATION: o local do job.

  • JOB_NAME: o nome do job.

  • PROVISIONING_MODEL: o modelo de provisionamento da sua opção de consumo: STANDARD para VMs sob demanda ou SPOT para VMs spot.

  • MACHINE_TYPE: um tipo de máquina de TPU de host único compatível, por exemplo, ct5p-hightpu-4t.

  • ALLOWED_LOCATIONS: os locais permitidos para o job. Para um job que usa TPUs, é necessário definir o campo allowedLocations[] como uma zona que ofereça o tipo de máquina de TPU. Por exemplo, zones/us-central1-a oferece o tipo de máquina de TPU ct5p-hightpu-4t.

Usar TPUs com opções de consumo vinculadas à reserva

É possível criar um job que use TPUs com uma opção de consumo para o modelo de provisionamento vinculado à reserva (RESERVATION_BOUND), seja reservas adiantadas para blocos de capacidade ou reservas adiantadas no modo calendário, usando a CLI gcloud ou a API Batch.

gcloud

  1. Crie um arquivo de configuração JSON que especifique uma reserva para um tipo de máquina de TPU de host único, defina o modelo de provisionamento como RESERVATION_BOUND e seja executado na zona da reserva.

    Por exemplo, para criar um job de script básico, crie um arquivo JSON com o conteúdo a seguir:

    {
  "taskGroups": [
    {
      "taskSpec": {
        "runnables": [
          {
            "script": {
              "text": "echo Hello world from TPU task ${BATCH_TASK_INDEX}."
            }
          }
        ]
      },
      "taskCount": 3,
      "parallelism": 1
    }
  ],
  "allocationPolicy": {
    "instances": [
      {
        "policy": {
          "provisioningModel": "RESERVATION_BOUND",
          "reservation": "RESERVATION",
          "machineType": "MACHINE_TYPE"
        }
      }
    ],
    "location": {
      "allowedLocations": [
        "ALLOWED_LOCATIONS"
      ]
    }
  },
  "logsPolicy": {
    "destination": "CLOUD_LOGGING"
  }
}

    Substitua:

    • RESERVATION: substitua por um dos seguintes caminhos para sua reserva:

      • Se você quiser usar uma reserva adiantada no modo calendário ou qualquer parte da reserva para uma reserva adiantada para blocos de capacidade, especifique a reserva:

        projects/PROJECT_ID/zones/ZONE/reservations/RESERVATION_NAME

      • Se você quiser usar um bloco específico de uma reserva adiantada para blocos de capacidade, especifique o bloco de reserva:

        projects/PROJECT_ID/zones/ZONE/reservations/RESERVATION_NAME/reservationBlocks/BLOCK_NAME

      • Se você quiser usar um sub-bloco específico de uma reserva adiantada para blocos de capacidade, especifique o sub-bloco de reserva:

        projects/PROJECT_ID/zones/ZONE/reservations/RESERVATION_NAME/reservationSubBlocks/SUB_BLOCK_NAME

    • MACHINE_TYPE: o tipo de máquina de TPU da reserva, por exemplo, ct5p-hightpu-4t.

    • ALLOWED_LOCATIONS: os locais permitidos para o job. Para um job que usa TPUs reservadas, defina o campo allowedLocations[] como a zona da sua reserva. Por exemplo, zones/us-central1-a oferece o tipo de máquina de TPU ct5p-hightpu-4t.

  2. Para criar e executar o job, use o comando gcloud batch jobs submit:

    gcloud batch jobs submit JOB_NAME \
    --location LOCATION \
    --config JSON_CONFIGURATION_FILE

    Substitua:

    • JOB_NAME: o nome do job.

    • LOCATION: o local do job.

    • JSON_CONFIGURATION_FILE: o caminho para um arquivo JSON com os detalhes de configuração do job.

API

Faça uma solicitação POST ao método jobs.create que especifique uma reserva para um tipo de máquina de TPU de host único, defina o modelo de provisionamento como RESERVATION_BOUND e seja executada na zona da reserva.

Por exemplo, para criar um job de script básico, faça a seguinte solicitação:

POST https://batch.googleapis.com/v1/projects/PROJECT_ID/locations/LOCATION/jobs?job_id=JOB_NAME

{
  "taskGroups": [
    {
      "taskSpec": {
        "runnables": [
          {
            "script": {
              "text": "echo Hello world from TPU task ${BATCH_TASK_INDEX}."
            }
          }
        ]
      },
      "taskCount": 3,
      "parallelism": 1
    }
  ],
  "allocationPolicy": {
    "instances": [
      {
        "policy": {
          "provisioningModel": "RESERVATION_BOUND",
          "reservation": "RESERVATION",
          "machineType": "MACHINE_TYPE"
        }
      }
    ],
    "location": {
      "allowedLocations": [
        "ALLOWED_LOCATIONS"
      ]
    }
  },
  "logsPolicy": {
    "destination": "CLOUD_LOGGING"
  }
}

Substitua:

  • PROJECT_ID: o ID do projeto.

  • LOCATION: o local do job.

  • JOB_NAME: o nome do job.

  • RESERVATION: substitua por um dos seguintes caminhos para sua reserva:

    • Se você quiser usar uma reserva adiantada no modo calendário ou qualquer parte da reserva para uma reserva adiantada para blocos de capacidade, especifique a reserva:

      projects/PROJECT_ID/zones/ZONE/reservations/RESERVATION_NAME

    • Se você quiser usar um bloco específico de uma reserva adiantada para blocos de capacidade, especifique o bloco de reserva:

      projects/PROJECT_ID/zones/ZONE/reservations/RESERVATION_NAME/reservationBlocks/BLOCK_NAME

    • Se você quiser usar um sub-bloco específico de uma reserva adiantada para blocos de capacidade, especifique o sub-bloco de reserva:

      projects/PROJECT_ID/zones/ZONE/reservations/RESERVATION_NAME/reservationSubBlocks/SUB_BLOCK_NAME

  • MACHINE_TYPE: o tipo de máquina de TPU da reserva, por exemplo, ct5p-hightpu-4t.

  • ALLOWED_LOCATIONS: os locais permitidos para o job. Para um job que usa TPUs reservadas, defina o campo allowedLocations[] como a zona da sua reserva. Por exemplo, zones/us-central1-a oferece o tipo de máquina de TPU ct5p-hightpu-4t.

Para mais informações sobre como criar jobs que consomem reservas, consulte Garantir a disponibilidade de recursos usando reservas de VM.

A seguir