Especifica la configuración de la máquina para un paso de canalización

Los componentes de canalización de Kubeflow son funciones de fábrica que crean pasos de canalización. Cada componente describe las entradas, salidas y su implementación. Por ejemplo, train_op es un componente en la siguiente muestra de código.

Por ejemplo, un componente de entrenamiento puede tomar un archivo CSV como entrada y usarlo para entrenar un modelo. Si configuras los parámetros del tipo de máquina en el paso de canalización, puedes administrar los requisitos de cada paso en tu canalización. Si tienes dos pasos de entrenamiento y un paso realiza el entrenamiento en un archivo de datos enorme y el segundo lo hace en 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á como un CustomJob de Vertex AI mediante una máquina e2-standard-4 con 4 CPU 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 Cloud enumerados en Tipos de máquinas, consulta Solicita recursos de máquinas de Google Cloud con Vertex AI Pipelines.

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

Nota: Si deseas especificar el espacio en disco en la configuración de la máquina, debes crear un trabajo de entrenamiento personalizado a partir de un componente mediante la solicitud de recursos de máquina de Google Cloud en su lugar.

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

Reemplaza lo siguiente:

• CPU_LIMIT: El límite máximo de CPU para este operador. Este valor de string puede ser un número (valor de número entero para la cantidad de CPU) o un número seguido de “m”, que significa 1/1000. Puedes especificar como máximo 96 CPU.

• MEMORY_LIMIT: El límite máximo de memoria para este operador. Este valor de string puede ser un número o un número seguido de “K” (kilobyte), “M” (megabyte) o “G” (gigabyte). Se admiten como máximo 624 GB.

  Nota: Vertex AI Pipelines no admite llamadas a set_memory_request en un operador. Debes usar set_memory_limit para solicitar una cantidad de memoria específica.

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

  Las siguientes restricciones están disponibles:

  • NVIDIA_H100_80GB
  • 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: Es el límite del acelerador (GPU o TPU) para el operador. Puedes especificar un número entero positivo. Para obtener más información sobre las GPU disponibles y cómo configurarlas, consulta GPU. Para obtener más información sobre las TPU disponibles y cómo configurarlas, consulta TPU.

  Nota: Las TPU solo están disponibles en ubicaciones específicas. Para obtener más información, consulta Usa aceleradores.

CustomJob admite tipos de máquinas específicos que te limitan a un máximo de 96 CPU y 624 GB de memoria. Según la configuración de CPU, memoria y acelerador que especifiques, Vertex AI Pipelines selecciona automáticamente la coincidencia más cercana de los tipos de máquinas compatibles.