Especifique a configuração da máquina para um passo do pipeline

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

Por exemplo, um componente de preparação pode usar um ficheiro CSV como entrada e usá-lo para preparar um modelo. Ao definir os parâmetros do tipo de máquina no passo do pipeline, pode gerir os requisitos de cada passo no pipeline. Se tiver dois passos de preparação e um passo for preparado num ficheiro de dados enorme e o segundo passo for preparado num ficheiro de dados pequeno, pode atribuir mais memória e CPU à primeira tarefa e menos recursos à segunda tarefa.

Por predefinição, o componente é executado como um Vertex AI CustomJob usando uma máquina e2-standard-4, com 4 CPUs de núcleo e 16 GB de memória. Para mais informações sobre a seleção de um dos recursos de máquinas específicos listados em Tipos de máquinas, consulte o artigo Peça Google Cloud recursos de máquinas com o Vertex AI Pipelines. Google Cloud

O exemplo seguinte mostra como definir as definições de configuração da CPU, da memória e da GPU para um passo:

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_type(ACCELERATOR_TYPE)
    .set_accelerator_limit(ACCELERATOR_LIMIT))

Substitua o seguinte:

  • CPU_LIMIT: o limite máximo da CPU para este operador. Este valor de cadeia pode ser um número (valor inteiro para o número de CPUs) ou um número seguido de "m", que significa 1/1000. Pode especificar, no máximo, 96 CPUs.

  • MEMORY_LIMIT: o limite máximo de memória para este operador. Este valor de string pode ser um número ou um número seguido de "K" (quilobyte), "M" (megabyte) ou "G" (gigabyte). É suportado um máximo de 624 GB.

  • SELECTOR_CONSTRAINT: cada restrição é uma etiqueta de par de chave-valor. Para que o contentor seja elegível para ser executado num nó, o nó tem de ter cada restrição como uma etiqueta. Por exemplo: 'cloud.google.com/gke-accelerator', 'NVIDIA_TESLA_T4'

    Estão disponíveis as seguintes restrições:

    • NVIDIA_GB200+ (inclui GPUDirect-RDMA)
    • NVIDIA_B200* (inclui GPUDirect-RDMA)
    • NVIDIA_H100_MEGA_80GB* (inclui GPUDirect-TCPXO)
    • NVIDIA_H100_80GB
    • NVIDIA_H200_141GB* (inclui GPUDirect-RDMA)
    • NVIDIA_A100_80GB
    • NVIDIA_TESLA_A100 (NVIDIA A100 40GB)
    • NVIDIA_TESLA_P4
    • NVIDIA_TESLA_P100
    • NVIDIA_TESLA_T4
    • NVIDIA_TESLA_V100
    • NVIDIA_L4

    • TPU_V2
    • TPU_V3

  • ACCELERATOR_TYPE: o tipo de acelerador. Para mais informações acerca das GPUs disponíveis e como as configurar, consulte GPUs. Para mais informações acerca dos tipos de TPUs disponíveis e como os configurar, consulte TPUs. Para mais informações, consulte a referência de parâmetros accelerator_type.

  • ACCELERATOR_LIMIT: o limite do acelerador (GPU ou TPU) para o operador. Pode especificar um número inteiro positivo. Para mais informações sobre as GPUs disponíveis e como as configurar, consulte o artigo GPUs. Para mais informações acerca das TPUs disponíveis e como as configurar, consulte TPUs.

O CustomJob suporta tipos de máquinas específicos que limitam a um máximo de 96 CPUs e 624 GB de memória. Com base na configuração da CPU, da memória e do acelerador que especificar, o Vertex AI Pipelines seleciona automaticamente a correspondência mais próxima dos tipos de máquinas suportados.