Repository: bentoml/BentoML
Pull request #3779 (2023-04-20 21:39:40)

fix(framework): adjust diffusers device_map default behavior
## What does this PR address?

In previous implementation we use "auto" as the default value of `device_map` when loading the pipeline, and then move the pipeline to gpu if available. In newer version of diffusers (after 0.14), loading a pipeline with `device_map=auto | balanced | etc.` will put the necessary parts of the pipeline to gpu if available, and enable sequential model loading. This means: 1. we don't need to put the pipeline to gpu again 2. we cannot put the pipeline to gpu again because enable sequential model loading disallowing putting the pipeline to gpu manually. (see https://github.com/huggingface/diffusers/issues/2782)

Another fix in this pr is if user pass `device_map=None`, we won't overwrite it with "auto" because `device_map=None` will let user to manually put pipeline to gpu for whatever reason they have. The default value is still "auto" though.

Fixes #(issue)

## 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?



Patch:
diff --git a/src/bentoml/_internal/frameworks/diffusers.py b/src/bentoml/_internal/frameworks/diffusers.py
index 28a7c2c6..5318480d 100644
--- a/src/bentoml/_internal/frameworks/diffusers.py
+++ b/src/bentoml/_internal/frameworks/diffusers.py
@@ -51,6 +51,7 @@ class DiffusersOptions(PartialKwargsModelOptions):
     pipeline_class: type[diffusers.pipelines.DiffusionPipeline] | None = None
     scheduler_class: type[diffusers.SchedulerMixin] | None = None
     torch_dtype: str | torch.dtype | None = None
+    device_map: str | dict[str, int | str | torch.device] | None = None
     custom_pipeline: str | None = None
     enable_xformers: bool | None = None
     enable_attention_slicing: int | str | None = None
