| """Configuration, reproducible IO and hardware checks shared by entrypoints.""" |
| import argparse |
| import hashlib |
| import json |
| import os |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
| import yaml |
|
|
|
|
| ROOT = Path(__file__).resolve().parents[1] |
|
|
|
|
| def arguments(description): |
| parser = argparse.ArgumentParser(description=description) |
| parser.add_argument('--config', type=Path, default=ROOT / 'config/config.yaml') |
| return parser |
|
|
|
|
| def load_config(path): |
| path = Path(path).resolve() |
| with path.open() as handle: |
| config = yaml.safe_load(handle) |
| return config, path.parents[1] |
|
|
|
|
| def output_path(root, value): |
| value = Path(value) |
| return value if value.is_absolute() else root / value |
|
|
|
|
| def save_json(path, value): |
| path = Path(path) |
| path.parent.mkdir(parents=True, exist_ok=True) |
| tmp = path.with_suffix(path.suffix + '.tmp') |
| tmp.write_text(json.dumps(value, indent=2, ensure_ascii=False, allow_nan=False) + '\n') |
| tmp.replace(path) |
|
|
|
|
| def sha256(path): |
| digest = hashlib.sha256() |
| with Path(path).open('rb') as handle: |
| for part in iter(lambda: handle.read(1024 * 1024), b''): |
| digest.update(part) |
| return digest.hexdigest() |
|
|
|
|
| def device_for(config, require_dcu=True): |
| device = torch.device(config['training']['device']) |
| if require_dcu: |
| if not os.environ.get('SLURM_JOB_ID'): |
| raise RuntimeError('This experiment must run in a SLURM DCU allocation.') |
| if device.type != 'cuda' or not torch.cuda.is_available() or not torch.version.hip: |
| raise RuntimeError('A HIP/DCU PyTorch device is required; CPU fallback is disabled.') |
| return device |
|
|
|
|
| def runtime_info(device): |
| return {'hostname': os.uname().nodename, 'slurm_job_id': os.environ.get('SLURM_JOB_ID'), |
| 'slurm_partition': os.environ.get('SLURM_JOB_PARTITION'), |
| 'torch': torch.__version__, 'hip': torch.version.hip, |
| 'device': str(device), 'device_name': torch.cuda.get_device_name(device), |
| 'device_memory_bytes': torch.cuda.get_device_properties(device).total_memory} |
|
|
|
|
| def errors(prediction, target): |
| difference = np.asarray(prediction, dtype=np.float64) - np.asarray(target, dtype=np.float64) |
| return {'MAE': np.mean(np.abs(difference), axis=0).tolist(), |
| 'RMSE': np.sqrt(np.mean(difference ** 2, axis=0)).tolist()} |
|
|
|
|
| def checkpoint_save(path, payload): |
| path.parent.mkdir(parents=True, exist_ok=True) |
| temporary = path.with_suffix('.tmp') |
| torch.save(payload, temporary) |
| temporary.replace(path) |
|
|
