chore: import upstream snapshot with attribution

python/ray/tests/test_gpu_provider_on_gpu.py

@@ -0,0 +1,89 @@
+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__]))