Repository: bentoml/BentoML
Pull request #3788 (2023-04-25 02:44:21)

chore(dispatcher): cancel jobs with a for loop
Cleans up a minor strange code segment that used `IndexError` to exit out of a loop popping the queue, which I don't think is necessary.

@bojiang does this look fine?


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..9d14f8a0 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)
 
@@ -145,12 +151,8 @@ class CorkDispatcher:
     def shutdown(self):
         if self._controller is not None:
             self._controller.cancel()
-        try:
-            while True:
-                fut = self._queue.pop().future
-                fut.cancel()
-        except IndexError:
-            pass
+        for job in self._queue:
+            job.future.cancel()
 
     @cached_property
     def _loop(self):
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/requirements/tests-requirements.txt b/requirements/tests-requirements.txt
index a47e24b7..61eaa636 100644
--- a/requirements/tests-requirements.txt
+++ b/requirements/tests-requirements.txt
@@ -7,12 +7,12 @@ black[jupyter]==22.12.0
 coverage[toml]==7.2.3
 setuptools>=63
 isort==5.11.4
-ruff==0.0.262
+ruff==0.0.263
 pytest-cov==4.0.0
 pytest==7.3.1
 pytest-xdist[psutil]==3.2.1
 pytest-asyncio==0.21.0
-imageio==2.27.0
+imageio==2.28.0
 pyarrow==11.0.0
 tritonclient[all]==2.32.0
 build[virtualenv]==0.10.0
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")