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