refactor: asign runner name (#2139)

* refactor: runner name * fix * fix for transformers
bentoml · Dec 24, 2021 · 5ab95a6 · 5ab95a6
1 parent 30edb7b
commit 5ab95a6
Show file tree

Hide file tree

Showing 25 changed files with 262 additions and 125 deletions.
diff --git a/bentoml/_internal/frameworks/catboost.py b/bentoml/_internal/frameworks/catboost.py
@@ -43,7 +43,7 @@
 
 
 def _get_model_info(
-    tag: t.Union[str, Tag],
+    tag: Tag,
     model_params: t.Optional[t.Dict[str, t.Union[str, int]]],
     model_store: "ModelStore",
 ) -> t.Tuple["Model", str, t.Dict[str, t.Any]]:
@@ -235,17 +235,18 @@ class _CatBoostRunner(Runner):
     @inject
     def __init__(
         self,
-        tag: t.Union[str, Tag],
+        tag: Tag,
         predict_fn_name: str,
         model_params: t.Optional[t.Dict[str, t.Union[str, int]]],
+        name: str,
         resource_quota: t.Optional[t.Dict[str, t.Any]],
         batch_options: t.Optional[t.Dict[str, t.Any]],
         model_store: "ModelStore" = Provide[BentoMLContainer.model_store],
     ):
         model_info, model_file, _model_params = _get_model_info(
             tag, model_params, model_store
         )
-        super().__init__(model_info.tag.name, resource_quota, batch_options)
+        super().__init__(name, resource_quota, batch_options)
         self._model_info = model_info
         self._model_file = model_file
         self._predict_fn_name = predict_fn_name
@@ -283,9 +284,10 @@ def load_runner(
     predict_fn_name: str = "predict",
     *,
     model_params: t.Union[None, t.Dict[str, t.Union[str, int]]] = None,
+    model_store: "ModelStore" = Provide[BentoMLContainer.model_store],
+    name: t.Optional[str] = None,
     resource_quota: t.Union[None, t.Dict[str, t.Any]] = None,
     batch_options: t.Union[None, t.Dict[str, t.Any]] = None,
-    model_store: "ModelStore" = Provide[BentoMLContainer.model_store],
 ) -> "_CatBoostRunner":
     """
     Runner represents a unit of serving logic that can be scaled horizontally to
@@ -320,11 +322,15 @@ def load_runner(
         runner = bentoml.catboost.load_runner("my_model:latest"")
         runner.run(cbt.Pool(input_data))
     """  # noqa
+    tag = Tag.from_taglike(tag)
+    if name is None:
+        name = tag.name
     return _CatBoostRunner(
         tag=tag,
         predict_fn_name=predict_fn_name,
         model_params=model_params,
+        model_store=model_store,
+        name=name,
         resource_quota=resource_quota,
         batch_options=batch_options,
-        model_store=model_store,
     )
diff --git a/bentoml/_internal/frameworks/detectron.py b/bentoml/_internal/frameworks/detectron.py
@@ -48,7 +48,7 @@
 
 @inject
 def load(
-    tag: t.Union[str, Tag],
+    tag: Tag,
     device: str = "cpu",
     model_store: "ModelStore" = Provide[BentoMLContainer.model_store],
 ) -> "torch.nn.Module":
@@ -201,19 +201,20 @@ class _DetectronRunner(Runner):
     # TODO add partial_kwargs @larme
     def __init__(
         self,
-        tag: t.Union[str, Tag],
+        tag: Tag,
         predict_fn_name: str,
+        name: str,
         resource_quota: t.Optional[t.Dict[str, t.Any]],
         batch_options: t.Optional[t.Dict[str, t.Any]],
         model_store: "ModelStore" = Provide[BentoMLContainer.model_store],
     ):
-        super().__init__(tag, resource_quota, batch_options)
+        super().__init__(name, resource_quota, batch_options)
         self._tag = tag
         self._predict_fn_name = predict_fn_name
         self._model_store = model_store
 
     @property
-    def required_models(self) -> t.List[str]:
+    def required_models(self) -> t.List[Tag]:
         return [self._tag]
 
     @property
@@ -263,6 +264,7 @@ def load_runner(
     tag: t.Union[str, Tag],
     predict_fn_name: str = "__call__",
     *,
+    name: t.Optional[str] = None,
     resource_quota: t.Union[None, t.Dict[str, t.Any]] = None,
     batch_options: t.Union[None, t.Dict[str, t.Any]] = None,
     model_store: "ModelStore" = Provide[BentoMLContainer.model_store],
@@ -290,10 +292,14 @@ def load_runner(
     Examples:
         TODO
     """  # noqa
+    tag = Tag.from_taglike(tag)
+    if name is None:
+        name = tag.name
     return _DetectronRunner(
         tag=tag,
         predict_fn_name=predict_fn_name,
+        model_store=model_store,
+        name=name,
         resource_quota=resource_quota,
         batch_options=batch_options,
-        model_store=model_store,
     )
diff --git a/bentoml/_internal/frameworks/easyocr.py b/bentoml/_internal/frameworks/easyocr.py
@@ -182,15 +182,16 @@ class _EasyOCRRunner(Runner):
     @inject
     def __init__(
         self,
-        tag: t.Union[str, Tag],
+        tag: Tag,
         predict_fn_name: str,
+        name: str,
         predict_params: t.Optional[t.Dict[str, t.Any]],
         resource_quota: t.Optional[t.Dict[str, t.Any]],
         batch_options: t.Optional[t.Dict[str, t.Any]],
         model_store: "ModelStore" = Provide[BentoMLContainer.model_store],
     ):
-        super().__init__(str(tag), resource_quota, batch_options)
-        self._tag = Tag.from_taglike(tag)
+        super().__init__(name, resource_quota, batch_options)
+        self._tag = tag
         self._predict_fn_name = predict_fn_name
         self._predict_params = predict_params
         self._model_store = model_store
@@ -227,6 +228,7 @@ def load_runner(
     tag: t.Union[str, Tag],
     predict_fn_name: str = "readtext_batched",
     *,
+    name: t.Optional[str] = None,
     predict_params: t.Union[None, t.Dict[str, t.Union[str, t.Any]]] = None,
     resource_quota: t.Union[None, t.Dict[str, t.Any]] = None,
     batch_options: t.Union[None, t.Dict[str, t.Any]] = None,
@@ -263,11 +265,15 @@ def load_runner(
         input_data = pd.from_csv("/path/to/csv")
         runner = bentoml.xgboost.load_runner("my_model:20201012_DE43A2")
         runner.run(xgb.DMatrix(input_data))
-    """  # noqa
+    """
+    tag = Tag.from_taglike(tag)
+    if name is None:
+        name = tag.name
     return _EasyOCRRunner(
         tag=tag,
         predict_fn_name=predict_fn_name,
         predict_params=predict_params,
+        name=name,
         resource_quota=resource_quota,
         batch_options=batch_options,
         model_store=model_store,

diff --git a/bentoml/_internal/frameworks/gluon.py b/bentoml/_internal/frameworks/gluon.py
@@ -38,7 +38,7 @@
 
 @inject
 def load(
-    tag: t.Union[str, Tag],
+    tag: Tag,
     mxnet_ctx: t.Optional[mxnet.context.Context] = None,
     model_store: "ModelStore" = Provide[BentoMLContainer.model_store],
 ) -> gluon.Block:
@@ -128,14 +128,15 @@ class _GluonRunner(Runner):
     @inject
     def __init__(
         self,
-        tag: t.Union[str, Tag],
+        tag: Tag,
         predict_fn_name: str,
+        name: str,
         resource_quota: t.Optional[t.Dict[str, t.Any]],
         batch_options: t.Optional[t.Dict[str, t.Any]],
         model_store: "ModelStore" = Provide[BentoMLContainer.model_store],
     ):
-        super().__init__(str(tag), resource_quota, batch_options)
-        self._tag = Tag.from_taglike(tag)
+        super().__init__(name, resource_quota, batch_options)
+        self._tag = tag
         self._predict_fn_name = predict_fn_name
         self._model_store = model_store
         self._ctx = None
@@ -190,6 +191,7 @@ def load_runner(
     predict_fn_name: str = "__call__",
     *,
     resource_quota: t.Union[None, t.Dict[str, t.Any]] = None,
+    name: t.Optional[str] = None,
     batch_options: t.Union[None, t.Dict[str, t.Any]] = None,
     model_store: "ModelStore" = Provide[BentoMLContainer.model_store],
 ) -> _GluonRunner:
@@ -216,6 +218,9 @@ def load_runner(
     Examples:
         TODO
     """  # noqa
+    tag = Tag.from_taglike(tag)
+    if name is None:
+        name = tag.name
     return _GluonRunner(
         tag=tag,
         predict_fn_name=predict_fn_name,

diff --git a/bentoml/_internal/frameworks/h2o.py b/bentoml/_internal/frameworks/h2o.py
@@ -39,7 +39,7 @@
 
 @inject
 def load(
-    tag: t.Union[str, Tag],
+    tag: Tag,
     init_params: t.Optional[t.Dict[str, t.Any]] = None,
     model_store: "ModelStore" = Provide[BentoMLContainer.model_store],
 ) -> h2o.model.model_base.ModelBase:
@@ -133,14 +133,15 @@ class _H2ORunner(Runner):
     @inject
     def __init__(
         self,
-        tag: t.Union[str, Tag],
+        tag: Tag,
         predict_fn_name: str,
         init_params: t.Optional[t.Dict[str, t.Union[str, t.Any]]],
+        name: str,
         resource_quota: t.Optional[t.Dict[str, t.Any]],
         batch_options: t.Optional[t.Dict[str, t.Any]],
         model_store: "ModelStore" = Provide[BentoMLContainer.model_store],
     ):
-        super().__init__(str(tag), resource_quota, batch_options)
+        super().__init__(name, resource_quota, batch_options)
 
         self._tag = Tag.from_taglike(tag)
         self._predict_fn_name = predict_fn_name
@@ -189,6 +190,7 @@ def load_runner(
     predict_fn_name: str = "predict",
     *,
     init_params: t.Optional[t.Dict[str, t.Union[str, t.Any]]],
+    name: t.Optional[str] = None,
     resource_quota: t.Optional[t.Dict[str, t.Any]] = None,
     batch_options: t.Optional[t.Dict[str, t.Any]] = None,
     model_store: "ModelStore" = Provide[BentoMLContainer.model_store],
@@ -220,10 +222,14 @@ def load_runner(
         TODO
 
     """  # noqa
+    tag = Tag.from_taglike(tag)
+    if name is None:
+        name = tag.name
     return _H2ORunner(
         tag=tag,
         predict_fn_name=predict_fn_name,
         init_params=init_params,
+        name=name,
         resource_quota=resource_quota,
         batch_options=batch_options,
         model_store=model_store,

diff --git a/bentoml/_internal/frameworks/keras.py b/bentoml/_internal/frameworks/keras.py
@@ -79,7 +79,7 @@ def get_session() -> "BaseSession":
 
 @inject
 def load(
-    tag: t.Union[str, Tag],
+    tag: Tag,
     model_store: "ModelStore" = Provide[BentoMLContainer.model_store],
 ) -> "keras.Model":
     """
@@ -198,10 +198,11 @@ class _KerasRunner(_TensorflowRunner):
     @inject
     def __init__(
         self,
-        tag: t.Union[str, Tag],
+        tag: Tag,
         predict_fn_name: str,
         device_id: str,
         predict_kwargs: t.Optional[t.Dict[str, t.Any]],
+        name: str,
         resource_quota: t.Optional[t.Dict[str, t.Any]],
         batch_options: t.Optional[t.Dict[str, t.Any]],
         model_store: "ModelStore" = Provide[BentoMLContainer.model_store],
@@ -210,6 +211,7 @@ def __init__(
             tag=tag,
             predict_fn_name=predict_fn_name,
             device_id=device_id,
+            name=name,
             partial_kwargs=predict_kwargs,
             resource_quota=resource_quota,
             batch_options=batch_options,
@@ -254,6 +256,7 @@ def load_runner(
     predict_fn_name: str = "predict",
     device_id: str = "CPU:0",
     predict_kwargs: t.Optional[t.Dict[str, t.Any]] = None,
+    name: t.Optional[str] = None,
     resource_quota: t.Union[None, t.Dict[str, t.Any]] = None,
     batch_options: t.Union[None, t.Dict[str, t.Any]] = None,
     model_store: "ModelStore" = Provide[BentoMLContainer.model_store],
@@ -284,10 +287,14 @@ def load_runner(
 
     Examples::
     """  # noqa: LN001
+    tag = Tag.from_taglike(tag)
+    if name is None:
+        name = tag.name
     return _KerasRunner(
         tag=tag,
         predict_fn_name=predict_fn_name,
         device_id=device_id,
+        name=name,
         predict_kwargs=predict_kwargs,
         resource_quota=resource_quota,
         batch_options=batch_options,

diff --git a/bentoml/_internal/frameworks/lightgbm.py b/bentoml/_internal/frameworks/lightgbm.py
@@ -73,7 +73,7 @@ def _get_model_info(
 
 @inject
 def load(
-    tag: t.Union[str, Tag],
+    tag: Tag,
     booster_params: t.Optional[t.Dict[str, t.Union[str, int]]] = None,
     model_store: "ModelStore" = Provide[BentoMLContainer.model_store],
 ) -> t.Union["lgb.basic.Booster", _LightGBMModelType]:
@@ -208,14 +208,15 @@ class _LightGBMRunner(Runner):
     @inject
     def __init__(
         self,
-        tag: t.Union[str, Tag],
+        tag: Tag,
         infer_api_callback: str,
         booster_params: t.Optional[t.Dict[str, t.Union[str, int]]],
+        name: str,
         resource_quota: t.Optional[t.Dict[str, t.Any]],
         batch_options: t.Optional[t.Dict[str, t.Any]],
         model_store: "ModelStore" = Provide[BentoMLContainer.model_store],
     ):
-        super().__init__(str(tag), resource_quota, batch_options)
+        super().__init__(name, resource_quota, batch_options)
         model_info, model_file, booster_params = _get_model_info(
             tag, booster_params, model_store
         )
@@ -269,6 +270,7 @@ def load_runner(
     infer_api_callback: str = "predict",
     *,
     booster_params: t.Optional[t.Dict[str, t.Union[str, int]]] = None,
+    name: t.Optional[str] = None,
     resource_quota: t.Union[None, t.Dict[str, t.Any]] = None,
     batch_options: t.Union[None, t.Dict[str, t.Any]] = None,
     model_store: "ModelStore" = Provide[BentoMLContainer.model_store],
@@ -303,9 +305,13 @@ def load_runner(
         runner = bentoml.lightgbm.load_runner("my_lightgbm_model:latest")
         runner.run_batch(X_test, num_iteration=gbm.best_iteration)
     """  # noqa
+    tag = Tag.from_taglike(tag)
+    if name is None:
+        name = tag.name
     return _LightGBMRunner(
         tag=tag,
         infer_api_callback=infer_api_callback,
+        name=name,
         booster_params=booster_params,
         resource_quota=resource_quota,
         batch_options=batch_options,