From e9f7a95ef707e581e4e5d481adcf5140c1e821ad Mon Sep 17 00:00:00 2001 From: hualxie Date: Thu, 17 Sep 2026 14:26:30 +0800 Subject: [PATCH] feat(sys): report NPU memory --- docs/commands/sys.md | 3 ++- src/winml/modelkit/commands/sys.py | 21 ++++++++------- tests/e2e/test_sys_e2e.py | 4 +-- tests/unit/commands/test_sys_device_info.py | 30 ++++++++++++--------- 4 files changed, 32 insertions(+), 26 deletions(-) diff --git a/docs/commands/sys.md b/docs/commands/sys.md index 8f3c6f588..754235ac3 100644 --- a/docs/commands/sys.md +++ b/docs/commands/sys.md @@ -36,7 +36,7 @@ display version, build, update build revision (UBR), build branch, and build lab It then probes PyTorch for CUDA availability and GPU device names. Backend availability checks use the installed runtime environment. GPU and NPU enumeration uses DXCore as the source of adapter identity and LUID, then enriches -those native rows with WMI/PnP driver and manufacturer details. GPU rows also +those native rows with WMI/PnP driver and manufacturer details. NPU and GPU rows also report DXCore's 64-bit dedicated adapter memory and shared system memory capacities in MiB. CPU enumeration uses WMI. Devices remain in NPU > GPU > CPU priority order, and EP enumeration @@ -91,6 +91,7 @@ ML Libraries Available Devices (priority order) #1 NPU Qualcomm(R) Hexagon NPU LUID: 0x00000000_0x00018393 | Driver: 1.0.0 | Manufacturer: Qualcomm + Dedicated memory: 0 MiB | Shared memory: 8192 MiB #2 GPU Qualcomm(R) Adreno GPU LUID: 0x00000000_0x00018394 | Driver: 1.0.0 | Manufacturer: Qualcomm Dedicated memory: 1024 MiB | Shared memory: 8192 MiB diff --git a/src/winml/modelkit/commands/sys.py b/src/winml/modelkit/commands/sys.py index f5e06cdf1..f6e2931e9 100644 --- a/src/winml/modelkit/commands/sys.py +++ b/src/winml/modelkit/commands/sys.py @@ -693,17 +693,18 @@ def _gather_device_info( "manufacturer": getattr(metadata_item, "manufacturer", None), "luid": item.luid if system_adapters else None, } - if device_label == "GPU": + if system_adapters: entry["details"].update( { - "dedicated_memory_mib": ( - item.dedicated_memory_mib - if system_adapters - else getattr(item, "vram_mib", None) - ), - "shared_memory_mib": ( - item.shared_memory_mib if system_adapters else None - ), + "dedicated_memory_mib": item.dedicated_memory_mib, + "shared_memory_mib": item.shared_memory_mib, + } + ) + elif device_label == "GPU": + entry["details"].update( + { + "dedicated_memory_mib": getattr(item, "vram_mib", None), + "shared_memory_mib": None, } ) elif device_label == "CPU": @@ -847,7 +848,7 @@ def _output_device_text(devices: list[dict[str, Any]]) -> None: if arch := details.get("architecture"): parts.append(f"Architecture: {arch}") console.print(f" {' | '.join(parts)}") - if dev["type"] == "GPU": + if dev["type"] in ("NPU", "GPU"): dedicated_memory = details.get("dedicated_memory_mib") shared_memory = details.get("shared_memory_mib") parts = [ diff --git a/tests/e2e/test_sys_e2e.py b/tests/e2e/test_sys_e2e.py index 8ba794bc6..9772af547 100644 --- a/tests/e2e/test_sys_e2e.py +++ b/tests/e2e/test_sys_e2e.py @@ -199,10 +199,10 @@ def test_default_json_devices_include_luid(self): r"0x[0-9A-F]{8}_0x[0-9A-F]{8}", luid ) - def test_default_json_gpus_include_memory(self): + def test_default_json_accelerators_include_memory(self): data = _run_sys_json() for device in data["devices"]: - if device["type"] != "GPU": + if device["type"] not in ("NPU", "GPU"): continue details = device["details"] for field in ("dedicated_memory_mib", "shared_memory_mib"): diff --git a/tests/unit/commands/test_sys_device_info.py b/tests/unit/commands/test_sys_device_info.py index aee672478..6aaeb44ee 100644 --- a/tests/unit/commands/test_sys_device_info.py +++ b/tests/unit/commands/test_sys_device_info.py @@ -252,24 +252,26 @@ def test_dxcore_is_primary_for_accelerator_identity( def test_dxcore_accelerator_survives_wmi_failure(self, monkeypatch: pytest.MonkeyPatch) -> None: """Native identity remains visible when descriptive enrichment fails.""" - native_gpu = DXCoreAdapterInfo( - device_type="GPU", - name="Example GPU", + native_npu = DXCoreAdapterInfo( + device_type="NPU", + name="Example NPU", luid="0x00000000_0x00000001", vendor_id=0x1234, device_id=0x5678, + dedicated_memory_mib=0, + shared_memory_mib=16384, ) monkeypatch.setattr( "winml.modelkit.sysinfo.enumerate_compute_adapters", - lambda: [native_gpu], + lambda: [native_npu], ) with ( - patch("winml.modelkit.sysinfo.NPU.get_all", return_value=[]), patch( - "winml.modelkit.sysinfo.GPU.get_all", + "winml.modelkit.sysinfo.NPU.get_all", side_effect=RuntimeError("WMI unavailable"), ), + patch("winml.modelkit.sysinfo.GPU.get_all", return_value=[]), patch("winml.modelkit.sysinfo.CPU.get_all", return_value=[]), ): result = _gather_device_info() @@ -277,14 +279,14 @@ def test_dxcore_accelerator_survives_wmi_failure(self, monkeypatch: pytest.Monke assert result == [ { "priority": 1, - "type": "GPU", - "name": "Example GPU", + "type": "NPU", + "name": "Example NPU", "details": { "driver": None, "manufacturer": None, "luid": "0x00000000_0x00000001", - "dedicated_memory_mib": None, - "shared_memory_mib": None, + "dedicated_memory_mib": 0, + "shared_memory_mib": 16384, }, } ] @@ -417,7 +419,9 @@ def test_compact_device_output_includes_luid( assert "CPU: Test CPU (LUID: N/A)" in output -def test_gpu_text_output_includes_dedicated_and_shared_memory( +@pytest.mark.parametrize("device_type", ["NPU", "GPU"]) +def test_accelerator_text_output_includes_dedicated_and_shared_memory( + device_type: str, capsys: pytest.CaptureFixture[str], ) -> None: from winml.modelkit.commands.sys import _output_device_text @@ -426,8 +430,8 @@ def test_gpu_text_output_includes_dedicated_and_shared_memory( [ { "priority": 1, - "type": "GPU", - "name": "Test GPU", + "type": device_type, + "name": f"Test {device_type}", "details": { "luid": "0x00000000_0x00018393", "driver": "1.0",