AI 런타임에서 데이터 로드

중요합니다

데이터 및 모델 자산은 대형 언어 모델(LLM) 및 비전 언어 모델(VLM)의 딥러닝 및 학습 후 작업에서 매우 중요합니다. AI 런타임을 사용하면 모든 데이터와 모델 자산이 Unity 카탈로그를 통해 접근할 수 있습니다:

Unity 카탈로그 볼륨: 주로 대규모 데이터셋과 비정형 파일(이미지, 오디오, 텍스트 포함)에 사용됩니다.
Unity 카탈로그 테이블: 구조화 및 표 형식 데이터에 사용되며, Spark Connect를 통해 접근합니다.

볼륨과 테이블은 Unity Catalog에 등록되어 있어야 하며, 사용자 또는 서비스 주체가 접근할 수 있어야 합니다.

비정형 데이터를 위한 Unity 카탈로그 볼륨

Unity 카탈로그 볼륨은 구조화, 반구조화, 비구조화 데이터를 포함한 모든 형식의 비표 데이터에 대한 거버넌스 접근을 제공합니다. AI 런타임에서 볼륨은 대규모 데이터셋, 텍스트, 모델 자산, 모델 체크포인트에 접근하는 주요 메커니즘입니다.

사용자는 로컬 디스크 내 파일을 다루는 것과 유사한 익숙한 파일 시스템 작업을 통해 Unity 카탈로그 볼륨에 파일을 나열, 읽기, 쓸 수 있습니다:

import os

dir_path = "/Volumes/<catalog-name>/<schema-name>/<volume-name>/sub-dir"
file_path = os.path.join(dir_path, "test_file")

os.makedirs(dir_path, exist_ok=True)

# Write to the file
with open(file_path, "w") as file:
    file.write("Hello, World!")

마찬가지로, 셸 연산도 동일하게 작동합니다:

%sh ls -l /Volumes/<catalog-name>/<schema-name>/<volume-name>
%sh mkdir -p /Volumes/<catalog-name>/<schema-name>/<volume-name>/sub-dir
%sh touch /Volumes/<catalog-name>/<schema-name>/<volume-name>/sub-dir/test_file

Unity 카탈로그 볼륨의 몇 가지 특징은 머신러닝 작업에 적합하게 만듭니다:

분산 스토리지: Unity 카탈로그는 분산 스토리지를 기반으로 하여, AI 런타임 워크로드가 노트북과 CLI 기반 워크로드 모두에서 데이터를 읽고 쓰고 자산을 모델링할 수 있습니다.
ML 접근 패턴에 최적화됨: 기본 저장 공간과 접근 경로는 특히 순차적인 읽기 및 쓰기가 가능한 대형 파일에서 일반적인 ML 워크로드에 최적화되어 있습니다. 이로 인해 Unity Catalog는 학습 데이터 로딩, 모델 자산 로딩, 모델 체크포인트 작성에 적합합니다.
파일 시스템과 유사한 접근: 사용자는 로컬 디스크의 파일을 다루는 것과 유사한 익숙한 파일 시스템 조작을 사용하여 Unity 카탈로그 볼륨 내 파일을 목록, 읽기, 쓸 수 있습니다.

자동 백그라운드 커밋 덕분에 사용자는 Unity 카탈로그 볼륨 데이터에 일관되게 접근할 수 있습니다:

쓰기: AI 런타임은 쓰기를 자동으로 커밋하므로 동일한 Unity Catalog 볼륨에 액세스하는 다른 애플리케이션과 워크로드에서도 변경 사항을 볼 수 있습니다.
읽기: AI 런타임은 사용자가 명시적으로 새로고침이나 동기화 작업을 하지 않아도 볼륨 변경 사항을 자동으로 감지합니다.

튠 볼륨 성능

앞서 언급했듯이, Unity 카탈로그 볼륨은 분산 저장소로 지원되며, 순차적인 읽기와 쓰기가 가능한 대용량 파일에 최적화되어 있습니다.

AI 런타임에서 최상의 성능을 얻는 데 도움이 될 몇 가지 팁이 있습니다:

데이터를 더 큰 파일로 연결하기: 가능하다면 데이터를 더 적고 큰 파일로 통합하세요. 대략 1 GiB에서 10 GiB 사이입니다. 이로 인해 AI 런타임은 데이터를 적극적으로 미리 가져와 거의 최적의 순차 읽기 성능을 자동으로 달성할 수 있습니다.

작은 파일 작업의 경우, 로컬 디스크를 사용하세요: 많은 작은 파일이 포함된 작업물은 파일을 처리하기 전에 병렬 복사를 사용해 로컬 디스크로 복사하는 것을 고려하세요. 이렇게 하면 볼륨을 통해 여러 개의 작은 파일에 반복적으로 접근하는 데 드는 오버헤드를 줄일 수 있습니다.

