Schemalägg GPU-arbetsbelastningar och skapa uppgifter

Important

Den här funktionen finns som allmänt tillgänglig förhandsversion. För att använda den måste en workspace-administratör aktivera AI Runtime-förhandsvisningen från sidan Förhandsgranskningar. Se Hantera förhandsversioner av Azure Databricks.

Använd Declarative Automation Bundles för att schemalägga en AI-runtime-GPU-arbetsbelastning och kombinera den med annat arbete, såsom upstream-dataförberedelse. Definiera arbetslasten och dess schema i YAML, och distribuera dem sedan som ett jobb. För en träningspipeline, kör förbehandling på CPU-resurser och starta GPU-träning efter att datan är klar.

Denna guide börjar med en arbetsbelastning som skriver ut "Hello world" och bygger sedan en schemalagd förbehandlings- och träningspipeline. Om du redan har en AI Runtime-arbetsbelastnings-YAML, se Konvertera en AI Runtime-arbetsbelastning till ett paket.

Så här fungerar det

  • Ett paket innehåller din kod och en databricks.yml konfiguration.
  • Ett jobb grupperar uppgifter och definierar när de körs.
  • En ai_runtime_task kör ditt kommando på serverlösa GPU-beräkningsresurser.

Fältet depends_on kopplar samman uppgifter i en riktad acyklisk graf (DAG), och schedule fältet schemalägger ditt jobb att köras i en rytm.

databricks bundle deploy laddar upp koden och skapar eller uppdaterar jobbet. databricks bundle run startar en jobbkörning omedelbart. Vid varje körning tillhandahåller AI Runtime GPU-beräkningen, kör ditt kommando och registrerar körningen i det namngivna MLflow-experimentet.

Requirements

Hello world-exempel med ai_runtime_task

Det här exemplet kör ett shell-kommando på ett A10-grafikkort.

  1. Skapa en katalog för paketet. I den katalogen, skapa command.sh med följande innehåll:

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

  2. Skapa databricks.yml i samma katalog:

    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. Från bundle-katalogen, validera, distribuera och kör jobbet:

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

  4. Öppna den kör-URL:en som skrivs ut av CLI:n och se hello uppgiftens utdata. Den innehåller Hello world. Körningen förekommer också i hello-ai-runtime MLflow-experimentet.

Mer komplext exempel: Schemalägg en förbearbetnings- och träningspipeline

Detta exempel förbereder en liten datamängd på CPU och tränar sedan en linjär modell på en A10-GPU. Båda uppgifterna använder samma volymfil för att skicka data mellan sig. Jobbet är schemalagt till 09:00 UTC varje dag, med schemat pausat tills du testar det.

Skapa projektet

  1. Skapa en separat katalog med följande layout:

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

  2. Skapa prep.py med följande innehåll. Ersätt <catalog>, <schema>, och <volume> med volymernas namn. Denna källnotebook normaliserar indatavärdena och skriver den förberedda datan till volymen:

    # 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. Skapa src/train.py med följande innehåll. Använd samma volymväg som i 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. Skapa command.sh för att köra träningskoden. Paketet nedan paketerar src/, så CODE_SOURCE_PATH pekar på den extraherade src katalogen:

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

  5. Skapa databricks.yml med följande innehåll. Paketet paketerar src/ för GPU-uppgiften. Notebooken prep körs på serverlös CPU-beräkningskapacitet och depends_on startar train först efter att prep har lyckats. max_concurrent_runs: 1 förhindrar att körningar av detta jobb skriver över varandras indatafiler:

    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

Testa och aktivera schemat

  1. Från scheduled-training/, distribuera paketet och starta en manuell körning:

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

  2. Öppna URL:en för körningen som skrivs ut av CLI. Bekräfta att prep lyckas innan train startar. Utdata prep rapporterar 10 förberedda rader, och train utdata rapporterar Trained on cuda:0 och förlusten.

  3. I databricks.yml, ändra schemats pause_status från PAUSED till UNPAUSED. Schemablocket är nu:

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

  4. Distribuera ändringen för att aktivera det dagliga schemat:

    databricks bundle deploy --target dev

    Att uttryckligen ange pause_status: UNPAUSED aktiverar detta schema även för ett utvecklingsmål. I Jobs & Pipelines, öppna det driftsatta jobbet och bekräfta att dess schema är aktivt för 09:00 UTC. Se Kör jobb enligt ett schema.

För att pausa schemalagda körningar, ställ in pause_status: PAUSED och deploya igen. För att ta bort exempeljobbet och filerna som paketet laddar upp, kör följande från dess bundle-katalogen:

databricks bundle destroy --target dev

De förberedda uppgifterna i din volym finns kvar. Ta bort datakatalogen scheduled-training när du inte längre behöver den.

Ytterligare resurser

  • Uppdaterades senast