From 2950cffd5b047f2a4013a8f1ea0da8a0957c7fed Mon Sep 17 00:00:00 2001 From: aarnphm-ec2-dev <29749331+aarnphm@users.noreply.github.com> Date: Tue, 11 Jul 2023 01:30:39 +0000 Subject: [PATCH] fix(save): set bento store and model store Signed-off-by: aarnphm-ec2-dev <29749331+aarnphm@users.noreply.github.com> --- src/openllm/_package.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/openllm/_package.py b/src/openllm/_package.py index eef14e0f..701f889d 100644 --- a/src/openllm/_package.py +++ b/src/openllm/_package.py @@ -53,6 +53,7 @@ if t.TYPE_CHECKING: from fs.base import FS import openllm + from bentoml._internal.bento import BentoStore logger = logging.getLogger(__name__) @@ -244,6 +245,7 @@ def create_bento( extra_dependencies: tuple[str, ...] | None = None, build_ctx: str | None = None, runtime: t.Literal["ggml", "transformers"] = "transformers", + _bento_store: BentoStore = Provide[BentoMLContainer.bento_store], _model_store: ModelStore = Provide[BentoMLContainer.model_store], ) -> bentoml.Bento: framework_envvar = llm.config["env"]["framework_value"] @@ -347,4 +349,4 @@ def create_bento( bento._fs.writetext(service_fs_path, script) - return bento.save() + return bento.save(bento_store=_bento_store, model_store=_model_store)