# Recommended using parallel copy (256 concurrency in this example, you can tune)
#
# This takes only 22 seconds to copy 15,375 150KiB small image files.
%sh cd /Volumes/<catalog-name>/<schema-name>/<volume-name>/sub-dir/ && find . -type f -print0 | xargs -0 -P 256 -I {} cp --parents "{}" /tmp/


# !!! Avoid doing this !!!
#
# Because the files are copied in serial, this copies the same 15,375 150KiB small image files much more slowly.
# %sh cp -r /Volumes/<catalog-name>/<schema-name>/<volume-name>/sub-dir/* /tmp

머신러닝 작업에도 사용할 UCVolumeDataset 수 있습니다. 위에서 설명한 최적화를 포함하여 Unity 카탈로그 볼륨에서 효율적인 데이터 접근 및 로딩을 제공합니다. 자세한 내용은 다음 섹션을 참조하세요.

비정형 데이터를 다음과 같이 불러오세요 `UCVolumeDataset`

Unity 카탈로그 볼륨에 저장된 이미지, 오디오, 텍스트 파일 등 비정형 데이터는 UCVolumeDataset 모듈의 databricks.air.data를 사용하세요. UCVolumeDataset 는 볼륨에서 첫 번째 액세스 시 빠른 로컬 캐시로 각 파일을 복사하고 캐시된 로컬 파일 경로를 생성하는 PyTorch IterableDataset 입니다. 그렇지 않으면 직접 처리해야 할 성능 및 배포 관련 문제를 처리합니다.

로컬 캐싱. 파일은 FUSE 탑재에서 첫 번째 액세스 시 로컬 캐시 디렉터리로 복사되고 나중에 캐시에서 제공되므로 다중 Epoch 학습은 볼륨을 다시 읽지 않습니다.
자동 분할. torch.distributed가 초기화되면 파일들은 rank 간에 분할되고, 이어서 DataLoader 작업자들에 걸쳐 추가로 분할되므로 각 (rank, worker) 쌍은 추가 설정 없이 서로 겹치지 않는 슬라이스를 할당받습니다.

메모

UCVolumeDataset 및 databricks.air.data.DataLoader은 databricks-sdk-air 패키지에서 제공됩니다. 추가 기능 data와 함께 설치하세요. 그러면 호환되는 torch도 함께 설치합니다:

%pip install "databricks-sdk-air[data]"

UCVolumeDataset 는 원시 로컬 파일 경로를 생성합니다. 이러한 파일을 텐서로 디코딩하려면 경로 스트림을 받아 파싱 로직을 적용하는 두 번째 IterableDataset로 감싸세요. 이렇게 하면 I/O 및 구문 분석 문제가 별도로 유지됩니다.

from databricks.air.data import UCVolumeDataset
from torch.utils.data import IterableDataset
from PIL import Image
import torchvision.transforms.functional as TF

class ImageDataset(IterableDataset):
    """Decodes each cached file path from UCVolumeDataset into a tensor."""

    def __init__(self, path_dataset: UCVolumeDataset):
        self._path_dataset = path_dataset

    def __iter__(self):
        for local_path in self._path_dataset:
            image = Image.open(local_path).convert("RGB")
            yield TF.to_tensor(image)

path_dataset = UCVolumeDataset("/Volumes/catalog/schema/my_volume/images")
dataset = ImageDataset(path_dataset)

래퍼는 이미 캐시된 로컬 경로를 받기 때문에 파싱 단계가 볼륨을 건드리지 않습니다. 확대, 토큰화 또는 필터링을 위해 추가 래퍼를 연결할 수 있습니다.

최적의 성능을 위해 기본 PyTorch UCVolumeDataset 대신 databricks.air.data.DataLoader와 DataLoader를 함께 사용하세요. AI 런타임 입출력 처리에 최적화되어 있으며, GPU가 계산하는 동안 파일을 동시에 가져오기와 캐시를 수행합니다.

볼륨 위의 체크포인트 모델

최신 스냅샷에서 훈련을 재개하거나 장애 발생 후 복구할 수 있도록 모델의 체크포인트를 저장하려면 Unity 카탈로그 볼륨을 로컬 파일 시스템처럼 사용할 수 있습니다.

Databricks는 단일 GPU 및 다중 GPU 작업 부하 모두에서 더 나은 성능을 위해 분산 체크포인트(DCP) 사용을 권장합니다. Databricks 엔지니어링 블로그의 AI Runtime에서의 빠르고 장애 허용 가능한 PyTorch 학습을 참조하세요.

