feat: inject build envs

Signed-off-by: Frost Ming <[email protected]>

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,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)
@@ -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
 

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(

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]

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))

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 %}

tests/unit/_internal/bento/test_bento.py

@@ -94,6 +94,11 @@ def test_bento_info(tmpdir: Path):
   module: model_b_module
   creation_time: '{model_creation_time}'
   alias: model_b_alias
+entry_service: ''
+services: []
+envs: []
+schema: {{}}
+spec: 1
 runners:
 - name: runner_a
   runnable_type: test_runnable_a
@@ -102,10 +107,6 @@ def test_bento_info(tmpdir: Path):
   - runner_a_model
   resource_config:
     cpu: 2
-entry_service: ''
-services: []
-envs: []
-schema: {{}}
 apis:
 - name: predict
   input_type: NumpyNdarray