@@ -151,13 +152,6 @@ def load_model(
 
     diffusion_model_dir = bento_model.path_of(DIFFUSION_MODEL_FOLDER)
 
-    if (
-        device_map is None
-        and is_torch_version(">=", "1.9.0")
-        and is_accelerate_available()
-    ):
-        device_map = "auto"
-
     if low_cpu_mem_usage is None:
         if is_torch_version(">=", "1.9.0") and is_accelerate_available():
             low_cpu_mem_usage = True
@@ -179,7 +173,10 @@ def load_model(
         pipeline.scheduler = scheduler
 
     if device_id is not None:
-        pipeline = pipeline.to(device_id)
+        # when device_map is not None, we should not move the pipeline to gpu again
+        # see https://github.com/huggingface/diffusers/issues/2782
+        if not (str(device_id).lower().startswith("cuda") and device_map is not None):
+            pipeline = pipeline.to(device_id)
 
     if enable_xformers:
         pipeline.enable_xformers_memory_efficient_attention()
@@ -405,19 +402,21 @@ def get_runnable(bento_model: bentoml.Model) -> t.Type[bentoml.Runnable]:
     Private API: use :obj:`~bentoml.Model.to_runnable` instead.
     """
 
-    partial_kwargs: t.Dict[str, t.Any] = bento_model.info.options.partial_kwargs  # type: ignore
+    bento_options = t.cast(DiffusersOptions, bento_model.info.options)
+    partial_kwargs: dict[str, t.Any] = bento_options.partial_kwargs  # type: ignore
     pipeline_class: type[diffusers.DiffusionPipeline] = (
-        bento_model.info.options.pipeline_class or diffusers.StableDiffusionPipeline
+        bento_options.pipeline_class or diffusers.StableDiffusionPipeline
     )
     scheduler_class: type[
         diffusers.SchedulerMixin
-    ] | None = bento_model.info.options.scheduler_class
-    custom_pipeline: str | None = bento_model.info.options.custom_pipeline
-    _enable_xformers: str | None = bento_model.info.options.enable_xformers
-    enable_attention_slicing: int | str | None = (
-        bento_model.info.options.enable_attention_slicing
-    )
-    _torch_dtype: str | torch.dtype | None = bento_model.info.options.torch_dtype
+    ] | None = bento_options.scheduler_class
+    custom_pipeline: str | None = bento_options.custom_pipeline
+    _enable_xformers: bool | None = bento_options.enable_xformers
+    enable_attention_slicing: int | str | None = bento_options.enable_attention_slicing
+    _torch_dtype: str | torch.dtype | None = bento_options.torch_dtype
+    device_map: str | dict[
+        str, int | str | torch.device
+    ] | None = bento_options.device_map
 
     class DiffusersRunnable(bentoml.Runnable):
         SUPPORTED_RESOURCES = ("nvidia.com/gpu", "cpu")
@@ -443,6 +442,7 @@ def get_runnable(bento_model: bentoml.Model) -> t.Type[bentoml.Runnable]:
             self.pipeline: diffusers.DiffusionPipeline = load_model(
                 bento_model,
                 device_id=device_id,
+                device_map=device_map,
                 pipeline_class=pipeline_class,
                 scheduler_class=scheduler_class,
                 torch_dtype=torch_dtype,
diff --git a/src/bentoml/_internal/marshal/dispatcher.py b/src/bentoml/_internal/marshal/dispatcher.py
index cad15209..e1f7a719 100644
--- a/src/bentoml/_internal/marshal/dispatcher.py
+++ b/src/bentoml/_internal/marshal/dispatcher.py
@@ -85,6 +85,12 @@ class Optimizer:
             self.trigger_refresh()
 
     def trigger_refresh(self):
+        if not self.o_stat:
+            logger.debug(
+                "o_stat is empty, skip dynamic batching optimizer params update"
+            )
+            return
+
         x = tuple((i, 1) for i, _, _ in self.o_stat)
         y = tuple(i for _, i, _ in self.o_stat)
 
diff --git a/src/bentoml/_internal/service/loader.py b/src/bentoml/_internal/service/loader.py
index 05f730d6..681bc2eb 100644
--- a/src/bentoml/_internal/service/loader.py
+++ b/src/bentoml/_internal/service/loader.py
@@ -102,7 +102,7 @@ def import_service(
                 '"<module>:<attribute>" or "<module>'
             )
 
-        if os.path.exists(import_path):
+        if os.path.isfile(import_path):
             import_path = os.path.realpath(import_path)
             # Importing from a module file path:
             if not import_path.startswith(working_dir):


Test patch:
diff --git a/tests/unit/_internal/service_loader/bento_service_in_module/__init__.py b/tests/unit/_internal/service_loader/bento_service_in_module/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/tests/unit/_internal/service_loader/bento_service_in_module/multi_service.py b/tests/unit/_internal/service_loader/bento_service_in_module/multi_service.py
new file mode 100644
index 00000000..b9547afb
--- /dev/null
+++ b/tests/unit/_internal/service_loader/bento_service_in_module/multi_service.py
@@ -0,0 +1,4 @@
+import bentoml
+
+svc_i = bentoml.Service("test-bento-service-in-module-i")
+svc_ii = bentoml.Service("test-bento-service-in-module-ii")
diff --git a/tests/unit/_internal/service_loader/bento_service_in_module/single_service.py b/tests/unit/_internal/service_loader/bento_service_in_module/single_service.py
new file mode 100644
index 00000000..81ed9e27
--- /dev/null
+++ b/tests/unit/_internal/service_loader/bento_service_in_module/single_service.py
@@ -0,0 +1,3 @@
+import bentoml
+
+svc = bentoml.Service("test-bento-service-in-module-single")
diff --git a/tests/unit/_internal/service_loader/multi_service_in_package/__init__.py b/tests/unit/_internal/service_loader/multi_service_in_package/__init__.py
new file mode 100644
index 00000000..79cabfab
--- /dev/null
+++ b/tests/unit/_internal/service_loader/multi_service_in_package/__init__.py
@@ -0,0 +1,4 @@
+import bentoml
+
+svc_i = bentoml.Service("test-multi-service-in-package-i")
+svc_ii = bentoml.Service("test-multi-service-in-package-ii")
diff --git a/tests/unit/_internal/service_loader/single_service_in_package/__init__.py b/tests/unit/_internal/service_loader/single_service_in_package/__init__.py
new file mode 100644
index 00000000..4ca53400
--- /dev/null
+++ b/tests/unit/_internal/service_loader/single_service_in_package/__init__.py
@@ -0,0 +1,3 @@
+import bentoml
+
+svc = bentoml.Service("test-bento-service-in-package")
diff --git a/tests/unit/_internal/service_loader/test_import_service.py b/tests/unit/_internal/service_loader/test_import_service.py
new file mode 100644
index 00000000..b7ed6a8f
--- /dev/null
+++ b/tests/unit/_internal/service_loader/test_import_service.py
@@ -0,0 +1,59 @@
+import pytest
+
+from bentoml.exceptions import ImportServiceError
+from bentoml._internal.service.loader import import_service
+
+
+@pytest.mark.usefixtures("change_test_dir")
+def test_load_by_service_file():
+    svc = import_service("bento_service_in_module/single_service.py")
+    assert svc.name == "test-bento-service-in-module-single"
+
+    svc = import_service("./bento_service_in_module/single_service.py:svc")
+    assert svc.name == "test-bento-service-in-module-single"
+
+    svc = import_service("./single_service_in_package/__init__.py")
+    assert svc.name == "test-bento-service-in-package"
+
+    svc = import_service("./single_service_in_package/__init__.py:svc")
+    assert svc.name == "test-bento-service-in-package"
+
+
+@pytest.mark.usefixtures("change_test_dir")
+def test_load_by_module_name():
+    svc = import_service("bento_service_in_module.single_service")
+    assert svc.name == "test-bento-service-in-module-single"
+
+    svc = import_service("bento_service_in_module.single_service:svc")
+    assert svc.name == "test-bento-service-in-module-single"
+
+    svc = import_service("single_service_in_package")
+    assert svc.name == "test-bento-service-in-package"
+
+    svc = import_service("single_service_in_package:svc")
+    assert svc.name == "test-bento-service-in-package"
+
+
+@pytest.mark.usefixtures("change_test_dir")
+def test_load_multi_service_module():
+    svc = import_service("bento_service_in_module.multi_service:svc_i")
+    assert svc.name == "test-bento-service-in-module-i"
+
+    svc = import_service("bento_service_in_module.multi_service:svc_ii")
+    assert svc.name == "test-bento-service-in-module-ii"
+
+    with pytest.raises(
+        ImportServiceError, match="Multiple Service instances found in module"
+    ):
+        import_service("bento_service_in_module.multi_service")
+
+    svc = import_service("multi_service_in_package:svc_i")
+    assert svc.name == "test-multi-service-in-package-i"
+
+    svc = import_service("multi_service_in_package:svc_ii")
+    assert svc.name == "test-multi-service-in-package-ii"
+
+    with pytest.raises(
+        ImportServiceError, match="Multiple Service instances found in module"
+    ):
+        import_service("multi_service_in_package")