Merge our fork with the opensource MLServer. All changes are to the huggingface runtime.

runtimes/huggingface/mlserver_huggingface/runtime.py

@@ -23,17 +23,16 @@ def __init__(self, settings: ModelSettings):
         super().__init__(settings)
 
     async def load(self) -> bool:
-        # Loading & caching pipeline in asyncio loop to avoid blocking
         logger.info(f"Loading model for task '{self.hf_settings.task_name}'...")
-        await asyncio.get_running_loop().run_in_executor(
-            None,
-            load_pipeline_from_settings,
-            self.hf_settings,
-            self.settings,
+        loop = asyncio.get_running_loop()
+        [self._model] = await asyncio.gather(
+            loop.run_in_executor(
+                None,
+                load_pipeline_from_settings,
+                self.hf_settings,
+                self.settings,
+            )
         )
-
-        # Now we load the cached model which should not block asyncio
-        self._model = load_pipeline_from_settings(self.hf_settings, self.settings)
         self._merge_metadata()
         return True
 

runtimes/huggingface/mlserver_huggingface/settings.py

@@ -83,10 +83,10 @@ class Config:
     runtime.
     """
 
-    device: int = -1
+    device: Optional[Union[int, str]] = None
     """
     Device in which this pipeline will be loaded (e.g., "cpu", "cuda:1", "mps",
-    or a GPU ordinal rank like 1).
+    or a GPU ordinal rank like 1). Default value of None becomes cpu.
     """
 
     inter_op_threads: Optional[int] = None

runtimes/huggingface/tests/test_common.py

@@ -2,7 +2,7 @@
 
 import pytest
 import torch
-from typing import Dict, Optional
+from typing import Dict, Optional, Union
 from optimum.onnxruntime.modeling_ort import ORTModelForQuestionAnswering
 from transformers.models.distilbert.modeling_distilbert import (
     DistilBertForQuestionAnswering,
@@ -13,6 +13,9 @@
 from mlserver_huggingface.runtime import HuggingFaceRuntime
 from mlserver_huggingface.settings import HuggingFaceSettings
 from mlserver_huggingface.common import load_pipeline_from_settings
+from mlserver.types import InferenceRequest, RequestInput
+from mlserver.types.dataplane import Parameters
+from mlserver_huggingface.codecs.base import MultiInputRequestCodec
 
 
 @pytest.mark.parametrize(
@@ -169,6 +172,43 @@ def test_pipeline_uses_model_kwargs(
     assert m.model.dtype == expected
 
 
+@pytest.mark.parametrize(
+    "pretrained_model, device, expected",
+    [
+        (
+            "hf-internal-testing/tiny-bert-for-token-classification",
+            None,
+            torch.device("cpu"),
+        ),
+        (
+            "hf-internal-testing/tiny-bert-for-token-classification",
+            -1,
+            torch.device("cpu"),
+        ),
+        (
+            "hf-internal-testing/tiny-bert-for-token-classification",
+            "cpu",
+            torch.device("cpu"),
+        ),
+    ],
+)
+def test_pipeline_cpu_device_set(
+    pretrained_model: str,
+    device: Optional[Union[str, int]],
+    expected: torch.device,
+):
+    hf_settings = HuggingFaceSettings(
+        pretrained_model=pretrained_model, task="token-classification", device=device
+    )
+    model_settings = ModelSettings(
+        name="foo",
+        implementation=HuggingFaceRuntime,
+    )
+    m = load_pipeline_from_settings(hf_settings, model_settings)
+
+    assert m.model.device == expected
+
+
 @pytest.mark.parametrize(
     "pretrained_model, task, input_batch_size, expected_batch_size",
     [

runtimes/huggingface/tests/test_settings.py

@@ -98,7 +98,7 @@ def test_merge_huggingface_settings_extra_raises(model_settings_extra_none):
                 pretrained_tokenizer=None,
                 framework=None,
                 optimum_model=False,
-                device=-1,
+                device=None,
                 inter_op_threads=None,
                 intra_op_threads=None,
             ),
@@ -113,7 +113,7 @@ def test_merge_huggingface_settings_extra_raises(model_settings_extra_none):
                 pretrained_tokenizer=None,
                 framework=None,
                 optimum_model=False,
-                device=-1,
+                device=None,
                 inter_op_threads=None,
                 intra_op_threads=None,
             ),
@@ -128,7 +128,7 @@ def test_merge_huggingface_settings_extra_raises(model_settings_extra_none):
                 pretrained_tokenizer=None,
                 framework=None,
                 optimum_model=False,
-                device=-1,
+                device=None,
                 inter_op_threads=None,
                 intra_op_threads=None,
             ),
@@ -143,7 +143,7 @@ def test_merge_huggingface_settings_extra_raises(model_settings_extra_none):
                 pretrained_tokenizer=None,
                 framework=None,
                 optimum_model=False,
-                device=-1,
+                device=None,
                 inter_op_threads=None,
                 intra_op_threads=None,
             ),