Planifiez les charges de travail GPU et composez les tâches

Utilisez Declarative Automation Bundles pour planifier une charge de travail GPU d’AI Runtime et la combiner avec d’autres tâches, comme la préparation des données en amont. Définissez la charge de travail et sa planification dans YAML, puis déployez-la en tant que travail. Pour un pipeline d’entraînement, lancez un prétraitement sur le calcul CPU et commencez l’entraînement GPU une fois les données prêtes.

Ce guide commence par une charge de travail qui imprime « Hello world », puis construit un pipeline de prétraitement et de formation planifié. Si vous avez déjà un fichier YAML de charge de travail AI Runtime, consultez Convertir une charge de travail AI Runtime en bundle.

Fonctionnement

  • Un bundle contient votre code et une databricks.yml configuration.
  • Un job regroupe les tâches et définit quand elles s’exécutent.
  • Un ai_runtime_task exécute votre commande sur AI Runtime.

Le champ depends_on relie les tâches dans un graphe acyclique orienté (DAG), et le champ schedule programme votre tâche pour s’exécuter selon une cadence définie.

databricks bundle deploy télécharge le code et crée ou met à jour le travail. databricks bundle run démarre immédiatement une exécution de tâche. À chaque exécution, AI Runtime provisionne les ressources de calcul GPU, exécute votre commande et enregistre l’exécution dans l’expérience MLflow nommée.

Requirements

Exemple de Hello World avec ai_runtime_task

Cet exemple exécute une commande shell sur un GPU A10.

  1. Créez un répertoire pour le bundle. Dans ce répertoire, créez command.sh avec le contenu suivant :

    #!/usr/bin/env bash
set -euo pipefail
echo "Hello world"

  2. Créez databricks.yml dans le même répertoire :

    bundle:
  name: hello-ai-runtime

resources:
  jobs:
    hello:
      name: hello-ai-runtime
      tasks:
        - task_key: hello
          environment_key: default
          ai_runtime_task:
            experiment: hello-ai-runtime
            deployments:
              - command_path: ${workspace.file_path}/command.sh
                compute:
                  accelerator_type: GPU_1xA10
                  accelerator_count: 1
      environments:
        - environment_key: default
          spec:
            environment_version: '6'

targets:
  dev:
    mode: development
    default: true

  3. Depuis le répertoire bundle, validez, déployez et exécutez la tâche :

    databricks bundle validate --target dev
databricks bundle deploy --target dev
databricks bundle run hello --target dev

  4. Ouvrez l’URL d’exécution imprimée par la CLI et consultez la sortie de la tâche hello. Il contient Hello world. Le run apparaît également dans l’expérience hello-ai-runtime MLflow.

Exemple plus complexe : planifier un pipeline de prétraitement et de formation

Cet exemple prépare un petit jeu de données sur le calcul CPU, puis entraîne un modèle linéaire sur un GPU A10. Les deux tâches utilisent le même fichier de volume pour transmettre les données entre elles. La tâche est programmée à 09h00 UTC tous les jours, avec le planning suspendu jusqu’à ce que vous le testiez.

Créer le projet

  1. Créez un répertoire séparé avec la disposition suivante :

    scheduled-training/
├── databricks.yml
├── prep.py
├── command.sh
└── src/
    └── train.py

  2. Créez un fichier prep.py comportant le contenu suivant. Remplacez <catalog>, <schema>, et <volume> par les noms de vos volumes. Ce carnet source normalise les valeurs d’entrée et écrit les données préparées dans le volume :

    # Databricks notebook source
import json
from pathlib import Path

data_path = Path("/Volumes/<catalog>/<schema>/<volume>/scheduled-training/data.json")
inputs = [value / 10 for value in range(10)]
data = {"x": inputs, "y": [2 * value + 1 for value in inputs]}
data_path.parent.mkdir(parents=True, exist_ok=True)
data_path.write_text(json.dumps(data))
print(f"Prepared {len(inputs)} rows at {data_path}")

  3. Créez src/train.py avec le contenu suivant. Utilisez le même chemin de volume que dans prep.py:

    import json
