Apache Beam으로 머신러닝 데이터 파이프라인 구축하기 2편 - 개발 및 최적화

Apache Beam으로 머신러닝 데이터 파이프라인을 개발하면서 맞닥뜨린 유지보수성과 성능 문제를 다룬다. Composite PTransform으로 로직을 모듈화하고, key·Fusion·batching을 조절해 Dataflow 병목을 푼다. 모델과 Faiss index처럼 큰 객체를 worker에서 공유하는 방법과 protobuf 2GB 제한을 우회하는 설계까지 이어진다.

핵심 포인트

DoFn을 무분별하게 늘리면 schema 변화 추적, 코드 재사용, 단위 테스트가 어려워지므로 여러 변환을 Composite PTransform으로 묶는다.
key space가 너무 좁으면 parallel degree가 낮아지고, high-fanout 변환이 Fusion되면 출력이 한 worker에 몰릴 수 있다.
`Reshuffle`로 불리한 Fusion을 끊고, 데이터 축소 변환을 앞에 배치하며, 외부 요청은 bundle 단위로 batch 처리한다.
Beam `Shared` 객체로 thread가 하나의 모델을 공유하고 `no_use_multiple_sdk_containers`로 worker당 process 수를 하나로 제한할 수 있다.
`BatchElements`는 실행 중 적절한 batch size를 찾아 GPU utilization을 높이며, TFX `RunInference`는 모델 공유와 batching을 묶어 제공한다.
Faiss용 embedding을 한 element로 합치면 protobuf 2GB 한계를 넘으므로 PCollection을 side input으로 전달해 index를 만든다.

상세 정리

Beam 추상화: PCollection은 분산 데이터셋을 메모리 컬렉션처럼 다루게 하고, PTransform은 입력 PCollection의 각 element를 처리해 새 PCollection을 만든다. ParDo의 사용자 로직은 DoFn `process()`에 구현한다.
유지보수 문제: 개별 DoFn이 난립하면 각 단계가 schema를 어떻게 바꾸는지 알기 어렵고 같은 처리가 여러 pipeline에 복제된다. 전체 흐름을 작은 단위로 잘라 테스트하기도 어려워진다.
Composite Transform: 여러 ParDo와 조건 분기를 하나의 PTransform `expand()` 안에 캡슐화했다. 입력별 중복을 없애고 condition만으로 출력 schema가 결정되게 해 가독성과 테스트 가능성을 높였다.
분산 key: GroupByKey뿐 아니라 ParDo도 key에 따라 bundle 분배 정도가 달라진다. key 종류가 적으면 random string을 붙이는 방식으로 key space를 늘릴 수 있다.
Fusion 역효과: Dataflow는 직렬화 비용을 줄이려고 변환을 합친다. 그러나 high-fanout A와 후속 B가 합쳐지면 A의 수백·수천 출력이 여러 worker로 퍼지지 못하고 한 worker에 머문다.
Fusion 차단: aggregation 전후에는 Fusion이 일어나지 않는 특성을 이용해 `Reshuffle`을 삽입한다. 데이터 자체는 바꾸지 않고 실행 graph의 분산 경계를 만든다.
일반 최적화: 데이터를 줄이는 변환은 앞쪽에 둔다. `process()`의 과도한 로그는 Cloud Logging 부하를 만들고, DB·서버 호출은 `start_bundle`과 `finish_bundle`에서 batch 처리한다.
모델 메모리: worker의 한 process가 수백 thread를 만들 수 있어 각 DoFn `setup()`에서 모델을 로드하면 GPU OOM이 난다. Beam 2.24 이후 `Shared.acquire()`로 thread가 모델 하나를 함께 쓰게 했다.
Process 한계: Shared는 thread 사이에서만 동작해 process마다 모델이 하나씩 생긴다. CPU 수를 조정하거나 `experiments=no_use_multiple_sdk_containers`로 worker당 process를 하나로 강제한다.
GPU 활용: element를 한 개씩 추론하면 작은 입력에서 GPU가 놀게 된다. `BatchElements(min_batch_size, max_batch_size)`는 실행 시 batch 크기를 탐색하며, SavedModel이면 TFX `RunInference`로 같은 최적화를 재사용할 수 있다.
Faiss 병목: 모든 embedding을 `ToList()`로 하나의 element에 모으면 worker 간 gRPC 직렬화에 쓰는 protobuf의 element당 2GB 제한을 넘는다.
Side input 해법: embedding PCollection을 `AsList` side input으로 넘겨 준비가 끝날 때까지 기다린 뒤 메모리에서 index를 만든다. Shared로 Faiss index도 한 번만 생성해 query transform들이 재사용한다.