| """Small environment smoke check; does not download data or train a PII model.""" |
| import importlib.metadata |
| import json |
| import platform |
| from pathlib import Path |
|
|
| import onnxruntime as ort |
| import torch |
|
|
| result = { |
| "platform": platform.platform(), |
| "python": platform.python_version(), |
| "versions": {name: importlib.metadata.version(name) for name in ( |
| "torch", "transformers", "datasets", "accelerate", "safetensors", |
| "scikit-learn", "onnx", "onnxscript", "onnxruntime", |
| )}, |
| "mps_built": torch.backends.mps.is_built(), |
| "mps_available": torch.backends.mps.is_available(), |
| "onnxruntime_providers": ort.get_available_providers(), |
| } |
| if result["mps_available"]: |
| torch.manual_seed(0) |
| model = torch.nn.Linear(8, 2).to("mps") |
| optimizer = torch.optim.AdamW(model.parameters()) |
| before = model.weight.detach().clone() |
| loss = model(torch.ones(4, 8, device="mps")).square().mean() |
| loss.backward() |
| finite_gradients = all(torch.isfinite(p.grad).all().item() for p in model.parameters()) |
| optimizer.step() |
| torch.mps.synchronize() |
| result["mps_smoke_check"] = { |
| "finite_loss": torch.isfinite(loss).item(), |
| "finite_gradients": finite_gradients, |
| "weights_changed": not torch.equal(before, model.weight), |
| } |
| assert all(result["mps_smoke_check"].values()), result |
| out = Path(__file__).resolve().parents[1] / "artifacts" / "environment-check.json" |
| out.parent.mkdir(exist_ok=True) |
| out.write_text(json.dumps(result, indent=2) + "\n") |
| print(json.dumps(result, indent=2)) |
|
|
