排程 GPU 工作負載並組合任務

Important

這項功能目前處於公開預覽版。使用時，工作區管理員必須從預覽頁面啟用 AI 執行時預覽。請參閱管理 Azure Databricks 預覽。

使用宣告式自動化套件來排程 AI 執行時 GPU 工作負載，並將其與其他工作結合，例如上游資料準備。在 YAML 中定義工作負載及其排程，然後將它們部署為作業。對於訓練管線，先對 CPU 運算執行預處理，資料準備好後再開始 GPU 訓練。

本指南從列印「Hello world」的工作負載開始，接著建立排程的預處理與訓練流程。如果你已經有 AI 執行時工作負載的 YAML，請參考「將 AI 執行時工作負載轉換成套件」。

運作原理

一個 bundle 包含你的程式碼和一個 databricks.yml 設定。
工作會將任務分組並定義執行時間。
ai_runtime_task 在無伺服器的 GPU 運算上執行你的指令。

欄位 depends_on 將任務連結成有向無環圖（DAG），並且 schedule 欄位會依固定排程執行你的工作。

databricks bundle deploy 上傳程式碼並建立或更新作業。 databricks bundle run 立即開始執行任務。每次執行時，AI Runtime 會配置 GPU 運算資源、執行你的指令，並在命名的 MLflow 實驗中記錄執行過程。

Requirements

一個位於支援區域的工作區。請參閱需求說明。
最新的 Databricks CLI，已通過驗證並連線到您的工作區。在使用這些範例前，先更新現有的安裝。
在預處理範例中，使用現有的 Unity Catalog 磁碟區。該工作的執行身分需要該磁碟區的 USE CATALOG、USE SCHEMA、READ VOLUME 及 WRITE VOLUME 權限。請參閱 Unity 目錄磁碟區的許可權。

含有 `ai_runtime_task` 的 Hello World 範例

這個例子是在一顆 A10 GPU 上執行 shell 指令。

為該套件建立一個目錄。在該目錄中，建立command.sh，內容如下：
```
#!/usr/bin/env bash
set -euo pipefail
echo "Hello world"
```

在同一目錄中建立 databricks.yml ：

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

從 bundle 目錄中驗證、部署並執行該工作：

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

打開 CLI 列印的執行網址，查看 hello 任務的輸出。它包含 Hello world。這個執行也出現在 hello-ai-runtime MLflow 實驗中。

更複雜的例子：排程預處理與訓練流程

這個範例是在 CPU 運算上準備一個小型資料集，然後在一顆 A10 GPU 上訓練一個線性模型。這兩個任務都使用同一個卷檔案來傳遞資料。該工作排在每天 09:00 UTC，且在你測試之前，此排程會維持暫停。

建立專案

建立一個獨立目錄，結構如下：

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

請使用以下內容建立 prep.py。將 <schema>、<volume> 和 <catalog> 替換為您的磁碟區名稱。此來源筆記本將輸入值正規化，並將已處理的資料寫入磁碟區：

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

請使用以下內容建立 src/train.py。使用與 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}")

建立 command.sh 以執行訓練程式碼。以下的套件組包含 src/，因此 CODE_SOURCE_PATH 指向擷取出的 src 目錄：
```
#!/usr/bin/env bash
set -euo pipefail
cd "$CODE_SOURCE_PATH"
python train.py
```

請使用以下內容建立 databricks.yml。此套件組會將 src/ 打包供 GPU 任務使用。 prep Notebook 在無伺服器 CPU 運算上執行，且 depends_on 只有在 train 成功後才啟動 prep。 max_concurrent_runs: 1 防止此作業的執行個體互相覆寫輸入檔案：

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

測試並啟用排程

在 scheduled-training/ 中，部署套件並開始手動執行：

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

打開 CLI 列印的執行網址。確認train已成功，再開始prep。 prep 輸出報告 10 筆已準備的列，而 train 輸出報告 Trained on cuda:0 和損失。
在 databricks.yml中，將排程的 pause_status 從 PAUSED 變更為 UNPAUSED。現在的時間表區塊如下：
```
schedule:
  quartz_cron_expression: '0 0 9 * * ?'
  timezone_id: UTC
  pause_status: UNPAUSED
```
部署變更以啟動每日排程：
```
databricks bundle deploy --target dev
```
明確設定 pause_status: UNPAUSED 會啟用此排程，即使是開發目標亦然。在 工作與管線中，開啟已部署的工作並確認其排程於 09：00 UTC 生效。請參閱依排程執行作業。

要暫停排程運行，請設定 pause_status: PAUSED，然後再次部署。要移除範例工作及由套件上傳的檔案，請從其套件目錄執行以下操作：

databricks bundle destroy --target dev

你的磁碟區中準備好的資料仍然存在。當你不再需要 scheduled-training 資料目錄時，就刪除它。

其他資源

將現有的 AI 執行時工作負載 YAML 轉換成一個套件。請參見「將 AI 執行時工作負載轉換為套件」。
配置程式碼封裝、運算、環境及任務參數。請參見「設定 AI 執行環境套件任務」。
在 GPU 上排程現有的 Notebook。請參考「使用 Jobs API 和宣告式自動化套件進行排程」。
追蹤訓練跑步參見實驗追蹤與可觀測性。
模型、優化器與資料管線狀態的檢查點。請參閱提升 AI 執行時的訓練效能與韌性。

意見反應

此頁面對您有幫助嗎？

Last updated on 2026-10-05

排程 GPU 工作負載並組合任務

運作原理

Requirements

含有 ai_runtime_task 的 Hello World 範例

更複雜的例子：排程預處理與訓練流程

建立專案

測試並啟用排程

其他資源

意見反應

其他資源

含有 `ai_runtime_task` 的 Hello World 範例