Repository: bentoml/BentoML

Issue:
feat: Runner to GPUs mapping
### Feature request

The default scheduling strategy implementation schedules the same number of runner (`nvidia.com/gpu` supported) instances as the number of available GPUs. If multiple types of runners are present in the service, they will each be scheduled with the same number of instances as the available GPUs. Ideally, we should enable device isolation and allow users to configure GPU device affinity to the runner types.

### Motivation

Allowing runner to GPUs mapping can reduce the number of runner instances required while maximizing GPU resource utilization.

### Other

_No response_


Discussion:

Resolving patch:
diff --git a/bentoml/_internal/resource.py b/bentoml/_internal/resource.py
--- a/bentoml/_internal/resource.py
+++ b/bentoml/_internal/resource.py
@@ -20,7 +20,9 @@
 T = t.TypeVar("T")
 
 
-def get_resource(resources: dict[str, t.Any], resource_kind: str) -> t.Any:
+def get_resource(
+    resources: dict[str, t.Any], resource_kind: str, validate: bool = True
+) -> t.Any:
     if resource_kind not in _RESOURCE_REGISTRY:
         raise BentoMLConfigException(f"Unknown resource kind '{resource_kind}'.")
 
@@ -31,7 +33,8 @@ def get_resource(resources: dict[str, t.Any], resource_kind: str) -> t.Any:
             return resource.from_system()
         else:
             res = resource.from_spec(resources[resource_kind])
-            resource.validate(res)
+            if validate:
+                resource.validate(res)
             return res
     else:
         return None
@@ -211,17 +214,23 @@ def query_os_cpu_count() -> int:
 #         raise ValueError(f"Invalid MEM resource limit '{mem}'")
 
 
-class NvidiaGpuResource(Resource[int], resource_id="nvidia.com/gpu"):
+class NvidiaGpuResource(Resource[t.List[int]], resource_id="nvidia.com/gpu"):
     @classmethod
-    def from_spec(cls, spec: t.Any) -> int:
-        if not isinstance(spec, (int, str)):
-            raise TypeError("NVidia GPU resource limit must be int or str")
-
-        if isinstance(spec, int):
-            return spec
+    def from_spec(cls, spec: int | str | list[int | str]) -> list[int]:
+        if not isinstance(spec, (int, str, list)):
+            raise TypeError(
+                "NVidia GPU device IDs must be int, str or a list specifing the exact GPUs to use."
+            )
 
         try:
