[Ray] 외부 소비자의 Object Store 사용량을 Resource Manager 예산에 반영

PR 링크: ray-project/ray#62117 상태: Merged | 변경: +310 / -36

들어가며

Ray Data의 ResourceManager는 각 오퍼레이터의 Object Store 메모리 사용량을 추정하여 예산을 할당합니다. 그런데 Iterator와 streaming_split의 프리페치가 sink 오퍼레이터의 출력 큐를 떠난 후에도 Object Store에 블록을 보유합니다. 이 메모리가 어떤 오퍼레이터의 사용량에도 반영되지 않아, allocator가 이를 "사용 가능"으로 잘못 판단할 수 있었습니다.

핵심 코드 분석

Before: 외부 소비자 바이트가 오퍼레이터 사용량에 미반영

def _estimate_object_store_memory_usage(self, op):
    if isinstance(op, InputDataBuffer):
        return 0  # 입력 버퍼는 0으로 처리

    # 외부 소비자 바이트가 어디에도 귀속되지 않음
    self._mem_op_outputs[op] = op_outputs_bytes + used_op_outputs_bytes
    return self._mem_op_outputs[op] + self._mem_op_internal[op]

After: executor sink에 외부 소비자 바이트 귀속

def _estimate_object_store_memory_usage(self, op):
    if isinstance(op, InputDataBuffer):
        if op is self._output_operator:
            # InputDataBuffer-only 파이프라인에서도 처리
            self._mem_op_internal[op] = 0
            self._mem_op_outputs[op] = self._external_consumer_bytes
            return self._external_consumer_bytes
        return 0

    # ... 기존 로직 ...
    self._mem_op_outputs[op] = op_outputs_bytes + used_op_outputs_bytes

    # executor sink에만 프리페치 바이트 추가
    if op is self._output_operator:
        self._mem_op_outputs[op] += self._external_consumer_bytes

    return self._mem_op_outputs[op] + self._mem_op_internal[op]

토폴로지에서 terminal operator를 자동으로 찾습니다:

def _terminal_operator_from_topology(self, topology):
    """output_dependencies가 비어있는 유일한 오퍼레이터를 반환."""
    sinks = [op for op in topology if not op.output_dependencies]
    if len(sinks) == 1:
        return sinks[0]
    raise ValueError(...)

왜 이게 좋은가

• 정확한 메모리 예산: 프리페치된 블록이 "자유 메모리"로 잘못 계산되어 과도한 할당이 발생하던 문제를 해결합니다.
• 단일 귀속 지점: 외부 소비자 바이트는 항상 executor sink 하나에만 귀속되므로 이중 계산이 불가능합니다.
• InputDataBuffer 처리: 프리페치만 있는 단순 파이프라인에서도 올바르게 동작합니다.
• 방어적 검증: set_external_consumer_bytes에 음수 방지 assert, 토폴로지에서 sink가 0개 또는 2개 이상이면 ValueError를 발생시킵니다.

참고 자료

• Ray Data 리소스 관리
• Object Store 메모리 관리