공용 서버 GPU 자원 관리하기

서버 GPU 어떻게 관리하지?

연구실같은 서버 환경에서는 여러 사용자가 동일한 GPU 자원을 공유하게 된다. 이때 자원을 효율적으로 사용하지 못하면 다른 사용자가 실험이나 작업을 수행할 때 불편을 겪거나, GPU 자원이 불필요하게 낭비될 수 있다.

 

특히 실험이 끝났는데 일부 프로세스가 종료되지 않은 경우, GPU 메모리를 계속 점유하면서 사용자가 GPU를 사용할 수 없는 경우가 생긴다. 

Out of memory

 

이를 수동으로 관리하려면 관리자가 매번 서버에 접속해 nvidia-smi 명령어를 사용하여 현재 상태를 확인하고, 불필요한 프로세스를 종료하는 번거로운 작업을 반복해야 했다... 어떻게 해결하지?

 

 

근데 스케줄링 해주는 프로그램 있지 않나요?

사실 GPU 자원을 스케줄링하는 툴이 이미 있다.

SLURM, Kubernetes(KubeFlow), NVIDIA DCGM(Data Center GPU Manager)같은 프로그램이 있는데

이걸로 GPU 자원 할당이랑 스케줄링 등을 효과적으로 할 수 있다.

 

이런 시스템은 대규모 클러스터 환경에서 유용한데

연구실이나 소규모 서버 환경에서는 도입과 설정이 번거롭고 과도한 솔루션인 것 같았다.

SLURM은 소잡는 칼 아닌가?

 

SLURM은 소잡는 칼 아닌가 해서 연구실 사람들이 GPU 자원을 효율적으로 관리할 수 있는 스크립트를 따로 만들었다.

 

Method

NVIDIA의 nvidia-smi는 GPU의 현재 상태를 확인할 수 있는 유용한 도구다.

nvidia-smi를 통해 GPU의 전력 사용량, GPU 사용률, 메모리 사용량을 확인할 수 있는데, 이를 기반으로 GPU가 실제로 사용 중인지, 유휴 상태인지 판단할 수 있다. 이걸 활용해 간단히 GPU 상태를 스케줄링 할 수 있게 만들어보고싶엇다.

 

만든 스크립트는 크게 2개 역할을 한다.

 

• GPU 유휴 상태 감지: 전력 사용량, 사용률, 메모리 사용량을 기준으로 GPU가 유휴 상태인지 확인
• 프로세스 종료: 가장 많은 메모리를 사용하고 있는 프로세스를 찾아 사용자 확인 후 종료할 수 있다.

#!/bin/bash

# GPU의 유휴 상태를 감지하고, 유휴 상태일 경우 가장 많은 메모리를 사용하는 프로세스를 종료하는 스크립트

# nvidia-smi로 GPU 상태를 확인
GPU_INFO=$(nvidia-smi --query-gpu=index,power.draw,power.limit,utilization.gpu,memory.used --format=csv,noheader,nounits)

# 파싱해서 GPU 정보 추출
GPU_INDEX=$(echo "$GPU_INFO" | awk -F',' '{print $1}')
POWER_DRAW=$(echo "$GPU_INFO" | awk -F',' '{print $2}')   # 현재 전력 사용량 (W)
POWER_LIMIT=$(echo "$GPU_INFO" | awk -F',' '{print $3}')  # 최대 전력 용량 (W)
GPU_UTIL=$(echo "$GPU_INFO" | awk -F',' '{print $4}')     # GPU 사용률 (%)
MEMORY_USED=$(echo "$GPU_INFO" | awk -F',' '{print $5}')  # 메모리 사용량 (MiB)

# 유휴 상태 기준 설정
IDLE_POWER_THRESHOLD=100    # 100W 이하일 경우 유휴 상태로 간주
IDLE_UTIL_THRESHOLD=20      # 20% 이하일 경우 유휴 상태로 간주
IDLE_MEMORY_THRESHOLD=8192  # 8192MiB 이하일 경우 유휴 상태로 간주

# 유휴 상태 감지 함수
check_gpu_idle() {
    if (( $(echo "$POWER_DRAW < $IDLE_POWER_THRESHOLD" | bc -l) )) && \
       (( "$GPU_UTIL" < "$IDLE_UTIL_THRESHOLD" )) && \
       (( "$MEMORY_USED" < "$IDLE_MEMORY_THRESHOLD" )); then
        return 0  # 유휴 상태로 간주
    else
        return 1  # 사용 중인 상태로 간주
    fi
}

