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")