GPU-workloads plannen en taken samenstellen

Gebruik Declarative Automation Bundles om een AI Runtime GPU-workload te plannen en deze te combineren met ander werk, zoals upstream datavoorbereiding. Definieer de workload en het schema ervan in YAML en implementeer ze vervolgens als een taak. Voor een trainingspijplijn voer je preprocessing uit op CPU-compute en start je de GPU-training nadat de data klaar is.

Deze gids begint met een werklast die "Hello world" print, en bouwt vervolgens een geplande preprocessing- en trainingspijplijn. Als je al een AI Runtime workload YAML hebt, zie dan Een AI Runtime-workload converteren naar een bundle.

Hoe werkt het?

  • Een bundel bevat je code en een databricks.yml configuratie.
  • Een job groepeert taken en definieert wanneer ze draaien.
  • ai_runtime_task voert je commando uit op AI Runtime.

Het depends_on veld verbindt taken in een gerichte acyclische graaf (DAG), en het schedule veld plant je taak in om volgens een frequentie te worden uitgevoerd.

databricks bundle deploy uploadt de code en maakt of werkt de taak bij. databricks bundle run start onmiddellijk een taakuitvoering. Bij elke run stelt AI Runtime GPU-rekenkracht beschikbaar, voert je commando uit en registreert de run in het benoemde MLflow-experiment.

Requirements

Hallo wereld voorbeeld met ai_runtime_task

Dit voorbeeld draait een shell-commando op één A10 GPU.

  1. Maak een map aan voor de bundel. Maak in die map het bestand command.sh aan met de volgende inhoud:

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

  2. Maak databricks.yml aan in dezelfde map:

    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. Vanuit de bundeldirectory valideer, deploy en voer de taak uit:

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

  4. Open de run URL die door de CLI is afgedrukt en bekijk de output van de hello taak. Het bevat Hello world. De run komt ook voor in het hello-ai-runtime MLflow-experiment.

Complexer voorbeeld: Een preprocessing- en trainingspipeline plannen

Dit voorbeeld bereidt een kleine dataset voor met CPU-rekenkracht en traint vervolgens een lineair model op één A10 GPU. Beide taken gebruiken hetzelfde volumebestand om data tussen hen door te geven. De taak staat gepland op 09:00 UTC elke dag, met het schema gepauzeerd totdat je het test.

Het project maken

  1. Maak een aparte map aan met de volgende lay-out:

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

  2. Maak prep.py met de volgende inhoud. Vervang <catalog>, <schema>, en <volume> door de namen van je volumes. Dit bronnotitieboek normaliseert de invoerwaarden en schrijft de voorbereide data naar het 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. Maak src/train.py met de volgende inhoud. Gebruik hetzelfde volumepad als in 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. Maak command.sh om de trainingscode uit te voeren. De bundel hieronder verpakt src/, dus CODE_SOURCE_PATH wijst naar de geëxtraheerde src directory:

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

  5. Maak databricks.yml met de volgende inhoud. De bundel verpakt src/ voor de GPU-taak. De prep notebook draait op serverless CPU-rekenkracht en depends_on start train pas nadat prep is geslaagd. max_concurrent_runs: 1 voorkomt dat runs van deze taak elkaars invoerbestand overschrijven:

    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

Test en activeer het schema

  1. Vanaf scheduled-training/, implementeer je de bundel en start je een handmatige run:

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

  2. Open de run URL die door de CLI is afgedrukt. Bevestig dat prep lukt voordat train begint. De prep output rapporteert 10 voorbereide rijen, en de train output rapporteert Trained on cuda:0 en het verlies.

  3. In databricks.yml, verander de pause_status van het schema van PAUSED naar UNPAUSED. Het schema is nu:

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

  4. Implementeer de wijziging om het dagelijkse schema te activeren:

    databricks bundle deploy --target dev

    Door pause_status: UNPAUSED expliciet in te stellen, wordt dit schema expliciet ingeschakeld, zelfs voor een ontwikkelingsdoel. Open in Jobs & Pipelines de geïmplementeerde taak en bevestig dat het schema actief is voor 09:00 UTC. Zie Taken uitvoeren volgens een planning.

Om geplande runs te pauzeren, stel pause_status: PAUSED in en deploy opnieuw. Om de voorbeeldtaak en de bestanden die door de bundel zijn geüpload te verwijderen, voer je het volgende uit vanuit de bundelmap:

databricks bundle destroy --target dev

De voorbereide gegevens in je volume blijven behouden. Verwijder de scheduled-training datamap wanneer je die niet meer nodig hebt.

Aanvullende bronnen

  • Laatst bijgewerkt om