Planifier des charges de travail GPU et composer des tâches

Utilisez les Declarative Automation Bundles pour planifier une charge de travail GPU AI Runtime et la combiner avec d’autres travaux, comme la préparation des données en amont. Définissez la charge de travail et son calendrier en YAML, puis déployez-les en tant que Job. Pour un pipeline d’entraînement, exécutez le prétraitement sur le CPU compute et start l’entraînement GPU une fois les données prêtes.

Ce guide commence avec une charge de travail qui imprime « Hello world », puis crée un pipeline de prétraitement et d’entraînement planifié. Si vous avez déjà un YAML de charge de travail AI Runtime, consultez Convertir une charge de travail AI Runtime en bundle.

Comment ça marche

Un bundle contient votre code et une configuration databricks.yml.
Un job regroupe des tâches et définit le moment de leur exécution.
Un ai_runtime_task exécute votre commande sur le Runtime IA.

Le champ depends_on connecte les tâches dans un graphe orienté acyclique (DAG), et le champ schedule programme votre job pour qu’il s’exécute selon un rythme défini.

databricks bundle deploy upload le code et crée ou met à jour le Job. databricks bundle run start immédiatement une exécution de Job. À chaque exécution, AI Runtime effectue le provisionnement du compute GPU, exécute votre commande et enregistre l'exécution dans l'Experimentation MLflow nommée.

Configuration requise

Un workspace dans une région prise en charge. Voir Exigences.
La dernière version de la CLI Databricks, authentifiée auprès de votre workspace. Mettez à jour une installation existante avant d'utiliser ces exemples.
Pour l'exemple de prétraitement, un volume Unity Catalog existant. L'identité d'exécution du job nécessite les privilèges USE CATALOG, USE SCHEMA, READ VOLUME et WRITE VOLUME pour ce volume. Voir Privilèges pour les volumes Unity Catalog.

Exemple Bonjour tout le monde avec `ai_runtime_task`

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

Créer un répertoire pour le bundle. Dans ce répertoire, créez command.sh avec les contenus suivants :
Bash
```
#!/usr/bin/env bash
set -euo pipefail
echo "Hello world"
```

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

YAML
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

Depuis le répertoire du bundle, validez, déployez et lancez le job :

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

Ouvrez l'URL d'exécution affichée par le CLI et affichez la sortie de la tâche hello. Il contient Hello world. L'exécution apparaît également dans l’Experimentation MLflow hello-ai-runtime.

Exemple plus complexe : Planifier un pipeline de prétraitement et d’entraînement

Cet exemple prépare un petit dataset sur un compute 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 des données entre elles. Le job est programmé quotidiennement à 09:00 UTC, avec le planning en pause jusqu'à ce que vous le testiez.

Créer le projet

Créez un répertoire séparé avec le layout suivant :

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

Créez prep.py avec les contenus suivants. Remplacez <catalog>, <schema> et <volume> par les noms de vos volumes. Ce notebook source normalise les valeurs d’entrée et écrit les données préparées dans le volume :

Python
# 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}")

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

Python
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}")

Créer command.sh pour exécuter le code d’entraînement. Le bundle ci-dessous package src/, donc CODE_SOURCE_PATH pointe vers le répertoire src extrait :
Bash
```
#!/usr/bin/env bash
set -euo pipefail
cd "$CODE_SOURCE_PATH"
python train.py
```

Créez databricks.yml avec le contenu suivant. Le bundle package src/ pour la tâche GPU. Le Notebook prep s'exécute sur le compute CPU Serverless, et depends_on start train seulement après que prep a réussi. max_concurrent_runs: 1 empêche les exécutions de ce job de remplacer leurs fichiers d'entrée respectifs :

YAML
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 calendrier

Depuis scheduled-training/, déployez le bundle et start une exécution manuelle :

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

Ouvrez l’URL d’exécution imprimée par le CLI. Confirmez que prep est réussi avant que train start. Le résultat prep indique 10 lignes préparées et le résultat train indique Trained on cuda:0 et la perte.
Dans databricks.yml, modifiez pause_status du planning de PAUSED à UNPAUSED. Le bloc de planning est maintenant :
YAML
```
schedule:
  quartz_cron_expression: '0 0 9 * * ?'
  timezone_id: UTC
  pause_status: UNPAUSED
```
Déployez la modification pour activer le planning quotidien :
Bash
```
databricks bundle deploy --target dev
```
La définition de pause_status: UNPAUSED active explicitement ce calendrier même pour une cible de développement. Dans Jobs et pipelines , ouvrez le job déployé et confirmez que son planning est actif pour 09:00 UTC. Voir Exécuter des jobs selon un calendrier.

Pour suspendre des exécutions programmées, définissez pause_status: PAUSED et déployez à nouveau. Pour supprimer le Job d’exemple et les fichiers upload par le bundle, exécutez la commande suivante depuis son répertoire :

Bash
databricks bundle destroy --target dev

Les données préparées dans votre volume sont conservées. Supprimez le répertoire de données scheduled-training lorsque vous n'en avez plus besoin.

Ressources supplémentaires

Convertir un YAML de charge de travail AI Runtime existant en bundle. Voir Convertir une charge de travail AI Runtime en bundle.
Configurez l’empaquetage du code, le compute, les environnements et les paramètres des tâches. Consultez Configurer des tâches de bundle AI Runtime.
Planifiez un notebook existant sur GPU. Consultez Créer et gérer des jobs de notebook planifiés.
Suivez les exécutions d'entraînement. Voir suivi de l’Experimentation et observabilité.
Checkpoint pour l'état du modèle, l'optimiseur et le pipeline de données. Consultez Améliorer les performances et la résilience de l’entraînement sur l’IA Runtime.

Comment ça marche​

Configuration requise​

Exemple Bonjour tout le monde avec ai_runtime_task​

Exemple plus complexe : Planifier un pipeline de prétraitement et d’entraînement​

Créer le projet​

Tester et activer le calendrier​

Ressources supplémentaires​