feat: Opt-in bento image support

src/_bentoml_impl/docker.py

@@ -0,0 +1,60 @@
+from __future__ import annotations
+
+import importlib
+import os
+import shlex
+import typing as t
+
+import attrs
+from fs.base import FS
+from jinja2 import Environment
+from jinja2 import FileSystemLoader
+
+from bentoml._internal.bento.bento import ImageInfo
+from bentoml._internal.container.generate import DEFAULT_BENTO_ENVS
+from bentoml._internal.container.generate import PREHEAT_PIP_PACKAGES
+from bentoml._internal.container.generate import expands_bento_path
+from bentoml._internal.container.generate import resolve_package_versions
+from bentoml._internal.container.generate import to_bento_field
+from bentoml._internal.container.generate import to_options_field
+
+
+def get_templates_variables(
+    image: ImageInfo, bento_fs: FS, **bento_envs: t.Any
+) -> dict[str, t.Any]:
+    from bentoml._internal.configuration.containers import BentoMLContainer
+
+    bento_envs = {**DEFAULT_BENTO_ENVS, **bento_envs}
+    options = attrs.asdict(image)
+    requirement_file = bento_fs.getsyspath("env/python/requirements.txt")
+    if os.path.exists(requirement_file):
+        python_packages = resolve_package_versions(requirement_file)
+    else:
+        python_packages = {}
+    pip_preheat_packages = [
+        python_packages[k] for k in PREHEAT_PIP_PACKAGES if k in python_packages
+    ]
+    return {
+        **{to_options_field(k): v for k, v in options.items()},
+        **{to_bento_field(k): v for k, v in bento_envs.items()},
+        "__prometheus_port__": BentoMLContainer.grpc.metrics.port.get(),
+        "__pip_preheat_packages__": pip_preheat_packages,
+    }
+
+
+def generate_dockerfile(
+    image: ImageInfo, bento_fs: FS, *, frontend: str = "dockerfile", **bento_envs: t.Any
+) -> str:
+    templates_path = importlib.import_module(
+        f"bentoml._internal.container.frontend.{frontend}.templates"
+    ).__path__
+    environment = Environment(
+        extensions=["jinja2.ext.do", "jinja2.ext.loopcontrols", "jinja2.ext.debug"],
+        trim_blocks=True,
+        lstrip_blocks=True,
+        loader=FileSystemLoader(templates_path, followlinks=True),
+    )
+    environment.filters["bash_quote"] = shlex.quote
+    environment.globals["expands_bento_path"] = expands_bento_path
+    template = environment.get_template("base_v2.j2")
+    return template.render(**get_templates_variables(image, bento_fs, **bento_envs))

src/_bentoml_sdk/images.py

