Dask

Dask는 기존 Python 및 PyData 생태계를 확장하는 병렬 및 분산 컴퓨팅 라이브러리입니다.

특히 Dask DataFrame를 사용하면 pandas 워크플로를 확장할 수 있습니다. Dask DataFrame은 pandas를 병렬화하여 대규모 표 형식 데이터를 처리합니다. pandas API를 밀접하게 반영하므로, 단일 데이터셋에서 테스트하던 것을 전체 데이터셋 처리로 전환하는 것이 간단합니다. Dask는 Hugging Face Datasets의 기본 형식인 Parquet와 특히 효과적으로 작동하며, 풍부한 데이터 유형, 효율적인 컬럼 필터링, 압축을 지원합니다.

Dask의 좋은 실용적 사용 사례는 데이터셋에 대해 데이터 처리 또는 모델 추론을 분산 방식으로 실행하는 것입니다. 예를 들어, Coiled의 Hugging Face + Dask를 활용한 AI 기반 데이터 처리 확장에 대한 훌륭한 블로그 게시물을 참고하세요.

읽기 및 쓰기

Dask은 원격 데이터를 읽고 쓰기 위해 fsspec를 사용하므로, Hugging Face 경로(hf://)를 사용하여 Hub에서 데이터를 읽고 쓸 수 있습니다.

먼저 Hugging Face 계정으로 로그인해야 합니다. 예를 들어 다음을 사용할 수 있습니다:

hf auth login

그 후 데이터셋 저장소를 생성할 수 있습니다. 예를 들어 다음을 사용할 수 있습니다:

from huggingface_hub import HfApi

HfApi().create_repo(repo_id="username/my_dataset", repo_type="dataset")

마지막으로, Dask에서 Hugging Face 경로를 사용할 수 있습니다. Dask DataFrame은 Hugging Face에 Parquet을 분산으로 기록하는 것을 지원하며, 이는 커밋을 사용하여 데이터셋 변경 사항을 추적합니다:

import dask.dataframe as dd

df.to_parquet("hf://datasets/username/my_dataset")

# or write in separate directories if the dataset has train/validation/test splits
df_train.to_parquet("hf://datasets/username/my_dataset/train")
df_valid.to_parquet("hf://datasets/username/my_dataset/validation")
df_test .to_parquet("hf://datasets/username/my_dataset/test")

이 방식은 파일별로 커밋을 하나씩 생성하므로, 업로드 후 히스토리를 스쿼시하는 것이 권장됩니다:

from huggingface_hub import HfApi

HfApi().super_squash_history(repo_id=repo_id, repo_type="dataset")

이 작업은 Parquet 형식의 Dask 데이터셋을 포함하는 데이터셋 저장소 username/my_dataset를 생성합니다. 나중에 다시 로드할 수 있습니다:

import dask.dataframe as dd

df = dd.read_parquet("hf://datasets/username/my_dataset")

# or read from separate directories if the dataset has train/validation/test splits
df_train = dd.read_parquet("hf://datasets/username/my_dataset/train")
df_valid = dd.read_parquet("hf://datasets/username/my_dataset/validation")
df_test  = dd.read_parquet("hf://datasets/username/my_dataset/test")

Hugging Face 경로 및 구현 방식에 대한 자세한 내용은 HfFileSystem에 대한 클라이언트 라이브러리 문서를 참고하세요.

데이터 처리

Dask를 사용하여 데이터셋을 병렬로 처리하려면, 먼저 pandas DataFrame 또는 Series에 대한 데이터 처리 함수를 정의한 다음, Dask의 map_partitions 함수를 사용하여 이 함수를 데이터셋의 모든 파티션에 병렬로 적용할 수 있습니다:

def dummy_count_words(texts):
    return pd.Series([len(text.split(" ")) for text in texts])

또는 pandas 문자열 메서드를 사용하는 유사한 함수(더 빠름):

def dummy_count_words(texts):
    return texts.str.count(" ")

pandas에서는 이 함수를 텍스트 열에 사용할 수 있습니다:

# pandas API
df["num_words"] = dummy_count_words(df.text)

그리고 Dask에서는 이 함수를 모든 파티션에 실행할 수 있습니다:

# Dask API: run the function on every partition
df["num_words"] = df.text.map_partitions(dummy_count_words, meta=int)

함수의 출력인 pandas Series 또는 DataFrame의 유형인 meta도 제공해야 한다는 점에 유의하세요. 이는 Dask DataFrame이 지연 평가(lazy) API를 사용하기 때문에 필요합니다. Dask는 .compute()가 호출될 때만 데이터 처리를 실행하므로, 그 사이에 새로운 열의 유형을 파악하기 위해 meta 인수가 필요합니다.

Predicate and Projection Pushdown

Hugging Face에서 Parquet 데이터를 읽을 때, Dask는 Parquet 파일의 메타데이터를 자동으로 활용하여 필요하지 않은 파일이나 row group 전체를 건너뜁니다. 예를 들어 Parquet 형식의 Hugging Face Dataset에 필터(predicate)를 적용하거나 컬럼의 부분 집합(projection)을 선택하는 경우, Dask는 Parquet 파일의 메타데이터를 읽어 필요하지 않은 부분을 다운로드하지 않고 폐기합니다.

Dask DataFrame API의 재구현을 통해 쿼리 최적화를 지원할 수 있으며, 이를 통해 Dask가 더 빠르고 견고해졌습니다.

예를 들어 FineWeb-Edu의 이 하위 집합에는 많은 Parquet 파일이 포함되어 있습니다. 최근 CC 덤프의 텍스트만 유지하도록 데이터셋을 필터링할 수 있다면, Dask는 대부분의 파일을 건너뛰고 필터에 일치하는 데이터만 다운로드합니다:

import dask.dataframe as dd

df = dd.read_parquet("hf://datasets/HuggingFaceFW/fineweb-edu/sample/10BT/*.parquet")

# Dask will skip the files or row groups that don't
# match the query without downloading them.
df = df[df.dump >= "CC-MAIN-2023"]

Dask는 계산에 필요한 열만 읽고 나머지는 건너뜁니다. 예를 들어, 코드 후반부에 특정 열을 제거하더라도 파이프라인 초기에 해당 열이 필요하지 않다면 로드하지 않습니다. 이는 열의 하위 집합을 조작하거나 분석을 수행할 때 유용합니다:

# Dask will download the 'dump' and 'token_count' needed
# for the filtering and computation and skip the other columns.
df.token_count.mean().compute()

클라이언트

dask의 대부분의 기능은 클러스터 또는 로컬 Client를 통해 병렬 계산을 시작하도록 최적화되어 있습니다:

import dask.dataframe as dd
from distributed import Client

if __name__ == "__main__":  # needed for creating new processes
    client = Client()
    df = dd.read_parquet(...)
    ...

로컬 사용의 경우, Client는 기본적으로 멀티프로세싱을 사용하는 Dask LocalCluster를 사용합니다. LocalCluster의 멀티프로세싱을 다음과 같이 수동으로 구성할 수 있습니다.

from dask.distributed import Client, LocalCluster
cluster = LocalCluster(n_workers=8, threads_per_worker=8)
client = Client(cluster)

Client 없이 로컬에서 기본 스레드 스케줄러를 사용하는 경우, 특정 작업 이후에 DataFrame이 더 느려질 수 있습니다(자세한 내용은 여기를 참조).

로컬 또는 클라우드 클러스터를 설정하는 방법에 대한 자세한 내용은 Dask 배포 문서에서 확인하세요.