71 lines
2.6 KiB
Python
71 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
|