scripts/gpu_utils.py

import datetime,json,os,signal
from pathlib import Path
import torch
STOP=False

def stop_request(signum,frame):
 global STOP
 STOP=True

def install_signals():
 signal.signal(signal.SIGTERM,stop_request);signal.signal(signal.SIGINT,stop_request)

def expired():
 from recovery_state import deadline
 return STOP or datetime.datetime.now(datetime.timezone.utc)>=deadline()

def require_cuda():
 if not torch.cuda.is_available():raise RuntimeError('CUDA is unavailable. Run scripts/run_gpu.sh from the normal Ubuntu SSH shell. CPU fallback is forbidden.')
 x=torch.ones((32,32),device='cuda');y=x@x;torch.cuda.synchronize();assert y[0,0].item()==32
 return torch.device('cuda')

def atomic_save(state,path):
 path=Path(path);temp=path.with_suffix(path.suffix+'.tmp');torch.save(state,temp);os.replace(temp,path)

def cpu_state(model):return {k:v.detach().cpu().clone() for k,v in model.state_dict().items()}