Buckets:
| """Hardware / platform detection. Stdlib only, safe to run anywhere.""" | |
| import os | |
| import platform | |
| import re | |
| import subprocess | |
| from .util import which, IS_WINDOWS, IS_MAC, IS_LINUX | |
| def _run_quiet(cmd): | |
| try: | |
| out = subprocess.run( | |
| cmd, stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, timeout=15 | |
| ) | |
| if out.returncode != 0: | |
| return None | |
| return out.stdout.decode("utf-8", errors="replace") | |
| except (FileNotFoundError, subprocess.TimeoutExpired, OSError): | |
| return None | |
| def detect_nvidia(): | |
| """Returns dict with gpus + driver info, or None if no working nvidia-smi.""" | |
| smi = which("nvidia-smi") | |
| if not smi: | |
| return None | |
| fields = "name,memory.total,driver_version,compute_cap" | |
| csv = _run_quiet([smi, "--query-gpu=" + fields, "--format=csv,noheader"]) | |
| if not csv: | |
| # older drivers don't know the compute_cap field | |
| csv = _run_quiet( | |
| [ | |
| smi, | |
| "--query-gpu=name,memory.total,driver_version", | |
| "--format=csv,noheader", | |
| ] | |
| ) | |
| if not csv: | |
| return None | |
| gpus = [] | |
| driver = None | |
| for line in csv.strip().splitlines(): | |
| parts = [p.strip() for p in line.split(",")] | |
| if len(parts) >= 3: | |
| gpu = {"name": parts[0], "memory": parts[1]} | |
| if len(parts) >= 4: | |
| gpu["compute_cap"] = parts[3] | |
| gpus.append(gpu) | |
| driver = parts[2] | |
| if not gpus: | |
| return None | |
| # Max CUDA version the driver supports only appears in the banner output. | |
| # Windows/WDDM labels it "CUDA UMD Version:" on recent drivers, Linux just | |
| # "CUDA Version:" — without the optional word we silently fall through to | |
| # the "driver present but version unknown, assume current" path and can | |
| # hand an old driver cu130 wheels it cannot run. | |
| banner = _run_quiet([smi]) or "" | |
| m = re.search(r"CUDA(?:\s+[A-Z]+)?\s+Version:\s*([0-9]+\.[0-9]+)", banner) | |
| cuda_version = m.group(1) if m else None | |
| return {"gpus": gpus, "driver": driver, "cuda_version": cuda_version} | |
| def detect_rocm(): | |
| """Returns dict if an AMD ROCm stack is present, else None.""" | |
| if not IS_LINUX: | |
| return None | |
| smi = which("rocm-smi") | |
| if not smi and not os.path.isdir("/opt/rocm"): | |
| return None | |
| gpus = [] | |
| out = _run_quiet([smi, "--showproductname"]) if smi else None | |
| if out: | |
| for m in re.finditer(r"Card [Ss]eries:\s*(.+)", out): | |
| gpus.append({"name": m.group(1).strip()}) | |
| return {"gpus": gpus} | |
| def detect(): | |
| """Full platform detection. Returns a plain dict (json-serializable).""" | |
| system = platform.system() # Linux / Darwin / Windows | |
| arch = platform.machine().lower() # x86_64 / amd64 / arm64 / aarch64 | |
| if arch == "amd64": | |
| arch = "x86_64" | |
| if arch == "arm64" and not IS_MAC: | |
| arch = "aarch64" | |
| result = { | |
| "os": {"Linux": "linux", "Darwin": "mac", "Windows": "windows"}.get( | |
| system, system.lower() | |
| ), | |
| "arch": arch, | |
| "python": platform.python_version(), | |
| "nvidia": None, | |
| "rocm": None, | |
| "backend": "cpu", | |
| } | |
| nvidia = detect_nvidia() | |
| if nvidia: | |
| result["nvidia"] = nvidia | |
| result["backend"] = "cuda" | |
| else: | |
| rocm = detect_rocm() | |
| if rocm: | |
| result["rocm"] = rocm | |
| result["backend"] = "rocm" | |
| if IS_MAC: | |
| result["backend"] = "mps" if arch == "arm64" else "cpu" | |
| # DGX OS / Grace (GB10, DGX Spark): NVIDIA GPU on aarch64 Linux | |
| result["is_dgx"] = bool(result["os"] == "linux" and arch == "aarch64" and nvidia) | |
| return result | |
Xet Storage Details
- Size:
- 3.71 kB
- Xet hash:
- d91d8d2d100d4637549d937052901508ac2989af27446f3f8c889a40280ee565
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.