Spécifier la configuration d'une machine pour une étape de pipeline

Les composants du pipeline Kubeflow sont des fonctions de fabrique qui créent des étapes de pipeline. Chaque composant décrit les entrées, les sorties et l'implémentation du composant. Par exemple, train_op est un composant de l'exemple de code suivant.

Par exemple, un composant d'entraînement peut utiliser un fichier CSV en entrée et l'utiliser pour entraîner un modèle. En définissant les paramètres de type de machine à l'étape du pipeline, vous pouvez gérer les exigences de chaque étape de votre pipeline. Si vous avez deux étapes d'entraînement, dont l'une sur un fichier de données volumineux et l'autre sur un petit fichier de données, vous pouvez allouer plus de mémoire et de processeurs à la première tâche, et moins de ressources à la deuxième.

Par défaut, le composant sera exécuté en tant qu'IA Vertex CustomJob à l'aide d'une machine e2-standard-4, avec des processeurs à 4 cœurs et 16 Go de mémoire. Pour en savoir plus sur la sélection de l'une des ressources machine spécifiques à Google Cloudrépertoriées dans la section Types de machines, consultez la page Demander des ressources machine Google Cloud avec Vertex AI Pipelines.

L'exemple suivant montre comment définir les paramètres de configuration du processeur, de la mémoire et du GPU pour une étape :

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

Remplacez les éléments suivants :

  • CPU_LIMIT : limite maximale du processeur pour cet opérateur. Cette valeur de chaîne peut être un nombre (valeur entière pour le nombre de processeurs) ou un nombre suivi de "m", qui signifie 1/1000. Vous pouvez spécifier au maximum 96 processeurs.

  • MEMORY_LIMIT : limite de mémoire maximale pour cet opérateur. Cette valeur de chaîne peut être un nombre, ou un nombre suivi de "K" (kilo-octet), "M" (mégaoctet) ou "G" (gigaoctet). La limite est fixée à 624 Go.

  • SELECTOR_CONSTRAINT: chaque contrainte est une étiquette de paire clé-valeur. Pour que le conteneur puisse être exécuté sur un nœud, chacune des contraintes doit y être définie sous forme d'étiquette. Exemple : 'cloud.google.com/gke-accelerator', 'NVIDIA_TESLA_T4'

    Les contraintes suivantes sont disponibles :

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

    • TPU_V2
    • TPU_V3

  • ACCELERATOR_LIMIT : limite d'accélérateur (GPU ou TPU) pour l'opérateur. Vous pouvez spécifier un entier positif. Pour en savoir plus sur les GPU disponibles et leur configuration, consultez la page GPU. Pour en savoir plus sur les TPU disponibles et leur configuration, consultez TPU.

CustomJob est compatible avec des types de machines spécifiques qui vous limitent à 96 processeurs et 624 Go de mémoire. En fonction de la configuration du processeur, de la mémoire et de l'accélérateur que vous spécifiez, Vertex AI Pipelines sélectionne automatiquement la correspondance la plus proche parmi les types de machines compatibles.