[Ray] RLlib 커넥터와 배치 유틸리티에 ndarray 빠른 경로 추가

PR #61320 - [RLlib] Optimize connector and batching utilities for ndarray fast paths

들어가며

Ray RLlib에서 학습 루프의 핫 패스(hot path)에서는 커넥터 배칭과 룩백 버퍼 접근이 빈번하게 발생합니다. 기존 코드는 모든 데이터 구조에 대해 tree.map_structure를 사용하여 재귀적 트리 순회를 수행했는데, 대부분의 경우 데이터가 단순한 numpy 배열임에도 불필요한 오버헤드가 발생했습니다. 이 PR은 numpy 배열인 경우를 감지하여 트리 순회를 건너뛰는 빠른 경로를 추가합니다.

핵심 코드 분석

Before (ConnectorV2.add_n_batch_items)

def _tag(s):
    return BatchedNdArray(s)

item_to_add = tree.map_structure(_tag, items_to_add)

After

if isinstance(items_to_add, np.ndarray):
    # Fast path: skip tree traversal for simple numpy arrays.
    tagged = items_to_add.view(BatchedNdArray)
else:
    def _tag(s):
        return BatchedNdArray(s)
    tagged = tree.map_structure(_tag, items_to_add)

Before (InfiniteLookbackBuffer.append)

self.data = tree.map_structure(
    lambda d, i: np.concatenate([d, [i]], axis=0), self.data, item
)

After

if isinstance(self.data, np.ndarray):
    self.data = np.concatenate(
        [self.data, np.asarray(item)[np.newaxis]], axis=0
    )
else:
    self.data = tree.map_structure(
        lambda d, i: np.concatenate([d, [i]], axis=0), self.data, item
    )

Before (batch 함수)

flat = tree.flatten(list_of_structs[0])
individual_items_already_have_batch_dim = isinstance(flat[0], BatchedNdArray)

After

if isinstance(first, BatchedNdArray):
    individual_items_already_have_batch_dim = True
elif is_nested:
    flat = tree.flatten(first)
    individual_items_already_have_batch_dim = isinstance(flat[0], BatchedNdArray)
else:
    individual_items_already_have_batch_dim = False

왜 이게 좋은가

1. 불필요한 트리 순회 제거: 단순 numpy 배열에 대해 tree.map_structure 호출을 건너뛰어 함수 호출 오버헤드를 제거합니다.
2. ndarray.view 활용: BatchedNdArray로의 변환을 복사 없이 view로 처리하여 메모리 할당을 줄입니다.
3. 학습 루프 핫패스 최적화: 초당 수천 번 실행되는 경로에서의 개선이므로 누적 효과가 큽니다.
4. 하위 호환성 유지: 중첩 구조(dict, tuple)에 대해서는 기존 로직이 그대로 동작합니다.

참고 자료

• Ray RLlib 공식 문서
• numpy ndarray.view