From 6176fa73bb23c665916594f7751eb5943b4b3c53 Mon Sep 17 00:00:00 2001 From: Matt Miller Date: Fri, 17 Jul 2026 17:40:49 -0700 Subject: [PATCH 1/2] feat: add cross-platform hardware block to comfy env --json (BE-3399) Add comfy_cli/hardware.py with a single detect_hardware() entry point that reports OS, CPU, RAM, and a GPU sub-block (vendor/model/VRAM/unified-memory) so agent surfaces can route weak machines away from local diffusion. - macOS: cpu via sysctl machdep.cpu.brand_string; Apple Silicon -> apple unified-memory GPU; Intel Mac -> unknown non-unified GPU (no system_profiler). - NVIDIA (any OS): nvidia-smi CSV, then ctypes libcuda fallback (reuses cuda_detect._load_libcuda). - AMD (Linux): best-effort rocm-smi --json. - Never raises, never blocks long: every probe wrapped, subprocesses timeout=5. Wire hardware into EnvChecker.fill_data() (JSON) and a Hardware row in fill_print_table() (pretty). Extend schemas/env.json with an OPTIONAL, fully nullable hardware property (not in required) so older/failed-probe payloads stay valid. Envelope schema unchanged (envelope/1). --- comfy_cli/env_checker.py | 34 ++++ comfy_cli/hardware.py | 279 +++++++++++++++++++++++++++++++ comfy_cli/schemas/env.json | 21 +++ tests/comfy_cli/test_hardware.py | 263 +++++++++++++++++++++++++++++ 4 files changed, 597 insertions(+) create mode 100644 comfy_cli/hardware.py create mode 100644 tests/comfy_cli/test_hardware.py diff --git a/comfy_cli/env_checker.py b/comfy_cli/env_checker.py index 2aea4cbd..60cfae45 100644 --- a/comfy_cli/env_checker.py +++ b/comfy_cli/env_checker.py @@ -9,6 +9,7 @@ from rich.console import Console from comfy_cli.config_manager import ConfigManager +from comfy_cli.hardware import detect_hardware from comfy_cli.utils import singleton console = Console() @@ -32,6 +33,36 @@ def format_python_version(version_info): return f"[bold red]{version_info.major}.{version_info.minor}.{version_info.micro}[/bold red]" +def _bytes_to_gb(value) -> str: + """Render a byte count as a whole-number GB string, or ``?`` if unknown.""" + if not isinstance(value, (int | float)): + return "?" + return f"{round(value / (1024**3))} GB" + + +def format_hardware_summary(hw: dict) -> str: + """One-line ``cpu / RAM GB / GPU model (VRAM GB or 'unified')`` summary. + + Used by ``fill_print_table`` (pretty mode). Tolerates missing/None fields so + it never raises on a partial (failed-probe) hardware block. + """ + cpu = hw.get("cpu") or "unknown CPU" + ram = _bytes_to_gb(hw.get("ram_bytes")) + + gpu = hw.get("gpu") + if not gpu: + gpu_part = "no GPU" + else: + model = gpu.get("model") or gpu.get("vendor") or "unknown GPU" + if gpu.get("unified_memory"): + gpu_part = f"{model} (unified)" + elif gpu.get("vram_bytes") is not None: + gpu_part = f"{model} ({_bytes_to_gb(gpu.get('vram_bytes'))} VRAM)" + else: + gpu_part = model + return f"{cpu} / {ram} RAM / {gpu_part}" + + def check_comfy_server_running(port=8188, host="localhost", timeout: float = 5.0): """ Checks if the Comfy server is running by making a GET request to the /history endpoint. @@ -106,6 +137,8 @@ def fill_print_table(self): config_data = ConfigManager().get_env_data() data.extend(config_data) + data.append(("Hardware", format_hardware_summary(detect_hardware()))) + if check_comfy_server_running(): data.append( ( @@ -139,4 +172,5 @@ def fill_data(self) -> dict: "running": server_running, "url": "http://localhost:8188" if server_running else None, }, + "hardware": detect_hardware(), } diff --git a/comfy_cli/hardware.py b/comfy_cli/hardware.py new file mode 100644 index 00000000..8a4d4b81 --- /dev/null +++ b/comfy_cli/hardware.py @@ -0,0 +1,279 @@ +"""Cross-platform hardware probing for ``comfy env --json``. + +A single entry point, :func:`detect_hardware`, returns a JSON-serializable dict +describing the machine's generation-relevant hardware (OS, CPU, RAM, GPU). Agent +surfaces use this to route weak machines away from local diffusion. + +Contract: **never raise, never block long.** Every probe is wrapped so a failure +yields ``None`` rather than an exception, and every subprocess is bounded by a +``timeout=5``. The RAM figure is the only value expected to always populate +(``psutil`` is a hard dependency); everything else degrades to ``None``. +""" + +from __future__ import annotations + +import ctypes +import logging +import platform +import subprocess + +import psutil + +from comfy_cli import cuda_detect + +logger = logging.getLogger(__name__) + +_SUBPROCESS_TIMEOUT = 5 + + +def _run(cmd: list[str]) -> str | None: + """Run ``cmd`` and return stripped stdout, or ``None`` on any failure. + + Bounded by ``timeout=5`` so a hung binary can never block the probe. + """ + try: + output = subprocess.check_output( + cmd, + text=True, + timeout=_SUBPROCESS_TIMEOUT, + stderr=subprocess.DEVNULL, + ) + except Exception: + logger.debug("hardware probe command failed: %s", cmd, exc_info=True) + return None + output = output.strip() + return output or None + + +def _detect_cpu(system: str) -> str | None: + """Best-effort human-readable CPU/chip name, or ``None`` if unknown.""" + try: + if system == "Darwin": + return _run(["sysctl", "-n", "machdep.cpu.brand_string"]) + + if system == "Linux": + try: + with open("/proc/cpuinfo", encoding="utf-8", errors="replace") as fh: + for line in fh: + if line.startswith("model name"): + _, _, value = line.partition(":") + value = value.strip() + if value: + return value + except OSError: + logger.debug("reading /proc/cpuinfo failed", exc_info=True) + + # Windows / Linux fallback: platform.processor() is often populated on + # Windows and empty on Linux (hence the /proc/cpuinfo pass above first). + return platform.processor() or None + except Exception: + logger.debug("CPU detection failed", exc_info=True) + return None + + +def _detect_gpu_nvidia_smi() -> dict | None: + """Query ``nvidia-smi`` for the first GPU's name and VRAM. + + Returns a gpu dict or ``None`` if the binary is absent / produced no usable + output. ``cuda_detect`` already parses nvidia-smi in production. + """ + output = _run( + [ + "nvidia-smi", + "--query-gpu=name,memory.total", + "--format=csv,noheader,nounits", + ] + ) + if not output: + return None + + # First line is the first GPU: "NVIDIA GeForce RTX 4090, 24564" + first = output.splitlines()[0] + parts = [p.strip() for p in first.split(",")] + if len(parts) < 2: + return None + + model = parts[0] or None + vram_bytes = None + try: + mib = int(parts[1]) + vram_bytes = mib * 1024 * 1024 + except (ValueError, TypeError): + logger.debug("could not parse nvidia-smi VRAM: %r", parts[1]) + + if not model and vram_bytes is None: + return None + + return { + "vendor": "nvidia", + "model": model, + "vram_bytes": vram_bytes, + "unified_memory": False, + } + + +def _detect_gpu_nvidia_ctypes() -> dict | None: + """ctypes fallback for NVIDIA when ``nvidia-smi`` is unavailable. + + Uses the CUDA driver API via ``cuda_detect._load_libcuda``: ``cuInit`` + + ``cuDeviceGet`` + ``cuDeviceGetName`` + ``cuDeviceTotalMem_v2`` — model and + VRAM without the binary. Returns a gpu dict or ``None``. + """ + try: + libcuda = cuda_detect._load_libcuda() + except OSError: + logger.debug("libcuda not loadable for GPU probe") + return None + + try: + if libcuda.cuInit(0) != 0: + return None + + device = ctypes.c_int() + if libcuda.cuDeviceGet(ctypes.byref(device), 0) != 0: + return None + + name_buf = ctypes.create_string_buffer(256) + model = None + if libcuda.cuDeviceGetName(name_buf, len(name_buf), device) == 0: + decoded = name_buf.value.decode("utf-8", errors="replace").strip() + model = decoded or None + + vram_bytes = None + total = ctypes.c_size_t() + total_mem = getattr(libcuda, "cuDeviceTotalMem_v2", None) or getattr(libcuda, "cuDeviceTotalMem", None) + if total_mem is not None and total_mem(ctypes.byref(total), device) == 0: + vram_bytes = int(total.value) + + if not model and vram_bytes is None: + return None + + return { + "vendor": "nvidia", + "model": model, + "vram_bytes": vram_bytes, + "unified_memory": False, + } + except Exception: + logger.debug("ctypes NVIDIA GPU probe failed", exc_info=True) + return None + + +def _detect_gpu_amd() -> dict | None: + """Best-effort AMD GPU probe on Linux via ``rocm-smi``. + + Minimal by design — nulls are acceptable. Returns a gpu dict (vendor "amd") + or ``None`` if ``rocm-smi`` is absent / unparseable. + """ + import json + + output = _run(["rocm-smi", "--showmeminfo", "vram", "--json"]) + if not output: + return None + + try: + payload = json.loads(output) + except (ValueError, TypeError): + logger.debug("could not parse rocm-smi JSON", exc_info=True) + return None + + if not isinstance(payload, dict) or not payload: + return None + + # rocm-smi keys cards as "card0", "card1", ...; take the first. + model = None + vram_bytes = None + for card in payload.values(): + if not isinstance(card, dict): + continue + for key, value in card.items(): + lowered = key.lower() + if model is None and "name" in lowered: + model = str(value).strip() or None + if vram_bytes is None and "total" in lowered and "vram" in lowered: + try: + vram_bytes = int(value) + except (ValueError, TypeError): + pass + break + + return { + "vendor": "amd", + "model": model, + "vram_bytes": vram_bytes, + "unified_memory": False, + } + + +def _detect_gpu(system: str, machine: str, cpu: str | None) -> dict | None: + """Resolve the GPU block, or ``None`` if no GPU is detected. + + Apple Silicon is unified-memory (the RAM figure IS the GPU budget). Every + other path tries NVIDIA (nvidia-smi then ctypes), then AMD. + """ + try: + if system == "Darwin" and machine == "arm64": + return { + "vendor": "apple", + "model": cpu, + "vram_bytes": None, + "unified_memory": True, + } + + gpu = _detect_gpu_nvidia_smi() + if gpu is not None: + return gpu + + gpu = _detect_gpu_nvidia_ctypes() + if gpu is not None: + return gpu + + if system == "Linux": + gpu = _detect_gpu_amd() + if gpu is not None: + return gpu + + return None + except Exception: + logger.debug("GPU detection failed", exc_info=True) + return None + + +def _detect_ram_bytes() -> int | None: + try: + return int(psutil.virtual_memory().total) + except Exception: + logger.debug("RAM detection failed", exc_info=True) + return None + + +def detect_hardware() -> dict: + """Detect generation-relevant hardware. Never raises. + + Returns a dict shaped like ``schemas/env.json``'s ``hardware`` block: + ``os``/``os_version``/``arch``/``cpu``/``ram_bytes`` plus a ``gpu`` sub-dict + (or ``None``). Any probe that fails contributes ``None`` rather than raising. + """ + try: + system = platform.system() + except Exception: + system = "" + try: + os_version = platform.release() + except Exception: + os_version = None + try: + machine = platform.machine() + except Exception: + machine = "" + + cpu = _detect_cpu(system) + + return { + "os": system or None, + "os_version": os_version or None, + "arch": machine or None, + "cpu": cpu, + "ram_bytes": _detect_ram_bytes(), + "gpu": _detect_gpu(system, machine, cpu), + } diff --git a/comfy_cli/schemas/env.json b/comfy_cli/schemas/env.json index 46afcccd..29b0eb72 100644 --- a/comfy_cli/schemas/env.json +++ b/comfy_cli/schemas/env.json @@ -51,6 +51,27 @@ "port": {"type": "integer"}, "pid": {"type": "integer"} } + }, + "hardware": { + "type": ["object", "null"], + "additionalProperties": true, + "properties": { + "os": {"type": ["string", "null"]}, + "os_version": {"type": ["string", "null"]}, + "arch": {"type": ["string", "null"]}, + "cpu": {"type": ["string", "null"]}, + "ram_bytes": {"type": ["integer", "null"]}, + "gpu": { + "type": ["object", "null"], + "additionalProperties": true, + "properties": { + "vendor": {"type": ["string", "null"]}, + "model": {"type": ["string", "null"]}, + "vram_bytes": {"type": ["integer", "null"]}, + "unified_memory": {"type": ["boolean", "null"]} + } + } + } } } } diff --git a/tests/comfy_cli/test_hardware.py b/tests/comfy_cli/test_hardware.py new file mode 100644 index 00000000..8057f13e --- /dev/null +++ b/tests/comfy_cli/test_hardware.py @@ -0,0 +1,263 @@ +"""Tests for :mod:`comfy_cli.hardware`. + +Each branch (macOS/apple, nvidia-smi, ctypes fallback, total failure) is mocked +at the module boundary so the tests run identically on any CI runner. The +overriding contract under test: ``detect_hardware`` never raises, always returns +RAM, and degrades every failed probe to ``None``. +""" + +from __future__ import annotations + +import ctypes +import json +from pathlib import Path +from unittest.mock import patch + +import jsonschema + +from comfy_cli import hardware +from comfy_cli.env_checker import EnvChecker, format_hardware_summary + +_EnvCheckerCls = EnvChecker.__closure__[0].cell_contents + +SCHEMA_PATH = Path(hardware.__file__).parent / "schemas" / "env.json" + + +def _env_schema() -> dict: + return json.loads(SCHEMA_PATH.read_text()) + + +class TestDetectHardwareMacOS: + def test_apple_silicon_unified_block(self): + with ( + patch.object(hardware.platform, "system", return_value="Darwin"), + patch.object(hardware.platform, "machine", return_value="arm64"), + patch.object(hardware.platform, "release", return_value="25.4.0"), + patch.object(hardware, "_run", return_value="Apple M4 Max"), + patch.object(hardware, "_detect_ram_bytes", return_value=68719476736), + ): + hw = hardware.detect_hardware() + + assert hw["os"] == "Darwin" + assert hw["os_version"] == "25.4.0" + assert hw["arch"] == "arm64" + assert hw["cpu"] == "Apple M4 Max" + assert hw["ram_bytes"] == 68719476736 + assert hw["gpu"] == { + "vendor": "apple", + "model": "Apple M4 Max", + "vram_bytes": None, + "unified_memory": True, + } + + def test_intel_mac_has_no_apple_gpu(self): + """Intel Macs are not unified-memory; they fall through to the NVIDIA/AMD + probes (which return nothing here) → gpu null, not an apple block.""" + with ( + patch.object(hardware.platform, "system", return_value="Darwin"), + patch.object(hardware.platform, "machine", return_value="x86_64"), + patch.object(hardware, "_run", return_value=None), + patch.object(hardware.cuda_detect, "_load_libcuda", side_effect=OSError), + patch.object(hardware, "_detect_ram_bytes", return_value=17179869184), + ): + hw = hardware.detect_hardware() + + assert hw["gpu"] is None + assert hw["ram_bytes"] == 17179869184 + + +class TestDetectHardwareNvidiaSmi: + def test_nvidia_smi_happy_path(self): + def fake_run(cmd): + if cmd[0] == "nvidia-smi": + return "NVIDIA GeForce RTX 4090, 24564" + return None + + with ( + patch.object(hardware.platform, "system", return_value="Linux"), + patch.object(hardware.platform, "machine", return_value="x86_64"), + patch.object(hardware, "_run", side_effect=fake_run), + patch.object(hardware, "_detect_ram_bytes", return_value=137438953472), + ): + hw = hardware.detect_hardware() + + assert hw["gpu"] == { + "vendor": "nvidia", + "model": "NVIDIA GeForce RTX 4090", + "vram_bytes": 24564 * 1024 * 1024, + "unified_memory": False, + } + + def test_nvidia_smi_multi_gpu_takes_first(self): + with patch.object( + hardware, + "_run", + return_value="NVIDIA GeForce RTX 4090, 24564\nNVIDIA GeForce RTX 3090, 24576", + ): + gpu = hardware._detect_gpu_nvidia_smi() + assert gpu["model"] == "NVIDIA GeForce RTX 4090" + assert gpu["vram_bytes"] == 24564 * 1024 * 1024 + + +class _FakeLibCuda: + """Minimal stand-in for a loaded ``libcuda`` CDLL for the ctypes path.""" + + def __init__(self, name: str, total_bytes: int): + self._name = name + self._total = total_bytes + + def cuInit(self, flags): + return 0 + + def cuDeviceGet(self, dev_ptr, ordinal): + dev_ptr._obj.value = 0 + return 0 + + def cuDeviceGetName(self, buf, length, dev): + buf.value = self._name.encode("utf-8") + return 0 + + def cuDeviceTotalMem_v2(self, size_ptr, dev): + size_ptr._obj.value = self._total + return 0 + + +class TestDetectHardwareCtypesFallback: + def test_nvidia_smi_absent_libcuda_used(self): + fake = _FakeLibCuda("NVIDIA A100-SXM4-80GB", 85899345920) + + def fake_run(cmd): + # nvidia-smi absent → None so we exercise the ctypes fallback. + return None + + with ( + patch.object(hardware.platform, "system", return_value="Linux"), + patch.object(hardware.platform, "machine", return_value="x86_64"), + patch.object(hardware, "_run", side_effect=fake_run), + patch.object(hardware.cuda_detect, "_load_libcuda", return_value=fake), + patch.object(hardware, "_detect_ram_bytes", return_value=270582939648), + ): + hw = hardware.detect_hardware() + + assert hw["gpu"] == { + "vendor": "nvidia", + "model": "NVIDIA A100-SXM4-80GB", + "vram_bytes": 85899345920, + "unified_memory": False, + } + + def test_byref_obj_contract(self): + """Guard the ctypes mock trick: byref exposes the wrapped object so the + fake driver can write results back through the pointer.""" + val = ctypes.c_int(7) + assert ctypes.byref(val)._obj is val + + +class TestDetectHardwareNeverRaises: + def test_all_probes_fail_yields_nulls_but_ram(self): + with ( + patch.object(hardware.platform, "system", return_value="Windows"), + patch.object(hardware.platform, "machine", return_value="AMD64"), + patch.object(hardware.platform, "processor", return_value=""), + patch.object(hardware, "_run", side_effect=TimeoutError("boom")), + patch.object(hardware.cuda_detect, "_load_libcuda", side_effect=OSError), + ): + # Must not raise even though every subprocess/ctypes probe blows up. + hw = hardware.detect_hardware() + + assert hw["os"] == "Windows" + assert hw["cpu"] is None + assert hw["gpu"] is None + # RAM comes from psutil (a hard dep) and is unaffected by probe failures. + assert isinstance(hw["ram_bytes"], int) + assert hw["ram_bytes"] > 0 + + def test_platform_calls_raising_do_not_escape(self): + with ( + patch.object(hardware.platform, "system", side_effect=RuntimeError), + patch.object(hardware.platform, "machine", side_effect=RuntimeError), + patch.object(hardware.platform, "release", side_effect=RuntimeError), + ): + hw = hardware.detect_hardware() + assert hw["os"] is None + assert hw["arch"] is None + assert hw["gpu"] is None + + +class TestFormatHardwareSummary: + def test_unified_memory_summary(self): + hw = { + "cpu": "Apple M4 Max", + "ram_bytes": 68719476736, + "gpu": {"vendor": "apple", "model": "Apple M4 Max", "vram_bytes": None, "unified_memory": True}, + } + summary = format_hardware_summary(hw) + assert "Apple M4 Max" in summary + assert "64 GB RAM" in summary + assert "unified" in summary + + def test_discrete_vram_summary(self): + hw = { + "cpu": "AMD Ryzen 9", + "ram_bytes": 137438953472, + "gpu": {"vendor": "nvidia", "model": "RTX 4090", "vram_bytes": 24 * 1024**3, "unified_memory": False}, + } + summary = format_hardware_summary(hw) + assert "RTX 4090" in summary + assert "24 GB VRAM" in summary + + def test_no_gpu_and_missing_fields(self): + summary = format_hardware_summary({"cpu": None, "ram_bytes": None, "gpu": None}) + assert "no GPU" in summary + assert "unknown CPU" in summary + + +class TestFillDataHardware: + @patch("comfy_cli.env_checker.check_comfy_server_running", return_value=False) + @patch("comfy_cli.env_checker.ConfigManager") + def test_fill_data_includes_hardware_and_validates(self, mock_cm, _mock_server): + mock_cm.return_value.get_data.return_value = {"path": "/tmp/config"} + + checker = _EnvCheckerCls.__new__(_EnvCheckerCls) + checker.virtualenv_path = None + checker.conda_env = None + + data = checker.fill_data() + assert "hardware" in data + hw = data["hardware"] + assert set(hw) >= {"os", "os_version", "arch", "cpu", "ram_bytes", "gpu"} + + # Complete the payload the way cmdline.py does (workspace is added there) + # and validate the whole thing against env.json. + data["workspace"] = {"path": None, "type": None} + jsonschema.Draft202012Validator(_env_schema()).validate(data) + + def test_fully_nulled_hardware_still_validates(self): + """A failed-probe payload (all nulls, gpu null) must stay schema-valid — + this is why `hardware` is optional and every leaf is nullable.""" + payload = { + "python": {"version": "3.12.1", "executable": "/usr/bin/python"}, + "config": {}, + "workspace": {"path": None, "type": None}, + "server": {"running": False}, + "hardware": { + "os": None, + "os_version": None, + "arch": None, + "cpu": None, + "ram_bytes": None, + "gpu": None, + }, + } + jsonschema.Draft202012Validator(_env_schema()).validate(payload) + + def test_payload_without_hardware_still_validates(self): + """Older consumers / pre-hardware payloads must remain valid — `hardware` + is not in `required`.""" + payload = { + "python": {"version": "3.12.1", "executable": "/usr/bin/python"}, + "config": {}, + "workspace": {"path": None, "type": None}, + "server": {"running": False}, + } + jsonschema.Draft202012Validator(_env_schema()).validate(payload) From 5e5b4f4ff7c918331bd1952369fe2a5329fd74b4 Mon Sep 17 00:00:00 2001 From: Matt Miller Date: Fri, 17 Jul 2026 18:01:05 -0700 Subject: [PATCH 2/2] fix(hardware): address Cursor review findings for env hardware block (BE-3399) - env_checker: escape untrusted cpu/gpu-model strings in the Rich-rendered hardware summary so markup-like content (e.g. "[/]") can't raise MarkupError and crash `comfy env` in pretty mode. - hardware(_detect_gpu_amd): return None (no phantom all-None AMD block) when nothing parses, matching the NVIDIA probes; gate the card loop on the "card" key prefix so a non-card metadata block isn't parsed as the GPU; exclude the "VRAM Total Used Memory" usage key so total capacity isn't understated. - hardware(_detect_gpu_nvidia_ctypes): pop/restore CUDA_VISIBLE_DEVICES around the ctypes probe (mirrors cuda_detect) so an exported ""/"-1" doesn't hide a present GPU. - tests: cover AMD total-vs-used key, metadata-block skip, all-None->None, and markup escaping. Co-Authored-By: Claude Opus 4.8 --- comfy_cli/env_checker.py | 9 +++- comfy_cli/hardware.py | 34 ++++++++++++--- tests/comfy_cli/test_hardware.py | 72 ++++++++++++++++++++++++++++++++ 3 files changed, 107 insertions(+), 8 deletions(-) diff --git a/comfy_cli/env_checker.py b/comfy_cli/env_checker.py index 60cfae45..c1ad8185 100644 --- a/comfy_cli/env_checker.py +++ b/comfy_cli/env_checker.py @@ -7,6 +7,7 @@ import requests from rich.console import Console +from rich.markup import escape from comfy_cli.config_manager import ConfigManager from comfy_cli.hardware import detect_hardware @@ -46,14 +47,18 @@ def format_hardware_summary(hw: dict) -> str: Used by ``fill_print_table`` (pretty mode). Tolerates missing/None fields so it never raises on a partial (failed-probe) hardware block. """ - cpu = hw.get("cpu") or "unknown CPU" + # cpu/model come from untrusted probe output (/proc/cpuinfo, nvidia-smi, + # rocm-smi, libcuda). This summary is rendered by fill_print_table via Rich + # (markup enabled), so escape those strings — a stray "[" (plausible on + # VMs/hypervisors) would otherwise raise MarkupError and crash `comfy env`. + cpu = escape(hw.get("cpu") or "unknown CPU") ram = _bytes_to_gb(hw.get("ram_bytes")) gpu = hw.get("gpu") if not gpu: gpu_part = "no GPU" else: - model = gpu.get("model") or gpu.get("vendor") or "unknown GPU" + model = escape(str(gpu.get("model") or gpu.get("vendor") or "unknown GPU")) if gpu.get("unified_memory"): gpu_part = f"{model} (unified)" elif gpu.get("vram_bytes") is not None: diff --git a/comfy_cli/hardware.py b/comfy_cli/hardware.py index 8a4d4b81..61aba2b3 100644 --- a/comfy_cli/hardware.py +++ b/comfy_cli/hardware.py @@ -14,6 +14,7 @@ import ctypes import logging +import os import platform import subprocess @@ -125,7 +126,15 @@ def _detect_gpu_nvidia_ctypes() -> dict | None: logger.debug("libcuda not loadable for GPU probe") return None + # Mirror cuda_detect.detect_cuda_driver_version: an exported + # CUDA_VISIBLE_DEVICES="" or "-1" (common on shared/CI hosts) hides all + # devices from cuDeviceGet, so a physically-present GPU would report as + # gpu:null. Drop it for the duration of the probe and restore it after. + saved_visible = os.environ.get("CUDA_VISIBLE_DEVICES") try: + if saved_visible is not None: + os.environ.pop("CUDA_VISIBLE_DEVICES", None) + if libcuda.cuInit(0) != 0: return None @@ -157,6 +166,9 @@ def _detect_gpu_nvidia_ctypes() -> dict | None: except Exception: logger.debug("ctypes NVIDIA GPU probe failed", exc_info=True) return None + finally: + if saved_visible is not None: + os.environ["CUDA_VISIBLE_DEVICES"] = saved_visible def _detect_gpu_amd() -> dict | None: @@ -180,23 +192,33 @@ def _detect_gpu_amd() -> dict | None: if not isinstance(payload, dict) or not payload: return None - # rocm-smi keys cards as "card0", "card1", ...; take the first. + # rocm-smi keys cards as "card0", "card1", ...; take the first card entry. + # Gate on the "card" prefix so a non-card metadata block (e.g. a "system" + # dict) that iterates first isn't mistaken for the GPU. model = None vram_bytes = None - for card in payload.values(): - if not isinstance(card, dict): + for key, card in payload.items(): + if not str(key).lower().startswith("card") or not isinstance(card, dict): continue - for key, value in card.items(): - lowered = key.lower() + for field, value in card.items(): + lowered = field.lower() if model is None and "name" in lowered: model = str(value).strip() or None - if vram_bytes is None and "total" in lowered and "vram" in lowered: + # Match the total-capacity key ("VRAM Total Memory (B)"), excluding + # the usage key ("VRAM Total Used Memory (B)") which also contains + # both "vram" and "total" and would otherwise understate capacity. + if vram_bytes is None and "vram" in lowered and "total" in lowered and "used" not in lowered: try: vram_bytes = int(value) except (ValueError, TypeError): pass break + # Report no GPU rather than a phantom all-None AMD block (which would spoof + # GPU presence for routing decisions), matching the NVIDIA probes. + if not model and vram_bytes is None: + return None + return { "vendor": "amd", "model": model, diff --git a/tests/comfy_cli/test_hardware.py b/tests/comfy_cli/test_hardware.py index 8057f13e..c496f990 100644 --- a/tests/comfy_cli/test_hardware.py +++ b/tests/comfy_cli/test_hardware.py @@ -184,7 +184,79 @@ def test_platform_calls_raising_do_not_escape(self): assert hw["gpu"] is None +class TestDetectGpuAmd: + def test_amd_happy_path(self): + payload = json.dumps( + {"card0": {"Card SKU": "gfx90a", "GPU Name": "AMD Instinct MI210", "VRAM Total Memory (B)": "68702699520"}} + ) + with patch.object(hardware, "_run", return_value=payload): + gpu = hardware._detect_gpu_amd() + assert gpu == { + "vendor": "amd", + "model": "AMD Instinct MI210", + "vram_bytes": 68702699520, + "unified_memory": False, + } + + def test_amd_total_key_beats_used_key(self): + """The 'VRAM Total Used Memory' key also contains 'vram'+'total'; the probe + must report the total-capacity key, not usage.""" + payload = json.dumps( + { + "card0": { + "VRAM Total Used Memory (B)": "1073741824", + "VRAM Total Memory (B)": "68702699520", + } + } + ) + with patch.object(hardware, "_run", return_value=payload): + gpu = hardware._detect_gpu_amd() + assert gpu["vram_bytes"] == 68702699520 + + def test_amd_skips_non_card_metadata_block(self): + """A leading non-card metadata dict must not be mistaken for the GPU.""" + payload = json.dumps( + { + "system": {"Driver version": "6.0.0"}, + "card0": {"GPU Name": "AMD Radeon RX 7900 XTX", "VRAM Total Memory (B)": "25757220864"}, + } + ) + with patch.object(hardware, "_run", return_value=payload): + gpu = hardware._detect_gpu_amd() + assert gpu["model"] == "AMD Radeon RX 7900 XTX" + assert gpu["vram_bytes"] == 25757220864 + + def test_amd_all_none_reports_no_gpu(self): + """An error/metadata-only payload with nothing parseable must yield None, + not a phantom {vendor: amd, model: None, vram_bytes: None} block.""" + payload = json.dumps({"card0": {"Something Unrelated": "x"}}) + with patch.object(hardware, "_run", return_value=payload): + assert hardware._detect_gpu_amd() is None + + class TestFormatHardwareSummary: + def test_markup_in_model_is_escaped_not_crashing(self): + """CPU/GPU strings are untrusted; markup-like content (e.g. a close tag + '[/]') must be escaped so rendering through Rich can't raise MarkupError.""" + import pytest + from rich.errors import MarkupError + from rich.text import Text + + raw_model = "GPU [/] X" + # Sanity: the raw model really would crash Rich markup parsing. + with pytest.raises(MarkupError): + Text.from_markup(raw_model) + + hw = { + "cpu": "Fake CPU", + "ram_bytes": 17179869184, + "gpu": {"vendor": "nvidia", "model": raw_model, "vram_bytes": 8 * 1024**3, "unified_memory": False}, + } + summary = format_hardware_summary(hw) + assert "\\[/]" in summary + # Rendering with markup enabled (as cmdline.py does) must not raise. + Text.from_markup(summary) + def test_unified_memory_summary(self): hw = { "cpu": "Apple M4 Max",