File size: 3,966 Bytes
e5a6e71 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 | """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
# A uniform state remains uniform for all learned transport plans.
uniform = torch.ones_like(initial)
assert torch.allclose(model(uniform), uniform, atol=2e-6)
# Every channel conserves its sum for arbitrary signed momentum.
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
# Independent stop-gradient check on an explicit two-state feedback graph.
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()
|