Repository: bentoml/BentoML
Pull request: feat: explicit GPU runner mappings
## What does this PR address?
<!--
Thanks for sending a pull request!

Congrats for making it this far! Here's a 🍱 for you. There are still a few steps ahead.

Please make sure to read the contribution guidelines, then fill out the blanks below before requesting a code review.

Name your Pull Request with one of the following prefixes, e.g. "feat: add support for PyTorch", to indicate the type of changes proposed. This is based on the [Conventional Commits specification](https://www.conventionalcommits.org/en/v1.0.0/#summary).
  - feat: (new feature for the user, not a new feature for build script)
  - fix: (bug fix for the user, not a fix to a build script)
  - docs: (changes to the documentation)
  - style: (formatting, missing semicolons, etc; no production code change)
  - refactor: (refactoring production code, eg. renaming a variable)
  - perf: (code changes that improve performance)
  - test: (adding missing tests, refactoring tests; no production code change)
  - chore: (updating grunt tasks etc; no production code change)
  - build: (changes that affect the build system or external dependencies)
  - ci: (changes to configuration files and scripts)
  - revert: (reverts a previous commit)

Describe your changes in detail. Attach screenshots here if appropriate.

Once you're done with this, someone from BentoML team or community member will help review your PR (see "Who can help review?" section for potential reviewers.). If no one has reviewed your PR after a week have passed, don't hesitate to post a new comment and ping @-the same person. Notifications sometimes get lost 🥲.
-->

<!-- Remove if not applicable -->
Fixes #2770 
With this change, the `NvidiaGpuResource` will accept and return a list that specifies the exact GPUs to use. This enables the user to configure exactly which GPU to map to each runner worker.

Runner currently supports:
```yaml
runners:
  resources: 
    nvidia.com/gpu: 4 => still valid but will be automatically converted to use GPUs [0, 1, 2, 3] internally.
```

This PR introduces:
```yaml
  resources: 
    nvidia.com/gpu: [2, 4] => only use GPU [2, 4] for runner worker 1 and 2
```

## Before submitting:
<!--- Go over all the following points, and put an `x` in all the boxes that apply. -->
<!--- If you're unsure about any of these, don't hesitate to ask. We're here to help! -->
<!--- If you plan to update documentation or tests in follow-up, please note -->
- [ ] Does the Pull Request follow [Conventional Commits specification](https://www.conventionalcommits.org/en/v1.0.0/#summary) naming? Here are [GitHub's
guide](https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/proposing-changes-to-your-work-with-pull-requests/creating-a-pull-request) on how to create a pull request.
- [ ] Does the code follow BentoML's code style, both `make format` and `make lint` script have passed ([instructions](https://github.com/bentoml/BentoML/blob/main/DEVELOPMENT.md#style-check-auto-formatting-type-checking))?
- [ ] Did you read through [contribution guidelines](https://github.com/bentoml/BentoML/blob/main/CONTRIBUTING.md#ways-to-contribute) and follow [development guidelines](https://github.com/bentoml/BentoML/blob/main/DEVELOPMENT.md#start-developing)?
- [ ] Did your changes require updates to the documentation? Have you updated
  those accordingly? Here are [documentation guidelines](https://github.com/bentoml/BentoML/tree/main/docs) and [tips on writting docs](https://github.com/bentoml/BentoML/tree/main/docs#writing-documentation).
- [ ] Did you write tests to cover your changes?

## Who can help review?

Feel free to tag members/contributors who can help review your PR.
<!--
Feel free to ping any of the BentoML members for help on your issue, but don't ping more than three people 😊.
If you know how to use git blame, that is probably the easiest way.

Team members that you can ping:
- @parano
- @yubozhao
- @bojiang
- @ssheng
- @aarnphm
- @sauyon
- @larme
- @yetone
- @jjmachan
-->


Problem statement:
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_

Hints from issue discussion:

Commit message: add an assert if you wannt check

Patch to review:
diff --git a/bentoml/_internal/resource.py b/bentoml/_internal/resource.py
index 8dbc4182d28..40580b4ceff 100644
--- 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
index f88524dd18f..028dcb6846b 100644
--- 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,
diff --git a/tests/unit/_internal/runner/test_strategy.py b/tests/unit/_internal/runner/test_strategy.py
new file mode 100644
index 00000000000..b852dde6b02
--- /dev/null
+++ b/tests/unit/_internal/runner/test_strategy.py
@@ -0,0 +1,31 @@
+import os
+import typing as t
+
+import bentoml
+from bentoml._internal.runner import strategy
+from bentoml._internal.resource import get_resource as resource_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 resource_get_resource(x, y, validate=False)
+
+
+def test_default_gpu_strategy(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
index a50eb61ed4d..3053cc37869 100644
--- a/tests/unit/_internal/test_configuration.py
+++ b/tests/unit/_internal/test_configuration.py
@@ -29,13 +29,18 @@
 """
 
 
-def test_bentoml_configuration_runner_override():
+def get_BentoMLConfiguration_from_str(config_str: str):
     tmpfile = NamedTemporaryFile(mode="w+", delete=False)
-    tmpfile.write(OVERRIDE_RUNNERS)
+    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():
+    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
index f5eaca6499d..d5214e35c68 100644
--- 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)


Code review comments:
diff_hunk: @@ -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 = """
path: tests/unit/_internal/test_configuration.py
text: @user1:
```suggestion
    GPU_INDEX = """\
```

diff_hunk: @@ -29,13 +29,18 @@
 """
 
 
-def test_bentoml_configuration_runner_override():
+def get_BentoMLConfiguration_from_str(config_str: str):
path: tests/unit/_internal/test_configuration.py
text: @user1:
nit style: keep it snake case

diff_hunk: @@ -0,0 +1,31 @@
+import os
path: tests/unit/_internal/runner/test_strategy.py
text: @user1:
also lets try to use `from __future__ import annotations`