Especifique a configuração da máquina para uma etapa do pipeline

Os componentes do pipeline do Kubeflow são funções de fábrica que criam etapas de pipeline. Cada componente descreve as entradas, saídas e a implementação do componente. Por exemplo, no exemplo de código abaixo, train_op é um componente.

Por exemplo, um componente de treinamento pode usar um arquivo CSV como entrada e usá-lo para treinar um modelo. Ao definir os parâmetros do tipo de máquina na etapa de pipeline, é possível gerenciar os requisitos de cada etapa no pipeline. Se você tem duas etapas de treinamento e uma etapa treina em um grande arquivo de dados e a segunda treina em um arquivo de dados pequeno, é possível alocar mais memória e CPU para a primeira tarefa e menos recursos para a segunda tarefa.

Por padrão, o componente será executado como uma Vertex AI CustomJob usando uma máquina e2-standard-4, com CPUs de quatro núcleos e 16 GB de memória. Para mais informações sobre como selecionar um dos recursos de máquina específicos do Google Cloud listados em Tipos de máquina, consulte Solicitar recursos de máquina do Google Cloud com os pipelines da Vertex AI.

O exemplo a seguir mostra como definir as configurações de CPU, memória e GPU para uma etapa:

from kfp import dsl

@dsl.pipeline(name='custom-container-pipeline')
def pipeline():
  generate = generate_op()
  train = (
    train_op(
      training_data=generate.outputs['training_data'],
      test_data=generate.outputs['test_data'],
      config_file=generate.outputs['config_file'])
    .set_cpu_limit('CPU_LIMIT')
    .set_memory_limit('MEMORY_LIMIT')
    .add_node_selector_constraint(SELECTOR_CONSTRAINT)
    .set_accelerator_limit(ACCELERATOR_LIMIT))

Substitua:

  • CPU_LIMIT: o limite máximo de CPU para este operador. Esse valor de string pode ser um número (valor inteiro para o número de CPUs) ou um número seguido de "m", o que significa 1/1000. É possível especificar, no máximo, 96 CPUs.

  • MEMORY_LIMIT: o limite máximo de memória para este operador. Esse valor de string pode ser um número ou um número seguido de "K" (kilobyte), "M" (megabyte) ou "G" (gigabyte). São aceitos no máximo 624 GB.

  • SELECTOR_CONSTRAINT: cada restrição é um rótulo de par de chave-valor. Para que o contêiner seja qualificado para execução em um nó, o nó precisa ter cada restrição como um rótulo. Por exemplo: 'cloud.google.com/gke-accelerator', 'NVIDIA_TESLA_K80'

    As seguintes restrições estão disponíveis:

    • NVIDIA_H100_80GB
    • NVIDIA_A100_80GB
    • NVIDIA_TESLA_A100 (NVIDIA A100 40GB)
    • NVIDIA_TESLA_K80
    • NVIDIA_TESLA_P4
    • NVIDIA_TESLA_P100
    • NVIDIA_TESLA_T4
    • NVIDIA_TESLA_V100
    • NVIDIA_L4

    • TPU_V2
    • TPU_V3

  • ACCELERATOR_LIMIT: o limite do acelerador (GPU ou TPU) para o operador. É possível especificar um número inteiro positivo. Para mais informações sobre as GPUs disponíveis e como configurá-las, consulte GPUs. Para mais informações sobre as TPUs disponíveis e como configurá-las, consulte TPUs.

No momento, o CustomJob oferece suporte a tipos de máquina específicos, que limitam você a um máximo de 96 CPUs e 624 GB de memória. Com base na configuração de CPU, memória e acelerador especificada, o Vertex AI Pipelines seleciona automaticamente a correspondência mais próxima dos tipos de máquina compatíveis.