본문으로 건너뛰기

[Ray RLlib] 커넥터 최적화: 벌크 데이터 추출과 리스트 연산 개선

PR 링크: ray-project/ray#60924 상태: Merged | 변경: +21 / -18

들어가며

Ray RLlib의 커넥터는 환경에서 수집한 에피소드 데이터를 학습 배치로 변환하는 역할을 합니다. 기존 구현에서는 두 가지 비효율적인 패턴이 있었습니다. BatchIndividualItems에서 매 타임스텝마다 리스트에 append하는 루프, 그리고 AddColumnsFromEpisodesToTrainBatch에서 타임스텝별로 개별 get_actions 호출을 수행하는 패턴입니다. 이 PR은 두 패턴을 모두 최적화합니다.

핵심 코드 분석

BatchIndividualItems: append 루프를 extend로

Before:

for (eps_id,) in column_data.keys():
    for item in column_data[(eps_id,)]:
        if column == Columns.OBS:
            memorized_map_structure.append(eps_id)
        list_to_be_batched.append(item)

After:

for (eps_id,) in column_data.keys():
    items = column_data[(eps_id,)]
    list_to_be_batched.extend(items)
    if column == Columns.OBS:
        num_samples = sum(
            len(item) if isinstance(item, BatchedNdArray) else 1
            for item in items
        )
        memorized_map_structure.extend([eps_id] * num_samples)

AddColumnsFromEpisodesToTrainBatch: 벌크 추출

Before:

self.add_n_batch_items(
    batch,
    Columns.ACTIONS,
    items_to_add=[
        sa_episode.get_actions(indices=ts)
        for ts in range(len(sa_episode))
    ],
    num_items=len(sa_episode),
    single_agent_episode=sa_episode,
)

After:

self.add_n_batch_items(
    batch,
    Columns.ACTIONS,
    items_to_add=sa_episode.get_actions(slice(0, len(sa_episode))),
    num_items=len(sa_episode),
    single_agent_episode=sa_episode,
)

동일한 패턴이 get_rewards와 get_extra_model_outputs에도 적용되었습니다.

왜 이게 좋은가

  • list.extend vs append 루프: extend는 C 레벨에서 동작하므로 Python 루프 오버헤드가 제거됩니다.
  • 벌크 데이터 추출: 1000 타임스텝 에피소드에서 get_actions를 1000번 호출하는 대신 한 번의 슬라이스 호출로 대체합니다.
  • NumPy 연산 최적화: 벌크 슬라이스는 내부적으로 NumPy의 벡터 연산을 활용할 수 있어, 개별 인덱싱보다 효율적입니다.
  • 실제 사용자 관찰 기반: 이 최적화는 실제 고객의 프로파일링 결과를 바탕으로 적용되었습니다.

참고 자료

댓글

관련 포스트

PR Analysis 의 다른글