Skills › Scientific Agent Skills › get-available-resources › snapshot_tools.py
snapshot_tools.py
504 lines · skills/get-available-resources/scripts/snapshot_tools.py at 5f2bc03
1#!/usr/bin/env python3
2"""Validate or deterministically diff resource snapshots."""
4from __future__ import annotations
6import argparse
7import copy
8import datetime as dt
9from typing import Any, Mapping, Sequence
11from _common import (
12 MAX_BYTES,
13 SCHEMA_VERSION,
14 ResourceToolError,
15 bounded_number,
16 cli_error,
17 emit_json,
18 read_json_file,
19)
21MAX_DIFFS = 512
24def _is_optional_nonnegative_number(value: Any) -> bool:
25 return value is None or bounded_number(value, maximum=float(MAX_BYTES)) is not None
28def _expect_mapping(
29 parent: Mapping[str, Any],
30 key: str,
31 path: str,
32 errors: list[dict[str, str]],
33) -> Mapping[str, Any]:
34 value = parent.get(key)
35 if not isinstance(value, dict):
36 errors.append(
37 {
38 "code": "TYPE",
39 "message": "must be an object",
40 "path": f"{path}.{key}",
41 }
42 )
43 return {}
44 return value
47def _check_optional_number(
48 parent: Mapping[str, Any],
49 key: str,
50 path: str,
51 errors: list[dict[str, str]],
52) -> None:
53 if key not in parent or not _is_optional_nonnegative_number(parent.get(key)):
54 errors.append(
55 {
56 "code": "NUMBER",
57 "message": "required field must be null or a finite nonnegative bounded number",
58 "path": f"{path}.{key}",
59 }
60 )
63def validate_snapshot(snapshot: Mapping[str, Any]) -> dict[str, Any]:
64 """Validate the stable 1.1 resource snapshot contract."""
65 errors: list[dict[str, str]] = []
66 notes: list[dict[str, str]] = []
67 if snapshot.get("schema_version") != SCHEMA_VERSION:
68 errors.append(
69 {
70 "code": "SCHEMA_VERSION",
71 "message": f"must equal {SCHEMA_VERSION}",
72 "path": "$.schema_version",
73 }
74 )
75 if snapshot.get("snapshot_kind") != "effective_resource_snapshot":
76 errors.append(
77 {
78 "code": "SNAPSHOT_KIND",
79 "message": "must equal effective_resource_snapshot",
80 "path": "$.snapshot_kind",
81 }
82 )
83 observed_at = snapshot.get("observed_at")
84 try:
85 valid_time = (isinstance(observed_at, str) and observed_at.endswith("Z")
86 and dt.datetime.fromisoformat(observed_at[:-1] + "+00:00").utcoffset()
87 == dt.timedelta(0))
88 except ValueError:
89 valid_time = False
90 if not valid_time:
91 errors.append(
92 {
93 "code": "TIMESTAMP",
94 "message": "must be a UTC string ending in Z",
95 "path": "$.observed_at",
96 }
97 )
99 platform = _expect_mapping(snapshot, "platform", "$", errors)
100 for key in ("system", "machine", "python_version"):
101 if not isinstance(platform.get(key), str) or not platform.get(key):
102 errors.append(
103 {
104 "code": "PLATFORM",
105 "message": "must be a nonempty string",
106 "path": f"$.platform.{key}",
107 }
108 )
110 cpu = _expect_mapping(snapshot, "cpu", "$", errors)
111 host_cpu = _expect_mapping(cpu, "host", "$.cpu", errors)
112 process_cpu = _expect_mapping(cpu, "process", "$.cpu", errors)
113 cpu_cgroup = _expect_mapping(cpu, "cgroup_v2", "$.cpu", errors)
114 effective_cpu = _expect_mapping(cpu, "effective", "$.cpu", errors)
115 for key in ("logical", "physical"):
116 _check_optional_number(host_cpu, key, "$.cpu.host", errors)
117 for key in ("affinity_logical", "python_available_logical"):
118 _check_optional_number(process_cpu, key, "$.cpu.process", errors)
119 for key in ("cpuset_logical", "quota_cores"):
120 _check_optional_number(cpu_cgroup, key, "$.cpu.cgroup_v2", errors)
121 _check_optional_number(
122 effective_cpu, "capacity_cores", "$.cpu.effective", errors
123 )
124 worker_ceiling = effective_cpu.get("worker_ceiling")
125 if (
126 isinstance(worker_ceiling, bool)
127 or not isinstance(worker_ceiling, int)
128 or not 1 <= worker_ceiling <= 1024
129 ):
130 errors.append(
131 {
132 "code": "WORKER_CEILING",
133 "message": "must be an integer between 1 and 1024",
134 "path": "$.cpu.effective.worker_ceiling",
135 }
136 )
137 if not isinstance(effective_cpu.get("limiting_sources"), list):
138 errors.append(
139 {
140 "code": "LIMIT_SOURCES",
141 "message": "must be an array",
142 "path": "$.cpu.effective.limiting_sources",
143 }
144 )
146 memory = _expect_mapping(snapshot, "memory", "$", errors)
147 host_memory = _expect_mapping(memory, "host", "$.memory", errors)
148 effective_memory = _expect_mapping(memory, "effective", "$.memory", errors)
149 memory_cgroup = _expect_mapping(memory, "cgroup_v2", "$.memory", errors)
150 swap = _expect_mapping(memory, "swap", "$.memory", errors)
151 for key in ("available_bytes", "total_bytes"):
152 _check_optional_number(host_memory, key, "$.memory.host", errors)
153 for key in (
154 "available_bytes",
155 "hard_limit_bytes",
156 "pressure_threshold_bytes",
157 ):
158 _check_optional_number(effective_memory, key, "$.memory.effective", errors)
159 for key in ("available_bytes", "current_bytes", "high_bytes", "max_bytes"):
160 _check_optional_number(memory_cgroup, key, "$.memory.cgroup_v2", errors)
161 for key in ("free_bytes", "total_bytes"):
162 _check_optional_number(swap, key, "$.memory.swap", errors)
163 available = effective_memory.get("available_bytes")
164 hard_limit = effective_memory.get("hard_limit_bytes")
165 if (
166 isinstance(available, (int, float))
167 and not isinstance(available, bool)
168 and isinstance(hard_limit, (int, float))
169 and not isinstance(hard_limit, bool)
170 and available > hard_limit
171 ):
172 errors.append(
173 {
174 "code": "MEMORY_ORDER",
175 "message": "available memory cannot exceed the effective hard limit",
176 "path": "$.memory.effective.available_bytes",
177 }
178 )
180 disk = _expect_mapping(snapshot, "disk", "$", errors)
181 for key in ("capacity_bytes", "free_bytes", "user_available_bytes"):
182 _check_optional_number(disk, key, "$.disk", errors)
183 if disk.get("writable") is not None and not isinstance(
184 disk.get("writable"), bool
185 ):
186 errors.append(
187 {
188 "code": "DISK_WRITABLE",
189 "message": "must be boolean or null",
190 "path": "$.disk.writable",
191 }
192 )
193 free = disk.get("free_bytes")
194 user_available = disk.get("user_available_bytes")
195 capacity = disk.get("capacity_bytes")
196 if all(
197 isinstance(value, (int, float)) and not isinstance(value, bool)
198 for value in (free, user_available, capacity)
199 ):
200 if not 0 <= user_available <= free <= capacity:
201 errors.append(
202 {
203 "code": "DISK_ORDER",
204 "message": "expected user_available <= free <= capacity",
205 "path": "$.disk",
206 }
207 )
209 accelerators = _expect_mapping(snapshot, "accelerators", "$", errors)
210 for key in ("candidate_counts", "candidate_upper_bounds"):
211 counts = _expect_mapping(accelerators, key, "$.accelerators", errors)
212 for backend in ("cuda", "metal", "rocm"):
213 value = counts.get(backend)
214 if (backend not in counts or
215 (value is None and key == "candidate_counts") or
216 (value is not None and (isinstance(value, bool) or not isinstance(value, int)
217 or not 0 <= value <= 256))):
218 errors.append({"code": "ACCELERATOR_COUNT", "message": "invalid candidate count",
219 "path": f"$.accelerators.{key}.{backend}"})
220 devices = accelerators.get("devices")
221 if not isinstance(devices, list) or len(devices) > 256:
222 errors.append(
223 {
224 "code": "ACCELERATOR_DEVICES",
225 "message": "must be an array with at most 256 entries",
226 "path": "$.accelerators.devices",
227 }
228 )
229 else:
230 for index, device in enumerate(devices):
231 if not isinstance(device, dict):
232 errors.append(
233 {
234 "code": "ACCELERATOR_DEVICE",
235 "message": "must be an object",
236 "path": f"$.accelerators.devices[{index}]",
237 }
238 )
239 continue
240 backend = device.get("backend_candidate")
241 if not isinstance(backend, str) or backend not in {"cuda", "metal", "rocm"}:
242 errors.append(
243 {
244 "code": "ACCELERATOR_BACKEND",
245 "message": "must be cuda, metal, or rocm",
246 "path": (
247 f"$.accelerators.devices[{index}].backend_candidate"
248 ),
249 }
250 )
251 if device.get("device_permission") != "not_tested":
252 notes.append(
253 {
254 "code": "PERMISSION_ASSERTION",
255 "message": (
256 "device permission differs from the detector's "
257 "conservative not_tested default"
258 ),
259 "path": (
260 f"$.accelerators.devices[{index}].device_permission"
261 ),
262 }
263 )
264 if device.get("runtime_compatibility") != "not_tested":
265 notes.append(
266 {
267 "code": "RUNTIME_ASSERTION",
268 "message": (
269 "runtime compatibility differs from the detector's "
270 "conservative not_tested default"
271 ),
272 "path": (
273 f"$.accelerators.devices[{index}].runtime_compatibility"
274 ),
275 }
276 )
278 scheduler = _expect_mapping(snapshot, "scheduler", "$", errors)
279 if not isinstance(scheduler.get("fields_read"), list):
280 errors.append(
281 {
282 "code": "SCHEDULER_FIELDS",
283 "message": "must be an array",
284 "path": "$.scheduler.fields_read",
285 }
286 )
287 privacy = _expect_mapping(snapshot, "privacy", "$", errors)
288 for key in (
289 "absolute_paths_included",
290 "environment_values_included",
291 "hostnames_included",
292 "identifiers_redacted_by_default",
293 ):
294 if not isinstance(privacy.get(key), bool):
295 errors.append(
296 {
297 "code": "PRIVACY",
298 "message": "must be boolean",
299 "path": f"$.privacy.{key}",
300 }
301 )
302 if any(
303 privacy.get(key) is True
304 for key in (
305 "absolute_paths_included",
306 "environment_values_included",
307 "hostnames_included",
308 )
309 ):
310 notes.append(
311 {
312 "code": "PRIVACY_EXPANDED",
313 "message": "snapshot includes fields normally redacted",
314 "path": "$.privacy",
315 }
316 )
318 warning_items = snapshot.get("warnings")
319 if not isinstance(warning_items, list) or len(warning_items) > 512:
320 errors.append(
321 {
322 "code": "WARNINGS",
323 "message": "must be an array with at most 512 entries",
324 "path": "$.warnings",
325 }
326 )
327 provenance = snapshot.get("provenance")
328 if not isinstance(provenance, list) or len(provenance) > 512:
329 errors.append(
330 {
331 "code": "PROVENANCE",
332 "message": "must be an array with at most 512 entries",
333 "path": "$.provenance",
334 }
335 )
337 errors.sort(key=lambda item: (item["path"], item["code"], item["message"]))
338 notes.sort(key=lambda item: (item["path"], item["code"], item["message"]))
339 return {
340 "errors": errors,
341 "kind": "snapshot_validation",
342 "notes": notes,
343 "schema_version": SCHEMA_VERSION,
344 "valid": not errors,
345 }
348def _strip_volatile(snapshot: Mapping[str, Any]) -> dict[str, Any]:
349 cleaned = copy.deepcopy(dict(snapshot))
350 cleaned.pop("observed_at", None)
351 return cleaned
354def _diff_values(
355 before: Any,
356 after: Any,
357 path: str,
358 changes: list[dict[str, Any]],
359) -> None:
360 if len(changes) >= MAX_DIFFS:
361 return
362 if isinstance(before, dict) and isinstance(after, dict):
363 for key in sorted(set(before) | set(after)):
364 child_path = f"{path}.{key}"
365 if key not in before:
366 changes.append(
367 {
368 "after": after[key],
369 "before": None,
370 "change": "added",
371 "path": child_path,
372 }
373 )
374 elif key not in after:
375 changes.append(
376 {
377 "after": None,
378 "before": before[key],
379 "change": "removed",
380 "path": child_path,
381 }
382 )
383 else:
384 _diff_values(before[key], after[key], child_path, changes)
385 if len(changes) >= MAX_DIFFS:
386 return
387 return
388 if isinstance(before, list) and isinstance(after, list):
389 maximum = max(len(before), len(after))
390 for index in range(maximum):
391 child_path = f"{path}[{index}]"
392 if index >= len(before):
393 changes.append(
394 {
395 "after": after[index],
396 "before": None,
397 "change": "added",
398 "path": child_path,
399 }
400 )
401 elif index >= len(after):
402 changes.append(
403 {
404 "after": None,
405 "before": before[index],
406 "change": "removed",
407 "path": child_path,
408 }
409 )
410 else:
411 _diff_values(before[index], after[index], child_path, changes)
412 if len(changes) >= MAX_DIFFS:
413 return
414 return
415 if before != after:
416 changes.append(
417 {
418 "after": after,
419 "before": before,
420 "change": "changed",
421 "path": path,
422 }
423 )
426def diff_snapshots(
427 before: Mapping[str, Any],
428 after: Mapping[str, Any],
429 *,
430 include_volatile: bool = False,
431) -> dict[str, Any]:
432 """Return a bounded, stable structural diff."""
433 before_validation = validate_snapshot(before)
434 after_validation = validate_snapshot(after)
435 if not before_validation["valid"] or not after_validation["valid"]:
436 raise ResourceToolError("both snapshots must validate before diffing")
437 try:
438 left = dict(before) if include_volatile else _strip_volatile(before)
439 right = dict(after) if include_volatile else _strip_volatile(after)
440 changes: list[dict[str, Any]] = []
441 _diff_values(left, right, "$", changes)
442 except RecursionError as exc:
443 raise ResourceToolError("snapshot nesting exceeds the diff safety bound") from exc
444 return {
445 "change_count": len(changes),
446 "changes": changes,
447 "ignored_paths": [] if include_volatile else ["$.observed_at"],
448 "kind": "snapshot_diff",
449 "schema_version": SCHEMA_VERSION,
450 "truncated": len(changes) >= MAX_DIFFS,
451 }
454def build_parser() -> argparse.ArgumentParser:
455 parser = argparse.ArgumentParser(
456 description="Validate or structurally diff resource snapshot JSON"
457 )
458 parser.add_argument(
459 "--output",
460 metavar="FILE.json",
461 help="write a private local JSON file instead of stdout",
462 )
463 parser.add_argument(
464 "--force",
465 action="store_true",
466 help="replace an existing explicit output file",
467 )
468 commands = parser.add_subparsers(dest="command", required=True)
469 validate = commands.add_parser("validate", help="validate one snapshot")
470 validate.add_argument("snapshot", help="snapshot JSON file")
471 diff = commands.add_parser("diff", help="diff two valid snapshots")
472 diff.add_argument("before", help="earlier snapshot JSON file")
473 diff.add_argument("after", help="later snapshot JSON file")
474 diff.add_argument(
475 "--include-volatile",
476 action="store_true",
477 help="include observed_at in the diff",
478 )
479 return parser
482def main(argv: Sequence[str] | None = None) -> int:
483 parser = build_parser()
484 args = parser.parse_args(argv)
485 if args.force and not args.output:
486 parser.error("--force requires --output")
487 try:
488 if args.command == "validate":
489 report = validate_snapshot(read_json_file(args.snapshot))
490 emit_json(report, args.output, force=args.force)
491 return 0 if report["valid"] else 1
492 report = diff_snapshots(
493 read_json_file(args.before),
494 read_json_file(args.after),
495 include_volatile=args.include_volatile,
496 )
497 emit_json(report, args.output, force=args.force)
498 return 0
499 except ResourceToolError as exc:
500 return cli_error(parser, exc)
503if __name__ == "__main__":
504 raise SystemExit(main())