from __future__ import annotations
import contextlib
import importlib
import logging
import os
import shutil
import sys
import tempfile
import types
import typing as t
from pathlib import Path
from typing import TYPE_CHECKING
from simple_di import Provide
from simple_di import inject
from ...exceptions import BentoMLException
from ...exceptions import InvalidArgument
from ..configuration.containers import BentoMLContainer
from ..utils.cattr import bentoml_cattr
from .base import OCIBuilder
from .generate import generate_containerfile
if TYPE_CHECKING:
from ..bento import Bento
from ..bento import BentoStore
from ..bento.build_config import BentoEnvSchema
from ..tag import Tag
from .base import Arguments
P = t.ParamSpec("P")
class DefaultBackendImpl(types.ModuleType):
BUILD_CMD: list[str] | None
ENV: dict[str, str] | None
BUILDKIT_SUPPORT: bool
def find_binary(self) -> str | None: ...
def construct_build_args(
self,
**kwargs: t.Any, # pylint: disable=unused-argument
) -> Arguments: ...
def health(self) -> bool: ...
DefaultBuilder: t.TypeAlias = t.Literal[
"docker", "podman", "buildah", "buildx", "nerdctl", "buildctl"
]
logger = logging.getLogger(__name__)
BUILDER_REGISTRY: dict[str, OCIBuilder] = {}
DEFAULT_BACKENDS = frozenset(
{"docker", "buildx", "buildah", "podman", "nerdctl", "buildctl"}
)
def register_default_backends():
for backend in DEFAULT_BACKENDS:
try:
module = t.cast(
"DefaultBackendImpl", importlib.import_module(f".{backend}", __name__)
)
register_backend(
backend,
health=module.health,
binary=module.find_binary(),
construct_build_args=module.construct_build_args,
buildkit_support=module.BUILDKIT_SUPPORT,
env=getattr(module, "ENV", None),
build_cmd=getattr(module, "BUILD_CMD", None),
)
except NotImplementedError:
logger.debug("%s is not yet implemented.", backend)
except Exception as err: # pylint: disable=broad-except
raise ValueError(f"Failed to register backend '{backend}: {err}'") from None
@inject
def determine_container_tag(
bento_tag: Tag | str,
image_tag: tuple[str] | None = None,
_bento_store: BentoStore = Provide[BentoMLContainer.bento_store],
):
# NOTE: for tags strategy, we will always generate a default tag from the bento:tag
# If '-t/--image-tag' is provided, we will use this tag provided by user.
bento = _bento_store.get(bento_tag)
tag = (str(bento.tag),)
if image_tag is not None:
assert isinstance(image_tag, tuple)
tag = image_tag
return tag
def enable_buildkit(
*, backend: str | None = None, builder: OCIBuilder | None = None
) -> bool:
# We will look for DOCKER_BUILDKIT in the environment variable.
# as this will be our entry point to enable BuildKit.
if "DOCKER_BUILDKIT" in os.environ:
return bool(int(os.environ["DOCKER_BUILDKIT"]))
# If the variable is not set, fallback to the default value
# provided by the builder.
if builder is not None:
return builder.enable_buildkit
elif backend is not None:
return get_backend(backend).enable_buildkit
else:
raise ValueError("Either backend or builder must be provided.")
def split_envs_by_stage(
envs: t.Sequence[BentoEnvSchema],
) -> tuple[list[BentoEnvSchema], list[BentoEnvSchema]]:
runtime_envs: list[BentoEnvSchema] = []
build_envs: list[BentoEnvSchema] = []
for env in envs:
if env.stage == "build":
build_envs.append(env)
else:
runtime_envs.append(env)
return runtime_envs, build_envs
# XXX: Sync with BentoML extra dependencies found in pyproject.toml
FEATURES = frozenset(
{
"tracing",
"grpc",
"grpc-reflection",
"grpc-channelz",
"monitor-otlp",
"triton",
"aws",
"all",
"io",
"io-file",
"io-image",
"io-pandas",
"io-json",
"tracing-zipkin",
"tracing-jaeger",
"tracing-otlp",
}
)
@contextlib.contextmanager
def construct_containerfile(
bento: Bento,
enable_buildkit: bool = True,
*,
features: t.Sequence[str] | None = None,
add_header: bool = False,
) -> t.Generator[tuple[str, str], None, None]:
from _bentoml_sdk.models import BentoModel
from _bentoml_sdk.models import HuggingFaceModel
from ..bento.bento import BaseBentoInfo
from ..bento.bento import BentoInfo
from ..bento.bento import BentoInfoV2
from ..bento.build_config import DockerOptions
dockerfile_path = "env/docker/Dockerfile"
instruction: list[str] = []
with tempfile.TemporaryDirectory() as tempdir:
with open(bento.path_of("bento.yaml"), "rb") as bento_yaml:
options = BaseBentoInfo.from_yaml_file(bento_yaml)
# tmpdir is our new build context.
shutil.copytree(bento.path, tempdir, dirs_exist_ok=True)
# copy models from model store
bento_model_dir = Path(tempdir, "models")
hf_model_dir = Path(tempdir, "hf-models")
for model_dir in (bento_model_dir, hf_model_dir):
model_dir.mkdir(parents=True, exist_ok=True)
for model in options.all_models:
if model.registry == "huggingface":
model_ref = HuggingFaceModel.from_info(model)
model_ref.resolve(hf_model_dir)
else:
model_ref = BentoModel.from_info(model)
model_ref.resolve(bento_model_dir)
# NOTE: dockerfile_template is already included in the
# Dockerfile inside bento, and it is not relevant to
# construct_containerfile. Hence it is safe to set it to None here.
# See https://github.com/bentoml/BentoML/issues/3399.
if isinstance(options, BentoInfo):
docker_attrs = bentoml_cattr.unstructure(options.docker)
if (
"dockerfile_template" in docker_attrs
and docker_attrs["dockerfile_template"] is not None
):
# NOTE: if users specify a dockerfile_template, we will
# save it to /env/docker/Dockerfile.template. This is necessary
# for the reconstruction of the Dockerfile.
docker_attrs["dockerfile_template"] = "env/docker/Dockerfile.template"
dockerfile = generate_containerfile(
docker=DockerOptions(**docker_attrs),
build_ctx=tempdir,
conda=options.conda,
bento_fs=Path(tempdir),
enable_buildkit=enable_buildkit,
add_header=add_header,
)
instruction.append(dockerfile)
if features is not None:
diff = set(features).difference(FEATURES)
if len(diff) > 0:
raise InvalidArgument(
f"Available features are: {FEATURES}. Invalid fields from provided: {diff}"
)
PIP_CACHE_MOUNT = (
"--mount=type=cache,target=/root/.cache/pip "
if enable_buildkit
else ""
)
instruction.append(
"RUN %spip install bentoml[%s]"
% (PIP_CACHE_MOUNT, ",".join(features))
)
else:
from _bentoml_impl.docker import generate_dockerfile
assert isinstance(options, BentoInfoV2)
runtime_envs, build_env = split_envs_by_stage(options.envs)
if build_env and not enable_buildkit:
raise BentoMLException(
"stage='build' environment variables require BuildKit. "
"Enable BuildKit for your backend (e.g. set DOCKER_BUILDKIT=1)."
)
dockerfile = generate_dockerfile(
options.image,
Path(tempdir),
enable_buildkit=enable_buildkit,
add_header=add_header,
envs=runtime_envs,
secret_envs=build_env,
)
instruction.append(dockerfile)
Path(tempdir, dockerfile_path).write_text("\n".join(instruction))
yield tempdir, os.path.join(tempdir, dockerfile_path)
@inject
def build(
bento_tag: Tag | str,
backend: str,
features: t.Sequence[str] | None = None,
_bento_store: BentoStore = Provide[BentoMLContainer.bento_store],
**kwargs: t.Any,
):
clean_context = contextlib.ExitStack()
bento = _bento_store.get(bento_tag)
builder = get_backend(backend)
buildkit_enabled = enable_buildkit(builder=builder)
_, build_stage_envs = split_envs_by_stage(bento.info.envs)
secret_specs: list[str] = []
if build_stage_envs:
if not buildkit_enabled:
raise BentoMLException(
"stage='build' environment variables require BuildKit. "
"Enable BuildKit for your backend (e.g. set DOCKER_BUILDKIT=1)."
)
for env in build_stage_envs:
secret_value = os.getenv(env.name, env.value)
if secret_value is None:
raise BentoMLException(
f"Environment variable '{env.name}' (stage='build') is required during image build."
)
secret_file = tempfile.NamedTemporaryFile("w", delete=False)
secret_file.write(secret_value)
secret_file.flush()
secret_file.close()
clean_context.callback(lambda path=secret_file.name: os.remove(path))
secret_specs.append(f"id={env.name},src={secret_file.name}")
context_path, dockerfile = clean_context.enter_context(
construct_containerfile(
bento,
features=features,
enable_buildkit=buildkit_enabled,
)
)
try:
if secret_specs:
existing_secret = kwargs.get("secret")
if existing_secret is None:
merged_secret: tuple[str, ...] = tuple(secret_specs)
elif isinstance(existing_secret, (tuple, list)):
merged_secret = (*existing_secret, *secret_specs)
else:
merged_secret = (existing_secret, *secret_specs)
kwargs["secret"] = merged_secret
kwargs.update({"file": dockerfile, "context_path": context_path})
return builder.build(**kwargs)
except Exception as e: # pylint: disable=broad-except
logger.error(
"\nEncountered exception while trying to building image: %s",
e,
exc_info=sys.exc_info(),
)
finally:
clean_context.close()
[docs]
def register_backend(
backend: str,
*,
buildkit_support: bool,
health: t.Callable[[], bool],
construct_build_args: t.Callable[P, Arguments],
binary: str | None = None,
build_cmd: t.Sequence[str] | None = None,
env: dict[str, str] | None = None,
):
"""
Register a custom backend, provided with a build and health check implementation.
Args:
backend: Name of the backend.
buildkit_support: Whether the backend has support for BuildKit.
health: Health check implementation. This is a callable that takes no
argument and returns a boolean.
construct_build_args: This is a callable that takes possible backend keyword arguments
and returns a list of command line arguments.
env: Default environment variables dict for this OCI builder implementation.
binary: Optional binary path. If not provided, the binary will use the backend name.
Make sure that the binary is on your ``PATH``.
build_cmd: Optional build command. If not provided, the command will be 'build'.
Examples:
.. code-block:: python
from bentoml.container import register_backend
register_backend(
"lima",
binary=shutil.which("limactl"),
buildkit_support=True,
health=buildx_health,
construct_build_args=buildx_build_args,
env={"DOCKER_BUILDKIT": "1"},
)
"""
if backend in BUILDER_REGISTRY:
raise ValueError(f"Backend {backend} already registered.")
if binary is None:
binary = backend
try:
BUILDER_REGISTRY[backend] = OCIBuilder.create(
binary,
env=env,
build_cmd=build_cmd,
enable_buildkit=buildkit_support,
construct_build_args=construct_build_args,
health=health,
)
except Exception as e: # pylint: disable=broad-except
logger.error(
"Failed to register backend %s: %s", backend, e, exc_info=sys.exc_info()
)
raise e from None
@t.overload
def health(backend: DefaultBuilder) -> bool: ...
@t.overload
def health(backend: str) -> bool: ...
[docs]
def health(backend: str) -> bool:
"""
Check if the backend is healthy.
Args:
backend: The name of the backend.
.. note::
If given backend is a type of OCIBuilder, and the backend is not registered,
make sure to register it with ``bentoml.container.register_backend``.
Returns:
True if the backend is healthy, False otherwise.
"""
return get_backend(backend).health()
@t.overload
def get_backend(backend: DefaultBuilder) -> OCIBuilder: ...
@t.overload
def get_backend(backend: str) -> OCIBuilder: ...
[docs]
def get_backend(backend: str) -> OCIBuilder:
"""
Get a given backend.
Raises:
``ValueError``: If given backend is not available in backend registry.
"""
if isinstance(backend, OCIBuilder):
if backend not in BUILDER_REGISTRY.values():
logger.warning(
"Backend '%s' not registered. To use with 'bentoml.container.build', make sure to regsiter it with 'bentoml.container.register_backend'",
backend,
)
return backend
if backend not in BUILDER_REGISTRY:
raise ValueError(
f"Backend {backend} not registered. Available backends: {REGISTERED_BACKENDS}."
)
return BUILDER_REGISTRY[backend]
register_default_backends()
REGISTERED_BACKENDS = list(BUILDER_REGISTRY.keys())
__all__ = [
"build",
"health",
"register_backend",
"get_backend",
"REGISTERED_BACKENDS",
"split_envs_by_stage",
]