113 lines
3.3 KiB
Python
113 lines
3.3 KiB
Python
"""Command-line entry point for safe Brain runtime diagnostics."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import platform
|
|
from collections.abc import Sequence
|
|
from typing import Any
|
|
|
|
from yovision_brain import __version__
|
|
|
|
|
|
def _load_torch() -> Any:
|
|
try:
|
|
import torch
|
|
except ImportError as exc:
|
|
raise RuntimeError(
|
|
"PyTorch runtime is not installed; install the pinned CPU or CUDA wheel "
|
|
"from Brain/README.md"
|
|
) from exc
|
|
return torch
|
|
|
|
|
|
def runtime_info() -> dict[str, object]:
|
|
"""Return non-sensitive interpreter and compute-runtime facts."""
|
|
info: dict[str, object] = {
|
|
"python": platform.python_version(),
|
|
"torch_installed": False,
|
|
"torch_version": None,
|
|
"cuda_build": None,
|
|
"cuda_available": False,
|
|
"cuda_device_count": 0,
|
|
}
|
|
try:
|
|
torch = _load_torch()
|
|
except RuntimeError:
|
|
return info
|
|
|
|
cuda_available = bool(torch.cuda.is_available())
|
|
info.update(
|
|
{
|
|
"torch_installed": True,
|
|
"torch_version": torch.__version__,
|
|
"cuda_build": torch.version.cuda,
|
|
"cuda_available": cuda_available,
|
|
"cuda_device_count": torch.cuda.device_count() if cuda_available else 0,
|
|
}
|
|
)
|
|
return info
|
|
|
|
|
|
def smoke(device: str) -> dict[str, object]:
|
|
"""Run a deterministic tensor operation on exactly the requested device."""
|
|
torch = _load_torch()
|
|
if device == "cuda" and not torch.cuda.is_available():
|
|
raise RuntimeError("CUDA was requested but PyTorch reports no available CUDA device")
|
|
|
|
tensor = torch.tensor([[1.0, 2.0], [3.0, 4.0]], device=device)
|
|
result = tensor @ tensor
|
|
expected = torch.tensor([[7.0, 10.0], [15.0, 22.0]], device=device)
|
|
if not torch.equal(result, expected):
|
|
raise RuntimeError("tensor smoke result did not match the expected value")
|
|
|
|
return {
|
|
"status": "ok",
|
|
"device": device,
|
|
"torch_version": torch.__version__,
|
|
"cuda_build": torch.version.cuda,
|
|
}
|
|
|
|
|
|
def build_parser() -> argparse.ArgumentParser:
|
|
parser = argparse.ArgumentParser(
|
|
prog="yovision-brain",
|
|
description="YoVision Brain runtime diagnostics",
|
|
)
|
|
parser.add_argument("--version", action="version", version=f"%(prog)s {__version__}")
|
|
action = parser.add_mutually_exclusive_group()
|
|
action.add_argument(
|
|
"--runtime-info",
|
|
action="store_true",
|
|
help="print non-sensitive Python/PyTorch/CUDA capability information",
|
|
)
|
|
action.add_argument(
|
|
"--smoke",
|
|
choices=("cpu", "cuda"),
|
|
help="run a tensor smoke test on exactly the selected device",
|
|
)
|
|
return parser
|
|
|
|
|
|
def main(argv: Sequence[str] | None = None) -> int:
|
|
args = build_parser().parse_args(argv)
|
|
try:
|
|
if args.runtime_info:
|
|
payload = runtime_info()
|
|
elif args.smoke:
|
|
payload = smoke(args.smoke)
|
|
else:
|
|
build_parser().print_help()
|
|
return 0
|
|
except RuntimeError as exc:
|
|
print(json.dumps({"status": "error", "message": str(exc)}, ensure_ascii=False))
|
|
return 2
|
|
|
|
print(json.dumps(payload, ensure_ascii=False, sort_keys=True))
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|