Files
vllm-omni/tests/edge/local/test_gpu_attribution.py
T
tzhouamandClaude Fable 5.1 4e3e8a1ebc [edge-infer] WIP snapshot: uncommitted M0/W1/W4/W5 edge work as of 2026-09-18
Verbatim copy of the edge/omni-edge working tree (tracked diffs + untracked
edge/local, tests/edge/local, docs/edge, pipeline files). Taken so that
vLLM-Omni v0.29.0rc1 can be merged on a branch without touching the
user's working tree. logs_edge/ and benchmark .npy outputs excluded.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
2026-09-18 09:47:29 +08:00

72 lines
2.6 KiB
Python

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""The placement report must measure per-process GPU attribution, not assume it.
[edge-infer W4] The W1 acceptance on native Windows carried a hardcoded "under
WSL2" explanation for an empty NVML process list. These tests pin the repaired
behaviour: the reason names the host it actually ran on, and a successful
attribution (NVML lists one of our PIDs) is reported as data, not a reason.
"""
from __future__ import annotations
import builtins
import sys
import types
import pytest
from vllm_omni.edge.local import engine
def test_host_kind_is_one_of_the_three_hosts():
assert engine._host_kind() in {"windows", "wsl2", "linux"}
def test_reason_names_the_host_when_nvml_is_unavailable(monkeypatch):
real_import = builtins.__import__
def no_pynvml(name, *args, **kwargs):
if name == "pynvml":
raise ImportError("pynvml is not installed in this venv")
return real_import(name, *args, **kwargs)
monkeypatch.setattr(builtins, "__import__", no_pynvml)
out = engine.gpu_attribution_for_tree()
assert out["per_process_gpu_attribution"] is None
assert engine._host_kind() in out["per_process_gpu_attribution_reason"]
assert "WSL2" not in out["per_process_gpu_attribution_reason"] or engine._host_kind() == "wsl2"
def _fake_pynvml(pids):
m = types.ModuleType("pynvml")
m.nvmlInit = lambda: None
m.nvmlShutdown = lambda: None
m.nvmlDeviceGetHandleByIndex = lambda i: object()
m.nvmlDeviceGetComputeRunningProcesses = lambda h: [
types.SimpleNamespace(pid=p, usedGpuMemory=1234) for p in pids
]
return m
def test_attribution_is_data_when_nvml_lists_our_pid(monkeypatch):
import os
monkeypatch.setitem(sys.modules, "pynvml", _fake_pynvml([os.getpid(), 1]))
out = engine.gpu_attribution_for_tree()
assert out["per_process_gpu_attribution"] == [
{"pid": os.getpid(), "used_gpu_memory": 1234}
]
assert out["per_process_gpu_attribution_reason"] is None
def test_reason_states_the_measured_count_when_none_are_ours(monkeypatch):
monkeypatch.setitem(sys.modules, "pynvml", _fake_pynvml([1, 2]))
out = engine.gpu_attribution_for_tree()
assert out["per_process_gpu_attribution"] is None
reason = out["per_process_gpu_attribution_reason"]
assert reason.startswith("nvmlDeviceGetComputeRunningProcesses returned 2 process(es)")
host = engine._host_kind()
expected = {"wsl2": "WSL2", "windows": "WDDM", "linux": "no compute process"}[host]
assert expected in reason