-            return int(spec)
+            if isinstance(spec, int):
+                if spec < -1:
+                    raise ValueError
+                return list(range(spec))
+            elif isinstance(spec, str):
+                return cls.from_spec(int(spec))
+            else:
+                return [int(x) for x in spec]
         except ValueError:
             raise BentoMLConfigException(
                 f"Invalid NVidia GPU resource limit '{spec}'. "
@@ -229,7 +238,7 @@ def from_spec(cls, spec: t.Any) -> int:
 
     @classmethod
     @functools.lru_cache(maxsize=1)
-    def from_system(cls) -> int:
+    def from_system(cls) -> list[int]:
         """
         query nvidia gpu count, available on Windows and Linux
         """
@@ -238,9 +247,10 @@ def from_system(cls) -> int:
         try:
             pynvml.nvmlInit()
             device_count = pynvml.nvmlDeviceGetCount()
-            return device_count
+            return list(range(device_count))
         except (pynvml.nvml.NVMLError, OSError):
-            return 0
+            logger.warning("GPU not detected. Unable to initialize pynvml lib.")
+            return []
         finally:
             try:
                 pynvml.nvmlShutdown()
@@ -248,12 +258,12 @@ def from_system(cls) -> int:
                 pass
 
     @classmethod
-    def validate(cls, val: int):
-        if val < -1:
-            raise BentoMLConfigException(f"Invalid negative GPU resource limit {val}.")
-        if val > cls.from_system():
+    def validate(cls, val: t.List[int]):
+        if any([gpu_index < 0 for gpu_index in val]):
+            raise BentoMLConfigException(f"Negative GPU device in {val}.")
+        if any([gpu_index >= len(cls.from_system()) for gpu_index in val]):
             raise BentoMLConfigException(
-                f"GPU resource limit {val} is greater than the system available: {cls.from_system()}"
+                f"GPU device index in {val} is greater than the system available: {cls.from_system()}"
             )
 
 
diff --git a/bentoml/_internal/runner/strategy.py b/bentoml/_internal/runner/strategy.py
--- a/bentoml/_internal/runner/strategy.py
+++ b/bentoml/_internal/runner/strategy.py
@@ -63,10 +63,10 @@ def get_worker_count(
         nvidia_gpus = get_resource(resource_request, "nvidia.com/gpu")
         if (
             nvidia_gpus is not None
-            and nvidia_gpus > 0
+            and len(nvidia_gpus) > 0
             and "nvidia.com/gpu" in runnable_class.SUPPORTED_RESOURCES
         ):
-            return math.ceil(nvidia_gpus)
+            return len(nvidia_gpus)
 
         # use CPU
         cpus = get_resource(resource_request, "cpu")
@@ -101,10 +101,10 @@ def setup_worker(
         nvidia_gpus = get_resource(resource_request, "nvidia.com/gpu")
         if (
             nvidia_gpus is not None
-            and nvidia_gpus > 0
+            and len(nvidia_gpus) > 0
             and "nvidia.com/gpu" in runnable_class.SUPPORTED_RESOURCES
         ):
-            os.environ["CUDA_VISIBLE_DEVICES"] = str(worker_index - 1)
+            os.environ["CUDA_VISIBLE_DEVICES"] = str(nvidia_gpus[worker_index - 1])
             logger.info(
                 "Setting up worker: set CUDA_VISIBLE_DEVICES to %s",
                 worker_index - 1,


Test patch:
diff --git a/tests/unit/_internal/runner/test_strategy.py b/tests/unit/_internal/runner/test_strategy.py
new file mode 100644
--- /dev/null
+++ b/tests/unit/_internal/runner/test_strategy.py
@@ -0,0 +1,36 @@
+from __future__ import annotations
+
+import os
+import typing as t
+
+if t.TYPE_CHECKING:
+    from _pytest.monkeypatch import MonkeyPatch
+
+import bentoml
+from bentoml._internal.runner import strategy
+from bentoml._internal.resource import get_resource
+from bentoml._internal.runner.strategy import DefaultStrategy
+
+
+class GPURunnable(bentoml.Runnable):
+    SUPPORTED_RESOURCES = ("nvidia.com/gpu",)
+
+
+def unvalidated_get_resource(x: t.Dict[str, t.Any], y: str):
+    return get_resource(x, y, validate=False)
+
+
+def test_default_gpu_strategy(monkeypatch: MonkeyPatch):
+    monkeypatch.setattr(strategy, "get_resource", unvalidated_get_resource)
+    assert DefaultStrategy.get_worker_count(GPURunnable, {"nvidia.com/gpu": 2}) == 2
+    assert DefaultStrategy.get_worker_count(GPURunnable, {"nvidia.com/gpu": 0}) == 1
+    assert (
+        DefaultStrategy.get_worker_count(GPURunnable, {"nvidia.com/gpu": [2, 7]}) == 2
+    )
+
+    DefaultStrategy.setup_worker(GPURunnable, {"nvidia.com/gpu": 2}, 1)
+    assert os.environ.get("CUDA_VISIBLE_DEVICES") == "0"
+    DefaultStrategy.setup_worker(GPURunnable, {"nvidia.com/gpu": 2}, 2)
+    assert os.environ.get("CUDA_VISIBLE_DEVICES") == "1"
+    DefaultStrategy.setup_worker(GPURunnable, {"nvidia.com/gpu": [2, 7]}, 2)
+    assert os.environ.get("CUDA_VISIBLE_DEVICES") == "7"
diff --git a/tests/unit/_internal/test_configuration.py b/tests/unit/_internal/test_configuration.py
--- a/tests/unit/_internal/test_configuration.py
+++ b/tests/unit/_internal/test_configuration.py
@@ -2,7 +2,19 @@
 
 from bentoml._internal.configuration.containers import BentoMLConfiguration
 
-OVERRIDE_RUNNERS = """
+
+def get_bentomlconfiguration_from_str(config_str: str):
+    tmpfile = NamedTemporaryFile(mode="w+", delete=False)
+    tmpfile.write(config_str)
+    tmpfile.flush()
+    tmpfile.close()
+
+    bentoml_cfg = BentoMLConfiguration(override_config_file=tmpfile.name).as_dict()
+    return bentoml_cfg
+
+
+def test_bentoml_configuration_runner_override():
+    OVERRIDE_RUNNERS = """\
 runners:
     batching:
         enabled: False
@@ -28,14 +40,7 @@
                 enabled: True
 """
 
-
-def test_bentoml_configuration_runner_override():
-    tmpfile = NamedTemporaryFile(mode="w+", delete=False)
-    tmpfile.write(OVERRIDE_RUNNERS)
-    tmpfile.flush()
-    tmpfile.close()
-
-    bentoml_cfg = BentoMLConfiguration(override_config_file=tmpfile.name).as_dict()
+    bentoml_cfg = get_bentomlconfiguration_from_str(OVERRIDE_RUNNERS)
     runner_cfg = bentoml_cfg["runners"]
 
     # test_runner_1
@@ -66,3 +71,22 @@ def test_bentoml_configuration_runner_override():
     assert test_runner_batching["batching"]["max_batch_size"] == 10
     assert test_runner_batching["logging"]["access"]["enabled"] is True
     assert test_runner_batching["resources"]["cpu"] == 4  # should use global
+
+
+def test_runner_gpu_configuration():
+    GPU_INDEX = """\
+runners:
+    resources:
+        nvidia.com/gpu: [1, 2, 4]
+"""
+    bentoml_cfg = get_bentomlconfiguration_from_str(GPU_INDEX)
+    assert bentoml_cfg["runners"]["resources"] == {"nvidia.com/gpu": [1, 2, 4]}
+
+    GPU_INDEX_WITH_STRING = """\
+runners:
+    resources:
+        nvidia.com/gpu: "[1, 2, 4]"
+"""
+    bentoml_cfg = get_bentomlconfiguration_from_str(GPU_INDEX_WITH_STRING)
+    # this behaviour can be confusing
+    assert bentoml_cfg["runners"]["resources"] == {"nvidia.com/gpu": "[1, 2, 4]"}
diff --git a/tests/unit/_internal/test_resource.py b/tests/unit/_internal/test_resource.py
--- a/tests/unit/_internal/test_resource.py
+++ b/tests/unit/_internal/test_resource.py
@@ -64,19 +64,31 @@ def test_CpuResource():
 
 
 def test_NvidiaGpuResource():
-    assert NvidiaGpuResource.from_system() >= 0  # TODO: real from_system tests
+    assert len(NvidiaGpuResource.from_system()) >= 0  # TODO: real from_system tests
 
     with pytest.raises(BentoMLConfigException):
-        NvidiaGpuResource.validate(NvidiaGpuResource.from_system() + 1)
+        NvidiaGpuResource.validate(NvidiaGpuResource.from_system() + [1])
     with pytest.raises(BentoMLConfigException):
-        NvidiaGpuResource.validate(-2)
+        NvidiaGpuResource.validate([-2])
+    with pytest.raises(BentoMLConfigException):
+        NvidiaGpuResource.validate([-1])
 
-    NvidiaGpuResource.validate(-1)
-    NvidiaGpuResource.validate(0)
+    NvidiaGpuResource.validate([])
     # NvidiaGpuResource.validate(1)  # TODO: work out how to skip this on systems with no GPU
 
-    assert NvidiaGpuResource.from_spec(1) == 1
-    assert NvidiaGpuResource.from_spec("5") == 5.0
+    assert NvidiaGpuResource.from_spec(1) == [0]
+    assert NvidiaGpuResource.from_spec("5") == [0, 1, 2, 3, 4]
+    assert NvidiaGpuResource.from_spec(1) == [0]
+    assert NvidiaGpuResource.from_spec(2) == [0, 1]
+    assert NvidiaGpuResource.from_spec("3") == [0, 1, 2]
+    assert NvidiaGpuResource.from_spec([1, 3]) == [1, 3]
+    assert NvidiaGpuResource.from_spec(["1", "3"]) == [1, 3]
+    assert NvidiaGpuResource.from_spec(-1) == []
+    assert NvidiaGpuResource.from_spec("-1") == []
+
+    with pytest.raises(BentoMLConfigException):
+        # Currently this is not supported and is considered invalid
+        assert NvidiaGpuResource.from_spec("[1, 2, 3]") == [1, 2, 3]
 
     with pytest.raises(TypeError):
         NvidiaGpuResource.from_spec((1, 2, 3))
@@ -84,5 +96,11 @@ def test_NvidiaGpuResource():
     with pytest.raises(BentoMLConfigException):
         NvidiaGpuResource.from_spec("100m")
 
+    with pytest.raises(BentoMLConfigException):
+        assert NvidiaGpuResource.from_spec(-2)
+
+    with pytest.raises(BentoMLConfigException):
+        assert NvidiaGpuResource.from_spec("-2")
+
     with pytest.raises(TypeError):
         NvidiaGpuResource.from_spec(1.5)