from pathlib import Path

import torch

data_path = Path("/Volumes/<catalog>/<schema>/<volume>/scheduled-training/data.json")
data = json.loads(data_path.read_text())
x = torch.tensor(data["x"], device="cuda").reshape(-1, 1)
y = torch.tensor(data["y"], device="cuda").reshape(-1, 1)
model = torch.nn.Linear(1, 1).to("cuda")
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)

for _ in range(200):
    optimizer.zero_grad()
    loss = torch.nn.functional.mse_loss(model(x), y)
    loss.backward()
    optimizer.step()

print(f"Trained on {x.device}: loss={loss.item():.4f}")

  4. Créez command.sh pour exécuter le code d’entraînement. Le bundle ci-dessous regroupe src/, donc CODE_SOURCE_PATH pointe vers le répertoire extrait src :

    #!/usr/bin/env bash
set -euo pipefail
cd "$CODE_SOURCE_PATH"
python train.py

  5. Créez databricks.yml avec le contenu suivant. Le lot regroupe src/ pour la tâche GPU. Le prep notebook fonctionne sur un CPU serverless, et depends_on ne démarre train qu’une fois que prep a réussi. max_concurrent_runs: 1 empêche les exécutions de cette tâche de remplacer le fichier d’entrée des autres :

    bundle:
  name: scheduled-training

artifacts:
  code:
    type: tgz
    path: .
    include: [src]
    files:
      - source: ./dist/code.tgz

resources:
  jobs:
    train_pipeline:
      name: scheduled-training
      max_concurrent_runs: 1
      schedule:
        quartz_cron_expression: '0 0 9 * * ?'
        timezone_id: UTC
        pause_status: PAUSED
      tasks:
        - task_key: prep
          notebook_task:
            notebook_path: ./prep.py
        - task_key: train
          depends_on:
            - task_key: prep
          environment_key: training
          ai_runtime_task:
            experiment: scheduled-training
            code_source_path: ./dist/code.tgz
            deployments:
              - command_path: ${workspace.file_path}/command.sh
                compute:
                  accelerator_type: GPU_1xA10
                  accelerator_count: 1
      environments:
        - environment_key: training
          spec:
            environment_version: '6'
            dependencies:
              - torch

targets:
  dev:
    mode: development
    default: true

Tester et activer le planning

  1. À partir de scheduled-training/, déployez le bundle et commencez une exécution manuelle :

    databricks bundle validate --target dev
databricks bundle deploy --target dev
databricks bundle run train_pipeline --target dev

  2. Ouvrez l’URL d’exécution imprimée par la CLI. Confirmez que prep réussit avant que train ne commence. La sortie prep signale 10 lignes préparées, et la sortie train signale Trained on cuda:0 et la perte.

  3. Dans databricks.yml, redéfinissez le pause_status du planning de PAUSED à UNPAUSED. Le bloc de planning est maintenant :

    schedule:
  quartz_cron_expression: '0 0 9 * * ?'
  timezone_id: UTC
  pause_status: UNPAUSED

  4. Déploie le changement pour activer le planning quotidien :

    databricks bundle deploy --target dev

    Le paramètre pause_status: UNPAUSED permet explicitement cette planification même pour un objectif de développement. Dans Jobs & Pipelines, ouvrez la tâche déployée et confirmez que son planning est actif à 09h00 UTC. Consultez Exécuter des travaux selon une planification.

Pour mettre en pause les exécutions programmées, définir pause_status: PAUSED et déployer à nouveau. Pour supprimer le job d’exemple et les fichiers téléversés par le bundle, exécutez ce qui suit depuis son répertoire bundle :

databricks bundle destroy --target dev

Les données préparées dans votre volume restent. Supprime le scheduled-training répertoire de données quand tu n’en auras plus besoin.

Ressources additionnelles

  • Dernière mise à jour :