feat: add cross-platform hardware block to comfy env --json (BE-3399)

comfy_cli/env_checker.py

@@ -7,8 +7,10 @@
 
 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
 from comfy_cli.utils import singleton
 
 console = Console()
@@ -32,6 +34,40 @@ 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/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 = 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:
+            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 +142,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 +177,5 @@ def fill_data(self) -> dict:
                 "running": server_running,
                 "url": "http://localhost:8188" if server_running else None,
             },
+            "hardware": detect_hardware(),
         }

comfy_cli/hardware.py

@@ -0,0 +1,301 @@
+"""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 os
+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
+
+    # 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
+
+        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
+    finally:
+        if saved_visible is not None:
+            os.environ["CUDA_VISIBLE_DEVICES"] = saved_visible
+
+
+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 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 key, card in payload.items():
+        if not str(key).lower().startswith("card") or not isinstance(card, dict):
+            continue
+        for field, value in card.items():
+            lowered = field.lower()
+            if model is None and "name" in lowered:
+                model = str(value).strip() or None
+            # 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,
+        "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),
+    }

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"]}
+          }
+        }
+      }
     }
   }
 }