wehub-resource-sync
90 lines
2.8 KiB
Python
90 lines
2.8 KiB
Python
import os
|
|
import sys
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from ray._common.test_utils import wait_for_condition
|
|
from ray.dashboard.modules.reporter.gpu_providers import (
|
|
NvidiaGpuProvider,
|
|
)
|
|
|
|
|
|
def _candidate_pids() -> set[int]:
|
|
# NVML may report PIDs from a different namespace when this runs in a
|
|
# container. /proc/self/status exposes the PID namespace chain, so accept
|
|
# any PID that can refer to this process.
|
|
pids = {os.getpid()}
|
|
try:
|
|
with open("/proc/self/status") as f:
|
|
for line in f:
|
|
if line.startswith("NSpid:"):
|
|
pids.update(int(pid) for pid in line.split()[1:])
|
|
break
|
|
except (OSError, ValueError):
|
|
pass
|
|
return pids
|
|
|
|
|
|
def _find_process_info(gpu_utilization, candidate_pids, tensor_size_mb):
|
|
fallback_match = None
|
|
|
|
for single_gpu_info in gpu_utilization:
|
|
procs = single_gpu_info["processes_pids"] or {}
|
|
for pid in candidate_pids:
|
|
if pid in procs:
|
|
return single_gpu_info, procs[pid]
|
|
|
|
for process_info in procs.values():
|
|
if process_info["gpu_memory_usage"] >= tensor_size_mb:
|
|
if (
|
|
fallback_match is None
|
|
or process_info["gpu_memory_usage"]
|
|
< fallback_match[1]["gpu_memory_usage"]
|
|
):
|
|
fallback_match = (single_gpu_info, process_info)
|
|
|
|
return fallback_match
|
|
|
|
|
|
def test_per_process_gpu_memory_usage_and_total_gpu_memory_usage():
|
|
"""
|
|
Allocate a large tensor on GPU, then verify the provider reports process
|
|
memory and overall gpu memory consistent with that.
|
|
"""
|
|
|
|
device = torch.device("cuda:0")
|
|
tensor_size_mb = 2048
|
|
num_elements = tensor_size_mb * 1024 * 1024
|
|
tensor = torch.zeros(num_elements, dtype=torch.int8, device=device)
|
|
torch.cuda.synchronize(device)
|
|
assert tensor.size(0) == num_elements
|
|
|
|
provider = NvidiaGpuProvider()
|
|
candidate_pids = _candidate_pids()
|
|
|
|
def check_utilization():
|
|
result = provider.get_gpu_utilization()
|
|
assert len(result) > 0
|
|
|
|
# Find the process and GPU that correspond to the tensor allocation.
|
|
match = _find_process_info(result, candidate_pids, tensor_size_mb)
|
|
assert match is not None
|
|
gpu_info, process_info = match
|
|
|
|
reported_proc_mb = process_info["gpu_memory_usage"]
|
|
# Proc memory usage should be at least the tensor size and within a
|
|
# reasonable amount of CUDA context overhead.
|
|
assert (
|
|
reported_proc_mb >= tensor_size_mb
|
|
and reported_proc_mb - tensor_size_mb < 1024
|
|
)
|
|
assert gpu_info["memory_used"] >= reported_proc_mb
|
|
return True
|
|
|
|
wait_for_condition(check_utilization)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(pytest.main(["-sv", __file__]))
|