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()}