[CI] Exclude current process memory from GPU idle check (#31571)
This commit is contained in:
@@ -768,6 +768,11 @@ def _subprocess_popen_with_outputs(
|
|||||||
env: Optional[dict],
|
env: Optional[dict],
|
||||||
return_stdout_stderr: Optional[tuple],
|
return_stdout_stderr: Optional[tuple],
|
||||||
) -> subprocess.Popen:
|
) -> subprocess.Popen:
|
||||||
|
# Release allocator-cached GPU memory to the driver before spawning a
|
||||||
|
# server: cached blocks stay cudaMalloc'd and shrink the child's memory.
|
||||||
|
if torch.cuda.is_initialized():
|
||||||
|
torch.cuda.empty_cache()
|
||||||
|
|
||||||
if not return_stdout_stderr:
|
if not return_stdout_stderr:
|
||||||
return subprocess.Popen(command, stdout=None, stderr=None, env=env)
|
return subprocess.Popen(command, stdout=None, stderr=None, env=env)
|
||||||
|
|
||||||
@@ -2201,6 +2206,7 @@ def _visible_gpu_indices(pynvml) -> List[int]:
|
|||||||
|
|
||||||
|
|
||||||
def _collect_busy_gpu_reports(pynvml, gpu_indices: List[int]) -> List[str]:
|
def _collect_busy_gpu_reports(pynvml, gpu_indices: List[int]) -> List[str]:
|
||||||
|
self_pid = os.getpid()
|
||||||
reports = []
|
reports = []
|
||||||
for index in gpu_indices:
|
for index in gpu_indices:
|
||||||
handle = pynvml.nvmlDeviceGetHandleByIndex(index)
|
handle = pynvml.nvmlDeviceGetHandleByIndex(index)
|
||||||
@@ -2209,10 +2215,25 @@ def _collect_busy_gpu_reports(pynvml, gpu_indices: List[int]) -> List[str]:
|
|||||||
continue
|
continue
|
||||||
try:
|
try:
|
||||||
procs = pynvml.nvmlDeviceGetComputeRunningProcesses(handle)
|
procs = pynvml.nvmlDeviceGetComputeRunningProcesses(handle)
|
||||||
proc_info = ", ".join(
|
|
||||||
f"pid={proc.pid} {_format_gib(proc.usedGpuMemory)}" for proc in procs
|
|
||||||
)
|
|
||||||
except pynvml.NVMLError:
|
except pynvml.NVMLError:
|
||||||
|
procs = None
|
||||||
|
if procs is not None:
|
||||||
|
# Discount our own usage: the caching allocator retains memory
|
||||||
|
# across test classes, and waiting on ourselves never succeeds.
|
||||||
|
self_used = sum(
|
||||||
|
proc.usedGpuMemory or 0 for proc in procs if proc.pid == self_pid
|
||||||
|
)
|
||||||
|
if used_bytes - self_used < _GPU_IDLE_USED_MEMORY_THRESHOLD:
|
||||||
|
continue
|
||||||
|
proc_info = ", ".join(
|
||||||
|
f"pid={proc.pid} {_format_gib(proc.usedGpuMemory)}"
|
||||||
|
for proc in procs
|
||||||
|
if proc.pid != self_pid
|
||||||
|
) or (
|
||||||
|
f"no other compute processes;"
|
||||||
|
f" self pid={self_pid} holds {_format_gib(self_used)}"
|
||||||
|
)
|
||||||
|
else:
|
||||||
proc_info = ""
|
proc_info = ""
|
||||||
reports.append(
|
reports.append(
|
||||||
f"GPU {index} uses {_format_gib(used_bytes)}"
|
f"GPU {index} uses {_format_gib(used_bytes)}"
|
||||||
|
|||||||
Reference in New Issue
Block a user