파이톨치

[프로젝트] auto GPU warning Program 본문

프로젝트

[프로젝트] auto GPU warning Program

파이톨치 2024. 11. 11. 13:11
728x90

auto GPU warning Program 

Python 코드에서 Discord Webhook을 사용해 GPU 메모리 상태를 Discord 채널로 알릴 수 있습니다. 이 방법을 사용하면 GPU 메모리 사용량이 일정 이상일 때 Discord 채널에 경고 메시지를 자동으로 전송할 수 있습니다.

우선, Discord 채널에 Webhook을 설정하고 URL을 확보해야 합니다. 이후 Python 코드에서 requests 라이브러리를 사용해 Webhook URL로 메시지를 전송할 수 있습니다.

단계 1: Discord Webhook 설정

  1. Discord에서 메시지를 전송할 채널로 이동합니다.
  2. 채널 설정 → "앱 통합" → "웹훅"으로 이동합니다.
  3. "새 웹훅"을 클릭하고 원하는 이름과 채널을 선택한 후 "웹훅 URL 복사" 버튼을 눌러 URL을 복사합니다.

단계 2: Python 코드 작성

아래는 pynvml과 requests 라이브러리를 이용해 GPU 메모리 사용량을 확인하고, 일정 이상일 경우 Discord Webhook으로 경고 메시지를 전송하는 코드입니다.

 

import pynvml
import time
import requests

# pynvml 초기화
pynvml.nvmlInit()
handle = pynvml.nvmlDeviceGetHandleByIndex(0)

# Discord Webhook URL (위에서 복사한 URL을 여기에 붙여넣기)
webhook_url = 'https://discord.com/api/webhooks/your_webhook_id/your_webhook_token'

def send_discord_message(content):
    """Discord Webhook으로 메시지 전송"""
    data = {"content": content}
    response = requests.post(webhook_url, json=data)
    if response.status_code == 204:
        print("메시지가 성공적으로 전송되었습니다.")
    else:
        print("메시지 전송 실패:", response.status_code)

while True:
    # GPU 메모리 정보 가져오기
    info = pynvml.nvmlDeviceGetMemoryInfo(handle)
    used_memory = info.used / info.total

    # 메모리 사용량이 50%를 넘으면 Discord로 경고 메시지 전송
    if used_memory > 0.5:
        message = f"🚨 경고: GPU 메모리가 50% 이상 사용 중입니다! 현재 사용량: {used_memory * 100:.2f}%"
        send_discord_message(message)
    
    # 주기적으로 체크 (예: 10초 간격)
    time.sleep(10)

 

코드 설명

  • send_discord_message 함수는 requests.post를 통해 Discord Webhook으로 메시지를 전송합니다.
  • used_memory > 0.5 조건이 참일 때 GPU 메모리 사용량에 대한 경고 메시지를 Discord 채널로 전송합니다.

이렇게 하면 GPU 메모리 사용량이 50% 이상일 때마다 Discord 채널로 자동 알림을 받을 수 있습니다. requests 라이브러리가 필요하므로, 설치되지 않은 경우 pip install requests 명령어로 설치할 수 있습니다.

 

테스트 코드 

 
GPU 메모리를 50% 이상 사용하도록 하려면 대량의 행렬 연산을 수행하는 코드가 도움이 됩니다. 아래는 PyTorch를 사용해 큰 크기의 텐서를 생성하고 연산을 반복하여 GPU 메모리 사용량을 빠르게 증가시키는 예시입니다.
 
import torch
import time

# GPU가 사용 가능한지 확인
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
if device.type == "cpu":
    print("GPU를 사용할 수 없습니다.")
else:
    print("GPU에서 메모리를 사용하여 테스트를 시작합니다.")

# 대용량 텐서를 생성하여 GPU 메모리를 채움
# 이 텐서의 크기나 수를 조절하여 메모리 사용량을 설정할 수 있음
try:
    tensors = []
    for _ in range(10):
        # 10240x10240 크기의 텐서를 생성하고 GPU로 이동
        tensor = torch.randn(10240, 10240, device=device)
        tensors.append(tensor)
        
        # 약간의 연산을 수행하여 GPU에 부담을 줌
        tensor = tensor @ tensor.T

        # 메모리 사용량 확인 (메모리 증가 확인용)
        print(f"현재 GPU 메모리 사용량: {torch.cuda.memory_allocated(device) / 1024**2:.2f} MB")

        # 약간 대기 시간
        time.sleep(1)

except RuntimeError as e:
    print("GPU 메모리가 부족합니다:", e)

finally:
    # 메모리 해제
    tensors.clear()
    torch.cuda.empty_cache()
    print("GPU 메모리가 해제되었습니다.")

 

코드 설명

  • 텐서 생성: 10240x10240 크기의 랜덤 텐서를 여러 개 생성하여 GPU 메모리를 빠르게 사용합니다.
  • 연산 수행: 생성한 텐서에 행렬 곱 연산을 수행하여 메모리와 GPU 연산 리소스를 함께 사용합니다.
  • 메모리 사용량 출력: torch.cuda.memory_allocated를 통해 현재 메모리 사용량을 출력합니다.
  • GPU 메모리 해제: tensors.clear()와 torch.cuda.empty_cache()를 통해 연산이 끝난 후 GPU 메모리를 해제합니다.

이 코드로 GPU 메모리를 빠르게 사용할 수 있으므로, 원하시는 사용량에 따라 텐서 크기와 반복 횟수를 조정하면 됩니다.

728x90