From 6ce6d7acfe9f64402506246d221de52f41f0aee8 Mon Sep 17 00:00:00 2001 From: Frost Ming Date: Wed, 23 Oct 2024 15:32:56 +0800 Subject: [PATCH] feat: inject build envs Signed-off-by: Frost Ming --- src/_bentoml_sdk/service/factory.py | 17 ++++++++++++++--- src/bentoml/_internal/bento/bento.py | 15 ++++++++++----- src/bentoml/_internal/bento/build_config.py | 2 +- src/bentoml/_internal/container/__init__.py | 1 + .../frontend/dockerfile/templates/base_v2.j2 | 7 +++++++ tests/unit/_internal/bento/test_bento.py | 1 + 6 files changed, 34 insertions(+), 9 deletions(-) diff --git a/src/_bentoml_sdk/service/factory.py b/src/_bentoml_sdk/service/factory.py index 91af95936b4..5654bcf3c9f 100644 --- a/src/_bentoml_sdk/service/factory.py +++ b/src/_bentoml_sdk/service/factory.py @@ -19,6 +19,7 @@ from bentoml import Runner from bentoml._internal.bento.bento import Bento +from bentoml._internal.bento.build_config import BentoEnvSchema from bentoml._internal.configuration.containers import BentoMLContainer from bentoml._internal.context import ServiceContext from bentoml._internal.models import Model as StoredModel @@ -62,6 +63,10 @@ def wrapper(self: Service[t.Any], *args: P.args, **kwargs: P.kwargs) -> R: return wrapper +def convert_envs(envs: t.List[t.Dict[str, t.Any]]) -> t.List[BentoEnvSchema]: + return [BentoEnvSchema(**env) for env in envs] + + @attrs.define class Service(t.Generic[T]): """A Bentoml service that can be served by BentoML server.""" @@ -69,7 +74,7 @@ class Service(t.Generic[T]): config: Config inner: type[T] image: t.Optional[Image] = None - + envs: t.List[BentoEnvSchema] = attrs.field(factory=list, converter=convert_envs) bento: t.Optional[Bento] = attrs.field(init=False, default=None) models: list[Model[t.Any]] = attrs.field(factory=list) apis: dict[str, APIMethod[..., t.Any]] = attrs.field(factory=dict) @@ -409,7 +414,12 @@ def service(inner: type[T], /) -> Service[T]: ... @t.overload def service( - inner: None = ..., /, *, image: Image | None = None, **kwargs: Unpack[Config] + inner: None = ..., + /, + *, + image: Image | None = None, + envs: list[dict[str, t.Any]] | None = None, + **kwargs: Unpack[Config], ) -> _ServiceDecorator: ... @@ -418,6 +428,7 @@ def service( /, *, image: Image | None = None, + envs: list[dict[str, t.Any]] | None = None, **kwargs: Unpack[Config], ) -> t.Any: """Mark a class as a BentoML service. @@ -435,7 +446,7 @@ def predict(self, input: str) -> str: def decorator(inner: type[T]) -> Service[T]: if isinstance(inner, Service): raise TypeError("service() decorator can only be applied once") - return Service(config=config, inner=inner, image=image) + return Service(config=config, inner=inner, image=image, envs=envs or []) return decorator(inner) if inner is not None else decorator diff --git a/src/bentoml/_internal/bento/bento.py b/src/bentoml/_internal/bento/bento.py index a9c8738c698..23b9846a90a 100644 --- a/src/bentoml/_internal/bento/bento.py +++ b/src/bentoml/_internal/bento/bento.py @@ -237,6 +237,7 @@ def create( platform: t.Optional[str] = None, bare: bool = False, reload: bool = False, + enabled_features: list[str] = Provide[BentoMLContainer.enabled_features], ) -> Bento: from _bentoml_sdk.images import Image from _bentoml_sdk.images import get_image_from_build_config @@ -250,6 +251,7 @@ def create( if build_ctx is None else os.path.realpath(os.path.expanduser(build_ctx)) ) + enable_image = "bento_image" in enabled_features if not os.path.isdir(build_ctx): raise InvalidArgument( f"Bento build context {build_ctx} does not exist or is not a directory." @@ -280,8 +282,10 @@ def create( if build_config.name is not None else to_snake_case(svc.name) ) - image = svc.image - if image is None: + build_config.envs.extend(svc.envs) + if enable_image: + image = svc.image + if image is None and enable_image: image = get_image_from_build_config(build_config) build_config = build_config.with_defaults() tag = Tag(bento_name, version) @@ -423,7 +427,7 @@ def append_model(model: BentoModelInfo) -> None: schema=svc.schema() if not is_legacy else {}, image=image.freeze(platform), ) - bento_info.image.write_to_bento(bento_fs, build_ctx) + bento_info.image.write_to_bento(bento_fs, build_config.envs) res = Bento(tag, bento_fs, bento_info) if bare: @@ -811,7 +815,7 @@ class ImageInfo: commands: t.List[str] = attr.field(factory=list) python_requirements: str = "" - def write_to_bento(self, bento_fs: FS, build_ctx: str) -> None: + def write_to_bento(self, bento_fs: FS, envs: list[BentoEnvSchema]) -> None: from importlib import resources from _bentoml_impl.docker import generate_dockerfile @@ -825,7 +829,8 @@ def write_to_bento(self, bento_fs: FS, build_ctx: str) -> None: bento_fs.makedirs(docker_folder, recreate=True) dockerfile_path = fs.path.join(docker_folder, "Dockerfile") bento_fs.writetext( - dockerfile_path, generate_dockerfile(self, bento_fs, enable_buildkit=False) + dockerfile_path, + generate_dockerfile(self, bento_fs, enable_buildkit=False, envs=envs), ) with resources.path( diff --git a/src/bentoml/_internal/bento/build_config.py b/src/bentoml/_internal/bento/build_config.py index 52ab382211e..ce55ea9f8ca 100644 --- a/src/bentoml/_internal/bento/build_config.py +++ b/src/bentoml/_internal/bento/build_config.py @@ -1013,4 +1013,4 @@ class FilledBentoBuildConfig(BentoBuildConfig): python: PythonOptions conda: CondaOptions models: t.List[ModelSpec] - envs: t.List[t.Dict[str, str]] + envs: t.List[BentoEnvSchema] diff --git a/src/bentoml/_internal/container/__init__.py b/src/bentoml/_internal/container/__init__.py index 0398d839b44..1746216d21c 100644 --- a/src/bentoml/_internal/container/__init__.py +++ b/src/bentoml/_internal/container/__init__.py @@ -219,6 +219,7 @@ def construct_containerfile( temp_fs, enable_buildkit=enable_buildkit, add_header=add_header, + envs=options.envs, ) instruction.append(dockerfile) temp_fs.writetext(dockerfile_path, "\n".join(instruction)) diff --git a/src/bentoml/_internal/container/frontend/dockerfile/templates/base_v2.j2 b/src/bentoml/_internal/container/frontend/dockerfile/templates/base_v2.j2 index eb89f606001..ec1936a5f4d 100644 --- a/src/bentoml/_internal/container/frontend/dockerfile/templates/base_v2.j2 +++ b/src/bentoml/_internal/container/frontend/dockerfile/templates/base_v2.j2 @@ -3,6 +3,7 @@ {% import '_macros.j2' as common %} {% set bento__entrypoint = expands_bento_path('env', 'docker', 'entrypoint.sh', bento_path=bento__path) %} {% set __enable_buildkit__ = bento__enable_buildkit | default(False) -%} +{% set __bento_envs__ = bento__envs | default([]) %} {% if __enable_buildkit__ %} # 1.2.1 is the current docker frontend that both buildkitd and kaniko supports. # syntax = {{ bento__buildkit_frontend }} @@ -43,6 +44,12 @@ ENV BENTO_PATH=$BENTO_PATH ENV BENTOML_HOME={{ bento__home }} ENV BENTOML_HF_CACHE_DIR={{ bento__path }}/hf-models +{% for env in __bento_envs__ %} +ARG {{ env.name }}{% if env.value %}={{ env.value }}{% endif %} + +ENV {{ env.name }}=${{ env.name }} +{% endfor %} + RUN mkdir $BENTO_PATH && chown {{ bento__user }}:{{ bento__user }} $BENTO_PATH -R WORKDIR $BENTO_PATH {% endblock %} diff --git a/tests/unit/_internal/bento/test_bento.py b/tests/unit/_internal/bento/test_bento.py index 09063b0d079..345ec328546 100644 --- a/tests/unit/_internal/bento/test_bento.py +++ b/tests/unit/_internal/bento/test_bento.py @@ -106,6 +106,7 @@ def test_bento_info(tmpdir: Path): services: [] envs: [] schema: {{}} +spec: 1 apis: - name: predict input_type: NumpyNdarray