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):
|
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":
|
if gpu_arg.lower() == "cpu":
|
||||||
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
|
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
|
||||||
return "cpu"
|
return "cpu"
|
||||||
if gpu_arg.lower() == "auto":
|
if gpu_arg.lower() == "auto":
|
||||||
return None # let ultralytics/torch decide
|
return None # let ultralytics/torch decide
|
||||||
os.environ["CUDA_VISIBLE_DEVICES"] = gpu_arg
|
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):
|
def ensure_directory(directory):
|
||||||
@@ -300,7 +311,7 @@ def main():
|
|||||||
from ultralytics import YOLO
|
from ultralytics import YOLO
|
||||||
import cv2
|
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 "
|
print("CUDA requested but not visible after pinning "
|
||||||
f"(--gpu {args.gpu}). Falling back to CPU.")
|
f"(--gpu {args.gpu}). Falling back to CPU.")
|
||||||
device = "cpu"
|
device = "cpu"
|
||||||
|
|||||||
+9
-1
@@ -40,8 +40,16 @@ def main():
|
|||||||
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
|
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
|
||||||
device = "cpu"
|
device = "cpu"
|
||||||
elif args.gpu.lower() != "auto":
|
elif args.gpu.lower() != "auto":
|
||||||
|
# Lock the pin with a warm CUDA alloc BEFORE ultralytics runs
|
||||||
|
# select_device(), which would otherwise rewrite
|
||||||
|
# CUDA_VISIBLE_DEVICES and un-pin us onto GPU0.
|
||||||
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
|
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
|
||||||
device = "cuda"
|
import torch
|
||||||
|
if torch.cuda.is_available():
|
||||||
|
torch.zeros(1, device="cuda")
|
||||||
|
device = "cuda:0"
|
||||||
|
else:
|
||||||
|
device = "cpu"
|
||||||
else:
|
else:
|
||||||
device = None
|
device = None
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user