본문으로 건너뛰기

[Ray] memory_full_info 호출 제거로 ML 학습 성능 저하 문제 해결

PR 링크: ray-project/ray#60000 상태: Merged | 변경: +25 / -13

들어가며

Pinterest 엔지니어링 팀이 ML 학습 과정에서 원인 불명의 성능 저하를 추적한 결과, Ray의 모니터링 에이전트가 주기적으로 호출하는 psutil의 memory_full_info()가 원인이었습니다. 이 함수는 Linux에서 /proc/[pid]/smaps를 파싱하여 USS(Unique Set Size)를 계산하는데, 수백~수천 개의 메모리 매핑을 가진 프로세스에서는 호출 비용이 매우 높습니다. 이 PR은 Linux와 Windows에서 memory_full_info() 호출을 제거하고, memory_info()에서 USS를 근사 추정합니다.

핵심 코드 분석

memory_full_info 수집 대상 축소

Before:

PSUTIL_PROCESS_ATTRS = [
    "pid", "create_time", "cpu_percent", "cpu_times",
    "cmdline", "memory_info", "memory_full_info",
] + (["num_fds"] if sys.platform != "win32" else [])

After:

PSUTIL_PROCESS_ATTRS = (
    ["pid", "create_time", "cpu_percent", "cpu_times",
     "cmdline", "memory_info"]
    + (["num_fds"] if sys.platform != "win32" else [])
    # Only collect memory_full_info in Mac OS X
    + (["memory_full_info"] if sys.platform == "darwin" else [])
)

USS 근사 추정

Before:

mem_full_info = stat.get("memory_full_info")
if mem_full_info is not None:
    total_uss += float(mem_full_info.uss) / 1.0e6

After:

mem_full_info = stat.get("memory_full_info")
if mem_full_info is not None:
    # For Mac OS X, directly get USS metric from memory_full_info
    total_uss += float(mem_full_info.uss) / 1.0e6
elif memory_info is not None:
    # For linux or windows, approximate USS from memory_info
    if hasattr(memory_info, "shared"):
        # Linux: USS ≈ RSS - shared
        total_uss += float(memory_info.rss - memory_info.shared) / 1.0e6
    elif hasattr(memory_info, "private"):
        # Windows: private IS USS
        total_uss += float(memory_info.private) / 1.0e6

왜 이게 좋은가

  1. ML 학습 성능 복원: memory_full_info()는 /proc/[pid]/smaps 전체를 파싱하므로, GPU 메모리 맵이 많은 ML 워크로드에서 특히 비용이 높습니다. 이 호출을 제거하면 학습 중 주기적인 지연이 사라집니다.
  2. 합리적 근사: Linux에서 USS ≈ RSS - shared는 정확한 USS는 아니지만, 모니터링 대시보드용으로는 충분한 정밀도를 제공합니다.
  3. 플랫폼별 최적화: macOS는 Mach 커널 API를 통해 USS를 저비용으로 제공하므로 기존 방식을 유지하고, 비용이 큰 Linux/Windows만 근사 추정으로 전환했습니다.
  4. 실제 프로덕션 검증: Pinterest가 프로덕션 ML 파이프라인에서 발견하고 디버깅한 실전 이슈에 대한 수정입니다.

참고 자료

댓글

관련 포스트

PR Analysis 의 다른글