Plánujte zatížení GPU a sestavujte úlohy

Použijte Declarative Automation Bundles k plánování AI Runtime GPU workloadu a kombinujte ho s dalšími pracemi, například přípravou dat v upstream fázi. Definujte pracovní zátěž a její harmonogram v YAML a pak je nasaďte jako úlohu. V rámci tréninkové pipeline spusťte předzpracování na CPU a trénování na GPU spusťte až po připravení dat.

Tento průvodce začíná úlohou, která vytiskne „Hello world“, poté vytvoří plánovanou pipeline pro předzpracování a trénování. Pokud už máte AI Runtime workload YAML, podívejte se na Převod workloadu AI Runtime na bundle.

Jak to funguje

  • Balíček obsahuje váš kód a databricks.yml konfiguraci.
  • Úloha seskupuje úkoly a definuje, kdy se spouštějí.
  • ai_runtime_task spustí váš příkaz v AI Runtime.

Pole depends_on propojuje úkoly do orientovaného acyklického grafu (DAG) a pole schedule plánuje pravidelné spouštění vaší úlohy.

databricks bundle deploy Nahraje kód a vytvoří nebo aktualizuje úlohu. databricks bundle run okamžitě spustí běh úlohy. Při každém spuštění AI Runtime zajistí GPU výpočetní prostředky, spustí váš příkaz a zaznamená běh v pojmenovaném experimentu MLflow.

Requirements

Příklad „Ahoj světe“ s ai_runtime_task

Tento příklad spouští shell příkaz na jednom GPU A10.

  1. Vytvořte adresář pro tento balíček. V tomto adresáři vytvořte command.sh s následujícím obsahem:

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

  2. Vytvořte databricks.yml ve stejném adresáři:

    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. Z adresáře balíčku ověřte, nasaďte a spusťte job:

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

  4. Otevřete adresu URL běhu vytištěnou rozhraním CLI a podívejte se na výstup hellodaného úkolu. Obsahuje Hello world. Běh se také objevuje v experimentu hello-ai-runtime MLflow.

Složitější příklad: Naplánujte pipeline pro předzpracování a trénování

Tento příklad připravuje malou datovou sadu pro výpočty na CPU a poté trénuje lineární model na jednom GPU A10. Obě úlohy používají stejný soubor volume pro přenos dat mezi sebou. Úkol je naplánován denně v 09:00 UTC, přičemž rozvrh je pozastaven, dokud ho neotestujete.

Vytvoření projektu

  1. Vytvořte samostatný adresář s následujícím rozložením:

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

  2. Vytvořte prep.py s následujícím obsahem. Nahraďte <catalog>, <schema> a <volume> názvy vašich svazků. Tento zdrojový poznámkový blok normalizuje vstupní hodnoty a zapisuje připravená data do úložiště:

    # 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. Vytvořte src/train.py s následujícím obsahem. Použijte stejnou cestu ke svazku jako v 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. Vytvořte command.sh, abyste spustili tréninkový kód. Níže uvedený balíček zahrnuje src/, takže CODE_SOURCE_PATH ukazuje na extrahovaný adresář src:

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

  5. Vytvořte databricks.yml s následujícím obsahem. Balíček obsahuje src/ pro úlohu GPU. Notebook prep běží na serverless CPU výpočtech a depends_on spouští se train teprve poté, co prep úspěšně proběhne. max_concurrent_runs: 1 zabraňuje, aby si běhy této úlohy navzájem přepisovaly vstupní soubor:

    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

Otestujte a povolte plán

  1. Z scheduled-training/ nasaďte bundle a spusťte ruční spuštění:

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

  2. Otevřete URL pro spuštění vypsanou rozhraním CLI. Ověřte, že prep proběhne úspěšně dříve než train začne. Výstup prep uvádí 10 připravených řádků a výstup train uvádí Trained on cuda:0 a ztrátu.

  3. V databricks.yml změňte u plánu položku pause_status z PAUSED na UNPAUSED. Blok rozpisu je nyní:

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

  4. Nasaďte změnu pro aktivaci denního rozvrhu:

    databricks bundle deploy --target dev

    Nastavení pause_status: UNPAUSED explicitně umožňuje tento plán i pro vývojový cíl. V Jobs & Pipelines otevřete nasazenou úlohu a potvrďte, že její plán je aktivní v 09:00 UTC. Viz Spuštění úloh podle plánu.

Chcete-li pozastavit plánovaná spuštění, nastavte pause_status: PAUSED a znovu nasaďte. Pro odstranění ukázkové úlohy a souborů nahraných balíčkem spusťte následující z jeho adresáře balíčku:

databricks bundle destroy --target dev

Připravená data ve vašem svazku zůstávají zachována. Smažte scheduled-training datový adresář, až ho už nebudete potřebovat.

Dodatečné zdroje

  • Poslední aktualizace v