본문으로 건너뛰기

[Ray Data] 리소스 매니저 스케줄링에 논리적 메모리 포함

PR 링크: ray-project/ray#60774 상태: Merged | 변경: +100 / -61

들어가며

Ray Data의 리소스 매니저는 오퍼레이터들의 리소스 사용량을 추적하여 태스크 제출 여부를 결정합니다. 기존에는 current_processor_usage라는 이름으로 CPU와 GPU만 추적했으며, memory 리소스는 고려하지 않았습니다. 이로 인해 메모리를 많이 요구하는 태스크가 가용 메모리를 초과하여 할당될 수 있었습니다. 이 PR은 메모리를 스케줄링 결정에 포함시킵니다.

핵심 코드 분석

메서드명 변경: processor -> logical

Before:

def current_processor_usage(self) -> ExecutionResources:
    """CPU/GPU 사용량만 반환"""
    return ExecutionResources(0, 0, 0)

After:

def current_logical_usage(self) -> ExecutionResources:
    """CPU, GPU, 메모리 사용량 반환"""
    return ExecutionResources.zero()

TaskPoolMapOperator에 메모리 추가

Before:

def current_processor_usage(self) -> ExecutionResources:
    num_active_workers = self.num_active_tasks()
    return ExecutionResources(
        cpu=self._ray_remote_args.get("num_cpus", 0) * num_active_workers,
        gpu=self._ray_remote_args.get("num_gpus", 0) * num_active_workers,
    )

After:

def current_logical_usage(self) -> ExecutionResources:
    num_active_workers = self.num_active_tasks()
    return ExecutionResources(
        cpu=self._ray_remote_args.get("num_cpus", 0) * num_active_workers,
        gpu=self._ray_remote_args.get("num_gpus", 0) * num_active_workers,
        memory=self._ray_remote_args.get("memory", 0) * num_active_workers,
    )

ActorPoolMapOperator도 동일하게 수정

def current_logical_usage(self) -> ExecutionResources:
    num_active_workers = self._actor_pool.current_size()
    return ExecutionResources(
        cpu=self._ray_remote_args.get("num_cpus", 0) * num_active_workers,
        gpu=self._ray_remote_args.get("num_gpus", 0) * num_active_workers,
        memory=self._ray_remote_args.get("memory", 0) * num_active_workers,
    )

왜 이게 좋은가

  • 메모리 초과 할당 방지: 리소스 매니저가 메모리도 고려하므로, 가용 메모리를 초과하는 태스크 제출을 차단합니다.
  • 일관된 네이밍: processor_usage에서 logical_usage로 변경하여, CPU/GPU/메모리를 포괄하는 의미를 정확히 전달합니다.
  • 테스트 추가: test_memory_limit_blocks_task_submission 테스트로 메모리 제한이 태스크 제출을 차단하는 동작을 검증합니다.
  • 점진적 개선: HashShuffle 오퍼레이터에도 메모리 추적이 추가되어, 모든 오퍼레이터 유형에서 일관되게 메모리를 추적합니다.

참고 자료

댓글

관련 포스트

PR Analysis 의 다른글