fix: disable GPU allocation if CUDA_VISIBLE_DEVICES is set (#5368)

* fix: disable GPU allocation if CUDA_VISIBLE_DEVICES is set

Signed-off-by: Frost Ming <me@frostming.com>

* fix: onnx tests

Signed-off-by: Frost Ming <me@frostming.com>

Author: frostming

noxfile.py

@@ -28,7 +28,7 @@
     ],
     "keras": ["keras>=3.4"],
     "lightgbm": ["lightgbm"],
-    "onnx": ["onnx", "onnxruntime", "skl2onnx"],
+    "onnx": ["onnx==1.17.0", "onnxruntime", "skl2onnx"],
     "picklable_model": [],
     "pytorch": [],
     "pytorch_lightning": ["lightning"],

src/_bentoml_impl/server/allocator.py

@@ -25,6 +25,13 @@ def __init__(self) -> None:
             for _ in range(self.remaining_gpus)
         ]
 
+    @staticmethod
+    def gpu_allocation_disabled() -> bool:
+        return (
+            DISABLE_GPU_ALLOCATION_ENV in os.environ
+            or "CUDA_VISIBLE_DEVICES" in os.environ
+        )
+
     def assign_gpus(self, count: float) -> list[int]:
         if count > self.remaining_gpus:
             warnings.warn(
@@ -101,7 +108,7 @@ def get_worker_env(
                 return num_workers, worker_env
             else:  # workers is a number
                 num_workers = workers
-        if num_gpus and DISABLE_GPU_ALLOCATION_ENV not in os.environ:
+        if num_gpus and not self.gpu_allocation_disabled():
             assigned = self.assign_gpus(num_gpus)
             # assign gpus to all workers
             worker_env = [