Skip to content

Commit

Permalink
feat: inject build envs
Browse files Browse the repository at this point in the history
Signed-off-by: Frost Ming <[email protected]>
  • Loading branch information
frostming committed Oct 23, 2024
1 parent 77375b4 commit 6ce6d7a
Show file tree
Hide file tree
Showing 6 changed files with 34 additions and 9 deletions.
17 changes: 14 additions & 3 deletions src/_bentoml_sdk/service/factory.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -62,14 +63,18 @@ 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."""

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)
Expand Down Expand Up @@ -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: ...


Expand All @@ -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.
Expand All @@ -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

Expand Down
15 changes: 10 additions & 5 deletions src/bentoml/_internal/bento/bento.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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."
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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
Expand All @@ -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(
Expand Down
2 changes: 1 addition & 1 deletion src/bentoml/_internal/bento/build_config.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]
1 change: 1 addition & 0 deletions src/bentoml/_internal/container/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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))
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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 }}
Expand Down Expand Up @@ -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 %}
Expand Down
1 change: 1 addition & 0 deletions tests/unit/_internal/bento/test_bento.py
Original file line number Diff line number Diff line change
Expand Up @@ -106,6 +106,7 @@ def test_bento_info(tmpdir: Path):
services: []
envs: []
schema: {{}}
spec: 1
apis:
- name: predict
input_type: NumpyNdarray
Expand Down

0 comments on commit 6ce6d7a

Please sign in to comment.