| """OneScience installer preflight and independent structural checks; no formal weights.""" |
| import ast |
| import importlib |
| import json |
| import os |
| import shutil |
| import time |
|
|
| import torch |
| from torch.utils.data import DataLoader |
|
|
| from common import (ROOT, Trajectories, Windows, build_model, device, load_config, |
| output_dir, parser, provenance, save_json, setup, validate) |
| from train import train_batch |
|
|
|
|
| def main(): |
| args = parser(__doc__).parse_args() |
| cfg = load_config(args.config) |
| setup(cfg["seed"]) |
| target_device = device() |
| for module in ("onescience", "torch", "numpy", "h5py", "yaml", "matplotlib"): |
| importlib.import_module(module) |
| print("DEPENDENCY_OK " + module, flush=True) |
| for path in ROOT.rglob("*.py"): |
| ast.parse(path.read_text()) |
| for command in ("sbatch", "squeue", "sacct"): |
| assert shutil.which(command), command |
| datasets = {split: Trajectories(cfg, split) for split in cfg["data"]["expected_counts"]} |
| out = output_dir(cfg) |
| save_json(out / "dataset_validation.json", {s: d.metadata for s, d in datasets.items()}) |
| model = build_model(cfg, target_device) |
| sample = datasets["train"].values[0, :6].unsqueeze(0).to(target_device) |
| model.eval() |
| with torch.no_grad(): |
| initial = sample[:, 0] |
| predicted = model(initial) |
| assert predicted.shape == initial.shape |
| assert torch.isfinite(predicted).all() |
| assert predicted[:, 0].min() >= 0 |
| drift = (predicted.double().sum((-2, -1)) - initial.double().sum((-2, -1))).abs() |
| scale = initial.double().abs().sum((-2, -1)).clamp_min(1.0) |
| assert (drift / scale).max() < 2e-6 |
| |
| uniform = torch.ones_like(initial) |
| assert torch.allclose(model(uniform), uniform, atol=2e-6) |
| |
| random_state = torch.rand_like(initial) |
| random_state[:, 1:] -= 0.5 |
| result = model(random_state) |
| random_drift = (result.double().sum((-2, -1)) - random_state.double().sum((-2, -1))).abs() |
| assert (random_drift / random_state.double().abs().sum((-2, -1)).clamp_min(1)).max() < 2e-6 |
| print("STRUCTURE_OK forward, positivity, uniform state, per-channel conservation", flush=True) |
| model.train() |
| optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-2) |
| first, terminal = train_batch(model, sample, optimizer, 5) |
| assert any(p.grad is not None for p in model.parameters()) |
| validation = validate(model, DataLoader(Windows(datasets["val"], 5), batch_size=2), target_device) |
| assert validation["mse"] >= 0 |
| |
| x = sample[:, 0].detach().requires_grad_(True) |
| intermediate = model(x) |
| final = model(intermediate.detach()) |
| assert torch.autograd.grad(final.sum(), x, allow_unused=True)[0] is None |
| props = torch.cuda.get_device_properties(target_device) |
| result = {"status": "passed", "preflight_passed": True, "execution_readiness": "ready", |
| "evidence": {"preflight": {"status": "passed"}}, "installer_reason": "preflight_validation", |
| "execution_channel": "local_slurm", "provenance": provenance(), |
| "gpu_memory_gb": props.total_memory / 2**30, |
| "model_parameters": sum(p.numel() for p in model.parameters()), |
| "structural_checks": ["shape", "finite", "positivity", "uniform_state", "conservation", "detached_feedback"], |
| "real_data_train_batch_loss": {"step1": first, "step5": terminal}, |
| "real_data_validation": validation, |
| "formal_training_completed": False, "formal_weights_written": False} |
| save_json(out / "preflight_validation.json", result) |
| print("PREFLIGHT_PASSED " + json.dumps(result), flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
