45 lines
1.1 KiB
Python
45 lines
1.1 KiB
Python
import torch
|
|
import time
|
|
import os
|
|
import signal
|
|
import sys
|
|
|
|
print(f"[Worker] PID: {os.getpid()}, PPID: {os.getppid()}")
|
|
|
|
# GPU 메모리 할당 (~2GB)
|
|
device = torch.device('cuda:0')
|
|
tensors = []
|
|
for i in range(8):
|
|
t = torch.randn(1024, 1024, 64, device=device)
|
|
tensors.append(t)
|
|
|
|
allocated = torch.cuda.memory_allocated() / 1024**3
|
|
print(f"[Worker] GPU 메모리 할당 완료: {allocated:.2f} GB")
|
|
|
|
# === 핵심: SIGTERM handler 등록 ===
|
|
def graceful_shutdown(signum, frame):
|
|
print(f"[Worker] Signal {signum} 수신, CUDA cleanup 시작...")
|
|
|
|
# 1. tensor 참조 해제
|
|
tensors.clear()
|
|
|
|
# 2. GPU 캐시 비우기
|
|
torch.cuda.empty_cache()
|
|
|
|
# 3. 동기화 후 종료
|
|
torch.cuda.synchronize()
|
|
|
|
after = torch.cuda.memory_allocated() / 1024**3
|
|
print(f"[Worker] cleanup 완료, 잔여 GPU 메모리: {after:.2f} GB")
|
|
|
|
sys.exit(0)
|
|
|
|
signal.signal(signal.SIGTERM, graceful_shutdown)
|
|
signal.signal(signal.SIGINT, graceful_shutdown)
|
|
|
|
# 무한 대기
|
|
while True:
|
|
time.sleep(5)
|
|
print(f"[Worker] alive, PID={os.getpid()}, PPID={os.getppid()}, "
|
|
f"GPU={torch.cuda.memory_allocated()/1024**3:.2f}GB")
|