[Ray] Ray Data에 cuDF 배치 포맷 추가

PR 링크: ray-project/ray#61329 상태: Merged | 변경: +469 / -35

들어가며

Ray Data는 map_batches()에서 pandas, pyarrow, numpy 등의 batch format을 지원한다. GPU 기반 데이터 처리 워크로드에서는 Arrow 테이블을 받아 cuDF로 변환하는 추가 단계가 필요했다. 이 PR은 batch_format="cudf"를 공식 지원해 Ray Data가 직접 cudf.DataFrame으로 배치를 제공하도록 한다.

핵심 코드 분석

BatchFormat enum에 CUDF 추가

Before:

# python/ray/data/block.py
class BatchFormat(str, Enum):
    ARROW = "pyarrow"
    PANDAS = "pandas"
    NUMPY = "numpy"

VALID_BATCH_FORMATS = ["pandas", "pyarrow", "numpy", None]

After:

class BatchFormat(str, Enum):
    ARROW = "pyarrow"
    PANDAS = "pandas"
    NUMPY = "numpy"
    CUDF = "cudf"

VALID_BATCH_FORMATS = ["pandas", "pyarrow", "numpy", "cudf", None]

"cudf"가 공식 batch format으로 추가되었다. map_batches(), iter_batches(), take_batch() 등 모든 배치 API에서 사용할 수 있다.

Arrow에서 cuDF로의 변환 경로

# python/ray/data/_internal/table_block.py
def to_cudf(self) -> Any:
    """Convert this block to a cudf.DataFrame."""
    from ray.data.util.data_batch_conversion import _lazy_import_cudf

    cudf = _lazy_import_cudf()
    if cudf is None:
        raise ValueError(
            "Attempted to convert data to cuDF DataFrame but cuDF "
            "is not installed. Please do `pip install cudf-cu12`."
        )
    return cudf.DataFrame.from_arrow(self.to_arrow())

내부적으로 Arrow 블록을 cudf.DataFrame.from_arrow()로 변환한다. Arrow는 columnar format이므로 cuDF로의 변환이 효율적인 zero-copy에 가깝다.

cuDF DataFrame을 블록으로 변환 (역방향)

Before:

def _validate_batch_output(batch: Block) -> None:
    if not isinstance(batch, (list, np.ndarray, pa.Table,
                              pd.core.frame.DataFrame, dict)):
        raise ValueError(...)

After:

def _validate_batch_output(batch: Block) -> None:
    allowed = isinstance(batch, (list, np.ndarray, pa.Table,
                                  pd.core.frame.DataFrame, dict)
    ) or _is_cudf_dataframe(batch)
    if not allowed:
        raise ValueError(...)

map_batches()의 UDF가 cudf.DataFrame을 반환하는 것도 허용된다. 반환된 cuDF DataFrame은 batch.to_arrow()로 Arrow 블록으로 변환되어 저장된다.

cuDF의 Mapping 프로토콜 처리

# cudf.DataFrame이 Mapping을 구현하므로 dict 경로보다 먼저 처리
elif _is_cudf_dataframe(batch):
    return batch.to_arrow()  # GPU -> CPU bulk transfer

elif isinstance(batch, collections.abc.Mapping):
    # dict 경로 (느린 컬럼별 처리)
    ...

cudf.DataFrame은 Python의 Mapping 프로토콜을 구현하고 있어서, dict 분기보다 먼저 체크해야 한다. 그렇지 않으면 느린 컬럼별 처리 경로로 빠진다. to_arrow()를 통한 bulk GPU-to-CPU transfer가 훨씬 효율적이다.

왜 이게 좋은가

1. GPU 네이티브 파이프라인: Arrow -> cuDF 변환을 사용자가 직접 할 필요 없이 batch_format="cudf"만 지정하면 된다. GPU 기반 전처리 코드가 간결해진다.
2. 성능 최적화: cuDF의 from_arrow()와 to_arrow()를 통해 Arrow columnar format과의 효율적인 변환을 활용한다.
3. Mapping 프로토콜 함정 회피: cuDF DataFrame이 Mapping을 구현하는 문제를 올바른 분기 순서로 해결했다.
4. CI/Test 격리: GPU 테스트를 cudf 태그로 분리해 CPU-only 환경의 테스트에서 제외한다. 전용 GPU CI 빌드(docgpubuild-py312)를 추가했다.

정리

이 PR은 Ray Data의 배치 포맷 체계에 cuDF를 추가해 GPU 데이터 처리 파이프라인의 진입 장벽을 낮췄다. 핵심은 Arrow 기반 내부 저장 포맷과 cuDF 간의 효율적인 양방향 변환이다. cuDF의 Mapping 프로토콜 구현으로 인한 분기 순서 문제를 올바르게 처리한 점, GPU 테스트를 CI에서 격리한 점 등 실용적인 엔지니어링 판단이 돋보인다.

참고 자료

• Ray Data batch_format 문서 -- Ray Data map_batches API
• cuDF GitHub -- RAPIDS cuDF GPU DataFrame 라이브러리

알림: 이 분석은 AI가 실제 코드 diff를 기반으로 작성했습니다.