From cdbc29c0ac51cc203561da7d71dbfa87f13369ce Mon Sep 17 00:00:00 2001
From: Frost Ming <me@frostming.com>
Date: Tue, 4 Feb 2025 13:21:24 +0800
Subject: [PATCH] feat: support running scripts in new image spec (#5196)

* feat: support running scripts in new image spec

Signed-off-by: Frost Ming <me@frostming.com>
---
 src/_bentoml_sdk/images.py           | 22 +++++++++++++++++-----
 src/bentoml/_internal/bento/bento.py |  3 +++
 2 files changed, 20 insertions(+), 5 deletions(-)

diff --git a/src/_bentoml_sdk/images.py b/src/_bentoml_sdk/images.py
index cde981612de..38b5bb6d7c3 100644
--- a/src/_bentoml_sdk/images.py
+++ b/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
diff --git a/src/bentoml/_internal/bento/bento.py b/src/bentoml/_internal/bento/bento.py
index fd59e06500d..b775cfe2a04 100644
--- a/src/bentoml/_internal/bento/bento.py
+++ b/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"