Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: support running scripts in new image spec #5196

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 17 additions & 5 deletions src/_bentoml_sdk/images.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
from __future__ import annotations

import hashlib
import logging
import platform
import subprocess
Expand Down Expand Up @@ -36,6 +37,7 @@ class Image:
lock_python_packages: bool = True
python_requirements: str = ""
post_commands: t.List[str] = attrs.field(factory=list)
scripts: t.Dict[str, str] = attrs.field(factory=dict, init=False)
_after_pip_install: bool = attrs.field(init=False, default=False, repr=False)

def requirements_file(self, file_path: str) -> t.Self:
Expand Down Expand Up @@ -79,6 +81,18 @@ def run(self, command: str) -> t.Self:
commands.append(command)
return self

def run_script(self, script: str) -> t.Self:
"""Run a script in the image. Supports chaining call."""
commands = self.post_commands if self._after_pip_install else self.commands
script = Path(script).resolve().as_posix()
# Files under /env/docker will be copied into the env image layer
target_script = (
f"./env/docker/script__{hashlib.md5(script.encode()).hexdigest()}"
)
commands.append(f"chmod +x {target_script} && {target_script}")
self.scripts[script] = target_script
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:
Expand All @@ -91,6 +105,7 @@ def freeze(self, platform_: str | None = None) -> ImageInfo:
commands=self.commands,
python_requirements=python_requirements,
post_commands=self.post_commands,
scripts=self.scripts,
)

def _freeze_python_requirements(self, platform_: str | None = None) -> str:
Expand Down Expand Up @@ -203,11 +218,6 @@ def get_image_from_build_config(build_config: BentoBuildConfig) -> Image | None:
"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
image_params = {}
if docker_options.base_image is not None:
image_params["base_image"] = docker_options.base_image
Expand Down Expand Up @@ -246,4 +256,6 @@ def get_image_from_build_config(build_config: BentoBuildConfig) -> Image | None:
image.requirements_file(python_options.requirements_txt)
if python_options.packages:
image.python_packages(*python_options.packages)
if docker_options.setup_script:
image.run_script(docker_options.setup_script)
return image
3 changes: 3 additions & 0 deletions src/bentoml/_internal/bento/bento.py
Original file line number Diff line number Diff line change
Expand Up @@ -839,6 +839,7 @@ class ImageInfo:
commands: t.List[str] = attr.field(factory=list)
python_requirements: str = ""
post_commands: t.List[str] = attr.field(factory=list)
scripts: t.Dict[str, str] = attr.field(factory=dict)

def write_to_bento(self, bento_fs: FS, envs: list[BentoEnvSchema]) -> None:
from importlib import resources
Expand All @@ -857,6 +858,8 @@ def write_to_bento(self, bento_fs: FS, envs: list[BentoEnvSchema]) -> None:
dockerfile_path,
generate_dockerfile(self, bento_fs, enable_buildkit=False, envs=envs),
)
for script_name, target_path in self.scripts.items():
copy_file_to_fs_folder(script_name, bento_fs, dst_filename=target_path)

with resources.path(
"bentoml._internal.container.frontend.dockerfile", "entrypoint.sh"
Expand Down
Loading