Especificar la configuración de la máquina de un paso de flujo de procesamiento

Los componentes de los flujos de procesamiento de Kubeflow son funciones de fábrica que crean pasos de flujos de procesamiento. En cada componente se describen las entradas, las salidas y la implementación del componente. Por ejemplo, train_op es un componente del siguiente código de muestra.

Por ejemplo, un componente de entrenamiento podría tomar un archivo CSV como entrada y usarlo para entrenar un modelo. Si defines los parámetros del tipo de máquina en el paso del flujo de procesamiento, puedes gestionar los requisitos de cada paso del flujo. Si tienes dos pasos de entrenamiento y uno de ellos se entrena con un archivo de datos enorme, mientras que el otro se entrena con un archivo de datos pequeño, puedes asignar más memoria y CPU a la primera tarea y menos recursos a la segunda.

De forma predeterminada, el componente se ejecutará en Vertex AI CustomJob con una máquina e2-standard-4, que tiene 4 CPUs principales y 16 GB de memoria. Para obtener más información sobre cómo seleccionar uno de los recursos de máquina específicos de Google Cloudque se indican en Tipos de máquinas, consulta Solicitar recursos de máquina con Vertex AI Pipelines. Google Cloud

En el siguiente ejemplo se muestra cómo definir la configuración de CPU, memoria y GPU de un paso:

.
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))

Haz los cambios siguientes:

  • CPU_LIMIT: límite máximo de CPU de este operador. Este valor de cadena puede ser un número (valor entero del número de CPUs) o un número seguido de "m", que significa 1/1000. Puedes especificar un máximo de 96 CPUs.

  • MEMORY_LIMIT: límite máximo de memoria de este operador. Este valor de cadena puede ser un número o un número seguido de "K" (kilobyte), "M" (megabyte) o "G" (gigabyte). Se admite un máximo de 624 GB.

  • SELECTOR_CONSTRAINT: cada restricción es una etiqueta de par clave-valor. Para que el contenedor pueda ejecutarse en un nodo, este debe tener cada restricción como etiqueta. Por ejemplo: 'cloud.google.com/gke-accelerator', 'NVIDIA_TESLA_T4'

    Están disponibles las siguientes restricciones:

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

    • TPU_V2
    • TPU_V3

  • ACCELERATOR_LIMIT: límite de aceleradores (GPU o TPU) del operador. Puedes especificar un número entero positivo. Para obtener más información sobre las GPUs disponibles y cómo configurarlas, consulta GPUs. Para obtener más información sobre las TPUs disponibles y cómo configurarlas, consulta el artículo sobre las TPUs.

    .

CustomJob admite tipos de máquinas específicos que te limitan a un máximo de 96 CPUs y 624 GB de memoria. En función de la configuración de CPU, memoria y acelerador que especifiques, Vertex AI Pipelines seleccionará automáticamente el tipo de máquina compatible que más se ajuste a tus necesidades.