[ray] Ray Object Manager의 Pull RPC 배치 처리 최적화 분석

PR 링크: ray-project/ray#64225 상태: Merged | 변경: +305 / -133

들어가며

분산 시스템에서 객체 전송은 성능의 핵심입니다. 기존 Ray의 PullManager는 객체를 가져올 때마다 개별적인 PullRequest RPC를 생성했습니다. 이는 객체 수가 많아질 경우 네트워크 오버헤드와 CPU 컨텍스트 스위칭 비용을 증가시키는 원인이 되었습니다. 본 PR은 동일한 노드로 향하는 여러 객체 요청을 하나의 RPC로 통합(Coalescing)하여 성능을 최적화했습니다.

코드 분석

1. src/ray/object_manager/object_manager.cc

기존에는 ObjectID 하나를 인자로 받던 SendPullRequest가 std::vector<ObjectID>를 받도록 변경되었습니다. Protobuf 정의 또한 repeated bytes object_ids를 사용하도록 업데이트되었습니다.

// Before
void ObjectManager::SendPullRequest(const ObjectID &object_id, const NodeID &client_id) {
  // ...
  pull_request.set_object_id(object_id.Binary());
  // ...
}

// After
void ObjectManager::SendPullRequest(const std::vector<ObjectID> &object_ids, const NodeID &client_id) {
  // ...
  for (const auto &oid : object_ids) {
    pull_request.add_object_ids(oid.Binary());
  }
  // ...
}

2. src/ray/object_manager/pull_manager.cc

TryToMakeObjectLocal이라는 단일 객체 처리 함수를 삭제하고, 배치 처리를 수행하는 TryToMakeObjectsLocal 함수로 통합했습니다. 이 함수는 absl::flat_hash_map을 사용하여 대상 노드별로 객체 ID들을 그룹화합니다.

void PullManager::TryToMakeObjectsLocal(const std::vector<ObjectID> &object_ids) {
  absl::flat_hash_map<NodeID, std::vector<ObjectID>> batch_by_node;
  // ... (로직 생략: 노드별로 그룹화)
  for (const auto &[node_id, oids] : batch_by_node) {
    send_pull_request_(oids, node_id);
  }
}

왜 이게 좋은가

이 최적화는 특히 작은 크기의 객체(4KiB 수준)를 대량으로 전송할 때 큰 효과를 발휘합니다. 벤치마크 결과에 따르면, 4096개의 객체를 전송할 때 Pusher 노드의 CPU 사용량이 약 35% 감소했으며, 전체적인 Pull 시간 또한 개선되었습니다.

핵심 교훈:

1. RPC Batching: 분산 시스템에서 작은 메시지를 빈번하게 보내는 것은 네트워크 스택과 CPU에 큰 부담을 줍니다. 연관된 요청을 배치로 묶는 것은 가장 효과적인 성능 개선 전략 중 하나입니다.
2. Protobuf 설계: reserved 키워드를 사용하여 기존 필드와의 충돌을 방지하고 하위 호환성을 고려하는 설계 방식은 프로덕션 환경에서 필수적입니다.
3. 테스트 전략: 단순히 기능을 구현하는 것에 그치지 않고, num_send_pull_request_rpcs_와 같은 카운터를 도입하여 배치 로직이 의도대로 동작하는지 검증하는 테스트 코드를 작성하는 것이 중요합니다.

리뷰어 피드백 반영

리뷰 과정에서 PullManager의 함수 문서를 Doxygen 스타일(/** **/)로 통일하고, 에러 로그에서 전체 객체 ID를 출력하도록 개선하여 디버깅 편의성을 높였습니다. 또한, 배치 처리와 단일 처리 로직을 하나로 통합하여 코드 유지보수성을 향상시켰습니다.

참고 자료

• https://grpc.io/docs/what-is-grpc/introduction/
• https://developers.google.com/protocol-buffers/docs/proto3

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