Skills › Scientific Agent Skills › get-available-resources › accelerator_diagnostics.py
accelerator_diagnostics.py
151 lines · skills/get-available-resources/scripts/accelerator_diagnostics.py at 5f2bc03
1#!/usr/bin/env python3
2"""Generate, but never execute, a read-only accelerator diagnostic plan."""
4from __future__ import annotations
6import argparse
7from typing import Any, Mapping, Sequence
9from _common import (
10 SCHEMA_VERSION,
11 ResourceToolError,
12 cli_error,
13 emit_json,
14 read_json_file,
15)
16from snapshot_tools import validate_snapshot
18FIXED_COMMANDS: dict[str, list[list[str]]] = {
19 "cuda": [
20 [
21 "nvidia-smi",
22 "--query-gpu=index,name,memory.total,memory.free,driver_version,compute_cap",
23 "--format=csv,noheader,nounits",
24 ]
25 ],
26 "metal": [["system_profiler", "SPDisplaysDataType", "-json"]],
27 "rocm": [
28 ["amd-smi", "static", "--asic", "--vram", "--json"],
29 [
30 "rocm-smi",
31 "--showproductname",
32 "--showmeminfo",
33 "vram",
34 "--json",
35 ],
36 ],
37}
40def build_diagnostic_plan(
41 snapshot: Mapping[str, Any],
42 *,
43 backend: str = "auto",
44) -> dict[str, Any]:
45 """Return fixed read-only probes and interpretation gates."""
46 validation = validate_snapshot(snapshot)
47 if not validation["valid"]:
48 raise ResourceToolError("snapshot does not satisfy schema 1.1")
49 if backend not in {"auto", "cuda", "metal", "rocm"}:
50 raise ResourceToolError("backend must be auto, cuda, metal, or rocm")
52 counts = snapshot["accelerators"].get("candidate_counts", {})
53 if backend == "auto":
54 selected = [
55 candidate
56 for candidate in ("cuda", "rocm", "metal")
57 if isinstance(counts.get(candidate), int) and counts[candidate] > 0
58 ]
59 else:
60 selected = [backend]
62 checks: list[dict[str, Any]] = []
63 for candidate in selected:
64 checks.append(
65 {
66 "backend": candidate,
67 "commands": FIXED_COMMANDS[candidate],
68 "management_visibility_gate": (
69 "A successful query confirms only management-tool visibility."
70 ),
71 "permission_gate": (
72 "Verify the allocation/container exposes the intended device; "
73 "do not broaden device permissions automatically."
74 ),
75 "runtime_gate": (
76 "Run the target framework's documented availability check in "
77 "the exact environment and verify driver/runtime compatibility."
78 ),
79 }
80 )
82 if not selected:
83 status = "no_management_visible_candidate"
84 elif any(
85 snapshot["accelerators"]
86 .get("candidate_upper_bounds", {})
87 .get(candidate)
88 == 0
89 for candidate in selected
90 ):
91 status = "visibility_or_allocation_restriction_present"
92 else:
93 status = "manual_runtime_validation_required"
94 return {
95 "checks": checks,
96 "execution": "not_performed",
97 "kind": "accelerator_diagnostic_plan",
98 "prohibited_actions": [
99 "driver installation or mutation",
100 "device reset, clock, power, or persistence changes",
101 "stress tests or large allocations",
102 "broad environment dumps",
103 ],
104 "schema_version": SCHEMA_VERSION,
105 "selected_backends": selected,
106 "status": status,
107 }
110def build_parser() -> argparse.ArgumentParser:
111 parser = argparse.ArgumentParser(
112 description="Produce a read-only accelerator diagnostic plan; run nothing"
113 )
114 parser.add_argument("snapshot", help="validated snapshot JSON file")
115 parser.add_argument(
116 "--backend",
117 choices=("auto", "cuda", "rocm", "metal"),
118 default="auto",
119 help="backend to plan; default: auto from snapshot candidates",
120 )
121 parser.add_argument(
122 "--output",
123 metavar="FILE.json",
124 help="write a private local JSON file instead of stdout",
125 )
126 parser.add_argument(
127 "--force",
128 action="store_true",
129 help="replace an existing explicit output file",
130 )
131 return parser
134def main(argv: Sequence[str] | None = None) -> int:
135 parser = build_parser()
136 args = parser.parse_args(argv)
137 if args.force and not args.output:
138 parser.error("--force requires --output")
139 try:
140 report = build_diagnostic_plan(
141 read_json_file(args.snapshot),
142 backend=args.backend,
143 )
144 emit_json(report, args.output, force=args.force)
145 except ResourceToolError as exc:
146 return cli_error(parser, exc)
147 return 0
150if __name__ == "__main__":
151 raise SystemExit(main())