import torch.distributed.checkpoint as dcp
from torch.distributed.checkpoint.state_dict import get_state_dict, set_state_dict

import databricks.air.data

checkpoint_path = "/Volumes/my-catalog/my-schema/my-volume/checkpoints/step_1000"

# Save
model_sd, optim_sd = get_state_dict(model, optimizer)
state_dict = {"model": model_sd, "optim": optim_sd, "step": 1000}
dcp.async_save(
    state_dict,
    storage_writer=databricks.air.data.UCVolumeWriter(checkpoint_path))

# Load
model_sd, optim_sd = get_state_dict(model, optimizer)
state_dict = {"model": model_sd, "optim": optim_sd}
dcp.load(
    state_dict,
    storage_reader=databricks.air.data.UCVolumeReader(checkpoint_path))

set_state_dict(
    model,
    optimizer,
    model_state_dict=state_dict["model"],
    optim_state_dict=state_dict["optim"],
)

단일형 torch.save 방식도 작동합니다.

단일 GPU 모델 체크포인트의 경우,

# The monolithic torch.save approach for single GPU chip

# Save
torch.save({"model": model.state_dict(), "opt": optimizer.state_dict()},
           "/Volumes/<catalog-name>/<schema-name>/<volume-name>/sub-dir/ckpt.pt")


# Load
ckpt = torch.load(
    "/Volumes/<catalog-name>/<schema-name>/<volume-name>/sub-dir/ckpt.pt",
    weights_only=True)
model.load_state_dict(ckpt["model"])
optimizer.load_state_dict(ckpt["opt"])

토치런을 통해 시작된 분산 교육을 위해,

# The monolithic torch.save approach for multi-GPU distributed training.
# This snippet assumes your launcher has already called
# dist.init_process_group(...).

import os
import torch.distributed as dist

# Save only on rank 0.
if dist.get_rank() == 0:
    torch.save({"model": model.state_dict(), "opt": optimizer.state_dict()},
               "/Volumes/<catalog-name>/<schema-name>/<volume-name>/sub-dir/ckpt.pt")

# Wait for rank 0 to finish writing before any rank reads.
dist.barrier()

# Load on ALL ranks (map to current rank's local GPU).
local_rank = int(os.environ["LOCAL_RANK"])
ckpt = torch.load(
    "/Volumes/<catalog-name>/<schema-name>/<volume-name>/sub-dir/ckpt.pt",
    map_location=f"cuda:{local_rank}",
    weights_only=True)
model.load_state_dict(ckpt["model"])
optimizer.load_state_dict(ckpt["opt"])

표 형식 데이터 로드

Spark Connect를 사용하여 델타 테이블에서 테이블 형식 기계 학습 데이터를 로드합니다.

단일 노드 학습의 경우 PySpark 메서드를 사용하여 Apache Spark DataFrames를 pandas DataFrames로 변환한 다음 필요에 따라 toPandas()to_numpy()를 사용하여 NumPy 형식으로 변환할 수 있습니다.

메모

Spark Connect는 분석 및 이름 확인을 실행 시간으로 연기하여 코드의 동작을 변경할 수 있습니다. Spark Connect와 Spark 클래식 비교를 참조하세요.

Spark Connect는 Spark SQL, Spark의 Pandas API, 구조적 스트리밍 및 MLlib(데이터 프레임 기반)를 비롯한 대부분의 PySpark API를 지원합니다. 지원되는 최신 API는 PySpark API 참조 설명서를 참조하세요.

다른 제한 사항은 서버리스 컴퓨팅 제한을 참조하세요.

Unity 카탈로그 볼륨을 사용하여 큰 델타 테이블을 불러오기

변환할 수 없는 큰 델타 테이블의 경우 데이터를 Unity 카탈로그 볼륨으로 toPandas()내보내고 PyTorch 또는 Hugging Face를 사용하여 직접 로드합니다.

# Step 1: Export the Delta table to Parquet files in a UC volume
output_path = "/Volumes/catalog/schema/my_volume/training_data"
spark.table("catalog.schema.my_table").write.mode("overwrite").parquet(output_path)

# Step 2: Load the exported data directly using Hugging Face datasets
from datasets import load_dataset

dataset = load_dataset("parquet", data_files="/Volumes/catalog/schema/my_volume/training_data/*.parquet")

이 방법은 학습 중에 Spark 오버헤드를 방지하고 단일 GPU 및 분산 학습 워크플로 모두에서 잘 작동합니다.

피드백

이 페이지가 도움이 되었나요?

2026-10-13 에 마지막으로 업데이트됨