Skills › Scientific Agent Skills › get-available-resources › detect_resources.py
detect_resources.py
1888 lines · skills/get-available-resources/scripts/detect_resources.py at 5f2bc03
1#!/usr/bin/env python3
2"""Collect a conservative, privacy-preserving resource snapshot.
4The script performs only bounded, read-only probes. It never stress-tests the
5machine, allocates a large buffer, changes affinity, or changes accelerator
6state. JSON is written to stdout unless a private local filename is explicitly
7requested.
8"""
10from __future__ import annotations
12import argparse
13import csv
14import datetime as dt
15import importlib
16import io
17import json
18import math
19import os
20import platform
21import re
22import shutil
23import subprocess
24import threading
25from pathlib import Path, PurePosixPath
26from typing import Any, Callable, Mapping, Sequence
28from _common import (
29 MAX_BYTES,
30 SCHEMA_VERSION,
31 ResourceToolError,
32 cli_error,
33 emit_json,
34)
36MAX_PROBE_OUTPUT = 65_536
37MAX_KERNEL_TEXT = 524_288
38MAX_CGROUP_LEVELS = 64
39MAX_CPU_ID = 1_048_575
40MAX_DEVICES = 256
41MIB = 1024**2
42GIB = 1024**3
43_AUTO = object()
45NVIDIA_QUERY = (
46 "nvidia-smi",
47 "--query-gpu=index,name,memory.total,memory.free,driver_version,compute_cap",
48 "--format=csv,noheader,nounits",
49)
50NVIDIA_QUERY_FALLBACK = (
51 "nvidia-smi",
52 "--query-gpu=index,name,memory.total,memory.free,driver_version",
53 "--format=csv,noheader,nounits",
54)
55AMD_SMI_QUERY = ("amd-smi", "static", "--asic", "--vram", "--json")
56ROCM_SMI_QUERY = (
57 "rocm-smi",
58 "--showproductname",
59 "--showmeminfo",
60 "vram",
61 "--json",
62)
63APPLE_DISPLAY_QUERY = ("system_profiler", "SPDisplaysDataType", "-json")
64APPLE_SYSCTL_QUERY = (
65 "sysctl",
66 "-n",
67 "hw.logicalcpu",
68 "hw.physicalcpu",
69 "hw.memsize",
70 "machdep.cpu.brand_string",
71)
73ACCELERATOR_ENV_KEYS = (
74 "CUDA_VISIBLE_DEVICES",
75 "HIP_VISIBLE_DEVICES",
76 "NVIDIA_VISIBLE_DEVICES",
77 "ROCR_VISIBLE_DEVICES",
78)
79SLURM_ENV_KEYS = (
80 "SLURM_JOB_ID",
81 "SLURM_JOBID",
82 "SLURM_CPUS_ON_NODE",
83 "SLURM_CPUS_PER_TASK",
84 "SLURM_JOB_CPUS_PER_NODE",
85 "SLURM_MEM_PER_CPU",
86 "SLURM_MEM_PER_NODE",
87 "SLURM_NTASKS",
88 "SLURM_NTASKS_PER_NODE",
89 "SLURM_TASKS_PER_NODE",
90 "SLURM_NODEID",
91 "SLURM_GPUS",
92 "SLURM_GPUS_ON_NODE",
93 "SLURM_GPUS_PER_TASK",
94 "SLURM_JOB_GPUS",
95 "SLURM_STEP_GPUS",
96)
99def _warning(
100 warnings: list[dict[str, str]],
101 code: str,
102 component: str,
103 message: str,
104 *,
105 severity: str = "warning",
106) -> None:
107 warnings.append(
108 {
109 "code": code,
110 "component": component,
111 "message": message,
112 "severity": severity,
113 }
114 )
117def _provenance(
118 records: list[dict[str, str]],
119 component: str,
120 source: str,
121 status: str,
122) -> None:
123 records.append(
124 {"component": component, "source": source, "status": status}
125 )
128def _safe_text(value: Any, maximum: int = 128) -> str | None:
129 if not isinstance(value, (str, int, float)):
130 return None
131 rendered = " ".join(str(value).split())
132 rendered = "".join(character for character in rendered if character.isprintable())
133 return rendered[:maximum] or None
136def _read_bounded_text(path: Path, maximum: int = MAX_KERNEL_TEXT) -> str:
137 with path.open("rb") as stream:
138 payload = stream.read(maximum + 1)
139 if len(payload) > maximum:
140 raise ValueError("bounded text input exceeded")
141 return payload.decode("utf-8", errors="replace")
144def _try_read(
145 read_text: Callable[[Path], str],
146 path: Path,
147) -> str | None:
148 try:
149 return read_text(path)
150 except (OSError, ValueError, UnicodeError):
151 return None
154def _bounded_nonnegative(value: str | int | None) -> int | None:
155 if isinstance(value, bool) or value is None:
156 return None
157 try:
158 parsed = int(str(value).strip(), 10)
159 except (TypeError, ValueError):
160 return None
161 return parsed if 0 <= parsed <= MAX_BYTES else None
164def parse_cpu_list(value: str) -> int | None:
165 """Count CPUs in a Linux range list without expanding it."""
166 if not isinstance(value, str) or not value.strip() or len(value) > 4096:
167 return None
168 intervals: list[tuple[int, int]] = []
169 for token in value.strip().split(","):
170 token = token.strip()
171 match = re.fullmatch(r"(\d+)(?:-(\d+))?", token)
172 if not match:
173 return None
174 start = int(match.group(1))
175 end = int(match.group(2) or start)
176 if start > end or end > MAX_CPU_ID:
177 return None
178 intervals.append((start, end))
179 intervals.sort()
180 total = 0
181 current_start, current_end = intervals[0]
182 for start, end in intervals[1:]:
183 if start <= current_end + 1:
184 current_end = max(current_end, end)
185 else:
186 total += current_end - current_start + 1
187 current_start, current_end = start, end
188 total += current_end - current_start + 1
189 return total if total <= MAX_CPU_ID + 1 else None
192def _parse_cpu_max(value: str) -> float | None:
193 fields = value.strip().split()
194 if len(fields) != 2 or fields[0] == "max":
195 return None
196 quota = _bounded_nonnegative(fields[0])
197 period = _bounded_nonnegative(fields[1])
198 if quota is None or period in (None, 0) or quota == 0:
199 return None
200 capacity = quota / period
201 if not math.isfinite(capacity) or capacity <= 0:
202 return None
203 return round(capacity, 6)
206def _run_bounded_command(
207 argv: Sequence[str],
208 *,
209 timeout: float = 4.0,
210 maximum: int = MAX_PROBE_OUTPUT,
211) -> dict[str, Any]:
212 """Run a fixed argv with bounded stdout/stderr and no shell."""
213 if not isinstance(argv, tuple) or not argv or not all(
214 isinstance(item, str) and item for item in argv
215 ):
216 raise ValueError("internal command must be a fixed nonempty tuple")
217 try:
218 process = subprocess.Popen(
219 list(argv),
220 stdin=subprocess.DEVNULL,
221 stdout=subprocess.PIPE,
222 stderr=subprocess.PIPE,
223 shell=False,
224 )
225 except FileNotFoundError:
226 return {"status": "not_found", "stdout": "", "stderr": ""}
227 except OSError:
228 return {"status": "start_error", "stdout": "", "stderr": ""}
230 buffers = {"stdout": bytearray(), "stderr": bytearray()}
231 truncated = threading.Event()
233 def drain(name: str, stream: Any) -> None:
234 try:
235 while True:
236 chunk = stream.read(4096)
237 if not chunk:
238 break
239 remaining = maximum - len(buffers[name])
240 if remaining > 0:
241 buffers[name].extend(chunk[:remaining])
242 if len(chunk) > remaining:
243 truncated.set()
244 try:
245 process.kill()
246 except OSError:
247 pass
248 finally:
249 try:
250 stream.close()
251 except OSError:
252 pass
254 threads = [
255 threading.Thread(
256 target=drain,
257 args=("stdout", process.stdout),
258 daemon=True,
259 ),
260 threading.Thread(
261 target=drain,
262 args=("stderr", process.stderr),
263 daemon=True,
264 ),
265 ]
266 for thread in threads:
267 thread.start()
269 timed_out = False
270 try:
271 return_code = process.wait(timeout=timeout)
272 except subprocess.TimeoutExpired:
273 timed_out = True
274 try:
275 process.kill()
276 except OSError:
277 pass
278 return_code = process.wait()
279 for thread in threads:
280 thread.join(timeout=1.0)
282 status = "ok" if return_code == 0 else "error"
283 if timed_out:
284 status = "timeout"
285 elif truncated.is_set():
286 status = "truncated"
287 return {
288 "returncode": return_code,
289 "status": status,
290 "stderr": buffers["stderr"].decode("utf-8", errors="replace"),
291 "stdout": buffers["stdout"].decode("utf-8", errors="replace"),
292 }
295def _load_psutil(
296 warnings: list[dict[str, str]],
297 provenance: list[dict[str, str]],
298) -> Any | None:
299 try:
300 module = importlib.import_module("psutil")
301 except (ImportError, OSError):
302 _warning(
303 warnings,
304 "PSUTIL_UNAVAILABLE",
305 "inventory",
306 "Optional psutil is unavailable; standard-library fallbacks were used.",
307 severity="info",
308 )
309 _provenance(provenance, "inventory.psutil", "optional_import", "unavailable")
310 return None
311 _provenance(provenance, "inventory.psutil", "optional_import", "ok")
312 return module
315def _linux_physical_cores(
316 read_text: Callable[[Path], str],
317) -> int | None:
318 payload = _try_read(read_text, Path("/proc/cpuinfo"))
319 if not payload:
320 return None
321 pairs: set[tuple[str, str]] = set()
322 for block in payload.split("\n\n"):
323 fields: dict[str, str] = {}
324 for line in block.splitlines():
325 if ":" in line:
326 key, value = line.split(":", 1)
327 fields[key.strip().lower()] = value.strip()
328 if "physical id" in fields and "core id" in fields:
329 pairs.add((fields["physical id"], fields["core id"]))
330 count = len(pairs)
331 return count if 0 < count <= MAX_CPU_ID + 1 else None
334def _mac_sysctl_values(
335 run_command: Callable[..., dict[str, Any]],
336 warnings: list[dict[str, str]],
337 provenance: list[dict[str, str]],
338) -> dict[str, Any]:
339 result = run_command(APPLE_SYSCTL_QUERY, timeout=2.0)
340 if result["status"] != "ok":
341 _warning(
342 warnings,
343 "APPLE_SYSCTL_UNAVAILABLE",
344 "platform",
345 "Apple sysctl inventory was unavailable; other sources were retained.",
346 severity="info",
347 )
348 _provenance(provenance, "platform.apple_sysctl", "sysctl", result["status"])
349 return {}
350 lines = result["stdout"].splitlines()
351 if len(lines) < 4:
352 _warning(
353 warnings,
354 "APPLE_SYSCTL_PARSE_FAILED",
355 "platform",
356 "Apple sysctl returned an unexpected bounded response.",
357 )
358 _provenance(provenance, "platform.apple_sysctl", "sysctl", "parse_error")
359 return {}
360 _provenance(provenance, "platform.apple_sysctl", "sysctl", "ok")
361 return {
362 "logical": _bounded_nonnegative(lines[0]),
363 "physical": _bounded_nonnegative(lines[1]),
364 "memory": _bounded_nonnegative(lines[2]),
365 "brand": _safe_text(lines[3]),
366 }
369def _detect_cpu_inventory(
370 *,
371 system: str,
372 psutil_module: Any | None,
373 mac_sysctl: Mapping[str, Any],
374 read_text: Callable[[Path], str],
375 warnings: list[dict[str, str]],
376 provenance: list[dict[str, str]],
377) -> dict[str, int | None]:
378 host_logical = None
379 logical_source = "os.cpu_count"
380 if psutil_module is not None:
381 try:
382 host_logical = psutil_module.cpu_count(logical=True)
383 logical_source = "psutil.cpu_count"
384 except (AttributeError, OSError, RuntimeError, ValueError,
385 getattr(psutil_module, "Error", OSError)):
386 pass
387 if host_logical is None and system == "Darwin":
388 host_logical = mac_sysctl.get("logical")
389 logical_source = "sysctl"
390 if host_logical is None:
391 try:
392 candidate = os.sysconf("SC_NPROCESSORS_ONLN")
393 if candidate > 0:
394 host_logical = candidate
395 logical_source = "os.sysconf"
396 except (AttributeError, OSError, ValueError):
397 pass
398 if host_logical is None:
399 host_logical = os.cpu_count()
400 logical_source = "os.cpu_count"
401 _warning(warnings, "PYTHON_CPU_COUNT_FALLBACK", "cpu",
402 "Host CPU fallback can reflect Python's configured CPU-count override.",
403 severity="info")
404 if (
405 not isinstance(host_logical, int)
406 or isinstance(host_logical, bool)
407 or not 1 <= host_logical <= MAX_CPU_ID + 1
408 ):
409 host_logical = None
410 _warning(
411 warnings,
412 "HOST_LOGICAL_CPU_UNKNOWN",
413 "cpu",
414 "Host logical CPU count could not be determined.",
415 )
416 _provenance(provenance, "cpu.host.logical", logical_source, "unavailable")
417 else:
418 _provenance(provenance, "cpu.host.logical", logical_source, "ok")
420 host_physical: int | None = None
421 if psutil_module is not None:
422 try:
423 candidate = psutil_module.cpu_count(logical=False)
424 if (
425 isinstance(candidate, int)
426 and not isinstance(candidate, bool)
427 and 0 < candidate <= MAX_CPU_ID + 1
428 ):
429 host_physical = candidate
430 _provenance(
431 provenance, "cpu.host.physical", "psutil.cpu_count", "ok"
432 )
433 except (AttributeError, OSError, RuntimeError, ValueError,
434 getattr(psutil_module, "Error", OSError)):
435 pass
436 if host_physical is None and system == "Darwin":
437 candidate = mac_sysctl.get("physical")
438 if isinstance(candidate, int) and 0 < candidate <= MAX_CPU_ID + 1:
439 host_physical = candidate
440 _provenance(provenance, "cpu.host.physical", "sysctl", "ok")
441 if host_physical is None and system == "Linux":
442 host_physical = _linux_physical_cores(read_text)
443 if host_physical is not None:
444 _provenance(provenance, "cpu.host.physical", "proc_cpuinfo", "ok")
445 if host_physical is None:
446 _warning(
447 warnings,
448 "HOST_PHYSICAL_CPU_UNKNOWN",
449 "cpu",
450 "Physical core count is unknown and was not inferred from logical CPUs.",
451 severity="info",
452 )
453 _provenance(provenance, "cpu.host.physical", "platform_fallbacks", "unavailable")
454 return {"logical": host_logical, "physical": host_physical}
457def _detect_process_cpu_count(
458 *,
459 psutil_module: Any | None,
460 warnings: list[dict[str, str]],
461 provenance: list[dict[str, str]],
462) -> dict[str, int | None]:
463 affinity_count: int | None = None
464 affinity_source: str | None = None
465 get_affinity = getattr(os, "sched_getaffinity", None)
466 if callable(get_affinity):
467 try:
468 affinity = get_affinity(0)
469 if affinity and len(affinity) <= MAX_CPU_ID + 1:
470 affinity_count = len(affinity)
471 affinity_source = "os.sched_getaffinity"
472 except (OSError, TypeError, ValueError):
473 pass
474 if affinity_count is None and psutil_module is not None:
475 try:
476 affinity = psutil_module.Process().cpu_affinity()
477 if affinity and len(affinity) <= MAX_CPU_ID + 1:
478 affinity_count = len(affinity)
479 affinity_source = "psutil.Process.cpu_affinity"
480 except (AttributeError, OSError, RuntimeError, ValueError,
481 getattr(psutil_module, "Error", OSError)):
482 pass
483 if affinity_count is not None:
484 _provenance(provenance, "cpu.process.affinity_logical", affinity_source or "", "ok")
485 else:
486 _provenance(
487 provenance, "cpu.process.affinity_logical", "affinity_apis", "unavailable"
488 )
490 process_count: int | None = None
491 process_cpu_count = getattr(os, "process_cpu_count", None)
492 if callable(process_cpu_count):
493 try:
494 candidate = process_cpu_count()
495 if (
496 isinstance(candidate, int)
497 and not isinstance(candidate, bool)
498 and 0 < candidate <= MAX_CPU_ID + 1
499 ):
500 process_count = candidate
501 _provenance(
502 provenance,
503 "cpu.process.python_available_logical",
504 "os.process_cpu_count",
505 "ok",
506 )
507 except (OSError, ValueError):
508 pass
509 if affinity_count is None and process_count is None:
510 _warning(
511 warnings,
512 "PROCESS_CPU_SCOPE_UNKNOWN",
513 "cpu",
514 "No process-aware CPU count API was available.",
515 severity="info",
516 )
517 return {
518 "affinity_logical": affinity_count,
519 "python_available_logical": process_count,
520 }
523def _cgroup_relative_path(payload: str) -> tuple[str, ...] | None:
524 for line in payload.splitlines():
525 fields = line.split(":", 2)
526 if len(fields) == 3 and fields[0] == "0" and fields[1] == "":
527 raw = PurePosixPath(fields[2])
528 if not raw.is_absolute():
529 return None
530 parts = tuple(part for part in raw.parts if part not in {"/", ""})
531 if any(part in {".", ".."} for part in parts):
532 return None
533 return parts
534 return None
537def _resolve_cgroup_mount(
538 payload: str, parts: tuple[str, ...]
539) -> tuple[Path, tuple[str, ...]] | None:
540 """Map namespace membership through a cgroup2 mount's root and mountpoint."""
541 matches = []
542 for line in payload.splitlines():
543 fields = line.split()
544 try:
545 separator = fields.index("-")
546 except ValueError:
547 continue
548 if len(fields) <= separator + 1 or fields[separator + 1] != "cgroup2" or separator < 6:
549 continue
550 decoded = [re.sub(r"\\([0-7]{3})", lambda m: chr(int(m[1], 8)), field)
551 for field in fields[3:5]]
552 mount_root, mountpoint = map(PurePosixPath, decoded)
553 if any(not p.is_absolute() or ".." in p.parts for p in (mount_root, mountpoint)):
554 continue
555 prefix = mount_root.parts[1:]
556 if parts[:len(prefix)] == prefix:
557 matches.append((len(prefix), Path(str(mountpoint)), parts[len(prefix):]))
558 if not matches:
559 return None
560 # A broader visible mount retains ancestors that a nested bind mount hides.
561 _, mountpoint, relative = min(matches, key=lambda item: item[0])
562 return mountpoint, relative
565def detect_cgroup_v2(
566 *,
567 read_text: Callable[[Path], str] = _read_bounded_text,
568 root: Path = Path("/sys/fs/cgroup"),
569 proc_self_cgroup: Path = Path("/proc/self/cgroup"),
570 proc_self_mountinfo: Path = Path("/proc/self/mountinfo"),
571 warnings: list[dict[str, str]] | None = None,
572 provenance: list[dict[str, str]] | None = None,
573) -> dict[str, Any]:
574 """Read current and ancestor cgroup v2 limits without exposing paths."""
575 warning_records = warnings if warnings is not None else []
576 provenance_records = provenance if provenance is not None else []
577 membership = _try_read(read_text, proc_self_cgroup)
578 parts = _cgroup_relative_path(membership or "")
579 mountinfo = _try_read(read_text, proc_self_mountinfo)
580 if parts is not None and mountinfo is not None:
581 resolved = _resolve_cgroup_mount(mountinfo, parts)
582 if resolved is not None:
583 root, parts = resolved
584 else:
585 parts = None
586 empty = {
587 "detected": False, "cpu_quota_cores": None, "cpuset_logical": None,
588 "memory_available_bytes": None, "memory_current_bytes": None,
589 "memory_high_bytes": None, "memory_max_bytes": None, "scope": "unknown",
590 }
591 if _try_read(read_text, root / "cgroup.controllers") is None:
592 _provenance(
593 provenance_records, "cgroup_v2", "cgroup.controllers", "unavailable"
594 )
595 _warning(warning_records, "CGROUP_V2_UNAVAILABLE", "cgroup",
596 "cgroup v2 limits are unavailable; cgroup v1 and hidden limits are not measured.")
597 return empty
599 if parts is None:
600 _warning(
601 warning_records,
602 "CGROUP_MEMBERSHIP_UNKNOWN",
603 "cgroup",
604 "cgroup v2 was detected but current membership could not be parsed.",
605 )
606 return {**empty, "detected": True}
607 current = root.joinpath(*parts)
608 if _try_read(read_text, current / "cgroup.controllers") is None:
609 _warning(warning_records, "CGROUP_CURRENT_UNAVAILABLE", "cgroup",
610 "Current cgroup could not be read; visible root limits were not substituted.")
611 return {**empty, "detected": True}
612 chain: list[Path] = []
613 cursor = current
614 for _ in range(MAX_CGROUP_LEVELS):
615 chain.append(cursor)
616 if cursor == root:
617 break
618 parent = cursor.parent
619 if parent == cursor or (parent != root and root not in parent.parents):
620 break
621 cursor = parent
622 else:
623 _warning(
624 warning_records,
625 "CGROUP_DEPTH_BOUNDED",
626 "cgroup",
627 "cgroup ancestor traversal reached its safety bound.",
628 )
630 quota_candidates: list[float] = []
631 memory_max_candidates: list[int] = []
632 memory_high_candidates: list[int] = []
633 memory_available_candidates: list[int] = []
634 current_memory: int | None = None
636 for index, directory in enumerate(chain):
637 cpu_max_text = _try_read(read_text, directory / "cpu.max")
638 if cpu_max_text is not None:
639 quota = _parse_cpu_max(cpu_max_text)
640 if quota is not None:
641 quota_candidates.append(quota)
642 elif not re.fullmatch(r"max\s+[1-9][0-9]*", cpu_max_text.strip()):
643 _warning(warning_records, "CGROUP_CPU_MAX_INVALID", "cgroup",
644 "A cgroup CPU quota could not be interpreted.")
646 memory_max_text = _try_read(read_text, directory / "memory.max")
647 memory_current_text = _try_read(read_text, directory / "memory.current")
648 memory_high_text = _try_read(read_text, directory / "memory.high")
649 memory_current = _bounded_nonnegative(memory_current_text)
650 if index == 0:
651 current_memory = memory_current
652 memory_max = (
653 None
654 if memory_max_text is None or memory_max_text.strip() == "max"
655 else _bounded_nonnegative(memory_max_text)
656 )
657 memory_high = (
658 None
659 if memory_high_text is None or memory_high_text.strip() == "max"
660 else _bounded_nonnegative(memory_high_text)
661 )
662 if memory_max is not None:
663 memory_max_candidates.append(memory_max)
664 if memory_current is not None:
665 memory_available_candidates.append(max(0, memory_max - memory_current))
666 if memory_high is not None:
667 memory_high_candidates.append(memory_high)
668 for text, parsed in ((memory_max_text, memory_max), (memory_high_text, memory_high)):
669 if text is not None and text.strip() != "max" and parsed is None:
670 _warning(warning_records, "CGROUP_MEMORY_LIMIT_INVALID", "cgroup",
671 "A cgroup memory boundary could not be interpreted.")
672 if memory_max is not None and memory_current is None:
673 _warning(warning_records, "CGROUP_MEMORY_USAGE_UNKNOWN", "cgroup",
674 "A finite memory limit has unknown usage; remaining memory is only an upper bound.")
676 cpuset_text = _try_read(read_text, current / "cpuset.cpus.effective")
677 cpuset_count = parse_cpu_list(cpuset_text) if cpuset_text else None
678 if cpuset_text and cpuset_count is None:
679 _warning(
680 warning_records,
681 "CGROUP_CPUSET_PARSE_FAILED",
682 "cgroup",
683 "cpuset.cpus.effective had an unexpected bounded value.",
684 )
686 _provenance(provenance_records, "cgroup_v2", "procfs_and_cgroupfs", "ok")
687 return {
688 "detected": True,
689 "cpu_quota_cores": min(quota_candidates) if quota_candidates else None,
690 "cpuset_logical": cpuset_count,
691 "memory_available_bytes": (
692 min(memory_available_candidates)
693 if memory_available_candidates
694 else None
695 ),
696 "memory_current_bytes": current_memory,
697 "memory_high_bytes": (
698 min(memory_high_candidates) if memory_high_candidates else None
699 ),
700 "memory_max_bytes": (
701 min(memory_max_candidates) if memory_max_candidates else None
702 ),
703 "scope": "non_root" if parts else "root",
704 }
707def _parse_slurm_count(value: str | None) -> int | None:
708 if value is None or not 0 < len(value) <= 128:
709 return None
710 stripped = value.strip()
711 if stripped.isdigit():
712 count = int(stripped)
713 return count if 0 <= count <= 1_000_000 else None
714 match = re.fullmatch(r"(?:[A-Za-z0-9_.+-]+:)+(\d+)", stripped)
715 if match:
716 count = int(match.group(1))
717 return count if count <= 1_000_000 else None
718 return None
721def _parse_slurm_memory(value: str | None) -> int | None:
722 if value is None or not 0 < len(value) <= 128:
723 return None
724 match = re.fullmatch(r"\s*(\d+)\s*([KkMmGgTt]?)\s*", value)
725 if not match:
726 return None
727 amount = int(match.group(1))
728 if amount == 0:
729 return None
730 suffix = match.group(2).upper()
731 multiplier = {"K": 1024, "M": MIB, "G": GIB, "T": 1024**4, "": MIB}[suffix]
732 result = amount * multiplier
733 return result if result <= MAX_BYTES else None
736def _parse_first_repeated_count(value: str | None) -> int | None:
737 return _parse_node_count(value, 0)
740def _parse_node_count(value: str | None, node_index: int | None) -> int | None:
741 """Resolve compressed Slurm node counts without expanding the list."""
742 if value is None or not value or len(value) > 4096:
743 return None
744 groups = []
745 for token in value.split(","):
746 match = re.fullmatch(r"\s*(\d+)(?:\(x(\d+)\))?\s*", token)
747 if not match:
748 return None
749 count, repeats = int(match.group(1)), int(match.group(2) or 1)
750 if count > 1_000_000 or not 1 <= repeats <= 1_000_000:
751 return None
752 groups.append((count, repeats))
753 if node_index is None:
754 counts = {count for count, _ in groups}
755 return next(iter(counts)) if len(counts) == 1 else None
756 for count, repeats in groups:
757 if node_index < repeats:
758 return count
759 node_index -= repeats
760 return None
763def detect_scheduler(
764 environ: Mapping[str, str],
765 *,
766 warnings: list[dict[str, str]] | None = None,
767 provenance: list[dict[str, str]] | None = None,
768) -> dict[str, Any]:
769 """Interpret only an allowlist of Slurm allocation variables."""
770 warning_records = warnings if warnings is not None else []
771 provenance_records = provenance if provenance is not None else []
772 present = {
773 key: environ[key]
774 for key in SLURM_ENV_KEYS
775 if key in environ and isinstance(environ[key], str)
776 }
777 detected = bool(present)
778 if not detected:
779 _provenance(provenance_records, "scheduler", "named_slurm_environment", "absent")
780 return {
781 "allocation": {
782 "cpu_per_process": None,
783 "cpus_on_node": None,
784 "gpus_per_process": None,
785 "gpus_on_node": None,
786 "memory_effective_bytes": None,
787 "memory_scope": None,
788 "tasks": None,
789 },
790 "detected": False,
791 "enforcement": "not_applicable",
792 "fields_read": [],
793 "kind": None,
794 }
796 cpus_per_task = _parse_slurm_count(present.get("SLURM_CPUS_PER_TASK"))
797 cpus_on_node = _parse_slurm_count(present.get("SLURM_CPUS_ON_NODE"))
798 tasks = _parse_slurm_count(present.get("SLURM_NTASKS"))
799 node_index = _parse_slurm_count(present.get("SLURM_NODEID"))
800 tasks_per_node = _parse_node_count(present.get("SLURM_TASKS_PER_NODE"), node_index)
801 if "SLURM_TASKS_PER_NODE" not in present:
802 tasks_per_node = _parse_slurm_count(present.get("SLURM_NTASKS_PER_NODE"))
803 first_job_cpus = _parse_first_repeated_count(
804 present.get("SLURM_JOB_CPUS_PER_NODE")
805 )
806 cpu_per_process = cpus_per_task
807 if cpu_per_process is None and cpus_on_node is not None and tasks_per_node == 1:
808 cpu_per_process = cpus_on_node
810 memory_per_cpu = _parse_slurm_memory(present.get("SLURM_MEM_PER_CPU"))
811 memory_per_node = _parse_slurm_memory(present.get("SLURM_MEM_PER_NODE"))
812 memory_effective: int | None = None
813 memory_scope: str | None = None
814 if memory_per_cpu is not None and cpu_per_process is not None:
815 memory_effective = min(MAX_BYTES, memory_per_cpu * cpu_per_process)
816 memory_scope = "per_task_from_per_cpu"
817 elif memory_per_node is not None:
818 memory_effective = memory_per_node
819 memory_scope = "shared_per_node"
820 if tasks_per_node is None or tasks_per_node > 1:
821 _warning(
822 warning_records,
823 "SLURM_MEMORY_SHARED",
824 "scheduler",
825 "Slurm per-node memory is shared; it is only an upper bound for this process.",
826 severity="info",
827 )
829 gpus_per_task = _parse_slurm_count(present.get("SLURM_GPUS_PER_TASK"))
830 gpus_on_node = _parse_slurm_count(present.get("SLURM_GPUS_ON_NODE"))
831 requested_gpus = _parse_slurm_count(present.get("SLURM_GPUS"))
833 _warning(
834 warning_records,
835 "SLURM_ENFORCEMENT_UNKNOWN",
836 "scheduler",
837 "Allocation variables do not prove task affinity or cgroup enforcement.",
838 severity="info",
839 )
840 if cpus_per_task is None and cpus_on_node is not None and tasks_per_node != 1:
841 _warning(
842 warning_records,
843 "SLURM_CPU_SCOPE_SHARED",
844 "scheduler",
845 "Node CPU allocation was not treated as a per-process limit.",
846 severity="info",
847 )
848 _provenance(provenance_records, "scheduler", "named_slurm_environment", "ok")
849 return {
850 "allocation": {
851 "cpu_per_process": cpu_per_process,
852 "cpus_on_node": cpus_on_node,
853 "first_job_cpus_per_node": first_job_cpus,
854 "gpus_per_process": gpus_per_task,
855 "gpus_on_node": gpus_on_node,
856 "gpus_requested": requested_gpus,
857 "memory_effective_bytes": memory_effective,
858 "memory_scope": memory_scope,
859 "tasks": tasks,
860 "tasks_per_node": tasks_per_node,
861 },
862 "detected": True,
863 "enforcement": "unknown",
864 "fields_read": sorted(present),
865 "kind": "slurm",
866 }
869def _effective_cpu(
870 host: Mapping[str, int | None],
871 process: Mapping[str, int | None],
872 cgroup: Mapping[str, Any],
873 scheduler: Mapping[str, Any],
874) -> dict[str, Any]:
875 candidates: list[tuple[str, float]] = []
876 for source, value in (
877 ("host_logical", host.get("logical")),
878 ("process_affinity", process.get("affinity_logical")),
879 ("python_process_count", process.get("python_available_logical")),
880 ("cgroup_cpuset", cgroup.get("cpuset_logical")),
881 ("cgroup_quota", cgroup.get("cpu_quota_cores")),
882 (
883 "scheduler_per_process",
884 scheduler.get("allocation", {}).get("cpu_per_process"),
885 ),
886 ):
887 if isinstance(value, (int, float)) and not isinstance(value, bool) and value > 0:
888 candidates.append((source, float(value)))
889 if not candidates:
890 return {
891 "capacity_cores": None,
892 "limiting_sources": [],
893 "worker_ceiling": 1,
894 }
895 capacity = min(value for _, value in candidates)
896 limiting = sorted(
897 source for source, value in candidates if abs(value - capacity) < 1e-9
898 )
899 return {
900 "capacity_cores": round(capacity, 6),
901 "limiting_sources": limiting,
902 "worker_ceiling": max(1, min(1024, math.floor(capacity))),
903 }
906def _parse_linux_meminfo(payload: str) -> dict[str, int | None]:
907 values: dict[str, int] = {}
908 for line in payload.splitlines():
909 match = re.fullmatch(r"([A-Za-z_()]+):\s+(\d+)\s+kB", line.strip())
910 if match:
911 values[match.group(1)] = int(match.group(2)) * 1024
912 return {
913 "available": values.get("MemAvailable"),
914 "swap_free": values.get("SwapFree"),
915 "swap_total": values.get("SwapTotal"),
916 "total": values.get("MemTotal"),
917 }
920def _detect_memory(
921 *,
922 system: str,
923 machine: str,
924 psutil_module: Any | None,
925 mac_sysctl: Mapping[str, Any],
926 cgroup: Mapping[str, Any],
927 scheduler: Mapping[str, Any],
928 read_text: Callable[[Path], str],
929 warnings: list[dict[str, str]],
930 provenance: list[dict[str, str]],
931) -> dict[str, Any]:
932 host_total: int | None = None
933 host_available: int | None = None
934 swap_total: int | None = None
935 swap_free: int | None = None
936 if psutil_module is not None:
937 try:
938 memory = psutil_module.virtual_memory()
939 host_total = _bounded_nonnegative(memory.total)
940 host_available = _bounded_nonnegative(memory.available)
941 _provenance(provenance, "memory.host", "psutil.virtual_memory", "ok")
942 except (AttributeError, OSError, RuntimeError, ValueError,
943 getattr(psutil_module, "Error", OSError)):
944 pass
945 try:
946 swap = psutil_module.swap_memory()
947 swap_total = _bounded_nonnegative(swap.total)
948 swap_free = _bounded_nonnegative(swap.free)
949 _provenance(provenance, "memory.swap", "psutil.swap_memory", "ok")
950 except (AttributeError, OSError, RuntimeError, ValueError,
951 getattr(psutil_module, "Error", OSError)):
952 pass
954 if system == "Linux" and (host_total is None or host_available is None):
955 payload = _try_read(read_text, Path("/proc/meminfo"))
956 if payload:
957 parsed = _parse_linux_meminfo(payload)
958 host_total = host_total or parsed["total"]
959 host_available = (
960 host_available
961 if host_available is not None
962 else parsed["available"]
963 )
964 swap_total = swap_total if swap_total is not None else parsed["swap_total"]
965 swap_free = swap_free if swap_free is not None else parsed["swap_free"]
966 _provenance(provenance, "memory.host", "proc_meminfo", "ok")
967 if system == "Darwin" and host_total is None:
968 candidate = mac_sysctl.get("memory")
969 if isinstance(candidate, int) and candidate > 0:
970 host_total = candidate
971 _provenance(provenance, "memory.host.total", "sysctl", "ok")
972 if host_total is None:
973 try:
974 pages = os.sysconf("SC_PHYS_PAGES")
975 page_size = os.sysconf("SC_PAGE_SIZE")
976 candidate = pages * page_size
977 if 0 < candidate <= MAX_BYTES:
978 host_total = candidate
979 _provenance(provenance, "memory.host.total", "os.sysconf", "ok")
980 except (AttributeError, OSError, TypeError, ValueError):
981 pass
983 if host_total is None:
984 _warning(
985 warnings,
986 "HOST_MEMORY_TOTAL_UNKNOWN",
987 "memory",
988 "Host memory total could not be determined.",
989 )
990 if host_available is None:
991 _warning(
992 warnings,
993 "HOST_MEMORY_AVAILABLE_UNKNOWN",
994 "memory",
995 "Host available memory could not be determined without an allocation probe.",
996 severity="info",
997 )
999 limit_candidates: list[tuple[str, int]] = []
1000 available_candidates: list[tuple[str, int]] = []
1001 for source, value in (
1002 ("host_total", host_total),
1003 ("cgroup_memory_max", cgroup.get("memory_max_bytes")),
1004 (
1005 "scheduler_memory",
1006 scheduler.get("allocation", {}).get("memory_effective_bytes"),
1007 ),
1008 ):
1009 if isinstance(value, int) and value >= 0:
1010 limit_candidates.append((source, value))
1011 for source, value in (
1012 ("host_available", host_available),
1013 ("cgroup_memory_remaining", cgroup.get("memory_available_bytes")),
1014 (
1015 "scheduler_memory_upper_bound",
1016 scheduler.get("allocation", {}).get("memory_effective_bytes"),
1017 ),
1018 ):
1019 if isinstance(value, int) and value >= 0:
1020 available_candidates.append((source, value))
1022 effective_limit = (
1023 min(value for _, value in limit_candidates) if limit_candidates else None
1024 )
1025 effective_available = (
1026 min(value for _, value in available_candidates)
1027 if available_candidates
1028 else None
1029 )
1030 if (
1031 effective_limit is not None
1032 and effective_available is not None
1033 and effective_available > effective_limit
1034 ):
1035 effective_available = effective_limit
1036 limiting_sources = (
1037 sorted(
1038 source
1039 for source, value in limit_candidates
1040 if value == effective_limit
1041 )
1042 if effective_limit is not None
1043 else []
1044 )
1045 available_sources = (
1046 sorted(
1047 source
1048 for source, value in available_candidates
1049 if value == effective_available
1050 )
1051 if effective_available is not None
1052 else []
1053 )
1054 unified = system == "Darwin" and (
1055 machine.lower() in {"arm64", "aarch64"}
1056 or str(mac_sysctl.get("brand", "")).startswith("Apple ")
1057 )
1058 return {
1059 "cgroup_v2": {
1060 "available_bytes": cgroup.get("memory_available_bytes"),
1061 "current_bytes": cgroup.get("memory_current_bytes"),
1062 "high_bytes": cgroup.get("memory_high_bytes"),
1063 "max_bytes": cgroup.get("memory_max_bytes"),
1064 },
1065 "effective": {
1066 "available_bytes": effective_available,
1067 "available_limiting_sources": available_sources,
1068 "hard_limit_bytes": effective_limit,
1069 "hard_limit_sources": limiting_sources,
1070 "pressure_threshold_bytes": cgroup.get("memory_high_bytes"),
1071 },
1072 "host": {
1073 "available_bytes": host_available,
1074 "total_bytes": host_total,
1075 },
1076 "model": "unified_cpu_gpu" if unified else "system_ram",
1077 "swap": {"free_bytes": swap_free, "total_bytes": swap_total},
1078 }
1081def _detect_disk(
1082 warnings: list[dict[str, str]],
1083 provenance: list[dict[str, str]],
1084) -> dict[str, Any]:
1085 try:
1086 usage = shutil.disk_usage(Path.cwd())
1087 total = _bounded_nonnegative(usage.total)
1088 free = _bounded_nonnegative(usage.free)
1089 if total is None or free is None or free > total:
1090 raise ValueError("invalid bounded disk values")
1091 _provenance(provenance, "disk", "shutil.disk_usage", "ok")
1092 except (OSError, ValueError):
1093 _warning(
1094 warnings,
1095 "DISK_CAPACITY_UNKNOWN",
1096 "disk",
1097 "Working filesystem capacity could not be read.",
1098 )
1099 _provenance(provenance, "disk", "shutil.disk_usage", "unavailable")
1100 return {
1101 "capacity_bytes": None,
1102 "free_bytes": None,
1103 "scope": "working_filesystem_path_redacted",
1104 "user_available_bytes": None,
1105 "writable": None,
1106 "writability_check": "not_performed",
1107 }
1109 user_available = free
1110 if hasattr(os, "statvfs"):
1111 try:
1112 statvfs = os.statvfs(Path.cwd())
1113 candidate = _bounded_nonnegative(
1114 statvfs.f_bavail * statvfs.f_frsize
1115 )
1116 filesystem_free = _bounded_nonnegative(statvfs.f_bfree * statvfs.f_frsize)
1117 if filesystem_free is not None and filesystem_free <= total:
1118 free = filesystem_free
1119 if candidate is not None:
1120 user_available = candidate
1121 _provenance(
1122 provenance, "disk.user_available", "os.statvfs", "ok"
1123 )
1124 except (AttributeError, OSError, TypeError, ValueError):
1125 pass
1126 writable = os.access(Path.cwd(), os.W_OK)
1127 if not writable:
1128 _warning(
1129 warnings,
1130 "WORKING_DIRECTORY_NOT_WRITABLE",
1131 "disk",
1132 "The working directory failed the non-writing access check.",
1133 )
1134 return {
1135 "capacity_bytes": total,
1136 "free_bytes": free,
1137 "scope": "working_filesystem_path_redacted",
1138 "user_available_bytes": max(0, min(free, user_available)),
1139 "writable": bool(writable),
1140 "writability_check": "os_access_only_no_write_probe",
1141 }
1144def _mib_to_bytes(value: str) -> int | None:
1145 try:
1146 amount = float(value)
1147 except ValueError:
1148 return None
1149 if not math.isfinite(amount) or not 0 <= amount <= MAX_BYTES / MIB:
1150 return None
1151 result = round(amount * MIB)
1152 return result if result <= MAX_BYTES else None
1155def parse_nvidia_csv(payload: str) -> list[dict[str, Any]]:
1156 devices: list[dict[str, Any]] = []
1157 try:
1158 rows = csv.reader(io.StringIO(payload))
1159 for row in rows:
1160 if len(row) not in {5, 6}:
1161 continue
1162 index = _bounded_nonnegative(row[0])
1163 if index is None or index > 4096:
1164 continue
1165 devices.append(
1166 {
1167 "backend_candidate": "cuda",
1168 "compute_capability": (
1169 _safe_text(row[5], 16) if len(row) == 6 else None
1170 ),
1171 "device_class": "gpu",
1172 "device_permission": "not_tested",
1173 "driver_version": _safe_text(row[4], 64),
1174 "local_index": index,
1175 "management_query": "visible",
1176 "memory": {
1177 "dedicated_free_bytes": _mib_to_bytes(row[3].strip()),
1178 "dedicated_total_bytes": _mib_to_bytes(row[2].strip()),
1179 "model": "dedicated",
1180 },
1181 "name": _safe_text(row[1]) or "NVIDIA GPU",
1182 "runtime_compatibility": "not_tested",
1183 "vendor": "nvidia",
1184 }
1185 )
1186 except csv.Error:
1187 return []
1188 return sorted(devices, key=lambda item: item["local_index"])
1191def _recursive_scalar(
1192 value: Any,
1193 accepted_keys: set[str],
1194) -> tuple[str, Any] | None:
1195 if isinstance(value, dict):
1196 for key, child in value.items():
1197 normalized = re.sub(r"[^a-z0-9]+", "_", str(key).lower()).strip("_")
1198 if normalized in accepted_keys and isinstance(child, (str, int, float)):
1199 return normalized, child
1200 for child in value.values():
1201 found = _recursive_scalar(child, accepted_keys)
1202 if found is not None:
1203 return found
1204 elif isinstance(value, list):
1205 for child in value:
1206 found = _recursive_scalar(child, accepted_keys)
1207 if found is not None:
1208 return found
1209 return None
1212def _amd_entries(data: Any) -> list[tuple[int, dict[str, Any]]]:
1213 entries: list[tuple[int, dict[str, Any]]] = []
1214 if isinstance(data, list):
1215 for index, item in enumerate(data):
1216 if isinstance(item, dict) and (
1217 "gpu" in {str(key).lower() for key in item}
1218 or "asic" in {str(key).lower() for key in item}
1219 ):
1220 gpu_value = next(
1221 (
1222 value
1223 for key, value in item.items()
1224 if str(key).lower() == "gpu"
1225 ),
1226 index,
1227 )
1228 gpu_index = _bounded_nonnegative(gpu_value)
1229 entries.append((gpu_index if gpu_index is not None else index, item))
1230 elif isinstance(data, dict):
1231 for key, item in data.items():
1232 normalized = re.sub(
1233 r"[^a-z0-9]+", "_", str(key).lower()
1234 ).strip("_")
1235 if normalized in {"gpu_data", "gpu_devices", "gpus"}:
1236 entries.extend(_amd_entries(item))
1237 for key, item in data.items():
1238 match = re.fullmatch(r"(?:card|gpu)\s*(\d+)", str(key), re.IGNORECASE)
1239 if match and isinstance(item, dict):
1240 entries.append((int(match.group(1)), item))
1241 if not entries and "gpu" in {str(key).lower() for key in data}:
1242 gpu_value = next(
1243 value for key, value in data.items() if str(key).lower() == "gpu"
1244 )
1245 if not isinstance(gpu_value, (dict, list)):
1246 gpu_index = _bounded_nonnegative(gpu_value)
1247 entries.append((gpu_index or 0, data))
1248 deduplicated: dict[int, dict[str, Any]] = {}
1249 for index, entry in entries:
1250 deduplicated.setdefault(index, entry)
1251 return sorted(deduplicated.items())
1254def _amd_memory_bytes(entry: Mapping[str, Any], *, in_vram: bool = False) -> int | None:
1255 for key, value in entry.items():
1256 normalized = re.sub(r"[^a-z0-9]+", "_", str(key).lower()).strip("_")
1257 if (
1258 ("vram" in normalized or in_vram)
1259 and ("total" in normalized or "size" in normalized)
1260 ):
1261 # AMD SMI JSON uses SIZE: {value: ..., unit: MB}; legacy ROCm
1262 # SMI reports a flat "VRAM Total Memory (B)" integer string.
1263 if isinstance(value, dict):
1264 fields = {str(key).lower(): item for key, item in value.items()}
1265 raw = fields.get("value")
1266 unit_context = str(fields.get("unit", "")).lower()
1267 else:
1268 raw = value
1269 unit_context = f"{normalized} {str(value).lower()}"
1270 match = re.fullmatch(r"\s*(\d+(?:\.\d+)?)\s*(?:[A-Za-z]+)?\s*", str(raw))
1271 if not match:
1272 continue
1273 amount = float(match.group(1))
1274 if not math.isfinite(amount):
1275 continue
1276 if "gib" in unit_context or "gb" in unit_context:
1277 multiplier = GIB
1278 elif "mib" in unit_context or "mb" in unit_context:
1279 multiplier = MIB
1280 elif "byte" in unit_context or unit_context == "b" or normalized.endswith("_b"):
1281 multiplier = 1
1282 else:
1283 continue
1284 if amount > MAX_BYTES / multiplier:
1285 continue
1286 result = round(amount * multiplier)
1287 if 0 <= result <= MAX_BYTES:
1288 return result
1289 if isinstance(value, dict):
1290 nested = _amd_memory_bytes(value, in_vram=in_vram or normalized == "vram")
1291 if nested is not None:
1292 return nested
1293 return None
1296def parse_amd_json(payload: str) -> list[dict[str, Any]]:
1297 try:
1298 data = json.loads(payload)
1299 entries = _amd_entries(data)
1300 except (json.JSONDecodeError, RecursionError):
1301 return []
1302 devices: list[dict[str, Any]] = []
1303 name_keys = {
1304 "asic_market_name",
1305 "card_model",
1306 "card_series",
1307 "market_name",
1308 "product_name",
1309 }
1310 for index, entry in entries:
1311 try:
1312 found = _recursive_scalar(entry, name_keys)
1313 total_memory = _amd_memory_bytes(entry)
1314 except RecursionError:
1315 continue
1316 name = _safe_text(found[1]) if found else None
1317 devices.append(
1318 {
1319 "backend_candidate": "rocm",
1320 "compute_capability": None,
1321 "device_class": "gpu",
1322 "device_permission": "not_tested",
1323 "driver_version": None,
1324 "local_index": index,
1325 "management_query": "visible",
1326 "memory": {
1327 "dedicated_free_bytes": None,
1328 "dedicated_total_bytes": None,
1329 "reported_total_bytes": total_memory,
1330 "model": "device_reported_pool",
1331 },
1332 "name": name or "AMD GPU",
1333 "runtime_compatibility": "not_tested",
1334 "vendor": "amd",
1335 }
1336 )
1337 return sorted(devices, key=lambda item: item["local_index"])
1340def parse_apple_profiler_json(
1341 payload: str,
1342 *,
1343 machine: str,
1344) -> list[dict[str, Any]]:
1345 try:
1346 data = json.loads(payload)
1347 except (json.JSONDecodeError, RecursionError):
1348 return []
1349 entries = data.get("SPDisplaysDataType") if isinstance(data, dict) else None
1350 if not isinstance(entries, list):
1351 return []
1352 devices: list[dict[str, Any]] = []
1353 for index, entry in enumerate(entries):
1354 if not isinstance(entry, dict):
1355 continue
1356 name = _safe_text(entry.get("sppci_model") or entry.get("_name"))
1357 vendor_text = f"{entry.get('spdisplays_vendor', '')} {name or ''}".lower()
1358 vendor = next((item for item in ("apple", "amd", "nvidia", "intel")
1359 if item in vendor_text), "unknown")
1360 unified = vendor == "apple" and (
1361 machine.lower() in {"arm64", "aarch64"} or (name or "").startswith("Apple M")
1362 )
1363 core_value = (
1364 entry.get("sppci_cores")
1365 or entry.get("spdisplays_gpu_cores")
1366 or entry.get("_spdisplays_gpu_cores")
1367 )
1368 devices.append(
1369 {
1370 "backend_candidate": "metal",
1371 "compute_capability": None,
1372 "device_class": "integrated_gpu" if unified else "gpu",
1373 "device_permission": "not_tested",
1374 "driver_version": None,
1375 "gpu_cores": _bounded_nonnegative(core_value),
1376 "local_index": index,
1377 "management_query": "visible",
1378 "memory": {
1379 "dedicated_free_bytes": None,
1380 "dedicated_total_bytes": None,
1381 "model": "unified" if unified else "unknown",
1382 },
1383 "name": name or "Apple display accelerator",
1384 "runtime_compatibility": "not_tested",
1385 "vendor": vendor,
1386 }
1387 )
1388 return devices
1391def _visibility_summary(environ: Mapping[str, str]) -> dict[str, Any]:
1392 summary: dict[str, Any] = {}
1393 for key in ACCELERATOR_ENV_KEYS:
1394 if key not in environ:
1395 summary[key] = {
1396 "entry_count": None,
1397 "set": False,
1398 "state": "unset",
1399 "value_redacted": True,
1400 }
1401 continue
1402 raw = environ[key] if isinstance(environ[key], str) else ""
1403 if len(raw) > 4096:
1404 summary[key] = {
1405 "entry_count": None,
1406 "set": True,
1407 "state": "invalid_or_oversized",
1408 "value_redacted": True,
1409 }
1410 continue
1411 stripped = raw.strip()
1412 lowered = stripped.lower()
1413 if lowered in {"", "-1", "none", "void"}:
1414 state, count = "none", 0
1415 elif lowered == "all":
1416 state, count = "all", None
1417 else:
1418 entries = {
1419 item.strip() for item in stripped.split(",") if item.strip()
1420 }
1421 state = "restricted" if entries else "none"
1422 count = min(len(entries), 4096)
1423 summary[key] = {
1424 "entry_count": count,
1425 "set": True,
1426 "state": state,
1427 "value_redacted": True,
1428 }
1429 return summary
1432def _accelerator_upper_bounds(
1433 devices: Sequence[Mapping[str, Any]],
1434 visibility: Mapping[str, Mapping[str, Any]],
1435 scheduler: Mapping[str, Any],
1436) -> dict[str, int | None]:
1437 result: dict[str, int | None] = {}
1438 scheduler_count = scheduler.get("allocation", {}).get("gpus_per_process")
1439 if scheduler_count is None:
1440 scheduler_count = scheduler.get("allocation", {}).get("gpus_on_node")
1441 environment_keys = {
1442 "cuda": ("CUDA_VISIBLE_DEVICES", "NVIDIA_VISIBLE_DEVICES"),
1443 "rocm": ("ROCR_VISIBLE_DEVICES", "HIP_VISIBLE_DEVICES", "CUDA_VISIBLE_DEVICES"),
1444 "metal": (),
1445 }
1446 for backend, keys in environment_keys.items():
1447 query_count = sum(
1448 1 for device in devices if device.get("backend_candidate") == backend
1449 )
1450 if query_count == 0:
1451 result[backend] = None
1452 continue
1453 limits = [query_count]
1454 if isinstance(scheduler_count, int):
1455 limits.append(scheduler_count)
1456 for key in keys:
1457 item = visibility.get(key, {})
1458 if item.get("state") == "none":
1459 limits.append(0)
1460 elif (
1461 item.get("state") == "restricted"
1462 and isinstance(item.get("entry_count"), int)
1463 ):
1464 limits.append(item["entry_count"])
1465 result[backend] = min(limits)
1466 return result
1469def _record_probe_failure(
1470 result: Mapping[str, Any],
1471 *,
1472 code: str,
1473 component: str,
1474 tool: str,
1475 warnings: list[dict[str, str]],
1476 provenance: list[dict[str, str]],
1477) -> None:
1478 status = str(result.get("status", "error"))
1479 _provenance(provenance, component, tool, status)
1480 if status != "not_found":
1481 _warning(
1482 warnings,
1483 code,
1484 component,
1485 f"{tool} read-only query did not complete successfully ({status}).",
1486 severity="info" if status in {"start_error", "error"} else "warning",
1487 )
1490def _detect_accelerators(
1491 *,
1492 system: str,
1493 machine: str,
1494 environ: Mapping[str, str],
1495 scheduler: Mapping[str, Any],
1496 run_command: Callable[..., dict[str, Any]],
1497 warnings: list[dict[str, str]],
1498 provenance: list[dict[str, str]],
1499) -> dict[str, Any]:
1500 devices: list[dict[str, Any]] = []
1502 nvidia = run_command(NVIDIA_QUERY, timeout=4.0)
1503 nvidia_devices = parse_nvidia_csv(nvidia.get("stdout", "")) if nvidia["status"] == "ok" else []
1504 if nvidia["status"] == "error":
1505 fallback = run_command(NVIDIA_QUERY_FALLBACK, timeout=4.0)
1506 if fallback["status"] == "ok":
1507 nvidia = fallback
1508 nvidia_devices = parse_nvidia_csv(fallback.get("stdout", ""))
1509 if nvidia["status"] == "ok":
1510 _provenance(provenance, "accelerators.nvidia", "nvidia-smi", "ok")
1511 if nvidia.get("stdout", "").strip() and not nvidia_devices:
1512 _warning(
1513 warnings,
1514 "NVIDIA_OUTPUT_PARSE_FAILED",
1515 "accelerators",
1516 "nvidia-smi returned an unexpected bounded CSV response.",
1517 )
1518 else:
1519 _record_probe_failure(
1520 nvidia,
1521 code="NVIDIA_QUERY_FAILED",
1522 component="accelerators.nvidia",
1523 tool="nvidia-smi",
1524 warnings=warnings,
1525 provenance=provenance,
1526 )
1527 devices.extend(nvidia_devices)
1529 amd = run_command(AMD_SMI_QUERY, timeout=4.0)
1530 amd_devices = parse_amd_json(amd.get("stdout", "")) if amd["status"] == "ok" else []
1531 amd_tool = "amd-smi"
1532 if not amd_devices and amd["status"] in {"not_found", "error", "start_error"}:
1533 amd = run_command(ROCM_SMI_QUERY, timeout=4.0)
1534 amd_tool = "rocm-smi"
1535 amd_devices = parse_amd_json(amd.get("stdout", "")) if amd["status"] == "ok" else []
1536 if amd["status"] == "ok":
1537 _provenance(provenance, "accelerators.amd", amd_tool, "ok")
1538 if amd.get("stdout", "").strip() and not amd_devices:
1539 _warning(
1540 warnings,
1541 "AMD_OUTPUT_PARSE_FAILED",
1542 "accelerators",
1543 f"{amd_tool} returned an unexpected bounded JSON response.",
1544 )
1545 else:
1546 _record_probe_failure(
1547 amd,
1548 code="AMD_QUERY_FAILED",
1549 component="accelerators.amd",
1550 tool=amd_tool,
1551 warnings=warnings,
1552 provenance=provenance,
1553 )
1554 devices.extend(amd_devices)
1556 if system == "Darwin":
1557 apple = run_command(APPLE_DISPLAY_QUERY, timeout=5.0)
1558 apple_devices = (
1559 parse_apple_profiler_json(apple.get("stdout", ""), machine=machine)
1560 if apple["status"] == "ok"
1561 else []
1562 )
1563 if apple["status"] == "ok":
1564 _provenance(
1565 provenance, "accelerators.apple", "system_profiler", "ok"
1566 )
1567 else:
1568 _record_probe_failure(
1569 apple,
1570 code="APPLE_ACCELERATOR_QUERY_FAILED",
1571 component="accelerators.apple",
1572 tool="system_profiler",
1573 warnings=warnings,
1574 provenance=provenance,
1575 )
1576 if (
1577 not apple_devices
1578 and machine.lower() in {"arm64", "aarch64"}
1579 ):
1580 apple_devices = [
1581 {
1582 "backend_candidate": "metal",
1583 "compute_capability": None,
1584 "device_class": "integrated_gpu",
1585 "device_permission": "not_tested",
1586 "driver_version": None,
1587 "gpu_cores": None,
1588 "local_index": 0,
1589 "management_query": "inferred_from_apple_silicon_platform",
1590 "memory": {
1591 "dedicated_free_bytes": None,
1592 "dedicated_total_bytes": None,
1593 "model": "unified",
1594 },
1595 "name": "Apple silicon integrated GPU",
1596 "runtime_compatibility": "not_tested",
1597 "vendor": "apple",
1598 }
1599 ]
1600 devices.extend(apple_devices)
1602 devices.sort(
1603 key=lambda item: (
1604 str(item.get("backend_candidate")),
1605 int(item.get("local_index", 0)),
1606 str(item.get("name")),
1607 )
1608 )
1609 if len(devices) > MAX_DEVICES:
1610 devices = devices[:MAX_DEVICES]
1611 _warning(
1612 warnings,
1613 "ACCELERATOR_DEVICE_LIST_BOUNDED",
1614 "accelerators",
1615 f"Accelerator output was limited to {MAX_DEVICES} devices.",
1616 )
1617 visibility = _visibility_summary(environ)
1618 if devices:
1619 _warning(
1620 warnings,
1621 "ACCELERATOR_RUNTIME_NOT_TESTED",
1622 "accelerators",
1623 "Management visibility does not prove permission or framework/runtime compatibility.",
1624 severity="info",
1625 )
1626 counts = {
1627 backend: sum(
1628 1 for device in devices if device.get("backend_candidate") == backend
1629 )
1630 for backend in ("cuda", "metal", "rocm")
1631 }
1632 return {
1633 "candidate_counts": counts,
1634 "candidate_upper_bounds": _accelerator_upper_bounds(
1635 devices, visibility, scheduler
1636 ),
1637 "devices": devices,
1638 "runtime_usable_devices": None,
1639 "visibility_environment": visibility,
1640 }
1643def _container_context(
1644 *,
1645 cgroup: Mapping[str, Any],
1646 exists: Callable[[Path], bool] = Path.exists,
1647) -> dict[str, Any]:
1648 evidence: list[str] = []
1649 for marker, label in (
1650 (Path("/.dockerenv"), "docker_marker"),
1651 (Path("/run/.containerenv"), "containerenv_marker"),
1652 ):
1653 try:
1654 if exists(marker):
1655 evidence.append(label)
1656 except OSError:
1657 pass
1658 if cgroup.get("detected") and any(
1659 cgroup.get(key) is not None
1660 for key in ("cpu_quota_cores", "cpuset_logical", "memory_max_bytes")
1661 ):
1662 evidence.append("cgroup_limit")
1663 return {
1664 "detected": bool(
1665 {"docker_marker", "containerenv_marker"}.intersection(evidence)
1666 ),
1667 "evidence": sorted(evidence),
1668 "runtime": (
1669 "docker_or_compatible"
1670 if "docker_marker" in evidence
1671 else "oci_compatible"
1672 if "containerenv_marker" in evidence
1673 else None
1674 ),
1675 }
1678def _observed_at() -> str:
1679 return (
1680 dt.datetime.now(dt.timezone.utc)
1681 .replace(microsecond=0)
1682 .isoformat()
1683 .replace("+00:00", "Z")
1684 )
1687def collect_snapshot(
1688 *,
1689 skip_accelerators: bool = False,
1690 observed_at: str | None = None,
1691 environ: Mapping[str, str] | None = None,
1692 psutil_module: Any = _AUTO,
1693 read_text: Callable[[Path], str] = _read_bounded_text,
1694 run_command: Callable[..., dict[str, Any]] = _run_bounded_command,
1695) -> dict[str, Any]:
1696 """Collect one snapshot; individual probe failures remain warnings."""
1697 warnings: list[dict[str, str]] = []
1698 provenance: list[dict[str, str]] = []
1699 environment = os.environ if environ is None else environ
1700 system = _safe_text(platform.system(), 64) or "Unknown"
1701 machine = _safe_text(platform.machine(), 64) or "unknown"
1703 if psutil_module is _AUTO:
1704 psutil_module = _load_psutil(warnings, provenance)
1705 elif psutil_module is None:
1706 _provenance(provenance, "inventory.psutil", "injected", "unavailable")
1708 mac_sysctl: dict[str, Any] = {}
1709 if system == "Darwin":
1710 mac_sysctl = _mac_sysctl_values(run_command, warnings, provenance)
1711 host_cpu = _detect_cpu_inventory(
1712 system=system,
1713 psutil_module=psutil_module,
1714 mac_sysctl=mac_sysctl,
1715 read_text=read_text,
1716 warnings=warnings,
1717 provenance=provenance,
1718 )
1719 process_cpu = _detect_process_cpu_count(
1720 psutil_module=psutil_module,
1721 warnings=warnings,
1722 provenance=provenance,
1723 )
1724 cgroup = (
1725 detect_cgroup_v2(
1726 read_text=read_text,
1727 warnings=warnings,
1728 provenance=provenance,
1729 )
1730 if system == "Linux"
1731 else {
1732 "detected": False,
1733 "cpu_quota_cores": None,
1734 "cpuset_logical": None,
1735 "memory_available_bytes": None,
1736 "memory_current_bytes": None,
1737 "memory_high_bytes": None,
1738 "memory_max_bytes": None,
1739 "scope": "not_applicable",
1740 }
1741 )
1742 scheduler = detect_scheduler(
1743 environment,
1744 warnings=warnings,
1745 provenance=provenance,
1746 )
1747 cpu = {
1748 "cgroup_v2": {
1749 "cpuset_logical": cgroup.get("cpuset_logical"),
1750 "quota_cores": cgroup.get("cpu_quota_cores"),
1751 },
1752 "effective": _effective_cpu(host_cpu, process_cpu, cgroup, scheduler),
1753 "host": host_cpu,
1754 "process": process_cpu,
1755 }
1756 memory = _detect_memory(
1757 system=system,
1758 machine=machine,
1759 psutil_module=psutil_module,
1760 mac_sysctl=mac_sysctl,
1761 cgroup=cgroup,
1762 scheduler=scheduler,
1763 read_text=read_text,
1764 warnings=warnings,
1765 provenance=provenance,
1766 )
1767 if skip_accelerators:
1768 accelerators = {
1769 "candidate_counts": {"cuda": 0, "metal": 0, "rocm": 0},
1770 "candidate_upper_bounds": {
1771 "cuda": None,
1772 "metal": None,
1773 "rocm": None,
1774 },
1775 "devices": [],
1776 "runtime_usable_devices": None,
1777 "visibility_environment": _visibility_summary(environment),
1778 }
1779 _provenance(provenance, "accelerators", "user_option", "skipped")
1780 else:
1781 accelerators = _detect_accelerators(
1782 system=system,
1783 machine=machine,
1784 environ=environment,
1785 scheduler=scheduler,
1786 run_command=run_command,
1787 warnings=warnings,
1788 provenance=provenance,
1789 )
1791 disk = _detect_disk(warnings, provenance)
1792 warnings.sort(
1793 key=lambda item: (
1794 item["component"],
1795 item["code"],
1796 item["message"],
1797 )
1798 )
1799 provenance.sort(
1800 key=lambda item: (item["component"], item["source"], item["status"])
1801 )
1802 completeness = (
1803 "partial"
1804 if any(item["severity"] == "warning" for item in warnings)
1805 else "complete_with_informational_notes"
1806 if warnings
1807 else "complete"
1808 )
1809 return {
1810 "accelerators": accelerators,
1811 "cgroup_v2": {
1812 "detected": cgroup.get("detected"),
1813 "scope": cgroup.get("scope"),
1814 },
1815 "completeness": completeness,
1816 "container": _container_context(cgroup=cgroup),
1817 "cpu": cpu,
1818 "disk": disk,
1819 "memory": memory,
1820 "observed_at": observed_at or _observed_at(),
1821 "platform": {
1822 "machine": machine,
1823 "python_version": _safe_text(platform.python_version(), 64) or "unknown",
1824 "system": system,
1825 },
1826 "privacy": {
1827 "absolute_paths_included": False,
1828 "environment_values_included": False,
1829 "hostnames_included": False,
1830 "identifiers_redacted_by_default": True,
1831 },
1832 "provenance": provenance,
1833 "scheduler": scheduler,
1834 "schema_version": SCHEMA_VERSION,
1835 "snapshot_kind": "effective_resource_snapshot",
1836 "warnings": warnings,
1837 }
1840def detect_all_resources(output_path: str | None = None) -> dict[str, Any]:
1841 """Compatibility API: collect a snapshot and optionally write it safely."""
1842 snapshot = collect_snapshot()
1843 if output_path is not None:
1844 emit_json(snapshot, output_path)
1845 return snapshot
1848def build_parser() -> argparse.ArgumentParser:
1849 parser = argparse.ArgumentParser(
1850 description=(
1851 "Report host inventory and effective process limits as deterministic JSON"
1852 )
1853 )
1854 parser.add_argument(
1855 "--output",
1856 metavar="FILE.json",
1857 help=(
1858 "write a private JSON file in the current directory; default: stdout"
1859 ),
1860 )
1861 parser.add_argument(
1862 "--force",
1863 action="store_true",
1864 help="replace an existing explicit output file",
1865 )
1866 parser.add_argument(
1867 "--skip-accelerators",
1868 action="store_true",
1869 help="skip bounded accelerator management queries",
1870 )
1871 return parser
1874def main(argv: Sequence[str] | None = None) -> int:
1875 parser = build_parser()
1876 args = parser.parse_args(argv)
1877 if args.force and not args.output:
1878 parser.error("--force requires --output")
1879 try:
1880 snapshot = collect_snapshot(skip_accelerators=args.skip_accelerators)
1881 emit_json(snapshot, args.output, force=args.force)
1882 except ResourceToolError as exc:
1883 return cli_error(parser, exc)
1884 return 0
1887if __name__ == "__main__":
1888 raise SystemExit(main())