# 가장 많은 메모리를 사용하는 프로세스를 종료하는 함수
terminate_largest_gpu_process() {
    LARGEST_INFO=$(nvidia-smi --query-compute-apps=pid,process_name,used_memory --format=csv,noheader,nounits | sort -k3 -nr | head -n 1)

    # 유효한 프로세스가 없는 경우를 체크
    if [ -z "$LARGEST_INFO" ]; then
        echo "현재 실행 중인 GPU 프로세스가 없습니다."
        return
    fi

    LARGEST_PID=$(echo $LARGEST_INFO | awk '{print $1}' | sed 's/,//')
    PROCESS_NAME=$(echo $LARGEST_INFO | awk '{print $2}')
    LARGEST_MEMORY=$(echo $LARGEST_INFO | awk '{print $3}')

    USER=$(ps -o user= -p $LARGEST_PID 2>/dev/null)

    # 종료 여부를 사용자에게 확인
    echo "가장 많은 메모리를 사용하는 프로세스 정보:"
    echo "사용자: ${USER:-알 수 없음}"
    echo "PID: $LARGEST_PID"
    echo "프로세스 이름: $PROCESS_NAME"
    echo "GPU 메모리 사용량: $LARGEST_MEMORY MiB"

    read -p "이 프로세스를 정말로 종료하시겠습니까? 종료 전 프로세스 사용자에게 확인받으세요. [y/N]: " CONFIRM
    if [[ "$CONFIRM" =~ ^[Yy]$ ]]; then
        echo "$LARGEST_PID번 프로세스(사용자: ${USER:-알 수 없음})를 종료합니다..."
        kill -9 $LARGEST_PID
        echo "$LARGEST_PID번 프로세스가 종료되었습니다."
    else
        echo "프로세스 종료가 취소되었습니다."
    fi
}

# 메인 스크립트 실행
if check_gpu_idle; then
    echo "$GPU_INDEX번 GPU가 유휴 상태로 감지되었습니다. (전력:${POWER_DRAW} /${POWER_LIMIT}W, 사용률:${GPU_UTIL}%, 메모리 사용량:${MEMORY_USED}MiB)"
    terminate_largest_gpu_process
else
    echo "$GPU_INDEX번 GPU가 사용 중입니다. (전력:${POWER_DRAW} /${POWER_LIMIT}W, 사용률:${GPU_UTIL}%, 메모리 사용량:${MEMORY_USED}MiB)"
    echo "GPU가 사용중인것으로 확인되므로, 프로세스 종료를 권장하지 않습니다."
    read -p "계속해서 프로세스 종료 작업을 수행합니까? [y/N]: " CONFIRM
    terminate_largest_gpu_process
fi

 

마지막 메인스크립트를 아래처럼 수정하면, GPU가 사용중이면 프로세스 종료를 하지 못하게 막을 수도 있다.

# ...생략...

# 메인 스크립트 실행
if check_gpu_idle; then
    echo "GPU $GPU_INDEX번 GPU가 유휴 상태로 감지되었습니다. (전력:${POWER_DRAW} /${POWER_LIMIT}W, 사용률:${GPU_UTIL}%, 메모리 사용량:${MEMORY_USED}MiB)"
    terminate_largest_gpu_process
else
    echo "$GPU_INDEX번 GPU가 사용 중입니다. (전력:${POWER_DRAW} /${POWER_LIMIT}W, 사용률:${GPU_UTIL}%, 메모리 사용량:${MEMORY_USED}MiB)"
    echo "GPU가 사용중이므로 프로세스를 종료할 수 없습니다."
fi

 

사용방법

https://github.com/godxxy1229/kill_largest_gpu_process.sh

GitHub - godxxy1229/kill_largest_gpu_process.sh: 가장 많은 GPU 메모리를 사용 중인 프로세스를 찾아 종료하

 

GPU 자원 관리가 필요한 소규모 연구실 환경에서는

이런 간단한 스크립트가 서버 관리자에게 큰 도움이 될 수 있을 것 같다.

 

시간이 된다면...
GUI를 통해 접속해있는 사용자 및 관리자에게 알림이나 로그가 남도록 하고

프로세스가 2개 이상 존재할 경우 원하는 프로세스를 선택해서 종료할 수 있도록 해야겠다.