Apache Beam으로 머신러닝 데이터 파이프라인 구축하기 3편 - RunInference로 모델 추론하기

Apache Beam과 GCP Dataflow에서 대규모 ML bulk inference를 실행할 때 직접 DoFn, Shared, BatchElements를 조합하던 구현을 공식 RunInference PTransform으로 단순화하는 방법을 설명한다. PyTorch 모델 로딩과 키 보존, 가변 길이 입력 패딩, GPU 할당과 추론 지표까지 실제 파이프라인 코드 수준으로 연결한다.

핵심 포인트

Beam은 데이터 저장소와 분산 실행을 추상화해 개발자가 추론 로직에 집중하면서 Dataflow가 데이터 양에 맞춰 worker를 확장하게 한다.
기존 구현은 worker마다 여러 process·thread가 모델을 중복 적재해 GPU OOM이 날 수 있어 Shared와 단일 SDK container 옵션이 필요했다.
Beam 2.40부터 RunInference가 모델 로딩, 입력 batching, 추론 실행과 결과 형식을 공식 API로 제공한다.
PyTorch의 단일 Tensor와 keyed Tensor 입력은 서로 다른 ModelHandler를 사용하며, 원본 키는 KeyedModelHandler로 유지한다.
언어 모델처럼 길이가 다른 Tensor는 기본 `torch.stack()`이 실패하므로 ModelHandler의 `run_inference()`를 확장해 padding해야 한다.

상세 정리

적용 상황: live inference 인프라와 다른 모델로 대량 데이터를 주기 처리할 때 매 파이프라인마다 별도 모델 서버를 만드는 비용을 줄이려 했다.
Beam 선택: DoFn으로 비즈니스 로직을 쓰고 Dataflow runner가 분산 처리와 worker 확장을 맡아 데이터 규모에 따른 인프라 관리를 숨긴다.
기존 OOM 원인: 한 worker에서 여러 process와 thread가 같은 DoFn을 실행하면 제한된 GPU 메모리에 모델이 반복 적재된다.
기존 해결: Beam 2.24의 Shared로 thread 간 모델 객체를 공유하고 `experiments=no_use_multiple_sdk_containers`로 worker당 process를 하나로 제한했다.
배치 처리: `BatchElements(min_batch_size=1024, max_batch_size=4096)`로 여러 element를 묶어 GPU utilization을 높였다.
GPU 할당: Dataflow resource hint에 `type:nvidia-tesla-t4;count:1;install-nvidia-driver`와 최소 RAM을 지정해 T4와 드라이버를 준비했다.
RunInference 구조: ModelHandler가 외부 저장소에서 모델을 읽고 batching·실제 추론을 담당하며, RunInference가 Beam PTransform으로 실행한다.
프레임워크 지원: PyTorch와 TensorFlow, Scikit-learn용 handler가 제공되고 지원되지 않는 프레임워크는 ModelHandler를 상속해 구현할 수 있다.
입력 유형: 단일 `torch.Tensor`는 PytorchModelHandlerTensor, `input_ids`와 `attention_mask` 같은 dict 입력은 KeyedTensor handler를 사용한다.
결과 형식: PredictionResult가 원본 example과 inference 결과를 함께 반환한다. element 식별자가 필요하면 handler를 KeyedModelHandler로 감싼다.
배치 조절: `batch_elements_kwargs()`를 재정의하면 RunInference 내부 BatchElements의 최소·최대 크기를 제어할 수 있다.
가변 길이 함정: 기본 handler는 Tensor 목록을 `torch.stack()`하므로 문장 길이가 다르면 실패한다. batch size를 1로 줄이면 GPU 효율이 크게 떨어진다.
커스텀 패딩: `run_inference()`를 재정의해 `pad_sequence`로 input_ids를 맞추고 `torch.inference_mode()`에서 추론한 뒤 PredictionResult로 변환했다.
향후 API: 당시 master에는 handler 생성 시 `inference_fn`을 주입하는 변경이 병합돼 추후 버전에서는 상속 없이 커스텀 추론 로직을 넣을 수 있었다.
관측성: Dataflow 패널에서 평균 batch size와 모델 추론 시간 등 RunInference가 기록한 지표를 바로 확인할 수 있다.