@@ -0,0 +1,237 @@
+from __future__ import annotations
+
+import logging
+import platform
+import subprocess
+import sys
+import typing as t
+from pathlib import Path
+
+import attrs
+
+from bentoml._internal.bento.bento import ImageInfo
+from bentoml._internal.bento.build_config import BentoBuildConfig
+from bentoml._internal.configuration import get_bentoml_requirement
+from bentoml._internal.configuration import get_debug_mode
+from bentoml._internal.configuration import get_quiet_mode
+from bentoml._internal.container.frontend.dockerfile import CONTAINER_METADATA
+from bentoml._internal.container.frontend.dockerfile import CONTAINER_SUPPORTED_DISTROS
+from bentoml.exceptions import BentoMLConfigException
+from bentoml.exceptions import BentoMLException
+
+logger = logging.getLogger("bentoml.build")
+
+DEFAULT_PYTHON_VERSION = f"{sys.version_info.major}.{sys.version_info.minor}"
+BENTOML_GIT_DEPENDENCY = "bentoml @ git+https://github.com/bentoml/bentoml.git"
+
+
+@attrs.define
+class Image:
+    """A class defining the environment requirements for bento."""
+
+    base_image: str
+    python_version: str = DEFAULT_PYTHON_VERSION
+    commands: t.List[str] = attrs.field(factory=list)
+    lock_python_packages: bool = True
+    python_requirements: str = ""
+
+    def requirements_file(self, file_path: str) -> t.Self:
+        """Add a requirements file to the image. Supports chaining call.
+
+        Example:
+
+        .. code-block:: python
+
+            image = Image("debian:latest").requirements_file("requirements.txt")
+        """
+        self.python_requirements += Path(file_path).read_text()
+        return self
+
+    def python_packages(self, *packages: str) -> t.Self:
+        """Add python dependencies to the image. Supports chaining call.
+
+        Example:
+
+        .. code-block:: python
+
+            image = Image("debian:latest")\
+                .python_packages("numpy", "pandas")\
+                .requirements_file("requirements.txt")
+        """
+        self.python_requirements += "\n".join(packages)
+        return self
+
+    def run(self, command: str) -> t.Self:
+        """Add a command to the image. Supports chaining call.
+
+        Example:
+
+        .. code-block:: python
+
+            image = Image("debian:latest").run("echo 'Hello, World!'")
+        """
+        self.commands.append(command)
+        return self
+
+    def freeze(self, platform_: str | None = None) -> ImageInfo:
+        """Freeze the image to an ImageInfo object for build."""
+        if not self.lock_python_packages:
+            python_requirements = self.python_requirements
+        else:
+            python_requirements = self._freeze_python_requirements(platform_)
+        return ImageInfo(
+            base_image=self.base_image,
+            python_version=self.python_version,
+            commands=self.commands,
+            python_requirements=python_requirements,
+        )
+
+    def _freeze_python_requirements(self, platform_: str | None = None) -> str:
+        from tempfile import TemporaryDirectory
+
+        from pip_requirements_parser import RequirementsFile
+
+        with TemporaryDirectory(prefix="bento-reqs-") as parent:
+            requirements_in = Path(parent).joinpath("requirements.in")
+            requirements_in.write_text(self.python_requirements)
+            # XXX: RequirementsFile.from_string() does not work due to bugs
+            requirements_file = RequirementsFile.from_file(str(requirements_in))
+            has_bentoml_req = any(
+                req.name and req.name.lower() == "bentoml"
+                for req in requirements_file.requirements
+            )
+            with requirements_in.open("w") as f:
+                f.write(requirements_file.dumps(preserve_one_empty_line=True))
+                if not has_bentoml_req:
+                    req = get_bentoml_requirement() or BENTOML_GIT_DEPENDENCY
+                    f.write(f"{req}\n")
+            lock_args = [
+                str(requirements_in),
+                "--allow-unsafe",
+                "--no-header",
+                f"--output-file={requirements_in.with_suffix('.lock')}",
+                "--emit-index-url",
+                "--emit-find-links",
+                "--no-annotate",
+            ]
+            if get_debug_mode():
+                lock_args.append("--verbose")
+            else:
+                lock_args.append("--quiet")
+            logger.info("Locking PyPI package versions.")
+            if platform_:
+                lock_args.extend(["--python-platform", platform_])
+            elif platform.system() != "Linux" or platform.machine() != "x86_64":
+                logger.info(
+                    "Locking packages for x86_64-unknown-linux-gnu. "
+                    "Pass `--platform` option to specify the platform."
+                )
+                lock_args.extend(["--python-platform", "linux"])
+            cmd = [sys.executable, "-m", "uv", "pip", "compile", *lock_args]
+            try:
+                subprocess.check_call(
+                    cmd,
+                    text=True,
+                    stderr=subprocess.DEVNULL if get_quiet_mode() else None,
+                    cwd=parent,
+                )
+            except subprocess.CalledProcessError as e:
+                raise BentoMLException(f"Failed to lock PyPI packages: {e}") from None
+            return requirements_in.with_suffix(".lock").read_text()
+
+
+@attrs.define
+class PythonImage(Image):
+    base_image: str = ""
+    distro: str = "debian"
+    _original_base_image: str = attrs.field(init=False, default="")
+
+    def __attrs_post_init__(self) -> None:
+        self._original_base_image = self.base_image
+        if not self.base_image:
+            if self.distro not in CONTAINER_METADATA:
+                raise BentoMLConfigException(
+                    f"Unsupported distro: {self.distro}, expected one of {CONTAINER_SUPPORTED_DISTROS}"
+                )
+
+            self.base_image = CONTAINER_METADATA[self.distro]["python"]["image"].format(
+                spec_version=self.python_version
+            )
+
+    def system_packages(self, *packages: str) -> t.Self:
+        if self._original_base_image:
+            raise BentoMLConfigException(
+                "system_packages() can only be used in default base image"
+            )
+        if self.distro not in CONTAINER_METADATA:
+            raise BentoMLConfigException(
+                f"Unsupported distro: {self.distro}, expected one of {CONTAINER_SUPPORTED_DISTROS}"
+            )
+        self.commands.append(
+            CONTAINER_METADATA[self.distro]["install_command"].format(
+                packages=" ".join(packages)
+            )
+        )
+        return self
+
+
+def get_image_from_build_config(build_config: BentoBuildConfig) -> Image | None:
+    if not build_config.conda.is_empty():
+        logger.warning(
+            "conda options are not supported by bento v2, fallback to bento v1"
+        )
+        return None
+    image = PythonImage()
+    docker_options = build_config.docker
+    if docker_options.cuda_version is not None:
+        logger.warning(
+            "docker.cuda_version is not supported by bento v2, fallback to bento v1"
+        )
+        return None
+    if docker_options.dockerfile_template is not None:
+        logger.warning(
+            "docker.dockerfile_template is not supported by bento v2, fallback to bento v1"
+        )
+        return None
+    if docker_options.setup_script is not None:
+        logger.warning(
+            "docker.setup_script is not supported by bento v2, fallback to bento v1"
+        )
+        return None
+    if docker_options.base_image is not None:
+        image.base_image = docker_options.base_image
+    if docker_options.distro is not None:
+        image.distro = docker_options.distro
+    if docker_options.python_version is not None:
+        image.python_version = docker_options.python_version
+    if docker_options.system_packages:
+        image.system_packages(*docker_options.system_packages)
+
+    python_options = build_config.python.with_defaults()
+    if python_options.wheels:
+        logger.warning(
+            "python.wheels is not supported by bento v2, fallback to bento v1"
+        )
+        return None
+    image.lock_python_packages = python_options.lock_packages
+    if python_options.index_url:
+        image.python_packages(f"--index-url {python_options.index_url}")
+    if python_options.no_index:
+        image.python_packages("--no-index")
+    if python_options.trusted_host:
+        image.python_packages(
+            *(f"--trusted-host {h}" for h in python_options.trusted_host)
+        )
+    if python_options.extra_index_url:
+        image.python_packages(
+            *(f"--extra-index-url {url}" for url in python_options.extra_index_url)
+        )
+    if python_options.find_links:
+        image.python_packages(
+            *(f"--find-links {link}" for link in python_options.find_links)
+        )
+    if python_options.requirements_txt:
+        image.requirements_file(python_options.requirements_txt)
+    if python_options.packages:
+        image.python_packages(*python_options.packages)
+    return image

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
@@ -27,6 +28,7 @@
 from bentoml.exceptions import BentoMLConfigException
 from bentoml.exceptions import BentoMLException
 
+from ..images import Image
 from ..method import APIMethod
 from ..models import BentoModel
 from ..models import HuggingFaceModel
@@ -61,13 +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)
@@ -406,10 +413,24 @@ def service(inner: type[T], /) -> Service[T]: ...
 
 
 @t.overload
-def service(inner: None = ..., /, **kwargs: Unpack[Config]) -> _ServiceDecorator: ...
-
-
-def service(inner: type[T] | None = None, /, **kwargs: Unpack[Config]) -> t.Any:
+def service(
+    inner: None = ...,
+    /,
+    *,
+    image: Image | None = None,
+    envs: list[dict[str, t.Any]] | None = None,
+    **kwargs: Unpack[Config],
+) -> _ServiceDecorator: ...
+
+
+def service(
+    inner: type[T] | None = None,
+    /,
+    *,
+    image: Image | None = None,
+    envs: list[dict[str, t.Any]] | None = None,
+    **kwargs: Unpack[Config],
+) -> t.Any:
     """Mark a class as a BentoML service.
 
     Example:
@@ -425,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)
+        return Service(config=config, inner=inner, image=image, envs=envs or [])
 
     return decorator(inner) if inner is not None else decorator