feat: support running scripts in new image spec (#5196)

* feat: support running scripts in new image spec

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

src/_bentoml_sdk/images.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import hashlib
 import logging
 import platform
 import subprocess
@@ -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:
@@ -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:
@@ -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:
@@ -207,11 +222,6 @@ def get_image_from_build_config(
             "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
@@ -252,4 +262,6 @@ def get_image_from_build_config(
         )
     if python_options.packages:
         image.python_packages(*python_options.packages)
+    if docker_options.setup_script:
+        image.run_script(docker_options.setup_script)
     return image

src/bentoml/_internal/bento/bento.py

@@ -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
@@ -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"