feat: add @bentoml.on_startup decorator (#5194)

* feat: add @bentoml.on_startup decorator

Added @bentoml.on_startup decorator to mirror existing lifecycle hook functionality:
- Added on_startup decorator in _bentoml_sdk/decorators.py
- Exposed decorator in bentoml/__init__.py
- Implemented startup hook handling in new-style services

Co-Authored-By: [email protected] <[email protected]>

* docs: add on_startup decorator doc and expose in _bentoml_sdk.__init__

Co-Authored-By: [email protected] <[email protected]>

---------

Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Co-authored-by: [email protected] <[email protected]>

docs/source/build-with-bentoml/lifecycle-hooks.rst

@@ -64,6 +64,57 @@ After the Service starts, you can see the following output on the server side in
     This runs on Service startup, once for each worker, so it runs 4 times.
     This runs on Service startup, once for each worker, so it runs 4 times.
 
+Startup hooks
+^^^^^^^^^^^^^
+
+Startup hooks are executed during Service initialization, after deployment hooks but before any API endpoints become available. These hooks run once per worker, making them ideal for worker-specific initialization tasks such as establishing database connections or loading resources.
+
+Use the ``@bentoml.on_startup`` decorator to specify a method as a startup hook. For example:
+
+.. code-block:: python
+
+    import bentoml
+
+    @bentoml.service(workers=4)
+    class HookService:
+        @bentoml.on_deployment
+        def prepare():
+            print("Global preparation, runs once before workers start.")
+
+        @bentoml.on_startup
+        def init_resources(self):
+            # This runs once per worker
+            print("Initializing resources for worker.")
+            self.db_connection = setup_database()
+
+        @bentoml.on_startup
+        async def init_async_resources(self):
+            # For async initialization tasks
+            print("Async resource initialization for worker.")
+            self.cache = await setup_cache()
+
+        @bentoml.api
+        def predict(self, text) -> str:
+            # Use initialized resources in API endpoints
+            return self.db_connection.query(text)
+
+When you start this Service, you'll see the following output:
+
+.. code-block:: bash
+
+    $ bentoml serve service:HookService
+
+    Global preparation, runs once before workers start. # on_deployment hook
+    2024-03-13T03:12:33+0000 [INFO] [cli] Starting production HTTP BentoServer from "service:HookService" listening on http://localhost:3000
+    Initializing resources for worker. # First worker's startup hooks
+    Async resource initialization for worker.
+    Initializing resources for worker. # Second worker's startup hooks
+    Async resource initialization for worker.
+    Initializing resources for worker. # Third worker's startup hooks
+    Async resource initialization for worker.
+    Initializing resources for worker. # Fourth worker's startup hooks
+    Async resource initialization for worker.
+
 Shutdown hooks
 ^^^^^^^^^^^^^^
 

src/_bentoml_impl/server/app.py

@@ -278,6 +278,19 @@ async def create_instance(self, app: Starlette) -> None:
         self._service_instance = self.service()
         self.service.gradio_app_startup_hook(max_concurrency=self.max_concurrency)
         logger.info("Service %s initialized", self.service.name)
+
+        # Call on_startup hook with optional ctx or context parameter
+        for name, member in vars(self.service.inner).items():
+            if callable(member) and getattr(member, "__bentoml_startup_hook__", False):
+                logger.info("Running startup hook: %s", name)
+                result = getattr(
+                    self._service_instance, name
+                )()  # call the bound method
+                if inspect.isawaitable(result):
+                    await result
+                    logger.info("Completed async startup hook: %s", name)
+                else:
+                    logger.info("Completed startup hook: %s", name)
         if deployment_url := os.getenv("BENTOCLOUD_DEPLOYMENT_URL"):
             proxy = RemoteProxy(
                 deployment_url, service=self.service, media_type="application/json"

src/_bentoml_sdk/__init__.py

@@ -14,7 +14,7 @@
 
 # ruff: noqa
 
-from .decorators import api, on_shutdown, asgi_app, on_deployment, task
+from .decorators import api, on_shutdown, on_startup, asgi_app, on_deployment, task
 from .service import get_current_service
 from .service import depends
 from .service import Service, ServiceConfig
@@ -36,6 +36,7 @@ def mount_asgi_app(
     "api",
     "task",
     "on_shutdown",
+    "on_startup",
     "on_deployment",
     "asgi_app",
     "mount_asgi_app",

src/_bentoml_sdk/decorators.py

@@ -27,6 +27,12 @@ def on_shutdown(func: F) -> F:
     return func
 
 
+def on_startup(func: F) -> F:
+    """Mark a method as a startup hook for the service."""
+    setattr(func, "__bentoml_startup_hook__", True)
+    return func
+
+
 def on_deployment(func: t.Callable[P, R] | staticmethod[P, R]) -> staticmethod[P, R]:
     inner = func.__func__ if isinstance(func, staticmethod) else func
     setattr(inner, "__bentoml_deployment_hook__", True)

src/bentoml/__init__.py

@@ -59,6 +59,7 @@
     "task": "_bentoml_sdk:task",
     "depends": "_bentoml_sdk:depends",
     "on_shutdown": "_bentoml_sdk:on_shutdown",
+    "on_startup": "_bentoml_sdk:on_startup",
     "on_deployment": "_bentoml_sdk:on_deployment",
     "asgi_app": "_bentoml_sdk:asgi_app",
     "mount_asgi_app": "_bentoml_sdk:mount_asgi_app",
@@ -159,6 +160,7 @@
     from _bentoml_sdk import mount_asgi_app
     from _bentoml_sdk import on_deployment
     from _bentoml_sdk import on_shutdown
+    from _bentoml_sdk import on_startup
     from _bentoml_sdk import runner_service
     from _bentoml_sdk import service
     from _bentoml_sdk import task
@@ -373,6 +375,7 @@ def __getattr__(name: str) -> Any:
     "task",
     "images",
     "on_shutdown",
+    "on_startup",
     "on_deployment",
     "depends",
     "IODescriptor",