GPU 작업 부하를 예약하고 작업을 구성

선언적 자동화 번들을 사용하여 AI 런타임 GPU 작업 부하를 예약하고, 이를 상위 데이터 준비와 같은 다른 작업과 결합하세요. YAML로 워크로드와 일정을 정의한 후, Job으로 배포하세요. 학습 파이프라인의 경우, CPU 컴퓨트에서 전처리를 실행하고 데이터가 준비된 후 GPU 훈련을 시작하세요.

이 가이드는 "Hello world"를 출력하는 작업으로 시작해, 예약된 전처리 및 훈련 파이프라인을 구축합니다. 이미 AI 런타임 워크로드 YAML이 있다면, ' AI 런타임 워크로드를 번들로 변환하기'를 참고하세요.

작동 방식

  • 번들에는 코드와 databricks.yml 구성이 포함되어 있습니다.
  • 작업은 작업을 그룹화하고 실행 시점을 정의합니다.
  • ai_runtime_task가 AI 런타임에서 당신의 명령을 실행합니다.

depends_on 필드는 작업을 방향성 비순환 그래프(DAG)로 연결하고, schedule 필드는 작업이 일정 간격으로 실행되도록 스케줄링합니다.

databricks bundle deploy 코드를 업로드하고 작업을 생성하거나 업데이트합니다. databricks bundle run 즉시 작업 실행을 시작합니다. 각 실행마다 AI 런타임은 GPU 컴퓨트를 프로비저닝하고, 명령어를 실행하며, 명명된 MLflow 실험에 실행을 기록합니다.

Requirements

ai_runtime_task가 포함된 안녕하세요 세계 예시

이 예시는 한 개의 A10 GPU에서 셸 명령을 실행합니다.

  1. 번들을 위한 디렉터리를 만드세요. 그 디렉터리에서 다음 내용으로 command.sh를 생성하세요:

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

  2. 같은 디렉터리에서 생성 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

  3. 번들 디렉터리에서 작업을 검증하고 배포한 후 실행하세요:

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

  4. CLI가 출력한 실행 URL을 열고 hello 작업의 출력을 확인하세요. Hello world를 포함합니다. 이 실행은 hello-ai-runtime MLflow 실험에도 나타납니다.

더 복잡한 예시: 전처리 및 훈련 파이프라인 예약

이 예시는 CPU 계산으로 작은 데이터셋을 준비한 후, 하나의 A10 GPU에서 선형 모델을 학습시킵니다. 두 작업 모두 동일한 볼륨 파일을 사용해 데이터를 주고받습니다. 작업은 매일 09:00 UTC에 예정되어 있으며, 테스트할 때까지 일정이 일시 중지된 상태입니다.

프로젝트를 만듭니다.

  1. 다음과 같은 레이아웃으로 별도의 디렉터리를 만드세요:

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

  2. 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}")

  3. 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}")

  4. 훈련 코드를 실행하기 위해 command.sh을 생성하세요. 아래 번들은 src/를 패키징하므로, CODE_SOURCE_PATH 즉 추출된 src 디렉터리를 가리킵니다:

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

  5. databricks.yml를 다음 내용으로 작성하세요. 번들에는 GPU 작업을 위해 src/가 포함됩니다. prep 노트북은 서버리스 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

스케줄을 테스트 및 활성화하세요

  1. scheduled-training/에서 번들을 배포하고 수동 실행을 시작합니다:

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

  2. CLI에 출력된 실행 URL을 여세요. train가 prep 시작하기 전에 성공하는지 확인하세요. prep 출력은 준비된 행 10개를 보고하고, train 출력은 Trained on cuda:0 및 손실을 보고합니다.

  3. databricks.yml에서 스케줄의 PAUSED을(를) UNPAUSED에서 pause_status(으)로 변경합니다. 현재 일정 블록은 다음과 같습니다:

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

  4. 변경 사항을 배포하여 일일 일정을 활성화하세요:

    databricks bundle deploy --target dev

    pause_status: UNPAUSED를 명시적으로 설정하면 개발 목표에도 이 일정을 활성화합니다. 작업 및 파이프라인에서 배포된 작업을 열고 09:00 UTC에 일정이 활성화되어 있는지 확인하세요. 일정에 따라 작업 실행을 참조하세요.

예약된 실행을 일시정지하려면 pause_status: PAUSED을 설정하고 다시 배포하세요. 예제 작업과 번들에서 업로드한 파일을 제거하려면 번들 디렉터리에서 다음 명령을 실행하세요:

databricks bundle destroy --target dev

볼륨에 준비된 데이터는 그대로 남아 있습니다. scheduled-training 데이터 디렉터리가 더 이상 필요하지 않을 때 삭제하세요.

추가 리소스

  • 에 마지막으로 업데이트됨