Lock GPU pin with warm CUDA alloc before ultralytics select_device runs
Ultralytics select_device() rewrites CUDA_VISIBLE_DEVICES on first predict() when the CUDA context is uninitialized: --gpu 1 silently became a two-GPU cuda:0,1 context that OOMed on GPU0 (Strata's card). Initializing torch and touching cuda once after setting CVD locks the pin; both scripts now use device cuda:0 within the pin.
This commit is contained in:
+14
-3
@@ -121,14 +121,25 @@ def parse_arguments():
|
||||
|
||||
|
||||
def resolve_gpu(gpu_arg):
|
||||
"""Pin CUDA before torch initializes. Returns torch device string."""
|
||||
"""Pin CUDA before torch initializes. Returns torch device string.
|
||||
|
||||
NOTE: ultralytics' select_device() rewrites CUDA_VISIBLE_DEVICES on
|
||||
first predict() if the CUDA context is not yet initialized, silently
|
||||
un-pinning us onto GPU0. We force-initialize a tiny CUDA allocation
|
||||
here so our CVD pin is locked in before ultralytics can touch it.
|
||||
"""
|
||||
if gpu_arg.lower() == "cpu":
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
|
||||
return "cpu"
|
||||
if gpu_arg.lower() == "auto":
|
||||
return None # let ultralytics/torch decide
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = gpu_arg
|
||||
return "cuda"
|
||||
import torch
|
||||
if torch.cuda.is_available():
|
||||
torch.zeros(1, device="cuda") # lock the context to our pin
|
||||
else:
|
||||
return "cpu"
|
||||
return "cuda:0" # within our CVD, physical GPU N is always cuda:0
|
||||
|
||||
|
||||
def ensure_directory(directory):
|
||||
@@ -300,7 +311,7 @@ def main():
|
||||
from ultralytics import YOLO
|
||||
import cv2
|
||||
|
||||
if device == "cuda" and not torch.cuda.is_available():
|
||||
if device == "cuda:0" and not torch.cuda.is_available():
|
||||
print("CUDA requested but not visible after pinning "
|
||||
f"(--gpu {args.gpu}). Falling back to CPU.")
|
||||
device = "cpu"
|
||||
|
||||
Reference in New Issue
Block a user